From 9c631643da56dd0c04588dbe522960854f782706 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Sat, 29 Aug 2026 06:31:02 +0900 Subject: [PATCH] fix(responses): gate combo recovery on availability Select a dispatchable combo target before encrypted task recovery and discard scoped cached plaintext when recovery cannot proceed. Strengthen the combo regression to verify the exact recovered assignment reaches the child. --- .../responses/agent-task-recovery-cache.ts | 8 ++ src/server/responses/agent-task-recovery.ts | 72 +++++++++++++----- src/server/responses/core.ts | 70 ++++++++++++------ tests/agent-task-recovery-combo.test.ts | 74 ++++++++++++++++++- 4 files changed, 182 insertions(+), 42 deletions(-) diff --git a/src/server/responses/agent-task-recovery-cache.ts b/src/server/responses/agent-task-recovery-cache.ts index b55767e18e..44e32c62bb 100644 --- a/src/server/responses/agent-task-recovery-cache.ts +++ b/src/server/responses/agent-task-recovery-cache.ts @@ -128,6 +128,10 @@ export async function resolveCachedAgentTaskRecovery( return flight ? waitForRecoveryFlight(flight, abortSignal) : null; } +export function discardCachedAgentTaskRecovery(key: string): void { + deleteRecoveryCacheEntry(key); +} + export function resetAgentTaskRecoveryCache(): void { for (const flight of RECOVERY_FLIGHTS.values()) { flight.controller.abort(new DOMException("Recovery state reset", "AbortError")); @@ -141,3 +145,7 @@ export function agentTaskRecoveryWaiterCountForTests(): number { for (const flight of RECOVERY_FLIGHTS.values()) count += flight.waiters; return count; } + +export function agentTaskRecoveryCacheSnapshotForTests(): { entries: number; bytes: number } { + return { entries: RECOVERY_CACHE.size, bytes: recoveryCacheBytes }; +} diff --git a/src/server/responses/agent-task-recovery.ts b/src/server/responses/agent-task-recovery.ts index 70003c116b..8b409e175b 100644 --- a/src/server/responses/agent-task-recovery.ts +++ b/src/server/responses/agent-task-recovery.ts @@ -5,6 +5,7 @@ import { readBoundedResponseBody } from "../../lib/bounded-body"; import { isApiAuthRequired, isProxyAdmissionSecret } from "../auth-cors"; import { structurallyValidFernetTokens } from "./encrypted-payload"; import { + discardCachedAgentTaskRecovery, resetAgentTaskRecoveryCache, resolveCachedAgentTaskRecovery, } from "./agent-task-recovery-cache"; @@ -266,6 +267,38 @@ function recoveryAdmission(req: Request, config: OcxConfig): RecoveryAdmission | return { headers, cacheScope }; } +interface AdmittedRecovery { + envelope: AgentEnvelope; + admission: RecoveryAdmission; + cacheKey: string; +} + +function admittedRecovery( + req: Request, + input: unknown, + config: OcxConfig, + parentThreadId?: string | null, +): AdmittedRecovery | null { + const envelope = findEnvelope(input); + if (!envelope) return null; + const admission = recoveryAdmission(req, config); + if (!admission) return null; + const cacheKey = createHash("sha256") + .update(admission.cacheScope) + .update("\0") + .update(parentThreadId ?? "") + .update("\0") + .update(envelope.messageType) + .update("\0") + .update(envelope.taskName) + .update("\0") + .update(envelope.sender) + .update("\0") + .update(envelope.ciphertext) + .digest("hex"); + return { envelope, admission, cacheKey }; +} + function recoveryPayload(envelope: AgentEnvelope, model: string): string { return JSON.stringify({ model, @@ -430,34 +463,33 @@ export async function recoverEncryptedAgentTask( config: OcxConfig, context: { parentThreadId?: string | null; abortSignal?: AbortSignal } = {}, ): Promise { - const envelope = findEnvelope(input); - if (!envelope) return false; // Admission is deliberately checked before cache access. A cache hit must not // turn this process into a plaintext oracle for an unauthenticated caller. - const admission = recoveryAdmission(req, config); - if (!admission) return false; - - const cacheKey = createHash("sha256") - .update(admission.cacheScope) - .update("\0") - .update(context.parentThreadId ?? "") - .update("\0") - .update(envelope.messageType) - .update("\0") - .update(envelope.taskName) - .update("\0") - .update(envelope.sender) - .update("\0") - .update(envelope.ciphertext) - .digest("hex"); + const admitted = admittedRecovery(req, input, config, context.parentThreadId); + if (!admitted) return false; + const { admission, cacheKey, envelope } = admitted; const assignment = await resolveCachedAgentTaskRecovery( cacheKey, options.cacheEntries ?? 200, signal => requestRecovery(admission, envelope, options, signal), context.abortSignal, ); - if (!assignment || context.abortSignal?.aborted) return false; - return injectAssignment(input, envelope, assignment); + if (!assignment) return false; + if (context.abortSignal?.aborted || !injectAssignment(input, envelope, assignment)) { + discardCachedAgentTaskRecovery(cacheKey); + return false; + } + return true; +} + +export function discardEncryptedAgentTaskRecovery( + req: Request, + input: unknown, + config: OcxConfig, + context: { parentThreadId?: string | null } = {}, +): void { + const admitted = admittedRecovery(req, input, config, context.parentThreadId); + if (admitted) discardCachedAgentTaskRecovery(admitted.cacheKey); } export function resetAgentTaskRecoveryState(): void { diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 785e90a295..e04ed3802d 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -276,6 +276,7 @@ import { } from "../relay"; import { agentTaskRecoveryConfig, + discardEncryptedAgentTaskRecovery, recoverEncryptedAgentTask, } from "./agent-task-recovery"; import { relaySseEagerBounded } from "../relay-eager"; @@ -2014,44 +2015,71 @@ export async function handleComboResponses( let comboPayloadReadable = false; const payloadEligible = (target: (typeof combo.targets)[number]): boolean => comboPayloadReadable || !unreadableEncryptedAgentTask || canDecryptUnreadableAgentTask(target); + const initialNow = Date.now(); + let pick: ReturnType = null; if (unreadableEncryptedAgentTask && !combo.targets.some(canDecryptUnreadableAgentTask)) { const recovery = agentTaskRecoveryConfig(config); - let recovered = false; if ( - (options.inboundWire ?? "responses") === "responses" - && isThreadSpawnRequest(req.headers) - && recovery - && !options.comboAttempt + (options.inboundWire ?? "responses") !== "responses" + || !isThreadSpawnRequest(req.headers) + || !recovery + || options.comboAttempt ) { - try { - recovered = await recoverEncryptedAgentTask( - req, - (body as { input?: unknown } | undefined)?.input, - recovery, - config, - { parentThreadId: inboundClientThreadId, abortSignal: options.abortSignal }, - ); - } catch { - recovered = false; - } + discardEncryptedAgentTaskRecovery( + req, + (body as { input?: unknown } | undefined)?.input, + config, + { parentThreadId: inboundClientThreadId }, + ); + return unreadableEncryptedAgentTaskResponse(); + } + pick = pickComboTarget(config, comboId, { + eligible: target => !isComboTargetInCooldown(comboId, target, initialNow), + }); + if (!pick) { + discardEncryptedAgentTaskRecovery( + req, + (body as { input?: unknown } | undefined)?.input, + config, + { parentThreadId: inboundClientThreadId }, + ); + return comboUnavailableResponse(`No available targets for combo: ${comboId}`); + } + let recovered = false; + try { + recovered = await recoverEncryptedAgentTask( + req, + (body as { input?: unknown } | undefined)?.input, + recovery, + config, + { parentThreadId: inboundClientThreadId, abortSignal: options.abortSignal }, + ); + } catch { + recovered = false; } // Recovery has the same in-place input mutation contract as the direct routed path. if ( !recovered || hasUnreadableEncryptedAgentTask((body as { input?: unknown } | undefined)?.input) ) { + discardEncryptedAgentTaskRecovery( + req, + (body as { input?: unknown } | undefined)?.input, + config, + { parentThreadId: inboundClientThreadId }, + ); return unreadableEncryptedAgentTaskResponse(); } comboPayloadReadable = true; comboReplaySnapshot.recoveredPlaintext = true; + } else { + pick = pickComboTarget(config, comboId, { + eligible: target => payloadEligible(target) + && !isComboTargetInCooldown(comboId, target, initialNow), + }); } - const initialNow = Date.now(); - let pick = pickComboTarget(config, comboId, { - eligible: target => payloadEligible(target) - && !isComboTargetInCooldown(comboId, target, initialNow), - }); if (!pick) { return comboUnavailableResponse(`No available targets for combo: ${comboId}`); } diff --git a/tests/agent-task-recovery-combo.test.ts b/tests/agent-task-recovery-combo.test.ts index ffb889959a..cb9e0dd33d 100644 --- a/tests/agent-task-recovery-combo.test.ts +++ b/tests/agent-task-recovery-combo.test.ts @@ -9,6 +9,7 @@ import { runPendingResponseStatePersistForTests, } from "../src/responses/state"; import { resetAgentTaskRecoveryState } from "../src/server/responses/agent-task-recovery"; +import { agentTaskRecoveryCacheSnapshotForTests } from "../src/server/responses/agent-task-recovery-cache"; import { codexHeaders, encryptedInput, @@ -67,7 +68,8 @@ describe("combo path encrypted agent task recovery", () => { test("recovers an all-third-party combo once without retaining plaintext continuation state", async () => { const assignment = "RECOVERED-COMBO-PLAINTEXT-SENTINEL"; const fetchedUrls: string[] = []; - globalThis.fetch = (async (input) => { + const forwardedBodies: string[] = []; + globalThis.fetch = (async (input, init) => { const url = String(input); fetchedUrls.push(url); if (url.includes("chatgpt.com")) { @@ -76,6 +78,7 @@ describe("combo path encrypted agent task recovery", () => { headers: { "content-type": "text/event-stream" }, }); } + forwardedBodies.push(typeof init?.body === "string" ? init.body : ""); return providerCompletion(); }) as typeof fetch; @@ -92,6 +95,10 @@ describe("combo path encrypted agent task recovery", () => { expect(typeof responsePayload.id).toBe("string"); expect(fetchedUrls).toHaveLength(2); expect(fetchedUrls[0]).toContain("chatgpt.com/backend-api/codex/responses"); + expect(forwardedBodies).toHaveLength(1); + expect(forwardedBodies[0]).toContain(assignment); + expect(forwardedBodies[0]).not.toContain(FERNET_TASK); + expect(forwardedBodies[0].match(/Message Type: NEW_TASK/g)).toHaveLength(1); expect(responseContinuationRetainedStoreSnapshot().count).toBe(0); const snapshotPath = join(home, "responses-state.json"); const snapshot = existsSync(snapshotPath) ? readFileSync(snapshotPath, "utf8") : ""; @@ -99,6 +106,71 @@ describe("combo path encrypted agent task recovery", () => { expect(snapshot).not.toContain(responsePayload.id!); }); + test("rejects an all-disabled combo before recovery creates or caches plaintext", async () => { + const assignment = "MUST-NOT-BE-PRODUCED-OR-CACHED"; + const config = comboConfig([{ provider: "xai", model: "grok-4.5" }]); + const headers = codexHeaders(); + config.providers.xai!.disabled = true; + let recoveryFetches = 0; + let providerFetches = 0; + globalThis.fetch = (async (input) => { + if (String(input).includes("chatgpt.com")) { + recoveryFetches += 1; + return new Response(recoverySse(assignment), { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + } + providerFetches += 1; + return providerCompletion(); + }) as typeof fetch; + + const coldResponse = await post( + config, + "combo/routed", + encryptedInput(), + headers, + ); + const coldRaw = await coldResponse.text(); + + expect(coldResponse.status).toBe(503); + expect(JSON.parse(coldRaw)).toMatchObject({ + error: { type: "server_error", code: "combo_unavailable" }, + }); + expect(agentTaskRecoveryCacheSnapshotForTests()).toEqual({ entries: 0, bytes: 0 }); + expect(recoveryFetches).toBe(0); + expect(providerFetches).toBe(0); + expect(coldRaw).not.toContain(assignment); + expect(coldRaw).not.toContain(FERNET_TASK); + + config.providers.xai!.disabled = false; + expect((await post( + config, + "combo/routed", + encryptedInput(), + headers, + )).status).toBe(200); + expect(agentTaskRecoveryCacheSnapshotForTests()).toEqual({ + entries: 1, + bytes: Buffer.byteLength(assignment), + }); + expect(recoveryFetches).toBe(1); + expect(providerFetches).toBe(1); + + config.providers.xai!.disabled = true; + const warmResponse = await post( + config, + "combo/routed", + encryptedInput(), + headers, + ); + + expect(warmResponse.status).toBe(503); + expect(agentTaskRecoveryCacheSnapshotForTests()).toEqual({ entries: 0, bytes: 0 }); + expect(recoveryFetches).toBe(1); + expect(providerFetches).toBe(1); + }); + test("keeps the canonical target bypass in a mixed combo without running recovery", async () => { const forwardedBodies: string[] = []; globalThis.fetch = (async (_input, init) => {