fix(agents): preserve Code Mode restart recovery

packages/gateway-protocol/src/schema/agent.ts

@@ -335,6 +335,7 @@ export const AgentParamsSchema = closedObject({
   // Host-owned recovery turns can force every Code Mode exec onto the
   // restart-safe path even if the model omits or clears the tool argument.
   forceRestartSafeTools: Type.Optional(Type.Boolean()),
+  forceCodeModeTools: Type.Optional(Type.Boolean()),
   voiceWakeTrigger: Type.Optional(Type.String()),
   idempotencyKey: NonEmptyString,
   label: Type.Optional(SessionLabelString),

src/agents/code-mode.ts

@@ -266,11 +266,12 @@ export function applyCodeModeCatalog(params: {
   toolHookContext?: HookContext;
   directToolNames?: Iterable<string>;
   codeModeSkills?: CodeModeToolContext["codeModeSkills"];
+  forceEnabled?: boolean;
 }) {
   const config = resolveCodeModeConfig(params.config, params.agentId);
   // Engagement (including "auto" per-model resolution) is decided by the run
   // gates before this is called; only a hard `false` may disable compaction.
-  if (config.enabled === false) {
+  if (config.enabled === false && params.forceEnabled !== true) {
     return applyToolCatalogCompaction({
       ...params,
       enabled: false,

src/agents/command/attempt-execution.ts

@@ -1040,6 +1040,7 @@ export function runAgentAttempt(params: {
     swarmCollector: params.opts.swarmCollector,
     swarmOutputSchema: params.opts.swarmOutputSchema,
     forceRestartSafeTools: params.opts.forceRestartSafeTools,
+    forceCodeModeTools: params.opts.forceCodeModeTools,
     streamParams: params.opts.streamParams,
     agentDir: params.agentDir,
     allowGatewaySubagentBinding: params.opts.allowGatewaySubagentBinding,

src/agents/command/types.ts

@@ -156,6 +156,7 @@ export type AgentCommandOpts = {
   swarmOutputSchema?: Record<string, unknown>;
   /** Restrict this reconstructed run to restart-safe tools. */
   forceRestartSafeTools?: boolean;
+  forceCodeModeTools?: boolean;
   /** Host-owned exact media set for a scoped automatic recovery delivery. */
   internalDeliveryMediaUrls?: string[];
   internalDeliverySuppressText?: boolean;

src/agents/embedded-agent-runner/run/attempt-tool-base-prepare.ts

@@ -90,6 +90,7 @@ export function prepareEmbeddedAttemptToolBase(params: {
     isRawModelRun,
     skillWorkshopProposalOnly: attempt.skillWorkshopProposalOnly,
     toolsAllow: attempt.toolsAllow,
+    forceCodeModeControls: attempt.forceCodeModeTools,
   });
   const effectiveToolsAllow =
     toolSearchControlsEnabledForRun && toolsAllowWithForcedRuntimeTools

src/agents/embedded-agent-runner/run/attempt-tool-catalog.ts

@@ -100,6 +100,7 @@ export function prepareEmbeddedAttemptToolCatalog(input: {
     codeModeControlsEnabled: codeModeControlsEnabledForRun,
     toolSearchConfig,
     forceDirectMessageTool: preparedToolBase.forceDirectMessageTool,
+    forceCodeModeControls: attempt.forceCodeModeTools,
     sessionId: attempt.sessionId,
     sessionKey: input.sandboxSessionKey,
     agentId: input.sessionAgentId,

src/agents/embedded-agent-runner/run/params.ts

@@ -150,6 +150,8 @@ export type RunEmbeddedAgentParams = {
   swarmOutputSchema?: Record<string, unknown>;
   /** Restrict this reconstructed run to restart-safe tools. */
   forceRestartSafeTools?: boolean;
+  /** Preserve Code Mode controls for a replay-safe restart recovery turn. */
+  forceCodeModeTools?: boolean;
   /** Internal one-shot model probe mode: no tools, no workspace/chat prompt policy. */
   modelRun?: boolean;
   /** Disable trajectory persistence for auxiliary runs with no durable session owner. */

src/agents/embedded-agent-runner/run/run-attempt-dispatch.ts

@@ -375,6 +375,7 @@ export async function dispatchEmbeddedRunAttempt(input: {
     swarmCollector: params.swarmCollector,
     swarmOutputSchema: params.swarmOutputSchema,
     forceRestartSafeTools: params.forceRestartSafeTools,
+    forceCodeModeTools: params.forceCodeModeTools,
     forceMessageTool: params.forceMessageTool,
     enableHeartbeatTool: params.enableHeartbeatTool,
     forceHeartbeatTool: params.forceHeartbeatTool,

src/agents/main-session-restart-dispatch.ts

@@ -353,6 +353,7 @@ export async function resumeMainSession(params: {
   sessionKey: string;
   pendingFinalDeliveryText?: string | null;
   forceRestartSafeTools?: boolean;
+  forceCodeModeTools?: boolean;
   sessionWorkAdmissionHandoffId?: string;
   lifecycleGeneration?: string;
   shouldContinue?: () => boolean;
@@ -482,6 +483,7 @@ export async function resumeMainSession(params: {
         ? { sourceReplyDeliveryMode: params.entry.restartRecoverySourceReplyDeliveryMode }
         : {}),
       ...(params.forceRestartSafeTools ? { forceRestartSafeTools: true } : {}),
+      ...(params.forceCodeModeTools ? { forceCodeModeTools: true } : {}),
       inputProvenance: {
         kind: "internal_system",
         sourceSessionKey: dispatchSessionKey,

src/agents/main-session-restart-recovery-resume-policy.ts

@@ -244,6 +244,7 @@ const LEGACY_RESTART_ABORT_ERROR_MESSAGES = new Set([
   "This operation was aborted",
   AGENT_RUN_RESTART_ABORT_ERROR,
 ]);
+const CODE_MODE_RESTART_ABORT_ERROR = "code mode execution aborted";
 
 function isRestartAbortAssistantMessage(message: unknown): boolean {
   if (!message || typeof message !== "object" || getMessageRole(message) !== "assistant") {
@@ -297,8 +298,7 @@ function isRestartAbortedWaitFailure(message: unknown): boolean {
   if (
     !details ||
     typeof details !== "object" ||
-    (details as { status?: unknown }).status !== "failed" ||
-    (details as { code?: unknown }).code !== "internal_error"
+    (details as { status?: unknown }).status !== "failed"
   ) {
     return false;
   }
@@ -315,6 +315,16 @@ function isRestartAbortedWaitFailure(message: unknown): boolean {
   const errorText =
     normalizeOptionalString((details as { error?: unknown }).error) ??
     normalizeOptionalString(contentText);
+  const code = normalizeOptionalString((details as { code?: unknown }).code);
+  if (code === "aborted") {
+    // Current Code Mode wait aborts use the runtime's explicit abort code and
+    // message. Recovery already owns the restart boundary, so this exact pair
+    // is lifecycle noise rather than a provider or tool failure.
+    return errorText === CODE_MODE_RESTART_ABORT_ERROR;
+  }
+  if (code !== "internal_error") {
+    return false;
+  }
   return /^(?:(?:Abort)?Error:\s*)?(?:The|This) operation was aborted\.?$/u.test(errorText ?? "");
 }
 
@@ -346,7 +356,11 @@ type MainSessionResumePolicy =
     }
   | { action: "complete"; reason: "handled-silent" }
   | { action: "fail"; reason: string }
-  | { action: "resume"; forceRestartSafeTools: boolean };
+  | {
+      action: "resume";
+      forceRestartSafeTools: boolean;
+      forceCodeModeTools?: true;
+    };
 
 export function resolveMainSessionResumePolicy(
   messages: unknown[],
@@ -417,7 +431,7 @@ export function resolveMainSessionResumePolicy(
     const waitCall = readCodeModeWaitCall(meaningfulMessages[1]);
     const checkpoint = readCodeModeCheckpoint(meaningfulMessages[2]);
     return waitCall && checkpoint?.replaySafe === true && checkpoint.runId === waitCall.runId
-      ? { action: "resume", forceRestartSafeTools: true }
+      ? { action: "resume", forceRestartSafeTools: true, forceCodeModeTools: true }
       : {
           action: "fail",
           reason: "failed Code Mode wait cannot be matched to a replay-safe checkpoint",
@@ -427,13 +441,13 @@ export function resolveMainSessionResumePolicy(
   if (waitCall) {
     const checkpoint = readCodeModeCheckpoint(meaningfulMessages[1]);
     return checkpoint?.replaySafe === true && checkpoint.runId === waitCall.runId
-      ? { action: "resume", forceRestartSafeTools: true }
+      ? { action: "resume", forceRestartSafeTools: true, forceCodeModeTools: true }
       : { action: "fail", reason: "Code Mode wait checkpoint is not replay-safe" };
   }
   const tailCheckpoint = readCodeModeCheckpoint(lastMeaningful);
   if (tailCheckpoint) {
     return tailCheckpoint.replaySafe
-      ? { action: "resume", forceRestartSafeTools: true }
+      ? { action: "resume", forceRestartSafeTools: true, forceCodeModeTools: true }
       : { action: "fail", reason: "Code Mode wait checkpoint is not replay-safe" };
   }
   // A tool call interrupted mid-execution resumes like the manual re-send the
@@ -455,8 +469,10 @@ export function resolveMainSessionResumePolicy(
   }
   // A later tool result can hide the checkpoint at the transcript tail; keep
   // the interrupted turn restricted without borrowing an earlier turn's state.
+  const forceCodeModeTools = hasReplaySafeCodeModeCheckpointInCurrentTurn(messages);
   return {
     action: "resume",
-    forceRestartSafeTools: hasReplaySafeCodeModeCheckpointInCurrentTurn(messages),
+    forceRestartSafeTools: forceCodeModeTools,
+    ...(forceCodeModeTools ? { forceCodeModeTools: true } : {}),
   };
 }

src/agents/main-session-restart-recovery-store.ts

@@ -558,6 +558,7 @@ export async function recoverStore(params: {
       sessionKey,
       forceRestartSafeTools:
         entry.restartRecoveryForceSafeTools === true || resumePolicy.forceRestartSafeTools,
+      forceCodeModeTools: resumePolicy.forceCodeModeTools === true,
       sessionWorkAdmissionHandoffId: params.sessionWorkAdmissionHandoffId,
       gatewayRuntime: params.gatewayRuntime,
     });

src/agents/main-session-restart-recovery.test.ts

@@ -4752,7 +4752,10 @@ describe("main-session-restart-recovery", () => {
     ]);
 
     await expectRecovery({ recovered: 1, failed: 0, skipped: 0 });
-    expect(gatewayParams()).toMatchObject({ forceRestartSafeTools: true });
+    expect(gatewayParams()).toMatchObject({
+      forceRestartSafeTools: true,
+      forceCodeModeTools: true,
+    });
   });
 
   it.each([
@@ -5178,7 +5181,55 @@ describe("main-session-restart-recovery", () => {
     ]);
 
     await expectRecovery({ recovered: 1, failed: 0, skipped: 0 });
-    expect(gatewayParams()).toMatchObject({ forceRestartSafeTools: true });
+    expect(gatewayParams()).toMatchObject({
+      forceRestartSafeTools: true,
+      forceCodeModeTools: true,
+    });
+  });
+
+  it("resumes through the current Code Mode abort persisted for an interrupted wait", async () => {
+    const sessionsDir = await makeSessionsDir();
+    await writeStore(sessionsDir, mainSessionStore());
+    await writeTranscript(sessionsDir, "main-session", [
+      { role: "user", content: "do the thing" },
+      codeModeCheckpointMessage(),
+      codeModeWaitCallMessage(),
+      {
+        role: "toolResult",
+        toolName: "wait",
+        toolCallId: "call-wait-1",
+        content: [
+          {
+            type: "text",
+            text: JSON.stringify({
+              status: "failed",
+              code: "aborted",
+              error: "code mode execution aborted",
+            }),
+          },
+        ],
+        details: {
+          status: "failed",
+          code: "aborted",
+          error: "code mode execution aborted",
+          replaySafe: true,
+        },
+        isError: true,
+      },
+      {
+        role: "assistant",
+        content: [],
+        stopReason: "aborted",
+        errorCode: "OPENCLAW_RESTART_ABORT",
+        errorMessage: "agent run aborted for restart",
+      },
+    ]);
+
+    await expectRecovery({ recovered: 1, failed: 0, skipped: 0 });
+    expect(gatewayParams()).toMatchObject({
+      forceRestartSafeTools: true,
+      forceCodeModeTools: true,
+    });
   });
 
   it("keeps an unmatched failed wait restricted when its checkpoint is replay-safe", async () => {
@@ -5203,7 +5254,10 @@ describe("main-session-restart-recovery", () => {
     ]);
 
     await expectRecovery({ recovered: 1, failed: 0, skipped: 0 });
-    expect(gatewayParams()).toMatchObject({ forceRestartSafeTools: true });
+    expect(gatewayParams()).toMatchObject({
+      forceRestartSafeTools: true,
+      forceCodeModeTools: true,
+    });
   });
 
   it.each([

src/agents/tool-surface-plan.test.ts

@@ -71,6 +71,20 @@ describe("resolveAgentToolSurfacePlan", () => {
     expect(plan.toolSearchControlsEnabled).toBe(expected.toolSearch);
     expect(plan.codeModeControlsEnabled && plan.toolSearchControlsEnabled).toBe(false);
   });
+
+  it("preserves Code Mode controls for a checkpoint-proven restart recovery", () => {
+    const config: OpenClawConfig = {
+      tools: { codeMode: false, toolSearch: true },
+    };
+    const plan = resolveAgentToolSurfacePlan({
+      ...basePlanParams,
+      config,
+      forceCodeModeControls: true,
+    });
+
+    expect(plan.codeModeControlsEnabled).toBe(true);
+    expect(plan.toolSearchControlsEnabled).toBe(false);
+  });
 });
 
 describe("applyAgentToolSurfaceCatalog", () => {
@@ -99,6 +113,34 @@ describe("applyAgentToolSurfaceCatalog", () => {
     expect(result.catalogToolCount).toBe(1);
   });
 
+  it("forces the Code Mode catalog for a checkpoint-proven restart recovery", () => {
+    const config: OpenClawConfig = {
+      tools: { codeMode: false, toolSearch: { enabled: true, mode: "directory" } },
+    };
+    const plan = resolveAgentToolSurfacePlan({
+      ...basePlanParams,
+      config,
+      forceCodeModeControls: true,
+    });
+    const catalogRef = createToolSearchCatalogRef();
+    const result = applyAgentToolSurfaceCatalog({
+      tools: [
+        ...createCodeModeTools({ config, catalogRef, executeTool }),
+        createStubTool("hidden_target"),
+      ],
+      config,
+      toolSearchRuntimeConfig: plan.toolSearchRuntimeConfig,
+      codeModeControlsEnabled: plan.codeModeControlsEnabled,
+      toolSearchConfig: plan.toolSearchConfig,
+      forceDirectMessageTool: false,
+      forceCodeModeControls: true,
+      catalogRef,
+    });
+
+    expect(result.tools.map((tool) => tool.name)).toEqual(["exec", "wait"]);
+    expect(result.catalogToolCount).toBe(1);
+  });
+
   it("uses the schema-directory catalog in directory mode", () => {
     const config: OpenClawConfig = {
       tools: { codeMode: false, toolSearch: { enabled: true, mode: "directory" } },

src/agents/tool-surface-plan.ts

@@ -24,6 +24,7 @@ type AgentToolSurfacePlanParams = {
   isRawModelRun: boolean;
   skillWorkshopProposalOnly?: boolean;
   toolsAllow?: readonly string[];
+  forceCodeModeControls?: boolean;
 };
 
 export function resolveAgentToolSurfacePlan(params: AgentToolSurfacePlanParams) {
@@ -45,7 +46,11 @@ export function resolveAgentToolSurfacePlan(params: AgentToolSurfacePlanParams)
     params.skillWorkshopProposalOnly !== true &&
     params.toolsAllow?.length !== 0;
   const codeModeControlsEnabled =
-    toolsAvailable && isCodeModeEngagedForModel(codeModeConfig, params.model);
+    toolsAvailable &&
+    // Restart recovery continues one provider turn. Keep its original control
+    // schema even when the reloaded config disables Code Mode for new turns.
+    (params.forceCodeModeControls === true ||
+      isCodeModeEngagedForModel(codeModeConfig, params.model));
   const toolSearchControlsEnabled =
     toolsAvailable && !codeModeControlsEnabled && toolSearchConfig.enabled;
   return {
@@ -65,6 +70,7 @@ type ApplyAgentToolSurfaceCatalogParams = Omit<CodeModeCatalogParams, "directToo
   codeModeControlsEnabled: boolean;
   toolSearchConfig: ToolSearchConfig;
   forceDirectMessageTool: boolean;
+  forceCodeModeControls?: boolean;
 };
 
 export function applyAgentToolSurfaceCatalog({
@@ -72,21 +78,27 @@ export function applyAgentToolSurfaceCatalog({
   toolSearchConfig,
   toolSearchRuntimeConfig,
   forceDirectMessageTool,
+  forceCodeModeControls,
   ...catalogParams
 }: ApplyAgentToolSurfaceCatalogParams) {
   // When the message tool is the only reply path it must stay directly visible
   // in every search mode; a hidden delivery tool can leave the run mute.
   const directToolNames = forceDirectMessageTool ? ["message"] : [];
-  const applyCatalog = codeModeControlsEnabled
-    ? applyCodeModeCatalog
-    : toolSearchConfig.mode === "directory"
+  if (codeModeControlsEnabled) {
+    return applyCodeModeCatalog({
+      ...catalogParams,
+      config: catalogParams.config,
+      directToolNames,
+      forceEnabled: forceCodeModeControls,
+    });
+  }
+  const applyCatalog =
+    toolSearchConfig.mode === "directory"
       ? applyToolSchemaDirectoryCatalog
       : applyToolSearchCatalog;
   return applyCatalog({
     ...catalogParams,
-    // Code mode reads the base config; tool-search modes read the run's
-    // resolved tool-search runtime config.
-    config: codeModeControlsEnabled ? catalogParams.config : toolSearchRuntimeConfig,
+    config: toolSearchRuntimeConfig,
     directToolNames,
   });
 }

src/gateway/server-methods/agent-request-preflight.test.ts

@@ -350,3 +350,55 @@ describe("agent request Swarm preflight", () => {
     expect(respond).not.toHaveBeenCalled();
   });
 });
+
+describe("agent request restart recovery preflight", () => {
+  function runRestartRecoveryPreflight(backend: boolean, sourceTool: string) {
+    const respond = vi.fn();
+    const result = prepareAgentRequestPreflight({
+      params: {
+        message: "continue",
+        idempotencyKey: "restart-recovery-run",
+        forceRestartSafeTools: true,
+        forceCodeModeTools: true,
+        inputProvenance: {
+          kind: "internal_system",
+          sourceSessionKey: "agent:main:main",
+          sourceTool,
+        },
+      },
+      respond,
+      context: {
+        getRuntimeConfig: () => ({}),
+        dedupe: new Map(),
+      },
+      client: backend
+        ? { connect: { client: { mode: "backend" }, scopes: ["operator.write"] } }
+        : undefined,
+    } as never);
+    return { respond, result };
+  }
+
+  it("accepts the Code Mode override only for backend restart recovery", () => {
+    const accepted = runRestartRecoveryPreflight(true, "main_session_restart_recovery");
+
+    expect(accepted.result).toBeDefined();
+    expect(accepted.respond).not.toHaveBeenCalled();
+  });
+
+  it.each([
+    { backend: false, sourceTool: "main_session_restart_recovery" },
+    { backend: true, sourceTool: "other_internal_source" },
+  ])("rejects an untrusted Code Mode override", ({ backend, sourceTool }) => {
+    const rejected = runRestartRecoveryPreflight(backend, sourceTool);
+
+    expect(rejected.result).toBeUndefined();
+    expect(rejected.respond).toHaveBeenCalledWith(
+      false,
+      undefined,
+      expect.objectContaining({
+        code: "INVALID_REQUEST",
+        message: "forceCodeModeTools is reserved for main-session restart recovery.",
+      }),
+    );
+  });
+});

src/gateway/server-methods/agent-request-preflight.ts

@@ -235,6 +235,19 @@ export function prepareAgentRequestPreflight(
     return undefined;
   }
   const inputProvenance = normalizeInputProvenance(request.inputProvenance);
+  const isRestartRecoveryResumeRun =
+    canUseInternalRuntimeHandoff && isMainSessionRestartRecoveryInputProvenance(inputProvenance);
+  if (request.forceCodeModeTools === true && !isRestartRecoveryResumeRun) {
+    params.respond(
+      false,
+      undefined,
+      errorShape(
+        ErrorCodes.INVALID_REQUEST,
+        "forceCodeModeTools is reserved for main-session restart recovery.",
+      ),
+    );
+    return undefined;
+  }
   const sessionEffects =
     isOneShotModelRun || requestedInternalSessionEffects ? "internal" : request.sessionEffects;
   const agentDedupeKeys = resolveAgentDedupeKeys({
@@ -293,8 +306,7 @@ export function prepareAgentRequestPreflight(
       groupSpace: request.groupSpace,
     }),
     inputProvenance,
-    isRestartRecoveryResumeRun:
-      canUseInternalRuntimeHandoff && isMainSessionRestartRecoveryInputProvenance(inputProvenance),
+    isRestartRecoveryResumeRun,
     preserveUserFacingSessionModelState:
       canUseInternalRuntimeHandoff &&
       shouldPreserveUserFacingSessionStateForInputProvenance(inputProvenance),

src/gateway/server-methods/agent-request-types.ts

@@ -47,6 +47,7 @@ export type AgentRunRequest = {
   swarmCollector?: boolean;
   swarmOutputSchema?: Record<string, unknown>;
   forceRestartSafeTools?: boolean;
+  forceCodeModeTools?: boolean;
   timeout?: number;
   bestEffortDeliver?: boolean;
   cleanupBundleMcpOnRunEnd?: boolean;

src/gateway/server-methods/agent-run-execution-phase.ts

@@ -376,6 +376,7 @@ export function startAgentRunExecution(params: {
           swarmCollector: params.request.swarmCollector,
           swarmOutputSchema: params.request.swarmOutputSchema,
           forceRestartSafeTools: params.request.forceRestartSafeTools,
+          forceCodeModeTools: params.request.forceCodeModeTools,
           internalDeliveryMediaUrls: params.client?.internal?.internalDeliveryMediaUrls,
           internalDeliverySuppressText: params.client?.internal?.internalDeliverySuppressText,
           suppressPromptPersistence: