From 0b3c45c948ad145b78644fd48554fc341f233307 Mon Sep 17 00:00:00 2001 From: Shoubhit Dash Date: Fri, 4 Sep 2026 19:30:02 +0530 Subject: [PATCH] feat(ai): add streamed compaction triggers (#47082) --- packages/ai/README.md | 39 ++- .../protocols/open-responses-continuation.ts | 7 + packages/ai/src/protocols/open-responses.ts | 10 +- packages/ai/src/protocols/openai-responses.ts | 46 +++- .../protocols/utils/responses-checkpoint.ts | 122 +++++++++ packages/ai/src/provider-package.ts | 4 +- packages/ai/src/providers/azure.ts | 7 +- packages/ai/src/providers/openai.ts | 13 +- packages/ai/src/providers/xai.ts | 13 +- packages/ai/src/route/client.ts | 153 ++++++++--- packages/ai/src/route/index.ts | 6 + packages/ai/src/schema/events.ts | 15 ++ packages/ai/src/schema/options.ts | 16 +- packages/ai/src/testing.ts | 50 +++- .../test/provider-options/checkpoint.types.ts | 129 ++++++++++ packages/ai/test/provider/checkpoint.test.ts | 243 ++++++++++++++++++ packages/ai/test/testing.test.ts | 26 ++ packages/ai/tsconfig.types.json | 2 +- packages/core/test/lib/websocket-server.ts | 2 + .../test/session-checkpoint-transport.test.ts | 172 +++++++++++++ 20 files changed, 993 insertions(+), 82 deletions(-) create mode 100644 packages/ai/src/protocols/utils/responses-checkpoint.ts create mode 100644 packages/ai/test/provider-options/checkpoint.types.ts create mode 100644 packages/ai/test/provider/checkpoint.test.ts create mode 100644 packages/core/test/session-checkpoint-transport.test.ts diff --git a/packages/ai/README.md b/packages/ai/README.md index 67c0c2e95b7..1ad0ab4610f 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -247,6 +247,8 @@ it does not repair or truncate them. For explicit compaction, script a `CompactionResponse` through `push`, `always`, or `serve`. Its `replacement` contains the next context window, including retained user messages. The client returns that result and usage directly, with the same lazy request recording and gates. Generation and compaction reject fixtures for the wrong operation instead of converting between response shapes. +For `compact(request, { mechanism: "trigger" })`, script a `CompactionCheckpointResponse` instead. It carries `checkpoint`, `responseID`, and optional `usage`. Endpoint and trigger calls reject each other's fixtures; both share the same queue, gates, lazy recording, and fallback controls. + The published legacy `Service`, `layer`, `clientLayer`, and module-level controls remain available as adapters over the same implementation, including the legacy live `requests` array. New tests should use `Test` and `testLayer`. @@ -259,7 +261,7 @@ This is different from prompt caching, server-side history storage, or truncatio ### Explicit compaction -`LLMClient.compact(request)` is the caller-controlled operation for OpenAI, Azure, and xAI Responses. It performs exactly one HTTP call to `/responses/compact`, using the selected route's endpoint, credentials, query, and HTTP middleware. It returns a `CompactionResponse` with `replacement: Message[]` and optional `usage`, not a normal generation response. +`LLMClient.compact(request)` (equivalently, `{ mechanism: "endpoint" }`) is the caller-controlled operation for OpenAI, Azure, and xAI Responses. It performs exactly one HTTP call to `/responses/compact`, using the selected route's endpoint, credentials, query, and HTTP middleware. It returns a `CompactionResponse` with `replacement: Message[]` and optional `usage`, not a normal generation response. This mechanism does not accept a WebSocket executor. Prefer this operation, where supported, when the application owns compaction policy and durable context updates. @@ -279,6 +281,41 @@ Generation-only body overlays such as `stream` and `store` are not sent to the c The input must still fit the model's context window. Explicit compaction is not an overflow-recovery operation. Anthropic does not expose this operation in this package; its in-band compaction remains available below. Compatible routes do not inherit an explicit compact endpoint simply because they use a Responses protocol. +### Streamed checkpoint compaction + +OpenAI Responses also exposes a separate, explicitly selected mechanism: + +```ts +const result = + yield * + LLMClient.compact(request, { + mechanism: "trigger", + webSocket, // Optional: without it, the request uses HTTP/SSE. + }) + +result.checkpoint // Successful encrypted CompactionPart. +result.responseID +result.usage +``` + +This appends a native `compaction_trigger` control item to the full input and sends a normal Responses request. It follows the [Codex V2 request shape](https://github.com/openai/codex/blob/728cb12/codex-rs/core/src/compact_remote_v2_attempt.rs), with tools and instructions retained, `stream: true`, `store: false`, and parallel tool calls enabled. It removes normal-answer text/output-format controls, forced tool choices, output-token/tool-call limits, and automatic `context_management`. Body overlays cannot replace `input` or supply `previous_response_id`/`conversation`; the complete canonical history is required for safe stateless replay. Session/cache identifiers, auth, headers, query parameters, service tier, and supported prompt-cache settings are preserved. + +Only a successful `response.completed` with a response ID and exactly one logical encrypted checkpoint succeeds. Repeated item events are correlated by ID/output slot, including ID-less checkpoints. Other output is ignored, not returned as assistant text or dispatched as tools. Failed, incomplete, malformed, and interrupted responses return errors rather than partial checkpoints. + +The result is **not a replacement window**. The caller selects retained history, combines it with `result.checkpoint`, and durably installs it before continuing. The operation does not choose a retention budget, prune messages, or modify the original request. + +The supplied WebSocket executor can reuse a compatible append baseline for the compaction request. On completion the protocol supplies no continuation checkpoint, clearing the old baseline so the next generation sends the newly installed window in full. Validation occurs before transport completion is acknowledged. There is no operation-level retry or fallback to `/responses/compact`; existing safe transport fallback may use SSE, with full history and no connection-local response ID. + +Trigger support is separate from endpoint support. Only the OpenAI Responses route advertises it; Azure, xAI, Chat, and compatible Responses routes do not inherit it. Untyped calls still fail before sending: missing route capabilities return `UnsupportedOperation`, while unknown mechanism names and invalid inputs return `InvalidRequest`. Dynamic callers must narrow for the selected mechanism: + +```ts +if (LLMClient.canCompact(request, { mechanism: "trigger" })) { + const result = yield * LLMClient.compact(request, { mechanism: "trigger" }) +} +``` + +This capability describes protocol implementation, **not universal availability on OpenAI API deployments**. The host application owns subscription/deployment eligibility, OAuth, endpoint selection, and deployment-specific headers. Local protocol/socket tests do not establish live provider support. + ### Advanced: in-band compaction `providerOptions.contextManagement` lets the provider decide when to compact during an ordinary `generate` or `stream` call. This is an advanced option for callers that own persistence and recovery: persist the complete assistant message, including its checkpoint, before continuing. Enabling the option does not provide durable checkpoint storage, interruption recovery, or model-switch policy. Keep the prior context until a successful checkpoint has been persisted. diff --git a/packages/ai/src/protocols/open-responses-continuation.ts b/packages/ai/src/protocols/open-responses-continuation.ts index 78c8e519729..27ae3a9f83e 100644 --- a/packages/ai/src/protocols/open-responses-continuation.ts +++ b/packages/ai/src/protocols/open-responses-continuation.ts @@ -42,6 +42,7 @@ const canonical = (value: unknown): string => { if (Array.isArray(value)) return `[${value.map(canonical).join(",")}]` if (!ProviderShared.isRecord(value)) return ProviderShared.encodeJson(value) return `{${Object.keys(value) + .filter((key) => value[key] !== undefined) .sort() .map((key) => `${ProviderShared.encodeJson(key)}:${canonical(value[key])}`) .join(",")}}` @@ -154,6 +155,12 @@ export const driver = (input: DriverInput): WebSocketChannelDriver => { if (rejection === "websocket_connection_limit_reached") return rejected(observation, "rotate-and-retry-full") } if (observation.type !== "completed") return observation + // A trigger installs a different context window. Clear the append baseline, retaining the socket. + if ( + Array.isArray(request.input) && + request.input.some((item) => ProviderShared.isRecord(item) && item.type === "compaction_trigger") + ) + return observation const responseID = event.response?.id if (!responseID || responseID.trim().length === 0) return observation return { diff --git a/packages/ai/src/protocols/open-responses.ts b/packages/ai/src/protocols/open-responses.ts index 15ff5ac466b..0dd32b4883c 100644 --- a/packages/ai/src/protocols/open-responses.ts +++ b/packages/ai/src/protocols/open-responses.ts @@ -923,7 +923,7 @@ const joinReasoningText = (parts: ReadonlyArray) => { return parts.filter((part) => part !== undefined).join("\n\n") } -const outputItemID = (state: ParserState, event: Event) => +const outputItemID = (state: Pick, event: Event) => event.output_index === undefined ? event.item_id : (state.outputItems[event.output_index] ?? event.item_id) const ITEM_ID_PREFIX: Readonly> = { @@ -935,7 +935,11 @@ const ITEM_ID_PREFIX: Readonly> = { // An item without an id adopts the id already open in its output slot, // otherwise it gets a locally minted one. -const resolveItem = (state: ParserState, item: StreamItem, index: number | undefined): OutputItem => ({ +const resolveItem = ( + state: Pick, + item: StreamItem, + index: number | undefined, +): OutputItem => ({ ...item, id: item.id ?? @@ -945,7 +949,7 @@ const resolveItem = (state: ParserState, item: StreamItem, index: number | undef // Registered output slots are authoritative for `item_id` routing, and items // are resolved here so everything downstream can rely on `item.id`. -export const normalize = (state: ParserState, input: Event): NormalizedEvent => ({ +export const normalize = (state: Pick, input: Event): NormalizedEvent => ({ ...input, item_id: input.item_id === undefined ? undefined : outputItemID(state, input), item: input.item ? resolveItem(state, input.item, input.output_index) : input.item, diff --git a/packages/ai/src/protocols/openai-responses.ts b/packages/ai/src/protocols/openai-responses.ts index 04ae45b438b..d848cde2fc0 100644 --- a/packages/ai/src/protocols/openai-responses.ts +++ b/packages/ai/src/protocols/openai-responses.ts @@ -5,7 +5,7 @@ import { Auth } from "../route/auth.js" import { Endpoint } from "../route/endpoint.js" import { Protocol } from "../route/protocol.js" import { HttpTransport } from "../route/transport/index.js" -import type { LLMRequest, JsonSchema, ToolDefinition, ToolEntry } from "../schema/index.js" +import { LLMRequest, mergeJsonRecords, type JsonSchema, type ToolDefinition, type ToolEntry } from "../schema/index.js" import { OpenResponses } from "./open-responses.js" import { JsonObject, optionalArray, optionalNull, ProviderShared } from "./shared.js" import { OpenAIImage } from "./utils/openai-image.js" @@ -13,6 +13,7 @@ import { ResponsesHostedTools } from "./utils/responses-hosted-tools.js" import { ToolSchemaProjection } from "./utils/tool-schema.js" import { OpenResponsesChannel } from "./open-responses-channel.js" import { ResponsesCompaction } from "./utils/responses-compaction.js" +import { ResponsesCheckpoint } from "./utils/responses-checkpoint.js" const ADAPTER = "openai-responses" const NAME = "OpenAI Responses" @@ -114,6 +115,18 @@ const OpenAIResponsesBody = Schema.Struct({ }) export type OpenAIResponsesBody = Schema.Schema.Type +/** Request control, never conversation content. */ +export const CompactionTrigger = Schema.Struct({ type: Schema.Literal("compaction_trigger") }) +const CheckpointBody = Schema.Struct({ + ...OpenAIResponsesBody.fields, + input: Schema.Array(Schema.Union([OpenResponses.InputItem, OpenAIResponsesHostedToolItem, CompactionTrigger])), + store: Schema.Literal(false), + prompt_cache_retention: optionalNull(Schema.String), + prompt_cache_options: optionalNull( + Schema.Struct({ mode: Schema.optional(Schema.String), ttl: Schema.optional(Schema.String) }), + ), +}) + const adapter = { id: ADAPTER, name: NAME, @@ -191,6 +204,35 @@ const fromRequest = Effect.fn("OpenAIResponses.fromRequest")(function* (request: }) }) +const checkpointBody = { + schema: CheckpointBody, + from: Effect.fn("OpenAIResponses.checkpointBody")(function* (request: LLMRequest) { + const native = yield* fromRequest(LLMRequest.update(request, { toolChoice: undefined })) + const overlay = request.http?.body + // Complete history is required for stateless replay and SSE recovery. Raw input overrides bypass that contract. + if ( + overlay?.input !== undefined || + overlay?.previous_response_id !== undefined || + overlay?.conversation !== undefined + ) + return yield* ProviderShared.invalidRequest( + "Trigger compaction requires complete canonical history, not an input or continuation override", + ) + return yield* ProviderShared.validateWith(Schema.decodeUnknownEffect(CheckpointBody))({ + ...mergeJsonRecords(native, overlay), + input: [...native.input, { type: "compaction_trigger" }], + stream: true, + store: false, + parallel_tool_calls: true, + tool_choice: undefined, + context_management: undefined, + text: undefined, + max_output_tokens: undefined, + max_tool_calls: undefined, + }) + }), +} + const hostedToolResult = Effect.fn("OpenAIResponses.hostedToolResult")(function* (item: ResponsesHostedTools.Item) { const isError = item.error !== undefined && item.error !== null if (item.type === "image_generation_call" && item.result) { @@ -268,7 +310,7 @@ export const transport = channelTransport({ }) export const route = Route.make({ - compact: ResponsesCompaction.make(adapter), + compact: { endpoint: ResponsesCompaction.make(adapter), trigger: ResponsesCheckpoint.make(checkpointBody) }, id: ADAPTER, provider: "openai", providerMetadataKey: "openai", diff --git a/packages/ai/src/protocols/utils/responses-checkpoint.ts b/packages/ai/src/protocols/utils/responses-checkpoint.ts new file mode 100644 index 00000000000..6231d657959 --- /dev/null +++ b/packages/ai/src/protocols/utils/responses-checkpoint.ts @@ -0,0 +1,122 @@ +import { Effect, Schema, Stream } from "effect" +import { Route, type RouteBody, type TriggerCompactOperation } from "../../route/client.js" +import { Protocol } from "../../route/protocol.js" +import { CompactionCheckpointResponse, HttpOptions, LLMEvent, LLMRequest } from "../../schema/index.js" +import { OpenResponses } from "../open-responses.js" +import { ProviderShared } from "../shared.js" + +interface State { + readonly parser: Pick + readonly checkpoints: Readonly> + readonly responseID?: string +} + +const onOutputItem = Effect.fn("ResponsesCheckpoint.onOutputItem")(function* ( + state: State, + input: OpenResponses.Event, +) { + const event = OpenResponses.normalize(state.parser, input) + const item = event.item + if (!item) return state + const parser = + event.output_index === undefined || state.parser.outputItems[event.output_index] === item.id + ? state.parser + : { ...state.parser, outputItems: { ...state.parser.outputItems, [event.output_index]: item.id } } + const next = parser === state.parser ? state : { ...state, parser } + if (event.type === "response.output_item.added" || item.type !== "compaction") return next + if ( + event.output_index !== undefined && + Object.entries(state.parser.outputItems).some( + ([index, id]) => id === item.id && Number(index) !== event.output_index, + ) + ) + return yield* ProviderShared.eventError(parser.id, "Compaction checkpoint appeared in multiple output slots") + if (!item.encrypted_content) + return yield* ProviderShared.eventError(parser.id, "Compaction output is missing its encrypted content") + const previous = state.checkpoints[item.id] + if (previous && previous.encrypted !== item.encrypted_content) + return yield* ProviderShared.eventError(parser.id, "Compaction output changed after completion") + if (previous) return next + return { + ...next, + checkpoints: { + ...state.checkpoints, + [item.id]: { type: "compaction", provider: parser.provider, id: item.id, encrypted: item.encrypted_content }, + }, + } satisfies State +}) + +/** Collect a trigger response before acknowledging transport completion. No generation output escapes. */ +export const make = (body: RouteBody): TriggerCompactOperation => + Effect.fn("ResponsesCheckpoint.execute")(function* (request, executor, options) { + const source = request.model.route + let result: CompactionCheckpointResponse | undefined + // Route registries erase the frame type. The codec validates that boundary before parsing. + const event: Schema.Codec = OpenResponses.protocol.stream.event + const protocol = Protocol.make({ + id: source.protocol, + body, + stream: { + event, + initial: (request: LLMRequest): State => ({ + parser: { id: source.id, provider: request.model.provider, outputItems: {} }, + checkpoints: {}, + }), + terminal: OpenResponses.terminal, + step: Effect.fn("ResponsesCheckpoint.step")(function* (state: State, event: OpenResponses.Event) { + if (event.response?.id && state.responseID && event.response.id !== state.responseID) + return yield* ProviderShared.eventError(source.id, "Compaction response ID changed during execution") + if (event.type === "response.created") return [{ ...state, responseID: event.response?.id }, []] as const + if (event.type === "error" || event.type === "response.failed") + return yield* OpenResponses.providerFailure(event, "Compaction request failed") + if (event.type === "response.incomplete") + return yield* ProviderShared.eventError(source.id, "Compaction response was incomplete") + if (event.type === "response.output_item.added" || event.type === "response.output_item.done") + return [yield* onOutputItem(state, event), []] as const + if (event.type !== "response.completed") return [state, []] as const + const responseID = event.response?.id + if (!responseID?.trim()) + return yield* ProviderShared.eventError(source.id, "Compaction response is missing its response ID") + if (event.response?.status !== undefined && event.response.status !== "completed") + return yield* ProviderShared.eventError(source.id, "Compaction response did not complete successfully") + let next = state + for (const [index, item] of (event.response?.output ?? []).entries()) { + next = yield* onOutputItem(next, { type: "response.output_item.done", output_index: index, item }) + } + const checkpoints = Object.values(next.checkpoints) + const checkpoint = checkpoints[0] + if (checkpoints.length !== 1 || !checkpoint) + return yield* ProviderShared.eventError( + source.id, + "Compaction response must contain exactly one checkpoint", + ) + result = new CompactionCheckpointResponse({ + checkpoint, + responseID, + usage: OpenResponses.mapUsage(event.response?.usage, OpenResponses.metadataKey(request.model)), + }) + return [next, [LLMEvent.finish({ reason: { normalized: "stop" } })]] as const + }), + }, + }) + const route = Route.make({ + id: source.id, + provider: source.provider, + providerMetadataKey: source.providerMetadataKey, + protocol, + endpoint: source.endpoint, + auth: source.auth, + transport: source.transport, + }) + const native = yield* body.from(request) + // The body builder already applied and validated overlays. Do not let transport reapply them. + const preparedRequest = LLMRequest.update(request, { + http: request.http === undefined ? undefined : new HttpOptions({ ...request.http, body: undefined }), + }) + const prepared = yield* route.prepareTransport(native, preparedRequest, options) + yield* route.streamPrepared(prepared, preparedRequest, { http: executor }, options).pipe(Stream.runDrain) + if (!result) return yield* ProviderShared.eventError(source.id, "Compaction response ended without a checkpoint") + return result + }) + +export * as ResponsesCheckpoint from "./responses-checkpoint.js" diff --git a/packages/ai/src/provider-package.ts b/packages/ai/src/provider-package.ts index 9595a0ec481..79f29084833 100644 --- a/packages/ai/src/provider-package.ts +++ b/packages/ai/src/provider-package.ts @@ -1,5 +1,5 @@ import type { LanguageModel, ProviderOptions } from "./schema/index.js" -import type { CompactOperation } from "./route/client.js" +import type { CompactionOperations } from "./route/client.js" export interface Settings extends Readonly> { readonly baseURL?: string @@ -10,7 +10,7 @@ export interface Settings extends Readonly> { export interface Definition< ProviderSettings extends Settings = Settings, Options extends ProviderOptions = ProviderOptions, - Compact extends CompactOperation | undefined = CompactOperation | undefined, + Compact extends CompactionOperations | undefined = CompactionOperations | undefined, > { readonly model: (modelID: string, settings: ProviderSettings) => LanguageModel } diff --git a/packages/ai/src/providers/azure.ts b/packages/ai/src/providers/azure.ts index 4bcdd6fad47..c1ae110ea64 100644 --- a/packages/ai/src/providers/azure.ts +++ b/packages/ai/src/providers/azure.ts @@ -1,7 +1,7 @@ import { Headers } from "effect/unstable/http" import { Auth } from "../route/auth.js" import { type AtLeastOne, type ProviderAuthOption } from "../route/auth-options.js" -import type { Route, RouteDefaultsInput, CompactOperation } from "../route/client.js" +import type { Route, RouteDefaultsInput, CompactionOperations } from "../route/client.js" import type { ProviderPackage } from "../provider-package.js" import { ProviderID, type ModelID } from "../schema/index.js" import * as OpenAIChat from "../protocols/openai-chat.js" @@ -39,6 +39,7 @@ export type Settings = ProviderPackage.Settings & const resourceBaseURL = (resourceName: string) => `https://${resourceName.trim()}.openai.azure.com/openai` const responsesRoute = OpenAIResponses.route.with({ + compact: { endpoint: OpenAIResponses.route.compact.endpoint }, id: "azure-openai-responses", provider: id, auth: routeAuth, @@ -102,7 +103,7 @@ const auth = (input: Config) => { ) } -const configuredRoute = ( +const configuredRoute = ( route: Route, input: Config, modelID: string | ModelID, @@ -168,7 +169,7 @@ const config = (settings: Settings): Config => { export const responsesModel: ProviderPackage.Definition< Settings, OpenAIProviderOptionsInput, - CompactOperation + typeof responsesRoute.compact >["model"] = (modelID, settings) => configure(config(settings)).responses(modelID) export const chatModel: ProviderPackage.Definition["model"] = ( modelID, diff --git a/packages/ai/src/providers/openai.ts b/packages/ai/src/providers/openai.ts index 1aa1dc09d7d..503dd9745df 100644 --- a/packages/ai/src/providers/openai.ts +++ b/packages/ai/src/providers/openai.ts @@ -1,5 +1,5 @@ import { AuthOptions, type ProviderAuthOption } from "../route/auth-options.js" -import type { Route, RouteDefaultsInput, CompactOperation } from "../route/client.js" +import type { Route, RouteDefaultsInput, CompactionOperations } from "../route/client.js" import type { ProviderPackage } from "../provider-package.js" import { HttpOptions, ProviderID, ToolDefinition, mergeHttpOptions, type ModelID } from "../schema/index.js" import * as OpenAIChat from "../protocols/openai-chat.js" @@ -73,7 +73,7 @@ const defaults = (input: Config) => { return rest } -const configuredRoute = ( +const configuredRoute = ( route: Route, input: Config, ) => @@ -132,10 +132,11 @@ const config = (settings: Settings): Config => { } } -export const model: ProviderPackage.Definition["model"] = ( - modelID, - settings, -) => { +export const model: ProviderPackage.Definition< + Settings, + OpenAIProviderOptionsInput, + typeof OpenAIResponses.route.compact +>["model"] = (modelID, settings) => { return configure(config(settings)).responses(modelID) } diff --git a/packages/ai/src/providers/xai.ts b/packages/ai/src/providers/xai.ts index a615ae2d6d9..213ba92c2b5 100644 --- a/packages/ai/src/providers/xai.ts +++ b/packages/ai/src/providers/xai.ts @@ -1,5 +1,5 @@ import { AuthOptions, type ProviderAuthOption } from "../route/auth-options.js" -import { Route, type RouteDefaultsInput, type CompactOperation } from "../route/client.js" +import { Route, type RouteDefaultsInput } from "../route/client.js" import { Endpoint } from "../route/endpoint.js" import { HttpOptions, ProviderID, type ModelID } from "../schema/index.js" import * as OpenAICompatibleProfiles from "./openai-compatible-profile.js" @@ -32,7 +32,7 @@ export type { XAIImageOptions } from "../protocols/xai-images.js" const RESPONSES_WEBSOCKET_ROTATE_AFTER_MS = 24 * 60 * 1000 const responsesRoute = Route.make({ - compact: XAIResponses.compact, + compact: { endpoint: XAIResponses.compact }, id: "openai-responses", provider: id, providerMetadataKey: "xai", @@ -103,10 +103,11 @@ export const configure = (input: LanguageModelOptions = {}) => { } export const provider = configure() -export const model: ProviderPackage.Definition["model"] = ( - modelID, - settings, -) => +export const model: ProviderPackage.Definition< + Settings, + XAIProviderOptionsInput, + typeof responsesRoute.compact +>["model"] = (modelID, settings) => configure({ apiKey: settings.apiKey, baseURL: settings.baseURL, diff --git a/packages/ai/src/route/client.ts b/packages/ai/src/route/client.ts index 4d2d07380fa..a53c9f58ee0 100644 --- a/packages/ai/src/route/client.ts +++ b/packages/ai/src/route/client.ts @@ -14,6 +14,7 @@ import type { ProtocolID, ProviderOptions } from "../schema/index.js" import { AIError, CompactionResponse, + CompactionCheckpointResponse, AIErrorReason, GenerationOptions, HttpOptions, @@ -38,7 +39,7 @@ export interface RouteBody { export interface Route< Body, Prepared = unknown, - Compact extends CompactOperation | undefined = CompactOperation | undefined, + Compact extends CompactionOperations | undefined = CompactionOperations | undefined, > { readonly compact: Compact readonly id: string @@ -53,7 +54,15 @@ export interface Route< readonly transport: Transport readonly defaults: RouteDefaults readonly body: RouteBody - readonly with: (patch: RoutePatch) => Route + readonly with: { + ( + patch: RoutePatch & { readonly compact: Next }, + ): Route + ( + patch: Omit, "compact"> & { readonly compact?: undefined }, + ): Route + (patch: RoutePatch): Route + } readonly model: ( input: RouteMappedLanguageModelInput, ) => LanguageModel @@ -74,7 +83,7 @@ export interface Route< // Normal call sites use `OpenAIChat.route`; callers only need body types // when preparing a request with a protocol-specific type assertion. // oxlint-disable-next-line typescript-eslint/no-explicit-any -export type AnyRoute = Route< +export type AnyRoute = Route< any, any, Compact @@ -101,6 +110,7 @@ export interface RouteDefaultsInput { } export interface RoutePatch extends RouteDefaultsInput { + readonly compact?: CompactionOperations readonly id?: string readonly provider?: string | ProviderID readonly providerMetadataKey?: string @@ -111,7 +121,7 @@ export interface RoutePatch extends RouteDefaultsInput { type RouteMappedLanguageModelInput = RouteLanguageModelInput | RouteRoutedLanguageModelInput -const makeRouteLanguageModel = ( +const makeRouteLanguageModel = ( route: AnyRoute, mapped: RouteMappedLanguageModelInput, ) => { @@ -162,10 +172,7 @@ export const httpOptions = (input: HttpOptionsInput | undefined) => { } export interface Interface { - readonly compact: ( - request: CompactionRequest, - options?: Pick, - ) => Effect.Effect + readonly compact: CompactMethod readonly stream: StreamMethod readonly generate: GenerateMethod } @@ -189,12 +196,64 @@ export type CompactOperation = ( options?: Pick, ) => Effect.Effect -export type CompactionRequest = LLMRequest & { - readonly model: LanguageModel +export type TriggerCompactOperation = ( + request: LLMRequest, + executor: RequestExecutor.Interface, + options: TriggerCompactOptions, +) => Effect.Effect + +/** Protocol capabilities, not deployment/model eligibility. */ +export interface CompactionOperations { + readonly endpoint?: CompactOperation + readonly trigger?: TriggerCompactOperation } -export const canCompact = (request: LLMRequest): request is CompactionRequest => - request.model.route.compact !== undefined +export interface EndpointCompactOptions extends Pick { + readonly mechanism?: "endpoint" + readonly webSocket?: never +} + +export interface TriggerCompactOptions extends StreamOptions { + readonly mechanism: "trigger" +} + +// Keep the required route shape explicit: the schema class's self type erases its model parameter in assignability. +export type CompactionRequest = LLMRequest & { + readonly model: LanguageModel +} +export type CheckpointRequest = LLMRequest & { + readonly model: LanguageModel +} + +export interface CompactMethod { + (request: CheckpointRequest, options: TriggerCompactOptions): Effect.Effect + (request: CompactionRequest, options?: EndpointCompactOptions): Effect.Effect +} + +export function canCompact( + request: LLMRequest, + options?: { readonly mechanism?: "endpoint" }, +): request is CompactionRequest +export function canCompact( + request: LLMRequest, + options: { readonly mechanism: "trigger" }, +): request is CheckpointRequest +export function canCompact(request: LLMRequest, options?: { readonly mechanism?: string }) { + if (options?.mechanism === "trigger") return request.model.route.compact?.trigger !== undefined + if (options?.mechanism !== undefined && options.mechanism !== "endpoint") return false + return request.model.route.compact?.endpoint !== undefined +} + +const unsupportedCompaction = (request: LLMRequest, mechanism: string | undefined) => { + if (mechanism !== undefined && mechanism !== "endpoint" && mechanism !== "trigger") + return ProviderShared.invalidRequest(`Unknown compaction mechanism: ${mechanism}`) + return ProviderShared.unsupportedOperation({ + operation: mechanism === "trigger" ? "compact.trigger" : "compact", + provider: request.model.provider, + route: request.model.route.id, + message: `${request.model.provider}/${request.model.route.id} does not support ${mechanism === "trigger" ? "trigger" : "explicit"} compaction`, + }) +} export class Service extends Context.Service()("@opencode/LLMClient") {} @@ -216,7 +275,7 @@ const resolveRequestOptions = (request: LLMRequest) => { } export interface MakeInput { - readonly compact?: CompactOperation + readonly compact?: CompactionOperations /** Route id used in diagnostics and prepared request metadata. */ readonly id: string /** Provider identity for route-owned model construction. */ @@ -238,7 +297,7 @@ export interface MakeInput { } export interface MakeTransportInput { - readonly compact?: CompactOperation + readonly compact?: CompactionOperations /** Route id used in diagnostics and prepared request metadata. */ readonly id: string /** Provider identity for route-owned model construction. */ @@ -326,9 +385,10 @@ function makeFromTransport( defaults: routeInput.defaults ?? {}, body: protocol.body, with: (patch: RoutePatch) => { - const { id, provider, providerMetadataKey, auth, transport, endpoint, ...defaults } = patch + const { compact, id, provider, providerMetadataKey, auth, transport, endpoint, ...defaults } = patch return build({ ...routeInput, + compact: "compact" in patch ? compact : routeInput.compact, id: id ?? routeInput.id, provider: provider ?? routeInput.provider, providerMetadataKey: @@ -343,7 +403,7 @@ function makeFromTransport( }) }, model: (input: RouteMappedLanguageModelInput) => - makeRouteLanguageModel(route, input), + makeRouteLanguageModel(route, input), prepareTransport: (body, request, options) => routeInput.transport.prepare({ body, @@ -440,12 +500,12 @@ function makeFromTransport( return build({ ...input, defaults: mergeRouteDefaults(undefined, input.defaults ?? {}) }) } -export function make( - input: MakeTransportInput & { readonly compact: CompactOperation }, -): Route -export function make( - input: MakeInput & { readonly compact: CompactOperation }, -): Route, CompactOperation> +export function make( + input: MakeTransportInput & { readonly compact: Compact }, +): Route +export function make( + input: MakeInput & { readonly compact: Compact }, +): Route, Compact> export function make( input: MakeTransportInput, ): Route @@ -561,14 +621,23 @@ export function generate(request: LLMRequest, options?: StreamOptions): Effect.E }) } -export const compact = ( +export function compact( + request: CheckpointRequest, + options: TriggerCompactOptions, +): Effect.Effect +export function compact( request: CompactionRequest, - options?: Pick, -): Effect.Effect => - Effect.gen(function* () { + options?: EndpointCompactOptions, +): Effect.Effect +export function compact(request: LLMRequest, options?: EndpointCompactOptions | TriggerCompactOptions) { + return Effect.gen(function* () { const client = yield* Service - return yield* client.compact(request, options) + if (options?.mechanism === "trigger" && canCompact(request, options)) return yield* client.compact(request, options) + if ((options?.mechanism === undefined || options.mechanism === "endpoint") && canCompact(request)) + return yield* client.compact(request, options) + return yield* unsupportedCompaction(request, options?.mechanism) }) +} export const streamRequest = (request: LLMRequest, options?: StreamOptions) => Stream.unwrap( @@ -582,21 +651,27 @@ export const layer: Layer.Layer = Layer Effect.gen(function* () { const executor = yield* RequestExecutor.Service const stream = streamRequestWith({ http: executor }) + function compact( + request: CompactionRequest, + options?: EndpointCompactOptions, + ): Effect.Effect + function compact( + request: CheckpointRequest, + options: TriggerCompactOptions, + ): Effect.Effect + function compact(request: LLMRequest, options?: EndpointCompactOptions | TriggerCompactOptions) { + return Effect.suspend((): Effect.Effect => { + if (options?.mechanism === "trigger" && canCompact(request, options)) + return request.model.route.compact.trigger(prepareRequest(request), executor, options) + if ((options?.mechanism === undefined || options.mechanism === "endpoint") && canCompact(request)) + return request.model.route.compact.endpoint(prepareRequest(request), executor, options) + return unsupportedCompaction(request, options?.mechanism) + }) + } return Service.of({ stream, generate: generateWith(stream), - compact: (request, options) => - Effect.suspend(() => { - const operation = request.model.route.compact - if (!operation) - return ProviderShared.unsupportedOperation({ - operation: "compact", - provider: request.model.provider, - route: request.model.route.id, - message: `${request.model.provider}/${request.model.route.id} does not support explicit compaction`, - }) - return operation(prepareRequest(request), executor, options) - }), + compact, }) }), ) diff --git a/packages/ai/src/route/index.ts b/packages/ai/src/route/index.ts index d794c5ae7e4..a3ecdc8b76d 100644 --- a/packages/ai/src/route/index.ts +++ b/packages/ai/src/route/index.ts @@ -9,6 +9,12 @@ export type { Interface as LLMClientShape, Service as LLMClientService, StreamOptions, + CompactMethod, + CompactionOperations, + CompactionRequest, + CheckpointRequest, + EndpointCompactOptions, + TriggerCompactOptions, } from "./client.js" export * from "./executor.js" export { Auth } from "./auth.js" diff --git a/packages/ai/src/schema/events.ts b/packages/ai/src/schema/events.ts index e2cce25c97a..6396319d5e1 100644 --- a/packages/ai/src/schema/events.ts +++ b/packages/ai/src/schema/events.ts @@ -97,6 +97,21 @@ export class CompactionResponse extends Schema.Class("LLM.Co usage: Schema.optional(Usage), }) {} +/** A checkpoint only; retained history and replacement-window construction belong to the caller. */ +export class CompactionCheckpointResponse extends Schema.Class( + "LLM.CompactionCheckpointResponse", +)({ + checkpoint: CompactionPart.pipe( + Schema.refine( + (part): part is CompactionPart & { readonly encrypted: string; readonly text?: never } => + part.encrypted !== undefined && part.encrypted.length > 0, + { message: "A checkpoint response requires encrypted compaction content" }, + ), + ), + responseID: Schema.String.check(Schema.isPattern(/\S/)), + usage: Schema.optional(Usage), +}) {} + export const StepStart = Schema.Struct({ type: Schema.tag("step-start"), index: Schema.Number, diff --git a/packages/ai/src/schema/options.ts b/packages/ai/src/schema/options.ts index 3cf9f6a3ee1..a1dffada699 100644 --- a/packages/ai/src/schema/options.ts +++ b/packages/ai/src/schema/options.ts @@ -1,6 +1,6 @@ import { Schema } from "effect" import { ModelID, ProviderID } from "./ids.js" -import type { AnyRoute, CompactOperation } from "../route/client.js" +import type { AnyRoute, CompactionOperations } from "../route/client.js" import { isRecord } from "../utils/record.js" export const JsonSchema = Schema.Record(Schema.String, Schema.Unknown) @@ -175,7 +175,7 @@ export namespace LanguageModelCompatibility { export class LanguageModel< Options extends ProviderOptions = ProviderOptions, - Compact extends CompactOperation | undefined = CompactOperation | undefined, + Compact extends CompactionOperations | undefined = CompactionOperations | undefined, > { declare protected readonly _ProviderOptions: Options readonly id: ModelID @@ -194,7 +194,7 @@ export class LanguageModel< static make< Options extends ProviderOptions = ProviderOptions, - Compact extends CompactOperation | undefined = CompactOperation | undefined, + Compact extends CompactionOperations | undefined = CompactionOperations | undefined, >(input: LanguageModel.Input) { return new LanguageModel({ id: ModelID.make(input.id), @@ -206,7 +206,7 @@ export class LanguageModel< }) } - static input( + static input( model: LanguageModel, ): LanguageModel.ConstructorInput { return { @@ -218,11 +218,11 @@ export class LanguageModel< } } - static update( + static update( model: LanguageModel, patch: Partial> & { readonly route: AnyRoute }, ): LanguageModel - static update( + static update( model: LanguageModel, patch: Partial> & { readonly route?: undefined }, ): LanguageModel @@ -241,7 +241,7 @@ export class LanguageModel< } export namespace LanguageModel { - export type ConstructorInput = { + export type ConstructorInput = { readonly id: ModelID readonly provider: ProviderID readonly route: AnyRoute @@ -249,7 +249,7 @@ export namespace LanguageModel { readonly compatibility?: LanguageModelCompatibility } - export type Input = Omit< + export type Input = Omit< ConstructorInput, "id" | "provider" | "defaults" | "compatibility" > & { diff --git a/packages/ai/src/testing.ts b/packages/ai/src/testing.ts index 3e697d4a8a9..420ff88282e 100644 --- a/packages/ai/src/testing.ts +++ b/packages/ai/src/testing.ts @@ -1,10 +1,17 @@ export * as TestLLM from "./testing.js" -import { LLMClient } from "./route/client.js" +import { + LLMClient, + type CompactionRequest, + type CheckpointRequest, + type EndpointCompactOptions, + type TriggerCompactOptions, +} from "./route/client.js" import { LLMEvent, LLMResponse, CompactionResponse, + CompactionCheckpointResponse, type FinishReasonDetails, type AIError, type LLMRequest, @@ -13,7 +20,11 @@ import { } from "./schema/index.js" import { Context, Deferred, Effect, Latch, Layer, Queue, Scope, Stream } from "effect" -export type Response = readonly LLMEvent[] | Stream.Stream | CompactionResponse +export type Response = + | readonly LLMEvent[] + | Stream.Stream + | CompactionResponse + | CompactionCheckpointResponse export type Gate = Readonly<{ started: Effect.Effect; release: Effect.Effect }> @@ -132,21 +143,38 @@ const make = (options: LayerOptions) => Stream.unwrap( take(request).pipe( Effect.map((response) => { - if (response instanceof CompactionResponse) + if (response instanceof CompactionResponse || response instanceof CompactionCheckpointResponse) return Stream.die("TestLLM generation requires an event response") return Stream.isStream(response) ? response : Stream.fromIterable(response) }), ), ) - const test = Test.of({ - compact: (request) => - take(request).pipe( - Effect.flatMap((response) => - response instanceof CompactionResponse + function compact( + request: CompactionRequest, + options?: EndpointCompactOptions, + ): Effect.Effect + function compact( + request: CheckpointRequest, + options: TriggerCompactOptions, + ): Effect.Effect + function compact( + request: LLMRequest, + options?: EndpointCompactOptions | TriggerCompactOptions, + ): Effect.Effect { + return take(request).pipe( + Effect.flatMap((response): Effect.Effect => { + if (options?.mechanism === "trigger") + return response instanceof CompactionCheckpointResponse ? Effect.succeed(response) - : Effect.die("TestLLM compaction requires a CompactionResponse"), - ), - ), + : Effect.die("TestLLM trigger compaction requires a CompactionCheckpointResponse") + return response instanceof CompactionResponse + ? Effect.succeed(response) + : Effect.die("TestLLM compaction requires a CompactionResponse") + }), + ) + } + const test = Test.of({ + compact, stream, generate: (request) => stream(request).pipe( diff --git a/packages/ai/test/provider-options/checkpoint.types.ts b/packages/ai/test/provider-options/checkpoint.types.ts new file mode 100644 index 00000000000..081bfc8484d --- /dev/null +++ b/packages/ai/test/provider-options/checkpoint.types.ts @@ -0,0 +1,129 @@ +import { Effect } from "effect" +import { + CompactionCheckpointResponse, + CompactionResponse, + LanguageModel, + LLM, + LLMClient, + LLMRequest, +} from "../../src/index.js" +import { + Anthropic, + Azure, + AmazonBedrock, + AmazonBedrockMantle, + OpenAI, + OpenAICompatibleResponses, + XAI, +} from "../../src/providers.js" +import type { WebSocketChannelExecutor } from "../../src/route.js" +import type { RoutePatch } from "../../src/route/client.js" +import type { OpenAIResponsesBody } from "../../src/protocols/openai-responses.js" +import type { Prepared } from "../../src/protocols/open-responses-channel.js" + +declare const webSocket: WebSocketChannelExecutor +const model = OpenAI.configure().responses("fixture") +const request = LLM.request({ model, prompt: "hello" }) +LLMClient.compact(request, { mechanism: "endpoint" }).pipe(Effect.map((result) => result satisfies CompactionResponse)) +LLMClient.compact(request, { mechanism: "trigger", webSocket }).pipe( + Effect.map((result) => { + result satisfies CompactionCheckpointResponse + result.checkpoint.encrypted satisfies string + result.responseID satisfies string + // @ts-expect-error A trigger does not return replacement history. + result.replacement + }), +) +// @ts-expect-error Endpoint compaction does not accept a WebSocket executor. +LLMClient.compact(request, { mechanism: "endpoint", webSocket }) +// @ts-expect-error Omitting mechanism selects the HTTP endpoint. +LLMClient.compact(request, { webSocket }) +// @ts-expect-error Unknown mechanisms do not have a permissive fallback overload. +LLMClient.compact(request, { mechanism: "other" }) + +for (const selected of [ + model, + OpenAI.model("fixture", {}), + model.route.with({ headers: { fixture: "test" } }).model({ id: "fixture" }), + LanguageModel.make(LanguageModel.input(model)), + LanguageModel.update(model, { defaults: { generation: { maxTokens: 100 } } }), +]) { + LLMClient.compact(LLM.request({ model: selected }), { mechanism: "trigger" }) +} +LLMClient.compact(new LLMRequest(LLMRequest.input(request)), { mechanism: "trigger" }) +LLMClient.compact(LLMRequest.update(request, { messages: [] }), { mechanism: "trigger" }) + +const azure = Azure.configure({ resourceName: "fixture" }).responses("fixture") +const xai = XAI.configure().responses("fixture") +// @ts-expect-error Azure must not inherit OpenAI's trigger operation. +LLMClient.compact(LLM.request({ model: azure }), { mechanism: "trigger" }) +LLMClient.compact(LLM.request({ model: Azure.responsesModel("fixture", { resourceName: "fixture" }) }), { + // @ts-expect-error Azure's package entrypoint must preserve its narrower capability. + mechanism: "trigger", +}) +// @ts-expect-error xAI endpoint support does not imply trigger support. +LLMClient.compact(LLM.request({ model: xai }), { mechanism: "trigger" }) +// @ts-expect-error xAI's package entrypoint must preserve its narrower capability. +LLMClient.compact(LLM.request({ model: XAI.model("fixture", {}) }), { mechanism: "trigger" }) + +const unsupported = { + bedrock: LLM.request({ model: AmazonBedrock.configure().model("fixture") }), + mantle: LLM.request({ model: AmazonBedrockMantle.configure().responses("fixture") }), + anthropic: LLM.request({ model: Anthropic.configure().model("fixture") }), + openai: LLM.request({ model: OpenAI.configure().chat("fixture") }), + azure: LLM.request({ model: Azure.configure({ resourceName: "fixture" }).chat("fixture") }), + xai: LLM.request({ model: XAI.configure().chat("fixture") }), + compatible: LLM.request({ + model: OpenAICompatibleResponses.configure({ baseURL: "https://example.com" }).model("fixture"), + }), +} +// @ts-expect-error Bedrock does not expose trigger compaction. +LLMClient.compact(unsupported.bedrock, { mechanism: "trigger" }) +// @ts-expect-error Mantle must not inherit trigger support from OpenAI's protocol. +LLMClient.compact(unsupported.mantle, { mechanism: "trigger" }) +// @ts-expect-error Anthropic does not expose trigger compaction. +LLMClient.compact(unsupported.anthropic, { mechanism: "trigger" }) +// @ts-expect-error OpenAI Chat does not expose trigger compaction. +LLMClient.compact(unsupported.openai, { mechanism: "trigger" }) +// @ts-expect-error Azure Chat does not expose trigger compaction. +LLMClient.compact(unsupported.azure, { mechanism: "trigger" }) +// @ts-expect-error xAI Chat does not expose trigger compaction. +LLMClient.compact(unsupported.xai, { mechanism: "trigger" }) +// @ts-expect-error Generic protocol compatibility does not grant trigger support. +LLMClient.compact(unsupported.compatible, { mechanism: "trigger" }) +// @ts-expect-error Changing the model replaces its capability. +LLMClient.compact(LLMRequest.update(request, { model: azure }), { mechanism: "trigger" }) +// @ts-expect-error Changing the route replaces its capability. +LLMClient.compact(LLM.request({ model: LanguageModel.update(model, { route: azure.route }) }), { mechanism: "trigger" }) +LLMClient.compact( + LLM.request({ + model: model.route.with({ compact: { endpoint: model.route.compact.endpoint } }).model({ id: "fixture" }), + }), + // @ts-expect-error Replacing route operations does not retain the old trigger capability. + { mechanism: "trigger" }, +) + +declare const dynamic: LLMRequest +declare const patch: Partial +declare const routePatch: RoutePatch +// @ts-expect-error A dynamic operation override cannot preserve trigger support. +LLMClient.compact(LLM.request({ model: model.route.with(routePatch).model({ id: "fixture" }) }), { + mechanism: "trigger", +}) +// @ts-expect-error Explicitly removing operations removes trigger support. +LLMClient.compact(LLM.request({ model: model.route.with({ compact: undefined }).model({ id: "fixture" }) }), { + mechanism: "trigger", +}) +// @ts-expect-error Dynamic models must be narrowed. +LLMClient.compact(dynamic, { mechanism: "trigger" }) +if (LLMClient.canCompact(dynamic, { mechanism: "trigger" })) { + LLMClient.compact(dynamic, { mechanism: "trigger" }) + LLMClient.Service.use((client) => client.compact(dynamic, { mechanism: "trigger" })) +} +if (LLMClient.canCompact(dynamic)) { + LLMClient.compact(dynamic) + // @ts-expect-error Endpoint narrowing does not grant trigger support. + LLMClient.compact(dynamic, { mechanism: "trigger" }) +} +// @ts-expect-error A dynamic model override cannot preserve trigger support. +LLMClient.compact(LLMRequest.update(request, patch), { mechanism: "trigger" }) diff --git a/packages/ai/test/provider/checkpoint.test.ts b/packages/ai/test/provider/checkpoint.test.ts new file mode 100644 index 00000000000..62636e5eb32 --- /dev/null +++ b/packages/ai/test/provider/checkpoint.test.ts @@ -0,0 +1,243 @@ +import { expect } from "bun:test" +import { Effect, Stream } from "effect" +import { LLM, LLMClient, LLMRequest } from "../../src/index.js" +import { Azure, OpenAI } from "../../src/providers.js" +import { testEffect } from "../lib/effect.js" +import { dynamicResponse, fixedResponse } from "../lib/http.js" +import { sseEvents } from "../lib/sse.js" + +const checkpoint = { type: "compaction", id: "cmp_1", encrypted_content: "opaque" } +const request = LLM.request({ model: OpenAI.configure({ apiKey: "fixture" }).responses("fixture"), prompt: "hello" }) +const trigger = { mechanism: "trigger" } as const + +testEffect( + dynamicResponse(({ request, text, respond }) => + Effect.sync(() => { + expect(new URL(request.url).pathname).toBe("/v1/responses") + expect(new URL(request.url).searchParams.get("trace")).toBe("request") + expect(request.headers.authorization).toBe("Bearer fixture") + expect(request.headers["chatgpt-account-id"]).toBe("fixture-account") + expect(request.headers["x-codex-beta-features"]).toBe("remote_compaction_v2") + const body = JSON.parse(text) + expect(body).toMatchObject({ + model: "fixture", + stream: true, + store: false, + instructions: "Keep the context", + parallel_tool_calls: true, + prompt_cache_key: "session-key", + service_tier: "priority", + reasoning: { effort: "high", summary: "auto" }, + prompt_cache_retention: "24h", + prompt_cache_options: { mode: "session", ttl: "1h" }, + input: [{ role: "user", content: [{ type: "input_text", text: "hello" }] }, { type: "compaction_trigger" }], + }) + expect(body.tools).toHaveLength(1) + expect(body.tools[0].name).toBe("lookup") + expect(body.tool_choice).toBeUndefined() + expect(body.context_management).toBeUndefined() + expect(body.text).toBeUndefined() + expect(body.max_output_tokens).toBeUndefined() + expect(body.previous_response_id).toBeUndefined() + return respond( + sseEvents({ + type: "response.completed", + response: { + id: "resp_1", + output: [checkpoint], + usage: { + input_tokens: 100, + input_tokens_details: { cached_tokens: 40 }, + output_tokens: 5, + total_tokens: 105, + }, + }, + }), + { headers: { "content-type": "text/event-stream" } }, + ) + }), + ), +).effect("trigger uses normal request preparation, configured deployment, and supplied subscription headers", () => + Effect.gen(function* () { + const calls: string[] = [] + const input = LLM.request({ + model: request.model, + system: "Keep the context", + prompt: "hello", + promptCacheKey: "session-key", + tools: [{ name: "lookup", description: "Lookup", inputSchema: { type: "object", properties: {} } }], + toolChoice: { type: "tool", name: "lookup" }, + generation: { maxTokens: 1 }, + providerOptions: { + store: true, + reasoningEffort: "high", + reasoningSummary: "auto", + contextManagement: [{ type: "compaction" }], + }, + http: { + headers: { "chatgpt-account-id": "fixture-account", "x-codex-beta-features": "remote_compaction_v2" }, + query: { trace: "request" }, + body: { + service_tier: "priority", + prompt_cache_retention: "24h", + prompt_cache_options: { mode: "session", ttl: "1h" }, + store: true, + stream: false, + text: { format: { type: "json_object" } }, + tool_choice: "required", + }, + }, + }) + const original = LLMRequest.input(input) + const result = yield* LLMClient.compact(input, { + ...trigger, + http: (request, next) => { + calls.push("http") + return next(request) + }, + }) + expect(result.checkpoint).toMatchObject({ + type: "compaction", + provider: "openai", + id: "cmp_1", + encrypted: "opaque", + }) + expect(result.responseID).toBe("resp_1") + expect(result.usage).toMatchObject({ + inputTokens: 100, + outputTokens: 5, + totalTokens: 105, + cacheReadInputTokens: 40, + }) + expect(LLMRequest.input(input)).toEqual(original) + expect(calls).toEqual(["http"]) + }), +) + +const idless = { type: "compaction", encrypted_content: "opaque" } +testEffect( + fixedResponse( + sseEvents( + { + type: "response.output_item.done", + output_index: 0, + item: { type: "function_call", id: "fc_1", name: "unexpected", arguments: "not JSON" }, + }, + { type: "response.output_text.delta", delta: "do not expose this" }, + { type: "response.output_item.added", output_index: 1, item: { type: "compaction" } }, + { type: "response.output_item.done", output_index: 1, item: idless }, + { type: "response.output_item.done", output_index: 1, item: idless }, + { + type: "response.completed", + response: { + id: "resp_1", + output: [{ type: "function_call", id: "fc_1", name: "unexpected", arguments: "not JSON" }, idless], + }, + }, + ), + ), +).effect("correlates an ID-less checkpoint across repeated events and ignores other output", () => + Effect.gen(function* () { + const result = yield* LLMClient.compact(request, trigger) + expect(result.checkpoint.encrypted).toBe("opaque") + expect(result.checkpoint.id).toBeString() + expect("message" in result).toBe(false) + }), +) + +for (const [name, events] of Object.entries({ + missing: [{ type: "response.completed", response: { id: "resp_1", output: [] } }], + multiple: [ + { type: "response.completed", response: { id: "resp_1", output: [checkpoint, { ...checkpoint, id: "cmp_2" }] } }, + ], + duplicateSlots: [{ type: "response.completed", response: { id: "resp_1", output: [checkpoint, checkpoint] } }], + malformed: [ + { type: "response.completed", response: { id: "resp_1", output: [{ type: "compaction", id: "cmp_1" }] } }, + ], + noResponseID: [{ type: "response.completed", response: { output: [checkpoint] } }], + incomplete: [ + { type: "response.output_item.done", item: checkpoint }, + { type: "response.incomplete", response: { id: "resp_1", incomplete_details: { reason: "max_output_tokens" } } }, + ], + failed: [ + { type: "response.output_item.done", item: checkpoint }, + { type: "response.failed", response: { id: "resp_1", error: { code: "server_error", message: "failed" } } }, + ], +})) { + const wire = events.map((event) => ({ ...event, fixture_extra: "preserved" })) + testEffect( + fixedResponse(sseEvents(...wire), { headers: { "content-type": "text/event-stream", "x-fixture": "preserved" } }), + ).effect(`rejects ${name} checkpoint response and preserves original error context`, () => + Effect.gen(function* () { + const error = yield* LLMClient.compact(request, trigger).pipe(Effect.flip) + expect(error.reason.body).toBe(JSON.stringify(wire.at(-1))) + expect(error.reason.http).toMatchObject({ status: 200, headers: { "x-fixture": "preserved" } }) + }), + ) +} + +testEffect(fixedResponse(sseEvents({ type: "response.output_item.done", item: checkpoint }))).effect( + "rejects clean EOF without response.completed", + () => + Effect.gen(function* () { + const error = yield* LLMClient.compact(request, trigger).pipe(Effect.flip) + expect(error.reason._tag).toBe("InvalidProviderOutput") + }), +) +for (const body of [{ input: [] }, { previous_response_id: "stale" }]) { + testEffect(dynamicResponse(() => Effect.die("Must reject before sending"))).effect( + `rejects caller-supplied ${Object.keys(body)[0]} before sending trigger`, + () => + Effect.gen(function* () { + const error = yield* LLMClient.compact(LLMRequest.update(request, { http: { body } }), trigger).pipe( + Effect.flip, + ) + expect(error.reason._tag).toBe("InvalidRequest") + }), + ) +} + +testEffect(dynamicResponse(() => Effect.die("Must reject before sending"))).effect( + "endpoint support does not grant trigger support, including for untyped callers", + () => + Effect.gen(function* () { + const client = yield* LLMClient.Service + const unsupported = LLM.request({ model: Azure.configure({ resourceName: "fixture" }).responses("fixture") }) + expect(LLMClient.canCompact(unsupported)).toBe(true) + expect(LLMClient.canCompact(unsupported, trigger)).toBe(false) + // @ts-expect-error Exercise untyped consumers; the service must reject before sending. + const serviceError = yield* client.compact(unsupported, trigger).pipe(Effect.flip) + expect(serviceError.reason._tag).toBe("UnsupportedOperation") + // @ts-expect-error Exercise an unknown mechanism supplied by JavaScript. + const error = yield* LLMClient.compact(request, { mechanism: "other" }).pipe(Effect.flip) + expect(error.reason._tag).toBe("InvalidRequest") + }), +) + +testEffect(fixedResponse("must not use HTTP")).effect( + "invalid checkpoint output is rejected before acknowledging channel completion", + () => + Effect.gen(function* () { + let completed = 0 + const operation = LLMClient.compact(request, { + mechanism: "trigger", + webSocket: { + execute: () => + Effect.succeed({ + frames: Stream.make( + JSON.stringify({ + type: "response.completed", + response: { id: "resp_1", output: [] }, + }), + ), + complete: Effect.sync(() => { + completed++ + }), + }), + }, + }) + const error = yield* operation.pipe(Effect.flip) + expect(error.reason._tag).toBe("InvalidProviderOutput") + expect(completed).toBe(0) + }), +) diff --git a/packages/ai/test/testing.test.ts b/packages/ai/test/testing.test.ts index f7ed7db7fe4..0839e15ed31 100644 --- a/packages/ai/test/testing.test.ts +++ b/packages/ai/test/testing.test.ts @@ -3,6 +3,7 @@ import { AIError, CompactionPart, CompactionResponse, + CompactionCheckpointResponse, LanguageModel, LLM, LLMClient, @@ -79,6 +80,31 @@ describe("TestLLM legacy client", () => { }) describe("TestLLM first-class client", () => { + it.effect("returns scripted checkpoints only for trigger compaction", () => + Effect.gen(function* () { + const client = yield* TestLLM.Test + const request = LLM.request({ model: OpenAI.configure().responses("fixture"), prompt: "hello" }) + const checkpoint = new CompactionCheckpointResponse({ + checkpoint: { type: "compaction", provider: ProviderID.make("openai"), encrypted: "opaque" }, + responseID: "resp_fixture", + }) + const endpoint = new CompactionResponse({ replacement: [] }) + yield* client.push(checkpoint, endpoint, checkpoint, checkpoint) + const operation = LLMClient.compact(request, { mechanism: "trigger" }) + expect(yield* client.requests()).toEqual([]) + expect(yield* operation).toBe(checkpoint) + expect(yield* client.compact(request, { mechanism: "trigger" }).pipe(Effect.catchDefect(Effect.succeed))).toBe( + "TestLLM trigger compaction requires a CompactionCheckpointResponse", + ) + expect(yield* client.compact(request).pipe(Effect.catchDefect(Effect.succeed))).toBe( + "TestLLM compaction requires a CompactionResponse", + ) + expect(yield* client.generate(request).pipe(Effect.catchDefect(Effect.succeed))).toBe( + "TestLLM generation requires an event response", + ) + }), + ) + it.effect("rejects response fixtures for the wrong operation", () => Effect.gen(function* () { const client = yield* TestLLM.Test diff --git a/packages/ai/tsconfig.types.json b/packages/ai/tsconfig.types.json index dc73644b340..bcd7878f6ba 100644 --- a/packages/ai/tsconfig.types.json +++ b/packages/ai/tsconfig.types.json @@ -5,5 +5,5 @@ "noEmit": true, "rootDir": "." }, - "include": ["test/**/*.types.ts", "test/testing.test.ts"] + "include": ["test/**/*.types.ts", "test/testing.test.ts", "test/provider/checkpoint.test.ts"] } diff --git a/packages/core/test/lib/websocket-server.ts b/packages/core/test/lib/websocket-server.ts index f1897759c0f..b67e67c72ba 100644 --- a/packages/core/test/lib/websocket-server.ts +++ b/packages/core/test/lib/websocket-server.ts @@ -19,6 +19,7 @@ export interface WebSocketServerFixture { } export interface WebSocketServerOptions { + readonly http?: (request: Request) => Response | Promise readonly upgrade?: (request: Request) => boolean readonly open?: (socket: Bun.ServerWebSocket) => void readonly message?: (socket: Bun.ServerWebSocket, message: string | Buffer) => void @@ -34,6 +35,7 @@ export const makeWebSocketServer = (options: WebSocketServerOptions = {}) => port: 0, fetch(request, server) { state.headers.push(Object.fromEntries(request.headers.entries())) + if (request.headers.get("upgrade") !== "websocket" && options.http) return options.http(request) if ((options.upgrade?.(request) ?? true) && server.upgrade(request, { data: { id: connection++ } })) return undefined return new Response("WebSocket upgrade required", { diff --git a/packages/core/test/session-checkpoint-transport.test.ts b/packages/core/test/session-checkpoint-transport.test.ts new file mode 100644 index 00000000000..b2b2a42c55a --- /dev/null +++ b/packages/core/test/session-checkpoint-transport.test.ts @@ -0,0 +1,172 @@ +import { expect } from "bun:test" +import { LLM, LLMRequest, Message } from "@opencode-ai/ai" +import { LLMClient, RequestExecutor } from "@opencode-ai/ai/route" +import { configure } from "@opencode-ai/ai/providers/openai" +import { SessionModelTransport } from "../src/session/model-transport" +import { WebSocketConstructor } from "../src/effect/websocket-constructor" +import { Session } from "@opencode-ai/schema/session" +import { Effect, Layer, Schema } from "effect" +import { FetchHttpClient } from "effect/unstable/http" +import { testEffect } from "./lib/effect" +import { makeWebSocketServer } from "./lib/websocket-server" + +const runtime = Layer.mergeAll( + LLMClient.layer.pipe(Layer.provide(RequestExecutor.layer), Layer.provide(FetchHttpClient.layer)), + SessionModelTransport.layer.pipe(Layer.provide(WebSocketConstructor.layer)), +) +const it = testEffect(runtime) +const sessionID = Session.ID.make("ses_checkpoint_transport") +const checkpoint = { type: "compaction", encrypted_content: "opaque" } +const fixture = (mode: "success" | "invalid" | "fallback") => + Effect.gen(function* () { + const requests: Array> = [] + const http: Array> = [] + let disconnect: (() => void) | undefined + const server = yield* makeWebSocketServer({ + upgrade: () => mode !== "fallback" || disconnect === undefined, + async http(request) { + http.push(Schema.decodeUnknownSync(Schema.Record(Schema.String, Schema.Unknown))(await request.json())) + return new Response( + `data: ${JSON.stringify({ + type: "response.completed", + response: { + id: "resp_http", + output: [checkpoint], + usage: { input_tokens: 70, output_tokens: 7 }, + }, + })}\n\n`, + { headers: { "content-type": "text/event-stream", "x-response": "http" } }, + ) + }, + open(socket) { + disconnect = () => socket.close() + }, + message(socket, message) { + const body = JSON.parse(message.toString()) + requests.push(body) + const id = `resp_${requests.length}` + const send = (event: unknown) => socket.send(JSON.stringify(event)) + if (body.input.at(-1)?.type !== "compaction_trigger") { + const item = { + type: "message", + id: `msg_${requests.length}`, + role: "assistant", + content: [{ type: "output_text", text: "Hello" }], + } + send({ type: "response.created", response: { id } }) + send({ type: "response.output_item.added", output_index: 0, item }) + send({ type: "response.output_text.delta", item_id: item.id, delta: "Hello" }) + send({ type: "response.output_item.done", output_index: 0, item }) + send({ type: "response.completed", response: { id, output: [item] } }) + return + } + send({ type: "response.created", response: { id } }) + send({ type: "response.output_item.done", output_index: 0, item: checkpoint }) + send({ + type: "response.completed", + response: { + id, + output: mode === "invalid" ? [checkpoint, { ...checkpoint, encrypted_content: "second" }] : [checkpoint], + usage: { input_tokens: 100, output_tokens: 10 }, + }, + }) + }, + }) + return { + requests, + headers: server.state.headers, + http, + opens: () => server.state.opens, + disconnect: () => disconnect?.(), + request: LLM.request({ + model: configure({ + apiKey: "fixture", + baseURL: server.url.replace(/^ws/, "http").replace(/responses$/, ""), + headers: { "chatgpt-account-id": "account", "x-codex-beta-features": "remote_compaction_v2" }, + providerOptions: { parallelToolCalls: true }, + }).responses("fixture"), + prompt: "First", + promptCacheKey: "session-key", + }), + } + }) + +it.live("trigger reuses the append baseline and clears it before the next generation", () => + Effect.gen(function* () { + const server = yield* fixture("success") + const transport = yield* SessionModelTransport.Service + const webSocket = transport.bind(sessionID) + const first = yield* LLMClient.generate(server.request, { webSocket }) + const compacted = yield* LLMClient.compact( + LLMRequest.update(server.request, { + messages: [...server.request.messages, first.message], + }), + { mechanism: "trigger", webSocket }, + ) + expect(server.requests[1]).toMatchObject({ + previous_response_id: "resp_1", + input: [{ type: "compaction_trigger" }], + }) + expect(compacted.responseID).toBe("resp_2") + expect(compacted.usage).toMatchObject({ inputTokens: 100, outputTokens: 10 }) + expect(compacted.checkpoint.encrypted).toBe("opaque") + const messages = [Message.assistant(compacted.checkpoint), Message.user("Continue")] + yield* LLMClient.generate(LLMRequest.update(server.request, { messages }), { webSocket }) + expect(server.requests[2]?.previous_response_id).toBeUndefined() + expect(server.requests[2]?.input).toMatchObject([ + { type: "compaction", encrypted_content: "opaque" }, + { role: "user", content: [{ type: "input_text", text: "Continue" }] }, + ]) + expect(server.requests[1]?.stream).toBeUndefined() + expect(server.requests[1]?.store).toBe(false) + expect(server.headers[0]).toMatchObject({ + authorization: "Bearer fixture", + "chatgpt-account-id": "account", + "x-codex-beta-features": "remote_compaction_v2", + }) + expect(server.opens()).toBe(1) + expect(server.http).toHaveLength(0) + }), +) + +it.live("invalid trigger output does not retry or commit a continuation checkpoint", () => + Effect.gen(function* () { + const server = yield* fixture("invalid") + const transport = yield* SessionModelTransport.Service + const webSocket = transport.bind(sessionID) + const first = yield* LLMClient.generate(server.request, { webSocket }) + const input = LLMRequest.update(server.request, { messages: [...server.request.messages, first.message] }) + const error = yield* LLMClient.compact(input, { mechanism: "trigger", webSocket }).pipe(Effect.flip) + expect(error.reason._tag).toBe("InvalidProviderOutput") + expect(server.requests).toHaveLength(2) + expect(server.http).toHaveLength(0) + yield* LLMClient.generate(input, { webSocket }) + expect(server.requests[2]?.previous_response_id).toBeUndefined() + expect(server.requests[2]?.input).toHaveLength(2) + }), +) + +it.live("trigger recovers over SSE with complete input after the old socket closes and reconnect is rejected", () => + Effect.gen(function* () { + const server = yield* fixture("fallback") + const transport = yield* SessionModelTransport.Service + const webSocket = transport.bind(sessionID) + const first = yield* LLMClient.generate(server.request, { webSocket }) + server.disconnect() + // Let the real close event reach the connector before the next send. + yield* Effect.sleep("20 millis") + const result = yield* LLMClient.compact( + LLMRequest.update(server.request, { + messages: [...server.request.messages, first.message], + }), + { mechanism: "trigger", webSocket }, + ) + expect(result.responseID).toBe("resp_http") + expect(result.usage).toMatchObject({ inputTokens: 70, outputTokens: 7 }) + expect(server.http).toHaveLength(1) + expect(server.http[0]?.previous_response_id).toBeUndefined() + expect(server.http[0]?.stream).toBe(true) + expect(server.http[0]?.input).toHaveLength(3) + expect(server.requests).toHaveLength(1) + }), +)