From f1b79c88d26804fb6e1ec0d063e49cc0a3abb078 Mon Sep 17 00:00:00 2001 From: Jeffrey Tseng Date: Thu, 1 Oct 2026 01:46:51 -0700 Subject: [PATCH] feat(models): enforce hosted gateway calls Route hosted turns, tool loops, compaction, summaries and child agents through one constrained Anthropic gateway. Persist private operation identity and the request digest before sending, reject transport and credential overrides, and suppress replacement generations after ambiguous failures so metering cannot be bypassed or duplicated. --- packages/ai/src/types.ts | 2 + packages/ai/src/utils/overflow.ts | 1 + packages/ai/src/utils/retry.ts | 1 + .../coding-agent/src/core/agent-session.ts | 1 + .../src/core/hosted-model-gateway.ts | 293 ++++++++++++++ .../coding-agent/src/core/model-runtime.ts | 81 +++- packages/coding-agent/src/index.ts | 1 + .../test/hosted-model-gateway.test.ts | 379 ++++++++++++++++++ packages/server/src/protocol.ts | 3 +- 9 files changed, 759 insertions(+), 3 deletions(-) create mode 100644 packages/coding-agent/src/core/hosted-model-gateway.ts create mode 100644 packages/coding-agent/test/hosted-model-gateway.test.ts 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" >