From 5d55f8fbd47423e4f72c056ce5f1a20da60b015f Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Tue, 22 Sep 2026 11:08:43 +0200 Subject: [PATCH] fix(ai): keep Codex request identity provider-scoped --- packages/agent/src/agent-loop.ts | 4 +- packages/agent/src/agent.ts | 27 ++----- packages/agent/src/proxy.ts | 2 - packages/agent/src/request-metadata.ts | 41 +++++++++++ packages/agent/src/types.ts | 3 - packages/agent/test/proxy.test.ts | 28 -------- packages/agent/test/request-identity.test.ts | 67 ++++++++++------- packages/ai/src/api/openai-codex-responses.ts | 72 +++++++++++++++---- packages/ai/src/api/simple-options.ts | 1 - packages/ai/src/index.ts | 7 +- packages/ai/src/types.ts | 14 ---- .../ai/test/openai-codex-attribution.test.ts | 14 ++-- .../coding-agent/src/core/agent-session.ts | 18 +---- .../core/compaction/branch-summarization.ts | 17 +++-- .../src/core/compaction/compaction.ts | 34 ++------- .../src/core/compaction/request-metadata.ts | 46 ++++++++++++ packages/coding-agent/src/core/sdk.ts | 23 +++--- .../test/branch-summarization.test.ts | 11 +-- packages/coding-agent/test/compaction.test.ts | 19 ++--- .../test/sdk-stream-options.test.ts | 24 ++++++- 20 files changed, 272 insertions(+), 200 deletions(-) create mode 100644 packages/agent/src/request-metadata.ts create mode 100644 packages/coding-agent/src/core/compaction/request-metadata.ts diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index eb01ffada..b0f7e394b 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -15,6 +15,7 @@ import { toToolDeclaration, validateToolArguments, } from "@earendil-works/pi-ai"; +import { rotateAgentRequestIdentity } from "./request-metadata.ts"; import { getDefaultStreamFn } from "./stream-fn.ts"; import type { AgentContext, @@ -301,8 +302,7 @@ async function runLoop( const followUpMessages = (await config.getFollowUpMessages?.()) || []; if (followUpMessages.length > 0) { // A follow-up is a new top-level turn, unlike tool and steering continuations. - const requestIdentity = config.createRequestIdentity?.(); - if (requestIdentity) config = { ...config, requestIdentity }; + rotateAgentRequestIdentity(config.metadata); explicitContinuation = false; pendingMessages = followUpMessages; continue; diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 4169bd436..c299d7c1a 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -1,5 +1,4 @@ import { - type AgentRequestIdentity, createInitialSystemMessage, getCurrentSystemMessage, getCurrentSystemPrompt, @@ -14,6 +13,7 @@ import { uuidv7, } from "@earendil-works/pi-ai"; import { runAgentLoop, runAgentLoopContinue } from "./agent-loop.ts"; +import { createAgentRequestMetadata } from "./request-metadata.ts"; import { getDefaultStreamFn } from "./stream-fn.ts"; import type { AfterToolCallContext, @@ -217,7 +217,7 @@ export class Agent { ) => Promise | AgentLoopTurnUpdate | undefined; private activeRun?: ActiveRun; private readonly attributionSessionId: string; - private activeRequestIdentity?: AgentRequestIdentity; + private activeRequestMetadata?: Record; /** Session identifier forwarded to providers for cache-aware backends. */ public sessionId?: string; /** Optional per-level thinking token budgets forwarded to the stream function. */ @@ -367,7 +367,7 @@ export class Agent { this._state.errorMessage = undefined; this.clearFollowUpQueue(); this.clearSteeringQueue(); - this.activeRequestIdentity = undefined; + this.activeRequestMetadata = undefined; } /** Start a new prompt from text, a single message, or a batch of messages. */ @@ -436,7 +436,7 @@ export class Agent { messages: AgentMessage[], options: { skipInitialSteeringPoll?: boolean } = {}, ): Promise { - this.activeRequestIdentity = this.createRequestIdentity(); + this.activeRequestMetadata = createAgentRequestMetadata(this.attributionSessionId); await this.runWithLifecycle(async (signal) => { await runAgentLoop( messages, @@ -450,7 +450,7 @@ export class Agent { } private async runContinuation(): Promise { - this.activeRequestIdentity ??= this.createRequestIdentity(); + this.activeRequestMetadata ??= createAgentRequestMetadata(this.attributionSessionId); await this.runWithLifecycle(async (signal) => { await runAgentLoopContinue( this.createContextSnapshot(), @@ -473,11 +473,7 @@ export class Agent { let skipInitialSteeringPoll = options.skipInitialSteeringPoll === true; return { model: this._state.model, - requestIdentity: this.activeRequestIdentity, - createRequestIdentity: () => { - this.activeRequestIdentity = this.createRequestIdentity(); - return this.activeRequestIdentity; - }, + metadata: this.activeRequestMetadata, reasoning: this._state.thinkingLevel === "off" ? undefined : this._state.thinkingLevel, sessionId: this.sessionId, onPayload: this.onPayload, @@ -513,17 +509,6 @@ export class Agent { }; } - /** Create an identity for a foreground turn or a provider-side compaction. */ - createRequestIdentity(requestKind: AgentRequestIdentity["requestKind"] = "turn"): AgentRequestIdentity { - return { - sessionId: this.attributionSessionId, - threadId: this.attributionSessionId, - turnId: uuidv7(), - requestKind, - startedAt: Date.now(), - }; - } - private async runWithLifecycle(executor: (signal: AbortSignal) => Promise): Promise { if (this.activeRun) { throw new Error("Agent is already processing."); diff --git a/packages/agent/src/proxy.ts b/packages/agent/src/proxy.ts index 17da289dd..0a4711c55 100644 --- a/packages/agent/src/proxy.ts +++ b/packages/agent/src/proxy.ts @@ -66,7 +66,6 @@ type ProxySerializableStreamOptions = Pick< | "reasoning" | "cacheRetention" | "sessionId" - | "requestIdentity" | "headers" | "metadata" | "transport" @@ -110,7 +109,6 @@ function buildProxyRequestOptions(options: ProxyStreamOptions): ProxySerializabl reasoning: options.reasoning, cacheRetention: options.cacheRetention, sessionId: options.sessionId, - requestIdentity: options.requestIdentity, headers: options.headers, metadata: options.metadata, transport: options.transport, diff --git a/packages/agent/src/request-metadata.ts b/packages/agent/src/request-metadata.ts new file mode 100644 index 000000000..20e8e3bb8 --- /dev/null +++ b/packages/agent/src/request-metadata.ts @@ -0,0 +1,41 @@ +import { uuidv7 } from "@earendil-works/pi-ai"; + +const REQUEST_IDENTITY_METADATA_KEY = "pi.requestIdentity"; + +interface AgentRequestIdentity { + sessionId: string; + threadId: string; + turnId: string; + requestKind: "turn"; + startedAt: number; +} + +export function createAgentRequestMetadata(sessionId: string): Record { + return { + [REQUEST_IDENTITY_METADATA_KEY]: { + sessionId, + threadId: sessionId, + turnId: uuidv7(), + requestKind: "turn", + startedAt: Date.now(), + } satisfies AgentRequestIdentity, + }; +} + +export function getAgentRequestIdentity( + metadata: Record | undefined, +): AgentRequestIdentity | undefined { + const value = metadata?.[REQUEST_IDENTITY_METADATA_KEY]; + if (!value || typeof value !== "object") return undefined; + return value as AgentRequestIdentity; +} + +export function rotateAgentRequestIdentity(metadata: Record | undefined): void { + const identity = getAgentRequestIdentity(metadata); + if (!metadata || !identity) return; + metadata[REQUEST_IDENTITY_METADATA_KEY] = { + ...identity, + turnId: uuidv7(), + startedAt: Date.now(), + } satisfies AgentRequestIdentity; +} diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index fe790a60b..8da843f56 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -1,5 +1,4 @@ import type { - AgentRequestIdentity, Api, AssistantMessage, AssistantMessageEvent, @@ -189,8 +188,6 @@ export interface PrepareNextTurnContext extends AgentTurnContext {} export interface AgentLoopConfig extends SimpleStreamOptions { model: Model; - /** Rotate the identity when a queued follow-up starts a new top-level turn. */ - createRequestIdentity?: () => AgentRequestIdentity; /** * Converts AgentMessage[] to LLM-compatible Message[] before each LLM call. diff --git a/packages/agent/test/proxy.test.ts b/packages/agent/test/proxy.test.ts index 66d273b57..43e5174c3 100644 --- a/packages/agent/test/proxy.test.ts +++ b/packages/agent/test/proxy.test.ts @@ -30,34 +30,6 @@ afterEach(() => { }); describe("streamProxy", () => { - it("serializes request identity for the proxy server", async () => { - let requestBody: { options?: { requestIdentity?: unknown } } | undefined; - vi.stubGlobal( - "fetch", - vi.fn(async (_input: string | URL | Request, init?: RequestInit) => { - requestBody = JSON.parse(String(init?.body)); - return new Response(`data: ${JSON.stringify({ type: "done", reason: "stop", usage })}\n\n`, { - status: 200, - }); - }), - ); - const requestIdentity = { - sessionId: "session", - threadId: "thread", - turnId: "turn", - requestKind: "turn" as const, - startedAt: 123, - }; - - await streamProxy(model, normalizeContext({ systemPrompt: "", messages: [] }), { - authToken: "test-token", - proxyUrl: "https://proxy.example.com", - requestIdentity, - }).result(); - - expect(requestBody?.options?.requestIdentity).toEqual(requestIdentity); - }); - it("preserves tool-call metadata received only on toolcall_end", async () => { const proxyEvents: ProxyAssistantMessageEvent[] = [ { type: "start" }, diff --git a/packages/agent/test/request-identity.test.ts b/packages/agent/test/request-identity.test.ts index d07ccf36f..1e13e00f4 100644 --- a/packages/agent/test/request-identity.test.ts +++ b/packages/agent/test/request-identity.test.ts @@ -1,15 +1,11 @@ -import { - type AgentRequestIdentity, - fauxAssistantMessage, - fauxToolCall, - registerFauxProvider, - streamSimple, -} from "@earendil-works/pi-ai/compat"; +import { fauxAssistantMessage, fauxToolCall, registerFauxProvider, streamSimple } from "@earendil-works/pi-ai/compat"; import { afterEach, describe, expect, it } from "vitest"; import { Agent } from "../src/agent.ts"; +import { getAgentRequestIdentity } from "../src/request-metadata.ts"; import { calculateTool } from "./utils/calculate.ts"; const registrations: Array> = []; +type RequestIdentity = NonNullable>; afterEach(() => { while (registrations.length > 0) registrations.pop()?.unregister(); @@ -20,20 +16,20 @@ describe("Agent request identity", () => { it("shares one identity across tools and rotates it for a follow-up", async () => { const faux = registerFauxProvider(); registrations.push(faux); - const identities: AgentRequestIdentity[] = []; + const identities: RequestIdentity[] = []; let agent: Agent; faux.setResponses([ (_context, options) => { - identities.push(options?.requestIdentity as AgentRequestIdentity); + identities.push(getAgentRequestIdentity(options?.metadata)!); return fauxAssistantMessage(fauxToolCall("calculate", { expression: "2 + 2" })); }, (_context, options) => { - identities.push(options?.requestIdentity as AgentRequestIdentity); + identities.push(getAgentRequestIdentity(options?.metadata)!); agent.followUp({ role: "user", content: "next", timestamp: Date.now() }); return fauxAssistantMessage("4"); }, (_context, options) => { - identities.push(options?.requestIdentity as AgentRequestIdentity); + identities.push(getAgentRequestIdentity(options?.metadata)!); return fauxAssistantMessage("done"); }, ]); @@ -58,16 +54,16 @@ describe("Agent request identity", () => { it("keeps steering within the active turn identity", async () => { const faux = registerFauxProvider(); registrations.push(faux); - const identities: AgentRequestIdentity[] = []; + const identities: RequestIdentity[] = []; let agent: Agent; faux.setResponses([ (_context, options) => { - identities.push(options?.requestIdentity as AgentRequestIdentity); + identities.push(getAgentRequestIdentity(options?.metadata)!); agent.steer({ role: "user", content: "redirect", timestamp: Date.now() }); return fauxAssistantMessage("first"); }, (_context, options) => { - identities.push(options?.requestIdentity as AgentRequestIdentity); + identities.push(getAgentRequestIdentity(options?.metadata)!); return fauxAssistantMessage("redirected"); }, ]); @@ -83,30 +79,51 @@ describe("Agent request identity", () => { }); // #9481 - it("preserves session and thread identity when recreating an agent", () => { - const first = new Agent({ sessionId: "session", streamFn: streamSimple }); - const second = new Agent({ sessionId: "session", streamFn: streamSimple }); + it("preserves session and thread identity when recreating an agent", async () => { + const faux = registerFauxProvider(); + registrations.push(faux); + const identities: RequestIdentity[] = []; + faux.setResponses([ + (_context, options) => { + identities.push(getAgentRequestIdentity(options?.metadata)!); + return fauxAssistantMessage("done"); + }, + (_context, options) => { + identities.push(getAgentRequestIdentity(options?.metadata)!); + return fauxAssistantMessage("done"); + }, + ]); + const first = new Agent({ + sessionId: "session", + streamFn: streamSimple, + initialState: { model: faux.getModel() }, + }); + const second = new Agent({ + sessionId: "session", + streamFn: streamSimple, + initialState: { model: faux.getModel() }, + }); - const firstIdentity = first.createRequestIdentity(); - const secondIdentity = second.createRequestIdentity(); + await first.prompt("first"); + await second.prompt("second"); - expect(firstIdentity).toMatchObject({ sessionId: "session", threadId: "session" }); - expect(secondIdentity).toMatchObject({ sessionId: "session", threadId: "session" }); - expect(secondIdentity.turnId).not.toBe(firstIdentity.turnId); + expect(identities[0]).toMatchObject({ sessionId: "session", threadId: "session" }); + expect(identities[1]).toMatchObject({ sessionId: "session", threadId: "session" }); + expect(identities[1].turnId).not.toBe(identities[0].turnId); }); // #9481 it("preserves identity when continuing a failed request", async () => { const faux = registerFauxProvider(); registrations.push(faux); - const identities: AgentRequestIdentity[] = []; + const identities: RequestIdentity[] = []; faux.setResponses([ (_context, options) => { - identities.push(options?.requestIdentity as AgentRequestIdentity); + identities.push(getAgentRequestIdentity(options?.metadata)!); return fauxAssistantMessage("", { stopReason: "error", errorMessage: "retry" }); }, (_context, options) => { - identities.push(options?.requestIdentity as AgentRequestIdentity); + identities.push(getAgentRequestIdentity(options?.metadata)!); return fauxAssistantMessage("done"); }, ]); diff --git a/packages/ai/src/api/openai-codex-responses.ts b/packages/ai/src/api/openai-codex-responses.ts index e3d37ca00..b8d399a34 100644 --- a/packages/ai/src/api/openai-codex-responses.ts +++ b/packages/ai/src/api/openai-codex-responses.ts @@ -9,7 +9,6 @@ import type { import { clampThinkingLevel } from "../models.ts"; import { registerSessionResourceCleanup } from "../session-resources.ts"; import type { - AgentRequestIdentity, Api, AssistantMessage, Model, @@ -64,6 +63,7 @@ const WEBSOCKET_MESSAGE_TOO_BIG_CLOSE_CODE = 1009; const WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE = "websocket_connection_limit_reached"; const PREVIOUS_RESPONSE_NOT_FOUND_CODE = "previous_response_not_found"; const TURN_STATE_HEADER = "x-codex-turn-state"; +const REQUEST_IDENTITY_METADATA_KEY = "pi.requestIdentity"; const WEBSOCKET_REQUEST_SCOPED_HEADERS = new Set([ "session-id", "thread-id", @@ -85,7 +85,7 @@ interface CodexTurnState { const codexTurnStates = new Map(); function getCodexTurnState( - identity: AgentRequestIdentity | undefined, + identity: OpenAICodexRequestIdentity | undefined, accountId: string, url: string, ): CodexTurnState | undefined { @@ -124,7 +124,20 @@ const CODEX_RESPONSE_STATUSES = new Set([ // Types // ============================================================================ +export interface OpenAICodexRequestIdentity { + sessionId: string; + threadId: string; + turnId: string; + requestKind: "turn" | "compaction"; + startedAt: number; + windowId?: string; + windowNumber?: number; + contextWindowId?: string; +} + export interface OpenAICodexResponsesOptions extends StreamOptions { + /** Stable identity shared by all Codex calls belonging to one logical agent request. */ + requestIdentity?: OpenAICodexRequestIdentity; reasoningEffort?: "none" | "minimal" | "low" | "medium" | "high" | "xhigh" | "max"; reasoningSummary?: "auto" | "concise" | "detailed" | "off" | "on" | null; serviceTier?: ResponseCreateParamsStreaming["service_tier"]; @@ -132,6 +145,11 @@ export interface OpenAICodexResponsesOptions extends StreamOptions { toolChoice?: "auto" | "none" | "required"; } +export interface OpenAICodexSimpleStreamOptions extends SimpleStreamOptions { + /** Stable identity shared by all Codex calls belonging to one logical agent request. */ + requestIdentity?: OpenAICodexRequestIdentity; +} + type CodexResponseStatus = "completed" | "incomplete" | "failed" | "cancelled" | "queued" | "in_progress"; interface RequestBody { @@ -156,7 +174,29 @@ interface RequestBody { type SuccessfulAssistantMessage = AssistantMessage & { stopReason: "stop" | "length" | "toolUse" }; -function buildCodexRequestMetadata(identity: AgentRequestIdentity | undefined): +function getMetadataRequestIdentity( + metadata: Record | undefined, +): OpenAICodexRequestIdentity | undefined { + const value = metadata?.[REQUEST_IDENTITY_METADATA_KEY]; + if (!value || typeof value !== "object") return undefined; + const identity = value as Partial; + if ( + typeof identity.sessionId !== "string" || + typeof identity.threadId !== "string" || + typeof identity.turnId !== "string" || + (identity.requestKind !== "turn" && identity.requestKind !== "compaction") || + typeof identity.startedAt !== "number" || + !Number.isFinite(identity.startedAt) || + (identity.windowId !== undefined && typeof identity.windowId !== "string") || + (identity.windowNumber !== undefined && typeof identity.windowNumber !== "number") || + (identity.contextWindowId !== undefined && typeof identity.contextWindowId !== "string") + ) { + return undefined; + } + return identity as OpenAICodexRequestIdentity; +} + +function buildCodexRequestMetadata(identity: OpenAICodexRequestIdentity | undefined): | { clientMetadata: Record; headers: Record; @@ -355,7 +395,8 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons const accountId = extractAccountId(apiKey); const codexUrl = resolveCodexUrl(model.baseUrl); const websocketUrl = resolveCodexWebSocketUrl(model.baseUrl); - const turnState = getCodexTurnState(options?.requestIdentity, accountId, codexUrl); + const requestIdentity = options?.requestIdentity ?? getMetadataRequestIdentity(options?.metadata); + const turnState = getCodexTurnState(requestIdentity, accountId, codexUrl); const grammarToolInputProperties = createGrammarToolInputProperties( getDeclaredTools(normalizedContext.messages), model.compat?.supportsOpenAIGrammarTools ?? false, @@ -366,8 +407,8 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons ? { key: JSON.stringify([ cacheSessionId, - options?.requestIdentity?.sessionId ?? cacheSessionId, - options?.requestIdentity?.threadId ?? cacheSessionId, + requestIdentity?.sessionId ?? cacheSessionId, + requestIdentity?.threadId ?? cacheSessionId, ]), resourceSessionId: cacheSessionId, } @@ -375,7 +416,7 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons const transport = options?.transport || "auto"; const websocketDisabledForSession = transport !== "sse" && isWebSocketSseFallbackActive(cacheSessionId); let body = buildRequestBody(model, normalizedContext, options, codexSessionId, grammarToolInputProperties); - const requestMetadata = buildCodexRequestMetadata(options?.requestIdentity); + const requestMetadata = buildCodexRequestMetadata(requestIdentity); if (requestMetadata) { body.client_metadata = { ...body.client_metadata, ...requestMetadata.clientMetadata }; } @@ -390,14 +431,14 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons client_metadata: { ...body.client_metadata, ...requestMetadata.clientMetadata }, }; } - const websocketRequestId = options?.requestIdentity?.threadId ?? codexSessionId ?? uuidv7(); + const websocketRequestId = requestIdentity?.threadId ?? codexSessionId ?? uuidv7(); const sseHeaders = buildSSEHeaders( model.headers, options?.headers, accountId, apiKey, codexSessionId, - options?.requestIdentity, + requestIdentity, ); const websocketHeaders = buildWebSocketHeaders( model.headers, @@ -405,7 +446,7 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons accountId, apiKey, websocketRequestId, - options?.requestIdentity, + requestIdentity, ); const websocketOwnerKey = getWebSocketOwnerKey(websocketUrl, websocketHeaders); const bodyJson = JSON.stringify(body); @@ -621,10 +662,10 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons return stream; }; -export const streamSimple: StreamFunction<"openai-codex-responses", SimpleStreamOptions> = ( +export const streamSimple: StreamFunction<"openai-codex-responses", OpenAICodexSimpleStreamOptions> = ( model: Model<"openai-codex-responses">, context: TranscriptContext, - options?: SimpleStreamOptions, + options?: OpenAICodexSimpleStreamOptions, ): AssistantMessageEventStream => { const apiKey = options?.apiKey; if (!apiKey) { @@ -633,6 +674,7 @@ export const streamSimple: StreamFunction<"openai-codex-responses", SimpleStream const base = { ...buildBaseOptions(model, context, options, apiKey), + requestIdentity: options?.requestIdentity, toolChoice: options?.toolChoice, } satisfies OpenAICodexResponsesOptions; const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined; @@ -1794,7 +1836,7 @@ function buildBaseCodexHeaders( additionalHeaders: ProviderHeaders | undefined, accountId: string, token: string, - requestIdentity?: AgentRequestIdentity, + requestIdentity?: OpenAICodexRequestIdentity, ): Headers { const headers = new Headers(initHeaders); for (const [key, value] of Object.entries(additionalHeaders || {})) { @@ -1823,7 +1865,7 @@ function buildSSEHeaders( accountId: string, token: string, sessionId?: string, - requestIdentity?: AgentRequestIdentity, + requestIdentity?: OpenAICodexRequestIdentity, ): Headers { const headers = buildBaseCodexHeaders(initHeaders, additionalHeaders, accountId, token, requestIdentity); headers.set("OpenAI-Beta", "responses=experimental"); @@ -1851,7 +1893,7 @@ function buildWebSocketHeaders( accountId: string, token: string, requestId: string, - requestIdentity?: AgentRequestIdentity, + requestIdentity?: OpenAICodexRequestIdentity, ): Headers { const headers = buildBaseCodexHeaders(initHeaders, additionalHeaders, accountId, token, requestIdentity); headers.delete("accept"); diff --git a/packages/ai/src/api/simple-options.ts b/packages/ai/src/api/simple-options.ts index 1a9e7b58c..63e34f3c7 100644 --- a/packages/ai/src/api/simple-options.ts +++ b/packages/ai/src/api/simple-options.ts @@ -39,7 +39,6 @@ export function buildBaseOptions( transport: options?.transport, cacheRetention: options?.cacheRetention, sessionId: options?.sessionId, - requestIdentity: options?.requestIdentity, headers: options?.headers, onPayload: options?.onPayload, onResponse: options?.onResponse, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 056280b9f..5fb508a88 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -14,7 +14,12 @@ export type { GoogleApiThinkingLevel, ResolvedGoogleThinkingLevel } from "./api/ export type { GoogleVertexOptions } from "./api/google-vertex.ts"; export * from "./api/lazy.ts"; export type { MistralOptions } from "./api/mistral-conversations.ts"; -export type { OpenAICodexResponsesOptions, OpenAICodexWebSocketDebugStats } from "./api/openai-codex-responses.ts"; +export type { + OpenAICodexRequestIdentity, + OpenAICodexResponsesOptions, + OpenAICodexSimpleStreamOptions, + OpenAICodexWebSocketDebugStats, +} from "./api/openai-codex-responses.ts"; export type { OpenAICompletionsOptions } from "./api/openai-completions.ts"; export type { OpenAIResponsesOptions } from "./api/openai-responses.ts"; export type { PiMessagesEvent, PiMessagesOptions, PiMessagesRewriteImpact } from "./api/pi-messages.ts"; diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 3c5e08228..7646a9c05 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -183,18 +183,6 @@ export interface ProviderRequestOptions> { maxRetryDelayMs?: number; } -/** Stable identity shared by all provider calls belonging to one agent turn. */ -export interface AgentRequestIdentity { - sessionId: string; - threadId: string; - turnId: string; - requestKind: "turn" | "compaction"; - startedAt: number; - windowId?: string; - windowNumber?: number; - contextWindowId?: string; -} - export interface StreamOptions extends ProviderRequestOptions> { /** * Optional callback invoked after an HTTP response is received and before @@ -227,8 +215,6 @@ export interface StreamOptions extends ProviderRequestOptions> { * session-aware features. Ignored by providers that don't support it. */ sessionId?: string; - /** Logical agent-turn identity, independent of provider cache identity. */ - requestIdentity?: AgentRequestIdentity; /** * WebSocket connect timeout in milliseconds for providers that support * WebSocket transports. This covers the connection/open handshake only; diff --git a/packages/ai/test/openai-codex-attribution.test.ts b/packages/ai/test/openai-codex-attribution.test.ts index 4e43628ce..98323d1bd 100644 --- a/packages/ai/test/openai-codex-attribution.test.ts +++ b/packages/ai/test/openai-codex-attribution.test.ts @@ -1,8 +1,8 @@ import { zstdDecompressSync } from "node:zlib"; import { afterEach, describe, expect, it, vi } from "vitest"; -import { streamSimple } from "../src/api/openai-codex-responses.ts"; +import { type OpenAICodexRequestIdentity, streamSimple } from "../src/api/openai-codex-responses.ts"; import { cleanupSessionResources } from "../src/session-resources.ts"; -import type { AgentRequestIdentity, Model } from "../src/types.ts"; +import type { Model } from "../src/types.ts"; import { normalizeContext } from "../src/utils/transcript.ts"; const model: Model<"openai-codex-responses"> = { @@ -18,7 +18,7 @@ const model: Model<"openai-codex-responses"> = { maxTokens: 128000, }; -const identity: AgentRequestIdentity = { +const identity: OpenAICodexRequestIdentity = { sessionId: "session-1", threadId: "thread-1", turnId: "turn-1", @@ -64,7 +64,7 @@ describe("Codex request attribution", () => { ); const context = normalizeContext({ messages: [{ role: "user", content: "hello", timestamp: 1 }] }); const request = async ( - requestIdentity: AgentRequestIdentity, + requestIdentity: OpenAICodexRequestIdentity, requestModel: Model<"openai-codex-responses"> = model, ) => streamSimple(requestModel, context, { @@ -132,7 +132,7 @@ describe("Codex request attribution", () => { apiKey: token(), transport: "sse", cacheRetention: "none", - requestIdentity: identity, + metadata: { "pi.requestIdentity": identity }, headers: { originator: "spoofed", "session-id": "spoofed", @@ -175,7 +175,7 @@ describe("Codex request attribution", () => { }), ); const context = normalizeContext({ messages: [] }); - const request = async (requestIdentity: AgentRequestIdentity) => + const request = async (requestIdentity: OpenAICodexRequestIdentity) => streamSimple(model, context, { apiKey: token(), transport: "sse", @@ -244,7 +244,7 @@ describe("Codex request attribution", () => { vi.stubGlobal("WebSocket", MockWebSocket); const context = normalizeContext({ messages: [{ role: "user", content: "hello", timestamp: 1 }] }); const request = async ( - requestIdentity: AgentRequestIdentity, + requestIdentity: OpenAICodexRequestIdentity, requestModel: Model<"openai-codex-responses"> = model, ) => streamSimple(requestModel, context, { diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index a74f6cd02..338d0e0d0 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -25,9 +25,8 @@ import type { PrepareNextTurnContext, ThinkingLevel, } from "@earendil-works/pi-agent-core"; -import { contentText, getCurrentSystemMessage, retryDelayMs, uuidv7 } from "@earendil-works/pi-ai"; +import { contentText, getCurrentSystemMessage, retryDelayMs } from "@earendil-works/pi-ai"; import type { - AgentRequestIdentity, AssistantMessage, AuthResult, ImageContent, @@ -2357,16 +2356,6 @@ export class AgentSession { // Compaction // ========================================================================= - private _createCompactionRequestIdentity(): AgentRequestIdentity { - return { - sessionId: this.sessionId, - threadId: this.sessionId, - turnId: uuidv7(), - requestKind: "compaction", - startedAt: Date.now(), - }; - } - /** Generate Pi's built-in compaction summary for manual and automatic compaction. */ private async _runDefaultCompaction( preparation: CompactionPreparation, @@ -2377,7 +2366,6 @@ export class AgentSession { signal: AbortSignal, env: Record | undefined, reason: "manual" | "threshold" | "overflow", - requestIdentity: AgentRequestIdentity, ): Promise { return compact( preparation, @@ -2392,7 +2380,6 @@ export class AgentSession { this.settingsManager.getRetrySettings(), this._summarizationRetryCallbacks({ source: "compaction", reason }), undefined, // sessionId - requestIdentity, ); } @@ -2497,7 +2484,6 @@ export class AgentSession { this._compactionAbortController.signal, env, "manual", - this._createCompactionRequestIdentity(), ); summary = result.summary; firstKeptEntryId = result.firstKeptEntryId; @@ -2840,7 +2826,6 @@ export class AgentSession { abortController.signal, env, reason, - this._createCompactionRequestIdentity(), ); summary = compactResult.summary; firstKeptEntryId = compactResult.firstKeptEntryId; @@ -3702,7 +3687,6 @@ export class AgentSession { streamFn: this.agent.streamFunction, retry: this.settingsManager.getRetrySettings(), callbacks: this._summarizationRetryCallbacks({ source: "branchSummary" }), - requestIdentity: this._createCompactionRequestIdentity(), }); if (result.aborted) { return { cancelled: true, aborted: true }; diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts index 81064fae4..8352cfa19 100644 --- a/packages/coding-agent/src/core/compaction/branch-summarization.ts +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -8,7 +8,7 @@ import type { AgentMessage, StreamFn } from "@earendil-works/pi-agent-core"; import type { RetryCallbacks, RetryPolicy } from "@earendil-works/pi-ai"; import { contentText, normalizeContext } from "@earendil-works/pi-ai"; -import type { AgentRequestIdentity, Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat"; +import type { Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat"; import { convertToLlm, createBranchSummaryMessage, @@ -17,6 +17,7 @@ import { } from "../messages.ts"; import type { ReadonlySessionManager, SessionEntry } from "../session-manager.ts"; import { completeSummarization, estimateTokens, getSummarizationFailure } from "./compaction.ts"; +import { withCompactionRequestMetadata } from "./request-metadata.ts"; import { computeFileLists, createFileOps, @@ -87,8 +88,6 @@ export interface GenerateBranchSummaryOptions { retry?: RetryPolicy; /** Optional callbacks for retry reporting (e.g. TUI retry indicators). */ callbacks?: RetryCallbacks; - /** Logical identity shared by the summary request and all of its retries. */ - requestIdentity?: AgentRequestIdentity; } // ============================================================================ @@ -308,7 +307,6 @@ export async function generateBranchSummary( streamFn, retry, callbacks, - requestIdentity, } = options; // Token budget = context window minus reserved space for prompt + response @@ -352,8 +350,15 @@ export async function generateBranchSummary( // without running through agent state/events. Retried via completeSummarization // so transient stream drops reuse the configured retry policy. const context = normalizeContext({ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }); - const requestOptions: SimpleStreamOptions = { apiKey, headers, env, signal, maxTokens, requestIdentity }; - const response = await completeSummarization(model, context, requestOptions, streamFn, retry, callbacks); + const requestOptions: SimpleStreamOptions = { apiKey, headers, env, signal, maxTokens }; + const response = await completeSummarization( + model, + context, + requestOptions, + withCompactionRequestMetadata(streamFn), + retry, + callbacks, + ); // Check if aborted or errored if (response.stopReason === "aborted") { diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 5944b83d3..6d9306f3e 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -16,7 +16,6 @@ import { uuidv7, } from "@earendil-works/pi-ai"; import type { - AgentRequestIdentity, AssistantMessage, Model, SimpleStreamOptions, @@ -34,6 +33,7 @@ import { type SessionProjection, sessionEntryToContextMessages, } from "../session-manager.ts"; +import { withCompactionRequestMetadata } from "./request-metadata.ts"; import { computeFileLists, createFileOps, @@ -624,9 +624,8 @@ function createSummarizationOptions( signal: AbortSignal | undefined, thinkingLevel: ThinkingLevel | undefined, sessionId: string | undefined, - requestIdentity: AgentRequestIdentity | undefined, ): SimpleStreamOptions { - const options: SimpleStreamOptions = { maxTokens, signal, apiKey, headers, env, sessionId, requestIdentity }; + const options: SimpleStreamOptions = { maxTokens, signal, apiKey, headers, env, sessionId }; if (model.reasoning && thinkingLevel && thinkingLevel !== "off") { options.reasoning = thinkingLevel; } @@ -681,7 +680,6 @@ export async function generateSummary( retry?: RetryPolicy, callbacks?: RetryCallbacks, sessionId?: string, - requestIdentity?: AgentRequestIdentity, ): Promise { return ( await generateSummaryWithUsage( @@ -699,7 +697,6 @@ export async function generateSummary( retry, callbacks, sessionId, - requestIdentity, ) ).text; } @@ -734,7 +731,6 @@ export async function generateSummaryWithUsage( retry?: RetryPolicy, callbacks?: RetryCallbacks, sessionId?: string, - requestIdentity?: AgentRequestIdentity, ): Promise<{ text: string; usage: Usage }> { const maxTokens = Math.min( Math.floor(0.8 * reserveTokens), @@ -768,7 +764,6 @@ export async function generateSummaryWithUsage( signal, thinkingLevel, sessionId, - requestIdentity, ); const response = await completeSummarization( @@ -989,7 +984,6 @@ Be concise. Focus on what's needed to understand the kept suffix.`; * @param preparation - Pre-calculated preparation from prepareCompaction() * @param customInstructions - Optional custom focus for the summary * @param sessionId - Optional routing session ID forwarded without enabling prompt caching - * @param requestIdentity - Logical identity shared by every summary call and retry in this compaction */ export async function compact( preparation: CompactionPreparation, @@ -1004,8 +998,8 @@ export async function compact( retry?: RetryPolicy, callbacks?: RetryCallbacks, sessionId?: string, - requestIdentity?: AgentRequestIdentity, ): Promise { + const attributedStreamFn = withCompactionRequestMetadata(streamFn); const { firstKeptEntryId, messagesToSummarize, @@ -1035,12 +1029,11 @@ export async function compact( customInstructions, previousSummary, thinkingLevel, - streamFn, + attributedStreamFn, env, retry, callbacks, sessionId, - requestIdentity, ); historyText = historyResult.text; historyUsage = historyResult.usage; @@ -1054,11 +1047,10 @@ export async function compact( env, signal, thinkingLevel, - streamFn, + attributedStreamFn, retry, callbacks, sessionId, - requestIdentity, ); // Merge into single summary summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.text}`; @@ -1075,12 +1067,11 @@ export async function compact( customInstructions, previousSummary, thinkingLevel, - streamFn, + attributedStreamFn, env, retry, callbacks, sessionId, - requestIdentity, ); summary = result.text; summaryUsage = result.usage; @@ -1119,7 +1110,6 @@ async function generateTurnPrefixSummary( retry?: RetryPolicy, callbacks?: RetryCallbacks, sessionId?: string, - requestIdentity?: AgentRequestIdentity, ): Promise<{ text: string; usage: Usage }> { const maxTokens = Math.min( Math.floor(0.5 * reserveTokens), @@ -1132,17 +1122,7 @@ async function generateTurnPrefixSummary( const response = await completeSummarization( model, buildSummarizationContext(promptText), - createSummarizationOptions( - model, - maxTokens, - apiKey, - headers, - env, - signal, - thinkingLevel, - sessionId, - requestIdentity, - ), + createSummarizationOptions(model, maxTokens, apiKey, headers, env, signal, thinkingLevel, sessionId), streamFn, retry, callbacks, diff --git a/packages/coding-agent/src/core/compaction/request-metadata.ts b/packages/coding-agent/src/core/compaction/request-metadata.ts new file mode 100644 index 000000000..bb0f6dcfc --- /dev/null +++ b/packages/coding-agent/src/core/compaction/request-metadata.ts @@ -0,0 +1,46 @@ +import type { StreamFn } from "@earendil-works/pi-agent-core"; +import { streamSimple, uuidv7 } from "@earendil-works/pi-ai/compat"; + +export const REQUEST_IDENTITY_METADATA_KEY = "pi.requestIdentity"; + +export interface InternalRequestIdentity { + sessionId: string; + threadId: string; + turnId: string; + requestKind: "turn" | "compaction"; + startedAt: number; + windowId?: string; + windowNumber?: number; + contextWindowId?: string; +} + +export function getRequestIdentityMetadata( + metadata: Record | undefined, +): InternalRequestIdentity | undefined { + const value = metadata?.[REQUEST_IDENTITY_METADATA_KEY]; + if (!value || typeof value !== "object") return undefined; + return value as InternalRequestIdentity; +} + +export function setRequestIdentityMetadata( + metadata: Record | undefined, + identity: InternalRequestIdentity, +): Record { + return { ...metadata, [REQUEST_IDENTITY_METADATA_KEY]: identity }; +} + +export function withCompactionRequestMetadata(streamFn: StreamFn = streamSimple): StreamFn { + const sessionId = uuidv7(); + const identity = { + sessionId, + threadId: sessionId, + turnId: uuidv7(), + requestKind: "compaction", + startedAt: Date.now(), + } satisfies InternalRequestIdentity; + return (model, context, options) => + streamFn(model, context, { + ...options, + metadata: setRequestIdentityMetadata(options?.metadata, identity), + }); +} diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 7dee6b6ae..9de48b04c 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -7,6 +7,7 @@ import { resolvePath } from "../utils/paths.ts"; import { AgentSession } from "./agent-session.ts"; import { formatNoModelsAvailableMessage } from "./auth-guidance.ts"; import { CacheWarmer } from "./cache-warmer.ts"; +import { getRequestIdentityMetadata, setRequestIdentityMetadata } from "./compaction/request-metadata.ts"; import { DEFAULT_THINKING_LEVEL } from "./defaults.ts"; import type { ExtensionRunner, LoadExtensionsResult, SessionStartEvent, ToolDefinition } from "./extensions/index.ts"; import { convertToLlm } from "./messages.ts"; @@ -316,20 +317,24 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} const httpIdleTimeoutMs = settingsManager.getHttpIdleTimeoutMs(); const effectiveTimeoutMs = httpIdleTimeoutMs === 0 ? 2147483647 : httpIdleTimeoutMs; const headerRunner = extensionRunnerRef.current; - const compactions = options.requestIdentity + const requestIdentity = getRequestIdentityMetadata(options.metadata); + const compactions = requestIdentity ? sessionManager.getBranch().filter((entry) => entry.type === "compaction") : []; - const requestIdentity = options.requestIdentity - ? { - ...options.requestIdentity, - windowId: `${options.requestIdentity.threadId}:${compactions.length}`, + const sessionId = sessionManager.getSessionId(); + const metadata = requestIdentity + ? setRequestIdentityMetadata(options.metadata, { + ...requestIdentity, + sessionId, + threadId: sessionId, + windowId: `${sessionId}:${compactions.length}`, windowNumber: compactions.length, - contextWindowId: compactions.at(-1)?.id ?? options.requestIdentity.threadId, - } - : undefined; + contextWindowId: compactions.at(-1)?.id ?? sessionId, + }) + : options.metadata; return { ...options, - requestIdentity, + metadata, timeoutMs: options.timeoutMs ?? providerRetrySettings.timeoutMs ?? effectiveTimeoutMs, websocketConnectTimeoutMs: options.websocketConnectTimeoutMs ?? settingsManager.getWebSocketConnectTimeoutMs(), maxRetries: options.maxRetries ?? providerRetrySettings.maxRetries, diff --git a/packages/coding-agent/test/branch-summarization.test.ts b/packages/coding-agent/test/branch-summarization.test.ts index ce60ce411..5d98cae2d 100644 --- a/packages/coding-agent/test/branch-summarization.test.ts +++ b/packages/coding-agent/test/branch-summarization.test.ts @@ -8,6 +8,7 @@ import { } from "@earendil-works/pi-ai"; import { describe, expect, it } from "vitest"; import { generateBranchSummary } from "../src/core/compaction/index.ts"; +import { getRequestIdentityMetadata } from "../src/core/compaction/request-metadata.ts"; import type { SessionEntry } from "../src/core/session-manager.ts"; const model: Model<"anthropic-messages"> = { @@ -55,23 +56,15 @@ describe("branch summarization", () => { return stream; }; - const requestIdentity = { - sessionId: "session", - threadId: "thread", - turnId: "summary", - requestKind: "compaction" as const, - startedAt: 123, - }; await generateBranchSummary(entries, { model, signal: new AbortController().signal, streamFn, - requestIdentity, }); expect(requestOptions?.maxTokens).toBe(4096); expect(requestOptions?.toolChoice).toBeUndefined(); - expect(requestOptions?.requestIdentity).toBe(requestIdentity); + expect(getRequestIdentityMetadata(requestOptions?.metadata)?.requestKind).toBe("compaction"); }); it("clamps the branch summary output cap to the model limit", async () => { diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index 4e0099a78..c9b93b300 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/packages/coding-agent/test/compaction.test.ts @@ -1,5 +1,5 @@ import type { AgentMessage, StreamFn } from "@earendil-works/pi-agent-core"; -import type { AgentRequestIdentity, AssistantMessage, Usage } from "@earendil-works/pi-ai/compat"; +import type { AssistantMessage, Usage } from "@earendil-works/pi-ai/compat"; import { createAssistantMessageEventStream, fauxAssistantMessage, getModel } from "@earendil-works/pi-ai/compat"; import { readFileSync } from "fs"; import { join } from "path"; @@ -15,6 +15,7 @@ import { prepareCompaction, shouldCompact, } from "../src/core/compaction/index.ts"; +import { getRequestIdentityMetadata, type InternalRequestIdentity } from "../src/core/compaction/request-metadata.ts"; import { buildSessionContext, type CompactionEntry, @@ -533,9 +534,9 @@ describe("prepareCompaction", () => { expect(preparation?.messagesToSummarize.length).toBeGreaterThan(0); expect(preparation?.turnPrefixMessages.length).toBeGreaterThan(0); - const identities: Array = []; + const identities: Array = []; const streamFn: StreamFn = (requestModel, _context, options) => { - identities.push(options?.requestIdentity); + identities.push(getRequestIdentityMetadata(options?.metadata)); const stream = createAssistantMessageEventStream(); queueMicrotask(() => { const message = { @@ -548,13 +549,6 @@ describe("prepareCompaction", () => { }); return stream; }; - const requestIdentity: AgentRequestIdentity = { - sessionId: "session", - threadId: "thread", - turnId: "compaction", - requestKind: "compaction", - startedAt: 123, - }; await compact( preparation!, getModel("anthropic", "claude-sonnet-4-5")!, @@ -568,10 +562,11 @@ describe("prepareCompaction", () => { undefined, undefined, undefined, - requestIdentity, ); - expect(identities).toEqual([requestIdentity, requestIdentity]); + expect(identities[0]).toBeDefined(); + expect(identities[0]?.requestKind).toBe("compaction"); + expect(identities[1]).toBe(identities[0]); }); }); diff --git a/packages/coding-agent/test/sdk-stream-options.test.ts b/packages/coding-agent/test/sdk-stream-options.test.ts index 601153bae..704b0e9f1 100644 --- a/packages/coding-agent/test/sdk-stream-options.test.ts +++ b/packages/coding-agent/test/sdk-stream-options.test.ts @@ -11,6 +11,7 @@ import { } from "@earendil-works/pi-ai"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; +import { getRequestIdentityMetadata, setRequestIdentityMetadata } from "../src/core/compaction/request-metadata.ts"; import { createAgentSession } from "../src/core/sdk.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { type Settings, SettingsManager } from "../src/core/settings-manager.ts"; @@ -203,7 +204,28 @@ describe("createAgentSession stream options", () => { it("does not synthesize a Codex identity for side requests", async () => { const options = await captureStreamOptions("openai-codex-responses", {}); - expect(options?.requestIdentity).toBeUndefined(); + expect(getRequestIdentityMetadata(options?.metadata)).toBeUndefined(); + }); + + it("adds session window data to Codex request identity metadata", async () => { + const options = await captureStreamOptions( + "openai-codex-responses", + {}, + { + metadata: setRequestIdentityMetadata(undefined, { + sessionId: "placeholder", + threadId: "placeholder", + turnId: "turn", + requestKind: "turn", + startedAt: 123, + }), + }, + ); + + const identity = getRequestIdentityMetadata(options?.metadata); + expect(identity).toMatchObject({ turnId: "turn", windowNumber: 0, contextWindowId: identity?.threadId }); + expect(identity?.sessionId).toBe(identity?.threadId); + expect(identity?.windowId).toBe(`${identity?.threadId}:0`); }); it("forwards httpIdleTimeoutMs as timeoutMs for OpenAI Codex", async () => {