refactor: split embedded run setup helpers

2026-03-29 02:41:07 +00:00 · 2026-03-27 00:44:56 +00:00
parent d4da878d64
commit aed6283faa
2 changed files with 159 additions and 81 deletions
--- a/src/agents/pi-embedded-runner/run.ts
+++ b/src/agents/pi-embedded-runner/run.ts
@@ -8,7 +8,6 @@ import {
 import { computeBackoff, sleepWithAbort } from "../../infra/backoff.js";
 import { getGlobalHookRunner } from "../../plugins/hook-runner-global.js";
 import { prepareProviderRuntimeAuth } from "../../plugins/provider-runtime.js";
-import type { PluginHookBeforeAgentStartResult } from "../../plugins/types.js";
 import { enqueueCommandInLane } from "../../process/command-queue.js";
 import { isMarkdownCapableMessageChannel } from "../../utils/message-channel.js";
 import { resolveOpenClawAgentDir } from "../agent-paths.js";
@@ -21,13 +20,7 @@ import {
  markAuthProfileUsed,
  resolveProfilesUnavailableReason,
 } from "../auth-profiles.js";
-import {
-  CONTEXT_WINDOW_HARD_MIN_TOKENS,
-  CONTEXT_WINDOW_WARN_BELOW_TOKENS,
-  evaluateContextWindowGuard,
-  resolveContextWindowInfo,
-} from "../context-window-guard.js";
-import { DEFAULT_CONTEXT_TOKENS, DEFAULT_MODEL, DEFAULT_PROVIDER } from "../defaults.js";
+import { DEFAULT_MODEL, DEFAULT_PROVIDER } from "../defaults.js";
 import {
  coerceToFailoverError,
  describeFailoverError,
@@ -90,6 +83,7 @@ import {
 } from "./run/helpers.js";
 import type { RunEmbeddedPiAgentParams } from "./run/params.js";
 import { buildEmbeddedRunPayloads } from "./run/payloads.js";
+import { resolveEffectiveRuntimeModel, resolveHookModelSelection } from "./run/setup.js";
 import {
  sessionLikelyHasOversizedToolResults,
  truncateOversizedToolResultsInSession,
@@ -152,14 +146,6 @@ export async function runEmbeddedPiAgent(
        sessionKey: params.sessionKey,
      });
      await ensureOpenClawModelsJson(params.config, agentDir);
-
-      // Run before_model_resolve hooks early so plugins can override the
-      // provider/model before resolveModel().
-      //
-      // Legacy compatibility: before_agent_start is also checked for override
-      // fields if present. New hook takes precedence when both are set.
-      let modelResolveOverride: { providerOverride?: string; modelOverride?: string } | undefined;
-      let legacyBeforeAgentStartResult: PluginHookBeforeAgentStartResult | undefined;
      const hookRunner = getGlobalHookRunner();
      const hookCtx = {
        agentId: workspaceResolution.agentId,
@@ -170,43 +156,17 @@ export async function runEmbeddedPiAgent(
        trigger: params.trigger,
        channelId: params.messageChannel ?? params.messageProvider ?? undefined,
      };
-      if (hookRunner?.hasHooks("before_model_resolve")) {
-        try {
-          modelResolveOverride = await hookRunner.runBeforeModelResolve(
-            { prompt: params.prompt },
-            hookCtx,
-          );
-        } catch (hookErr) {
-          log.warn(`before_model_resolve hook failed: ${String(hookErr)}`);
-        }
-      }
-      if (hookRunner?.hasHooks("before_agent_start")) {
-        try {
-          legacyBeforeAgentStartResult = await hookRunner.runBeforeAgentStart(
-            { prompt: params.prompt },
-            hookCtx,
-          );
-          modelResolveOverride = {
-            providerOverride:
-              modelResolveOverride?.providerOverride ??
-              legacyBeforeAgentStartResult?.providerOverride,
-            modelOverride:
-              modelResolveOverride?.modelOverride ?? legacyBeforeAgentStartResult?.modelOverride,
-          };
-        } catch (hookErr) {
-          log.warn(
-            `before_agent_start hook (legacy model resolve path) failed: ${String(hookErr)}`,
-          );
-        }
-      }
-      if (modelResolveOverride?.providerOverride) {
-        provider = modelResolveOverride.providerOverride;
-        log.info(`[hooks] provider overridden to ${provider}`);
-      }
-      if (modelResolveOverride?.modelOverride) {
-        modelId = modelResolveOverride.modelOverride;
-        log.info(`[hooks] model overridden to ${modelId}`);
-      }
+
+      const hookSelection = await resolveHookModelSelection({
+        prompt: params.prompt,
+        provider,
+        modelId,
+        hookRunner,
+        hookContext: hookCtx,
+      });
+      provider = hookSelection.provider;
+      modelId = hookSelection.modelId;
+      const legacyBeforeAgentStartResult = hookSelection.legacyBeforeAgentStartResult;

      const { model, error, authStorage, modelRegistry } = await resolveModelAsync(
        provider,
@@ -223,38 +183,14 @@ export async function runEmbeddedPiAgent(
      }
      let runtimeModel = model;

-      const ctxInfo = resolveContextWindowInfo({
+      const resolvedRuntimeModel = resolveEffectiveRuntimeModel({
        cfg: params.config,
        provider,
        modelId,
-        modelContextWindow: runtimeModel.contextWindow,
-        defaultTokens: DEFAULT_CONTEXT_TOKENS,
+        runtimeModel,
      });
-      // Apply contextTokens cap to model so pi-coding-agent's auto-compaction
-      // threshold uses the effective limit, not the native context window.
-      let effectiveModel =
-        ctxInfo.tokens < (runtimeModel.contextWindow ?? Infinity)
-          ? { ...runtimeModel, contextWindow: ctxInfo.tokens }
-          : runtimeModel;
-      const ctxGuard = evaluateContextWindowGuard({
-        info: ctxInfo,
-        warnBelowTokens: CONTEXT_WINDOW_WARN_BELOW_TOKENS,
-        hardMinTokens: CONTEXT_WINDOW_HARD_MIN_TOKENS,
-      });
-      if (ctxGuard.shouldWarn) {
-        log.warn(
-          `low context window: ${provider}/${modelId} ctx=${ctxGuard.tokens} (warn<${CONTEXT_WINDOW_WARN_BELOW_TOKENS}) source=${ctxGuard.source}`,
-        );
-      }
-      if (ctxGuard.shouldBlock) {
-        log.error(
-          `blocked model (context window too small): ${provider}/${modelId} ctx=${ctxGuard.tokens} (min=${CONTEXT_WINDOW_HARD_MIN_TOKENS}) source=${ctxGuard.source}`,
-        );
-        throw new FailoverError(
-          `Model context window too small (${ctxGuard.tokens} tokens). Minimum is ${CONTEXT_WINDOW_HARD_MIN_TOKENS}.`,
-          { reason: "unknown", provider, model: modelId },
-        );
-      }
+      const ctxInfo = resolvedRuntimeModel.ctxInfo;
+      let effectiveModel = resolvedRuntimeModel.effectiveModel;

      const authStore = ensureAuthProfileStore(agentDir, {
        allowKeychainPrompt: false,
--- a/src/agents/pi-embedded-runner/run/setup.ts
+++ b/src/agents/pi-embedded-runner/run/setup.ts
@@ -0,0 +1,142 @@
+import type { Api, Model } from "@mariozechner/pi-ai";
+import type { OpenClawConfig } from "../../../config/config.js";
+import type { PluginHookBeforeAgentStartResult } from "../../../plugins/types.js";
+import {
+  CONTEXT_WINDOW_HARD_MIN_TOKENS,
+  CONTEXT_WINDOW_WARN_BELOW_TOKENS,
+  evaluateContextWindowGuard,
+  resolveContextWindowInfo,
+} from "../../context-window-guard.js";
+import { DEFAULT_CONTEXT_TOKENS } from "../../defaults.js";
+import { FailoverError } from "../../failover-error.js";
+import { log } from "../logger.js";
+
+type HookContext = {
+  agentId?: string;
+  sessionKey?: string;
+  sessionId: string;
+  workspaceDir: string;
+  messageProvider?: string;
+  trigger?: string;
+  channelId?: string;
+};
+
+type HookRunnerLike = {
+  hasHooks(hookName: string): boolean;
+  runBeforeModelResolve(
+    input: { prompt: string },
+    context: HookContext,
+  ): Promise<{ providerOverride?: string; modelOverride?: string } | undefined>;
+  runBeforeAgentStart(
+    input: { prompt: string },
+    context: HookContext,
+  ): Promise<PluginHookBeforeAgentStartResult | undefined>;
+};
+
+export async function resolveHookModelSelection(params: {
+  prompt: string;
+  provider: string;
+  modelId: string;
+  hookRunner?: HookRunnerLike | null;
+  hookContext: HookContext;
+}) {
+  let provider = params.provider;
+  let modelId = params.modelId;
+  let modelResolveOverride: { providerOverride?: string; modelOverride?: string } | undefined;
+  let legacyBeforeAgentStartResult: PluginHookBeforeAgentStartResult | undefined;
+  const hookRunner = params.hookRunner;
+
+  // Run before_model_resolve hooks early so plugins can override the
+  // provider/model before resolveModel().
+  //
+  // Legacy compatibility: before_agent_start is also checked for override
+  // fields if present. New hook takes precedence when both are set.
+  if (hookRunner?.hasHooks("before_model_resolve")) {
+    try {
+      modelResolveOverride = await hookRunner.runBeforeModelResolve(
+        { prompt: params.prompt },
+        params.hookContext,
+      );
+    } catch (hookErr) {
+      log.warn(`before_model_resolve hook failed: ${String(hookErr)}`);
+    }
+  }
+
+  if (hookRunner?.hasHooks("before_agent_start")) {
+    try {
+      legacyBeforeAgentStartResult = await hookRunner.runBeforeAgentStart(
+        { prompt: params.prompt },
+        params.hookContext,
+      );
+      modelResolveOverride = {
+        providerOverride:
+          modelResolveOverride?.providerOverride ?? legacyBeforeAgentStartResult?.providerOverride,
+        modelOverride:
+          modelResolveOverride?.modelOverride ?? legacyBeforeAgentStartResult?.modelOverride,
+      };
+    } catch (hookErr) {
+      log.warn(`before_agent_start hook (legacy model resolve path) failed: ${String(hookErr)}`);
+    }
+  }
+
+  if (modelResolveOverride?.providerOverride) {
+    provider = modelResolveOverride.providerOverride;
+    log.info(`[hooks] provider overridden to ${provider}`);
+  }
+  if (modelResolveOverride?.modelOverride) {
+    modelId = modelResolveOverride.modelOverride;
+    log.info(`[hooks] model overridden to ${modelId}`);
+  }
+
+  return {
+    provider,
+    modelId,
+    legacyBeforeAgentStartResult,
+  };
+}
+
+export function resolveEffectiveRuntimeModel(params: {
+  cfg: OpenClawConfig | undefined;
+  provider: string;
+  modelId: string;
+  runtimeModel: Model<Api>;
+}) {
+  const ctxInfo = resolveContextWindowInfo({
+    cfg: params.cfg,
+    provider: params.provider,
+    modelId: params.modelId,
+    modelContextWindow: params.runtimeModel.contextWindow,
+    defaultTokens: DEFAULT_CONTEXT_TOKENS,
+  });
+
+  // Apply contextTokens cap to model so pi-coding-agent's auto-compaction
+  // threshold uses the effective limit, not the native context window.
+  const effectiveModel =
+    ctxInfo.tokens < (params.runtimeModel.contextWindow ?? Infinity)
+      ? { ...params.runtimeModel, contextWindow: ctxInfo.tokens }
+      : params.runtimeModel;
+  const ctxGuard = evaluateContextWindowGuard({
+    info: ctxInfo,
+    warnBelowTokens: CONTEXT_WINDOW_WARN_BELOW_TOKENS,
+    hardMinTokens: CONTEXT_WINDOW_HARD_MIN_TOKENS,
+  });
+  if (ctxGuard.shouldWarn) {
+    log.warn(
+      `low context window: ${params.provider}/${params.modelId} ctx=${ctxGuard.tokens} (warn<${CONTEXT_WINDOW_WARN_BELOW_TOKENS}) source=${ctxGuard.source}`,
+    );
+  }
+  if (ctxGuard.shouldBlock) {
+    log.error(
+      `blocked model (context window too small): ${params.provider}/${params.modelId} ctx=${ctxGuard.tokens} (min=${CONTEXT_WINDOW_HARD_MIN_TOKENS}) source=${ctxGuard.source}`,
+    );
+    throw new FailoverError(
+      `Model context window too small (${ctxGuard.tokens} tokens). Minimum is ${CONTEXT_WINDOW_HARD_MIN_TOKENS}.`,
+      { reason: "unknown", provider: params.provider, model: params.modelId },
+    );
+  }
+
+  return {
+    ctxInfo,
+    effectiveModel,
+  };
+}