From 5f9e4bc342e16f0ba6ed241c073facabc64c499a Mon Sep 17 00:00:00 2001 From: limityan Date: Mon, 24 Aug 2026 19:38:03 +0800 Subject: [PATCH] feat(sdk): bind local model for text queries Require one process-lifetime model configuration, keep its credentials runtime-only, and exercise the same SDK build through Node and Bun against the native Host. --- Cargo.lock | 1 + sdk/typescript/README.md | 28 +- sdk/typescript/package.json | 5 +- sdk/typescript/scripts/generate-wire.mjs | 2 + .../scripts/generated-wire-runtime.test.mjs | 5 +- sdk/typescript/src/client.ts | 99 +++- sdk/typescript/src/index.ts | 2 + sdk/typescript/src/internal/client.ts | 10 +- .../src/internal/wire-validation.ts | 7 +- sdk/typescript/src/session.ts | 7 +- sdk/typescript/src/types.ts | 21 +- sdk/typescript/test/client.test.ts | 140 +++++- sdk/typescript/test/fixtures/host.mjs | 9 +- .../test/lifecycle-timeouts.test.ts | 9 +- sdk/typescript/test/managed-host.test.ts | 109 ++++- sdk/typescript/test/real-host-smoke.mjs | 456 ++++++++++++++++++ src/apps/sdk-host/Cargo.toml | 1 + src/apps/sdk-host/src/main.rs | 7 + src/apps/sdk-host/src/model_source.rs | 208 ++++++++ src/apps/sdk-host/src/transport.rs | 16 +- src/apps/sdk-host/tests/stdio_process.rs | 61 ++- src/apps/sdk-host/tests/stdio_transport.rs | 80 ++- .../src/agentic/session/session_manager.rs | 44 +- .../src/infrastructure/ai/client_factory.rs | 103 ++-- .../core/src/service/config/service.rs | 159 ++++++ src/crates/interfaces/sdk-host/src/host.rs | 145 +++++- .../interfaces/sdk-host/src/protocol.rs | 45 +- .../sdk-host/tests/host_lifecycle.rs | 290 ++++++++++- .../sdk-host/tests/protocol_contracts.rs | 73 ++- 29 files changed, 2000 insertions(+), 142 deletions(-) create mode 100644 sdk/typescript/test/real-host-smoke.mjs create mode 100644 src/apps/sdk-host/src/model_source.rs diff --git a/Cargo.lock b/Cargo.lock index 1f0f7b84c2..dbc807b104 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1508,6 +1508,7 @@ dependencies = [ "tokio-util", "tracing", "tracing-subscriber", + "url", ] [[package]] diff --git a/sdk/typescript/README.md b/sdk/typescript/README.md index 95e3f5e159..eb3c9f18e8 100644 --- a/sdk/typescript/README.md +++ b/sdk/typescript/README.md @@ -26,15 +26,23 @@ existing `agent-runtime::sdk` API. ## Repository usage -Build `bitfun-sdk-host`, then pass its absolute path while the platform-native -package layout is still pending: +Build `bitfun-sdk-host`, then pass its absolute path and one process-lifetime +model configuration while the platform-native package layout is still pending. +The Host path must be explicit in this slice: ```typescript import { AgentClient } from "@bitfun/agent-sdk"; +const apiKey = await trustedSecretStore.read("openai"); await using client = await AgentClient.start({ cwd: process.cwd(), - hostPath: process.env.BITFUN_SDK_HOST_PATH, + hostPath: "/absolute/path/to/bitfun-sdk-host", + model: { + provider: "openai", + model: "gpt-5.4", + apiKey, + baseUrl: "https://api.openai.com/v1", + }, }); await using query = await client.query({ prompt: "Summarize this repository" }); @@ -46,15 +54,23 @@ for await (const item of query) { const result = await query.result(); ``` -`BITFUN_SDK_HOST_PATH` is also read directly when `hostPath` is omitted. The -eventual installable package must bundle or resolve a matching signed Host; it -must not require a separately installed BitFun CLI. +This repository-local package is private and unpublished. Node 24.14.1 and Bun +1.4.0 are the locally verified runners for this slice; they are not bundled +executables or a final minimum-version policy. The eventual installable package +must bundle or resolve a matching signed Host; it must not require a separately +installed BitFun CLI. + +Browser and mobile runtimes cannot launch the local native Host. Custom +functions, permission and user-input callbacks, structured output, usage, +Session resume, Python support, and native package staging remain deferred. ## Development ```bash pnpm --dir sdk/typescript test pnpm --dir sdk/typescript type-check +pnpm --dir sdk/typescript smoke:node +pnpm --dir sdk/typescript smoke:bun ``` The internal TypeScript wire bindings are generated from the Rust SDK Host diff --git a/sdk/typescript/package.json b/sdk/typescript/package.json index 735786447e..c6dd7d992c 100644 --- a/sdk/typescript/package.json +++ b/sdk/typescript/package.json @@ -4,9 +4,6 @@ "private": true, "description": "Internal TypeScript vertical slice for the BitFun Agent SDK", "type": "module", - "engines": { - "node": ">=22.12.0" - }, "exports": { ".": { "types": "./dist/sdk/typescript/src/index.d.ts", @@ -23,6 +20,8 @@ "scripts": { "build": "pnpm run generate:wire && tsc -p tsconfig.json", "generate:wire": "node scripts/generate-wire.mjs", + "smoke:bun": "bun test/real-host-smoke.mjs", + "smoke:node": "node test/real-host-smoke.mjs", "test": "pnpm run build && node --test scripts/*.test.mjs dist/src/crates/adapters/transport/typescript/test/**/*.test.js dist/sdk/typescript/test/**/*.test.js", "type-check": "pnpm run generate:wire && tsc -p tsconfig.json --noEmit" }, diff --git a/sdk/typescript/scripts/generate-wire.mjs b/sdk/typescript/scripts/generate-wire.mjs index bd94f6b863..04b4a9f56b 100644 --- a/sdk/typescript/scripts/generate-wire.mjs +++ b/sdk/typescript/scripts/generate-wire.mjs @@ -60,6 +60,8 @@ const requiredTypes = [ "SessionCreateParams", "SessionCreateResult", "ShutdownResult", + "TemporaryModelConfig", + "TemporaryModelProvider", ]; const missing = requiredTypes.filter((type) => !files.includes(type)); if (missing.length > 0) { diff --git a/sdk/typescript/scripts/generated-wire-runtime.test.mjs b/sdk/typescript/scripts/generated-wire-runtime.test.mjs index fa8306cf86..784a4ac7a0 100644 --- a/sdk/typescript/scripts/generated-wire-runtime.test.mjs +++ b/sdk/typescript/scripts/generated-wire-runtime.test.mjs @@ -43,7 +43,7 @@ test("Rust wire export produces executable validators for every type", async () } const initializeResult = { - protocolVersion: 1, + protocolVersion: 2, runtimeVersion: "0.1.0", stability: "not_delivered", capabilities: { @@ -61,8 +61,11 @@ test("Rust wire export produces executable validators for every type", async () mcpConfiguration: false, prestartedTransport: false, }, + modelId: "sdk:openai:resolved", }; assert.equal(validators.isInitializeResult(initializeResult), true); + const { modelId: _modelId, ...initializeResultWithoutModel } = initializeResult; + assert.equal(validators.isInitializeResult(initializeResultWithoutModel), false); assert.equal( validators.isInitializeResult({ ...initializeResult, unexpected: true }), false, diff --git a/sdk/typescript/src/client.ts b/sdk/typescript/src/client.ts index e73ef997e5..bd8e4d0d33 100644 --- a/sdk/typescript/src/client.ts +++ b/sdk/typescript/src/client.ts @@ -1,12 +1,28 @@ +import { isAbsolute } from "node:path"; + import type { InitializeResult, QueryStartParams, QueryStartResult } from "./internal/wire/index.js"; import type { JsonRpcConnection } from "./internal/json-rpc.js"; +import { SdkError } from "./errors.js"; import { Query } from "./query.js"; import { Session, Sessions } from "./session.js"; -import type { AgentCapabilities, AgentClientOptions, QueryInput } from "./types.js"; +import type { + AgentCapabilities, + AgentClientOptions, + AgentModelOptions, + QueryInput, +} from "./types.js"; + +const SUPPORTED_MODEL_PROVIDERS = new Set([ + "openai", + "responses", + "anthropic", + "gemini", +]); export class AgentClient { readonly #connection: JsonRpcConnection; - readonly #options: AgentClientOptions; + readonly #cwd: string; + readonly #modelId: string; readonly capabilities: AgentCapabilities; readonly sessions: Sessions; readonly #queries = new Set(); @@ -15,14 +31,14 @@ export class AgentClient { #closePromise?: Promise; static async start(options: AgentClientOptions): Promise { - const hostPath = options.hostPath ?? process.env.BITFUN_SDK_HOST_PATH; - if (hostPath === undefined || hostPath.length === 0) { - const { SdkError } = await import("./errors.js"); - throw new SdkError("SDK Host executable is unavailable", { - code: "not_found", + validateModelOptions(options.model as unknown); + const hostPath = options.hostPath; + if (typeof hostPath !== "string" || !isAbsolute(hostPath)) { + throw new SdkError("SDK Host path must be an explicit absolute path", { + code: "invalid_request", stage: "initialize", retryable: false, - correlationId: "local:host_start", + correlationId: "local:host_validation", outcomeCertainty: "not_started", }); } @@ -44,11 +60,12 @@ export class AgentClient { private constructor( connection: JsonRpcConnection, - options: AgentClientOptions, + options: Pick, initialized: InitializeResult, ) { this.#connection = connection; - this.#options = options; + this.#cwd = options.cwd; + this.#modelId = initialized.modelId; this.capabilities = Object.freeze({ query: initialized.capabilities.query, sessions: initialized.capabilities.sessionCreate, @@ -56,7 +73,8 @@ export class AgentClient { }); this.sessions = Sessions.forClient( connection, - options.cwd, + this.#cwd, + this.#modelId, (query) => this.#trackQuery(query), (session) => this.#trackSession(session), () => this.#ensureOpen(), @@ -66,7 +84,7 @@ export class AgentClient { /** @internal */ static create( connection: JsonRpcConnection, - options: AgentClientOptions, + options: Pick, initialized: InitializeResult, ): AgentClient { return new AgentClient(connection, options, initialized); @@ -79,8 +97,8 @@ export class AgentClient { sessionId: null, sessionName: null, agent: input.agent ?? null, - cwd: this.#options.cwd, - model: input.model ?? null, + cwd: this.#cwd, + model: this.#modelId, }; const started = await this.#connection.request( "query/start", @@ -152,3 +170,56 @@ export class AgentClient { session.onClosed(() => this.#ownedSessions.delete(session)); } } + +function validateModelOptions(value: unknown): asserts value is AgentModelOptions { + if (typeof value !== "object" || value === null || Array.isArray(value)) { + invalidModel("model is required"); + } + const model = value as Record; + if ( + typeof model.provider !== "string" || + !SUPPORTED_MODEL_PROVIDERS.has(model.provider) + ) { + invalidModel("model.provider is unsupported"); + } + if (typeof model.model !== "string" || model.model.trim().length === 0) { + invalidModel("model.model is required"); + } + if (typeof model.apiKey !== "string" || model.apiKey.trim().length === 0) { + invalidModel("model.apiKey is required"); + } + if (model.baseUrl === undefined) { + return; + } + const invalidBaseUrl = + "model.baseUrl must be an absolute http or https URL without credentials, query, or fragment"; + if (typeof model.baseUrl !== "string") { + invalidModel(invalidBaseUrl); + } + let url: URL; + try { + url = new URL(model.baseUrl); + } catch { + invalidModel(invalidBaseUrl); + } + if ( + (url.protocol !== "http:" && url.protocol !== "https:") || + url.hostname.length === 0 || + url.username.length > 0 || + url.password.length > 0 || + url.search.length > 0 || + url.hash.length > 0 + ) { + invalidModel(invalidBaseUrl); + } +} + +function invalidModel(message: string): never { + throw new SdkError(message, { + code: "invalid_request", + stage: "initialize", + retryable: false, + correlationId: "local:model_validation", + outcomeCertainty: "not_started", + }); +} diff --git a/sdk/typescript/src/index.ts b/sdk/typescript/src/index.ts index c3d6901c6b..54e043a9b8 100644 --- a/sdk/typescript/src/index.ts +++ b/sdk/typescript/src/index.ts @@ -5,6 +5,8 @@ export { Session, Sessions } from "./session.js"; export type { AgentCapabilities, AgentClientOptions, + AgentModelOptions, + AgentModelProvider, AssistantTextDelta, OutcomeCertainty, QueryInput, diff --git a/sdk/typescript/src/internal/client.ts b/sdk/typescript/src/internal/client.ts index 4199e3c539..5e0dc4ac8e 100644 --- a/sdk/typescript/src/internal/client.ts +++ b/sdk/typescript/src/internal/client.ts @@ -4,18 +4,24 @@ import { JsonRpcConnection } from "./json-rpc.js"; import type { HostTransport } from "./transport.js"; import type { InitializeParams, InitializeResult } from "./wire/index.js"; -const PROTOCOL_VERSION = 1; +const PROTOCOL_VERSION = 2; const DEFAULT_INITIALIZE_TIMEOUT_MS = 30_000; export async function createAgentClient( transport: HostTransport, - options: AgentClientOptions, + options: Pick, ): Promise { const connection = new JsonRpcConnection(transport); const params: InitializeParams = { protocolVersion: PROTOCOL_VERSION, clientInfo: { name: "@bitfun/agent-sdk", version: "0.0.0" }, capabilities: { serverNotifications: true }, + model: { + provider: options.model.provider, + model: options.model.model, + apiKey: options.model.apiKey, + baseUrl: options.model.baseUrl, + }, }; const initialized = await connection.request( "initialize", diff --git a/sdk/typescript/src/internal/wire-validation.ts b/sdk/typescript/src/internal/wire-validation.ts index 7b390f41dc..acea726f89 100644 --- a/sdk/typescript/src/internal/wire-validation.ts +++ b/sdk/typescript/src/internal/wire-validation.ts @@ -93,8 +93,11 @@ function validateInitializeResult(value: unknown): InitializeResult { value, "initialize result", ); - if (!Number.isSafeInteger(result.protocolVersion)) { - throw new Error("SDK Host initialize protocol version is invalid"); + if ( + !Number.isSafeInteger(result.protocolVersion) || + !isNonEmptyString(result.modelId) + ) { + throw new Error("SDK Host initialize protocol version or model id is invalid"); } return result; } diff --git a/sdk/typescript/src/session.ts b/sdk/typescript/src/session.ts index c6b865e9b7..d063792dab 100644 --- a/sdk/typescript/src/session.ts +++ b/sdk/typescript/src/session.ts @@ -14,6 +14,7 @@ import type { SessionCreateInput, SessionLifetime, TurnInput } from "./types.js" export class Sessions { readonly #connection: JsonRpcConnection; readonly #cwd: string; + readonly #modelId: string; readonly #onQuery: (query: Query) => Query; readonly #onSession: (session: Session) => void; readonly #ensureClientOpen: () => void; @@ -22,6 +23,7 @@ export class Sessions { static forClient( connection: JsonRpcConnection, cwd: string, + modelId: string, onQuery: (query: Query) => Query, onSession: (session: Session) => void, ensureClientOpen: () => void, @@ -29,6 +31,7 @@ export class Sessions { return new Sessions( connection, cwd, + modelId, onQuery, onSession, ensureClientOpen, @@ -38,12 +41,14 @@ export class Sessions { private constructor( connection: JsonRpcConnection, cwd: string, + modelId: string, onQuery: (query: Query) => Query, onSession: (session: Session) => void, ensureClientOpen: () => void, ) { this.#connection = connection; this.#cwd = cwd; + this.#modelId = modelId; this.#onQuery = onQuery; this.#onSession = onSession; this.#ensureClientOpen = ensureClientOpen; @@ -55,7 +60,7 @@ export class Sessions { sessionName: input.sessionName ?? null, agent: input.agent ?? null, cwd: this.#cwd, - model: input.model ?? null, + model: this.#modelId, }; const created = await this.#connection.request( "session/create", diff --git a/sdk/typescript/src/types.ts b/sdk/typescript/src/types.ts index cd460482c4..0fad3f955a 100644 --- a/sdk/typescript/src/types.ts +++ b/sdk/typescript/src/types.ts @@ -1,9 +1,24 @@ +export type AgentModelProvider = + | "openai" + | "responses" + | "anthropic" + | "gemini"; + +export interface AgentModelOptions { + provider: AgentModelProvider; + model: string; + apiKey: string; + baseUrl?: string; +} + export interface AgentClientOptions { cwd: string; - /** Native `bitfun-sdk-host` path. Platform packages will provide this later. */ - hostPath?: string; + /** Explicit absolute path to the native `bitfun-sdk-host`. */ + hostPath: string; /** Deadline for the SDK Host initialize handshake. */ initializeTimeoutMs?: number; + /** Process-lifetime model credentials installed into this Host connection. */ + model: AgentModelOptions; } export interface AgentCapabilities { @@ -66,13 +81,11 @@ export type SessionLifetime = "connection"; export interface QueryInput { prompt: string; agent?: string; - model?: string; } export interface SessionCreateInput { sessionName?: string; agent?: string; - model?: string; } export interface TurnInput { diff --git a/sdk/typescript/test/client.test.ts b/sdk/typescript/test/client.test.ts index b9b9106aac..f059ad545e 100644 --- a/sdk/typescript/test/client.test.ts +++ b/sdk/typescript/test/client.test.ts @@ -5,11 +5,42 @@ import test from "node:test"; import { AgentClient, SdkError } from "../src/index.js"; import { createAgentClient } from "../src/internal/client.js"; +import type { + AgentClientOptions, + QueryInput, + SessionCreateInput, +} from "../src/types.js"; + +const clientOptions = { + cwd: "D:/workspace/project", + hostPath: process.execPath, + model: { + provider: "openai" as const, + model: "fixture-model", + apiKey: "fixture-secret", + baseUrl: "http://127.0.0.1:43123/v1", + }, +} satisfies AgentClientOptions; + +// @ts-expect-error An explicit native Host path is required until platform packages exist. +const missingHostPathOptions: AgentClientOptions = { + cwd: "D:/workspace/project", + model: clientOptions.model, +}; + +// @ts-expect-error Query model selection is bound at AgentClient.start. +const queryModelOverride: QueryInput = { prompt: "hello", model: "attempted-override" }; +// @ts-expect-error Session model selection is bound at AgentClient.start. +const sessionModelOverride: SessionCreateInput = { model: "attempted-override" }; +void queryModelOverride; +void sessionModelOverride; +void missingHostPathOptions; test("a Query streams ordered events and returns the Host terminal Result", async () => { const clientToHost = new PassThrough(); const hostToClient = new PassThrough(); - const host = runFixtureHost(clientToHost, hostToClient); + const initializeRequests: unknown[] = []; + const host = runFixtureHost(clientToHost, hostToClient, initializeRequests); const client = await createAgentClient( { readable: hostToClient, @@ -19,11 +50,26 @@ test("a Query streams ordered events and returns the Host terminal Result", asyn await host; }, }, - { cwd: "D:/workspace/project" }, + clientOptions, ); assert.ok(client instanceof AgentClient); - const query = await client.query({ prompt: "hello" }); + assert.equal(initializeRequests.length, 1); + assert.deepEqual(initializeRequests[0], { + protocolVersion: 2, + clientInfo: { name: "@bitfun/agent-sdk", version: "0.0.0" }, + capabilities: { serverNotifications: true }, + model: { + provider: "openai", + model: "fixture-model", + apiKey: "fixture-secret", + baseUrl: "http://127.0.0.1:43123/v1", + }, + }); + const query = await client.query({ + prompt: "hello", + model: "attempted-override", + } as QueryInput); assert.equal(query.id, "query-1"); assert.equal(query.operationId, "operation-1"); assert.deepEqual(query.turn, { id: "turn-1", sessionId: "session-1" }); @@ -73,10 +119,13 @@ test("an explicit Session starts Turns on the existing client connection", async await host; }, }, - { cwd: "D:/workspace/project" }, + clientOptions, ); - const session = await client.sessions.create({ agent: "agentic" }); + const session = await client.sessions.create({ + agent: "agentic", + model: "attempted-override", + } as SessionCreateInput); assert.equal(session.id, "session-explicit"); assert.equal(session.agent, "agentic"); @@ -111,7 +160,7 @@ test("Query cancel and close are idempotent and the Host Result remains authorit await host; }, }, - { cwd: "D:/workspace/project" }, + clientOptions, ); const query = await client.query({ prompt: "wait" }); @@ -146,7 +195,7 @@ test("leaving Query iteration early cancels and settles the Turn", async () => { await host; }, }, - { cwd: "D:/workspace/project" }, + clientOptions, ); const query = await client.query({ prompt: "stream" }); @@ -175,7 +224,7 @@ test("Host loss rejects an accepted Query with unknown outcome instead of fabric await host; }, }, - { cwd: "D:/workspace/project" }, + clientOptions, ); const query = await client.query({ prompt: "may have side effects" }); @@ -192,6 +241,36 @@ test("Host loss rejects an accepted Query with unknown outcome instead of fabric assert.equal(transportClosed, true); }); +test("an empty initialized model id fails the connection closed", async () => { + const clientToHost = new PassThrough(); + const hostToClient = new PassThrough(); + let transportClosed = false; + const host = runEmptyModelIdFixtureHost(clientToHost, hostToClient); + + await assert.rejects( + createAgentClient( + { + readable: hostToClient, + writable: clientToHost, + close: async () => { + transportClosed = true; + clientToHost.end(); + await host; + }, + }, + clientOptions, + ), + (error: unknown) => { + assert.ok(error instanceof SdkError); + assert.equal(error.code, "process_lost"); + assert.equal(error.stage, "protocol"); + assert.equal(error.outcomeCertainty, "unknown"); + return true; + }, + ); + assert.equal(transportClosed, true); +}); + test("AgentClient.close settles owned Queries before shutting down its connection", async () => { const clientToHost = new PassThrough(); const hostToClient = new PassThrough(); @@ -206,7 +285,7 @@ test("AgentClient.close settles owned Queries before shutting down its connectio await host; }, }, - { cwd: "D:/workspace/project" }, + clientOptions, ); const query = await client.query({ prompt: "still running" }); @@ -234,7 +313,7 @@ test("Host operation errors preserve stable SDK error facts", async () => { await host; }, }, - { cwd: "D:/workspace/project" }, + clientOptions, ); await assert.rejects(client.query({ prompt: "requires auth" }), (error: unknown) => { @@ -264,7 +343,7 @@ test("unknown Host error facts fail the protocol closed", async () => { await host; }, }, - { cwd: "D:/workspace/project" }, + clientOptions, ); await assert.rejects(client.query({ prompt: "invalid error" }), (error: unknown) => { @@ -289,7 +368,7 @@ test("ambiguous JSON-RPC response envelopes fail the protocol closed", async () await host; }, }, - { cwd: "D:/workspace/project" }, + clientOptions, ); await assert.rejects(client.query({ prompt: "reject ambiguous response" }), (error: unknown) => { @@ -356,7 +435,7 @@ test("unknown Query event and Result status fail the protocol closed", async (co await host; }, }, - { cwd: "D:/workspace/project" }, + clientOptions, ); try { @@ -377,6 +456,7 @@ test("unknown Query event and Result status fail the protocol closed", async (co async function runFixtureHost( requests: PassThrough, responses: PassThrough, + initializeRequests: unknown[], ): Promise { const lines = createInterface({ input: requests, crlfDelay: Infinity }); for await (const line of lines) { @@ -386,11 +466,12 @@ async function runFixtureHost( params: Record; }; if (request.method === "initialize") { + initializeRequests.push(request.params); write(responses, { jsonrpc: "2.0", id: request.id, result: { - protocolVersion: 1, + protocolVersion: 2, runtimeVersion: "0.2.17", stability: "not_delivered", capabilities: { @@ -408,11 +489,20 @@ async function runFixtureHost( mcpConfiguration: false, prestartedTransport: false, }, + modelId: "sdk:openai:resolved", }, }); continue; } if (request.method === "query/start") { + assert.deepEqual(request.params, { + prompt: "hello", + sessionId: null, + sessionName: null, + agent: null, + cwd: "D:/workspace/project", + model: "sdk:openai:resolved", + }); write(responses, { jsonrpc: "2.0", id: request.id, @@ -495,7 +585,7 @@ async function runSessionFixtureHost( sessionName: null, agent: "agentic", cwd: "D:/workspace/project", - model: null, + model: "sdk:openai:resolved", }); write(responses, { jsonrpc: "2.0", @@ -807,6 +897,23 @@ async function runClientCloseFixtureHost( } } +async function runEmptyModelIdFixtureHost( + requests: PassThrough, + responses: PassThrough, +): Promise { + const lines = createInterface({ input: requests, crlfDelay: Infinity }); + for await (const line of lines) { + const request = JSON.parse(line) as { id: number; method: string }; + assert.equal(request.method, "initialize"); + const response = initializeResponse(request.id) as { + result: { modelId: string }; + }; + response.result.modelId = ""; + write(responses, response); + } + responses.end(); +} + async function runOperationErrorFixtureHost( requests: PassThrough, responses: PassThrough, @@ -989,7 +1096,7 @@ function initializeResponse(id: number): unknown { jsonrpc: "2.0", id, result: { - protocolVersion: 1, + protocolVersion: 2, runtimeVersion: "0.2.17", stability: "not_delivered", capabilities: { @@ -1007,6 +1114,7 @@ function initializeResponse(id: number): unknown { mcpConfiguration: false, prestartedTransport: false, }, + modelId: "sdk:openai:resolved", }, }; } diff --git a/sdk/typescript/test/fixtures/host.mjs b/sdk/typescript/test/fixtures/host.mjs index 0fea572501..6a56b2a07f 100644 --- a/sdk/typescript/test/fixtures/host.mjs +++ b/sdk/typescript/test/fixtures/host.mjs @@ -4,11 +4,17 @@ const lines = createInterface({ input: process.stdin, crlfDelay: Infinity }); for await (const line of lines) { const request = JSON.parse(line); if (request.method === "initialize") { + if ( + request.params?.protocolVersion !== 2 || + request.params?.model?.apiKey !== "fixture-secret" + ) { + throw new Error("Invalid initialize request"); + } write({ jsonrpc: "2.0", id: request.id, result: { - protocolVersion: 1, + protocolVersion: 2, runtimeVersion: "fixture", stability: "not_delivered", capabilities: { @@ -26,6 +32,7 @@ for await (const line of lines) { mcpConfiguration: false, prestartedTransport: false, }, + modelId: "sdk:openai:resolved", }, }); continue; diff --git a/sdk/typescript/test/lifecycle-timeouts.test.ts b/sdk/typescript/test/lifecycle-timeouts.test.ts index b63903a481..2436ba3f4a 100644 --- a/sdk/typescript/test/lifecycle-timeouts.test.ts +++ b/sdk/typescript/test/lifecycle-timeouts.test.ts @@ -13,6 +13,13 @@ import type { import { Query } from "../src/query.js"; import { Session } from "../src/session.js"; +const model = { + provider: "openai" as const, + model: "fixture-model", + apiKey: "fixture-secret", + baseUrl: "http://127.0.0.1:43123/v1", +}; + test("client initialization aborts the Host connection after its startup deadline", async () => { const clientToHost = new PassThrough(); const hostToClient = new PassThrough(); @@ -27,7 +34,7 @@ test("client initialization aborts the Host connection after its startup deadlin hostToClient.end(); }, }, - { cwd: "D:/workspace/project", initializeTimeoutMs: 20 }, + { cwd: "D:/workspace/project", initializeTimeoutMs: 20, model }, ); await assert.rejects(withTestDeadline(starting), (error: unknown) => { diff --git a/sdk/typescript/test/managed-host.test.ts b/sdk/typescript/test/managed-host.test.ts index 9667f012ba..150b18c500 100644 --- a/sdk/typescript/test/managed-host.test.ts +++ b/sdk/typescript/test/managed-host.test.ts @@ -1,11 +1,20 @@ import assert from "node:assert/strict"; import { once } from "node:events"; +import { basename, dirname } from "node:path"; import { fileURLToPath } from "node:url"; import test from "node:test"; import { AgentClient, SdkError } from "../src/index.js"; import { createAgentClient } from "../src/internal/client.js"; import { forceKillTree, startManagedHost } from "../src/internal/managed-host.js"; +import type { AgentClientOptions } from "../src/types.js"; + +const model = { + provider: "openai" as const, + model: "fixture-model", + apiKey: "fixture-secret", + baseUrl: "http://127.0.0.1:43123/v1", +}; test("the managed transport owns one child Host process", async () => { const fixture = fileURLToPath( @@ -16,7 +25,7 @@ test("the managed transport owns one child Host process", async () => { args: [fixture], cwd: process.cwd(), }); - const client = await createAgentClient(transport, { cwd: process.cwd() }); + const client = await createAgentClient(transport, { cwd: process.cwd(), model }); assert.ok(client instanceof AgentClient); await client.close(); @@ -30,6 +39,7 @@ test("AgentClient.start reports a missing Host before an operation begins", asyn hostPath: fileURLToPath( new URL("../../../../test/fixtures/missing-host", import.meta.url), ), + model, }), (error: unknown) => { assert.ok(error instanceof SdkError); @@ -41,6 +51,103 @@ test("AgentClient.start reports a missing Host before an operation begins", asyn ); }); +test("AgentClient.start rejects missing and relative Host paths before spawning", async (context) => { + const previousHostPath = process.env.BITFUN_SDK_HOST_PATH; + delete process.env.BITFUN_SDK_HOST_PATH; + try { + const cases: Array<{ name: string; options: unknown }> = [ + { + name: "missing path", + options: { cwd: process.cwd(), model }, + }, + { + name: "relative path", + options: { + cwd: dirname(process.execPath), + hostPath: basename(process.execPath), + initializeTimeoutMs: 100, + model, + }, + }, + ]; + + for (const fixture of cases) { + await context.test(fixture.name, async () => { + await assert.rejects( + AgentClient.start(fixture.options as AgentClientOptions), + (error: unknown) => { + assert.ok(error instanceof SdkError); + assert.equal(error.code, "invalid_request"); + assert.equal(error.stage, "initialize"); + assert.equal(error.outcomeCertainty, "not_started"); + assert.doesNotMatch(String(error.stack), /fixture-secret/); + return true; + }, + ); + }); + } + } finally { + if (previousHostPath === undefined) { + delete process.env.BITFUN_SDK_HOST_PATH; + } else { + process.env.BITFUN_SDK_HOST_PATH = previousHostPath; + } + } +}); + +test("AgentClient.start rejects invalid model options before spawning a Host", async (context) => { + const missingHost = fileURLToPath( + new URL("../../../../test/fixtures/missing-host", import.meta.url), + ); + const cases: Array<{ name: string; options: unknown }> = [ + { + name: "missing model", + options: { cwd: process.cwd(), hostPath: missingHost }, + }, + { + name: "blank model", + options: { cwd: process.cwd(), hostPath: missingHost, model: { ...model, model: " " } }, + }, + { + name: "blank API key", + options: { cwd: process.cwd(), hostPath: missingHost, model: { ...model, apiKey: " " } }, + }, + { + name: "unsupported provider", + options: { + cwd: process.cwd(), + hostPath: missingHost, + model: { ...model, provider: "unsupported" }, + }, + }, + { + name: "invalid base URL", + options: { + cwd: process.cwd(), + hostPath: missingHost, + model: { ...model, baseUrl: "not an absolute URL" }, + }, + }, + ]; + + for (const fixture of cases) { + await context.test(fixture.name, async () => { + await assert.rejects( + AgentClient.start(fixture.options as AgentClientOptions), + (error: unknown) => { + assert.ok(error instanceof SdkError); + assert.equal(error.code, "invalid_request"); + assert.equal(error.stage, "initialize"); + assert.equal(error.retryable, false); + assert.equal(error.outcomeCertainty, "not_started"); + assert.doesNotMatch(String(error.stack), /fixture-secret/); + return true; + }, + ); + }); + } +}); + test("forced managed Host cleanup reclaims its descendant process tree", async () => { const fixture = fileURLToPath( new URL("../../../../test/fixtures/unresponsive-host.mjs", import.meta.url), diff --git a/sdk/typescript/test/real-host-smoke.mjs b/sdk/typescript/test/real-host-smoke.mjs new file mode 100644 index 0000000000..c38d225456 --- /dev/null +++ b/sdk/typescript/test/real-host-smoke.mjs @@ -0,0 +1,456 @@ +import assert from "node:assert/strict"; +import { spawn } from "node:child_process"; +import { randomBytes } from "node:crypto"; +import { mkdir, mkdtemp, readFile, readdir, rm } from "node:fs/promises"; +import { createServer } from "node:http"; +import { tmpdir } from "node:os"; +import { dirname, join, resolve } from "node:path"; + +import { AgentClient, SdkError } from "../dist/sdk/typescript/src/index.js"; + +const scriptPath = resolve(process.argv[1]); +const packageRoot = resolve(dirname(scriptPath), ".."); +const MAX_CAPTURED_OUTPUT_BYTES = 1024 * 1024; +const WORKER_TIMEOUT_MS = 120_000; + +if (process.argv.includes("--worker")) { + await runWorker(); +} else { + await runParent(); +} + +async function runParent() { + const isolatedRoot = await mkdtemp(join(tmpdir(), "bitfun-sdk-real-host-")); + const workspace = join(isolatedRoot, "workspace"); + const userRoot = join(isolatedRoot, "user-root"); + const home = join(isolatedRoot, "home"); + const configRoot = join(isolatedRoot, "config-root"); + await Promise.all( + [workspace, userRoot, home, configRoot].map((directory) => + mkdir(directory, { recursive: true }), + ), + ); + + const hostPath = process.env.BITFUN_SDK_HOST_PATH; + assert.equal(typeof hostPath, "string", "BITFUN_SDK_HOST_PATH is required"); + assert.notEqual(hostPath.length, 0, "BITFUN_SDK_HOST_PATH is required"); + + const apiKey = `bitfun-sdk-${randomBytes(24).toString("hex")}`; + let requestCount = 0; + const requestTraces = []; + let fixtureFailure; + const server = createServer(async (request, response) => { + requestCount += 1; + try { + if (request.method !== "POST" || request.url !== "/v1/chat/completions") { + throw new Error("SSE fixture received an unexpected request target"); + } + if (request.headers.authorization !== `Bearer ${apiKey}`) { + throw new Error("SSE fixture received invalid authorization"); + } + const body = await readRequestJson(request); + if (body.model !== "fixture-model" || body.stream !== true) { + throw new Error("SSE fixture received an invalid model request"); + } + requestTraces.push(summarizeModelRequest(requestCount, request.url, body)); + + response.writeHead(200, { + "content-type": "text/event-stream", + "cache-control": "no-cache", + }); + response.end( + [ + 'data: {"id":"fixture","object":"chat.completion.chunk","model":"fixture-model","choices":[{"index":0,"delta":{"role":"assistant","content":"BitFun SDK "},"finish_reason":null}]}', + 'data: {"id":"fixture","object":"chat.completion.chunk","model":"fixture-model","choices":[{"index":0,"delta":{"content":"fixture response"},"finish_reason":null}]}', + 'data: {"id":"fixture","object":"chat.completion.chunk","model":"fixture-model","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + "", + ].join("\n\n"), + ); + } catch (error) { + fixtureFailure ??= error; + response.writeHead(400, { "content-type": "text/plain" }); + response.end("SSE fixture rejected the request"); + } + }); + + let worker; + try { + const address = await listenLocalhost(server); + const workerEnvironment = { + BITFUN_SDK_SMOKE_BASE_URL: `http://127.0.0.1:${String(address.port)}/v1`, + BITFUN_SDK_HOST_PATH: hostPath, + BITFUN_SDK_SMOKE_WORKSPACE: workspace, + BITFUN_E2E_STORAGE_GUARD: "1", + BITFUN_E2E_USER_ROOT: userRoot, + BITFUN_E2E_HOME: home, + APPDATA: configRoot, + XDG_CONFIG_HOME: configRoot, + HOME: home, + USERPROFILE: home, + }; + assert.equal(Object.values(workerEnvironment).includes(apiKey), false); + + worker = spawn(process.execPath, [scriptPath, "--worker"], { + cwd: packageRoot, + env: workerEnvironment, + shell: false, + stdio: ["pipe", "pipe", "pipe"], + windowsHide: true, + }); + const output = captureWorkerOutput(worker); + worker.stdin.end(`${apiKey}\n`); + let exit; + try { + exit = await waitForWorker(worker, WORKER_TIMEOUT_MS); + } catch (error) { + const captured = output(); + const phase = lastWorkerPhase(captured.stdout); + throw new Error( + [ + `real Host smoke worker failed at ${phase} after ${String(requestCount)} model requests`, + formatRequestTraces(requestTraces, apiKey), + ].join("\n"), + { cause: error }, + ); + } + const captured = output(); + + assert.equal(exit.signal, null, "real Host smoke worker was terminated"); + assert.equal(captured.stdout.includes(apiKey), false, "API key leaked to worker stdout"); + assert.equal(captured.stderr.includes(apiKey), false, "API key leaked to worker stderr"); + assert.equal( + exit.code, + 0, + `real Host smoke worker failed: ${captured.stderr}`, + ); + if (fixtureFailure !== undefined) { + throw fixtureFailure; + } + assert.equal(requestCount, 1, "real Host smoke must issue exactly one model request"); + await assertTreeDoesNotContain(isolatedRoot, apiKey); + process.stdout.write("real-host-smoke: PASS\n"); + } finally { + if (worker !== undefined && worker.exitCode === null) { + worker.kill(); + } + await closeServer(server); + await rm(isolatedRoot, { + recursive: true, + force: true, + maxRetries: 20, + retryDelay: 100, + }); + } +} + +async function runWorker() { + const apiKey = await readApiKeyFromStdin(); + const workspace = requiredEnvironment("BITFUN_SDK_SMOKE_WORKSPACE"); + const hostPath = requiredEnvironment("BITFUN_SDK_HOST_PATH"); + const baseUrl = requiredEnvironment("BITFUN_SDK_SMOKE_BASE_URL"); + const missingHost = join(workspace, "missing-bitfun-sdk-host"); + + const validModel = { + provider: "openai", + model: "fixture-model", + apiKey, + baseUrl, + }; + const invalidOptions = [ + { cwd: workspace, hostPath: missingHost }, + { cwd: workspace, hostPath: missingHost, model: { ...validModel, model: " " } }, + { cwd: workspace, hostPath: missingHost, model: { ...validModel, apiKey: " " } }, + { + cwd: workspace, + hostPath: missingHost, + model: { ...validModel, provider: "unsupported" }, + }, + { + cwd: workspace, + hostPath: missingHost, + model: { ...validModel, baseUrl: "not an absolute URL" }, + }, + ]; + for (const options of invalidOptions) { + await assert.rejects(AgentClient.start(options), (error) => { + assert.ok(error instanceof SdkError); + assert.equal(error.code, "invalid_request"); + assert.equal(error.stage, "initialize"); + assert.equal(error.outcomeCertainty, "not_started"); + assert.equal(renderError(error).includes(apiKey), false); + return true; + }); + } + process.stdout.write("phase:validation_complete\n"); + + const client = await AgentClient.start({ + cwd: workspace, + hostPath, + model: validModel, + }); + process.stdout.write("phase:client_started\n"); + let query; + try { + query = await client.query({ prompt: "Return the fixture response" }); + process.stdout.write("phase:query_started\n"); + const items = []; + for await (const item of query) { + items.push(item); + } + const result = await query.result(); + const deltas = items.filter((item) => item.type === "assistant_text_delta"); + const terminalResults = items.filter((item) => item.type === "result"); + + assert.equal(deltas.map((item) => item.text).join(""), "BitFun SDK fixture response"); + assert.equal(terminalResults.length, 1); + assert.deepEqual(terminalResults[0], result); + assert.equal(result.status, "completed"); + assert.equal(result.outputText, "BitFun SDK fixture response"); + assert.equal(JSON.stringify(items).includes(apiKey), false); + process.stdout.write("phase:result_received\n"); + } finally { + if (query !== undefined) { + await query.close(); + process.stdout.write("phase:query_closed\n"); + } + await client.close(); + process.stdout.write("phase:client_closed\n"); + } +} + +async function readRequestJson(request) { + const chunks = []; + let bytes = 0; + for await (const chunk of request) { + const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk); + bytes += buffer.byteLength; + if (bytes > 4 * 1024 * 1024) { + throw new Error("SSE fixture request exceeded its size limit"); + } + chunks.push(buffer); + } + return JSON.parse(Buffer.concat(chunks).toString("utf8")); +} + +async function listenLocalhost(server) { + await new Promise((resolvePromise, rejectPromise) => { + server.once("error", rejectPromise); + server.listen(0, "127.0.0.1", () => { + server.off("error", rejectPromise); + resolvePromise(); + }); + }); + const address = server.address(); + if (address === null || typeof address === "string") { + throw new Error("SSE fixture did not bind a TCP address"); + } + return address; +} + +function captureWorkerOutput(worker) { + let stdout = ""; + let stderr = ""; + let captureFailure; + worker.stdout.setEncoding("utf8"); + worker.stderr.setEncoding("utf8"); + worker.stdout.on("data", (chunk) => { + try { + stdout = appendBounded(stdout, chunk); + } catch (error) { + captureFailure ??= error; + worker.kill(); + } + }); + worker.stderr.on("data", (chunk) => { + try { + stderr = appendBounded(stderr, chunk); + } catch (error) { + captureFailure ??= error; + worker.kill(); + } + }); + return () => { + if (captureFailure !== undefined) { + throw captureFailure; + } + return { stdout, stderr }; + }; +} + +function appendBounded(current, chunk) { + const next = current + chunk; + if (Buffer.byteLength(next, "utf8") > MAX_CAPTURED_OUTPUT_BYTES) { + throw new Error("real Host smoke worker output exceeded its size limit"); + } + return next; +} + +async function waitForWorker(worker, timeoutMs) { + return new Promise((resolvePromise, rejectPromise) => { + let timedOut = false; + let terminationTimeout; + const timeout = setTimeout(() => { + timedOut = true; + worker.kill(); + terminationTimeout = setTimeout(() => { + rejectPromise(new Error("real Host smoke worker did not exit after its deadline")); + }, 5_000); + }, timeoutMs); + worker.once("error", (error) => { + clearTimeout(timeout); + clearTimeout(terminationTimeout); + rejectPromise(error); + }); + worker.once("exit", (code, signal) => { + clearTimeout(timeout); + clearTimeout(terminationTimeout); + if (timedOut) { + rejectPromise(new Error("real Host smoke worker exceeded its deadline")); + } else { + resolvePromise({ code, signal }); + } + }); + }); +} + +function lastWorkerPhase(stdout) { + const phases = [...stdout.matchAll(/^phase:([a-z_]+)$/gm)]; + return phases.at(-1)?.[1] ?? "worker_start"; +} + +function summarizeModelRequest(index, path, body) { + const messages = Array.isArray(body.messages) ? body.messages : []; + const lastMessage = messages.at(-1); + const tools = Array.isArray(body.tools) ? body.tools : []; + return { + index, + path, + model: typeof body.model === "string" ? body.model : typeof body.model, + lastMessage: { + role: + typeof lastMessage?.role === "string" + ? lastMessage.role + : typeof lastMessage?.role, + text: summarizeMessageContent(lastMessage?.content), + }, + tools: { + count: tools.length, + names: tools.map((tool) => + typeof tool?.function?.name === "string" + ? tool.function.name + : typeof tool?.name === "string" + ? tool.name + : "", + ), + }, + toolChoice: summarizeToolChoice(body.tool_choice), + stream: body.stream, + }; +} + +function summarizeMessageContent(content) { + if (typeof content === "string") { + return summarizeText(content); + } + if (!Array.isArray(content)) { + return `<${typeof content}>`; + } + return summarizeText( + content + .map((part) => + typeof part?.text === "string" + ? part.text + : typeof part?.content === "string" + ? part.content + : `<${String(part?.type ?? typeof part)}>`, + ) + .join(" "), + ); +} + +function summarizeText(value) { + const compact = value.replace(/\s+/g, " ").trim(); + return compact.length <= 160 ? compact : `${compact.slice(0, 157)}...`; +} + +function summarizeToolChoice(toolChoice) { + if (toolChoice === undefined) { + return ""; + } + if (typeof toolChoice === "string") { + return toolChoice; + } + const functionName = toolChoice?.function?.name; + if (typeof functionName === "string") { + return { type: toolChoice.type ?? "function", functionName }; + } + return `<${typeof toolChoice}>`; +} + +function formatRequestTraces(traces, apiKey) { + return `secret-safe request trace:\n${JSON.stringify(traces, null, 2).replaceAll(apiKey, "[redacted]")}`; +} + +async function readApiKeyFromStdin() { + let value = ""; + for await (const chunk of process.stdin) { + value += Buffer.isBuffer(chunk) ? chunk.toString("utf8") : String(chunk); + if (Buffer.byteLength(value, "utf8") > 512) { + throw new Error("real Host smoke credential exceeded its size limit"); + } + } + const apiKey = value.trimEnd(); + if (apiKey.length === 0) { + throw new Error("real Host smoke credential is unavailable"); + } + return apiKey; +} + +function requiredEnvironment(name) { + const value = process.env[name]; + if (value === undefined || value.length === 0) { + throw new Error(`real Host smoke environment is missing ${name}`); + } + return value; +} + +function renderError(error) { + const values = []; + let current = error; + for (let depth = 0; depth < 8 && current instanceof Error; depth += 1) { + values.push(current.name, current.message, current.stack ?? ""); + current = current.cause; + } + return values.join("\n"); +} + +async function assertTreeDoesNotContain(directory, secret) { + const entries = await readdir(directory, { withFileTypes: true }); + const secretBytes = Buffer.from(secret, "utf8"); + for (const entry of entries) { + const path = join(directory, entry.name); + assert.equal(path.includes(secret), false, "API key leaked to an isolated path"); + if (entry.isDirectory()) { + await assertTreeDoesNotContain(path, secret); + } else if (entry.isFile()) { + const contents = await readFile(path); + assert.equal(contents.includes(secretBytes), false, "API key leaked to an isolated file"); + } + } +} + +async function closeServer(server) { + if (!server.listening) { + return; + } + await new Promise((resolvePromise, rejectPromise) => { + server.close((error) => { + if (error === undefined) { + resolvePromise(); + } else { + rejectPromise(error); + } + }); + }); +} diff --git a/src/apps/sdk-host/Cargo.toml b/src/apps/sdk-host/Cargo.toml index f513aaffab..a219866951 100644 --- a/src/apps/sdk-host/Cargo.toml +++ b/src/apps/sdk-host/Cargo.toml @@ -43,6 +43,7 @@ tokio = { workspace = true, features = ["io-std", "io-util", "macros", "rt-multi tokio-util = { workspace = true, features = ["codec"] } tracing = { workspace = true } tracing-subscriber = { workspace = true } +url = { workspace = true } [dev-dependencies] tempfile = "3" diff --git a/src/apps/sdk-host/src/main.rs b/src/apps/sdk-host/src/main.rs index 3e02d28870..e993af5117 100644 --- a/src/apps/sdk-host/src/main.rs +++ b/src/apps/sdk-host/src/main.rs @@ -1,6 +1,8 @@ +mod model_source; mod runtime; use anyhow::{Context, Result}; +use std::sync::Arc; async fn run_host() -> Result<()> { tracing_subscriber::fmt() @@ -20,6 +22,10 @@ async fn run_host() -> Result<()> { bitfun_core::infrastructure::ai::AIClientFactory::initialize_global() .await .context("Failed to initialize global AI client factory")?; + let config_service = bitfun_core::service::config::get_global_config_service().await?; + let installer = Arc::new(model_source::ConfigTemporaryModelInstaller::new( + config_service, + )); let host = runtime::SdkHostRuntime::build(&workspace_root) .await @@ -27,6 +33,7 @@ async fn run_host() -> Result<()> { bitfun_sdk_host_app::transport::serve_stdio( host.agent_runtime().clone(), host.workspace_root().to_string_lossy().into_owned(), + installer, ) .await .context("Agent SDK Host transport failed") diff --git a/src/apps/sdk-host/src/model_source.rs b/src/apps/sdk-host/src/model_source.rs new file mode 100644 index 0000000000..d6bca74eae --- /dev/null +++ b/src/apps/sdk-host/src/model_source.rs @@ -0,0 +1,208 @@ +use std::sync::Arc; + +use bitfun_core::infrastructure::ai::AIClientFactory; +use bitfun_core::service::config::{ + model_runtime_binding_fingerprint, AIModelConfig, ConfigService, ModelCapability, ModelCategory, +}; +use bitfun_sdk_host::host::{TemporaryModelInstallError, TemporaryModelInstaller}; +use bitfun_sdk_host::protocol::{TemporaryModelConfig, TemporaryModelProvider}; + +pub(crate) struct ConfigTemporaryModelInstaller { + config: Arc, +} + +impl ConfigTemporaryModelInstaller { + pub(crate) fn new(config: Arc) -> Self { + Self { config } + } +} + +fn resolve_temporary_model( + model: TemporaryModelConfig, +) -> Result { + let model_name = model.model.trim().to_string(); + if model_name.is_empty() || model.api_key.trim().is_empty() { + return Err(TemporaryModelInstallError::InvalidModel); + } + + let (provider_id, default_base_url) = match model.provider { + TemporaryModelProvider::Openai => ("openai", "https://api.openai.com/v1"), + TemporaryModelProvider::Responses => ("responses", "https://api.openai.com/v1"), + TemporaryModelProvider::Anthropic => ("anthropic", "https://api.anthropic.com"), + TemporaryModelProvider::Gemini => { + ("gemini", "https://generativelanguage.googleapis.com/v1beta") + } + }; + let base_url = model + .base_url + .unwrap_or_else(|| default_base_url.to_string()); + let parsed = + url::Url::parse(&base_url).map_err(|_| TemporaryModelInstallError::InvalidBaseUrl)?; + if !matches!(parsed.scheme(), "http" | "https") + || parsed.host_str().is_none() + || !parsed.username().is_empty() + || parsed.password().is_some() + || parsed.query().is_some() + || parsed.fragment().is_some() + { + return Err(TemporaryModelInstallError::InvalidBaseUrl); + } + + let mut config = AIModelConfig { + id: String::new(), + name: model_name.clone(), + provider: provider_id.to_string(), + model_name, + base_url, + api_key: model.api_key, + enabled: true, + category: ModelCategory::GeneralChat, + capabilities: vec![ModelCapability::TextChat], + ..AIModelConfig::default() + }; + let fingerprint = model_runtime_binding_fingerprint(&config); + config.id = format!("sdk:{provider_id}:{}", &fingerprint[..24]); + Ok(config) +} + +#[async_trait::async_trait] +impl TemporaryModelInstaller for ConfigTemporaryModelInstaller { + async fn install( + &self, + model: TemporaryModelConfig, + ) -> Result { + let config = resolve_temporary_model(model)?; + let model_id = config.id.clone(); + self.config + .install_runtime_ai_model(config) + .await + .map_err(|_| TemporaryModelInstallError::Internal)?; + Ok(model_id) + } + + async fn remove(&self, model_id: &str) { + self.config.remove_runtime_ai_model(model_id).await; + if let Ok(factory) = AIClientFactory::get_global().await { + factory.invalidate_model(model_id); + } + } +} + +#[cfg(test)] +mod tests { + use bitfun_core::service::config::{ModelCapability, ModelCategory}; + use bitfun_sdk_host::host::TemporaryModelInstallError; + use bitfun_sdk_host::protocol::{TemporaryModelConfig, TemporaryModelProvider}; + + use super::resolve_temporary_model; + + fn temporary_model( + provider: TemporaryModelProvider, + api_key: &str, + base_url: Option<&str>, + ) -> TemporaryModelConfig { + TemporaryModelConfig { + provider, + model: "fixture-model".to_string(), + api_key: api_key.to_string(), + base_url: base_url.map(str::to_string), + } + } + + #[test] + fn provider_defaults_and_minimal_model_fields_are_resolved() { + let cases = [ + ( + TemporaryModelProvider::Openai, + "openai", + "https://api.openai.com/v1", + ), + ( + TemporaryModelProvider::Responses, + "responses", + "https://api.openai.com/v1", + ), + ( + TemporaryModelProvider::Anthropic, + "anthropic", + "https://api.anthropic.com", + ), + ( + TemporaryModelProvider::Gemini, + "gemini", + "https://generativelanguage.googleapis.com/v1beta", + ), + ]; + + for (provider, provider_id, default_url) in cases { + let model = + resolve_temporary_model(temporary_model(provider, "fixture-secret", None)).unwrap(); + assert!(model.id.starts_with(&format!("sdk:{provider_id}:"))); + assert_eq!(model.id.len(), "sdk::".len() + provider_id.len() + 24); + assert_eq!(model.name, "fixture-model"); + assert_eq!(model.provider, provider_id); + assert_eq!(model.model_name, "fixture-model"); + assert_eq!(model.base_url, default_url); + assert_eq!(model.api_key, "fixture-secret"); + assert!(model.enabled); + assert!(matches!(model.category, ModelCategory::GeneralChat)); + assert_eq!(model.capabilities, vec![ModelCapability::TextChat]); + assert!(model.request_url.is_none()); + assert!(model.context_window.is_none()); + assert!(model.custom_headers.is_none()); + assert!(model.custom_request_body.is_none()); + } + } + + #[test] + fn model_id_is_deterministic_across_api_key_rotation() { + let first = resolve_temporary_model(temporary_model( + TemporaryModelProvider::Openai, + "fixture-secret-one", + Some("http://127.0.0.1:43123/v1"), + )) + .unwrap(); + let second = resolve_temporary_model(temporary_model( + TemporaryModelProvider::Openai, + "fixture-secret-two", + Some("http://127.0.0.1:43123/v1"), + )) + .unwrap(); + + assert_eq!(first.id, second.id); + assert!(!first.id.contains("fixture-secret")); + } + + #[test] + fn invalid_model_values_fail_without_echoing_input() { + for base_url in [ + "not-a-url", + "ftp://example.com/v1", + "https://user:password@example.com/v1", + "https://example.com/v1?secret=value", + "https://example.com/v1#fragment", + ] { + assert!(matches!( + resolve_temporary_model(temporary_model( + TemporaryModelProvider::Openai, + "fixture-secret", + Some(base_url), + )), + Err(TemporaryModelInstallError::InvalidBaseUrl) + )); + } + + for (model, api_key) in [("", "fixture-secret"), ("fixture-model", " ")] { + let mut temporary = temporary_model( + TemporaryModelProvider::Openai, + api_key, + Some("https://example.com/v1"), + ); + temporary.model = model.to_string(); + assert!(matches!( + resolve_temporary_model(temporary), + Err(TemporaryModelInstallError::InvalidModel) + )); + } + } +} diff --git a/src/apps/sdk-host/src/transport.rs b/src/apps/sdk-host/src/transport.rs index cb1002d48a..e0757efa53 100644 --- a/src/apps/sdk-host/src/transport.rs +++ b/src/apps/sdk-host/src/transport.rs @@ -13,7 +13,9 @@ use tokio::time::{timeout, Instant}; use tokio_util::codec::{FramedRead, LinesCodec, LinesCodecError}; use tokio_util::sync::CancellationToken; -use bitfun_sdk_host::host::{ConnectionControl, HostOutput, SdkHostConfig, SdkHostConnection}; +use bitfun_sdk_host::host::{ + ConnectionControl, HostOutput, SdkHostConfig, SdkHostConnection, TemporaryModelInstaller, +}; use bitfun_sdk_host::protocol::{ JsonRpcErrorResponse, JsonRpcRequest, RequestId, METHOD_INITIALIZE, METHOD_QUERY_CANCEL, METHOD_SESSION_CLOSE, METHOD_SHUTDOWN, @@ -91,6 +93,7 @@ where pub async fn serve_streams( runtime: AgentRuntime, default_cwd: impl Into, + temporary_model_installer: Arc, reader: Reader, writer: Writer, config: SdkHostTransportConfig, @@ -108,8 +111,13 @@ where config.max_output_line_bytes, )); let output_failed = output.failure_token(); - let connection = - SdkHostConnection::with_output(runtime, default_cwd, output.clone(), config.host); + let connection = SdkHostConnection::with_output( + runtime, + default_cwd, + output.clone(), + config.host, + temporary_model_installer, + ); let connection_failed = connection.connection_failed_token(); let mut lines = FramedRead::new( reader, @@ -372,10 +380,12 @@ async fn drain_requests(requests: &mut JoinSet, drain_timeout pub async fn serve_stdio( runtime: AgentRuntime, default_cwd: impl Into, + temporary_model_installer: Arc, ) -> Result<(), std::io::Error> { serve_streams( runtime, default_cwd, + temporary_model_installer, tokio::io::stdin(), tokio::io::stdout(), SdkHostTransportConfig::default(), diff --git a/src/apps/sdk-host/tests/stdio_process.rs b/src/apps/sdk-host/tests/stdio_process.rs index 3709a7a6ab..a26a081cc9 100644 --- a/src/apps/sdk-host/tests/stdio_process.rs +++ b/src/apps/sdk-host/tests/stdio_process.rs @@ -6,6 +6,7 @@ use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; #[tokio::test] async fn standalone_sdk_host_negotiates_and_shuts_down_without_cli() { + const FIXTURE_KEY: &str = "bitfun-sdk-fixture-key-7f6b1d"; let temp = tempfile::tempdir().expect("isolated SDK Host environment"); let workspace = temp.path().join("workspace"); let user_root = temp.path().join("user-root"); @@ -42,15 +43,24 @@ async fn standalone_sdk_host_negotiates_and_shuts_down_without_cli() { 1, "initialize", json!({ - "protocolVersion": 1, + "protocolVersion": 2, "clientInfo": { "name": "standalone-process-fixture", "version": "0.1.0" }, - "capabilities": { "serverNotifications": true } + "capabilities": { "serverNotifications": true }, + "model": { + "provider": "openai", + "model": "fixture-model", + "apiKey": FIXTURE_KEY, + "baseUrl": "http://127.0.0.1:43123/v1" + } }), ) .await; let initialized = read_response(&mut stdout, "initialize").await; assert_eq!(initialized["id"], 1); - assert_eq!(initialized["result"]["protocolVersion"], 1); + assert_eq!(initialized["result"]["protocolVersion"], 2); + assert!(initialized["result"]["modelId"] + .as_str() + .is_some_and(|model_id| model_id.starts_with("sdk:openai:"))); send_request(&mut stdin, 2, "shutdown", json!({})).await; let shutdown = read_response(&mut stdout, "shutdown").await; @@ -67,11 +77,54 @@ async fn standalone_sdk_host_negotiates_and_shuts_down_without_cli() { .read_to_end(&mut stderr_output) .await .expect("read SDK Host stderr"); + let mut stdout_remainder = Vec::new(); + stdout + .read_to_end(&mut stdout_remainder) + .await + .expect("read remaining SDK Host stdout"); assert!( status.success(), "SDK Host failed: {}", String::from_utf8_lossy(&stderr_output) ); + + let mut captured = serde_json::to_vec(&initialized).unwrap(); + captured.extend(serde_json::to_vec(&shutdown).unwrap()); + captured.extend(stdout_remainder); + captured.extend(stderr_output); + assert!(!contains_bytes(&captured, FIXTURE_KEY.as_bytes())); + for root in [&user_root, &home_root, &config_root] { + for contents in read_regular_files_recursively(root) { + assert!( + !contains_bytes(&contents, FIXTURE_KEY.as_bytes()), + "isolated SDK Host storage contained fixture credentials" + ); + } + } +} + +fn contains_bytes(haystack: &[u8], needle: &[u8]) -> bool { + !needle.is_empty() + && haystack + .windows(needle.len()) + .any(|window| window == needle) +} + +fn read_regular_files_recursively(root: &std::path::Path) -> Vec> { + let mut contents = Vec::new(); + let mut pending = vec![root.to_path_buf()]; + while let Some(path) = pending.pop() { + for entry in std::fs::read_dir(&path).expect("read isolated SDK Host storage directory") { + let entry = entry.expect("read isolated SDK Host storage entry"); + let file_type = entry.file_type().expect("read SDK Host storage file type"); + if file_type.is_dir() { + pending.push(entry.path()); + } else if file_type.is_file() { + contents.push(std::fs::read(entry.path()).expect("read SDK Host storage file")); + } + } + } + contents } async fn send_request( @@ -106,5 +159,5 @@ async fn read_response( .expect("read SDK Host stdout"); assert_ne!(bytes, 0, "SDK Host stdout closed during {operation}"); serde_json::from_str(&line) - .unwrap_or_else(|error| panic!("SDK Host stdout was not JSON: {error}: {line}")) + .unwrap_or_else(|error| panic!("SDK Host stdout was not JSON: {error}")) } diff --git a/src/apps/sdk-host/tests/stdio_transport.rs b/src/apps/sdk-host/tests/stdio_transport.rs index b30a8c5625..d6f7215729 100644 --- a/src/apps/sdk-host/tests/stdio_transport.rs +++ b/src/apps/sdk-host/tests/stdio_transport.rs @@ -1,5 +1,5 @@ use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use async_trait::async_trait; use bitfun_agent_runtime::sdk::{ @@ -9,6 +9,8 @@ use bitfun_agent_runtime::sdk::{ AgentSessionWorkspaceRequest, AgentSubmissionPort, AgentSubmissionRequest, AgentSubmissionResult, AgentTransientSessionDiscardRequest, PortResult, }; +use bitfun_sdk_host::host::{TemporaryModelInstallError, TemporaryModelInstaller}; +use bitfun_sdk_host::protocol::TemporaryModelConfig; use bitfun_sdk_host_app::transport::{serve_streams, SdkHostTransportConfig}; use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; use tokio::sync::Notify; @@ -16,6 +18,29 @@ use tokio::time::{timeout, Duration}; struct MinimalOwner; +#[derive(Default)] +struct FakeTemporaryModelInstaller { + removed: Mutex>, +} + +#[async_trait] +impl TemporaryModelInstaller for FakeTemporaryModelInstaller { + async fn install( + &self, + _model: TemporaryModelConfig, + ) -> Result { + Ok("sdk:openai:transport".to_string()) + } + + async fn remove(&self, model_id: &str) { + self.removed.lock().unwrap().push(model_id.to_string()); + } +} + +fn fake_installer() -> Arc { + Arc::new(FakeTemporaryModelInstaller::default()) +} + fn created_session_result( session_id: impl Into, request: AgentSessionCreateRequest, @@ -185,9 +210,11 @@ async fn stdio_transport_serves_initialize_and_shutdown_without_non_protocol_std let (client, server) = tokio::io::duplex(16 * 1024); let (client_read, mut client_write) = tokio::io::split(client); let (server_read, server_write) = tokio::io::split(server); + let installer = Arc::new(FakeTemporaryModelInstaller::default()); let task = tokio::spawn(serve_streams( runtime, "D:/workspace/project", + installer.clone(), server_read, server_write, SdkHostTransportConfig::default(), @@ -195,7 +222,7 @@ async fn stdio_transport_serves_initialize_and_shutdown_without_non_protocol_std client_write .write_all( concat!( - "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":2,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true},\"model\":{\"provider\":\"openai\",\"model\":\"fixture-model\",\"apiKey\":\"fixture-secret\"}}}\n", "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"shutdown\",\"params\":{}}\n" ) .as_bytes(), @@ -211,11 +238,16 @@ async fn stdio_transport_serves_initialize_and_shutdown_without_non_protocol_std serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); assert_eq!(initialized["id"], 1); - assert_eq!(initialized["result"]["protocolVersion"], 1); + assert_eq!(initialized["result"]["protocolVersion"], 2); + assert_eq!(initialized["result"]["modelId"], "sdk:openai:transport"); assert_eq!(shutdown["id"], 2); assert_eq!(shutdown["result"]["accepted"], true); assert!(lines.next_line().await.unwrap().is_none()); task.await.unwrap().unwrap(); + assert_eq!( + installer.removed.lock().unwrap().as_slice(), + &["sdk:openai:transport".to_string()] + ); } #[tokio::test] @@ -232,6 +264,7 @@ async fn stdio_transport_executes_json_rpc_notifications_without_replying() { let task = tokio::spawn(serve_streams( runtime, "D:/workspace/project", + fake_installer(), server_read, server_write, SdkHostTransportConfig::default(), @@ -239,7 +272,7 @@ async fn stdio_transport_executes_json_rpc_notifications_without_replying() { client_write .write_all( concat!( - "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":2,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true},\"model\":{\"provider\":\"openai\",\"model\":\"fixture-model\",\"apiKey\":\"fixture-secret\"}}}\n", "{\"jsonrpc\":\"2.0\",\"method\":\"shutdown\",\"params\":{}}\n" ) .as_bytes(), @@ -274,6 +307,7 @@ async fn malformed_and_oversized_lines_fail_closed_with_standard_parse_errors() let task = tokio::spawn(serve_streams( runtime, "D:/workspace/project", + fake_installer(), server_read, server_write, SdkHostTransportConfig { @@ -324,6 +358,7 @@ async fn transport_accepts_input_while_an_owner_call_is_pending_and_bounds_reque let task = tokio::spawn(serve_streams( runtime, "D:/workspace/project", + fake_installer(), server_read, server_write, SdkHostTransportConfig { @@ -337,7 +372,7 @@ async fn transport_accepts_input_while_an_owner_call_is_pending_and_bounds_reque client_write .write_all( concat!( - "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":2,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true},\"model\":{\"provider\":\"openai\",\"model\":\"fixture-model\",\"apiKey\":\"fixture-secret\"}}}\n", "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"session/create\",\"params\":{}}\n", "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"session/create\",\"params\":{}}\n" ) @@ -395,6 +430,7 @@ async fn shutdown_remains_available_when_the_data_request_budget_is_exhausted() let task = tokio::spawn(serve_streams( runtime, "D:/workspace/project", + fake_installer(), server_read, server_write, SdkHostTransportConfig { @@ -408,7 +444,7 @@ async fn shutdown_remains_available_when_the_data_request_budget_is_exhausted() )); client_write .write_all( - b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":2,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true},\"model\":{\"provider\":\"openai\",\"model\":\"fixture-model\",\"apiKey\":\"fixture-secret\"}}}\n", ) .await .unwrap(); @@ -463,6 +499,7 @@ async fn duplicate_initialize_does_not_abort_an_in_flight_request() { let task = tokio::spawn(serve_streams( runtime, "D:/workspace/project", + fake_installer(), server_read, server_write, SdkHostTransportConfig::default(), @@ -470,9 +507,9 @@ async fn duplicate_initialize_does_not_abort_an_in_flight_request() { client_write .write_all( concat!( - "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":2,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true},\"model\":{\"provider\":\"openai\",\"model\":\"fixture-model\",\"apiKey\":\"fixture-secret\"}}}\n", "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"session/create\",\"params\":{}}\n", - "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n" + "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"initialize\",\"params\":{\"protocolVersion\":2,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true},\"model\":{\"provider\":\"openai\",\"model\":\"fixture-model\",\"apiKey\":\"fixture-secret\"}}}\n" ) .as_bytes(), ) @@ -515,9 +552,11 @@ async fn connection_eof_cleans_a_session_created_after_its_request_is_aborted() let (client, server) = tokio::io::duplex(16 * 1024); let (client_read, mut client_write) = tokio::io::split(client); let (server_read, server_write) = tokio::io::split(server); + let installer = Arc::new(FakeTemporaryModelInstaller::default()); let mut task = tokio::spawn(serve_streams( runtime, "D:/workspace/project", + installer.clone(), server_read, server_write, SdkHostTransportConfig { @@ -528,7 +567,7 @@ async fn connection_eof_cleans_a_session_created_after_its_request_is_aborted() client_write .write_all( concat!( - "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":2,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true},\"model\":{\"provider\":\"openai\",\"model\":\"fixture-model\",\"apiKey\":\"fixture-secret\"}}}\n", "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"session/create\",\"params\":{}}\n" ) .as_bytes(), @@ -555,6 +594,10 @@ async fn connection_eof_cleans_a_session_created_after_its_request_is_aborted() .unwrap(); assert_eq!(owner.deleted.load(Ordering::Acquire), 1); assert!(lines.next_line().await.unwrap().is_none()); + assert_eq!( + installer.removed.lock().unwrap().as_slice(), + &["sdk:openai:transport".to_string()] + ); } #[tokio::test] @@ -572,6 +615,7 @@ async fn explicit_shutdown_bounds_request_drain_and_transient_cleanup_together() let mut task = tokio::spawn(serve_streams( runtime, "D:/workspace/project", + fake_installer(), server_read, server_write, SdkHostTransportConfig { @@ -582,7 +626,7 @@ async fn explicit_shutdown_bounds_request_drain_and_transient_cleanup_together() client_write .write_all( concat!( - "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":2,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true},\"model\":{\"provider\":\"openai\",\"model\":\"fixture-model\",\"apiKey\":\"fixture-secret\"}}}\n", "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"session/create\",\"params\":{}}\n" ) .as_bytes(), @@ -630,6 +674,7 @@ async fn requests_before_a_successful_initialize_cannot_cross_the_handshake() { let task = tokio::spawn(serve_streams( runtime, "D:/workspace/project", + fake_installer(), server_read, server_write, SdkHostTransportConfig::default(), @@ -637,9 +682,9 @@ async fn requests_before_a_successful_initialize_cannot_cross_the_handshake() { client_write .write_all( concat!( - "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":999,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":999,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true},\"model\":{\"provider\":\"openai\",\"model\":\"fixture-model\",\"apiKey\":\"fixture-secret\"}}}\n", "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"session/create\",\"params\":{}}\n", - "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n" + "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"initialize\",\"params\":{\"protocolVersion\":2,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true},\"model\":{\"provider\":\"openai\",\"model\":\"fixture-model\",\"apiKey\":\"fixture-secret\"}}}\n" ) .as_bytes(), ) @@ -658,7 +703,8 @@ async fn requests_before_a_successful_initialize_cannot_cross_the_handshake() { assert_eq!(pre_initialize["id"], 2); assert_eq!(pre_initialize["error"]["data"]["code"], "not_initialized"); assert_eq!(initialized["id"], 3); - assert_eq!(initialized["result"]["protocolVersion"], 1); + assert_eq!(initialized["result"]["protocolVersion"], 2); + assert_eq!(initialized["result"]["modelId"], "sdk:openai:transport"); client_write .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":4,\"method\":\"shutdown\",\"params\":{}}\n") @@ -684,9 +730,11 @@ async fn blocked_output_times_out_and_ends_the_connection() { let (client, server) = tokio::io::duplex(64); let (_client_read, mut client_write) = tokio::io::split(client); let (server_read, server_write) = tokio::io::split(server); + let installer = Arc::new(FakeTemporaryModelInstaller::default()); let task = tokio::spawn(serve_streams( runtime, "D:/workspace/project", + installer.clone(), server_read, server_write, SdkHostTransportConfig { @@ -696,7 +744,7 @@ async fn blocked_output_times_out_and_ends_the_connection() { )); client_write .write_all( - b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":2,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true},\"model\":{\"provider\":\"openai\",\"model\":\"fixture-model\",\"apiKey\":\"fixture-secret\"}}}\n", ) .await .unwrap(); @@ -706,4 +754,8 @@ async fn blocked_output_times_out_and_ends_the_connection() { .expect("blocked SDK Host output must have a deadline") .unwrap(); assert_eq!(result.unwrap_err().kind(), std::io::ErrorKind::BrokenPipe); + assert_eq!( + installer.removed.lock().unwrap().as_slice(), + &["sdk:openai:transport".to_string()] + ); } diff --git a/src/crates/assembly/core/src/agentic/session/session_manager.rs b/src/crates/assembly/core/src/agentic/session/session_manager.rs index 5bee32f8e8..6985088dd5 100644 --- a/src/crates/assembly/core/src/agentic/session/session_manager.rs +++ b/src/crates/assembly/core/src/agentic/session/session_manager.rs @@ -660,7 +660,13 @@ impl SessionManager { } let config_service = get_global_config_service().await.ok()?; - config_service.get_config(Some("ai")).await.ok() + Self::load_effective_ai_config_from_service(config_service.as_ref()).await + } + + async fn load_effective_ai_config_from_service( + config_service: &crate::service::config::ConfigService, + ) -> Option { + config_service.get_effective_ai_config().await.ok() } pub(crate) async fn resolve_effective_reasoning_preset_for_turn( @@ -9393,6 +9399,7 @@ mod tests { model_runtime_binding_fingerprint as service_model_runtime_binding_fingerprint, AIConfig as ServiceAIConfig, AIModelConfig as ServiceAIModelConfig, }; + use crate::service::config::{ConfigManagerSettings, ConfigService}; use crate::service::session::{ DialogTurnData, DialogTurnKind, DialogTurnRecoveryStatus, ModelRoundData, SessionContextUsage, SessionContextUsageSource, SessionKind, SessionMetadata, @@ -9413,6 +9420,41 @@ mod tests { use std::time::{Duration, SystemTime}; use uuid::Uuid; + #[tokio::test] + async fn runtime_model_is_visible_to_turn_admission_config() { + let dir = tempfile::tempdir().expect("temporary config directory"); + let config = ConfigService::with_settings(ConfigManagerSettings { + path_manager: Some(Arc::new(PathManager::with_user_root_for_tests( + dir.path().join("runtime-turn-admission"), + ))), + auto_save: true, + backup_count: 0, + }) + .await + .expect("test ConfigService"); + config + .install_runtime_ai_model(ServiceAIModelConfig { + id: "sdk:openai:fixture".to_string(), + name: "SDK fixture".to_string(), + provider: "openai".to_string(), + model_name: "fixture-model".to_string(), + base_url: "http://127.0.0.1:43123/v1".to_string(), + api_key: "fixture-secret".to_string(), + enabled: true, + ..ServiceAIModelConfig::default() + }) + .await + .unwrap(); + + let ai_config = SessionManager::load_effective_ai_config_from_service(&config) + .await + .expect("turn admission should see the runtime model"); + assert_eq!( + ai_config.resolve_model_reference("sdk:openai:fixture"), + Some("sdk:openai:fixture".to_string()) + ); + } + struct TestWorkspace { path: PathBuf, } diff --git a/src/crates/assembly/core/src/infrastructure/ai/client_factory.rs b/src/crates/assembly/core/src/infrastructure/ai/client_factory.rs index d595314ad3..dc6378b06b 100644 --- a/src/crates/assembly/core/src/infrastructure/ai/client_factory.rs +++ b/src/crates/assembly/core/src/infrastructure/ai/client_factory.rs @@ -144,12 +144,11 @@ impl AIClientFactory { } async fn resolve_model_id(&self, model_id: &str) -> Result { - let global_config: crate::service::config::GlobalConfig = - self.config_service.get_config(None).await?; + let ai_config = self.config_service.get_effective_ai_config().await?; resolve_required_model_selector( model_id, - |selector| global_config.ai.resolve_model_selection(selector), - |model_ref| global_config.ai.resolve_model_reference(model_ref), + |selector| ai_config.resolve_model_selection(selector), + |model_ref| ai_config.resolve_model_reference(model_ref), ) .map_err(|error| anyhow!(error.to_string())) } @@ -166,10 +165,8 @@ impl AIClientFactory { else { return Ok(client); }; - let global_config: crate::service::config::GlobalConfig = - self.config_service.get_config(None).await?; - let model = global_config - .ai + let ai_config = self.config_service.get_effective_ai_config().await?; + let model = ai_config .models .iter() .find(|model| model.id == model_id) @@ -236,35 +233,31 @@ impl AIClientFactory { if normalized_model_id.is_empty() { return Err(anyhow!("Model configuration id is empty")); } - if global_config - .ai - .models - .iter() - .filter(|model| model.id == normalized_model_id) - .nth(1) - .is_some() - { - return Err(anyhow!( - "Multiple model configurations use the same ID: {}", - normalized_model_id - )); - } - debug!("Creating new AI client: model_id={}", normalized_model_id); - let mut matching_models = global_config - .ai - .models - .iter() - .filter(|m| m.id == normalized_model_id); - let model_config = matching_models - .next() - .ok_or_else(|| anyhow!("Model configuration not found: {}", normalized_model_id))?; - if matching_models.next().is_some() { - return Err(anyhow!( - "Multiple model configurations use the same ID: {}", - normalized_model_id - )); - } + let model_config = if let Some(runtime_model) = self + .config_service + .get_runtime_ai_model(&normalized_model_id) + .await + { + runtime_model + } else { + let mut matching_models = global_config + .ai + .models + .iter() + .filter(|model| model.id == normalized_model_id); + let model = matching_models + .next() + .cloned() + .ok_or_else(|| anyhow!("Model configuration not found: {}", normalized_model_id))?; + if matching_models.next().is_some() { + return Err(anyhow!( + "Multiple model configurations use the same ID: {}", + normalized_model_id + )); + } + model + }; if !model_config.enabled { return Err(anyhow!( @@ -274,7 +267,7 @@ impl AIClientFactory { )); } - let configuration_fingerprint = model_runtime_binding_fingerprint(model_config); + let configuration_fingerprint = model_runtime_binding_fingerprint(&model_config); if expected_configuration_fingerprint .is_some_and(|expected| expected != configuration_fingerprint) { @@ -286,7 +279,7 @@ impl AIClientFactory { let models_dev = load_models_dev_reasoning_catalog().await; let reasoning_projection = - project_model_reasoning_catalog(model_config, models_dev.catalog.as_deref()); + project_model_reasoning_catalog(&model_config, models_dev.catalog.as_deref()); let default_reasoning_preset = resolve_default_reasoning_preset(&reasoning_projection).cloned(); @@ -328,7 +321,7 @@ impl AIClientFactory { #[cfg(not(feature = "subscription-auth"))] let _ = credential_expires_at; - let stream_options = build_stream_options_for_model(&global_config.ai, Some(model_config)); + let stream_options = build_stream_options_for_model(&global_config.ai, Some(&model_config)); let client = apply_default_reasoning_preset( AIClient::new_with_runtime_options(ai_config, proxy_config, stream_options), &reasoning_projection, @@ -578,12 +571,16 @@ pub async fn list_subscription_accounts() -> Vec>, + runtime_ai_models: Arc>>, } /// Configuration import/export format. @@ -59,6 +61,7 @@ impl ConfigService { let service = Self { manager: Arc::new(RwLock::new(manager)), + runtime_ai_models: Arc::new(RwLock::new(BTreeMap::new())), }; let recovered_with_defaults = service @@ -91,6 +94,43 @@ impl ConfigService { } } + pub async fn install_runtime_ai_model(&self, model: AIModelConfig) -> BitFunResult<()> { + if model.id.trim().is_empty() { + return Err(BitFunError::validation( + "Runtime model id is required".to_string(), + )); + } + self.runtime_ai_models + .write() + .await + .insert(model.id.clone(), model); + Ok(()) + } + + pub async fn get_runtime_ai_model(&self, model_id: &str) -> Option { + self.runtime_ai_models.read().await.get(model_id).cloned() + } + + pub async fn get_effective_ai_config(&self) -> BitFunResult { + let mut ai: AIConfig = self.get_config(Some("ai")).await?; + for runtime_model in self.runtime_ai_models.read().await.values() { + if let Some(model) = ai + .models + .iter_mut() + .find(|model| model.id == runtime_model.id) + { + *model = runtime_model.clone(); + } else { + ai.models.push(runtime_model.clone()); + } + } + Ok(ai) + } + + pub async fn remove_runtime_ai_model(&self, model_id: &str) { + self.runtime_ai_models.write().await.remove(model_id); + } + /// Sets a configuration value (supports dot-paths). /// /// When the path touches AI models / default model slots / agent-model @@ -620,6 +660,21 @@ mod tests { } } + fn runtime_model(id: &str, key: &str) -> AIModelConfig { + AIModelConfig { + id: id.to_string(), + name: "SDK fixture".to_string(), + provider: "openai".to_string(), + model_name: "fixture-model".to_string(), + base_url: "http://127.0.0.1:43123/v1".to_string(), + api_key: key.to_string(), + enabled: true, + category: ModelCategory::GeneralChat, + capabilities: vec![ModelCapability::TextChat], + ..AIModelConfig::default() + } + } + async fn test_service(name: &str) -> (ConfigService, tempfile::TempDir) { let dir = tempfile::tempdir().expect("tempdir"); let user_root = dir.path().join(name); @@ -636,6 +691,110 @@ mod tests { (service, dir) } + #[tokio::test] + async fn runtime_ai_model_is_effective_but_never_persisted() { + let (service, _dir) = test_service("runtime-overlay-test").await; + + service + .install_runtime_ai_model(runtime_model("sdk:openai:fixture", "fixture-secret")) + .await + .unwrap(); + let runtime = service + .get_runtime_ai_model("sdk:openai:fixture") + .await + .unwrap(); + assert_eq!(runtime.api_key, "fixture-secret"); + let effective_ai = service.get_effective_ai_config().await.unwrap(); + assert!(effective_ai + .models + .iter() + .any(|model| model.id == "sdk:openai:fixture")); + let persisted: GlobalConfig = service.get_config(None).await.unwrap(); + assert!(!persisted + .ai + .models + .iter() + .any(|model| model.id == "sdk:openai:fixture")); + let persisted_models: Vec = + service.get_config(Some("ai.models")).await.unwrap(); + assert!(!persisted_models + .iter() + .any(|model| model.id == "sdk:openai:fixture")); + let persisted_ai: AIConfig = service.get_config(Some("ai")).await.unwrap(); + assert!(!persisted_ai + .models + .iter() + .any(|model| model.id == "sdk:openai:fixture")); + + let export = service.export_config().await.unwrap(); + let export_json = serde_json::to_string(&export).unwrap(); + assert!(!export_json.contains("sdk:openai:fixture")); + assert!(!export_json.contains("fixture-secret")); + + service + .reconcile_models("runtime-overlay-test") + .await + .unwrap(); + service + .add_ai_model(model("persisted", true, ModelCategory::GeneralChat)) + .await + .unwrap(); + let app_file = service + .get_statistics() + .await + .config_directory + .join("app.json"); + let disk = tokio::fs::read_to_string(app_file).await.unwrap(); + assert!(!disk.contains("sdk:openai:fixture")); + assert!(!disk.contains("fixture-secret")); + + let backup = service.create_backup().await.unwrap(); + let backup_text = tokio::fs::read_to_string(backup).await.unwrap(); + assert!(!backup_text.contains("sdk:openai:fixture")); + assert!(!backup_text.contains("fixture-secret")); + + service.remove_runtime_ai_model("sdk:openai:fixture").await; + let effective = service.get_effective_ai_config().await.unwrap(); + assert!(!effective + .models + .iter() + .any(|model| model.id == "sdk:openai:fixture")); + assert!(service + .get_runtime_ai_model("sdk:openai:fixture") + .await + .is_none()); + } + + #[tokio::test] + async fn runtime_ai_model_overlays_a_persisted_duplicate_for_effective_reads() { + let (service, _dir) = test_service("runtime-overlay-duplicate-test").await; + service + .add_ai_model(runtime_model("sdk:openai:duplicate", "persisted-secret")) + .await + .unwrap(); + service + .install_runtime_ai_model(runtime_model("sdk:openai:duplicate", "runtime-secret")) + .await + .unwrap(); + + let effective = service.get_effective_ai_config().await.unwrap(); + let matches = effective + .models + .iter() + .filter(|model| model.id == "sdk:openai:duplicate") + .collect::>(); + assert_eq!(matches.len(), 1); + assert_eq!(matches[0].api_key, "runtime-secret"); + + let persisted: AIConfig = service.get_config(Some("ai")).await.unwrap(); + let persisted = persisted + .models + .iter() + .find(|model| model.id == "sdk:openai:duplicate") + .unwrap(); + assert_eq!(persisted.api_key, "persisted-secret"); + } + #[tokio::test] async fn review_team_policy_config_survives_service_restart() { let dir = tempfile::tempdir().expect("tempdir"); diff --git a/src/crates/interfaces/sdk-host/src/host.rs b/src/crates/interfaces/sdk-host/src/host.rs index 6f709c74c3..150e3eb896 100644 --- a/src/crates/interfaces/sdk-host/src/host.rs +++ b/src/crates/interfaces/sdk-host/src/host.rs @@ -27,10 +27,10 @@ use crate::protocol::{ QueryCancelParams, QueryCancelResult, QueryEvent, QueryEventParams, QueryOutput, QueryResultError, QueryResultParams, QueryStartParams, QueryStartResult, QueryTerminalStatus, RecoveryAction, RequestId, SessionCloseParams, SessionCloseResult, SessionCreateParams, - SessionCreateResult, SessionLifetime, ShutdownParams, ShutdownResult, JSON_RPC_VERSION, - METHOD_INITIALIZE, METHOD_QUERY_CANCEL, METHOD_QUERY_START, METHOD_SESSION_CLOSE, - METHOD_SESSION_CREATE, METHOD_SHUTDOWN, NOTIFICATION_QUERY_EVENT, NOTIFICATION_QUERY_RESULT, - PROTOCOL_VERSION, + SessionCreateResult, SessionLifetime, ShutdownParams, ShutdownResult, TemporaryModelConfig, + JSON_RPC_VERSION, METHOD_INITIALIZE, METHOD_QUERY_CANCEL, METHOD_QUERY_START, + METHOD_SESSION_CLOSE, METHOD_SESSION_CREATE, METHOD_SHUTDOWN, NOTIFICATION_QUERY_EVENT, + NOTIFICATION_QUERY_RESULT, PROTOCOL_VERSION, }; const DEFAULT_SESSION_NAME: &str = "BitFun SDK query"; @@ -82,6 +82,22 @@ pub trait HostOutput: Send + Sync { async fn send(&self, value: serde_json::Value) -> Result<(), ()>; } +#[async_trait::async_trait] +pub trait TemporaryModelInstaller: Send + Sync { + async fn install( + &self, + model: TemporaryModelConfig, + ) -> Result; + async fn remove(&self, model_id: &str); +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TemporaryModelInstallError { + InvalidModel, + InvalidBaseUrl, + Internal, +} + struct ChannelHostOutput(mpsc::Sender); #[async_trait::async_trait] @@ -96,6 +112,7 @@ struct ConnectionInner { runtime_version: &'static str, default_cwd: String, output: Arc, + temporary_model_installer: Arc, state: Arc>, request_budget: Arc, control_request_budget: Arc, @@ -107,7 +124,8 @@ struct ConnectionInner { #[derive(Default)] struct ConnectionState { - initialized: bool, + initialization: InitializationState, + model_id: Option, shutting_down: bool, cleanup_failed: bool, sessions: HashMap, @@ -120,6 +138,14 @@ struct ConnectionState { untracked_transient_cleanups: HashMap, } +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +enum InitializationState { + #[default] + Uninitialized, + Installing, + Initialized, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum QueryReservationError { Unavailable, @@ -176,12 +202,14 @@ impl SdkHostConnection { default_cwd: impl Into, output: mpsc::Sender, config: SdkHostConfig, + temporary_model_installer: Arc, ) -> Self { Self::with_output( runtime, default_cwd, Arc::new(ChannelHostOutput(output)), config, + temporary_model_installer, ) } @@ -190,6 +218,7 @@ impl SdkHostConnection { default_cwd: impl Into, output: Arc, config: SdkHostConfig, + temporary_model_installer: Arc, ) -> Self { Self { inner: Arc::new(ConnectionInner { @@ -197,6 +226,7 @@ impl SdkHostConnection { runtime_version: env!("CARGO_PKG_VERSION"), default_cwd: default_cwd.into(), output, + temporary_model_installer, state: Arc::new(Mutex::new(ConnectionState::default())), request_budget: Arc::new(Semaphore::new(config.max_in_flight_requests.max(1))), control_request_budget: Arc::new(Semaphore::new( @@ -278,7 +308,7 @@ impl SdkHostConnection { let (initialized, shutting_down, cleanup_failed) = { let state = self.inner.state.lock().await; ( - state.initialized, + state.initialization == InitializationState::Initialized, state.shutting_down, state.cleanup_failed || !state.untracked_transient_cleanups.is_empty(), ) @@ -376,7 +406,7 @@ impl SdkHostConnection { /// Reports whether this connection has completed the required initialize /// handshake so a transport can serialize re-initialization safely. pub async fn is_initialized(&self) -> bool { - self.inner.state.lock().await.initialized + self.inner.state.lock().await.initialization == InitializationState::Initialized } async fn reap_finished_pending_session_tasks(&self) { @@ -555,6 +585,17 @@ impl SdkHostConnection { } } } + let model_id = { + let mut state = self.inner.state.lock().await; + let model_id = state.model_id.take(); + if model_id.is_some() { + state.initialization = InitializationState::Uninitialized; + } + model_id + }; + if let Some(model_id) = model_id { + self.inner.temporary_model_installer.remove(&model_id).await; + } cleanup_complete } @@ -729,26 +770,98 @@ impl SdkHostConnection { .await; return; } - { + let initialization_error = { let mut state = self.inner.state.lock().await; - if state.initialized { + if state.shutting_down { + Some((ErrorCode::Cancelled, "SDK Host connection is shutting down")) + } else if state.initialization != InitializationState::Uninitialized { + Some(( + ErrorCode::AlreadyInitialized, + "SDK Host connection is already initialized", + )) + } else { + state.initialization = InitializationState::Installing; + None + } + }; + if let Some((code, message)) = initialization_error { + self.send_error( + request.id.clone(), + code, + ErrorStage::Initialize, + false, + None, + message, + ) + .await; + return; + } + let model_id = match self + .inner + .temporary_model_installer + .install(params.model) + .await + { + Ok(model_id) => model_id, + Err(error) => { + let mut state = self.inner.state.lock().await; + if state.initialization == InitializationState::Installing { + state.initialization = InitializationState::Uninitialized; + } drop(state); + let (code, message) = match error { + TemporaryModelInstallError::InvalidModel => ( + ErrorCode::InvalidRequest, + "model provider, model, and apiKey are required", + ), + TemporaryModelInstallError::InvalidBaseUrl => ( + ErrorCode::InvalidRequest, + "baseUrl must be an absolute http or https URL without credentials, query, or fragment", + ), + TemporaryModelInstallError::Internal => ( + ErrorCode::Internal, + "SDK Host could not install the temporary model", + ), + }; self.send_error( request.id.clone(), - ErrorCode::AlreadyInitialized, + code, ErrorStage::Initialize, false, None, - "SDK Host connection is already initialized", + message, ) .await; return; } - state.initialized = true; + }; + let remove_after_shutdown = { + let mut state = self.inner.state.lock().await; + if state.shutting_down { + state.initialization = InitializationState::Uninitialized; + true + } else { + state.model_id = Some(model_id.clone()); + state.initialization = InitializationState::Initialized; + false + } + }; + if remove_after_shutdown { + self.inner.temporary_model_installer.remove(&model_id).await; + self.send_error( + request.id.clone(), + ErrorCode::Cancelled, + ErrorStage::Initialize, + false, + None, + "SDK Host connection is shutting down", + ) + .await; + return; } self.send_success( request.id.clone(), - InitializeResult::current(self.inner.runtime_version), + InitializeResult::current(self.inner.runtime_version, model_id), ) .await; } @@ -773,6 +886,7 @@ impl SdkHostConnection { return; }; let workspace_path = params.cwd.unwrap_or_else(|| self.inner.default_cwd.clone()); + let model_id = self.inner.state.lock().await.model_id.clone(); let result = self .create_leased_session( AgentSessionCreateRequest { @@ -786,7 +900,7 @@ impl SdkHostConnection { workspace_id: None, remote_connection_id: None, remote_ssh_host: None, - model_id: params.model, + model_id, metadata: serde_json::Map::new(), }, workspace_path, @@ -895,6 +1009,7 @@ impl SdkHostConnection { .cwd .clone() .unwrap_or_else(|| self.inner.default_cwd.clone()); + let model_id = self.inner.state.lock().await.model_id.clone(); match self .create_leased_session( AgentSessionCreateRequest { @@ -912,7 +1027,7 @@ impl SdkHostConnection { workspace_id: None, remote_connection_id: None, remote_ssh_host: None, - model_id: params.model.clone(), + model_id, metadata: serde_json::Map::new(), }, workspace_path, diff --git a/src/crates/interfaces/sdk-host/src/protocol.rs b/src/crates/interfaces/sdk-host/src/protocol.rs index 8c72312203..0a6cb3983a 100644 --- a/src/crates/interfaces/sdk-host/src/protocol.rs +++ b/src/crates/interfaces/sdk-host/src/protocol.rs @@ -6,7 +6,7 @@ use serde::de::DeserializeOwned; use serde::{Deserialize, Serialize}; pub const JSON_RPC_VERSION: &str = "2.0"; -pub const PROTOCOL_VERSION: u32 = 1; +pub const PROTOCOL_VERSION: u32 = 2; pub const METHOD_INITIALIZE: &str = "initialize"; pub const METHOD_SESSION_CREATE: &str = "session/create"; @@ -33,7 +33,7 @@ impl RequestId { } } -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Clone, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct JsonRpcRequest { pub jsonrpc: String, @@ -50,6 +50,18 @@ pub struct JsonRpcRequest { pub params: serde_json::Value, } +impl std::fmt::Debug for JsonRpcRequest { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("JsonRpcRequest") + .field("jsonrpc", &self.jsonrpc) + .field("id", &self.id) + .field("method", &self.method) + .field("params", &"") + .finish() + } +} + impl JsonRpcRequest { pub fn params_as(&self) -> Result { serde_json::from_value(self.params.clone()) @@ -164,13 +176,36 @@ impl JsonRpcNotification { } } -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS), ts(export))] #[serde(deny_unknown_fields, rename_all = "camelCase")] pub struct InitializeParams { pub protocol_version: u32, pub client_info: ClientInfo, pub capabilities: ClientCapabilities, + pub model: TemporaryModelConfig, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS), ts(export))] +#[serde(rename_all = "snake_case")] +pub enum TemporaryModelProvider { + Openai, + Responses, + Anthropic, + Gemini, +} + +#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS), ts(export))] +#[serde(deny_unknown_fields, rename_all = "camelCase")] +pub struct TemporaryModelConfig { + pub provider: TemporaryModelProvider, + pub model: String, + pub api_key: String, + #[cfg_attr(feature = "ts", ts(optional = nullable))] + #[serde(default)] + pub base_url: Option, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -196,15 +231,17 @@ pub struct InitializeResult { pub runtime_version: String, pub stability: Stability, pub capabilities: HostCapabilities, + pub model_id: String, } impl InitializeResult { - pub fn current(runtime_version: impl Into) -> Self { + pub fn current(runtime_version: impl Into, model_id: impl Into) -> Self { Self { protocol_version: PROTOCOL_VERSION, runtime_version: runtime_version.into(), stability: Stability::NotDelivered, capabilities: HostCapabilities::current(), + model_id: model_id.into(), } } } diff --git a/src/crates/interfaces/sdk-host/tests/host_lifecycle.rs b/src/crates/interfaces/sdk-host/tests/host_lifecycle.rs index d66dd608da..fcc070a4e4 100644 --- a/src/crates/interfaces/sdk-host/tests/host_lifecycle.rs +++ b/src/crates/interfaces/sdk-host/tests/host_lifecycle.rs @@ -1,3 +1,4 @@ +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -20,8 +21,11 @@ use bitfun_runtime_ports::{ ClockPort, PermissionAuditRecord, PermissionAuditStorePort, PermissionGrant, PermissionReplyStorePort, RuntimeServiceCapability, RuntimeServicePort, }; -use bitfun_sdk_host::host::{ConnectionControl, HostOutput, SdkHostConfig, SdkHostConnection}; -use bitfun_sdk_host::protocol::{JsonRpcRequest, PROTOCOL_VERSION}; +use bitfun_sdk_host::host::{ + ConnectionControl, HostOutput, SdkHostConfig, SdkHostConnection, TemporaryModelInstallError, + TemporaryModelInstaller, +}; +use bitfun_sdk_host::protocol::{JsonRpcRequest, TemporaryModelConfig, PROTOCOL_VERSION}; use tokio::sync::{mpsc, Notify}; use tokio::time::timeout; @@ -29,6 +33,7 @@ use tokio::time::timeout; struct FakeOwner { queue: Mutex>>, created_session_ids: Mutex>, + created_model_ids: Mutex>>, cancel_requests: Mutex>, discard_requests: Mutex>, settlement_requests: Mutex>, @@ -199,6 +204,58 @@ impl FakeOwner { } } +#[derive(Default)] +struct FakeTemporaryModelInstaller { + installed: Mutex>, + removed: Mutex>, + block_first_install: AtomicBool, + fail_first_install: AtomicBool, + install_started: Notify, + release_install: Notify, +} + +impl FakeTemporaryModelInstaller { + fn blocking_first_install() -> Self { + Self { + block_first_install: AtomicBool::new(true), + ..Self::default() + } + } + + fn failing_first_install() -> Self { + Self { + fail_first_install: AtomicBool::new(true), + ..Self::default() + } + } +} + +#[async_trait] +impl TemporaryModelInstaller for FakeTemporaryModelInstaller { + async fn install( + &self, + model: TemporaryModelConfig, + ) -> Result { + self.installed.lock().unwrap().push(model); + if self.fail_first_install.swap(false, Ordering::AcqRel) { + return Err(TemporaryModelInstallError::InvalidModel); + } + if self.block_first_install.swap(false, Ordering::AcqRel) { + self.install_started.notify_one(); + self.release_install.notified().await; + } + Ok("sdk:openai:resolved".to_string()) + } + + async fn remove(&self, model_id: &str) { + self.removed.lock().unwrap().push(model_id.to_string()); + } +} + +fn fake_installer() -> Arc { + Arc::new(FakeTemporaryModelInstaller::default()) +} + #[async_trait] impl AgentSubmissionPort for FakeOwner { async fn create_session( @@ -210,6 +267,10 @@ impl AgentSubmissionPort for FakeOwner { self.release_session_create.notified().await; } let session_id = "session-fixture".to_string(); + self.created_model_ids + .lock() + .unwrap() + .push(request.model_id.clone()); self.created_session_ids .lock() .unwrap() @@ -234,6 +295,10 @@ impl AgentSubmissionPort for FakeOwner { .lock() .unwrap() .push(session_id.clone()); + self.created_model_ids + .lock() + .unwrap() + .push(request.model_id.clone()); if self.panic_after_session_create { panic!("fixture panics after creating the transient Session"); } @@ -561,6 +626,7 @@ async fn host_with_query_limit( max_active_queries, ..SdkHostConfig::default() }, + fake_installer(), ), owner, receiver, @@ -663,6 +729,16 @@ async fn host() -> ( SdkHostConnection, Arc, mpsc::Receiver, +) { + host_with_temporary_model_installer(fake_installer()).await +} + +async fn host_with_temporary_model_installer( + installer: Arc, +) -> ( + SdkHostConnection, + Arc, + mpsc::Receiver, ) { let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); let owner = Arc::new(FakeOwner::with_queue(queue.clone())); @@ -684,6 +760,7 @@ async fn host() -> ( "D:/workspace/project", output, SdkHostConfig::default(), + installer, ), owner, receiver, @@ -698,13 +775,193 @@ async fn initialize(host: &SdkHostConnection, output: &mut mpsc::Receiver JsonRpcRequest { + request(serde_json::json!({ + "jsonrpc": "2.0", + "id": id, + "method": "initialize", + "params": { + "protocolVersion": PROTOCOL_VERSION, + "clientInfo": { "name": "fixture", "version": "0.1.0" }, + "capabilities": { "serverNotifications": true }, + "model": { + "provider": "openai", + "model": "fixture-model", + "apiKey": "fixture-secret" + } + } + })) +} + +#[tokio::test] +async fn temporary_model_is_connection_scoped_and_cannot_be_overridden() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::with_queue(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let installer = Arc::new(FakeTemporaryModelInstaller::default()); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + installer.clone(), + ); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "initialize-model", + "method": "initialize", + "params": { + "protocolVersion": PROTOCOL_VERSION, + "clientInfo": { "name": "fixture", "version": "0.1.0" }, + "capabilities": { "serverNotifications": true }, + "model": { + "provider": "openai", + "model": "fixture-model", + "apiKey": "fixture-secret" + } + } + }))) + .await; + let initialized = output.recv().await.unwrap(); + assert_eq!(initialized["result"]["modelId"], "sdk:openai:resolved"); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "create-with-override", + "method": "session/create", + "params": { "model": "attempted-override" } + }))) + .await; + assert_eq!(output.recv().await.unwrap()["id"], "create-with-override"); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-with-override", + "method": "query/start", + "params": { + "prompt": "hello", + "model": "attempted-query-override" + } + }))) + .await; + assert_eq!(output.recv().await.unwrap()["id"], "query-with-override"); + assert_eq!(installer.installed.lock().unwrap().len(), 1); + assert_eq!( + owner.created_model_ids.lock().unwrap().as_slice(), + &[ + Some("sdk:openai:resolved".to_string()), + Some("sdk:openai:resolved".to_string()), + ] + ); + + host.shutdown_connection().await; + host.shutdown_connection().await; + assert_eq!( + installer.removed.lock().unwrap().as_slice(), + &["sdk:openai:resolved".to_string()] + ); +} + +#[tokio::test] +async fn concurrent_initialize_installs_temporary_model_once() { + let installer = Arc::new(FakeTemporaryModelInstaller::blocking_first_install()); + let (host, _owner, mut output) = host_with_temporary_model_installer(installer.clone()).await; + let first_host = host.clone(); + let first = tokio::spawn(async move { + first_host + .handle_request(temporary_model_initialize_request("initialize-first")) + .await + }); + installer.install_started.notified().await; + + host.handle_request(temporary_model_initialize_request("initialize-second")) + .await; + let second_response = output.recv().await.unwrap(); + installer.release_install.notify_waiters(); + first.await.unwrap(); + let first_response = output.recv().await.unwrap(); + + assert_eq!(second_response["id"], "initialize-second"); + assert_eq!( + second_response["error"]["data"]["code"], + "already_initialized" + ); + assert_eq!(first_response["id"], "initialize-first"); + assert_eq!(first_response["result"]["modelId"], "sdk:openai:resolved"); + assert_eq!(installer.installed.lock().unwrap().len(), 1); + host.shutdown_connection().await; +} + +#[tokio::test] +async fn initialize_finishing_after_shutdown_removes_the_installed_model() { + let installer = Arc::new(FakeTemporaryModelInstaller::blocking_first_install()); + let (host, _owner, mut output) = host_with_temporary_model_installer(installer.clone()).await; + let initialize_host = host.clone(); + let initialize = tokio::spawn(async move { + initialize_host + .handle_request(temporary_model_initialize_request( + "initialize-during-shutdown", + )) + .await + }); + installer.install_started.notified().await; + + host.shutdown_connection().await; + installer.release_install.notify_waiters(); + initialize.await.unwrap(); + let response = output.recv().await.unwrap(); + + assert_eq!(response["id"], "initialize-during-shutdown"); + assert_eq!(response["error"]["data"]["code"], "cancelled"); + assert!(!host.is_initialized().await); + assert_eq!( + installer.removed.lock().unwrap().as_slice(), + &["sdk:openai:resolved".to_string()] + ); +} + +#[tokio::test] +async fn failed_temporary_model_install_rolls_back_for_retry() { + let installer = Arc::new(FakeTemporaryModelInstaller::failing_first_install()); + let (host, _owner, mut output) = host_with_temporary_model_installer(installer.clone()).await; + + host.handle_request(temporary_model_initialize_request("initialize-invalid")) + .await; + let rejected = output.recv().await.unwrap(); + assert_eq!(rejected["error"]["data"]["code"], "invalid_request"); + + host.handle_request(temporary_model_initialize_request("initialize-retry")) + .await; + let initialized = output.recv().await.unwrap(); + assert_eq!(initialized["result"]["modelId"], "sdk:openai:resolved"); + assert_eq!(installer.installed.lock().unwrap().len(), 2); + host.shutdown_connection().await; +} + #[tokio::test] async fn resource_lifecycle_notifications_do_not_create_unaddressable_sessions() { let (host, owner, mut output) = host().await; @@ -756,7 +1013,12 @@ async fn initialize_is_required_and_version_mismatch_fails_closed() { "params": { "protocolVersion": 99, "clientInfo": { "name": "fixture", "version": "0.1.0" }, - "capabilities": { "serverNotifications": true } + "capabilities": { "serverNotifications": true }, + "model": { + "provider": "openai", + "model": "fixture-model", + "apiKey": "fixture-secret" + } } }))) .await; @@ -829,6 +1091,7 @@ async fn escaped_query_output_fails_before_exceeding_the_wire_budget() { "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -912,6 +1175,7 @@ async fn cancellation_timeout_reports_unknown_outcome_for_the_exact_operation() "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -1024,6 +1288,7 @@ async fn dialog_session_identity_mismatch_releases_the_requested_session_reserva "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -1143,6 +1408,7 @@ async fn uncertain_session_close_cleanup_requires_host_restart() { "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; host.handle_request(request(serde_json::json!({ @@ -1237,6 +1503,7 @@ async fn cancellation_remains_available_when_data_request_capacity_is_exhausted( max_in_flight_control_requests: 1, ..SdkHostConfig::default() }, + fake_installer(), ); initialize(&host, &mut output).await; @@ -1350,6 +1617,7 @@ async fn visible_session_create_response_is_exposed_before_shutdown_cleanup() { release_response: release_response.clone(), }), SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -1403,6 +1671,7 @@ async fn shutdown_waits_for_in_flight_session_creation_then_cleans_it() { "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -1457,6 +1726,7 @@ async fn shutdown_compensates_a_session_creation_task_that_panics_after_creation "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -1499,6 +1769,7 @@ async fn shutdown_reports_failure_when_post_panic_session_compensation_fails() { "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -1540,6 +1811,7 @@ async fn a_later_request_registers_panicked_session_cleanup_for_shutdown() { "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -1597,6 +1869,7 @@ async fn shutdown_does_not_forget_cleanup_registered_by_a_later_request() { "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -1699,6 +1972,7 @@ async fn failed_implicit_query_submission_deletes_the_unexposed_session() { "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -1738,6 +2012,7 @@ async fn shutdown_takes_over_failed_query_start_cleanup_within_its_total_budget( "D:/workspace/project", Arc::new(FailQueryStartOutput { output: sender }), SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -1790,6 +2065,7 @@ async fn failed_unexposed_session_cleanup_poison_connection_and_allows_shutdown( "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -1900,6 +2176,7 @@ async fn uncertain_turn_settlement_fails_the_connection_without_a_result() { "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); let connection_failed = host.connection_failed_token(); initialize(&host, &mut output).await; @@ -2024,6 +2301,7 @@ async fn queued_query_is_accepted_and_tracked_by_its_exact_turn() { "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; @@ -2123,6 +2401,7 @@ async fn session_close_rejects_while_query_start_is_in_flight() { "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; host.handle_request(request(serde_json::json!({ @@ -2192,6 +2471,7 @@ async fn query_start_rejects_if_session_close_finishes_before_reservation() { "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; host.handle_request(request(serde_json::json!({ @@ -2262,6 +2542,7 @@ async fn permission_without_callback_is_rejected_and_finishes_action_required() "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; host.handle_request(request(serde_json::json!({ @@ -2370,6 +2651,7 @@ async fn stalled_permission_rejection_is_bounded_and_cancels_the_exact_turn() { "D:/workspace/project", sender, SdkHostConfig::default(), + fake_installer(), ); initialize(&host, &mut output).await; host.handle_request(request(serde_json::json!({ diff --git a/src/crates/interfaces/sdk-host/tests/protocol_contracts.rs b/src/crates/interfaces/sdk-host/tests/protocol_contracts.rs index 93eb25860e..acb686c849 100644 --- a/src/crates/interfaces/sdk-host/tests/protocol_contracts.rs +++ b/src/crates/interfaces/sdk-host/tests/protocol_contracts.rs @@ -2,29 +2,65 @@ use bitfun_sdk_host::protocol::{ ErrorCode, ErrorData, ErrorStage, HostCapabilities, InitializeParams, InitializeResult, JsonRpcErrorResponse, JsonRpcRequest, JsonRpcSuccessResponse, OutcomeCertainty, QueryEvent, QueryOutput, QueryResultError, QueryResultParams, QueryTerminalStatus, RecoveryAction, - RequestId, SessionLifetime, Stability, PROTOCOL_VERSION, + RequestId, SessionLifetime, Stability, TemporaryModelConfig, TemporaryModelProvider, + PROTOCOL_VERSION, }; #[test] -fn initialize_contract_is_versioned_and_uses_familiar_capability_names() { +fn initialize_contract_is_versioned_and_binds_one_temporary_model() { let request: JsonRpcRequest = serde_json::from_value(serde_json::json!({ "jsonrpc": "2.0", "id": 1, "method": "initialize", "params": { - "protocolVersion": 1, + "protocolVersion": 2, "clientInfo": { "name": "fixture", "version": "0.1.0" }, - "capabilities": { "serverNotifications": true } + "capabilities": { "serverNotifications": true }, + "model": { + "provider": "openai", + "model": "fixture-model", + "apiKey": "fixture-secret", + "baseUrl": "http://127.0.0.1:43123/v1" + } } })) .unwrap(); let params: InitializeParams = request.params_as().unwrap(); assert_eq!(request.id, Some(RequestId::Number(1))); + assert_eq!(PROTOCOL_VERSION, 2); assert_eq!(params.protocol_version, PROTOCOL_VERSION); assert!(params.capabilities.server_notifications); + assert_eq!(params.model.provider, TemporaryModelProvider::Openai); + assert_eq!(params.model.model, "fixture-model"); + assert_eq!(params.model.api_key, "fixture-secret"); + assert_eq!( + params.model.base_url.as_deref(), + Some("http://127.0.0.1:43123/v1") + ); + + for provider in ["openai", "responses", "anthropic", "gemini"] { + let model: TemporaryModelConfig = serde_json::from_value(serde_json::json!({ + "provider": provider, + "model": "fixture-model", + "apiKey": "fixture-secret" + })) + .unwrap(); + assert_eq!( + serde_json::to_value(model.provider).unwrap(), + serde_json::json!(provider) + ); + } + assert!( + serde_json::from_value::(serde_json::json!({ + "provider": "unknown", + "model": "fixture-model", + "apiKey": "fixture-secret" + })) + .is_err() + ); - let result = InitializeResult::current("0.2.13"); + let result = InitializeResult::current("0.2.13", "sdk:openai:fixture"); assert_eq!(result.protocol_version, PROTOCOL_VERSION); assert_eq!(result.stability, Stability::NotDelivered); assert_eq!( @@ -45,6 +81,9 @@ fn initialize_contract_is_versioned_and_uses_familiar_capability_names() { prestarted_transport: false, } ); + let result_json = serde_json::to_string(&result).unwrap(); + assert!(result_json.contains("\"modelId\":\"sdk:openai:fixture\"")); + assert!(!result_json.contains("fixture-secret")); } #[test] @@ -192,3 +231,27 @@ fn request_correlation_ids_preserve_json_rpc_id_type() { "request:string:1" ); } + +#[test] +fn json_rpc_request_debug_redacts_temporary_model_secret() { + let request: JsonRpcRequest = serde_json::from_value(serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": 2, + "clientInfo": { "name": "fixture", "version": "0.1.0" }, + "capabilities": { "serverNotifications": true }, + "model": { + "provider": "openai", + "model": "fixture-model", + "apiKey": "bitfun-sdk-debug-secret-31d4" + } + } + })) + .unwrap(); + + let debug = format!("{request:?}"); + assert!(debug.contains("initialize")); + assert!(!debug.contains("bitfun-sdk-debug-secret-31d4")); +}