From 121ce7faed7eec21480f062f72b5354a7b2e4983 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Sat, 8 Aug 2026 19:11:29 -0700 Subject: [PATCH] fix(worker): preserve long Responses sessions across cloud handoff (#120803) * fix(worker): preserve Responses replay across handoff Long OpenAI Responses cloud-worker sessions retain authoritative compaction replay through windowing and fail visibly when it cannot be preserved. * test(worker): align replay fixture with frame limit --- .../src/schema/worker-admission.ts | 6 +- .../inference-runtime.test.ts | 24 +- .../worker-environments/inference-runtime.ts | 51 ++-- .../inference-terminal-message.ts | 7 +- .../worker-turn-launcher.test.ts | 82 ++++++ .../worker-turn-launcher.ts | 50 +++- .../worker-turn-payload.test.ts | 271 +++++++++++++++--- .../worker-turn-payload.ts | 83 ++++-- .../embedded-agent-transcript.runtime.ts | 119 ++++---- src/worker/embedded-agent.runtime.ts | 30 +- src/worker/replay-message-window.ts | 45 +++ src/worker/transcript-message.test.ts | 47 +-- src/worker/transcript-message.ts | 83 +++--- src/worker/worker-rpc-clients.test.ts | 13 +- src/worker/worker.fault-injection.test.ts | 43 ++- src/worker/worker.runtime.test.ts | 60 +++- .../worker-provider-replay-roundtrip.ts | 170 +++++++++++ 17 files changed, 922 insertions(+), 262 deletions(-) create mode 100644 src/worker/replay-message-window.ts create mode 100644 test/helpers/worker-provider-replay-roundtrip.ts diff --git a/packages/gateway-protocol/src/schema/worker-admission.ts b/packages/gateway-protocol/src/schema/worker-admission.ts index 8965649f1cb9..735b2a327d42 100644 --- a/packages/gateway-protocol/src/schema/worker-admission.ts +++ b/packages/gateway-protocol/src/schema/worker-admission.ts @@ -49,9 +49,9 @@ export const WORKER_PROTOCOL_MAX_FEATURE_LENGTH = 128; export const WORKER_TRANSCRIPT_MAX_BATCH_MESSAGES = 64; export const WORKER_TRANSCRIPT_MAX_CONTENT_PARTS = 128; export const WORKER_TRANSCRIPT_MAX_JSON_DEPTH = 32; -// Reserve at least 12 KiB for the complete transcript commit envelope and -// assistant content. A 60 KiB replay leaves effectively no safe frame headroom. -export const WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES = 48 * 1024; +// Replay is opaque and cannot be truncated. Transcript projection separately +// verifies that the complete commit frame fits the protocol payload ceiling. +export const WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES = WORKER_PROTOCOL_MAX_PAYLOAD_BYTES; const WorkerCredentialSchema = Type.String({ minLength: 16, maxLength: 256 }); const WorkerProtocolFeatureSchema = Type.String({ diff --git a/src/gateway/worker-environments/inference-runtime.test.ts b/src/gateway/worker-environments/inference-runtime.test.ts index 775b2134e481..b5db463b2556 100644 --- a/src/gateway/worker-environments/inference-runtime.test.ts +++ b/src/gateway/worker-environments/inference-runtime.test.ts @@ -20,7 +20,10 @@ import { onTrustedInternalDiagnosticEvent } from "../../infra/diagnostic-events. import { bindModelLlmRuntime } from "../../llm/model-runtime-binding.js"; import type { AssistantMessage, Model, StreamFn, Usage } from "../../llm/types.js"; import { createAssistantMessageEventStream } from "../../llm/utils/event-stream.js"; -import { isWorkerTranscriptMessageFrameSafe } from "../../worker/transcript-message.js"; +import { + isWorkerTranscriptMessageFrameSafe, + WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE, +} from "../../worker/transcript-message.js"; import type { WorkerConnectionIdentity } from "./connection-identity.js"; import { createWorkerInferenceExecutor, @@ -499,7 +502,7 @@ describe("worker inference provider runtime", () => { }); }); - it("keeps inference successful while omitting over-budget replay with a redacted diagnostic", async () => { + it("returns a typed error when authoritative replay cannot be persisted", async () => { const runtime = setup(); const message = finalMessage(); message.providerReplay = { @@ -520,8 +523,12 @@ describe("worker inference provider runtime", () => { const outcome = await runtime.executor(params(request(), vi.fn())).finally(unsubscribe); - expect(outcome).toMatchObject({ type: "done", message: { stopReason: "stop" } }); - expect(outcome.type === "done" ? outcome.message.providerReplay : undefined).toBeUndefined(); + expect(outcome).toMatchObject({ + type: "error", + reason: "provider-error", + message: WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE, + usage: message.usage, + }); expect(payloadEvents).toEqual([ expect.objectContaining({ type: "payload.large", @@ -535,13 +542,14 @@ describe("worker inference provider runtime", () => { expect(JSON.stringify(payloadEvents)).not.toContain(message.providerReplay.data); }); - it("keeps a maximum-budget replay frame-safe through the terminal projection", async () => { + it("keeps a maximum fitting replay exact through the terminal projection", async () => { const runtime = setup(); const message = finalMessage(); + const ciphertext = `cipher-${"x".repeat(60 * 1024)}-€`; message.providerReplay = { v: 1, type: "openai-responses-compaction", - data: "x".repeat(WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES), + data: ciphertext, provider: "openai", api: "openai-responses", model: MODEL, @@ -554,9 +562,7 @@ describe("worker inference provider runtime", () => { if (outcome.type !== "done") { throw new Error("expected successful worker inference"); } - expect(outcome.message.providerReplay?.data).toHaveLength( - WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES, - ); + expect(outcome.message.providerReplay?.data).toBe(ciphertext); expect(isWorkerTranscriptMessageFrameSafe(outcome.message)).toBe(true); }); diff --git a/src/gateway/worker-environments/inference-runtime.ts b/src/gateway/worker-environments/inference-runtime.ts index cddfbf095643..18e21c754f9e 100644 --- a/src/gateway/worker-environments/inference-runtime.ts +++ b/src/gateway/worker-environments/inference-runtime.ts @@ -69,6 +69,7 @@ import type { } from "../../llm/types.js"; import { resolveProviderModelRoutes } from "../../plugins/provider-model-routes.js"; import { estimateUsageCost, resolveModelCostConfig } from "../../utils/usage-format.js"; +import { WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE } from "../../worker/transcript-message.js"; import { projectWorkerInferenceTerminalMessage, type WorkerInferenceModelIdentity, @@ -146,11 +147,12 @@ const ERROR_MESSAGES = { function inferenceError( reason: Extract["reason"], usage?: Usage, + message: string = ERROR_MESSAGES[reason], ): WorkerInferenceTerminalOutcome { return { type: "error", reason, - message: ERROR_MESSAGES[reason], + message, ...(usage ? { usage: structuredClone(usage) } : {}), }; } @@ -766,28 +768,31 @@ export function createWorkerInferenceExecutor( if (!toolCalls.matchesTerminal(event.message)) { return inferenceError("provider-error"); } - return { - type: "done", - message: projectWorkerInferenceTerminalMessage({ - message: event.message, - modelIdentity, - stopReason: event.reason, - onProviderReplayOmitted: ({ bytes, limitBytes, reason }) => { - if (!isDiagnosticsEnabled(approved.config)) { - return; - } - emitTrustedDiagnosticEvent({ - type: "payload.large", - surface: "worker.provider-replay", - action: "rejected", - bytes, - limitBytes, - reason, - trace: freezeDiagnosticTraceContext(trace), - }); - }, - }), - }; + const terminal = projectWorkerInferenceTerminalMessage({ + message: event.message, + modelIdentity, + stopReason: event.reason, + }); + if (terminal.kind === "provider-replay-unavailable") { + if (isDiagnosticsEnabled(approved.config)) { + const { bytes, limitBytes, reason } = terminal.details; + emitTrustedDiagnosticEvent({ + type: "payload.large", + surface: "worker.provider-replay", + action: "rejected", + bytes, + limitBytes, + reason, + trace: freezeDiagnosticTraceContext(trace), + }); + } + return inferenceError( + "provider-error", + event.message.usage, + WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE, + ); + } + return { type: "done", message: terminal.message }; } if (event.type === "error") { recordUsage(event.error.usage); diff --git a/src/gateway/worker-environments/inference-terminal-message.ts b/src/gateway/worker-environments/inference-terminal-message.ts index c57af082d815..6c2a9ff94247 100644 --- a/src/gateway/worker-environments/inference-terminal-message.ts +++ b/src/gateway/worker-environments/inference-terminal-message.ts @@ -2,7 +2,7 @@ import type { WorkerInferenceTerminalOutcome } from "../../../packages/gateway-p import type { AssistantMessage } from "../../llm/types.js"; import { projectWorkerProviderReplay, - type WorkerProviderReplayOmission, + type WorkerMessageProjection, } from "../../worker/transcript-message.js"; export type WorkerInferenceModelIdentity = { @@ -15,8 +15,7 @@ export function projectWorkerInferenceTerminalMessage(params: { message: AssistantMessage; modelIdentity: WorkerInferenceModelIdentity; stopReason: Extract; - onProviderReplayOmitted?: (omission: WorkerProviderReplayOmission) => void; -}): Extract["message"] { +}): WorkerMessageProjection["message"]> { const content = params.message.content.map((part) => { switch (part.type) { case "text": @@ -88,6 +87,6 @@ export function projectWorkerInferenceTerminalMessage(params: { return projectWorkerProviderReplay({ message: projected, providerReplay: params.message.providerReplay, - onOmitted: params.onProviderReplayOmitted, + purpose: "transcript", }); } diff --git a/src/gateway/worker-environments/worker-turn-launcher.test.ts b/src/gateway/worker-environments/worker-turn-launcher.test.ts index 31d4f4684a26..9abf53ec04b6 100644 --- a/src/gateway/worker-environments/worker-turn-launcher.test.ts +++ b/src/gateway/worker-environments/worker-turn-launcher.test.ts @@ -3,6 +3,7 @@ import fs from "node:fs/promises"; import os from "node:os"; import path from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { WORKER_PROTOCOL_MAX_INFERENCE_PAYLOAD_BYTES } from "../../../packages/gateway-protocol/src/schema/worker-inference.js"; import { createEmbeddedRunLaneController } from "../../agents/embedded-agent-runner/run/lane-controller.js"; import { installSessionPlacementAdmissionProvider } from "../../agents/session-placement-admission.js"; import { SessionManager } from "../../agents/sessions/session-manager.js"; @@ -37,6 +38,7 @@ import { parseWorkerLaunchDescriptor, type WorkerLaunchDescriptor, } from "../../worker/launch-descriptor.js"; +import { WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE } from "../../worker/transcript-message.js"; import type { MintedWorkerCredential } from "./credential.js"; import { createWorkerSessionPlacementStore, @@ -1116,6 +1118,86 @@ describe("worker turn launcher", () => { expect(placements.get(SESSION_ID)).toMatchObject({ state: "active", turnClaim: null }); }); + it("fails impossible replay before handoff and keeps the active placement reusable", async () => { + seedActivePlacement(); + const manager = openSessionManager(); + manager.appendMessage( + makeAgentAssistantMessage({ + content: [{ type: "toolCall", id: "call-replay", name: "read", arguments: {} }], + model: "gpt-test", + providerReplay: { + v: 1, + type: "openai-responses-compaction", + data: "gAAAAlauncherReplayCiphertext", + provider: "openai", + api: "openai-responses", + model: "gpt-test", + baseUrlHash: "ozhevd1smnk8s", + }, + stopReason: "toolUse", + timestamp: 1, + }), + ); + manager.appendMessage({ + role: "toolResult", + toolCallId: "call-replay", + toolName: "read", + content: [{ type: "text", text: "result" }], + details: { payload: "x".repeat(WORKER_PROTOCOL_MAX_INFERENCE_PAYLOAD_BYTES) }, + isError: false, + timestamp: 2, + }); + const runWorkspaceCommand = vi.fn(async (): Promise => { + throw new Error("unexpected worker handoff"); + }); + const acknowledgeCredentialDelivery = vi.fn(() => true); + const startTunnel = vi.fn( + async (): Promise => ({ + environmentId: ENVIRONMENT_ID, + ownerEpoch: OWNER_EPOCH, + remoteSocketPath: "/worker/gateway.sock", + quiesceWorkspace: vi.fn(), + runWorkspaceCommand, + syncWorkspace: vi.fn(), + reconcileWorkspace: vi.fn(), + stop: vi.fn(async () => {}), + }), + ); + const stopTunnel = vi.fn(async () => {}); + const destroy = vi.fn(async () => attachedEnvironment()); + const environments: WorkerTurnEnvironmentService = { + get: vi.fn(() => attachedEnvironment()), + acquireTurnCredential: vi.fn(async () => credential()), + acknowledgeCredentialDelivery, + startTunnel, + stopTunnel, + destroy, + }; + const provider = createWorkerSessionTurnPlacementProvider({ environments, placements }); + const runLocal = vi.fn(async () => ({ meta: { durationMs: 1 } })); + + await expect( + provider.executeTurn( + { + sessionId: SESSION_ID, + sessionKey: SESSION_KEY, + agentId: "main", + runId: "run-replay-local-fallback", + }, + turn("run-replay-local-fallback"), + runLocal, + ), + ).rejects.toThrow(WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE); + + expect(startTunnel).toHaveBeenCalledOnce(); + expect(runWorkspaceCommand).not.toHaveBeenCalled(); + expect(runLocal).not.toHaveBeenCalled(); + expect(acknowledgeCredentialDelivery).not.toHaveBeenCalled(); + expect(stopTunnel).not.toHaveBeenCalled(); + expect(destroy).not.toHaveBeenCalled(); + expect(placements.get(SESSION_ID)).toMatchObject({ state: "active", turnClaim: null }); + }); + it("preserves a terminal workspace result when the worker child later exits nonzero", async () => { seedActivePlacement(); const destroy = vi.fn(async () => attachedEnvironment()); diff --git a/src/gateway/worker-environments/worker-turn-launcher.ts b/src/gateway/worker-environments/worker-turn-launcher.ts index 17c089c6a6a1..bbaa06c6a9e6 100644 --- a/src/gateway/worker-environments/worker-turn-launcher.ts +++ b/src/gateway/worker-environments/worker-turn-launcher.ts @@ -13,6 +13,7 @@ import { emitTrustedDiagnosticEvent, isDiagnosticsEnabled } from "../../infra/di import { formatErrorMessage } from "../../infra/errors.js"; import { redactSensitiveText } from "../../logging/redact.js"; import { parseWorkerLaunchDescriptor } from "../../worker/launch-descriptor.js"; +import { WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE } from "../../worker/transcript-message.js"; import type { WorkerSessionPlacementRecord, WorkerSessionPlacementStore, @@ -81,6 +82,20 @@ type WorkerTurnLauncherOptions = { class WorkerTurnExecutionError extends Error {} class WorkerWorkspaceReconciliationError extends Error {} +function emitProviderReplayRejected( + config: SessionPlacementTurnParams["config"], + details: { reason: string; bytes?: number; limitBytes?: number; count?: number }, +): void { + if (isDiagnosticsEnabled(config)) { + emitTrustedDiagnosticEvent({ + type: "payload.large", + surface: "worker.provider-replay", + action: "rejected", + ...details, + }); + } +} + async function executeLocalTurn(params: { claim: LocalTurnPlacementClaim; placements: WorkerSessionPlacementStore; @@ -245,24 +260,20 @@ async function executeWorkerTurn(params: { turn.userTurnTranscriptRecorder?.hasPersisted() === true; const contextMessages = convertToLlm(manager.buildSessionContext().messages); const leaf = manager.getLeafEntry(); - const initialMessages = windowInitialMessages( + const initialMessagePlan = windowInitialMessages( userMessageAlreadyPersisted && leaf?.type === "message" && leaf.message.role === "user" ? contextMessages.slice(0, -1) : contextMessages, - ({ bytes, limitBytes, reason }) => { - if (!isDiagnosticsEnabled(turn.config)) { - return; - } - emitTrustedDiagnosticEvent({ - type: "payload.large", - surface: "worker.provider-replay", - action: "rejected", - bytes, - limitBytes, - reason, - }); - }, ); + if (initialMessagePlan.kind === "provider-replay-unavailable") { + const details = initialMessagePlan.details; + emitProviderReplayRejected( + turn.config, + "bytes" in details ? details : { count: details.messageCount, reason: details.reason }, + ); + throw new WorkerTurnExecutionError(WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE); + } + const initialMessages = initialMessagePlan.messages; let baseLeafId = manager.getLeafId(); if (!userMessageAlreadyPersisted) { const persisted = turn.userTurnTranscriptRecorder @@ -308,7 +319,7 @@ async function executeWorkerTurn(params: { }); const reasoning = mapThinkingLevelForProvider(turn.thinkLevel); const toolAuthority = resolveWorkerToolAuthority({ modelRef, turn }); - const descriptor = fitLaunchDescriptor( + const launchPlan = fitLaunchDescriptor( (windowedMessages) => parseWorkerLaunchDescriptor({ version: 2, @@ -344,6 +355,15 @@ async function executeWorkerTurn(params: { }), initialMessages, ); + if (launchPlan.kind === "local-fallback") { + emitProviderReplayRejected(turn.config, { + bytes: launchPlan.bytes, + limitBytes: launchPlan.limitBytes, + reason: launchPlan.reason, + }); + throw new WorkerTurnExecutionError(WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE); + } + const descriptor = launchPlan.descriptor; turn.userTurnTranscriptRecorder?.markSentToProvider?.(); turn.onExecutionPhase?.({ phase: "attempt_dispatch", backend: "cloud-worker" }); const handoffAbort = new AbortController(); diff --git a/src/gateway/worker-environments/worker-turn-payload.test.ts b/src/gateway/worker-environments/worker-turn-payload.test.ts index cc7e728a0eb6..8d775557cfa4 100644 --- a/src/gateway/worker-environments/worker-turn-payload.test.ts +++ b/src/gateway/worker-environments/worker-turn-payload.test.ts @@ -1,8 +1,97 @@ -import { describe, expect, it, vi } from "vitest"; -import { WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES } from "../../../packages/gateway-protocol/src/schema/worker-admission.js"; +import { describe, expect, it } from "vitest"; +import { + WORKER_PROTOCOL_MAX_INFERENCE_PAYLOAD_BYTES, + WORKER_INFERENCE_MAX_CONTEXT_MESSAGES, +} from "../../../packages/gateway-protocol/src/schema/worker-inference.js"; import type { AgentMessage } from "../../agents/runtime/index.js"; import type { SessionPlacementTurnParams } from "../../agents/session-placement-admission.js"; -import { assertSupportedTurn, windowInitialMessages } from "./worker-turn-payload.js"; +import type { WorkerLaunchDescriptor } from "../../worker/launch-descriptor.js"; +import { + assertSupportedTurn, + fitLaunchDescriptor, + windowInitialMessages, +} from "./worker-turn-payload.js"; + +const PROVIDER_REPLAY = { + v: 1 as const, + type: "openai-responses-compaction", + data: "opaque-worker-replay", + provider: "openai", + api: "openai-responses", + model: "gpt-5.6-luna", + baseUrlHash: "ozhevd1smnk8s", +}; + +function userMessage(text: string, timestamp: number): AgentMessage { + return { role: "user", content: [{ type: "text", text }], timestamp }; +} + +function assistantMessage(timestamp: number, replay = false): AgentMessage { + return { + role: "assistant", + content: [{ type: "text", text: "visible" }], + api: "openai-responses", + provider: "openai", + model: "gpt-5.6-luna", + ...(replay ? { providerReplay: structuredClone(PROVIDER_REPLAY) } : {}), + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp, + }; +} + +function toolResultMessage(details: unknown, timestamp: number): AgentMessage { + return { + role: "toolResult", + toolCallId: "call-replay", + toolName: "read", + content: [{ type: "text", text: "result" }], + details, + isError: false, + timestamp, + }; +} + +function buildDescriptor( + initialMessages: WorkerLaunchDescriptor["assignment"]["initialMessages"], +): WorkerLaunchDescriptor { + return { + version: 2, + socketPath: "/tmp/worker.sock", + admission: { + environmentId: "environment", + credential: "worker-fixture-credential", + sessionId: "session", + ownerEpoch: 1, + rpcSetVersion: 1, + handshake: { + bundleHash: "a".repeat(64), + openclawVersion: "test", + protocolFeatures: [], + }, + }, + assignment: { + runId: "run", + turnId: "turn", + prompt: "prompt", + suppressPromptTranscript: true, + workspaceDir: "/tmp/workspace", + modelRef: { provider: "openai", model: "gpt-5.6-luna" }, + inferenceOptions: {}, + initialMessages, + transcript: { baseLeafId: null, nextSeq: 1 }, + liveEvents: { ackedSeq: 0, nextSeq: 1 }, + toolAuthority: { allowedToolNames: [] }, + }, + }; +} describe("assertSupportedTurn", () => { it("accepts scheduled authority for the worker launch envelope", () => { @@ -36,46 +125,146 @@ describe("assertSupportedTurn", () => { }); describe("windowInitialMessages", () => { - it("reports replay omitted from the worker launch envelope", () => { - const onOmitted = vi.fn(); - const messages = windowInitialMessages( - [ - { - role: "assistant", - content: [{ type: "text", text: "visible" }], - api: "openai-responses", - provider: "openai", - model: "gpt-5.5", - providerReplay: { - v: 1, - type: "openai-responses-compaction", - data: "x".repeat(WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES + 1), - provider: "openai", - api: "openai-responses", - model: "gpt-5.5", - baseUrlHash: "ozhevd1smnk8s", - }, - usage: { - input: 1, - output: 1, - cacheRead: 0, - cacheWrite: 0, - totalTokens: 2, - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, - }, - stopReason: "stop", - timestamp: 1, - } as AgentMessage, - ], - onOmitted, + it("pins the newest replay carrier when the normal cutoff would pass it", () => { + const history = [userMessage("old", 1), assistantMessage(2, true)]; + history.push( + ...Array.from({ length: WORKER_INFERENCE_MAX_CONTEXT_MESSAGES - 2 }, (_value, index) => + userMessage(`suffix-${index}`, index + 3), + ), ); - expect(messages).toHaveLength(1); - expect(messages[0]).not.toHaveProperty("providerReplay"); - expect(onOmitted).toHaveBeenCalledWith({ - bytes: WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES + 1, - limitBytes: WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES, - reason: "provider-replay-data-budget", + const result = windowInitialMessages(history); + + expect(result.kind).toBe("complete"); + if (result.kind !== "complete") { + throw new Error("expected complete window"); + } + expect(result.messages).toHaveLength(WORKER_INFERENCE_MAX_CONTEXT_MESSAGES - 1); + expect(result.messages[0]).toMatchObject({ + role: "assistant", + providerReplay: PROVIDER_REPLAY, + }); + }); + + it("reserves one context slot for the current prompt", () => { + const history = Array.from({ length: WORKER_INFERENCE_MAX_CONTEXT_MESSAGES }, (_value, index) => + userMessage(`history-${index}`, index + 1), + ); + + const result = windowInitialMessages(history); + + expect(result.kind).toBe("complete"); + if (result.kind !== "complete") { + throw new Error("expected complete window"); + } + expect(result.messages).toHaveLength(WORKER_INFERENCE_MAX_CONTEXT_MESSAGES - 1); + expect(result.messages[0]).toMatchObject({ + role: "user", + content: [{ type: "text", text: "history-1" }], + }); + }); + + it("keeps historical replay that fits launch inference but not a transcript commit frame", () => { + const message = assistantMessage(1, true); + if (message.role !== "assistant" || !message.providerReplay) { + throw new Error("expected replay carrier"); + } + const ciphertext = "\0".repeat(12_000); + message.providerReplay = { + ...message.providerReplay, + id: "i".repeat(65_536), + data: ciphertext, + }; + + const result = windowInitialMessages([message]); + + expect(result.kind).toBe("complete"); + if (result.kind !== "complete") { + throw new Error("expected complete window"); + } + expect(result.messages[0]).toMatchObject({ + role: "assistant", + providerReplay: { id: "i".repeat(65_536), data: ciphertext }, + }); + }); + + it("returns a typed degraded result instead of slicing past replay", () => { + const history = [assistantMessage(1, true)]; + history.push( + ...Array.from({ length: WORKER_INFERENCE_MAX_CONTEXT_MESSAGES - 1 }, (_value, index) => + userMessage(`suffix-${index}`, index + 2), + ), + ); + + expect(windowInitialMessages(history)).toEqual({ + kind: "provider-replay-unavailable", + details: { + reason: "provider-replay-message-limit", + messageCount: WORKER_INFERENCE_MAX_CONTEXT_MESSAGES, + limitMessages: WORKER_INFERENCE_MAX_CONTEXT_MESSAGES - 1, + }, + }); + }); +}); + +describe("fitLaunchDescriptor", () => { + it("drops complete old turns while retaining the replay anchor", () => { + const large = "x".repeat(13 * 1024 * 1024); + const projected = windowInitialMessages([ + userMessage(large, 1), + userMessage(large, 2), + assistantMessage(3, true), + ]); + if (projected.kind !== "complete") { + throw new Error("expected complete projection"); + } + + const plan = fitLaunchDescriptor(buildDescriptor, projected.messages); + + expect(plan.kind).toBe("launch"); + if (plan.kind !== "launch") { + throw new Error("expected launch plan"); + } + expect(plan.descriptor.assignment.initialMessages).toHaveLength(2); + expect(plan.descriptor.assignment.initialMessages[1]).toMatchObject({ + role: "assistant", + providerReplay: PROVIDER_REPLAY, + }); + }); + + it("drops a non-user prefix directly to the replay owner", () => { + const projected = windowInitialMessages([ + toolResultMessage({ payload: "x".repeat(WORKER_PROTOCOL_MAX_INFERENCE_PAYLOAD_BYTES) }, 1), + assistantMessage(2, true), + ]); + if (projected.kind !== "complete") { + throw new Error("expected complete projection"); + } + + const plan = fitLaunchDescriptor(buildDescriptor, projected.messages); + + expect(plan.kind).toBe("launch"); + if (plan.kind !== "launch") { + throw new Error("expected launch plan"); + } + expect(plan.descriptor.assignment.initialMessages).toEqual([ + expect.objectContaining({ role: "assistant", providerReplay: PROVIDER_REPLAY }), + ]); + }); + + it("requires local fallback when the replay unit cannot fit the descriptor", () => { + const projected = windowInitialMessages([ + assistantMessage(1, true), + toolResultMessage({ payload: "x".repeat(WORKER_PROTOCOL_MAX_INFERENCE_PAYLOAD_BYTES) }, 2), + ]); + if (projected.kind !== "complete") { + throw new Error("expected complete projection"); + } + + expect(fitLaunchDescriptor(buildDescriptor, projected.messages)).toMatchObject({ + kind: "local-fallback", + reason: "provider-replay-launch-payload-limit", + limitBytes: WORKER_PROTOCOL_MAX_INFERENCE_PAYLOAD_BYTES, }); }); }); diff --git a/src/gateway/worker-environments/worker-turn-payload.ts b/src/gateway/worker-environments/worker-turn-payload.ts index d45519675939..49a20d8a49c3 100644 --- a/src/gateway/worker-environments/worker-turn-payload.ts +++ b/src/gateway/worker-environments/worker-turn-payload.ts @@ -22,53 +22,84 @@ import type { SessionPlacementTurnParams } from "../../agents/session-placement- import { resolveEffectiveAgentRuntime } from "../../agents/thinking-runtime.js"; import { hasNonzeroUsage, normalizeUsage } from "../../agents/usage.js"; import type { WorkerLaunchDescriptor } from "../../worker/launch-descriptor.js"; +import { + windowWorkerReplayMessages, + type WorkerReplayMessageWindowUnavailable, +} from "../../worker/replay-message-window.js"; import { toWorkerTranscriptMessage, - type WorkerProviderReplayOmission, + type WorkerProviderReplayUnavailable, } from "../../worker/transcript-message.js"; import type { WorkerRuntimeResult } from "../../worker/worker.runtime.js"; -export function windowInitialMessages( - messages: AgentMessage[], - onProviderReplayOmitted?: (omission: WorkerProviderReplayOmission) => void, -): WorkerTranscriptMessage[] { - const projected = messages.flatMap((message) => { - const value = toWorkerTranscriptMessage(message, { onProviderReplayOmitted }); - return value ? [value] : []; - }); - if (projected.length <= WORKER_INFERENCE_MAX_CONTEXT_MESSAGES) { - return projected; +type WorkerInitialMessagePlan = + | { kind: "complete"; messages: WorkerTranscriptMessage[] } + | { + kind: "provider-replay-unavailable"; + details: WorkerProviderReplayUnavailable | WorkerReplayMessageWindowUnavailable; + }; + +export function windowInitialMessages(messages: AgentMessage[]): WorkerInitialMessagePlan { + const windowed = windowWorkerReplayMessages(messages, WORKER_INFERENCE_MAX_CONTEXT_MESSAGES - 1); + if (windowed.kind === "provider-replay-unavailable") { + return windowed; } - const minimumStart = projected.length - WORKER_INFERENCE_MAX_CONTEXT_MESSAGES; - const completeTurnStart = projected.findIndex( - (message, index) => index >= minimumStart && message.role === "user", - ); - if (completeTurnStart < 0) { - throw new Error("Worker turn transcript has no complete context window"); + const projected: WorkerTranscriptMessage[] = []; + for (const message of windowed.messages) { + const result = toWorkerTranscriptMessage(message, "inference"); + if (!result) { + continue; + } + if (result.kind === "provider-replay-unavailable") { + return result; + } + projected.push(result.message); } - return projected.slice(completeTurnStart); + return { kind: "complete", messages: projected }; } +type WorkerLaunchPlan = + | { kind: "launch"; descriptor: WorkerLaunchDescriptor } + | { + kind: "local-fallback"; + reason: "provider-replay-launch-payload-limit"; + bytes: number; + limitBytes: number; + }; + export function fitLaunchDescriptor( build: (initialMessages: WorkerTranscriptMessage[]) => WorkerLaunchDescriptor, messages: WorkerTranscriptMessage[], -): WorkerLaunchDescriptor { +): WorkerLaunchPlan { let initialMessages = messages; while (true) { const descriptor = build(initialMessages); - if ( - Buffer.byteLength(JSON.stringify(descriptor), "utf8") <= - WORKER_PROTOCOL_MAX_INFERENCE_PAYLOAD_BYTES - ) { - return descriptor; + const bytes = Buffer.byteLength(JSON.stringify(descriptor), "utf8"); + if (bytes <= WORKER_PROTOCOL_MAX_INFERENCE_PAYLOAD_BYTES) { + return { kind: "launch", descriptor }; + } + const replayIndex = initialMessages.findLastIndex( + (message) => message.role === "assistant" && message.providerReplay !== undefined, + ); + if (replayIndex === 0) { + return { + kind: "local-fallback", + reason: "provider-replay-launch-payload-limit", + bytes, + limitBytes: WORKER_PROTOCOL_MAX_INFERENCE_PAYLOAD_BYTES, + }; } const nextTurn = initialMessages.findIndex( (message, index) => index > 0 && message.role === "user", ); - if (nextTurn < 0) { + // A replay owner is a valid context start because its checkpoint replaces + // the discarded prefix; never advance past it to reach a later user turn. + const nextStart = + replayIndex > 0 && (nextTurn < 0 || nextTurn > replayIndex) ? replayIndex : nextTurn; + if (nextStart < 0) { throw new Error("Worker turn context exceeds the launch descriptor payload limit"); } - initialMessages = initialMessages.slice(nextTurn); + initialMessages = initialMessages.slice(nextStart); } } diff --git a/src/worker/embedded-agent-transcript.runtime.ts b/src/worker/embedded-agent-transcript.runtime.ts index 87edbd83c900..830412f4ef85 100644 --- a/src/worker/embedded-agent-transcript.runtime.ts +++ b/src/worker/embedded-agent-transcript.runtime.ts @@ -5,12 +5,17 @@ import { WORKER_INFERENCE_MAX_CONTEXT_MESSAGES } from "../../packages/gateway-pr import type { AgentMessage } from "../agents/runtime/index.js"; import type { AgentSessionWriteLockRunner } from "../agents/sessions/agent-session.js"; import type { Context, Message } from "../llm/types.js"; +import { + windowWorkerReplayMessages, + type WorkerReplayMessageWindowUnavailable, +} from "./replay-message-window.js"; import { cloneImageContent, cloneTextContent, - cloneUsage, isWorkerTranscriptMessageFrameSafe, toWorkerTranscriptMessage, + type WorkerMessageProjection, + type WorkerProviderReplayUnavailable, } from "./transcript-message.js"; export function toAgentMessage(message: WorkerTranscriptMessage): Message { @@ -36,61 +41,73 @@ export function toAgentMessage(message: WorkerTranscriptMessage): Message { timestamp: message.timestamp, }; } - return { - ...cloneUsage(message), - diagnostics: message.diagnostics?.map((diagnostic) => structuredClone(diagnostic)), - }; + return structuredClone(message); } -function toWorkerInferenceMessage(message: Message): WorkerInferenceContext["messages"][number] { +function toWorkerInferenceMessage( + message: Message, +): WorkerMessageProjection { if (message.role === "user") { return { - role: "user", - content: - typeof message.content === "string" - ? message.content - : message.content.map((part) => - part.type === "text" ? cloneTextContent(part) : cloneImageContent(part), - ), - timestamp: message.timestamp, - ...(message.runtimeContextCarrier ? { runtimeContextCarrier: true } : {}), + kind: "complete", + message: { + role: "user", + content: + typeof message.content === "string" + ? message.content + : message.content.map((part) => + part.type === "text" ? cloneTextContent(part) : cloneImageContent(part), + ), + timestamp: message.timestamp, + ...(message.runtimeContextCarrier ? { runtimeContextCarrier: true } : {}), + }, }; } - const projected = toWorkerTranscriptMessage(message); + const projected = toWorkerTranscriptMessage(message, "inference"); if (!projected) { throw new Error(`Unsupported inference message role: ${message.role}`); } return projected; } -function windowWorkerInferenceMessages(messages: Context["messages"]): Context["messages"] { - if (messages.length <= WORKER_INFERENCE_MAX_CONTEXT_MESSAGES) { - return messages; - } - const minimumStart = messages.length - WORKER_INFERENCE_MAX_CONTEXT_MESSAGES; - // Start at a user turn when possible so truncation cannot orphan a tool result - // from the assistant tool call that owns it. - for (let index = minimumStart; index < messages.length; index += 1) { - if (messages[index]?.role === "user") { - return messages.slice(index); - } - } - throw new Error("Worker inference context has no complete user turn within the message limit."); -} +type WorkerInferenceContextProjection = + | { kind: "complete"; context: WorkerInferenceContext } + | { + kind: "provider-replay-unavailable"; + details: WorkerProviderReplayUnavailable | WorkerReplayMessageWindowUnavailable; + }; -export function toWorkerInferenceContext(context: Context): WorkerInferenceContext { +export function toWorkerInferenceContext(context: Context): WorkerInferenceContextProjection { + const windowed = windowWorkerReplayMessages( + context.messages, + WORKER_INFERENCE_MAX_CONTEXT_MESSAGES, + ); + if (windowed.kind === "provider-replay-unavailable") { + return windowed; + } + const messages: WorkerInferenceContext["messages"] = []; + for (const message of windowed.messages) { + const projected = toWorkerInferenceMessage(message); + if (projected.kind === "provider-replay-unavailable") { + return projected; + } + messages.push(projected.message); + } return { - ...(context.systemPrompt === undefined ? {} : { systemPrompt: context.systemPrompt }), - messages: windowWorkerInferenceMessages(context.messages).map(toWorkerInferenceMessage), - ...(context.tools - ? { - tools: context.tools.map((tool) => ({ - name: tool.name, - description: tool.description, - parameters: structuredClone(tool.parameters), - })), - } - : {}), + kind: "complete", + context: { + ...(context.systemPrompt === undefined ? {} : { systemPrompt: context.systemPrompt }), + messages, + ...(context.tools + ? { + tools: context.tools.map((tool) => ({ + name: tool.name, + description: tool.description, + parameters: structuredClone(tool.parameters), + })), + } + : {}), + }, }; } @@ -108,13 +125,19 @@ export function createWorkerTranscriptRuntime( ): WorkerTranscriptRuntime { const pendingTranscriptMessages: WorkerTranscriptMessage[] = []; const onMessagePersisted = (message: AgentMessage) => { - const projected = toWorkerTranscriptMessage(message); - if (projected) { - if (!isWorkerTranscriptMessageFrameSafe(projected)) { - throw new Error("Worker transcript message exceeds the protocol payload limit."); - } - pendingTranscriptMessages.push(projected); + const projected = toWorkerTranscriptMessage(message, "transcript"); + if (!projected) { + return; } + if (projected.kind === "provider-replay-unavailable") { + throw new Error( + `Worker transcript cannot persist authoritative provider replay: ${projected.details.reason}.`, + ); + } + if (!isWorkerTranscriptMessageFrameSafe(projected.message)) { + throw new Error("Worker transcript message exceeds the protocol payload limit."); + } + pendingTranscriptMessages.push(projected.message); }; const flushTranscript = async () => { while (pendingTranscriptMessages.length > 0) { diff --git a/src/worker/embedded-agent.runtime.ts b/src/worker/embedded-agent.runtime.ts index c6bbc42398c4..9adfc696b3a3 100644 --- a/src/worker/embedded-agent.runtime.ts +++ b/src/worker/embedded-agent.runtime.ts @@ -33,7 +33,7 @@ import { toWorkerInferenceContext, } from "./embedded-agent-transcript.runtime.js"; import { WORKER_LOCAL_TOOL_NAMES, type WorkerLocalToolName } from "./tool-authority.js"; -import { toWorkerTranscriptMessage } from "./transcript-message.js"; +import { WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE } from "./transcript-message.js"; function toError(value: unknown, fallback: string): Error { return value instanceof Error ? value : new Error(fallback, { cause: value }); @@ -79,15 +79,9 @@ type RunWorkerEmbeddedTurnParams = { signal?: AbortSignal; }; -type RunWorkerEmbeddedTurnResult = { - messages: WorkerTranscriptMessage[]; -}; - const WORKER_TOOL_CONFIG = { plugins: { enabled: false } } satisfies OpenClawConfig; -export async function runWorkerEmbeddedTurn( - params: RunWorkerEmbeddedTurnParams, -): Promise { +export async function runWorkerEmbeddedTurn(params: RunWorkerEmbeddedTurnParams): Promise { const model = createNativeModelOwnedRuntimeModel({ provider: params.modelRef.provider, modelId: params.modelRef.model, @@ -203,13 +197,20 @@ export async function runWorkerEmbeddedTurn( }); session.agent.sessionId = params.sessionId; session.setActiveToolsByName([...activeToolNames]); - session.agent.streamFn = (_model, context, options) => - params.inference.stream({ + session.agent.streamFn = (_model, context, options) => { + const projected = toWorkerInferenceContext(context); + if (projected.kind === "provider-replay-unavailable") { + throw new Error( + `${WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE} (${projected.details.reason})`, + ); + } + return params.inference.stream({ modelRef: params.modelRef, - context: toWorkerInferenceContext(context), + context: projected.context, options: structuredClone(params.inferenceOptions ?? {}), ...(options?.signal ? { signal: options.signal } : {}), }); + }; const liveRuntime = createWorkerLiveRuntime(params.live); const unsubscribe = session.subscribe(liveRuntime.handleSessionEvent); @@ -270,11 +271,4 @@ export async function runWorkerEmbeddedTurn( if (finalTranscriptFailure !== undefined) { throw finalTranscriptFailure; } - - return { - messages: session.agent.state.messages.flatMap((message) => { - const projected = toWorkerTranscriptMessage(message); - return projected ? [projected] : []; - }), - }; } diff --git a/src/worker/replay-message-window.ts b/src/worker/replay-message-window.ts new file mode 100644 index 000000000000..fe4c17a2ed97 --- /dev/null +++ b/src/worker/replay-message-window.ts @@ -0,0 +1,45 @@ +export type WorkerReplayMessageWindowUnavailable = { + reason: "provider-replay-message-limit"; + messageCount: number; + limitMessages: number; +}; + +type WorkerReplayMessageWindow = + | { kind: "complete"; messages: T[] } + | { kind: "provider-replay-unavailable"; details: WorkerReplayMessageWindowUnavailable }; + +type ReplayWindowMessage = { role: string; providerReplay?: unknown }; + +export function windowWorkerReplayMessages( + messages: T[], + limitMessages: number, +): WorkerReplayMessageWindow { + if (messages.length <= limitMessages) { + return { kind: "complete", messages }; + } + const minimumStart = messages.length - limitMessages; + // Replay owner plus suffix is one authoritative unit. Starting after the + // owner leaves a context-blind suffix, so fail instead of trimming through it. + const replayIndex = messages.findLastIndex((message) => message.providerReplay !== undefined); + if (replayIndex >= 0 && messages.length - replayIndex > limitMessages) { + return { + kind: "provider-replay-unavailable", + details: { + reason: "provider-replay-message-limit", + messageCount: messages.length - replayIndex, + limitMessages, + }, + }; + } + const completeTurnStart = messages.findIndex( + (message, index) => index >= minimumStart && message.role === "user", + ); + const start = + replayIndex >= 0 && (completeTurnStart < 0 || completeTurnStart > replayIndex) + ? replayIndex + : completeTurnStart; + if (start < 0) { + throw new Error("Worker context has no complete user turn within the message limit."); + } + return { kind: "complete", messages: messages.slice(start) }; +} diff --git a/src/worker/transcript-message.test.ts b/src/worker/transcript-message.test.ts index 218df84640c1..40473032d4c2 100644 --- a/src/worker/transcript-message.test.ts +++ b/src/worker/transcript-message.test.ts @@ -52,11 +52,12 @@ describe("worker transcript provider replay", () => { const message = assistantWithReplay(); Object.assign(message.providerReplay!, { providerScratch: "private" }); - const projected = toWorkerTranscriptMessage(message); - expect(projected?.role).toBe("assistant"); - if (!projected || projected.role !== "assistant") { + const result = toWorkerTranscriptMessage(message, "transcript"); + expect(result?.kind).toBe("complete"); + if (!result || result.kind !== "complete" || result.message.role !== "assistant") { throw new Error("expected projected assistant message"); } + const projected = result.message; expect(projected.providerReplay).toEqual(providerReplay); expect(JSON.stringify(projected)).not.toContain("providerScratch"); expect(isWorkerTranscriptMessageFrameSafe(projected)).toBe(true); @@ -71,43 +72,51 @@ describe("worker transcript provider replay", () => { expect(toAgentMessage(projected)).toMatchObject({ providerReplay }); }); - it("keeps the maximum replay data budget inside a complete commit frame", () => { + it("keeps replay above 48 KiB whole when the complete commit frame fits", () => { + const ciphertext = `cipher-${"x".repeat(60 * 1024)}-€`; const message = assistantWithReplay({ ...providerReplay, - data: "x".repeat(WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES), + data: ciphertext, }); - const projected = toWorkerTranscriptMessage(message); + const result = toWorkerTranscriptMessage(message, "transcript"); - expect(projected?.role).toBe("assistant"); - expect(projected?.role === "assistant" ? projected.providerReplay?.data.length : 0).toBe( - WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES, - ); - expect(projected && isWorkerTranscriptMessageFrameSafe(projected)).toBe(true); + expect(result?.kind).toBe("complete"); + if (!result || result.kind !== "complete" || result.message.role !== "assistant") { + throw new Error("expected projected assistant message"); + } + expect(result.message.providerReplay?.data).toBe(ciphertext); + expect(isWorkerTranscriptMessageFrameSafe(result.message)).toBe(true); }); it.each([ { - name: "raw UTF-8 data over budget", + name: "raw UTF-8 data over the replay field budget", replay: { ...providerReplay, data: "x".repeat(WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES + 1) }, + reason: "provider-replay-data-budget" as const, }, { - name: "multibyte data over its byte budget", - replay: { ...providerReplay, data: "€".repeat(20_000) }, + name: "multibyte data whose complete frame is over budget", + replay: { ...providerReplay, data: "€".repeat(21_845) }, + reason: "transcript-commit-frame-budget" as const, }, { name: "JSON-escaped data over the complete frame budget", replay: { ...providerReplay, data: "\0".repeat(12_000) }, + reason: "transcript-commit-frame-budget" as const, }, { name: "a schema-valid id over the complete frame budget", replay: { ...providerReplay, id: "i".repeat(65_536), data: "opaque" }, + reason: "transcript-commit-frame-budget" as const, }, - ])("omits the entire replay for $name", ({ replay }) => { - const projected = toWorkerTranscriptMessage(assistantWithReplay(replay)); + ])("degrades without ciphertext for $name", ({ replay, reason }) => { + const result = toWorkerTranscriptMessage(assistantWithReplay(replay), "transcript"); - expect(projected?.role).toBe("assistant"); - expect(projected?.role === "assistant" ? projected.providerReplay : undefined).toBeUndefined(); - expect(projected && isWorkerTranscriptMessageFrameSafe(projected)).toBe(true); + if (!result || result.kind !== "provider-replay-unavailable") { + throw new Error("expected degraded replay projection"); + } + expect(result.details).toMatchObject({ reason }); + expect(JSON.stringify(result.details)).not.toContain(replay.data); }); }); diff --git a/src/worker/transcript-message.ts b/src/worker/transcript-message.ts index 79760fbe39a0..95b56fc7d30b 100644 --- a/src/worker/transcript-message.ts +++ b/src/worker/transcript-message.ts @@ -12,11 +12,22 @@ import type { AssistantMessage, ProviderReplayState } from "../llm/types.js"; const SIZE_FRAME_ID = "00000000-0000-4000-8000-000000000000"; type WorkerTranscriptAssistantMessage = Extract; -export type WorkerProviderReplayOmission = { +export type WorkerProviderReplayUnavailable = { bytes: number; limitBytes: number; reason: "provider-replay-data-budget" | "transcript-commit-frame-budget"; }; +type WorkerProviderReplayUnavailableProjection = { + kind: "provider-replay-unavailable"; + details: WorkerProviderReplayUnavailable; +}; +export type WorkerMessageProjection = + | { kind: "complete"; message: T } + | WorkerProviderReplayUnavailableProjection; +type WorkerMessageProjectionPurpose = "inference" | "transcript"; +export const WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE = + "Cloud worker could not preserve authoritative provider replay. " + + "Stop or reclaim the cloud worker, then retry locally."; export function cloneTextContent(part: { type: "text"; text: string; textSignature?: string }) { return { @@ -30,6 +41,12 @@ export function cloneImageContent(part: { type: "image"; data: string; mimeType: return { type: "image" as const, data: part.data, mimeType: part.mimeType }; } +function providerReplayUnavailable( + details: WorkerProviderReplayUnavailable, +): WorkerProviderReplayUnavailableProjection { + return { kind: "provider-replay-unavailable", details }; +} + function cloneProviderReplay(state: ProviderReplayState): ProviderReplayState { return { v: state.v, @@ -70,41 +87,39 @@ export function projectWorkerProviderReplay< >(params: { message: TMessage; providerReplay: ProviderReplayState | undefined; - onOmitted?: (omission: WorkerProviderReplayOmission) => void; -}): TMessage { + purpose: WorkerMessageProjectionPurpose; +}): WorkerMessageProjection { if (!params.providerReplay) { - return params.message; + return { kind: "complete", message: params.message }; } const dataBytes = Buffer.byteLength(params.providerReplay.data, "utf8"); if (dataBytes > WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES) { - params.onOmitted?.({ + return providerReplayUnavailable({ bytes: dataBytes, limitBytes: WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES, reason: "provider-replay-data-budget", }); - return params.message; } const candidate = { ...params.message, providerReplay: cloneProviderReplay(params.providerReplay), }; + if (params.purpose === "inference") { + return { kind: "complete", message: candidate }; + } const frameBytes = workerTranscriptMessageFrameBytes(candidate); if (frameBytes === undefined || frameBytes > WORKER_PROTOCOL_MAX_PAYLOAD_BYTES) { - params.onOmitted?.({ + return providerReplayUnavailable({ bytes: frameBytes ?? WORKER_PROTOCOL_MAX_PAYLOAD_BYTES + 1, limitBytes: WORKER_PROTOCOL_MAX_PAYLOAD_BYTES, reason: "transcript-commit-frame-budget", }); - return params.message; } - return candidate; + return { kind: "complete", message: candidate }; } -export function cloneUsage( - message: AssistantMessage, - onProviderReplayOmitted?: (omission: WorkerProviderReplayOmission) => void, -): WorkerTranscriptMessage & { role: "assistant" } { - const projected: WorkerTranscriptAssistantMessage = { +function toWorkerAssistantMessage(message: AssistantMessage): WorkerTranscriptAssistantMessage { + return { role: "assistant", content: message.content.map((part) => { if (part.type === "text") { @@ -176,19 +191,12 @@ export function cloneUsage( ...(message.errorBody ? { errorBody: message.errorBody } : {}), timestamp: message.timestamp, }; - return projectWorkerProviderReplay({ - message: projected, - providerReplay: message.providerReplay, - onOmitted: onProviderReplayOmitted, - }); } export function toWorkerTranscriptMessage( message: AgentMessage, - options?: { - onProviderReplayOmitted?: (omission: WorkerProviderReplayOmission) => void; - }, -): WorkerTranscriptMessage | undefined { + purpose: WorkerMessageProjectionPurpose, +): WorkerMessageProjection | undefined { if (message.role === "user") { const content = typeof message.content === "string" @@ -196,22 +204,29 @@ export function toWorkerTranscriptMessage( : message.content.map((part) => part.type === "text" ? cloneTextContent(part) : cloneImageContent(part), ); - return { role: "user", content, timestamp: message.timestamp }; + return { kind: "complete", message: { role: "user", content, timestamp: message.timestamp } }; } if (message.role === "assistant") { - return cloneUsage(message, options?.onProviderReplayOmitted); + return projectWorkerProviderReplay({ + message: toWorkerAssistantMessage(message), + providerReplay: message.providerReplay, + purpose, + }); } if (message.role === "toolResult") { return { - role: "toolResult", - toolCallId: message.toolCallId, - toolName: message.toolName, - content: message.content.map((part) => - part.type === "text" ? cloneTextContent(part) : cloneImageContent(part), - ), - ...(message.details === undefined ? {} : { details: structuredClone(message.details) }), - isError: message.isError, - timestamp: message.timestamp, + kind: "complete", + message: { + role: "toolResult", + toolCallId: message.toolCallId, + toolName: message.toolName, + content: message.content.map((part) => + part.type === "text" ? cloneTextContent(part) : cloneImageContent(part), + ), + ...(message.details === undefined ? {} : { details: structuredClone(message.details) }), + isError: message.isError, + timestamp: message.timestamp, + }, }; } return undefined; diff --git a/src/worker/worker-rpc-clients.test.ts b/src/worker/worker-rpc-clients.test.ts index a98e001205a3..131d763e981b 100644 --- a/src/worker/worker-rpc-clients.test.ts +++ b/src/worker/worker-rpc-clients.test.ts @@ -1,9 +1,8 @@ import { describe, expect, it, vi } from "vitest"; -import { - WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES, - type WorkerHelloOk, - type WorkerLiveEvent, - type WorkerTranscriptMessage, +import type { + WorkerHelloOk, + WorkerLiveEvent, + WorkerTranscriptMessage, } from "../../packages/gateway-protocol/src/schema/worker-admission.js"; import type { WorkerInferenceEventFrame, @@ -288,7 +287,7 @@ describe("worker transcript commit client", () => { }); }); - it("commits a terminal assistant message at the provider replay budget", async () => { + it("commits a terminal assistant message with replay near the frame ceiling", async () => { const harness = connectionHarness(); harness.requestTranscriptCommit.mockResolvedValueOnce({ type: "res", @@ -309,7 +308,7 @@ describe("worker transcript commit client", () => { providerReplay: { v: 1, type: "openai-responses-compaction", - data: "x".repeat(WORKER_PROVIDER_REPLAY_MAX_DATA_BYTES), + data: "x".repeat(60 * 1024), provider: "openai", api: "openai-responses", model: "gpt-5.6-sol", diff --git a/src/worker/worker.fault-injection.test.ts b/src/worker/worker.fault-injection.test.ts index bb2077f91160..b210eb95a5af 100644 --- a/src/worker/worker.fault-injection.test.ts +++ b/src/worker/worker.fault-injection.test.ts @@ -13,6 +13,7 @@ import type { WorkerInferenceStartParams, WorkerInferenceTerminalOutcome, } from "../../packages/gateway-protocol/src/schema/worker-inference.js"; +import { runWorkerProviderReplayRoundTrip } from "../../test/helpers/worker-provider-replay-roundtrip.js"; import { SessionManager } from "../agents/sessions/session-manager.js"; import { resolveSessionTranscriptRuntimeTarget, @@ -180,7 +181,7 @@ type TranscriptGate = { }; type ProviderPlan = - | { kind: "immediate"; text: string } + | { kind: "immediate"; text: string; outcome?: WorkerInferenceTerminalOutcome } | { kind: "partitioned"; firstRelease: Deferred; @@ -197,6 +198,15 @@ type WorkerClients = { inference: WorkerInferenceProxyClient; }; +type WorkerClientOptions = { + admissionProof?: string; + epoch?: number; + baseLeafId?: string | null; + initialSeq?: number; + initialAckedSeq?: number; + runId?: string; +}; + class ComposedGatewayHarness { readonly root: string; readonly stateDir: string; @@ -316,19 +326,10 @@ class ComposedGatewayHarness { this.faults.push(rule); } - createClients( - params: { - admissionProof?: string; - epoch?: number; - baseLeafId?: string | null; - initialSeq?: number; - initialAckedSeq?: number; - runId?: string; - } = {}, - ): WorkerClients { + createDescriptor(params: WorkerClientOptions = {}): WorkerLaunchDescriptor { const epoch = params.epoch ?? this.epoch; const credential = params.admissionProof ?? CREDENTIAL; - const descriptor: WorkerLaunchDescriptor = { + return { version: 2, socketPath: this.socketPath, admission: { @@ -358,6 +359,11 @@ class ComposedGatewayHarness { }, }, }; + } + + createClients(params: WorkerClientOptions = {}): WorkerClients { + const descriptor = this.createDescriptor(params); + const epoch = descriptor.admission.ownerEpoch; const connection = createWorkerConnection({ socketPath: this.socketPath, connectParams: buildWorkerConnectParams(descriptor), @@ -559,7 +565,7 @@ class ComposedGatewayHarness { } const plan = this.providerPlan; if (plan.kind === "immediate") { - return doneOutcome(plan.text); + return structuredClone(plan.outcome ?? doneOutcome(plan.text)); } if (plan.kind === "pending") { plan.started.resolve(); @@ -715,6 +721,17 @@ describe("cloud worker milestone 2 fault injection", () => { await harness.close(); }); + it("replays captured compaction exactly after worker commit and canonical reopen", async () => { + await runWorkerProviderReplayRoundTrip({ + createDescriptor: (options) => harness.createDescriptor(options), + requestParams: (method) => harness.requestParams(method), + sessionTarget: harness.sessionTarget, + setOutcome: (outcome) => { + harness.providerPlan = { kind: "immediate", text: "roundtrip", outcome }; + }, + }); + }); + it("survives repeated tunnel partitions without transcript duplication, live replay, or rebilling", async () => { const current = harness.createClients(); clients.push(current); diff --git a/src/worker/worker.runtime.test.ts b/src/worker/worker.runtime.test.ts index e62619d8cbab..f8eddfc3cc54 100644 --- a/src/worker/worker.runtime.test.ts +++ b/src/worker/worker.runtime.test.ts @@ -34,6 +34,7 @@ import { import { listRunningSessions } from "../agents/bash-process-registry.js"; import { rawDataToString } from "../infra/ws.js"; import { buildWorkerConnectParams, type WorkerLaunchDescriptor } from "./launch-descriptor.js"; +import { WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE } from "./transcript-message.js"; import { WorkerAdmissionDeadlineExceededError } from "./worker-connection-contract.js"; import { createWorkerConnection, WorkerConnectionStoppedError } from "./worker-connection.js"; import { @@ -54,6 +55,15 @@ const SESSION_ID = "worker-session"; const RUN_ID = "worker-run"; const OWNER_EPOCH = 4; const MODEL_REF = { provider: "openai", model: "gpt-5.6-luna" } as const; +const WORKER_LOOP_REPLAY = { + v: 1 as const, + type: "openai-responses-compaction", + data: "opaque-worker-loop-replay", + provider: "openai", + api: "openai-responses", + model: MODEL_REF.model, + baseUrlHash: "ozhevd1smnk8s", +}; const BUNDLE_HASH = Array.from({ length: 64 }, () => "a").join(""); const CREDENTIAL = ["worker", "fixture", "admission"].join("-"); @@ -1148,16 +1158,20 @@ describe("worker runtime", () => { ).toBe(true); }); - it("windows near-limit history for every local tool-loop inference", async () => { + it("keeps a pinned replay anchor through repeated local tool-loop inference", async () => { const { gateway, launch } = await setup({ inferencePlans: ["tool", "text"] }); launch.assignment.initialMessages = Array.from( - { length: WORKER_INFERENCE_MAX_CONTEXT_MESSAGES }, + { length: WORKER_INFERENCE_MAX_CONTEXT_MESSAGES - 2 }, (_value, index): WorkerTranscriptMessage => ({ role: "user", content: [{ type: "text", text: `history-${index}` }], timestamp: index + 1, }), ); + launch.assignment.initialMessages[2] = { + ...assistantMessage([{ type: "text", text: "checkpoint suffix" }], "stop"), + providerReplay: structuredClone(WORKER_LOOP_REPLAY), + }; await expect(runWorkerDescriptor(launch)).resolves.toMatchObject({ status: "completed" }); @@ -1167,6 +1181,11 @@ describe("worker runtime", () => { WORKER_INFERENCE_MAX_CONTEXT_MESSAGES, ); expect(request.context.messages[0]?.role).toBe("user"); + expect( + request.context.messages.find( + (message) => message.role === "assistant" && message.providerReplay, + ), + ).toMatchObject({ providerReplay: WORKER_LOOP_REPLAY }); } expect( gateway.inferenceRequests[1]?.context.messages.some( @@ -1182,6 +1201,43 @@ describe("worker runtime", () => { .map((message) => message.role), ).toEqual(["user", "assistant", "toolResult", "assistant"]); }); + + it("fails before a second inference when the replay unit outgrows the window", async () => { + const { gateway, launch } = await setup({ inferencePlans: ["tool", "text"] }); + launch.assignment.initialMessages = Array.from( + { length: WORKER_INFERENCE_MAX_CONTEXT_MESSAGES - 1 }, + (_value, index): WorkerTranscriptMessage => ({ + role: "user", + content: [{ type: "text", text: `history-${index}` }], + timestamp: index + 1, + }), + ); + launch.assignment.initialMessages[0] = { + ...assistantMessage([{ type: "text", text: "checkpoint suffix" }], "stop"), + providerReplay: structuredClone(WORKER_LOOP_REPLAY), + }; + + await expect(runWorkerDescriptor(launch)).resolves.toEqual({ + status: "failed", + reason: "turn-failed", + }); + + expect(gateway.inferenceRequests).toHaveLength(1); + expect(gateway.inferenceRequests[0]?.context.messages).toHaveLength( + WORKER_INFERENCE_MAX_CONTEXT_MESSAGES, + ); + expect(gateway.inferenceRequests[0]?.context.messages[0]).toMatchObject({ + providerReplay: WORKER_LOOP_REPLAY, + }); + const terminal = gateway.transcriptRequests + .flatMap((request) => request.messages) + .toReversed() + .find((message) => message.role === "assistant"); + expect(terminal).toMatchObject({ + stopReason: "error", + errorMessage: `${WORKER_PROVIDER_REPLAY_LOCAL_RETRY_MESSAGE} (provider-replay-message-limit)`, + }); + }); }); describe("worker reconnect clients", () => { diff --git a/test/helpers/worker-provider-replay-roundtrip.ts b/test/helpers/worker-provider-replay-roundtrip.ts new file mode 100644 index 000000000000..de6cf979511c --- /dev/null +++ b/test/helpers/worker-provider-replay-roundtrip.ts @@ -0,0 +1,170 @@ +import { expect } from "vitest"; +import { convertResponsesMessages } from "../../packages/ai/src/providers/openai-responses-shared.js"; +import { captureOpenAIResponsesCompaction } from "../../packages/ai/src/transports/openai-responses-compaction-replay.js"; +import type { WorkerTranscriptMessage } from "../../packages/gateway-protocol/src/schema/worker-admission.js"; +import type { + WorkerInferenceStartParams, + WorkerInferenceTerminalOutcome, +} from "../../packages/gateway-protocol/src/schema/worker-inference.js"; +import { validateWorkerInferenceTerminalOutcome } from "../../packages/gateway-protocol/src/schema/worker-inference.js"; +import type { Context, Model } from "../../packages/llm-core/src/types.js"; +import { SessionManager } from "../../src/agents/sessions/session-manager.js"; +import { windowInitialMessages } from "../../src/gateway/worker-environments/worker-turn-payload.js"; +import type { WorkerLaunchDescriptor } from "../../src/worker/launch-descriptor.js"; +import { runWorkerDescriptor } from "../../src/worker/worker.runtime.js"; + +type DescriptorOptions = { + baseLeafId?: string | null; + initialSeq?: number; + initialAckedSeq?: number; + runId?: string; +}; + +type RoundTripHarness = { + createDescriptor(options?: DescriptorOptions): WorkerLaunchDescriptor; + requestParams(method: string): unknown[]; + sessionTarget: Parameters[0]; + setOutcome(outcome: WorkerInferenceTerminalOutcome): void; +}; +type WorkerDoneMessage = Extract["message"]; + +function doneMessage( + model: Model<"openai-responses">, + content: WorkerDoneMessage["content"], +): WorkerDoneMessage { + return { + role: "assistant", + content, + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 1, + }; +} + +export async function runWorkerProviderReplayRoundTrip(harness: RoundTripHarness): Promise { + const baseDescriptor = harness.createDescriptor({ runId: "replay-run-1" }); + const model = { + id: baseDescriptor.assignment.modelRef.model, + name: "Fault replay model", + api: "openai-responses", + provider: baseDescriptor.assignment.modelRef.provider, + baseUrl: "https://api.openai.com/v1", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1_000_000, + maxTokens: 8_192, + } satisfies Model<"openai-responses">; + const ciphertext = "gAAAAworkerReplayCiphertextExact_123"; + const captured = doneMessage(model, [ + { type: "text", text: "before checkpoint" }, + { type: "text", text: "after checkpoint" }, + ]); + captureOpenAIResponsesCompaction( + captured, + { type: "compaction", id: "cmp_worker_roundtrip", encrypted_content: ciphertext }, + 1, + model, + ); + expect(validateWorkerInferenceTerminalOutcome({ type: "done", message: captured })).toBe(true); + harness.setOutcome({ type: "done", message: captured }); + baseDescriptor.assignment.turnId = "replay-turn-1"; + baseDescriptor.assignment.prompt = "capture replay"; + + const first = await runWorkerDescriptor(baseDescriptor); + expect(first.status).toBe("completed"); + if (first.status !== "completed") { + throw new Error("expected completed first worker turn"); + } + const committed = ( + harness.requestParams("worker.transcript.commit") as Array<{ + messages: WorkerTranscriptMessage[]; + }> + ).flatMap((request) => request.messages); + const committedReplay = committed.find( + (message) => message.role === "assistant" && message.providerReplay, + ); + expect( + committedReplay?.role === "assistant" ? committedReplay.providerReplay?.data : undefined, + ).toBe(ciphertext); + + const canonicalMessages = SessionManager.open(harness.sessionTarget) + .getBranch() + .flatMap((entry) => (entry.type === "message" ? [entry.message] : [])); + const replayOwner = canonicalMessages.find( + (message) => message.role === "assistant" && message.providerReplay, + ); + expect(replayOwner?.role === "assistant" ? replayOwner.providerReplay?.data : undefined).toBe( + ciphertext, + ); + const initial = windowInitialMessages(canonicalMessages); + if (initial.kind !== "complete") { + throw new Error("expected replayable canonical history"); + } + + harness.setOutcome({ + type: "done", + message: doneMessage(model, [{ type: "text", text: "next reply" }]), + }); + const liveAckedSeq = Math.max( + ...(harness.requestParams("worker.live-event") as Array<{ seq: number }>).map( + (request) => request.seq, + ), + ); + const secondDescriptor = harness.createDescriptor({ + runId: "replay-run-2", + baseLeafId: first.transcriptLeafId, + initialSeq: first.transcriptNextSeq, + initialAckedSeq: liveAckedSeq, + }); + secondDescriptor.assignment.turnId = "replay-turn-2"; + secondDescriptor.assignment.prompt = "next worker turn"; + secondDescriptor.assignment.initialMessages = initial.messages; + await expect(runWorkerDescriptor(secondDescriptor)).resolves.toMatchObject({ + status: "completed", + }); + + const requests = harness.requestParams("worker.inference.start") as WorkerInferenceStartParams[]; + const nextContext = requests[1]?.context; + if (!nextContext) { + throw new Error("missing next worker inference context"); + } + const converted = convertResponsesMessages( + model, + nextContext as Context, + new Set([model.provider]), + ); + const compactionIndex = converted.findIndex((item) => item.type === "compaction"); + expect(converted[compactionIndex]).toEqual({ + type: "compaction", + id: "cmp_worker_roundtrip", + encrypted_content: ciphertext, + }); + const suffixIndex = converted.findIndex( + (item) => + item.type === "message" && + item.role === "assistant" && + Array.isArray(item.content) && + item.content.some( + (content) => content.type === "output_text" && content.text === "after checkpoint", + ), + ); + expect(suffixIndex).toBeGreaterThan(compactionIndex); + expect(converted).toContainEqual( + expect.objectContaining({ + type: "message", + role: "user", + content: [expect.objectContaining({ type: "input_text", text: "next worker turn" })], + }), + ); +}