diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 7806a3f10..bb325c171 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -423,6 +423,8 @@ export interface AssistantMessage { stopReason: StopReason; deferred?: DeferredHandle; errorMessage?: string; + /** False when starting a replacement request could duplicate a billable operation. */ + retryable?: boolean; rawStopReason?: string; timestamp: number; // Unix timestamp in milliseconds } diff --git a/packages/ai/src/utils/overflow.ts b/packages/ai/src/utils/overflow.ts index c1e6457ff..533892ed1 100644 --- a/packages/ai/src/utils/overflow.ts +++ b/packages/ai/src/utils/overflow.ts @@ -132,6 +132,7 @@ const NON_OVERFLOW_PATTERNS = [ * @returns true if the message indicates a context overflow */ export function isContextOverflow(message: AssistantMessage, contextWindow?: number): boolean { + if (message.retryable === false) return false; // Case 1: Check error message patterns if (message.stopReason === "error" && message.errorMessage) { // Skip messages matching known non-overflow patterns (e.g. throttling / rate-limit) diff --git a/packages/ai/src/utils/retry.ts b/packages/ai/src/utils/retry.ts index 907d19504..cdb0a598b 100644 --- a/packages/ai/src/utils/retry.ts +++ b/packages/ai/src/utils/retry.ts @@ -227,6 +227,7 @@ export function isRetryableProviderErrorMessage(errorMessage: string): boolean { */ export function isRetryableAssistantError(message: AssistantMessage): boolean { return ( + message.retryable !== false && message.stopReason === "error" && message.errorMessage !== undefined && isRetryableProviderErrorMessage(message.errorMessage) diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 9cb758052..0826753bd 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -3506,6 +3506,7 @@ export class AgentSession { signal?: AbortSignal, deferCommit?: (transaction: CompactionAppendTransaction) => void, ): Promise { + if (assistantMessage.retryable === false) return false; const settings = this.settingsManager.getCompactionSettings(); if (!settings.enabled) return false; diff --git a/packages/coding-agent/src/core/hosted-model-gateway.ts b/packages/coding-agent/src/core/hosted-model-gateway.ts new file mode 100644 index 000000000..546742577 --- /dev/null +++ b/packages/coding-agent/src/core/hosted-model-gateway.ts @@ -0,0 +1,293 @@ +import { createHash, randomUUID } from "node:crypto"; +import { constants } from "node:fs"; +import { lstat, mkdir, open, realpath, unlink } from "node:fs/promises"; +import { dirname, isAbsolute, join, resolve } from "node:path"; +import { + type Api, + type AssistantMessage, + type AssistantMessageEventStream, + createAssistantMessageEventStream, + type Model, + type ModelsRequestTransforms, + type ProviderRequestOptions, + type StreamOptions, +} from "@diffexai/diffex-ai"; + +export interface HostedModelGatewayOptions { + baseUrl: string; + credential: string; + provider: string; + modelId: string; + maxOutputTokens: number; + /** Persistent runtime directory containing private operation metadata. */ + stateDirectory: string; +} + +/** A gateway failure must never trigger a new billable generation automatically. */ +export class HostedModelGatewayError extends Error { + readonly code: string; + readonly operationId: string | undefined; + + constructor(code: string, operationId?: string) { + super(`gateway_${code}${operationId ? ` (operation ${operationId})` : ""}`); + this.name = "HostedModelGatewayError"; + this.code = code; + this.operationId = operationId; + } +} + +async function syncDirectory(directory: string): Promise { + const handle = await open(directory, constants.O_RDONLY | constants.O_NOFOLLOW); + try { + await handle.sync(); + } finally { + await handle.close(); + } +} + +export class HostedModelGateway { + readonly options: Readonly; + private readonly directory: string; + + private constructor(options: HostedModelGatewayOptions, directory: string) { + this.options = Object.freeze({ ...options }); + this.directory = directory; + } + + static async create(options: HostedModelGatewayOptions): Promise { + const url = new URL(options.baseUrl); + const loopback = ["127.0.0.1", "[::1]"].includes(url.hostname); + if ( + (url.protocol !== "https:" && !(url.protocol === "http:" && loopback)) || + url.username || + url.password || + url.search || + url.hash || + !/^\/v1\/runtime\/sessions\/[0-9a-f]{8}-[0-9a-f]{4}-[1-8][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}\/model\/?$/iu.test( + url.pathname, + ) || + !/^[\x21-\x7e]{1,4096}$/u.test(options.credential) || + options.provider !== "anthropic" || + !options.modelId || + !Number.isSafeInteger(options.maxOutputTokens) || + options.maxOutputTokens <= 0 || + options.maxOutputTokens > 2_147_483_647 || + !isAbsolute(options.stateDirectory) + ) { + throw new HostedModelGatewayError("invalid_configuration"); + } + const directory = resolve(options.stateDirectory, "model-operations"); + await mkdir(directory, { recursive: true, mode: 0o700 }); + const info = await lstat(directory); + if ( + !info.isDirectory() || + info.isSymbolicLink() || + (info.mode & 0o077) !== 0 || + (await realpath(directory)) !== directory + ) { + throw new HostedModelGatewayError("invalid_state_directory"); + } + // Sync all ancestors even when an interrupted earlier startup created them. + for (let current = directory; ; current = dirname(current)) { + await syncDirectory(current); + if (dirname(current) === current) break; + } + return new HostedModelGateway({ ...options, baseUrl: url.toString().replace(/\/$/u, "") }, directory); + } + + assertModel(model: Model): void { + if ( + model.provider !== this.options.provider || + model.id !== this.options.modelId || + model.api !== "anthropic-messages" || + (model.baseUrl !== this.options.baseUrl && model.baseUrl !== "https://api.anthropic.com") || + model.headers !== undefined + ) { + throw new HostedModelGatewayError("model_override_forbidden"); + } + } + + stream( + model: Model, + requestOptions: (StreamOptions & ModelsRequestTransforms & { client?: unknown }) | undefined, + start: (options: ProviderRequestOptions) => AssistantMessageEventStream, + ): AssistantMessageEventStream { + const outer = createAssistantMessageEventStream(); + const operationId = randomUUID(); + const pendingPath = join(this.directory, `${operationId}.json`); + let sent = false; + let recorded = false; + let localFailureCode: string | undefined; + const fetch: typeof globalThis.fetch = async (input, init) => { + if (sent) { + localFailureCode = "automatic_retry_forbidden"; + throw new HostedModelGatewayError(localFailureCode, operationId); + } + const request = new Request(input, init); + if (request.url !== `${this.options.baseUrl}/v1/messages` || request.method !== "POST") { + localFailureCode = "endpoint_override_forbidden"; + throw new HostedModelGatewayError(localFailureCode, operationId); + } + const body = await request.text(); + const payload: unknown = JSON.parse(body); + if (!payload || typeof payload !== "object" || Array.isArray(payload)) { + localFailureCode = "invalid_request"; + throw new HostedModelGatewayError(localFailureCode, operationId); + } + const parameters = payload as Record; + const thinking = parameters.thinking as Record | undefined; + if ( + parameters.model !== this.options.modelId || + parameters.stream !== true || + !Number.isSafeInteger(parameters.max_tokens) || + (parameters.max_tokens as number) <= 0 || + (parameters.max_tokens as number) > this.options.maxOutputTokens || + (thinking?.type === "enabled" && + (!Number.isSafeInteger(thinking.budget_tokens) || + (thinking.budget_tokens as number) >= (parameters.max_tokens as number))) + ) { + localFailureCode = "invalid_request"; + throw new HostedModelGatewayError(localFailureCode, operationId); + } + try { + const handle = await open( + pendingPath, + constants.O_WRONLY | constants.O_CREAT | constants.O_EXCL | constants.O_NOFOLLOW, + 0o600, + ); + try { + await handle.writeFile( + JSON.stringify({ + version: 1, + operationId, + model: this.options.modelId, + requestDigest: createHash("sha256").update(body).digest("hex"), + createdAt: new Date().toISOString(), + }), + ); + await handle.sync(); + } finally { + await handle.close(); + } + await syncDirectory(this.directory); + recorded = true; + } catch { + localFailureCode = "local_journal_unavailable"; + throw new HostedModelGatewayError(localFailureCode, operationId); + } + request.signal.throwIfAborted(); + sent = true; + return globalThis.fetch(request.url, { + method: "POST", + body, + signal: request.signal, + redirect: "error", + headers: { + "content-type": "application/json", + accept: "text/event-stream", + "x-api-key": this.options.credential, + "x-diffex-operation-id": operationId, + }, + }); + }; + const fail = (message: AssistantMessage, code: string): void => { + message.retryable = false; + message.errorMessage = new HostedModelGatewayError(code, operationId).message; + message.diagnostics = [ + ...(message.diagnostics ?? []), + { + type: "model_gateway_failure", + timestamp: Date.now(), + error: { name: "HostedModelGatewayError", message: message.errorMessage, code }, + details: { operationId }, + }, + ]; + }; + void (async () => { + try { + this.assertModel(model); + if (requestOptions?.maxTokens !== undefined && requestOptions.maxTokens > this.options.maxOutputTokens) { + throw new HostedModelGatewayError("invalid_request", operationId); + } + if ( + requestOptions?.fetch || + requestOptions?.client || + requestOptions?.env || + (requestOptions?.apiKey !== undefined && requestOptions.apiKey !== this.options.credential) + ) { + throw new HostedModelGatewayError("transport_override_forbidden", operationId); + } + const inner = start({ + ...requestOptions, + // This is an SDK auth placeholder; only fetch() handles the gateway credential. + apiKey: "diffex-hosted-model-gateway", + headers: undefined, + fetch, + maxRetries: 0, + }); + for await (const event of inner) { + if (event.type === "start") { + event.partial.diagnostics = [ + ...(event.partial.diagnostics ?? []), + { + type: "model_gateway_operation", + timestamp: Date.now(), + details: { operationId }, + }, + ]; + } else if (event.type === "error") { + const code = + localFailureCode ?? + event.error.errorMessage?.match(/gateway_([a-z_]+)/u)?.[1] ?? + "transport_or_stream_interrupted"; + fail(event.error, code); + } else if (event.type === "done" && recorded) { + try { + await unlink(pendingPath); + await syncDirectory(this.directory); + } catch { + // A cleanup failure cannot invalidate an already completed generation. + event.message.diagnostics = [ + ...(event.message.diagnostics ?? []), + { + type: "model_gateway_pending_cleanup_failed", + timestamp: Date.now(), + error: { + name: "HostedModelGatewayError", + code: "local_journal_cleanup_unavailable", + message: "Pending operation metadata cleanup failed", + }, + details: { operationId }, + }, + ]; + } + } + outer.push(event); + } + outer.end(await inner.result()); + } catch (error) { + const message: AssistantMessage = { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: requestOptions?.signal?.aborted ? "aborted" : "error", + timestamp: Date.now(), + }; + fail(message, error instanceof HostedModelGatewayError ? error.code : "transport_or_stream_interrupted"); + outer.push({ type: "error", reason: message.stopReason as "error" | "aborted", error: message }); + outer.end(message); + } + })(); + return outer; + } +} diff --git a/packages/coding-agent/src/core/model-runtime.ts b/packages/coding-agent/src/core/model-runtime.ts index da58971b9..2982f7d8d 100644 --- a/packages/coding-agent/src/core/model-runtime.ts +++ b/packages/coding-agent/src/core/model-runtime.ts @@ -37,9 +37,11 @@ import { type StreamOptions, } from "@diffexai/diffex-ai"; import * as builtinProviderCatalog from "@diffexai/diffex-ai/providers/all"; +import { anthropicProvider } from "@diffexai/diffex-ai/providers/anthropic"; import { getAgentDir } from "../config.ts"; import { operationSignal, raceWithAbortSignal } from "../utils/abort.ts"; import { AuthStorage as DefaultAuthStorage } from "./auth-storage.ts"; +import { HostedModelGateway, type HostedModelGatewayOptions } from "./hosted-model-gateway.ts"; import { ModelConfig } from "./model-config.ts"; import { FileModelsStore, InMemoryCodingAgentModelsStore } from "./models-store.ts"; import { @@ -64,6 +66,8 @@ interface ModelRuntimeSnapshot { } export interface CreateModelRuntimeOptions { + /** Route every hosted model call through the authenticated metering gateway. */ + hostedGateway?: HostedModelGatewayOptions; /** Credential storage. Defaults to the file at authPath. */ credentials?: CredentialStore; authPath?: string; @@ -137,6 +141,7 @@ export class ModelRuntime implements Models { private readonly compositionErrors = new Map(); private readonly modelsPath: string | undefined; private readonly modelNetworkEnabled: boolean; + private readonly hostedGateway: HostedModelGateway | undefined; private config: ModelConfig; private snapshot: ModelRuntimeSnapshot = { all: [], @@ -158,11 +163,13 @@ export class ModelRuntime implements Models { modelsStore: ModelsStore, providers: readonly Provider[], modelNetworkEnabled: boolean, + hostedGateway?: HostedModelGateway, ) { this.credentials = credentials; this.config = config; this.modelsPath = modelsPath; this.modelNetworkEnabled = modelNetworkEnabled; + this.hostedGateway = hostedGateway; this.defaultBuiltins = new Map(providers.map((provider) => [provider.id, provider])); for (const [providerId, provider] of this.defaultBuiltins) this.builtins.set(providerId, provider); this.models = createModels({ credentials, modelsStore }); @@ -170,7 +177,9 @@ export class ModelRuntime implements Models { } static async create(options: CreateModelRuntimeOptions = {}): Promise { + const gateway = options.hostedGateway ? await HostedModelGateway.create(options.hostedGateway) : undefined; const credentials = new RuntimeCredentials(options.credentials ?? DefaultAuthStorage.create(options.authPath)); + if (gateway) credentials.setRuntimeApiKey(gateway.options.provider, gateway.options.credential); const modelsPath = options.modelsPath === null ? undefined : (options.modelsPath ?? join(getAgentDir(), "models.json")); const config = await ModelConfig.load(modelsPath); @@ -180,20 +189,46 @@ export class ModelRuntime implements Models { ? new FileModelsStore(options.modelsStorePath ?? join(dirname(modelsPath), "models-store.json")) : new InMemoryCodingAgentModelsStore()); const builtinModelDataGeneratedAt = builtinProviderCatalog.getBuiltinModelDataGeneratedAt(); - const providers = builtinProviderCatalog + let providers = builtinProviderCatalog .builtinProviders() .map((provider) => provider.id === "radius" ? provider : withRemoteCatalog(provider, options.catalogBaseUrl, builtinModelDataGeneratedAt), ); + if (gateway) { + const provider = anthropicProvider(); + const model = provider.getModels().find((candidate) => candidate.id === gateway.options.modelId); + if (!model) throw new Error("Hosted gateway model is not in the Anthropic catalog"); + providers = [ + { + ...provider, + getModels: () => [ + { + ...model, + maxTokens: Math.min(model.maxTokens, gateway.options.maxOutputTokens), + input: ["text"], + baseUrl: gateway.options.baseUrl, + headers: undefined, + compat: { + ...model.compat, + supportsEagerToolInputStreaming: false, + supportsStrictTools: false, + supportsToolReferences: false, + }, + }, + ], + }, + ]; + } const runtime = new ModelRuntime( credentials, config, modelsPath, modelsStore, providers, - process.env.DIFFEX_OFFLINE === undefined, + !gateway && process.env.DIFFEX_OFFLINE === undefined, + gateway, ); runtime.configureRadiusProviders(); runtime.rebuildProviders(); @@ -243,6 +278,12 @@ export class ModelRuntime implements Models { } private recomposeProvider(providerId: string): void { + if (this.hostedGateway) { + const provider = this.defaultBuiltins.get(providerId); + if (provider) this.models.setProvider(provider); + else this.models.deleteProvider(providerId); + return; + } const base = this.nativeExtensionProviders.get(providerId) ?? this.builtins.get(providerId); const extension = this.extensionProviders.get(providerId); if (!base && !this.config.getProvider(providerId) && !extension) { @@ -473,6 +514,14 @@ export class ModelRuntime implements Models { providerOrModel: string | Model, overrides: ModelRuntimeAuthOverrides = {}, ): Promise { + if (this.hostedGateway) { + if (typeof providerOrModel !== "string") this.hostedGateway.assertModel(providerOrModel); + else if (providerOrModel !== this.hostedGateway.options.provider) return undefined; + return { + auth: { apiKey: this.hostedGateway.options.credential, baseUrl: this.hostedGateway.options.baseUrl }, + source: "hosted model gateway", + }; + } if (typeof providerOrModel === "string") return this.models.getAuth(providerOrModel, overrides); const resolution = await this.models.getAuth(providerOrModel, overrides); if (!resolution) return undefined; @@ -534,6 +583,7 @@ export class ModelRuntime implements Models { } setRuntimeApiKey(providerId: string, apiKey: string, options: AuthOperationOptions = {}): Promise { + if (this.hostedGateway) return Promise.reject(new Error("Hosted gateway does not allow credential overrides")); const signal = operationSignal(options.signal); return this.enqueueCredentialOperation(providerId, signal, async () => { this.credentials.setRuntimeApiKey(providerId, apiKey); @@ -547,6 +597,7 @@ export class ModelRuntime implements Models { } removeRuntimeApiKey(providerId: string, options: AuthOperationOptions = {}): Promise { + if (this.hostedGateway) return Promise.reject(new Error("Hosted gateway does not allow credential overrides")); const signal = operationSignal(options.signal); return this.enqueueCredentialOperation(providerId, signal, async () => { this.credentials.removeRuntimeApiKey(providerId); @@ -612,6 +663,17 @@ export class ModelRuntime implements Models { context: Context, options?: ModelsApiStreamOptions, ): AssistantMessageEventStream { + if (this.hostedGateway) { + const provider = this.defaultBuiltins.get("anthropic"); + if (!provider) throw new Error("Hosted Anthropic provider is missing"); + return this.hostedGateway.stream(model, options, (requestOptions) => + provider.stream( + { ...model, baseUrl: this.hostedGateway?.options.baseUrl ?? model.baseUrl }, + context, + requestOptions as ApiStreamOptions, + ), + ); + } return lazyStream( model, async () => { @@ -638,6 +700,17 @@ export class ModelRuntime implements Models { } streamSimple(model: Model, context: Context, options?: ModelsSimpleStreamOptions): AssistantMessageEventStream { + if (this.hostedGateway) { + const provider = this.defaultBuiltins.get("anthropic"); + if (!provider) throw new Error("Hosted Anthropic provider is missing"); + return this.hostedGateway.stream(model, options, (requestOptions) => + provider.streamSimple( + { ...model, baseUrl: this.hostedGateway?.options.baseUrl ?? model.baseUrl }, + context, + requestOptions as SimpleStreamOptions, + ), + ); + } return lazyStream( model, async () => { @@ -683,6 +756,7 @@ export class ModelRuntime implements Models { } login(providerId: string, type: AuthType, interaction: AuthInteraction): Promise { + if (this.hostedGateway) return Promise.reject(new Error("Hosted gateway does not allow credential overrides")); const signal = operationSignal(interaction.signal); return this.enqueueCredentialOperation(providerId, signal, async () => { const credential = await this.models.login(providerId, type, { ...interaction, signal }); @@ -692,6 +766,7 @@ export class ModelRuntime implements Models { } logout(providerId: string, options: AuthOperationOptions = {}): Promise { + if (this.hostedGateway) return Promise.reject(new Error("Hosted gateway does not allow credential overrides")); const signal = operationSignal(options.signal); return this.enqueueCredentialOperation(providerId, signal, async () => { await this.models.logout(providerId, { signal }); @@ -743,6 +818,7 @@ export class ModelRuntime implements Models { } registerNativeProvider(provider: Provider): void { + if (this.hostedGateway) throw new Error("Hosted gateway does not allow provider overrides"); if (!provider.id.trim()) throw new Error("Provider id must not be empty."); this.extensionProviders.delete(provider.id); this.nativeExtensionProviders.set(provider.id, provider); @@ -752,6 +828,7 @@ export class ModelRuntime implements Models { } registerProvider(providerId: string, config: ProviderConfigInput): void { + if (this.hostedGateway) throw new Error("Hosted gateway does not allow provider overrides"); // Validate the incoming registration on its own, like the legacy registry: // a broken re-registration must throw without touching the stored config. validateExtensionProvider(providerId, this.builtins.get(providerId), this.config.getProvider(providerId), config); diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 4faec9167..d2bc23022 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -172,6 +172,7 @@ export { } from "./core/extensions/index.ts"; // Footer data provider (git branch + extension statuses - data not otherwise available to extensions) export type { ReadonlyFooterDataProvider } from "./core/footer-data-provider.ts"; +export { HostedModelGatewayError, type HostedModelGatewayOptions } from "./core/hosted-model-gateway.ts"; export { convertToLlm } from "./core/messages.ts"; export { ModelRegistry } from "./core/model-registry.ts"; export { diff --git a/packages/coding-agent/test/hosted-model-gateway.test.ts b/packages/coding-agent/test/hosted-model-gateway.test.ts new file mode 100644 index 000000000..0c09e2596 --- /dev/null +++ b/packages/coding-agent/test/hosted-model-gateway.test.ts @@ -0,0 +1,379 @@ +import { chmod, mkdtemp, readdir, readFile, realpath, rm, stat, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { type Context, isRetryableAssistantError, retryAssistantCall } from "@diffexai/diffex-ai"; +import { afterEach, describe, expect, test, vi } from "vitest"; +import { AuthStorage } from "../src/core/auth-storage.ts"; +import { generateBranchSummary } from "../src/core/compaction/branch-summarization.ts"; +import { generateSummaryWithUsage } from "../src/core/compaction/compaction.ts"; +import { ModelRuntime } from "../src/core/model-runtime.ts"; +import { createAgentSession } from "../src/core/sdk.ts"; +import { SessionManager } from "../src/core/session-manager.ts"; +import { SettingsManager } from "../src/core/settings-manager.ts"; +import { captureSubagentSessionConfig, createSubagentSession } from "../src/core/subagents/services.ts"; + +const sessionId = "2f1c9e00-aaaa-4000-8000-000000000001"; +const baseUrl = `https://cloud.fixture/v1/runtime/sessions/${sessionId}/model`; +const modelId = "claude-sonnet-4-5"; +const context: Context = { messages: [{ role: "user", content: "Say hello", timestamp: 1 }] }; +const roots: string[] = []; + +afterEach(async () => { + vi.restoreAllMocks(); + await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true }))); +}); + +async function createRuntime() { + const directory = await realpath(await mkdtemp(join(tmpdir(), "gateway-runtime-test-"))); + roots.push(directory); + const runtime = await ModelRuntime.create({ + credentials: AuthStorage.inMemory(), + modelsPath: null, + hostedGateway: { + baseUrl, + credential: "gateway-fixture-credential", + provider: "anthropic", + modelId, + stateDirectory: directory, + maxOutputTokens: 8192, + }, + }); + const model = runtime.getModel("anthropic", modelId); + if (!model) throw new Error("Missing fixture model"); + return { runtime, model, directory }; +} + +function completion(): Response { + const events = [ + { + type: "message_start", + message: { + id: "msg_fixture", + type: "message", + role: "assistant", + model: modelId, + content: [], + stop_reason: null, + stop_sequence: null, + usage: { input_tokens: 8, output_tokens: 0, cache_read_input_tokens: 0, cache_creation_input_tokens: 0 }, + }, + }, + { type: "content_block_start", index: 0, content_block: { type: "text", text: "" } }, + { type: "content_block_delta", index: 0, delta: { type: "text_delta", text: "Hello" } }, + { type: "content_block_stop", index: 0 }, + { + type: "message_delta", + delta: { stop_reason: "end_turn", stop_sequence: null, stop_details: null }, + usage: { output_tokens: 2 }, + }, + { type: "message_stop" }, + ]; + return new Response(events.map((event) => `event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`).join(""), { + headers: { "content-type": "text/event-stream" }, + }); +} + +describe("hosted model gateway", () => { + test("persists a private operation before send, constrains headers, and uses a unique ID for identical calls", async () => { + const { runtime, model, directory } = await createRuntime(); + const ids: string[] = []; + const fetch = vi.spyOn(globalThis, "fetch").mockImplementation(async (input, init) => { + const request = new Request(input, init); + expect(request.url).toBe(`${baseUrl}/v1/messages`); + expect(request.headers.get("x-api-key")).toBe("gateway-fixture-credential"); + expect(request.headers.get("authorization")).toBeNull(); + expect(request.headers.get("x-untrusted")).toBeNull(); + expect(request.redirect).toBe("error"); + const id = request.headers.get("x-diffex-operation-id")!; + expect(id).toMatch(/^[0-9a-f-]{36}$/u); + ids.push(id); + const pending = join(directory, "model-operations", `${id}.json`); + const record = JSON.parse(await readFile(pending, "utf8")); + expect(record).toMatchObject({ version: 1, operationId: id, model: modelId }); + expect(record.requestDigest).toMatch(/^[0-9a-f]{64}$/u); + expect(JSON.stringify(record)).not.toContain("Say hello"); + expect((await stat(pending)).mode & 0o777).toBe(0o600); + const payload = await request.json(); + expect(payload).toMatchObject({ model: modelId, stream: true, max_tokens: 16 }); + return completion(); + }); + for (let i = 0; i < 2; i++) { + const result = await runtime.completeSimple(model, context, { + maxTokens: 16, + headers: { "x-untrusted": "ignored" }, + maxRetries: 10, + }); + expect(result.stopReason).toBe("stop"); + expect(result.usage).toMatchObject({ input: 8, output: 2 }); + } + expect(fetch).toHaveBeenCalledTimes(2); + expect(ids[0]).not.toBe(ids[1]); + expect(await readdir(join(directory, "model-operations"))).toEqual([]); + expect(runtime.getAvailableSnapshot().map((entry) => entry.id)).toEqual([modelId]); + }); + + test.each(["in_progress", "unresolved", "completed_output_unavailable", "unavailable"])( + "does not retry gateway %s at provider or summary retry layers", + async (code) => { + const { runtime, model, directory } = await createRuntime(); + const fetch = vi + .spyOn(globalThis, "fetch") + .mockResolvedValue( + Response.json( + { type: "error", error: { type: "api_error", message: `gateway_${code}` } }, + { status: code === "unavailable" ? 503 : 409 }, + ), + ); + const produce = vi.fn(() => runtime.completeSimple(model, context, { maxRetries: 4 })); + const result = await retryAssistantCall(produce, { enabled: true, maxRetries: 4, baseDelayMs: 1 }, undefined); + expect(result.stopReason).toBe("error"); + expect(result.retryable).toBe(false); + expect(result.errorMessage).toContain(`gateway_${code}`); + expect(result.diagnostics?.at(-1)).toMatchObject({ type: "model_gateway_failure", error: { code } }); + expect(isRetryableAssistantError(result)).toBe(false); + expect(produce).toHaveBeenCalledTimes(1); + expect(fetch).toHaveBeenCalledTimes(1); + expect(await readdir(join(directory, "model-operations"))).toHaveLength(1); + }, + ); + + test.each(["socket hang up", "Request timed out", "fetch failed"])( + "retains pending evidence and never retries ambiguous %s", + async (message) => { + const { runtime, model, directory } = await createRuntime(); + const fetch = vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error(message)); + const result = await retryAssistantCall( + () => runtime.completeSimple(model, context, { maxRetries: 10 }), + { enabled: true, maxRetries: 10, baseDelayMs: 1 }, + undefined, + ); + expect(result.retryable).toBe(false); + expect(fetch).toHaveBeenCalledTimes(1); + expect(await readdir(join(directory, "model-operations"))).toHaveLength(1); + }, + ); + + test("rejects endpoint, model, direct credentials and custom transports before send", async () => { + const { runtime, model } = await createRuntime(); + const fetch = vi.spyOn(globalThis, "fetch"); + for (const override of [ + { ...model, baseUrl: "https://api.other-provider.test" }, + { ...model, id: "another" }, + ]) { + expect((await runtime.completeSimple(override, context)).errorMessage).toContain( + "gateway_model_override_forbidden", + ); + } + expect((await runtime.completeSimple(model, context, { apiKey: "direct-provider-key" })).errorMessage).toContain( + "gateway_transport_override_forbidden", + ); + expect((await runtime.completeSimple(model, context, { fetch: globalThis.fetch })).errorMessage).toContain( + "gateway_transport_override_forbidden", + ); + expect(() => runtime.registerNativeProvider(runtime.getProvider("anthropic")!)).toThrow("provider overrides"); + expect(fetch).not.toHaveBeenCalled(); + }); + + test("does not replay pending calls when the runtime restarts", async () => { + const { runtime, model, directory } = await createRuntime(); + const fetch = vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("socket hang up")); + await runtime.completeSimple(model, context); + const files = await readdir(join(directory, "model-operations")); + await ModelRuntime.create({ + credentials: AuthStorage.inMemory(), + modelsPath: null, + hostedGateway: { + baseUrl, + credential: "replacement-gateway-credential", + provider: "anthropic", + modelId, + stateDirectory: directory, + maxOutputTokens: 8192, + }, + }); + expect(fetch).toHaveBeenCalledTimes(1); + expect(await readdir(join(directory, "model-operations"))).toEqual(files); + }); + + test("keeps dropped terminal streams non-retryable and enforces the server output cap before send", async () => { + const { runtime, model, directory } = await createRuntime(); + const fetch = vi.spyOn(globalThis, "fetch").mockImplementation(async () => { + const body = await completion().text(); + return new Response(body.replace('event: message_stop\ndata: {"type":"message_stop"}\n\n', ""), { + headers: { "content-type": "text/event-stream" }, + }); + }); + const result = await retryAssistantCall( + () => runtime.completeSimple(model, context), + { enabled: true, maxRetries: 4, baseDelayMs: 1 }, + undefined, + ); + expect(result.retryable).toBe(false); + expect(result.errorMessage).toContain("gateway_transport_or_stream_interrupted"); + expect(fetch).toHaveBeenCalledTimes(1); + expect(await readdir(join(directory, "model-operations"))).toHaveLength(1); + const excessive = await runtime.completeSimple(model, context, { maxTokens: 9000 }); + expect(excessive.errorMessage).toContain("gateway_invalid_request"); + expect(fetch).toHaveBeenCalledTimes(1); + expect(await readdir(join(directory, "model-operations"))).toHaveLength(1); + }); + + test("routes actual parent, compaction, branch summary and child agent calls through one gateway", async () => { + const { runtime, model, directory } = await createRuntime(); + const requests: Array<{ id: string; payload: Record }> = []; + const fetch = vi.spyOn(globalThis, "fetch").mockImplementation(async (input, init) => { + const request = new Request(input, init); + expect(request.url).toBe(`${baseUrl}/v1/messages`); + expect(request.headers.get("x-api-key")).toBe("gateway-fixture-credential"); + requests.push({ + id: request.headers.get("x-diffex-operation-id")!, + payload: (await request.json()) as Record, + }); + return completion(); + }); + const { session: parent } = await createAgentSession({ + cwd: directory, + agentDir: join(directory, "agent"), + modelRuntime: runtime, + model, + sessionManager: SessionManager.inMemory(directory), + settingsManager: SettingsManager.inMemory({ + compaction: { enabled: false }, + retry: { enabled: true, maxRetries: 4, baseDelayMs: 1 }, + }), + }); + let child: Awaited> | undefined; + try { + await parent.prompt("A parent request"); + const summary = await generateSummaryWithUsage( + parent.agent.state.messages, + model, + 16384, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + parent.agent.streamFunction, + ); + expect(summary.text).toBe("Hello"); + const branch = await generateBranchSummary(parent.sessionManager.getBranch(), { + model, + signal: new AbortController().signal, + streamFn: parent.agent.streamFunction, + }); + expect(branch.summary).toContain("Hello"); + child = await createSubagentSession(captureSubagentSessionConfig(parent, join(directory, "agent")), { + childName: "gateway_child", + }); + expect(child.modelRuntime).toBe(runtime); + await child.prompt("A child request"); + expect(fetch).toHaveBeenCalledTimes(4); + expect(new Set(requests.map((request) => request.id)).size).toBe(4); + for (const request of requests) expect(request.payload).toMatchObject({ model: modelId, stream: true }); + expect(requests[0]?.payload.max_tokens).toBe(8192); + expect(requests[1]?.payload.max_tokens).toBeLessThanOrEqual(8192); + expect(await readdir(join(directory, "model-operations"))).toEqual([]); + } finally { + await child?.dispose(); + await parent.dispose(); + } + }); + + test.skipIf(process.getuid?.() === 0)( + "preserves completed output and usage when pending metadata cannot be removed", + async () => { + const { runtime, model, directory } = await createRuntime(); + const pendingDirectory = join(directory, "model-operations"); + vi.spyOn(globalThis, "fetch").mockImplementation(async () => { + await chmod(pendingDirectory, 0o500); + return completion(); + }); + try { + const result = await runtime.completeSimple(model, context); + expect(result.stopReason).toBe("stop"); + expect(result.content).toEqual([{ type: "text", text: "Hello" }]); + expect(result.usage).toMatchObject({ input: 8, output: 2 }); + expect(result.diagnostics?.at(-1)?.type).toBe("model_gateway_pending_cleanup_failed"); + expect(await readdir(pendingDirectory)).toHaveLength(1); + } finally { + await chmod(pendingDirectory, 0o700); + } + }, + ); + + test("gives each generation in an actual tool loop a distinct durable operation ID", async () => { + const { runtime, model, directory } = await createRuntime(); + await writeFile(join(directory, "tool-loop-fixture.txt"), "gateway-tool-loop"); + const requests: Array<{ id: string; payload: Record }> = []; + vi.spyOn(globalThis, "fetch").mockImplementation(async (input, init) => { + const request = new Request(input, init); + requests.push({ + id: request.headers.get("x-diffex-operation-id")!, + payload: (await request.json()) as Record, + }); + if (requests.length > 1) return completion(); + const events = [ + { + type: "message_start", + message: { + id: "msg_tool_fixture", + type: "message", + role: "assistant", + model: modelId, + content: [], + stop_reason: null, + stop_sequence: null, + usage: { + input_tokens: 8, + output_tokens: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }, + }, + { + type: "content_block_start", + index: 0, + content_block: { type: "tool_use", id: "tool_fixture", name: "read", input: {} }, + }, + { + type: "content_block_delta", + index: 0, + delta: { type: "input_json_delta", partial_json: '{"path":"tool-loop-fixture.txt"}' }, + }, + { type: "content_block_stop", index: 0 }, + { + type: "message_delta", + delta: { stop_reason: "tool_use", stop_sequence: null, stop_details: null }, + usage: { output_tokens: 2 }, + }, + { type: "message_stop" }, + ]; + return new Response( + events.map((event) => `event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`).join(""), + { headers: { "content-type": "text/event-stream" } }, + ); + }); + const { session } = await createAgentSession({ + cwd: directory, + agentDir: join(directory, "agent"), + modelRuntime: runtime, + model, + sessionManager: SessionManager.inMemory(directory), + settingsManager: SettingsManager.inMemory({ compaction: { enabled: false } }), + }); + try { + await session.prompt("Execute the harmless fixture tool and then answer."); + expect(requests).toHaveLength(2); + expect(requests[0]?.id).not.toBe(requests[1]?.id); + expect(JSON.stringify(requests[1]?.payload.messages)).toContain('"type":"tool_result"'); + expect(JSON.stringify(requests[1]?.payload.messages)).toContain("gateway-tool-loop"); + expect(session.agent.state.messages.at(-1)).toMatchObject({ role: "assistant", stopReason: "stop" }); + expect(await readdir(join(directory, "model-operations"))).toEqual([]); + } finally { + await session.dispose(); + } + }); +}); diff --git a/packages/server/src/protocol.ts b/packages/server/src/protocol.ts index 7a0a4a024..51d9d7fd5 100644 --- a/packages/server/src/protocol.ts +++ b/packages/server/src/protocol.ts @@ -31,7 +31,7 @@ type _AiModelInputsFitProtocol = Assert; /** * Enumerate mapped and intentionally omitted diffex-ai fields so additions fail compilation here. - * Provider replay metadata, diagnostics, cache-write retention splits, model transport settings, + * Provider replay metadata, retry policy, diagnostics, cache-write retention splits, model transport settings, * model sampling defaults, pricing tiers, and deferred-tool availability remain intentionally * server-side. */ @@ -93,6 +93,7 @@ type _AiAssistantMessageFieldsAccountedFor = Assert< | "stopReason" | "deferred" | "errorMessage" + | "retryable" | "rawStopReason" | "timestamp" >