From e30297f87eee787ed9fcfd23c2ff960a668e6de9 Mon Sep 17 00:00:00 2001 From: Vincent Koc Date: Fri, 31 Jul 2026 08:56:55 +0800 Subject: [PATCH] fix(agents): preserve Code Mode restart recovery --- packages/gateway-protocol/src/schema/agent.ts | 1 + src/agents/code-mode.ts | 3 +- src/agents/command/attempt-execution.ts | 1 + src/agents/command/types.ts | 1 + .../run/attempt-tool-base-prepare.ts | 1 + .../run/attempt-tool-catalog.ts | 1 + .../embedded-agent-runner/run/params.ts | 2 + .../run/run-attempt-dispatch.ts | 1 + src/agents/main-session-restart-dispatch.ts | 2 + ...-session-restart-recovery-resume-policy.ts | 30 +++++++--- .../main-session-restart-recovery-store.ts | 1 + .../main-session-restart-recovery.test.ts | 60 ++++++++++++++++++- src/agents/tool-surface-plan.test.ts | 42 +++++++++++++ src/agents/tool-surface-plan.ts | 26 +++++--- .../agent-request-preflight.test.ts | 52 ++++++++++++++++ .../server-methods/agent-request-preflight.ts | 16 ++++- .../server-methods/agent-request-types.ts | 1 + .../agent-run-execution-phase.ts | 1 + 18 files changed, 222 insertions(+), 20 deletions(-) diff --git a/packages/gateway-protocol/src/schema/agent.ts b/packages/gateway-protocol/src/schema/agent.ts index f9ca6aeba2ae..5c40074860a6 100644 --- a/packages/gateway-protocol/src/schema/agent.ts +++ b/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), diff --git a/src/agents/code-mode.ts b/src/agents/code-mode.ts index dda368c0d6e5..70335e89bea6 100644 --- a/src/agents/code-mode.ts +++ b/src/agents/code-mode.ts @@ -266,11 +266,12 @@ export function applyCodeModeCatalog(params: { toolHookContext?: HookContext; directToolNames?: Iterable; 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, diff --git a/src/agents/command/attempt-execution.ts b/src/agents/command/attempt-execution.ts index c3d2da36ba4c..c6fde499f25a 100644 --- a/src/agents/command/attempt-execution.ts +++ b/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, diff --git a/src/agents/command/types.ts b/src/agents/command/types.ts index cc672db14472..0055710af2e4 100644 --- a/src/agents/command/types.ts +++ b/src/agents/command/types.ts @@ -156,6 +156,7 @@ export type AgentCommandOpts = { swarmOutputSchema?: Record; /** 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; diff --git a/src/agents/embedded-agent-runner/run/attempt-tool-base-prepare.ts b/src/agents/embedded-agent-runner/run/attempt-tool-base-prepare.ts index e5b359518fc4..026aa011113b 100644 --- a/src/agents/embedded-agent-runner/run/attempt-tool-base-prepare.ts +++ b/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 diff --git a/src/agents/embedded-agent-runner/run/attempt-tool-catalog.ts b/src/agents/embedded-agent-runner/run/attempt-tool-catalog.ts index cd886c36411e..04292e967b87 100644 --- a/src/agents/embedded-agent-runner/run/attempt-tool-catalog.ts +++ b/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, diff --git a/src/agents/embedded-agent-runner/run/params.ts b/src/agents/embedded-agent-runner/run/params.ts index 421ec8feaeda..e3dc5fea4c32 100644 --- a/src/agents/embedded-agent-runner/run/params.ts +++ b/src/agents/embedded-agent-runner/run/params.ts @@ -150,6 +150,8 @@ export type RunEmbeddedAgentParams = { swarmOutputSchema?: Record; /** 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. */ diff --git a/src/agents/embedded-agent-runner/run/run-attempt-dispatch.ts b/src/agents/embedded-agent-runner/run/run-attempt-dispatch.ts index 0ffa51e79796..d5bced61a11c 100644 --- a/src/agents/embedded-agent-runner/run/run-attempt-dispatch.ts +++ b/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, diff --git a/src/agents/main-session-restart-dispatch.ts b/src/agents/main-session-restart-dispatch.ts index a17c8ac7533b..35569ca45bb5 100644 --- a/src/agents/main-session-restart-dispatch.ts +++ b/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, diff --git a/src/agents/main-session-restart-recovery-resume-policy.ts b/src/agents/main-session-restart-recovery-resume-policy.ts index 2694d76dcef2..fa21d5f1f588 100644 --- a/src/agents/main-session-restart-recovery-resume-policy.ts +++ b/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 } : {}), }; } diff --git a/src/agents/main-session-restart-recovery-store.ts b/src/agents/main-session-restart-recovery-store.ts index aa512d8c53c7..45ab917e7746 100644 --- a/src/agents/main-session-restart-recovery-store.ts +++ b/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, }); diff --git a/src/agents/main-session-restart-recovery.test.ts b/src/agents/main-session-restart-recovery.test.ts index f32da632be73..7ff4efd9f911 100644 --- a/src/agents/main-session-restart-recovery.test.ts +++ b/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([ diff --git a/src/agents/tool-surface-plan.test.ts b/src/agents/tool-surface-plan.test.ts index ee0e234582c9..51bde1a54b16 100644 --- a/src/agents/tool-surface-plan.test.ts +++ b/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" } }, diff --git a/src/agents/tool-surface-plan.ts b/src/agents/tool-surface-plan.ts index 5669b3c6cd3b..0bb252b9118d 100644 --- a/src/agents/tool-surface-plan.ts +++ b/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