diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index 41daa0cd488..677b218c64c 100644 --- a/scripts/test-layout/layout.json +++ b/scripts/test-layout/layout.json @@ -204,6 +204,7 @@ "acl-error-classification.test.ts": "lib", "active-registry-admission.test.ts": "codex-integration", "adapter-buffered-tool-conformance.test.ts": "adapters", + "adapter-dispatch-local-terminal.test.ts": "responses", "adapter-error-inline.test.ts": "adapters", "adapter-event-oauth-failover.test.ts": "oauth", "adapter-inner-send-budget-wiring.test.ts": "adapters", diff --git a/src/providers/reasoning-metadata.ts b/src/providers/reasoning-metadata.ts index f22ef8363ea..302b04dd24f 100644 --- a/src/providers/reasoning-metadata.ts +++ b/src/providers/reasoning-metadata.ts @@ -175,7 +175,7 @@ function modelLadderValue( * so both sides still bind learned refusals to the same wire credential, and a rotation behind * a stable reference starts clean instead of inheriting the previous credential's refusals. */ -function credentialIdentity(provider: OcxProviderConfig): string | undefined { +export function credentialIdentity(provider: OcxProviderConfig): string | undefined { const resolved = provider._apiKeyAttempt?.reference !== undefined ? provider.apiKey : resolveProviderApiKey(provider.apiKey); diff --git a/src/responses/turn-termination.ts b/src/responses/turn-termination.ts index 03834591eac..fd8a4003760 100644 --- a/src/responses/turn-termination.ts +++ b/src/responses/turn-termination.ts @@ -9,7 +9,7 @@ interface DeliveredFinalAnswerRecord { createdAt: number; } -const scopesByRequest = new WeakMap(); +const scopesByRequest = new WeakMap string | undefined)>(); const deliveredFinalAnswers = new Map(); function pruneDeliveredFinalAnswers(at = Date.now()): void { @@ -58,16 +58,32 @@ function deliveredFinalAnswerText(response: unknown): string | undefined { } /** Bind the normalized per-conversation digest without adding proxy-private fields to the wire body. */ -export function bindTurnTerminationScope(parsed: OcxParsedRequest, scope: string | undefined): void { - // Only the normalized log-conversation digest may key this process-wide map. Refusing any raw - // fallback prevents a future caller from retaining a client header or account identifier here. +export function bindTurnTerminationScope( + parsed: OcxParsedRequest, + scope: string | (() => string | undefined) | undefined, +): void { + // Only the normalized log-conversation digest may key this process-wide map. A bound string must + // already BE that digest; a resolver is deferred to read time so a credential rotation after this + // bind lands the record under the identity that actually served the answer. Both paths enforce the + // digest shape -- refusing any raw value prevents a future caller from retaining a client header + // or account identifier here. + if (typeof scope === "function") { + scopesByRequest.set(parsed, scope); + return; + } if (!scope || !/^[0-9a-f]{32}$/.test(scope)) return; scopesByRequest.set(parsed, scope); } +function recordedScope(parsed: OcxParsedRequest): string | undefined { + const bound = scopesByRequest.get(parsed); + const scope = typeof bound === "function" ? bound() : bound; + return typeof scope === "string" && /^[0-9a-f]{32}$/.test(scope) ? scope : undefined; +} + /** Remember only a final-answer message the proxy actually emitted for this exact conversation. */ export function rememberDeliveredFinalAnswer(parsed: OcxParsedRequest, response: unknown): void { - const scope = scopesByRequest.get(parsed); + const scope = recordedScope(parsed); if (!scope) return; const text = deliveredFinalAnswerText(response); if (!text) return; @@ -88,7 +104,7 @@ export function hasRecordedTrailingDeliveredFinalAnswer( parsed: OcxParsedRequest, messages: readonly OcxMessage[], ): boolean { - const scope = scopesByRequest.get(parsed); + const scope = recordedScope(parsed); if (!scope) return false; pruneDeliveredFinalAnswers(); const record = deliveredFinalAnswers.get(scope); diff --git a/src/server/request-log-conversation.ts b/src/server/request-log-conversation.ts index c69f5773ba1..0f6541e8a6a 100644 --- a/src/server/request-log-conversation.ts +++ b/src/server/request-log-conversation.ts @@ -77,12 +77,22 @@ export function sessionIdHeaderFromRequest(headers: Headers): string | null { * pair separates siblings while still keeping one conversation's overlapping turns together. */ export function sessionLaneIdFromRequest(headers: Headers): string | undefined { + return sessionSpecificLaneIdFromRequest(headers) + ?? normalizeLogConversationId(headers.get("x-codex-parent-thread-id")); +} + +/** + * The lane's child-specific half only: the thread or session a request actually owns, still + * qualified by its parent when both exist. A bare parent is a coalescing group, never an exact + * conversation, so callers that need one child's identity — the Kiro termination scope — must + * take this and refuse parent-only requests rather than collapse siblings into a shared key. + */ +export function sessionSpecificLaneIdFromRequest(headers: Headers): string | undefined { + const specific = normalizeLogConversationId(headers.get("thread-id")) + ?? normalizeLogConversationId(sessionIdHeaderFromRequest(headers)); + if (!specific) return undefined; const parent = normalizeLogConversationId(headers.get("x-codex-parent-thread-id")); - const thread = normalizeLogConversationId(headers.get("thread-id")); - const session = normalizeLogConversationId(sessionIdHeaderFromRequest(headers)); - const specific = thread ?? session; - if (parent && specific) return `${parent}\u0000${specific}`; - return specific ?? parent; + return parent ? `${parent}\u0000${specific}` : specific; } function firstSanitizedConversationId( diff --git a/src/server/responses/adapter-dispatch.ts b/src/server/responses/adapter-dispatch.ts index 6609d5e28e1..532ad9e1391 100644 --- a/src/server/responses/adapter-dispatch.ts +++ b/src/server/responses/adapter-dispatch.ts @@ -122,6 +122,9 @@ export async function prepareAdapterExchange( | "genericFailovers" | "applyFailoverSnapshot" | "noteRoutedAttemptSend" + | "selectionIsCurrent" + | "adapterBindings" + | "refreshDispatchAdapter" >, responseEffects: Pick, sendBudgetState: Pick< @@ -148,6 +151,9 @@ export async function prepareAdapterExchange( anthropicSessionKey, commitResolvedOAuthSelection, applyFailoverSnapshot, + selectionIsCurrent, + adapterBindings, + refreshDispatchAdapter, } = transportState; const { parsed, @@ -194,7 +200,28 @@ export async function prepareAdapterExchange( // sendCount stays 0), and crucially no empty-completion guard, which treats an outputless // terminal as a failed turn and re-invokes the identical request. Routing this through the // ordinary event path would therefore reinstate the loop it exists to end. - const localTerminal = transportState.activeAdapter.localTerminal?.(parsed); + // A bound termination scope resolves its serving credential lazily, so it must read the same + // selection the upcoming send would use. Selection is mutable while a request waits: run the + // dispatch binding's revalidation before evaluating a local terminal, or a stale-credential hit + // suppresses work the send path would have moved onto the newly selected credential. + let selectionCurrent = true; + if ( + transportState.activeAdapter.localTerminal + && !selectionIsCurrent(adapterBindings.get(transportState.activeAdapter)) + ) { + try { + await refreshDispatchAdapter(parsed); + bindRouteReasoningReplayScope({ parsed, providerName: route.providerName, provider: route.provider, + adapterName: transportState.activeAdapter.name, + oauthCredentialSnapshot: transportState.replayOAuthCredentialSnapshot }); + } catch { + // A failed refresh leaves the route stale; a scope evaluated under it could still hit a + // record the replaced credential made and return success where the credential error + // belongs. The send path re-validates and surfaces that error on its own terms. + selectionCurrent = false; + } + } + const localTerminal = selectionCurrent ? transportState.activeAdapter.localTerminal?.(parsed) : undefined; if (localTerminal) { logCtx.localTerminalReason = localTerminal.reason; // Mark the physical attempt too, not just the parent row. `finishRequestAttempt` finalizes the diff --git a/src/server/responses/request-prepare.ts b/src/server/responses/request-prepare.ts index 867f816414d..c4a0f5756e4 100644 --- a/src/server/responses/request-prepare.ts +++ b/src/server/responses/request-prepare.ts @@ -51,7 +51,7 @@ import { buildToolBridgeMaps } from "./collaboration"; import { parseRequest } from "../../responses/parser"; import { anthropicSessionKeyFromParts } from "../../oauth/anthropic-routing"; import { isTranslatorBudgetExceededError } from "../../lib/translator-budget"; -import { bindTurnTerminationScope, rememberDeliveredFinalAnswer } from "../../responses/turn-termination"; +import { rememberDeliveredFinalAnswer } from "../../responses/turn-termination"; import { observeCacheDiagnosticInbound, rebindCacheDiagnosticBody, requestLogSpeedLabel, readConfiguredCodexServiceTier } from "../request-log"; import type { RouteResult } from "../../router"; import { @@ -388,7 +388,6 @@ export async function prepareResponsesRequest( threadIdHeader: req.headers.get("thread-id"), cursorConversationId: parsed._cursorConversationId, }); - bindTurnTerminationScope(parsed, resolvedConversationId); const rememberKiroDeliveredFinalAnswer = (adapterName: string, response: unknown): void => { if (adapterName === "kiro") rememberDeliveredFinalAnswer(parsed, response); }; @@ -785,7 +784,6 @@ export async function prepareResponsesRequest( (reparsed as unknown as Record)[key] = parsed[key]; } } - bindTurnTerminationScope(reparsed, resolvedConversationId); parsed = reparsed; // The recovery mutated `body.input` in place, so `_rawBody` now carries decrypted task // text. Bar it from the continuation cache before any recording path can reach it — diff --git a/src/server/responses/request-transport.ts b/src/server/responses/request-transport.ts index e6834c7e941..ef86db699cf 100644 --- a/src/server/responses/request-transport.ts +++ b/src/server/responses/request-transport.ts @@ -34,6 +34,7 @@ import { noteGenericPoolSelection, } from "../../oauth/generic-account-failover"; import { stampOAuthAccountLabel, usesApiKeyAccount } from "../../providers/label"; +import { credentialIdentity } from "../../providers/reasoning-metadata"; import { resolveProviderTransport } from "../../providers/xai-transport"; import { resolveCopilotApiBaseUrl } from "../../oauth/github-copilot"; import { @@ -48,7 +49,9 @@ import { recordAnthropicAccountQuotaFromHeaders, hasPassiveAccountQuota } from " import { checkOutboundBodySize, describeOutboundBodyRefusal } from "./outbound-body-guard"; import { formatErrorResponse } from "../../bridge"; import { bindRouteReasoningReplayScope } from "./core-replay"; -import { sessionIdHeaderFromRequest, normalizeLogConversationId } from "../request-log-conversation"; +import { sessionIdHeaderFromRequest, sessionLaneIdFromRequest, sessionSpecificLaneIdFromRequest, normalizeLogConversationId } from "../request-log-conversation"; +import { bindTurnTerminationScope } from "../../responses/turn-termination"; +import { codexLogAccountId } from "./core-codex-account"; import { redactSecretString } from "../../lib/redact"; import { selectProactiveApiKeyTransport } from "../../providers/key-failover"; import { isCanonicalOpenAiForwardProvider } from "../../providers/openai-tiers"; @@ -659,6 +662,42 @@ export async function prepareResponsesTransport( delete logCtx.accountLogLabel; } adapter = resolveSelectionAdapter(adapterProvider, config.cacheRetention); + if (adapter.name === "kiro") { + // A log conversation deliberately coalesces a parent's parallel subagents, but a delivered + // final answer may suppress work only for the exact child and serving identity that emitted it. + // Conversation and admission are bound to the admitted request and captured eagerly; the serving + // credential is resolved lazily at check/record time instead. Kiro key-pool and OAuth failover + // can swap the physical transport after this bind -- a record pinned to the credential that + // failed would replay or suppress under the wrong identity, and `route.provider` is the field + // every rotation site rewrites, so reading it late tracks the credential that actually served. + // Hash the composite before binding so no caller, account, key, or route identifier is retained. + // A parent-only lane is a coalescing group, not a child: two siblings that send only + // `x-codex-parent-thread-id` would share one scope and suppress each other, the exact leak + // this guard exists to close. Only a thread/session id or the Cursor conversation counts. + const exactConversation = sessionSpecificLaneIdFromRequest(req.headers) + ?? normalizeLogConversationId(parsed._cursorConversationId); + const admissionIdentity = options.admission?.kind === "configured" + ? `configured:${options.admission.keyId}` + : options.admission?.kind; + bindTurnTerminationScope(parsed, exactConversation + ? () => { + const provider = route.provider; + // credentialIdentity hashes the resolved wire credential, not the configured + // reference: an env:keychain-backed key can rotate behind a stable expression, + // and the scope must follow what was actually sent, never the retained secret. + const servingAccount = replayOAuthCredentialSnapshot?.accountId + ?? (usesApiKeyAccount(provider) ? credentialIdentity(provider) : undefined) + ?? codexLogAccountId(admissionState.authCtx); + return normalizeLogConversationId(JSON.stringify([ + exactConversation, + admissionIdentity, + route.providerName, + route.modelId, + servingAccount, + ])); + } + : undefined); + } bindRouteReasoningReplayScope({ parsed, providerName: route.providerName, @@ -813,6 +852,7 @@ export async function prepareResponsesTransport( applyFailoverSnapshot, selectionIsCurrent, resolveSelectionAdapter, + refreshDispatchAdapter, refreshRunTurnAdapter, oauthDispatch, noteRoutedAttemptSend, diff --git a/structure/INDEX.md b/structure/INDEX.md index bfd932df79a..7ab843e3a81 100644 --- a/structure/INDEX.md +++ b/structure/INDEX.md @@ -55,7 +55,7 @@ Per-vendor contracts and the adapter authority that constructs them. | [`providers/openai-tiers.md`](providers/openai-tiers.md) | Pool/Direct account modes, API-key separation, wire identity, and quota evidence. | | [`providers/cursor.md`](providers/cursor.md) | Cursor native exec, parameterized models, checkpoints, and active-context usage. | | [`providers/google.md`](providers/google.md) | Gemini thought-text, response parts, thought-signature replay, and adjacency repair. | -| [`providers/kiro.md`](providers/kiro.md) | Kiro parallel-tool hints, Responses text controls, and reasoning round-trip. | +| [`providers/kiro.md`](providers/kiro.md) | Kiro parallel-tool hints, Responses text controls, reasoning round-trip, and delivered final-answer termination scope. | | [`providers/xai-grok.md`](providers/xai-grok.md) | Grok Build contract parity and hardening. | | [`providers/chat-compat.md`](providers/chat-compat.md) | Cross-vendor Chat Completions behavior: reasoning, tool results, structured output, parallel tools. | | [`adapters/registry.md`](adapters/registry.md) | The single adapter construction authority and contract inheritance. | diff --git a/structure/manifest.json b/structure/manifest.json index 4f5dc0b3da8..845e5d5ff2d 100644 --- a/structure/manifest.json +++ b/structure/manifest.json @@ -265,7 +265,7 @@ "path": "providers/kiro.md", "tier": 4, "title": "Kiro Provider", - "scope": "Kiro parallel-tool hints, Responses text controls, and reasoning round-trip.", + "scope": "Kiro parallel-tool hints, Responses text controls, reasoning round-trip, and delivered final-answer termination scope.", "documents": [ "src/responses/" ] diff --git a/structure/providers/kiro.md b/structure/providers/kiro.md index b1d0b503d55..c33a7a10f74 100644 --- a/structure/providers/kiro.md +++ b/structure/providers/kiro.md @@ -96,6 +96,53 @@ positive value overwrites an earlier one. Spend arrives in `meteringEvent` as **credits, not tokens**. No captured response carried `tokenUsage` on any event, which is why Kiro usage stays estimated; `meteringEvent` is currently ignored because a credit is not a token count. +## Delivered final-answer termination scope + +A delivered `final_answer` may close a Kiro turn only for the exact request that emitted it. +`src/responses/turn-termination.ts` keeps a process-wide map of delivered-answer fingerprints +keyed by a bound scope rather than by the parsed request's fields. The scope is bound in +`src/server/responses/request-transport.ts` after the final adapter is resolved, and only when +that adapter is `kiro`: the digest covers the conversation lane (`sessionSpecificLaneIdFromRequest`, +or the normalized Cursor conversation id when no lane headers exist), the admission identity, the +routed provider and model, and the serving credential. The composite is hashed before binding, so +no caller, account, key or route identifier is retained, and a request with no conversation +identity binds no scope at all. The lane must name a child: a request carrying only +`x-codex-parent-thread-id` gets the coalescing group as its lane, which every sibling under that +parent shares, so a parent-only request binds no scope rather than collapsing the group into one +conversation. + +Conversation, admission, and route are captured eagerly — they belong to the admitted request — +but the serving credential resolves lazily at check/record time. Kiro key-pool and OAuth failover +can swap the physical transport after the bind, and every rotation site rewrites +`route.provider`, so a deferred resolver lands the record under the credential that actually +served instead of the one that failed. For key-authenticated routes the credential element is +`credentialIdentity` — the digest of the resolved wire credential, not the configured +`env:`/`keychain:` reference, so a rotation behind a stable expression terminates the old +identity's scope; an OAuth snapshot account id wins when one is bound, and the Codex auth +context is the last resort — different keys therefore never share a scope. + +Because the credential resolves lazily, the check must read the same selection the upcoming send +would use, and selection is mutable while a request waits. `prepareAdapterExchange` therefore +runs the dispatch binding's staleness check (`selectionIsCurrent`, `refreshDispatchAdapter`) before +evaluating `localTerminal`, so a record a since-replaced credential made cannot suppress work the +send path would have moved onto the live one. A failed refresh leaves the route stale, so the +terminal is skipped entirely and the ordinary send path surfaces the credential error instead +of a stale-credential hit masquerading as success. + +A bare log-conversation digest is too wide here: it deliberately coalesces a parent's parallel +subagents, and a scope that coarse would let one child's delivered answer suppress a sibling's +unfinished work. Nothing in `request-prepare.ts` binds the scope — including encrypted-task +recovery, whose reparsed body reaches the same transport binding — so the transport write is the +only write of the scope. + +`rememberDeliveredFinalAnswer` records the trailing `final_answer` text fingerprint at delivery +(`adapter-delivery.ts`, `run-turn-execution.ts`, `sidecar-execution.ts`), holding it for one hour +across at most 1,024 scopes. `hasTrailingDeliveredFinalAnswer` in +`src/adapters/kiro/conversation.ts` matches a later turn only while its trailing assistant text is +still the recorded answer, so `src/adapters/kiro/payload.ts` withholds the completion tool and +emits the neutral acknowledgement instead of a continuation prompt +(`tests/server/server-kiro-completion-e2e.test.ts`). + ## Remote image references Kiro's wire inlines base64 bytes only, so a remote `https` image reference cannot be diff --git a/structure/transports/responses.md b/structure/transports/responses.md index b1e28958fa1..d44df38260d 100644 --- a/structure/transports/responses.md +++ b/structure/transports/responses.md @@ -1100,7 +1100,7 @@ is composed from the following owners in `src/server/responses/`; none is a gene | Owner | Responsibility | | --- | --- | | `request-prepare.ts` | Body parsing, combo handoff, final route, encrypted-task recovery and initial admission. | -| `request-transport.ts` | Live credential selection, dispatch bindings, adapter replacement and same-target request identity. | +| `request-transport.ts` | Live credential selection, dispatch bindings, adapter replacement, the Kiro turn-termination scope binding and same-target request identity. | | `request-sidecar-auth.ts` | Sidecar credential resolution and vision preprocessing. | | `response-effects.ts` | Completion notification, replay publication and live request-tool aliases. | | `request-send-budget.ts` | Request-wide send accounting, remaining allowance, the pending recovery permit and the shared ambiguous-resend grant. | diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index a7a7fcee72f..68fd3d13f67 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -36,6 +36,7 @@ "acl-error-classification.test.ts": "lib", "active-registry-admission.test.ts": "codex-integration", "adapter-buffered-tool-conformance.test.ts": "adapters", + "adapter-dispatch-local-terminal.test.ts": "responses", "adapter-error-inline.test.ts": "adapters", "adapter-event-oauth-failover.test.ts": "oauth", "adapter-inner-send-budget-wiring.test.ts": "adapters", diff --git a/tests/responses/adapter-dispatch-local-terminal.test.ts b/tests/responses/adapter-dispatch-local-terminal.test.ts new file mode 100644 index 00000000000..c66b503df5b --- /dev/null +++ b/tests/responses/adapter-dispatch-local-terminal.test.ts @@ -0,0 +1,185 @@ +import { describe, expect, test } from "bun:test"; +import { prepareAdapterExchange } from "../../src/server/responses/adapter-dispatch"; +import { createResponsesSendBudget } from "../../src/server/responses/request-send-budget"; +import { createTranslatorBudget } from "../../src/lib/translator-budget"; +import { + bindTurnTerminationScope, + hasRecordedTrailingDeliveredFinalAnswer, + rememberDeliveredFinalAnswer, +} from "../../src/responses/turn-termination"; +import type { ProviderAdapter } from "../../src/adapters/base"; +import type { OcxMessage, OcxParsedRequest, OcxProviderConfig } from "../../src/types"; + +// The termination scope resolves its serving credential lazily, so a credential switch between +// binding and the localTerminal check leaves the bound selection stale. The dispatch pipeline must +// revalidate and refresh the live selection BEFORE evaluating the terminal — otherwise a record +// made by the stale credential suppresses work the send path would have moved to the new one. +describe("adapter dispatch local-terminal selection revalidation", () => { + const ANSWER = "The answer is 42."; + const STALE_SCOPE = "a".repeat(32); + const FRESH_SCOPE = "b".repeat(32); + const BUILD_SENTINEL = "reached-build-request"; + + const staleProvider = { authMode: "key", apiKey: "kiro-key-a", baseUrl: "https://example.test" } as unknown as OcxProviderConfig; + const freshProvider = { authMode: "key", apiKey: "kiro-key-b", baseUrl: "https://example.test" } as unknown as OcxProviderConfig; + + const messages = [ + { role: "user", content: [{ type: "text", text: "question" }] }, + { role: "assistant", content: [{ type: "text", text: ANSWER }] }, + ] as OcxMessage[]; + + const recordResponse = { + output: [{ + type: "message", + role: "assistant", + phase: "final_answer", + content: [{ type: "output_text", text: ANSWER }], + }], + }; + + const makeParsed = (): OcxParsedRequest => + ({ stream: false, modelId: "gpt-5.6-sol" } as unknown as OcxParsedRequest); + + const seedRecord = (scope: string): void => { + const seeded = makeParsed(); + bindTurnTerminationScope(seeded, scope); + rememberDeliveredFinalAnswer(seeded, recordResponse); + }; + + const makeAdapter = (): ProviderAdapter => ({ + name: "kiro", + localTerminal: parsed => hasRecordedTrailingDeliveredFinalAnswer(parsed, messages) + ? { reason: "kiro_final_answer_already_delivered" } + : undefined, + buildRequest: () => { throw new Error(BUILD_SENTINEL); }, + async* parseStream() { /* unreachable on this path */ }, + }); + + const makeExchange = (args: { + parsed: OcxParsedRequest; + route: { providerName: string; modelId: string; provider: OcxProviderConfig }; + current: boolean; + // What the real refresh does besides rebuild: move the route onto the live selection. + onRefresh?: () => void; + }) => { + const calls = { refresh: 0 }; + const adapter = makeAdapter(); + const transportState = { + adapter, + activeAdapter: adapter, + sameTargetRequest: undefined, + sameTargetParsed: undefined, + sameTargetToken: 0, + transportToken: 0, + oauthDispatch: () => undefined, + imageTierBias: 0, + isOAuth401ReplayProvider: false, + sentOAuthSnapshot: undefined, + refreshResolvedOAuthSelection: async () => null, + replayOAuthCredentialSnapshot: undefined, + invalidateSameTargetRequest: () => {}, + resolveSelectionAdapter: () => adapter, + anthropicPoolAccountId: null, + anthropicPoolFailovers: 0, + anthropicSessionKey: null, + commitResolvedOAuthSelection: async () => null, + genericFailoverAccountId: null, + genericFailovers: 0, + applyFailoverSnapshot: async () => null, + noteRoutedAttemptSend: () => {}, + selectionIsCurrent: () => args.current, + adapterBindings: new WeakMap([[adapter, { kind: "api-key" as const, provider: args.route.provider }]]), + refreshDispatchAdapter: async () => { calls.refresh++; args.onRefresh?.(); return adapter; }, + }; + const requestContext = { + options: {}, + config: {}, + logCtx: {}, + req: new Request("http://localhost/v1/responses", { method: "POST" }), + }; + return { + calls, + run: () => prepareAdapterExchange( + requestContext as never, + { pendingHostAdmissionLease: null, authCtx: {} } as never, + { + parsed: args.parsed, + toolBridgeMaps: { toolNsMap: new Map(), freeformToolNames: new Set(), toolSearchToolNames: new Set() }, + translatorBudget: createTranslatorBudget(), + selectedForwardHeaders: new Headers(), + route: args.route, + inboundWire: "responses", + clientRequestedStream: false, + subagentQuotaFailureModel: "gpt-5.6-sol", + subagentFallbackAccountId: null, + } as never, + transportState as never, + { cancelResponseCompletion: () => {}, notifyResponseComplete: () => {}, refreshRequestToolAliases: () => {} } as never, + createResponsesSendBudget(requestContext as never) as never, + ), + }; + }; + + test("a stale binding refreshes the selection before evaluating the terminal", async () => { + // The record was made under the STALE credential. After the refresh moves the route to the + // live key, the same trailing answer is a different serving identity's work and must send. + seedRecord(STALE_SCOPE); + const route = { providerName: "kiro-test", modelId: "gpt-5.6-sol", provider: staleProvider }; + const parsed = makeParsed(); + bindTurnTerminationScope(parsed, () => (route.provider === staleProvider ? STALE_SCOPE : FRESH_SCOPE)); + const exchange = makeExchange({ parsed, route, current: false, onRefresh: () => { route.provider = freshProvider; } }); + + const response = await exchange.run(); + // The refreshed selection has no record, so the terminal misses and the turn reaches + // buildRequest — which the pipeline maps to an error response carrying the sentinel. + expect(response.status).toBe(400); + expect(await response.text()).toContain(BUILD_SENTINEL); + expect(exchange.calls.refresh).toBe(1); + }); + + test("a post-refresh terminal still suppresses under the live credential", async () => { + // Same setup, but the record belongs to the LIVE credential: refreshing first keeps the + // legitimate suppression instead of evaluating it against the stale selection. + seedRecord(FRESH_SCOPE); + const route = { providerName: "kiro-test", modelId: "gpt-5.6-sol", provider: staleProvider }; + const parsed = makeParsed(); + bindTurnTerminationScope(parsed, () => (route.provider === staleProvider ? STALE_SCOPE : FRESH_SCOPE)); + const exchange = makeExchange({ parsed, route, current: false, onRefresh: () => { route.provider = freshProvider; } }); + + const response = await exchange.run(); + expect(response.status).toBe(200); + expect(exchange.calls.refresh).toBe(1); + const json = await response.json() as { output?: unknown[] }; + expect(json.output ?? []).toHaveLength(0); + }); + + test("a failed refresh skips the terminal check entirely", async () => { + // The record was made under the stale credential and would still match under it: evaluating + // the terminal after a failed refresh is what turns the credential error into a fake success. + seedRecord(STALE_SCOPE); + const route = { providerName: "kiro-test", modelId: "gpt-5.6-sol", provider: staleProvider }; + const parsed = makeParsed(); + bindTurnTerminationScope(parsed, () => (route.provider === staleProvider ? STALE_SCOPE : FRESH_SCOPE)); + const exchange = makeExchange({ + parsed, route, current: false, + onRefresh: () => { throw new Error("credential unavailable"); }, + }); + + const response = await exchange.run(); + expect(response.status).toBe(400); + expect(await response.text()).toContain(BUILD_SENTINEL); + expect(exchange.calls.refresh).toBe(1); + }); + + test("a current binding is not re-resolved before the terminal check", async () => { + seedRecord(STALE_SCOPE); + const route = { providerName: "kiro-test", modelId: "gpt-5.6-sol", provider: staleProvider }; + const parsed = makeParsed(); + bindTurnTerminationScope(parsed, () => (route.provider === staleProvider ? STALE_SCOPE : FRESH_SCOPE)); + const exchange = makeExchange({ parsed, route, current: true }); + + const response = await exchange.run(); + expect(response.status).toBe(200); + expect(exchange.calls.refresh).toBe(0); + }); +}); diff --git a/tests/server/server-agent-task-recovery-replay.test.ts b/tests/server/server-agent-task-recovery-replay.test.ts index 9c660b2e292..86e50d6c17c 100644 --- a/tests/server/server-agent-task-recovery-replay.test.ts +++ b/tests/server/server-agent-task-recovery-replay.test.ts @@ -1,13 +1,32 @@ import { afterEach, beforeEach, expect, spyOn, test } from "bun:test"; -import { createKiroAdapter } from "../../src/adapters/kiro"; +import { KIRO_COMPLETION_TOOL_NAME } from "../../src/adapters/kiro-constants"; import { ADAPTER_REGISTRY } from "../../src/adapters/registry"; -import { parseRequest } from "../../src/responses/parser"; -import { bindTurnTerminationScope, rememberDeliveredFinalAnswer } from "../../src/responses/turn-termination"; -import { conversationIdFromResponsesRequest } from "../../src/server/request-log-conversation"; -import type { OcxParsedRequest } from "../../src/types"; +import { encodeMessage } from "../../src/lib/eventstream-decoder"; +import type { OcxConfig, OcxParsedRequest } from "../../src/types"; import { recoverEncryptedAgentTask, resetAgentTaskRecoveryState, restoreCachedEncryptedAgentTasks } from "../../src/server/responses/agent-task-recovery"; import { codexHeaders, encryptedInput, fakeChatGptJwt, FERNET_TASK, SECOND_FERNET_TASK, originalFetch, recoverySse, routedConfig } from "../helpers/agent-task-recovery"; import { acquireOwnedSpendHome } from "../helpers/owned-spend-home"; + +const kiroEventEncoder = new TextEncoder(); +function kiroCompletionStream(answer: string): ReadableStream { + const input = JSON.stringify({ answer }); + const frames = [ + { name: KIRO_COMPLETION_TOOL_NAME, toolUseId: "completion-1" }, + { name: KIRO_COMPLETION_TOOL_NAME, toolUseId: "completion-1", input }, + { name: KIRO_COMPLETION_TOOL_NAME, toolUseId: "completion-1", stop: true }, + ].map(payload => encodeMessage( + { ":message-type": "event", ":event-type": "toolUseEvent" }, + kiroEventEncoder.encode(JSON.stringify(payload)), + )); + let index = 0; + return new ReadableStream({ + pull(controller) { + if (index < frames.length) controller.enqueue(frames[index++]!); + else controller.close(); + }, + }); +} + afterEach(() => { globalThis.fetch = originalFetch; resetAgentTaskRecoveryState(); }); // Direct handler dispatch never takes the writer lease that startServer would take, so it is refused. @@ -359,69 +378,63 @@ test("Responses handler restores known history and recovers only the new MESSAGE }); test("cached-history reparse preserves recorded final-answer scope without suppressing a user follow-up", async () => { - const { post, providerResponse } = await import("../helpers/agent-task-recovery"); + const { post } = await import("../helpers/agent-task-recovery"); const sessionId = `recovery-final-replay-${crypto.randomUUID()}`; const headers = codexHeaders("acct-caller", { session_id: sessionId }); const config = routedConfig({ enabled: true }); + config.providers["kiro-test"] = { + adapter: "kiro", + baseUrl: "https://kiro.test", + authMode: "key", + apiKey: "synthetic-key", + liveModels: false, + models: ["gpt-5.6-sol"], + } as OcxConfig["providers"][string]; const deliveredAnswer = "The assignment is complete."; - const recorded = parseRequest({ model: "xai/grok-4.5", input: "Earlier turn" }); - bindTurnTerminationScope(recorded, conversationIdFromResponsesRequest({ sessionIdHeader: sessionId })); - rememberDeliveredFinalAnswer(recorded, { output: [{ - type: "message", role: "assistant", phase: "final_answer", - content: [{ type: "output_text", text: deliveredAnswer }], - }] }); + const tools = [{ type: "function", name: "bash", description: "Run a command", parameters: { type: "object" } }]; let recoveries = 0; - const providerBodies: string[] = []; + const upstreamBodies: string[] = []; globalThis.fetch = (async (url: unknown, init?: RequestInit) => { if (String(url).includes("chatgpt.com")) { recoveries++; return new Response(recoverySse("Read the assignment.")); } - providerBodies.push(String(init?.body)); - return providerResponse(); + upstreamBodies.push(String(init?.body)); + return new Response(kiroCompletionStream(deliveredAnswer), { + headers: { "content-type": "application/vnd.amazon.eventstream" }, + }); }) as typeof fetch; const req = new Request("http://localhost/v1/responses", { headers }); expect(await recoverEncryptedAgentTask(req, encryptedInput(), {}, config)).toBe(true); - // Keep the ordinary transport fixture, but exercise Kiro's real pre-send termination hook. - // The remembered record above belongs to a different parsed object: only core can bind - // the new object produced by recovery reparse to the same conversation. - const kiro = createKiroAdapter({ adapter: "kiro", baseUrl: "https://kiro.test", authMode: "key", apiKey: "synthetic-key" }); - const createChat = ADAPTER_REGISTRY["openai-chat"].create; - const inspectedBodies: string[] = []; - const factory = spyOn(ADAPTER_REGISTRY["openai-chat"], "create").mockImplementation((provider, context) => ({ - ...createChat(provider, context), - localTerminal(parsed: OcxParsedRequest) { - inspectedBodies.push(JSON.stringify(parsed._rawBody)); - return kiro.localTerminal?.(parsed); - }, - })); + // Deliver the answer through the real Kiro adapter so the record lands under the same + // transport-bound scope a reparsed body receives on replay. + const first = await post(config, "kiro-test/gpt-5.6-sol", [ + { type: "message", role: "user", content: [{ type: "input_text", text: "Earlier turn" }] }, + ], headers, undefined, { tools }); + expect(first.status).toBe(200); + await first.text(); + expect(upstreamBodies).toHaveLength(1); + const finalMessage = { type: "message", role: "assistant", content: [{ type: "output_text", text: deliveredAnswer }] }; - try { - for (let attempt = 0; attempt < 2; attempt++) { - const response = await post(config, "xai/grok-4.5", [...encryptedInput(), finalMessage], headers); - expect(response.status).toBe(200); - expect((await response.json() as { output: unknown[] }).output).toEqual([]); - expect(providerBodies).toHaveLength(0); - } - const followUp = await post(config, "xai/grok-4.5", [ - ...encryptedInput(), finalMessage, - { type: "message", role: "user", content: "Now explain your result." }, - ], headers); - expect(followUp.status).toBe(200); - await followUp.text(); - expect(providerBodies).toHaveLength(1); - expect(providerBodies[0]).toContain("Now explain your result."); - expect(inspectedBodies).toHaveLength(3); - for (const inspected of inspectedBodies) { - expect(inspected).toContain("Read the assignment."); - expect(inspected).not.toContain(FERNET_TASK); - } - expect(recoveries).toBe(1); - } finally { - factory.mockRestore(); + for (let attempt = 0; attempt < 2; attempt++) { + const response = await post(config, "kiro-test/gpt-5.6-sol", [...encryptedInput(), finalMessage], headers, undefined, { tools }); + expect(response.status).toBe(200); + expect((await response.json() as { output: unknown[] }).output).toEqual([]); + expect(upstreamBodies).toHaveLength(1); } + const followUp = await post(config, "kiro-test/gpt-5.6-sol", [ + ...encryptedInput(), finalMessage, + { type: "message", role: "user", content: [{ type: "input_text", text: "Now explain your result." }] }, + ], headers, undefined, { tools }); + expect(followUp.status).toBe(200); + await followUp.text(); + expect(upstreamBodies).toHaveLength(2); + expect(upstreamBodies[1]).toContain("Now explain your result."); + expect(upstreamBodies[1]).toContain("Read the assignment."); + expect(upstreamBodies[1]).not.toContain(FERNET_TASK); + expect(recoveries).toBe(1); }); test("fresh recovery only handles the current tail, leaving uncached history unchanged", async () => { diff --git a/tests/server/server-kiro-completion-e2e.test.ts b/tests/server/server-kiro-completion-e2e.test.ts index 01327b482c8..71578d799fa 100644 --- a/tests/server/server-kiro-completion-e2e.test.ts +++ b/tests/server/server-kiro-completion-e2e.test.ts @@ -89,16 +89,17 @@ function kiroConfig(baseUrl: string): OcxConfig { } as OcxConfig; } -function scriptedKiroUpstream(attempts: Uint8Array[][]) { +function scriptedKiroUpstream(attempts: Array) { const requests: Array> = []; const server = Bun.serve({ hostname: "127.0.0.1", port: 0, async fetch(req) { requests.push(await req.json() as Record); - const frames = attempts.shift(); - if (!frames) return new Response("unexpected extra Kiro attempt", { status: 500 }); - return new Response(streamOf(frames), { + const attempt = attempts.shift(); + if (!attempt) return new Response("unexpected extra Kiro attempt", { status: 500 }); + if (!Array.isArray(attempt)) return new Response(attempt.body ?? "", { status: attempt.status }); + return new Response(streamOf(attempt), { headers: { "content-type": "application/vnd.amazon.eventstream" }, }); }, @@ -440,6 +441,101 @@ describe("Kiro completion through public server endpoints", () => { } }); + test("a proxy-recorded final answer does not suppress a sibling conversation", async () => { + const deliveredAnswer = "Code mode runs JavaScript that calls tools."; + const upstream = scriptedKiroUpstream([ + completionFrames(deliveredAnswer, "completion-a"), + completionFrames("Sibling work completed.", "completion-b"), + ]); + saveConfig(kiroConfig(upstream.server.url.toString())); + const proxy = startServer(0); + const headers = { + "content-type": "application/json", + "x-codex-parent-thread-id": "shared-parent", + }; + try { + const first = await originalFetch(new URL("/v1/responses", proxy.url), { + method: "POST", + headers: { ...headers, session_id: "child-a" }, + body: JSON.stringify({ model: "kiro-test/gpt-5.6-sol", stream: false, input: "first task" }), + }); + expect(first.status).toBe(200); + await first.text(); + + const sibling = await originalFetch(new URL("/v1/responses", proxy.url), { + method: "POST", + headers: { ...headers, session_id: "child-b" }, + body: JSON.stringify({ + model: "kiro-test/gpt-5.6-sol", + stream: false, + input: [ + { type: "message", role: "user", content: [{ type: "input_text", text: "different task" }] }, + { type: "message", role: "assistant", content: [{ type: "output_text", text: deliveredAnswer }] }, + ], + }), + }); + + expect(sibling.status).toBe(200); + await sibling.text(); + // The shared parent is only a correlation qualifier. A cache hit here would return before + // build/send and leave this at one request, silently terminating the sibling's work. + expect(upstream.requests).toHaveLength(2); + } finally { + await proxy.stop(true); + upstream.server.stop(true); + } + }); + + // Sibling isolation above still let a bare parent id stand in for the lane: a request sending + // only `x-codex-parent-thread-id` got the shared parent digest as its scope, so every sibling + // in the group collapsed into one conversation and suppressed each other. + test("a parent-only lane is not an exact conversation scope", async () => { + const deliveredAnswer = "Code mode runs JavaScript that calls tools."; + const upstream = scriptedKiroUpstream([ + completionFrames(deliveredAnswer, "completion-a"), + completionFrames("Parent-only sibling work completed.", "completion-b"), + ]); + saveConfig(kiroConfig(upstream.server.url.toString())); + const proxy = startServer(0); + // No session/thread id at all: the shared parent header is the only identity either request + // carries, which is exactly what a coalescing group sends for every child under it. + const headers = { + "content-type": "application/json", + "x-codex-parent-thread-id": "shared-parent", + }; + try { + const first = await originalFetch(new URL("/v1/responses", proxy.url), { + method: "POST", + headers, + body: JSON.stringify({ model: "kiro-test/gpt-5.6-sol", stream: false, input: "first task" }), + }); + expect(first.status).toBe(200); + await first.text(); + expect(upstream.requests).toHaveLength(1); + + const sibling = await originalFetch(new URL("/v1/responses", proxy.url), { + method: "POST", + headers, + body: JSON.stringify({ + model: "kiro-test/gpt-5.6-sol", + stream: false, + input: [ + { type: "message", role: "user", content: [{ type: "input_text", text: "different task" }] }, + { type: "message", role: "assistant", content: [{ type: "output_text", text: deliveredAnswer }] }, + ], + }), + }); + expect(sibling.status).toBe(200); + await sibling.text(); + // The parent is a group, not a conversation. Accepted as the scope, the first request's + // record would suppress this sibling and leave upstream at one request. + expect(upstream.requests).toHaveLength(2); + } finally { + await proxy.stop(true); + upstream.server.stop(true); + } + }); + test("a new user request after a proxy-recorded final answer is not suppressed", async () => { const deliveredAnswer = "Code mode runs JavaScript that calls tools."; const upstream = scriptedKiroUpstream([ @@ -576,6 +672,167 @@ describe("Kiro completion through public server endpoints", () => { upstream.server.stop(true); } }); + + // The scope must name the credential that SERVED, not the credential selected when the request + // was bound. Key-pool failover commits a new key mid-request; a record keyed to the failed key + // (or to no key at all, as it was for key-authenticated routes) lets one key's delivered answer + // suppress a replay that a different key must answer for itself. + test("a recorded final answer follows the serving credential through key failover", async () => { + const deliveredAnswer = "Code mode runs JavaScript that calls tools."; + const upstream = scriptedKiroUpstream([ + completionFrames(deliveredAnswer, "completion-a"), + { status: 429, body: JSON.stringify({ message: "quota exceeded" }) }, + completionFrames("A rotated key answered this one.", "completion-b"), + completionFrames("This replay is a different credential's work.", "completion-c"), + ]); + const config = kiroConfig(upstream.server.url.toString()); + config.providers["kiro-test"].apiKey = "kiro-key-a"; + config.providers["kiro-test"].apiKeyPool = [ + { id: "key-a", key: "kiro-key-a" }, + { id: "key-b", key: "kiro-key-b" }, + ]; + saveConfig(config); + const proxy = startServer(0); + const tools = [{ type: "function", name: "bash", description: "Run a command", parameters: { type: "object" } }]; + const replayBody = JSON.stringify({ + model: "kiro-test/gpt-5.6-sol", + stream: false, + input: [ + { type: "message", role: "user", content: [{ type: "input_text", text: "what is code mode" }] }, + { type: "message", role: "assistant", content: [{ type: "output_text", text: deliveredAnswer }] }, + ], + tools, + }); + try { + const first = await originalFetch(new URL("/v1/responses", proxy.url), { + method: "POST", + headers: { "content-type": "application/json", session_id: "kiro-key-boundary-thread" }, + body: JSON.stringify({ + model: "kiro-test/gpt-5.6-sol", + stream: false, + input: "what is code mode", + tools, + }), + }); + expect(first.status).toBe(200); + await first.text(); + expect(upstream.requests).toHaveLength(1); + + // Control: replaying under the SAME committed key still suppresses — the record is only + // meaningful if same-credential replays keep short-circuiting. + const sameKeyReplay = await originalFetch(new URL("/v1/responses", proxy.url), { + method: "POST", + headers: { "content-type": "application/json", session_id: "kiro-key-boundary-thread" }, + body: replayBody, + }); + expect(sameKeyReplay.status).toBe(200); + expect((await sameKeyReplay.json() as { output?: unknown[] }).output ?? []).toHaveLength(0); + expect(upstream.requests).toHaveLength(1); + + // Exhaust key-a on an unrelated conversation so the committed key rotates to key-b. The + // insufficient-quota body is what skips Kiro's transient-throttle probe and reaches the + // pool rotation -- a bare 429 would just retry key-a inside fetchKiroWithRetry. + const rotated = await originalFetch(new URL("/v1/responses", proxy.url), { + method: "POST", + headers: { "content-type": "application/json", session_id: "kiro-key-boundary-other-thread" }, + body: JSON.stringify({ + model: "kiro-test/gpt-5.6-sol", + stream: false, + input: "unrelated task", + tools, + }), + }); + expect(rotated.status).toBe(200); + await rotated.text(); + expect(upstream.requests).toHaveLength(3); + + // The original conversation replays again, now routed onto key-b. Its delivered answer was + // recorded under key-a, so this is a different serving identity's work and must send. Under + // the old scope (serving account always null for key routes, or pinned at bind time) this + // turn was suppressed — exactly the cross-credential leak the review flagged. + const rotatedKeyReplay = await originalFetch(new URL("/v1/responses", proxy.url), { + method: "POST", + headers: { "content-type": "application/json", session_id: "kiro-key-boundary-thread" }, + body: replayBody, + }); + expect(rotatedKeyReplay.status).toBe(200); + await rotatedKeyReplay.text(); + expect(upstream.requests).toHaveLength(4); + } finally { + await proxy.stop(true); + upstream.server.stop(true); + } + }); + + test("a delivered final answer does not follow an env-backed key through rotation", async () => { + // The configured apiKey is an env reference: the attempt's reference stays constant across a + // rotation while the wire credential changes. Identity must follow the resolved value, or the + // rotated key inherits the exhausted key's suppression record. + const deliveredAnswer = "Code mode runs JavaScript that calls tools."; + const upstream = scriptedKiroUpstream([ + completionFrames(deliveredAnswer, "completion-a"), + completionFrames("This replay is a different credential's work.", "completion-b"), + ]); + const config = kiroConfig(upstream.server.url.toString()); + config.providers["kiro-test"].apiKey = "$OCX_KIRO_E2E_ROTATED_KEY"; + saveConfig(config); + const previousEnv = process.env.OCX_KIRO_E2E_ROTATED_KEY; + process.env.OCX_KIRO_E2E_ROTATED_KEY = "kiro-key-a"; + const proxy = startServer(0); + const tools = [{ type: "function", name: "bash", description: "Run a command", parameters: { type: "object" } }]; + const replayBody = JSON.stringify({ + model: "kiro-test/gpt-5.6-sol", + stream: false, + input: [ + { type: "message", role: "user", content: [{ type: "input_text", text: "what is code mode" }] }, + { type: "message", role: "assistant", content: [{ type: "output_text", text: deliveredAnswer }] }, + ], + tools, + }); + try { + const first = await originalFetch(new URL("/v1/responses", proxy.url), { + method: "POST", + headers: { "content-type": "application/json", session_id: "kiro-env-rotation-thread" }, + body: JSON.stringify({ + model: "kiro-test/gpt-5.6-sol", + stream: false, + input: "what is code mode", + tools, + }), + }); + expect(first.status).toBe(200); + await first.text(); + expect(upstream.requests).toHaveLength(1); + + // Control: replaying under the SAME resolved key still suppresses. + const sameKeyReplay = await originalFetch(new URL("/v1/responses", proxy.url), { + method: "POST", + headers: { "content-type": "application/json", session_id: "kiro-env-rotation-thread" }, + body: replayBody, + }); + expect(sameKeyReplay.status).toBe(200); + expect((await sameKeyReplay.json() as { output?: unknown[] }).output ?? []).toHaveLength(0); + expect(upstream.requests).toHaveLength(1); + + // Rotate the credential behind the stable env reference. The record belongs to key-a's + // serving identity, so under the resolved-credential scope this is new work and must send — + // a reference-keyed label would suppress it, which is the leak the review flagged. + process.env.OCX_KIRO_E2E_ROTATED_KEY = "kiro-key-b"; + const rotatedReplay = await originalFetch(new URL("/v1/responses", proxy.url), { + method: "POST", + headers: { "content-type": "application/json", session_id: "kiro-env-rotation-thread" }, + body: replayBody, + }); + expect(rotatedReplay.status).toBe(200); + await rotatedReplay.text(); + expect(upstream.requests).toHaveLength(2); + } finally { + if (previousEnv === undefined) delete process.env.OCX_KIRO_E2E_ROTATED_KEY; + else process.env.OCX_KIRO_E2E_ROTATED_KEY = previousEnv; + await proxy.stop(true); + upstream.server.stop(true); + } + }); }); // The behavioural tests above prove nothing was SENT. This one proves the request log says so.