diff --git a/docs/.generated/plugin-sdk-api-baseline.sha256 b/docs/.generated/plugin-sdk-api-baseline.sha256 index f9b4f3c24fb..0c78286d3d5 100644 --- a/docs/.generated/plugin-sdk-api-baseline.sha256 +++ b/docs/.generated/plugin-sdk-api-baseline.sha256 @@ -145,7 +145,7 @@ d3102bd433eea32ca2a761030788fe18294aa34266f2fb2d01d74268cbf8fb54 module/inbound ee34dd840075bd687624ead228d3aa906fc52965e1708a5a335bb3ad3b671949 module/json-unsafe-integers 8f37bca66178f4d77303fdd3ded0d445c9358b4f094c07a2f8e32b1721e953e4 module/keyed-async-queue d8ca27a737f235e09f1a9f8cd4b82ac0cde96f10c662347b6ac82e27b70b4a40 module/lazy-runtime -15af70a301d18c36aef7dc3cd75350308f7a020738cc8fe1e118cfc90172729d module/llm +11fc218065ff47d3b2c0a2ca36ee66a1effac4e24e2981c9859f17b7b4cb82ae module/llm 6ae9b0f1f55ec50fa3e0ebc02052bfd04eca5f34c4a936692634b36a4cfa2035 module/lmstudio dd166527916dc5b9694512e3526e774da0b2427ba7fc8316d2fd799c06238e80 module/lmstudio-runtime fd00374a91ab5b393152bbbaf90010f9140cf15827bde360d0c6c8724ed34e2e module/logging-core diff --git a/extensions/amazon-bedrock/index.test.ts b/extensions/amazon-bedrock/index.test.ts index d06f2d718aa..8fb2b031ce8 100644 --- a/extensions/amazon-bedrock/index.test.ts +++ b/extensions/amazon-bedrock/index.test.ts @@ -307,6 +307,17 @@ describe("amazon-bedrock provider plugin", () => { vi.resetModules(); }); + it("publishes its stream through the provider lifecycle", async () => { + const provider = await registerSingleProviderPlugin(amazonBedrockPlugin); + + expect( + provider.createStreamFn?.({ model: { api: "bedrock-converse-stream" } } as never), + ).toBeTypeOf("function"); + expect( + provider.createStreamFn?.({ model: { api: "anthropic-messages" } } as never), + ).toBeUndefined(); + }); + it("marks Claude 4.6 Bedrock models as adaptive by default", async () => { const provider = await registerSingleProviderPlugin(amazonBedrockPlugin); diff --git a/extensions/amazon-bedrock/register.sync.runtime.ts b/extensions/amazon-bedrock/register.sync.runtime.ts index 4f9c195246e..ae94aef7a04 100644 --- a/extensions/amazon-bedrock/register.sync.runtime.ts +++ b/extensions/amazon-bedrock/register.sync.runtime.ts @@ -5,7 +5,6 @@ import type { BedrockClient } from "@aws-sdk/client-bedrock"; import type { StreamFn } from "openclaw/plugin-sdk/agent-core"; import type { OpenClawConfig } from "openclaw/plugin-sdk/config-contracts"; -import { registerApiProvider, streamSimple } from "openclaw/plugin-sdk/llm"; import { resolvePluginConfigObject } from "openclaw/plugin-sdk/plugin-config-runtime"; import type { OpenClawPluginApi, @@ -25,7 +24,7 @@ import { supportsBedrockPromptCaching } from "./bedrock-options.js"; import { loadBedrockControlPlaneSdk, runBedrockControlPlaneRequest } from "./control-plane.js"; import { mergeImplicitBedrockProvider, resolveBedrockConfigApiKey } from "./discovery-shared.js"; import { bedrockMemoryEmbeddingProviderAdapter } from "./memory-embedding-adapter.js"; -import { streamBedrock, streamSimpleBedrock } from "./stream.runtime.js"; +import { streamSimpleBedrock } from "./stream.runtime.js"; import { isLatestAdaptiveBedrockModelRef, isOpus47OrNewerBedrockModelRef, @@ -95,8 +94,16 @@ function isAnthropicBedrockModel(modelId: string): boolean { return false; } +const bedrockStreamFn: StreamFn = (model, context, options) => { + if (model.api !== "bedrock-converse-stream") { + throw new Error(`Amazon Bedrock stream received unsupported API: ${model.api}`); + } + // The API check narrows the generic host model to the transport contract. + return streamSimpleBedrock(model as Parameters[0], context, options); +}; + function createBedrockNoCacheWrapper(baseStreamFn: StreamFn | undefined): StreamFn { - const underlying = baseStreamFn ?? streamSimple; + const underlying = baseStreamFn ?? bedrockStreamFn; return (model, context, options) => underlying(model, context, { ...options, @@ -382,15 +389,6 @@ export function registerAmazonBedrockPlugin(api: OpenClawPluginApi): void { }); const startupPluginConfig = (api.pluginConfig ?? {}) as AmazonBedrockPluginConfig; - registerApiProvider( - { - api: "bedrock-converse-stream", - stream: streamBedrock, - streamSimple: streamSimpleBedrock, - }, - `plugin:${providerId}`, - ); - function resolveCurrentPluginConfig( config: OpenClawConfig | undefined, ): AmazonBedrockPluginConfig | undefined { @@ -543,6 +541,8 @@ export function registerAmazonBedrockPlugin(api: OpenClawPluginApi): void { }, resolveConfigApiKey: ({ env }) => resolveBedrockConfigApiKey(env), normalizeResolvedModel: normalizeBedrockResolvedModel, + createStreamFn: ({ model }) => + model.api === "bedrock-converse-stream" ? bedrockStreamFn : undefined, ...anthropicByModelReplayHooks, wrapStreamFn: ({ modelId, config, model, streamFn, thinkingLevel, extraParams }) => { const currentPluginConfig = resolveCurrentPluginConfig(config); diff --git a/extensions/amazon-bedrock/stream.runtime.test.ts b/extensions/amazon-bedrock/stream.runtime.test.ts index 1b1dcbe59db..ca48649d1df 100644 --- a/extensions/amazon-bedrock/stream.runtime.test.ts +++ b/extensions/amazon-bedrock/stream.runtime.test.ts @@ -6,7 +6,8 @@ import { } from "@aws-sdk/client-bedrock-runtime"; import { onLlmRequestActivity } from "openclaw/plugin-sdk/provider-stream-shared"; import { afterEach, describe, expect, it, vi } from "vitest"; -import { streamBedrock, streamSimpleBedrock } from "./stream.runtime.js"; +import type { BedrockOptions } from "./bedrock-options.js"; +import { streamSimpleBedrock } from "./stream.runtime.js"; import { streamTesting as testing } from "./test-support.js"; function bedrockModel(overrides: Record) { @@ -52,9 +53,17 @@ async function* streamEvents(events: unknown[]) { } } +function streamBedrockForTest( + model: Parameters[0], + context: Parameters[1], + options: BedrockOptions = {}, +) { + return streamSimpleBedrock(model, context, options as never); +} + async function captureClientRegion( - model: Parameters[0], - options: Parameters[2] = {}, + model: Parameters[0], + options: BedrockOptions = {}, ): Promise { const send = vi.spyOn(BedrockRuntimeClient.prototype, "send").mockResolvedValue({ $metadata: { httpStatusCode: 200 }, @@ -64,7 +73,7 @@ async function captureClientRegion( ]), } as never); - await streamBedrock( + await streamBedrockForTest( model, { messages: [{ role: "user", content: "Hello", timestamp: 0 }] } as never, options, @@ -275,6 +284,7 @@ describe("Bedrock profile endpoint resolution", () => { ])( "resolves $name to $expectedRegion", async ({ modelId, ambientRegion, fallbackRegion, explicitRegion, expectedRegion }) => { + vi.stubEnv("AWS_PROFILE", ""); vi.stubEnv("AWS_REGION", ambientRegion); if (fallbackRegion !== undefined) { vi.stubEnv("AWS_DEFAULT_REGION", fallbackRegion); @@ -305,7 +315,7 @@ describe("Bedrock stop reasons", () => { ]), } as never); - const result = await streamBedrock(bedrockModel({}), { + const result = await streamBedrockForTest(bedrockModel({}), { messages: [{ role: "user", content: "Hello", timestamp: 0 }], } as never).result(); @@ -579,7 +589,7 @@ describe("Bedrock Fable contract", () => { ]), } as never); - const stream = streamBedrock(fableModel(), context(), { + const stream = streamBedrockForTest(fableModel(), context(), { reasoning: "high", temperature: 0.2, toolChoice: "any", @@ -614,7 +624,7 @@ describe("Bedrock Fable contract", () => { ]), } as never); - const stream = streamBedrock(fableModel(), context(), { + const stream = streamBedrockForTest(fableModel(), context(), { reasoning: "high", toolChoice: "none", }); diff --git a/extensions/amazon-bedrock/stream.runtime.ts b/extensions/amazon-bedrock/stream.runtime.ts index ffea69ebaf4..5a17b19e0eb 100644 --- a/extensions/amazon-bedrock/stream.runtime.ts +++ b/extensions/amazon-bedrock/stream.runtime.ts @@ -126,7 +126,7 @@ function resolveAdaptiveBedrockMaxTokens( } /** Stream a Bedrock Converse request using Bedrock-specific options. */ -export const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOptions> = ( +const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOptions> = ( model: Model<"bedrock-converse-stream">, context: Context, options: BedrockOptions = {}, @@ -400,7 +400,11 @@ function resolveSimpleBedrockOptions( model: Model<"bedrock-converse-stream">, options?: SimpleStreamOptions, ): BedrockOptions { - const base = buildBaseOptions(model, options, undefined); + const bedrockOptions = options as BedrockOptions | undefined; + const base = { + ...bedrockOptions, + ...buildBaseOptions(model, options, undefined), + }; if (requiresMandatoryAdaptiveThinking(model)) { return { ...base, diff --git a/packages/ai/src/internal/default-runtime.test.ts b/packages/ai/src/internal/default-runtime.test.ts new file mode 100644 index 00000000000..4da976ba341 --- /dev/null +++ b/packages/ai/src/internal/default-runtime.test.ts @@ -0,0 +1,47 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { createApiRegistry } from "../api-registry.js"; +import { createLlmRuntime } from "../stream.js"; + +const DEFAULT_RUNTIME_KEY = Symbol.for("openclaw.ai.defaultRuntime"); +const globalStore = globalThis as Record; +const originalDefaultRuntime = globalStore[DEFAULT_RUNTIME_KEY]; + +afterEach(() => { + if (originalDefaultRuntime === undefined) { + delete globalStore[DEFAULT_RUNTIME_KEY]; + } else { + globalStore[DEFAULT_RUNTIME_KEY] = originalDefaultRuntime; + } + vi.resetModules(); +}); + +describe("default LLM runtime compatibility state", () => { + it("keeps opaque legacy registrations out of lifecycle publications", async () => { + const registry = createApiRegistry(); + const stream = () => ({}) as never; + registry.registerApiProvider( + { api: "test-legacy-plugin", stream, streamSimple: stream }, + "plugin:test-legacy", + ); + globalStore[DEFAULT_RUNTIME_KEY] = { + registry, + runtime: createLlmRuntime(registry), + }; + vi.resetModules(); + + const runtime = await import("./default-runtime.js"); + + expect(runtime.defaultApiRegistry).toBe(registry); + expect(runtime.getPublishedApiProviders()).toEqual([]); + + runtime.registerApiProvider( + { api: "test-current-plugin", stream, streamSimple: stream }, + "plugin:test-current", + ); + + expect(runtime.getApiProvider("test-legacy-plugin")).toBeDefined(); + expect(runtime.getPublishedApiProviders().map((provider) => provider.api)).toEqual([ + "test-current-plugin", + ]); + }); +}); diff --git a/packages/ai/src/internal/default-runtime.ts b/packages/ai/src/internal/default-runtime.ts index d6c6a7c3146..1a5295ec70c 100644 --- a/packages/ai/src/internal/default-runtime.ts +++ b/packages/ai/src/internal/default-runtime.ts @@ -1,13 +1,20 @@ +import type { Api, StreamOptions } from "@openclaw/llm-core"; // Process-default registry/runtime retained for the OpenClaw compatibility // facade (src/llm). Deliberately not part of the public package API: external // consumers create isolated runtimes via createLlmRuntime(); exporting these // from the root barrel would reintroduce the mutable process-global registry. -import { createApiRegistry, type ApiRegistry } from "../api-registry.js"; +import { + createApiRegistry, + type ApiProvider, + type ApiRegistry, + type RegisteredApiProvider, +} from "../api-registry.js"; import { createLlmRuntime, type LlmRuntime } from "../stream.js"; type DefaultRuntimeState = { registry: ApiRegistry; runtime: LlmRuntime; + publishedRegistry: ApiRegistry; }; const DEFAULT_RUNTIME_KEY = Symbol.for("openclaw.ai.defaultRuntime"); @@ -15,11 +22,19 @@ const DEFAULT_RUNTIME_KEY = Symbol.for("openclaw.ai.defaultRuntime"); function resolveDefaultRuntime(): DefaultRuntimeState { const globalStore = globalThis as Record; if (Object.hasOwn(globalStore, DEFAULT_RUNTIME_KEY)) { - return globalStore[DEFAULT_RUNTIME_KEY] as DefaultRuntimeState; + const existing = globalStore[DEFAULT_RUNTIME_KEY] as Omit< + DefaultRuntimeState, + "publishedRegistry" + > & + Partial>; + // Keep the legacy facade registry intact. Its entries are opaque, so only + // registrations made through this host version enter lifecycle snapshots. + existing.publishedRegistry ??= createApiRegistry(); + return existing as DefaultRuntimeState; } const registry = createApiRegistry(); const runtime = createLlmRuntime(registry); - const state = { registry, runtime }; + const state = { registry, runtime, publishedRegistry: createApiRegistry() }; globalStore[DEFAULT_RUNTIME_KEY] = state; return state; } @@ -29,12 +44,29 @@ const defaultRuntime = resolveDefaultRuntime(); export const defaultApiRegistry = defaultRuntime.registry; export const defaultLlmRuntime = defaultRuntime.runtime; -export const { - registerApiProvider, - getApiProvider, - getApiProviders, - unregisterApiProviders, - clearApiProviders, -} = defaultApiRegistry; +export function registerApiProvider( + provider: ApiProvider, + sourceId?: string, +): void { + defaultApiRegistry.registerApiProvider(provider, sourceId); + defaultRuntime.publishedRegistry.registerApiProvider(provider, sourceId); +} + +export const { getApiProvider, getApiProviders } = defaultApiRegistry; + +/** Returns only explicit compatibility registrations, excluding request-generated aliases. */ +export function getPublishedApiProviders(): RegisteredApiProvider[] { + return defaultRuntime.publishedRegistry.getApiProviders(); +} + +export function unregisterApiProviders(sourceId: string): void { + defaultApiRegistry.unregisterApiProviders(sourceId); + defaultRuntime.publishedRegistry.unregisterApiProviders(sourceId); +} + +export function clearApiProviders(): void { + defaultApiRegistry.clearApiProviders(); + defaultRuntime.publishedRegistry.clearApiProviders(); +} export const { stream, complete, streamSimple, completeSimple } = defaultLlmRuntime; diff --git a/src/agents/btw.test.ts b/src/agents/btw.test.ts index 21ad058eaf1..f1ebbe627b4 100644 --- a/src/agents/btw.test.ts +++ b/src/agents/btw.test.ts @@ -20,6 +20,7 @@ const buildSessionContextMock = vi.fn(); const ensureOpenClawModelsJsonMock = vi.fn(); const discoverAuthStorageMock = vi.fn(); const discoverModelsMock = vi.fn(); +const getModelRegistryRuntimeMock = vi.fn(); const resolveModelWithRegistryMock = vi.fn(); const ensureAuthProfileStoreMock = vi.fn(); const ensureAuthProfileStoreWithoutExternalProfilesMock = vi.fn(); @@ -82,6 +83,10 @@ vi.mock("./agent-model-discovery.js", () => ({ discoverModels: (...args: unknown[]) => discoverModelsMock(...args), })); +vi.mock("./sessions/model-registry-runtime.js", () => ({ + getModelRegistryRuntime: (...args: unknown[]) => getModelRegistryRuntimeMock(...args), +})); + vi.mock("./embedded-agent-runner/model.js", () => ({ resolveModelAsync: (...args: unknown[]) => resolveModelAsyncMock(...args), resolveModelWithRegistry: (...args: unknown[]) => resolveModelWithRegistryMock(...args), @@ -532,6 +537,11 @@ describe("runBtwSideQuestion", () => { ensureOpenClawModelsJsonMock.mockReset(); discoverAuthStorageMock.mockReset(); discoverModelsMock.mockReset(); + getModelRegistryRuntimeMock.mockReset(); + getModelRegistryRuntimeMock.mockReturnValue({ + apiRegistry: {}, + llmRuntime: { streamSimple: streamSimpleMock }, + }); resolveModelAsyncMock.mockReset(); resolveModelWithRegistryMock.mockReset(); ensureAuthProfileStoreMock.mockReset(); diff --git a/src/agents/btw.ts b/src/agents/btw.ts index 5188e13e651..8fd9646997c 100644 --- a/src/agents/btw.ts +++ b/src/agents/btw.ts @@ -11,7 +11,6 @@ import type { ChatType } from "../channels/chat-type.js"; import type { SessionEntry as StoredSessionEntry } from "../config/sessions.js"; import type { OpenClawConfig } from "../config/types.openclaw.js"; import { streamWithPayloadPatch } from "../llm/providers/stream-wrappers/stream-payload-utils.js"; -import { streamSimple } from "../llm/stream.js"; import type { AssistantMessageEvent, ImageContent, @@ -81,6 +80,7 @@ import { import type { AgentRuntimeAuthPlan } from "./runtime-plan/types.js"; import { resolveSessionRuntimeOverrideForProvider } from "./session-runtime-compat.js"; import { stripToolResultDetails } from "./session-transcript-repair.js"; +import { getModelRegistryRuntime } from "./sessions/model-registry-runtime.js"; import { resolveAgentTimeoutMs } from "./timeout.js"; import { sanitizeImageBlocks } from "./tool-images.js"; @@ -1145,6 +1145,7 @@ export async function runBtwSideQuestion( } } runtimeModel = applySecretRefHeaderSentinels(runtimeModel, params.cfg); + const modelRegistryRuntime = getModelRegistryRuntime(modelRegistry); // Use the provider's own stream fn so providers like Ollama (which build // `/api/chat` or `/v1/chat/completions` paths based on api mode) construct @@ -1156,9 +1157,11 @@ export async function runBtwSideQuestion( agentDir: params.agentDir, workspaceDir, env: process.env, + apiRegistry: modelRegistryRuntime.apiRegistry, }); const streamFn = resolveEmbeddedAgentStreamFn({ - currentStreamFn: streamSimple, + llmRuntime: modelRegistryRuntime.llmRuntime, + currentStreamFn: modelRegistryRuntime.llmRuntime.streamSimple, providerStreamFn, sessionId, signal: params.opts?.abortSignal, diff --git a/src/agents/custom-api-registry.test.ts b/src/agents/custom-api-registry.test.ts index da4db95e090..bd0b43f4e8d 100644 --- a/src/agents/custom-api-registry.test.ts +++ b/src/agents/custom-api-registry.test.ts @@ -1,19 +1,15 @@ -import { - clearApiProviders, - defaultApiRegistry, - getApiProvider, - registerApiProvider, - unregisterApiProviders, -} from "@openclaw/ai/internal/runtime"; -import { registerBuiltInApiProviders, resetApiProviders } from "@openclaw/ai/providers"; +import { createApiRegistry, type ApiRegistry } from "@openclaw/ai"; +import { resetApiProviders } from "@openclaw/ai/providers"; // Covers dynamic registration of custom model API providers. -import { afterEach, describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import { createAssistantMessageEventStream } from "../llm/utils/event-stream.js"; import { ensureCustomApiRegistered } from "./custom-api-registry.js"; import { buildAssistantMessage, buildUsageWithNoCost } from "./stream-message-shared.js"; +let registry: ApiRegistry; + function getRegisteredTestProvider() { - const provider = getApiProvider("test-custom-api"); + const provider = registry.getApiProvider("test-custom-api"); if (!provider) { throw new Error("expected test-custom-api provider to be registered"); } @@ -21,9 +17,8 @@ function getRegisteredTestProvider() { } describe("ensureCustomApiRegistered", () => { - afterEach(() => { - clearApiProviders(); - registerBuiltInApiProviders(defaultApiRegistry); + beforeEach(() => { + registry = createApiRegistry(); }); it("registers a custom api provider once", () => { @@ -31,8 +26,8 @@ describe("ensureCustomApiRegistered", () => { // replace provider entries or create duplicate sources. const streamFn = vi.fn(() => createAssistantMessageEventStream()); - expect(ensureCustomApiRegistered("test-custom-api", streamFn)).toBe(true); - expect(ensureCustomApiRegistered("test-custom-api", streamFn)).toBe(false); + expect(ensureCustomApiRegistered(registry, "test-custom-api", streamFn)).toBe(true); + expect(ensureCustomApiRegistered(registry, "test-custom-api", streamFn)).toBe(false); const provider = getRegisteredTestProvider(); expect(typeof provider.stream).toBe("function"); @@ -42,7 +37,7 @@ describe("ensureCustomApiRegistered", () => { it("delegates both stream entrypoints to the provided stream function", () => { const stream = createAssistantMessageEventStream(); const streamFn = vi.fn(() => stream); - ensureCustomApiRegistered("test-custom-api", streamFn); + ensureCustomApiRegistered(registry, "test-custom-api", streamFn); const provider = getRegisteredTestProvider(); @@ -68,7 +63,7 @@ describe("ensureCustomApiRegistered", () => { stream.push({ type: "done", reason: "stop", message }); return stream; }); - ensureCustomApiRegistered("test-custom-api", streamFn); + ensureCustomApiRegistered(registry, "test-custom-api", streamFn); const provider = getRegisteredTestProvider(); const stream = provider.stream( @@ -85,7 +80,7 @@ describe("ensureCustomApiRegistered", () => { const streamFn = vi.fn(async () => { throw new Error("factory failed"); }); - ensureCustomApiRegistered("test-custom-api", streamFn); + ensureCustomApiRegistered(registry, "test-custom-api", streamFn); const provider = getRegisteredTestProvider(); const stream = provider.stream( @@ -107,7 +102,7 @@ describe("ensureCustomApiRegistered", () => { const api = "test-reset-plugin-api"; const streamFn = vi.fn(() => createAssistantMessageEventStream()); const streamSimpleFn = vi.fn(() => createAssistantMessageEventStream()); - registerApiProvider( + registry.registerApiProvider( { api, stream: streamFn, @@ -116,11 +111,11 @@ describe("ensureCustomApiRegistered", () => { sourceId, ); - resetApiProviders(defaultApiRegistry); + resetApiProviders(registry); - expect(getApiProvider(api)).toBeDefined(); - expect(getApiProvider("openai-responses")).toBeDefined(); + expect(registry.getApiProvider(api)).toBeDefined(); + expect(registry.getApiProvider("openai-responses")).toBeDefined(); - unregisterApiProviders(sourceId); + registry.unregisterApiProviders(sourceId); }); }); diff --git a/src/agents/custom-api-registry.ts b/src/agents/custom-api-registry.ts index 97b06fecb76..e0ad35cb65e 100644 --- a/src/agents/custom-api-registry.ts +++ b/src/agents/custom-api-registry.ts @@ -1,7 +1,7 @@ /** * Registers caller-supplied custom API stream functions with the LLM registry. */ -import { getApiProvider, registerApiProvider } from "@openclaw/ai/internal/runtime"; +import type { ApiRegistry } from "@openclaw/ai"; import type { Api, AssistantMessageEventStreamContract, @@ -47,12 +47,16 @@ function adaptCustomStream( } /** Registers a custom API stream function when no provider already owns it. */ -export function ensureCustomApiRegistered(api: Api, streamFn: StreamFn): boolean { - if (getApiProvider(api)) { +export function ensureCustomApiRegistered( + registry: ApiRegistry, + api: Api, + streamFn: StreamFn, +): boolean { + if (registry.getApiProvider(api)) { return false; } - registerApiProvider( + registry.registerApiProvider( { api, stream: (model, context, options) => diff --git a/src/agents/embedded-agent-runner/compact.hooks.harness.ts b/src/agents/embedded-agent-runner/compact.hooks.harness.ts index c21980df69b..bf0d8e9fea7 100644 --- a/src/agents/embedded-agent-runner/compact.hooks.harness.ts +++ b/src/agents/embedded-agent-runner/compact.hooks.harness.ts @@ -208,6 +208,10 @@ export const buildEmbeddedSystemPromptMock = vi.fn(() => ""); export const resolveEmbeddedAgentStreamFnMock: Mock< (params?: unknown) => MockEmbeddedAgentStreamFn > = vi.fn((_params?: unknown) => vi.fn()); +const getModelRegistryRuntimeMock = vi.fn(() => ({ + apiRegistry: {}, + llmRuntime: { streamSimple: vi.fn() }, +})); export const getApiKeyForModelMock: Mock< (params?: { profileId?: string; allowAuthProfileFallback?: boolean }) => Promise<{ apiKey: string; @@ -457,6 +461,11 @@ export function resetCompactSessionStateMocks(): void { })); resolveEmbeddedAgentStreamFnMock.mockReset(); resolveEmbeddedAgentStreamFnMock.mockImplementation((_params?: unknown) => vi.fn()); + getModelRegistryRuntimeMock.mockReset(); + getModelRegistryRuntimeMock.mockReturnValue({ + apiRegistry: {}, + llmRuntime: { streamSimple: vi.fn() }, + }); getApiKeyForModelMock.mockReset(); getApiKeyForModelMock.mockImplementation(async (params?: { profileId?: string }) => ({ apiKey: "test", @@ -677,6 +686,10 @@ export async function loadCompactHooksHarness(): Promise<{ registerProviderStreamForModel: registerProviderStreamForModelMock, })); + vi.doMock("../sessions/model-registry-runtime.js", () => ({ + getModelRegistryRuntime: getModelRegistryRuntimeMock, + })); + vi.doMock("../../hooks/internal-hooks.js", async () => { const actual = await vi.importActual( "../../hooks/internal-hooks.js", diff --git a/src/agents/embedded-agent-runner/compact.hooks.test.ts b/src/agents/embedded-agent-runner/compact.hooks.test.ts index 1c622998c10..033b060ac8f 100644 --- a/src/agents/embedded-agent-runner/compact.hooks.test.ts +++ b/src/agents/embedded-agent-runner/compact.hooks.test.ts @@ -856,6 +856,7 @@ describe("compactEmbeddedAgentSessionDirect hooks", () => { await compactTesting.prepareCompactionSessionAgent({ session: session as never, + llmRuntime: { streamSimple: vi.fn() } as never, providerStreamFn: vi.fn(), sessionId: "session-1", signal: new AbortController().signal, @@ -921,6 +922,7 @@ describe("compactEmbeddedAgentSessionDirect hooks", () => { agent: { streamFn: vi.fn() }, messages: [{ role: "user", content: "hello" }], } as never, + llmRuntime: { streamSimple: vi.fn() } as never, providerStreamFn: vi.fn(), sessionId: "session-1", signal: new AbortController().signal, @@ -2297,6 +2299,7 @@ describe("compactEmbeddedAgentSessionDirect hooks", () => { config: undefined, agentDir: "/tmp", effectiveWorkspace: "/tmp", + apiRegistry: {} as never, }); expect(result).toBe(streamFn); diff --git a/src/agents/embedded-agent-runner/compact.ts b/src/agents/embedded-agent-runner/compact.ts index 5d261454fe3..7c164fd75ba 100644 --- a/src/agents/embedded-agent-runner/compact.ts +++ b/src/agents/embedded-agent-runner/compact.ts @@ -3,6 +3,7 @@ */ import fs from "node:fs/promises"; import os from "node:os"; +import type { ApiRegistry } from "@openclaw/ai"; import { isAcpRuntimeSpawnAvailable } from "../../acp/runtime/availability.js"; import type { ThinkLevel } from "../../auto-reply/thinking.js"; import { resolveAgentModelFallbackValues } from "../../config/model-input.js"; @@ -158,6 +159,7 @@ import { resolveSessionWriteLockOptions, } from "../session-write-lock.js"; import { createAgentSession, estimateTokens, SessionManager } from "../sessions/index.js"; +import { getModelRegistryRuntime } from "../sessions/model-registry-runtime.js"; import { detectRuntimeShell } from "../shell-utils.js"; import { filterProviderNormalizableTools, @@ -257,12 +259,14 @@ function resolveCompactionProviderStream(params: { config?: OpenClawConfig; agentDir: string; effectiveWorkspace: string; + apiRegistry: ApiRegistry; }) { return registerProviderStreamForModel({ model: params.effectiveModel, cfg: params.config, agentDir: params.agentDir, workspaceDir: params.effectiveWorkspace, + apiRegistry: params.apiRegistry, }); } @@ -1499,6 +1503,7 @@ async function compactEmbeddedAgentSessionDirectOnce( config: params.config, agentDir, effectiveWorkspace, + apiRegistry: getModelRegistryRuntime(modelRegistry).apiRegistry, }); while (true) { // Rebuild the compaction session on retry so provider wrappers, payload @@ -1527,6 +1532,7 @@ async function compactEmbeddedAgentSessionDirectOnce( // through the same transport/payload shaping stack as normal turns. await prepareCompactionSessionAgent({ session, + llmRuntime: getModelRegistryRuntime(modelRegistry).llmRuntime, providerStreamFn, sessionId: params.sessionId, signal: runAbortController.signal, diff --git a/src/agents/embedded-agent-runner/compaction-session-agent.ts b/src/agents/embedded-agent-runner/compaction-session-agent.ts index 4a95f91c1d1..51d7cc040ac 100644 --- a/src/agents/embedded-agent-runner/compaction-session-agent.ts +++ b/src/agents/embedded-agent-runner/compaction-session-agent.ts @@ -1,3 +1,4 @@ +import type { LlmRuntime } from "@openclaw/ai"; import type { ThinkLevel } from "../../auto-reply/thinking.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; import type { ProviderRuntimeModel } from "../../plugins/provider-runtime-model.types.js"; @@ -14,6 +15,7 @@ import { mapThinkingLevelForProvider } from "./utils.js"; export async function prepareCompactionSessionAgent(params: { session: { agent: { streamFn?: unknown } }; + llmRuntime: LlmRuntime; providerStreamFn: unknown; sessionId: string; signal: AbortSignal; @@ -58,6 +60,7 @@ export async function prepareCompactionSessionAgent(params: { }) : params.resolvedApiKey; params.session.agent.streamFn = resolveEmbeddedAgentStreamFn({ + llmRuntime: params.llmRuntime, currentStreamFn: resolveEmbeddedAgentBaseStreamFn({ session: params.session as never }), providerStreamFn: params.providerStreamFn as never, sessionId: params.sessionId, diff --git a/src/agents/embedded-agent-runner/extra-params.cache-retention-default.test.ts b/src/agents/embedded-agent-runner/extra-params.cache-retention-default.test.ts index 1f4e3946a22..7a411f7c741 100644 --- a/src/agents/embedded-agent-runner/extra-params.cache-retention-default.test.ts +++ b/src/agents/embedded-agent-runner/extra-params.cache-retention-default.test.ts @@ -1,7 +1,6 @@ // Coverage for cache-retention defaults and overrides in extra params. import type { StreamFn } from "openclaw/plugin-sdk/agent-core"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { createLlmStreamSimpleMock } from "../../../test/helpers/agents/llm-stream-simple-mock.js"; import { applyExtraParamsToAgent } from "./extra-params.js"; import { testing as extraParamsTesting } from "./extra-params.test-support.js"; import { log } from "./logger.js"; @@ -16,7 +15,7 @@ function applyAndExpectWrapped(params: { }) { // Wrapping is the observable signal that cache-retention handling was enabled // without requiring a real provider stream call. - const agent: { streamFn?: StreamFn } = {}; + const agent: { streamFn?: StreamFn } = { streamFn: vi.fn() as StreamFn }; applyExtraParamsToAgent( agent, @@ -43,8 +42,6 @@ vi.mock("./logger.js", () => ({ }, })); -vi.mock("../../llm/stream.js", () => createLlmStreamSimpleMock()); - beforeEach(() => { vi.mocked(log.warn).mockClear(); extraParamsTesting.setProviderRuntimeDepsForTest({ diff --git a/src/agents/embedded-agent-runner/extra-params.ts b/src/agents/embedded-agent-runner/extra-params.ts index f72b8f459e1..c469ad767f8 100644 --- a/src/agents/embedded-agent-runner/extra-params.ts +++ b/src/agents/embedded-agent-runner/extra-params.ts @@ -20,7 +20,6 @@ import { } from "../../llm/providers/stream-wrappers/openai.js"; import { createOpenRouterSystemCacheWrapper } from "../../llm/providers/stream-wrappers/proxy.js"; import { streamWithPayloadPatch } from "../../llm/providers/stream-wrappers/stream-payload-utils.js"; -import { streamSimple } from "../../llm/stream.js"; import type { SimpleStreamOptions } from "../../llm/types.js"; import { createDeepSeekV4OpenAICompatibleThinkingWrapper, @@ -43,6 +42,13 @@ import type { StreamFn } from "../runtime/index.js"; import type { SettingsManager } from "../sessions/index.js"; import { log } from "./logger.js"; import { parseCacheRetention, resolveCacheRetention } from "./prompt-cache-retention.js"; + +function requireBaseStreamFn(streamFn: StreamFn | undefined): StreamFn { + if (!streamFn) { + throw new Error("Cannot apply stream policy without a lifecycle-owned base stream."); + } + return streamFn; +} import type { ProviderThinkLevel } from "./utils.js"; const defaultProviderRuntimeDeps = { @@ -558,7 +564,7 @@ function createStreamFnWithExtraParams( log.debug(`creating streamFn wrapper with params: ${JSON.stringify(debugParams)}`); } - const underlying = baseStreamFn ?? streamSimple; + const underlying = requireBaseStreamFn(baseStreamFn); const wrappedStreamFn: StreamFn = (callModel, context, options) => { const cacheRetention = resolveCacheRetention( extraParams, @@ -661,7 +667,7 @@ function createParallelToolCallsWrapper( baseStreamFn: StreamFn | undefined, enabled: boolean, ): StreamFn { - const underlying = baseStreamFn ?? streamSimple; + const underlying = requireBaseStreamFn(baseStreamFn); return (model, context, options) => { if (!supportsGptParallelToolCallsPayload(model.api)) { return underlying(model, context, options); @@ -695,7 +701,7 @@ function shouldStripOpenAICompletionsStore(model: ProviderRuntimeModel): boolean } function createOpenAICompletionsStoreCompatWrapper(baseStreamFn: StreamFn | undefined): StreamFn { - const underlying = baseStreamFn ?? streamSimple; + const underlying = requireBaseStreamFn(baseStreamFn); return (model, context, options) => { if (!shouldStripOpenAICompletionsStore(model as ProviderRuntimeModel)) { return underlying(model, context, options); @@ -751,7 +757,7 @@ function createOpenAICompletionsChatTemplateKwargsWrapper(params: { baseStreamFn: StreamFn | undefined; configured: Record; }): StreamFn { - const underlying = params.baseStreamFn ?? streamSimple; + const underlying = requireBaseStreamFn(params.baseStreamFn); return (model, context, options) => { if (model.api !== "openai-completions") { return underlying(model, context, options); @@ -774,7 +780,7 @@ function createOpenAICompletionsExtraBodyWrapper( baseStreamFn: StreamFn | undefined, extraBody: Record, ): StreamFn { - const underlying = baseStreamFn ?? streamSimple; + const underlying = requireBaseStreamFn(baseStreamFn); return (model, context, options) => { if (model.api !== "openai-completions") { return underlying(model, context, options); diff --git a/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.test-support.ts b/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.test-support.ts index f5e8921f197..74fc63f3524 100644 --- a/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.test-support.ts +++ b/src/agents/embedded-agent-runner/run/attempt.spawn-workspace.test-support.ts @@ -18,6 +18,7 @@ import type { IngestResult, } from "../../../context-engine/types.js"; import { formatErrorMessage } from "../../../infra/errors.js"; +import { bindStreamLlmRuntime } from "../../../llm/model-runtime-binding.js"; import type { Model } from "../../../llm/types.js"; import type { PluginMetadataSnapshot } from "../../../plugins/plugin-metadata-snapshot.js"; import { createLazyPromise } from "../../../shared/lazy-runtime.js"; @@ -27,6 +28,10 @@ import type { MessagingToolSourceReplyPayload, } from "../../embedded-agent-messaging.types.js"; import type { AgentMessage } from "../../runtime/index.js"; +import { + getModelRegistryRuntime, + initializeModelRegistryRuntime, +} from "../../sessions/model-registry-runtime.js"; import type { WorkspaceBootstrapFile } from "../../workspace.js"; type SubscribeEmbeddedAgentSessionFn = @@ -1342,14 +1347,21 @@ export async function createContextEngineAttemptRunner(params: { .mockReset() .mockReturnValue({ messages: params.sessionMessagesAfterRepair ?? seedMessages }); - hoisted.createAgentSessionMock.mockImplementation(async () => ({ - session: + const modelRegistry = {}; + initializeModelRegistryRuntime(modelRegistry); + const modelRuntime = getModelRegistryRuntime(modelRegistry).llmRuntime; + hoisted.createAgentSessionMock.mockImplementation(async () => { + const session = params.createSession?.() ?? createDefaultEmbeddedSession({ initialMessages: seedMessages, prompt: params.sessionPrompt, - }), - })); + }); + if (session.agent.streamFn) { + bindStreamLlmRuntime(session.agent.streamFn, modelRuntime); + } + return { session }; + }); const previousTrajectoryEnv = process.env.OPENCLAW_TRAJECTORY; const previousTrajectoryDirEnv = process.env.OPENCLAW_TRAJECTORY_DIR; @@ -1378,7 +1390,7 @@ export async function createContextEngineAttemptRunner(params: { model: testModel, authStorage: testAuthStorage as never, authProfileStore: { version: 1, profiles: {} }, - modelRegistry: {} as never, + modelRegistry: modelRegistry as never, thinkLevel: "off", disableTools: true, disableMessageTool: true, diff --git a/src/agents/embedded-agent-runner/run/attempt.test.ts b/src/agents/embedded-agent-runner/run/attempt.test.ts index fc20b8cc0a3..cb4406634fc 100644 --- a/src/agents/embedded-agent-runner/run/attempt.test.ts +++ b/src/agents/embedded-agent-runner/run/attempt.test.ts @@ -5,6 +5,8 @@ import { streamSimple } from "../../../llm/stream.js"; vi.mock("../context-engine-capabilities.js", () => ({ resolveContextEngineCapabilities: async () => ({ llm: undefined }), })); +import type { LlmRuntime } from "@openclaw/ai"; +import { defaultLlmRuntime } from "@openclaw/ai/internal/runtime"; import { SYSTEM_PROMPT_CACHE_BOUNDARY } from "@openclaw/ai/internal/shared"; import type { OpenClawConfig } from "../../../config/config.js"; import { addSession } from "../../bash-process-registry.js"; @@ -15,7 +17,7 @@ import { buildAgentSystemPrompt } from "../../system-prompt.js"; import type { NormalizedUsage } from "../../usage.js"; import { resolveEmbeddedAgentBaseStreamFn, - resolveEmbeddedAgentStreamFn, + resolveEmbeddedAgentStreamFn as resolveEmbeddedAgentStreamFnImpl, } from "../stream-resolution.js"; import { buildContextEnginePromptCacheInfo } from "./attempt.context-engine-helpers.js"; import { @@ -36,6 +38,17 @@ import { } from "./attempt.tool-call-normalization.js"; import { buildEmbeddedAttemptToolRunContext } from "./attempt.tool-run-context.js"; +const llmRuntime = { + ...defaultLlmRuntime, + streamSimple, +} as LlmRuntime; + +function resolveEmbeddedAgentStreamFn( + params: Omit[0], "llmRuntime">, +) { + return resolveEmbeddedAgentStreamFnImpl({ ...params, llmRuntime }); +} + type FakeWrappedStream = { result: () => Promise; [Symbol.asyncIterator]: () => AsyncIterator; diff --git a/src/agents/embedded-agent-runner/stream-resolution.test.ts b/src/agents/embedded-agent-runner/stream-resolution.test.ts index cd6c71a7e32..2eff34f7f90 100644 --- a/src/agents/embedded-agent-runner/stream-resolution.test.ts +++ b/src/agents/embedded-agent-runner/stream-resolution.test.ts @@ -1,16 +1,18 @@ -import { getApiProvider } from "@openclaw/ai/internal/runtime"; +import type { LlmRuntime } from "@openclaw/ai"; +import { defaultLlmRuntime, getApiProvider } from "@openclaw/ai/internal/runtime"; import { SYSTEM_PROMPT_CACHE_BOUNDARY } from "@openclaw/ai/internal/shared"; // Stream resolution tests cover how embedded runs choose provider, boundary, // native Codex, or custom stream functions and pass auth/cache/signal options. import type { StreamFn } from "openclaw/plugin-sdk/agent-core"; import { afterEach, describe, expect, it, vi } from "vitest"; +import { bindStreamLlmRuntime } from "../../llm/model-runtime-binding.js"; import { streamSimple } from "../../llm/stream.js"; import { mintSecretSentinel } from "../../secrets/sentinel.js"; import * as providerTransportStream from "../provider-transport-stream.js"; import { - describeEmbeddedAgentStreamStrategy, + describeEmbeddedAgentStreamStrategy as describeEmbeddedAgentStreamStrategyImpl, resolveEmbeddedAgentApiKey, - resolveEmbeddedAgentStreamFn, + resolveEmbeddedAgentStreamFn as resolveEmbeddedAgentStreamFnImpl, } from "./stream-resolution.js"; const streamMocks = vi.hoisted(() => ({ @@ -38,6 +40,23 @@ vi.mock("../provider-transport-stream.js", async (importOriginal) => { }; }); +const llmRuntime = { + ...defaultLlmRuntime, + streamSimple: streamSimple as StreamFn, +} as LlmRuntime; + +function describeEmbeddedAgentStreamStrategy( + params: Omit[0], "llmRuntime">, +) { + return describeEmbeddedAgentStreamStrategyImpl({ ...params, llmRuntime }); +} + +function resolveEmbeddedAgentStreamFn( + params: Omit[0], "llmRuntime">, +) { + return resolveEmbeddedAgentStreamFnImpl({ ...params, llmRuntime }); +} + const overrideBoundaryAwareStreamFnOnce = (streamFn: StreamFn): void => { // Boundary wrapping remains real by default; individual cases replace only // the inner stream when they need to inspect forwarded options. @@ -75,6 +94,21 @@ afterEach(() => { }); describe("describeEmbeddedAgentStreamStrategy", () => { + it("recovers the lifecycle owner from a prepared session stream", () => { + bindStreamLlmRuntime(streamSimple, llmRuntime); + + expect( + describeEmbeddedAgentStreamStrategyImpl({ + currentStreamFn: streamSimple, + model: { + api: "openai-responses", + provider: "openai", + id: "gpt-5.4", + } as never, + }), + ).toBe("boundary-aware:openai-responses"); + }); + it("describes provider-owned stream paths explicitly", () => { expect( describeEmbeddedAgentStreamStrategy({ diff --git a/src/agents/embedded-agent-runner/stream-resolution.ts b/src/agents/embedded-agent-runner/stream-resolution.ts index e1f31d8a784..73560404968 100644 --- a/src/agents/embedded-agent-runner/stream-resolution.ts +++ b/src/agents/embedded-agent-runner/stream-resolution.ts @@ -1,9 +1,9 @@ /** * Resolves provider stream functions and API keys for embedded agents. */ -import { getApiProvider } from "@openclaw/ai/internal/runtime"; +import type { LlmRuntime } from "@openclaw/ai"; import { stripSystemPromptCacheBoundary } from "@openclaw/ai/internal/shared"; -import { streamSimple } from "../../llm/stream.js"; +import { getStreamLlmRuntime } from "../../llm/model-runtime-binding.js"; import { createAnthropicVertexStreamFnForModel } from "../anthropic-vertex-stream.js"; import { createBoundaryAwareStreamFnForModel } from "../provider-transport-stream.js"; import type { StreamFn } from "../runtime/index.js"; @@ -18,28 +18,53 @@ type EmbeddedStreamOptions = Parameters[2] & { export function resolveEmbeddedAgentBaseStreamFn(params: { session: { agent: { streamFn?: StreamFn } }; -}): StreamFn | undefined { +}): StreamFn { const cached = embeddedAgentBaseStreamFnCache.get(params.session); if (cached !== undefined || embeddedAgentBaseStreamFnCache.has(params.session)) { + if (!cached) { + throw new Error("Agent session has no lifecycle-owned base stream."); + } return cached; } const baseStreamFn = params.session.agent.streamFn; embeddedAgentBaseStreamFnCache.set(params.session, baseStreamFn); + if (!baseStreamFn) { + throw new Error("Agent session has no lifecycle-owned base stream."); + } return baseStreamFn; } +type EmbeddedStreamRuntimeOwner = + | { + llmRuntime: LlmRuntime; + currentStreamFn: StreamFn | undefined; + } + | { + llmRuntime?: never; + currentStreamFn: StreamFn; + }; + +function resolveEmbeddedStreamRuntime(owner: EmbeddedStreamRuntimeOwner): LlmRuntime { + const runtime = owner.llmRuntime ?? getStreamLlmRuntime(owner.currentStreamFn); + if (!runtime) { + throw new Error("Embedded stream has no lifecycle runtime owner."); + } + return runtime; +} + function isDefaultOpenClawStreamFnForModel( model: EmbeddedRunAttemptParams["model"], streamFn: StreamFn | undefined, + llmRuntime: LlmRuntime, ): boolean { - if (!streamFn || streamFn === streamSimple) { + if (!streamFn || streamFn === llmRuntime.streamSimple) { return true; } const api = typeof model.api === "string" ? model.api.trim() : ""; if (!api) { return false; } - const provider = getApiProvider(api as never); + const provider = llmRuntime.registry.getApiProvider(api as never); return streamFn === provider?.streamSimple || streamFn === provider?.stream; } @@ -54,22 +79,25 @@ function isOpenAICodexResponsesModel(model: EmbeddedRunAttemptParams["model"]): function resolveOpenClawNativeCodexResponsesStreamFn(params: { model: EmbeddedRunAttemptParams["model"]; currentStreamFn: StreamFn | undefined; + llmRuntime: LlmRuntime; }): StreamFn | undefined { if (!isOpenAICodexResponsesModel(params.model)) { return undefined; } - if (!isDefaultOpenClawStreamFnForModel(params.model, params.currentStreamFn)) { + if (!isDefaultOpenClawStreamFnForModel(params.model, params.currentStreamFn, params.llmRuntime)) { return undefined; } - return params.currentStreamFn ?? streamSimple; + return params.currentStreamFn ?? params.llmRuntime.streamSimple; } -export function describeEmbeddedAgentStreamStrategy(params: { - currentStreamFn: StreamFn | undefined; - providerStreamFn?: StreamFn; - model: EmbeddedRunAttemptParams["model"]; - resolvedApiKey?: string; -}): string { +export function describeEmbeddedAgentStreamStrategy( + params: EmbeddedStreamRuntimeOwner & { + providerStreamFn?: StreamFn; + model: EmbeddedRunAttemptParams["model"]; + resolvedApiKey?: string; + }, +): string { + const llmRuntime = resolveEmbeddedStreamRuntime(params); if (params.providerStreamFn) { return "provider"; } @@ -80,11 +108,12 @@ export function describeEmbeddedAgentStreamStrategy(params: { resolveOpenClawNativeCodexResponsesStreamFn({ model: params.model, currentStreamFn: params.currentStreamFn, + llmRuntime, }) ) { return "openclaw-native-codex-responses"; } - if (isDefaultOpenClawStreamFnForModel(params.model, params.currentStreamFn)) { + if (isDefaultOpenClawStreamFnForModel(params.model, params.currentStreamFn, llmRuntime)) { return createBoundaryAwareStreamFnForModel(params.model) ? `boundary-aware:${params.model.api}` : "stream-simple"; @@ -110,18 +139,20 @@ export async function resolveEmbeddedAgentApiKey(params: { return params.authStorage ? await params.authStorage.getApiKey(params.provider) : undefined; } -export function resolveEmbeddedAgentStreamFn(params: { - currentStreamFn: StreamFn | undefined; - providerStreamFn?: StreamFn; - sessionId: string; - promptCacheKey?: string; - signal?: AbortSignal; - model: EmbeddedRunAttemptParams["model"]; - resolvedApiKey?: string; - transportAuthAvailable?: boolean; - authProfileId?: string; - authStorage?: { getApiKey(provider: string): Promise }; -}): StreamFn { +export function resolveEmbeddedAgentStreamFn( + params: EmbeddedStreamRuntimeOwner & { + providerStreamFn?: StreamFn; + sessionId: string; + promptCacheKey?: string; + signal?: AbortSignal; + model: EmbeddedRunAttemptParams["model"]; + resolvedApiKey?: string; + transportAuthAvailable?: boolean; + authProfileId?: string; + authStorage?: { getApiKey(provider: string): Promise }; + }, +): StreamFn { + const llmRuntime = resolveEmbeddedStreamRuntime(params); if (params.providerStreamFn) { return wrapEmbeddedAgentStreamFn(params.providerStreamFn, { runSignal: params.signal, @@ -140,7 +171,7 @@ export function resolveEmbeddedAgentStreamFn(params: { }); } - const currentStreamFn = params.currentStreamFn ?? streamSimple; + const currentStreamFn = params.currentStreamFn ?? llmRuntime.streamSimple; if (params.model.provider === "anthropic-vertex") { return createAnthropicVertexStreamFnForModel(params.model); } @@ -148,6 +179,7 @@ export function resolveEmbeddedAgentStreamFn(params: { const openClawNativeCodexResponsesStreamFn = resolveOpenClawNativeCodexResponsesStreamFn({ model: params.model, currentStreamFn: params.currentStreamFn, + llmRuntime, }); if (openClawNativeCodexResponsesStreamFn) { return wrapEmbeddedAgentStreamFn(openClawNativeCodexResponsesStreamFn, { @@ -169,7 +201,7 @@ export function resolveEmbeddedAgentStreamFn(params: { } if ( - isDefaultOpenClawStreamFnForModel(params.model, params.currentStreamFn) || + isDefaultOpenClawStreamFnForModel(params.model, params.currentStreamFn, llmRuntime) || hasResolvedRuntimeApiKey(params.resolvedApiKey) || params.transportAuthAvailable || // Proxied anthropic-messages providers (provider !== "anthropic", e.g. pioneer) diff --git a/src/agents/embedded-agent-runner/tool-call-argument-decoding.ts b/src/agents/embedded-agent-runner/tool-call-argument-decoding.ts index ba80b8da986..eb689c4adc0 100644 --- a/src/agents/embedded-agent-runner/tool-call-argument-decoding.ts +++ b/src/agents/embedded-agent-runner/tool-call-argument-decoding.ts @@ -1,7 +1,6 @@ /** * Decodes HTML-entity escaped tool-call arguments in stream wrappers. */ -import { streamSimple } from "../../llm/stream.js"; import { decodeHtmlEntities } from "../../shared/html-entities.js"; import { visitObjectContentBlocks } from "../../shared/message-content-blocks.js"; import type { StreamFn } from "../runtime/index.js"; @@ -91,12 +90,9 @@ function wrapStreamMessageObjects( } /** Wraps a stream function so tool-call arguments are decoded before consumers inspect them. */ -export function createHtmlEntityToolCallArgumentDecodingWrapper( - baseStreamFn: StreamFn | undefined, -): StreamFn { - const underlying = baseStreamFn ?? streamSimple; +export function createHtmlEntityToolCallArgumentDecodingWrapper(baseStreamFn: StreamFn): StreamFn { return (model, context, options) => { - const maybeStream = underlying(model, context, options); + const maybeStream = baseStreamFn(model, context, options); if (maybeStream && typeof maybeStream === "object" && "then" in maybeStream) { return Promise.resolve(maybeStream).then((stream) => wrapStreamMessageObjects(stream, decodeToolCallArgumentsHtmlEntitiesInMessage), diff --git a/src/agents/google-simple-completion-stream.test.ts b/src/agents/google-simple-completion-stream.test.ts index 1b70be48a7c..5a87a10c139 100644 --- a/src/agents/google-simple-completion-stream.test.ts +++ b/src/agents/google-simple-completion-stream.test.ts @@ -1,3 +1,4 @@ +import type { ApiRegistry } from "@openclaw/ai"; // Verifies the Google simple-completion wrapper and thinking-payload sanitizer hook. import { beforeEach, describe, expect, it, vi } from "vitest"; import type { Model } from "../llm/types.js"; @@ -5,6 +6,9 @@ import type { Model } from "../llm/types.js"; const streamSimple = vi.fn(); const sanitizeGoogleThinkingPayload = vi.fn(); const ensureCustomApiRegistered = vi.fn(); +const apiRegistry = { + getApiProvider: vi.fn(() => ({ streamSimple })), +} as unknown as ApiRegistry; vi.mock("../llm/stream.js", () => ({ streamSimple, @@ -72,7 +76,7 @@ describe("prepareGoogleSimpleCompletionModel", () => { api: "openai-responses", } as unknown as Model<"openai-responses">; - const result = prepareGoogleSimpleCompletionModel(model); + const result = prepareGoogleSimpleCompletionModel(apiRegistry, model); expect(result).toBe(model); expect(ensureCustomApiRegistered).not.toHaveBeenCalled(); @@ -81,22 +85,23 @@ describe("prepareGoogleSimpleCompletionModel", () => { it("registers an OpenClaw-owned Google simple-completion api alias", () => { const model = makeGoogleModel(); - const result = prepareGoogleSimpleCompletionModel(model); + const result = prepareGoogleSimpleCompletionModel(apiRegistry, model); expect(result).toEqual({ ...model, api: GOOGLE_SIMPLE_COMPLETION_API, }); expect(ensureCustomApiRegistered).toHaveBeenCalledTimes(1); - expect(ensureCustomApiRegistered.mock.calls[0]?.[0]).toBe(GOOGLE_SIMPLE_COMPLETION_API); + expect(ensureCustomApiRegistered.mock.calls[0]?.[0]).toBe(apiRegistry); + expect(ensureCustomApiRegistered.mock.calls[0]?.[1]).toBe(GOOGLE_SIMPLE_COMPLETION_API); }); it.each(["off", "low", "medium", "high", "adaptive"] as const)( "sanitizes outbound thinking payload for gemini-flash-latest with reasoning=%s", async (reasoning) => { const model = makeGoogleModel(); - const wrapped = prepareGoogleSimpleCompletionModel(model); - const streamFn = ensureCustomApiRegistered.mock.calls[0]?.[1] as ( + const wrapped = prepareGoogleSimpleCompletionModel(apiRegistry, model); + const streamFn = ensureCustomApiRegistered.mock.calls[0]?.[2] as ( ...args: unknown[] ) => unknown; @@ -130,8 +135,8 @@ describe("prepareGoogleSimpleCompletionModel", () => { payload.generationConfig.thinkingConfig.thinkingLevel = "MINIMAL"; }); const model = makeGoogleModel(); - prepareGoogleSimpleCompletionModel(model); - const streamFn = ensureCustomApiRegistered.mock.calls[0]?.[1] as ( + prepareGoogleSimpleCompletionModel(apiRegistry, model); + const streamFn = ensureCustomApiRegistered.mock.calls[0]?.[2] as ( ...args: unknown[] ) => unknown; @@ -181,8 +186,8 @@ describe("prepareGoogleSimpleCompletionModel", () => { max: null, }, }); - const wrapped = prepareGoogleSimpleCompletionModel(model); - const streamFn = ensureCustomApiRegistered.mock.calls[0]?.[1] as ( + const wrapped = prepareGoogleSimpleCompletionModel(apiRegistry, model); + const streamFn = ensureCustomApiRegistered.mock.calls[0]?.[2] as ( ...args: unknown[] ) => unknown; diff --git a/src/agents/google-simple-completion-stream.ts b/src/agents/google-simple-completion-stream.ts index e7c598835e5..3648ef83d2e 100644 --- a/src/agents/google-simple-completion-stream.ts +++ b/src/agents/google-simple-completion-stream.ts @@ -1,3 +1,4 @@ +import type { ApiRegistry } from "@openclaw/ai"; /** * Google simple-completion stream adapter. * @@ -5,7 +6,6 @@ * backend but sanitizes unsupported thinking payload options for simple models. */ import { clampThinkingLevel } from "@openclaw/ai/internal/runtime"; -import { streamSimple } from "../llm/stream.js"; import type { Api, Model, ModelThinkingLevel } from "../llm/types.js"; import { sanitizeGoogleThinkingPayload, @@ -40,11 +40,15 @@ function resolveGoogleSimpleThinkingLevel( } } -function buildGoogleSimpleCompletionStreamFn(): StreamFn { +function buildGoogleSimpleCompletionStreamFn(registry: ApiRegistry): StreamFn { return (model, context, options) => { const googleModel: Model = { ...model, api: SOURCE_API }; + const sourceProvider = registry.getApiProvider(SOURCE_API); + if (!sourceProvider) { + throw new Error(`No API provider registered for api: ${SOURCE_API}`); + } return streamWithPayloadPatch( - streamSimple as unknown as StreamFn, + sourceProvider.streamSimple as StreamFn, googleModel, context, options, @@ -63,10 +67,17 @@ function buildGoogleSimpleCompletionStreamFn(): StreamFn { } /** Rewrites Google generative-ai models to the simple-completion adapter when needed. */ -export function prepareGoogleSimpleCompletionModel(model: Model): Model { +export function prepareGoogleSimpleCompletionModel( + registry: ApiRegistry, + model: Model, +): Model { if (model.api !== SOURCE_API) { return model; } - ensureCustomApiRegistered(GOOGLE_SIMPLE_COMPLETION_API, buildGoogleSimpleCompletionStreamFn()); + ensureCustomApiRegistered( + registry, + GOOGLE_SIMPLE_COMPLETION_API, + buildGoogleSimpleCompletionStreamFn(registry), + ); return { ...model, api: GOOGLE_SIMPLE_COMPLETION_API }; } diff --git a/src/agents/model-scan.ts b/src/agents/model-scan.ts index ebb40744e6c..de9be83dc12 100644 --- a/src/agents/model-scan.ts +++ b/src/agents/model-scan.ts @@ -1,5 +1,7 @@ +import { createLlmRuntime, type LlmRuntime } from "@openclaw/ai"; import type { OpenAICompletionsOptions } from "@openclaw/ai/internal/openai"; import { getEnvApiKey } from "@openclaw/ai/internal/runtime"; +import { registerBuiltInApiProviders } from "@openclaw/ai/providers"; import { normalizeProviderId } from "@openclaw/model-catalog-core/provider-id"; import { asDateTimestampMs, @@ -20,7 +22,7 @@ import { formatErrorMessage } from "../infra/errors.js"; * Scans remote provider model catalogs for configured providers. */ import { readResponseWithLimit } from "../infra/http-body.js"; -import { complete } from "../llm/stream.js"; +import "../llm/ai-transport-host.js"; import type { Context, Model, Tool } from "../llm/types.js"; import { inferParamBFromIdOrName } from "../shared/model-param-b.js"; @@ -299,6 +301,7 @@ async function probeTool( model: OpenAIModel, apiKey: string, timeoutMs: number, + complete: LlmRuntime["complete"], ): Promise { const context: Context = { messages: [ @@ -345,6 +348,7 @@ async function probeImage( model: OpenAIModel, apiKey: string, timeoutMs: number, + complete: LlmRuntime["complete"], ): Promise { const context: Context = { messages: [ @@ -433,6 +437,8 @@ export async function scanOpenRouterModels( const providerFilter = normalizeProviderId(options.providerFilter ?? ""); const catalog = await fetchOpenRouterModels(fetchImpl, timeoutMs); + const llmRuntime = createLlmRuntime(); + registerBuiltInApiProviders(llmRuntime.registry); const now = Date.now(); const filtered = catalog.filter((entry) => { @@ -504,9 +510,9 @@ export async function scanOpenRouterModels( reasoning: baseModel.reasoning, }; - const toolResult = await probeTool(model, apiKey, timeoutMs); + const toolResult = await probeTool(model, apiKey, timeoutMs, llmRuntime.complete); const imageResult = model.input?.includes("image") - ? await probeImage(ensureImageInput(model), apiKey, timeoutMs) + ? await probeImage(ensureImageInput(model), apiKey, timeoutMs, llmRuntime.complete) : { ok: false, latencyMs: null, skipped: true }; result = buildOpenRouterScanResult({ diff --git a/src/agents/models.profiles.live.test.ts b/src/agents/models.profiles.live.test.ts index 5768eebe45b..b3bf5d428a2 100644 --- a/src/agents/models.profiles.live.test.ts +++ b/src/agents/models.profiles.live.test.ts @@ -1,5 +1,6 @@ // Live-sweeps discovered model profiles with optional provider/model filters and probes. import { writeSync } from "node:fs"; +import { defaultApiRegistry } from "@openclaw/ai/internal/runtime"; import { normalizeProviderId } from "@openclaw/model-catalog-core/provider-id"; import { expectDefined } from "@openclaw/normalization-core"; import { type Api, completeSimple, type Model } from "openclaw/plugin-sdk/llm"; @@ -347,6 +348,7 @@ async function ensureLiveProviderApisRegistered(params: { const providerConfig = params.config.models?.providers?.ollama; const providerBaseUrl = readConfiguredOllamaBaseUrl(providerConfig) || OLLAMA_DEFAULT_BASE_URL; ensureCustomApiRegistered( + defaultApiRegistry, "ollama", createLiveOllamaRuntimeStreamFn({ createConfiguredOllamaStreamFn, @@ -1436,6 +1438,7 @@ async function completeSimpleWithTimeout( }); try { const completionModel = prepareModelForSimpleCompletion({ + apiRegistry: defaultApiRegistry, model, cfg: activeLiveCompletionConfig, }); diff --git a/src/agents/provider-stream.lifecycle.test.ts b/src/agents/provider-stream.lifecycle.test.ts new file mode 100644 index 00000000000..3c21aebe173 --- /dev/null +++ b/src/agents/provider-stream.lifecycle.test.ts @@ -0,0 +1,39 @@ +import { createApiRegistry, createLlmRuntime } from "@openclaw/ai"; +import { describe, expect, it, vi } from "vitest"; +import { bindModelLlmRuntime } from "../llm/model-runtime-binding.js"; +import { createAssistantMessageEventStream } from "../llm/utils/event-stream.js"; +import { registerProviderStreamForModel } from "./provider-stream.js"; + +const { providerStream } = vi.hoisted(() => ({ + providerStream: vi.fn(), +})); + +vi.mock("../plugins/provider-runtime.js", () => ({ + resolveProviderStreamFn: () => providerStream, +})); + +describe("provider stream lifecycle registration", () => { + it("registers provider streams into the prepared model runtime", () => { + providerStream.mockReturnValue(createAssistantMessageEventStream()); + const apiRegistry = createApiRegistry(); + const llmRuntime = createLlmRuntime(apiRegistry); + const model = bindModelLlmRuntime( + { + api: "test-lifecycle-provider", + provider: "test-provider", + id: "test-model", + name: "Test Model", + baseUrl: "https://example.test", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1024, + maxTokens: 512, + }, + llmRuntime, + ); + + expect(registerProviderStreamForModel({ model })).toBeTypeOf("function"); + expect(apiRegistry.getApiProvider("test-lifecycle-provider")).toBeDefined(); + }); +}); diff --git a/src/agents/provider-stream.ts b/src/agents/provider-stream.ts index 8a7af02e403..b071ee6a52b 100644 --- a/src/agents/provider-stream.ts +++ b/src/agents/provider-stream.ts @@ -3,7 +3,9 @@ * Resolves plugin-owned or transport-aware stream functions and registers the * model API once a concrete stream implementation exists. */ +import type { ApiRegistry } from "@openclaw/ai"; import type { OpenClawConfig } from "../config/types.openclaw.js"; +import { getModelLlmRuntime } from "../llm/model-runtime-binding.js"; import type { Api, Model } from "../llm/types.js"; import { resolveProviderStreamFn } from "../plugins/provider-runtime.js"; import { ensureCustomApiRegistered } from "./custom-api-registry.js"; @@ -23,7 +25,7 @@ export function registerProviderStreamForModel(params: { workspaceDir?: string; env?: NodeJS.ProcessEnv; allowRuntimePluginLoad?: boolean; - registerStream?: boolean; + apiRegistry?: ApiRegistry; }): StreamFn | undefined { // Plugin stream factories may capture model headers, so construction is the // last safe boundary for providers that do not expose the host fetch seam. @@ -67,8 +69,9 @@ export function registerProviderStreamForModel(params: { } // Register custom APIs only after a concrete stream exists, so later callers // can route by model.api without reloading provider runtime hooks. - if (params.registerStream !== false) { - ensureCustomApiRegistered(params.model.api, streamFn); + const apiRegistry = params.apiRegistry ?? getModelLlmRuntime(params.model)?.registry; + if (apiRegistry) { + ensureCustomApiRegistered(apiRegistry, params.model.api, streamFn); } return streamFn; } diff --git a/src/agents/sessions/agent-session-base.ts b/src/agents/sessions/agent-session-base.ts index a16b00e7b98..3cb22f1490a 100644 --- a/src/agents/sessions/agent-session-base.ts +++ b/src/agents/sessions/agent-session-base.ts @@ -1,5 +1,4 @@ import { cleanupSessionResources } from "@openclaw/ai/internal/runtime"; -import { streamSimple } from "../../llm/stream.js"; import type { AssistantMessage, Model } from "../../llm/types.js"; import type { Agent, @@ -36,6 +35,7 @@ import { type TurnStartEvent, } from "./extensions/index.js"; import type { BashExecutionMessage, CustomMessage } from "./messages.js"; +import { getModelRegistryRuntime } from "./model-registry-runtime.js"; import type { ModelRegistry } from "./model-registry.js"; import type { PromptTemplate } from "./prompt-templates.js"; import type { ResourceLoader } from "./resource-loader.js"; @@ -183,7 +183,10 @@ export abstract class AgentSessionBase { apiKey?: string; headers?: Record; }> { - if (this.agent.streamFn === streamSimple) { + if ( + this.agent.streamFn === + getModelRegistryRuntime(this.sessionModelRegistry).llmRuntime.streamSimple + ) { return this.getRequiredRequestAuth(model); } diff --git a/src/agents/sessions/agent-session-compaction.ts b/src/agents/sessions/agent-session-compaction.ts index dbf60f76d11..f9783cf3952 100644 --- a/src/agents/sessions/agent-session-compaction.ts +++ b/src/agents/sessions/agent-session-compaction.ts @@ -1,5 +1,4 @@ import { isContextOverflow } from "@openclaw/ai/internal/runtime"; -import { streamSimple } from "../../llm/stream.js"; import type { AssistantMessage, Model } from "../../llm/types.js"; import { calculateContextTokens, @@ -14,6 +13,7 @@ import { AgentSessionInspection } from "./agent-session-inspection.js"; import { unwrapCoreResult } from "./agent-session-utils.js"; import { formatNoModelSelectedMessage } from "./auth-guidance.js"; import { preflightManualSessionCompaction } from "./manual-compaction-preflight.js"; +import { getModelRegistryRuntime } from "./model-registry-runtime.js"; import { getLatestCompactionEntry, type CompactionEntry } from "./session-manager.js"; import type { SettingsManager } from "./settings-manager.js"; @@ -109,7 +109,10 @@ export abstract class AgentSessionCompaction extends AgentSessionInspection { } | undefined > { - if (this.agent.streamFn !== streamSimple) { + if ( + this.agent.streamFn !== + getModelRegistryRuntime(this.sessionModelRegistry).llmRuntime.streamSimple + ) { return this.getCompactionRequestAuth(model); } diff --git a/src/agents/sessions/agent-session-extensions.ts b/src/agents/sessions/agent-session-extensions.ts index 2a1b67685db..f6684bdc0f5 100644 --- a/src/agents/sessions/agent-session-extensions.ts +++ b/src/agents/sessions/agent-session-extensions.ts @@ -1,6 +1,4 @@ import { basename, dirname } from "node:path"; -import { defaultApiRegistry } from "@openclaw/ai/internal/runtime"; -import { resetApiProviders } from "@openclaw/ai/providers"; import { AgentSessionCompaction } from "./agent-session-compaction.js"; import type { ExtensionBindings } from "./agent-session-types.js"; import { ExtensionRunner, type ToolDefinition, wrapRegisteredTools } from "./extensions/index.js"; @@ -401,8 +399,8 @@ export abstract class AgentSessionExtensions extends AgentSessionCompaction { await this.settingsManager.reload(); this.agent.steeringMode = this.settingsManager.getSteeringMode(); this.agent.followUpMode = this.settingsManager.getFollowUpMode(); - resetApiProviders(defaultApiRegistry); await this.sessionResourceLoader.reload(); + this.sessionModelRegistry.refresh(); this.buildRuntime({ activeToolNames: this.getActiveToolNames(), flagValues: previousFlagValues, diff --git a/src/agents/sessions/agent-session-loop-correctness.test.ts b/src/agents/sessions/agent-session-loop-correctness.test.ts index 571c458b5b0..5630db8a63d 100644 --- a/src/agents/sessions/agent-session-loop-correctness.test.ts +++ b/src/agents/sessions/agent-session-loop-correctness.test.ts @@ -12,10 +12,6 @@ const streamMocks = vi.hoisted(() => ({ streamSimple: vi.fn(), })); -vi.mock("../../llm/stream.js", () => ({ - streamSimple: streamMocks.streamSimple, -})); - import type { AgentTool } from "../runtime/index.js"; import type { AgentSessionEvent } from "./agent-session-types.js"; import { AgentSession } from "./agent-session.js"; @@ -172,6 +168,11 @@ async function createTestSession( retry: { enabled: false }, }); const sessionManager = options.sessionManager ?? SessionManager.inMemory(); + const modelRegistry = ModelRegistry.inMemory(authStorage); + modelRegistry.registerProvider(model.provider, { + api: model.api, + streamSimple: streamMocks.streamSimple, + }); const result = await createAgentSession({ model, noTools: "builtin", @@ -179,7 +180,7 @@ async function createTestSession( resourceLoader: options.resourceLoader ?? createResourceLoader(), sessionManager, settingsManager, - modelRegistry: ModelRegistry.inMemory(authStorage), + modelRegistry, }); sessions.push(result.session); return { ...result, settingsManager, sessionManager }; diff --git a/src/agents/sessions/model-registry-runtime.ts b/src/agents/sessions/model-registry-runtime.ts new file mode 100644 index 00000000000..504a8550bca --- /dev/null +++ b/src/agents/sessions/model-registry-runtime.ts @@ -0,0 +1,50 @@ +import { + createApiRegistry, + createLlmRuntime, + type ApiRegistry, + type LlmRuntime, +} from "@openclaw/ai"; +import { getPublishedApiProviders } from "@openclaw/ai/internal/runtime"; +import { registerBuiltInApiProviders } from "@openclaw/ai/providers"; +import "../../llm/ai-transport-host.js"; +import { bindStreamLlmRuntime } from "../../llm/model-runtime-binding.js"; + +type ModelRegistryRuntime = { + apiRegistry: ApiRegistry; + llmRuntime: LlmRuntime; +}; + +const modelRegistryRuntimes = new WeakMap(); + +function resetApiRegistry(runtime: ModelRegistryRuntime): void { + runtime.apiRegistry.clearApiProviders(); + registerBuiltInApiProviders(runtime.apiRegistry); + // The Plugin SDK registry is a shipped compatibility facade. Snapshot it at + // lifecycle publication so request-time routing never depends on mutable global state. + for (const provider of getPublishedApiProviders()) { + runtime.apiRegistry.registerApiProvider(provider); + } +} + +/** Creates the runtime facts owned by one model-registry lifecycle. */ +export function initializeModelRegistryRuntime(owner: object): void { + const apiRegistry = createApiRegistry(); + const llmRuntime = createLlmRuntime(apiRegistry); + const runtime = { apiRegistry, llmRuntime }; + bindStreamLlmRuntime(llmRuntime.streamSimple, llmRuntime); + resetApiRegistry(runtime); + modelRegistryRuntimes.set(owner, runtime); +} + +/** Returns the prepared runtime facts for one model-registry lifecycle. */ +export function getModelRegistryRuntime(owner: object): ModelRegistryRuntime { + const runtime = modelRegistryRuntimes.get(owner); + if (!runtime) { + throw new Error("Model registry runtime is not initialized"); + } + return runtime; +} + +export function resetModelRegistryRuntime(owner: object): void { + resetApiRegistry(getModelRegistryRuntime(owner)); +} diff --git a/src/agents/sessions/model-registry.test.ts b/src/agents/sessions/model-registry.test.ts index da7d3413fd4..9e84831e47d 100644 --- a/src/agents/sessions/model-registry.test.ts +++ b/src/agents/sessions/model-registry.test.ts @@ -3,9 +3,16 @@ import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { dirname, join } from "node:path"; -import { afterEach, describe, expect, it } from "vitest"; +import { + defaultApiRegistry, + getApiProvider, + registerApiProvider, + unregisterApiProviders, +} from "@openclaw/ai/internal/runtime"; +import { afterEach, describe, expect, it, vi } from "vitest"; import { PLUGIN_MODEL_CATALOG_GENERATED_BY } from "../plugin-model-catalog.js"; import { AuthStorage } from "./auth-storage.js"; +import { getModelRegistryRuntime } from "./model-registry-runtime.js"; import { ModelRegistry, type ProviderConfigInput } from "./model-registry.js"; const PLUGIN_MODEL_CATALOG_FILE = "catalog.json"; @@ -595,3 +602,58 @@ describe("ModelRegistry OAuth provider ownership", () => { ).toBe("Anthropic (Claude Pro/Max)"); }); }); + +describe("ModelRegistry API provider ownership", () => { + it("keeps stream registrations isolated across registry refreshes", () => { + const sessionA = ModelRegistry.inMemory(AuthStorage.inMemory()); + const sessionB = ModelRegistry.inMemory(AuthStorage.inMemory()); + const streamA = vi.fn(() => ({}) as never); + const streamB = vi.fn(() => ({}) as never); + + sessionA.registerProvider("session-a", { + api: "test-session-api", + streamSimple: streamA, + }); + sessionB.registerProvider("session-b", { + api: "test-session-api", + streamSimple: streamB, + }); + const runtimeA = getModelRegistryRuntime(sessionA); + const runtimeB = getModelRegistryRuntime(sessionB); + + expect(runtimeA.apiRegistry.getApiProvider("test-session-api")?.streamSimple).not.toBe( + runtimeB.apiRegistry.getApiProvider("test-session-api")?.streamSimple, + ); + expect(getApiProvider("test-session-api")).toBeUndefined(); + + sessionB.unregisterProvider("session-b"); + + expect(runtimeA.apiRegistry.getApiProvider("test-session-api")).toBeDefined(); + expect(runtimeB.apiRegistry.getApiProvider("test-session-api")).toBeUndefined(); + }); + + it("imports published SDK providers without copying request-generated aliases", () => { + const publishedSource = "plugin:test-published-api"; + const requestSource = "custom-api:test-request-api"; + const stream = vi.fn(() => ({}) as never); + registerApiProvider( + { api: "test-published-api", stream, streamSimple: stream }, + publishedSource, + ); + defaultApiRegistry.registerApiProvider( + { api: "test-request-api", stream, streamSimple: stream }, + requestSource, + ); + + try { + const session = ModelRegistry.inMemory(AuthStorage.inMemory()); + const runtime = getModelRegistryRuntime(session); + + expect(runtime.apiRegistry.getApiProvider("test-published-api")).toBeDefined(); + expect(runtime.apiRegistry.getApiProvider("test-request-api")).toBeUndefined(); + } finally { + unregisterApiProviders(publishedSource); + defaultApiRegistry.unregisterApiProviders(requestSource); + } + }); +}); diff --git a/src/agents/sessions/model-registry.ts b/src/agents/sessions/model-registry.ts index bfeabc84ee1..e30d20a7045 100644 --- a/src/agents/sessions/model-registry.ts +++ b/src/agents/sessions/model-registry.ts @@ -4,8 +4,6 @@ import { existsSync, readFileSync } from "node:fs"; import { dirname, join } from "node:path"; -import { defaultApiRegistry, registerApiProvider } from "@openclaw/ai/internal/runtime"; -import { resetApiProviders } from "@openclaw/ai/providers"; import { type Static, Type } from "typebox"; import { Compile } from "typebox/compile"; import type { TLocalizedValidationError } from "typebox/error"; @@ -31,6 +29,11 @@ import { } from "../plugin-model-catalog.js"; import { getAuthStorageOAuthProviderRegistry } from "./auth-storage-oauth-registry.js"; import type { AuthStatus, AuthStorage } from "./auth-storage.js"; +import { + getModelRegistryRuntime, + initializeModelRegistryRuntime, + resetModelRegistryRuntime, +} from "./model-registry-runtime.js"; import { BUILT_IN_PROVIDER_DISPLAY_NAMES } from "./provider-display-names.js"; import { clearConfigValueCache, @@ -322,6 +325,7 @@ export class ModelRegistry { options: ModelRegistryOptions = {}, ) { this.authStorage = authStorage; + initializeModelRegistryRuntime(this); this.modelsJsonPath = modelsJsonPath; this.pluginMetadataSnapshot = resolveModelPluginMetadataSnapshot({ ...(options.pluginMetadataSnapshot @@ -354,8 +358,8 @@ export class ModelRegistry { this.modelRequestHeaders.clear(); this.loadError = undefined; - // Ensure dynamic API/OAuth registrations are rebuilt from current provider state. - resetApiProviders(defaultApiRegistry); + // Rebuild this lifecycle's API/OAuth registrations from current provider state. + resetModelRegistryRuntime(this); getAuthStorageOAuthProviderRegistry(this.authStorage).reset(); this.loadModels(); @@ -859,7 +863,7 @@ export class ModelRegistry { if (config.streamSimple) { const streamSimple = config.streamSimple; - registerApiProvider( + getModelRegistryRuntime(this).apiRegistry.registerApiProvider( { api: config.api!, stream: (model, context, options) => diff --git a/src/agents/sessions/sdk.test.ts b/src/agents/sessions/sdk.test.ts index 11e2849dff9..10640b1948a 100644 --- a/src/agents/sessions/sdk.test.ts +++ b/src/agents/sessions/sdk.test.ts @@ -3,6 +3,7 @@ import { createAssistantMessageEventStream, type AssistantMessage } from "opencl // session write-lock behavior. import { Type } from "typebox"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { getStreamLlmRuntime } from "../../llm/model-runtime-binding.js"; import type { Model, SimpleStreamOptions } from "../../llm/types.js"; import { createUserTurnTranscriptRecorder } from "../../sessions/user-turn-transcript.js"; import { createTestUserTurnTranscriptTarget } from "../../sessions/user-turn-transcript.test-support.js"; @@ -24,6 +25,7 @@ import { takeRuntimeUserTurnTranscriptContext } from "../../sessions/user-turn-t import { AuthStorage } from "./auth-storage.js"; import { createExtensionRuntime } from "./extensions/loader.js"; import type { LoadExtensionsResult, ToolDefinition } from "./extensions/types.js"; +import { getModelRegistryRuntime } from "./model-registry-runtime.js"; import { ModelRegistry } from "./model-registry.js"; import type { ResourceLoader } from "./resource-loader.js"; import { createAgentSession } from "./sdk.js"; @@ -44,6 +46,23 @@ const testModel: Model = { maxTokens: 1000, }; +describe("createAgentSession runtime ownership", () => { + it("binds the installed stream wrapper to the model-registry lifecycle", async () => { + const modelRegistry = createTestModelRegistry(); + const { session } = await createAgentSession({ + model: testModel, + resourceLoader: createEmptyResourceLoader(), + sessionManager: SessionManager.inMemory(), + settingsManager: SettingsManager.inMemory(), + modelRegistry, + }); + + expect(getStreamLlmRuntime(session.agent.streamFn)).toBe( + getModelRegistryRuntime(modelRegistry).llmRuntime, + ); + }); +}); + function createModelWithoutBaseUrl(overrides: Partial): Model { const { baseUrl: _baseUrl, ...model } = { ...testModel, ...overrides }; return model as unknown as Model; @@ -87,6 +106,17 @@ function createEmptyResourceLoader(): ResourceLoader { return createResourceLoaderWithHandlers(new Map()); } +function createTestModelRegistry(authStorage = AuthStorage.inMemory()): ModelRegistry { + const modelRegistry = ModelRegistry.inMemory(authStorage); + for (const api of ["openai-responses", "bedrock-converse-stream"] as const) { + modelRegistry.registerProvider(`test-${api}`, { + api, + streamSimple: streamMocks.streamSimple, + }); + } + return modelRegistry; +} + function createResourceLoaderWithHandlers( handlers: Map Promise>>, ): ResourceLoader { @@ -130,7 +160,7 @@ async function createSessionAndStreamModel(model: Model): Promise { settingsManager: SettingsManager.inMemory({ retry: retry ?? { baseDelayMs: 0, maxRetries: 1 }, }), - modelRegistry: ModelRegistry.inMemory(authStorage), + modelRegistry: createTestModelRegistry(authStorage), }); } diff --git a/src/agents/sessions/sdk.ts b/src/agents/sessions/sdk.ts index 772f99fa81b..4da8604a39c 100644 --- a/src/agents/sessions/sdk.ts +++ b/src/agents/sessions/sdk.ts @@ -9,7 +9,7 @@ import { resolveThinkingDefaultForModel, type ThinkingCatalogEntry, } from "../../auto-reply/thinking.js"; -import { streamSimple } from "../../llm/stream.js"; +import { bindStreamLlmRuntime } from "../../llm/model-runtime-binding.js"; import type { Message, Model } from "../../llm/types.js"; import { getAgentDir } from "../config.js"; import { @@ -30,6 +30,7 @@ import type { ToolDefinition, } from "./extensions/index.js"; import { convertToLlm } from "./messages.js"; +import { getModelRegistryRuntime } from "./model-registry-runtime.js"; import { ModelRegistry } from "./model-registry.js"; import { findInitialModel } from "./model-resolver.js"; import { DefaultResourceLoader, type ResourceLoader } from "./resource-loader.js"; @@ -293,6 +294,7 @@ export async function createAgentSession( if (!resourceLoader) { resourceLoader = new DefaultResourceLoader({ cwd, agentDir, settingsManager }); await resourceLoader.reload(); + modelRegistry.refresh(); } // Check if session has existing data to restore @@ -438,6 +440,7 @@ export async function createAgentSession( const runWithSessionWriteLock = async (run: () => Promise | T): Promise => options.withSessionWriteLock ? await options.withSessionWriteLock(run) : await run(); + const modelRegistryRuntime = getModelRegistryRuntime(modelRegistry); const agent: Agent = new Agent({ initialState: { systemPrompt: "", @@ -453,7 +456,7 @@ export async function createAgentSession( } const providerRetrySettings = settingsManager.getProviderRetrySettings(); const attributionHeaders = getAttributionHeaders(modelResult, settingsManager); - return streamSimple(modelResult, context, { + return modelRegistryRuntime.llmRuntime.streamSimple(modelResult, context, { ...optionsLocal, apiKey: auth.apiKey, timeoutMs: optionsLocal?.timeoutMs ?? providerRetrySettings.timeoutMs, @@ -506,6 +509,9 @@ export async function createAgentSession( thinkingBudgets: settingsManager.getThinkingBudgets(), maxRetryDelayMs: settingsManager.getProviderRetrySettings().maxRetryDelayMs, }); + if (agent.streamFn) { + bindStreamLlmRuntime(agent.streamFn, modelRegistryRuntime.llmRuntime); + } // Restore messages if session has existing data if (hasExistingSession) { diff --git a/src/agents/simple-completion-runtime.test.ts b/src/agents/simple-completion-runtime.test.ts index cea29396078..89a1e42a4f1 100644 --- a/src/agents/simple-completion-runtime.test.ts +++ b/src/agents/simple-completion-runtime.test.ts @@ -31,6 +31,20 @@ vi.mock("../llm/stream.js", () => ({ completeSimple: hoisted.completeMock, })); +vi.mock("./sessions/model-registry-runtime.js", () => ({ + getModelRegistryRuntime: () => { + const apiRegistry = {}; + return { + apiRegistry, + llmRuntime: { + registry: apiRegistry, + completeSimple: (...args: unknown[]) => hoisted.completeMock(...args), + streamSimple: vi.fn(), + }, + }; + }, +})); + vi.mock("./embedded-agent-runner/model.js", () => ({ resolveModel: hoisted.resolveModelMock, resolveModelAsync: hoisted.resolveModelAsyncMock, @@ -880,7 +894,11 @@ describe("completeWithPreparedSimpleCompletionModel", () => { }, }); - expect(hoisted.prepareModelForSimpleCompletionMock).toHaveBeenCalledWith({ model, cfg }); + expect(hoisted.prepareModelForSimpleCompletionMock).toHaveBeenCalledWith({ + apiRegistry: expect.anything(), + model, + cfg, + }); expect(hoisted.completeMock).toHaveBeenCalledWith( preparedModel, { diff --git a/src/agents/simple-completion-runtime.ts b/src/agents/simple-completion-runtime.ts index 1314842508a..85ca9bb5c06 100644 --- a/src/agents/simple-completion-runtime.ts +++ b/src/agents/simple-completion-runtime.ts @@ -1,4 +1,5 @@ import { supportsOpenAIReasoningEffort } from "@openclaw/ai/internal/openai"; +import { defaultApiRegistry } from "@openclaw/ai/internal/runtime"; import { resolveClaudeSonnet5ModelIdentity } from "@openclaw/llm-core"; /** * Simple completion runtime preparation. @@ -8,6 +9,7 @@ import { resolveClaudeSonnet5ModelIdentity } from "@openclaw/llm-core"; import type { ThinkLevel } from "../auto-reply/thinking.js"; import type { OpenClawConfig } from "../config/types.openclaw.js"; import { formatErrorMessage } from "../infra/errors.js"; +import { bindModelLlmRuntime, getModelLlmRuntime } from "../llm/model-runtime-binding.js"; import { completeSimple } from "../llm/stream.js"; import type { AssistantMessage, @@ -49,6 +51,7 @@ import { applyPreparedRuntimeAuthToModel } from "./provider-request-config.js"; import { protectPreparedProviderRuntimeAuth } from "./provider-secret-egress.js"; import { buildAgentRuntimeAuthPlan } from "./runtime-plan/auth.js"; import { materializePreparedRuntimeModel } from "./runtime-plan/materialize-model.js"; +import { getModelRegistryRuntime } from "./sessions/model-registry-runtime.js"; import { resolveSimpleCompletionModelResolverWorkspace } from "./simple-completion-scope.js"; import { prepareModelForSimpleCompletion } from "./simple-completion-transport.js"; import { resolveUtilityModelRefForAgent } from "./utility-model.js"; @@ -445,11 +448,15 @@ export async function prepareSimpleCompletionModel(params: { }) : fingerprintResolvedProviderAuth(auth) : undefined; + const modelRuntime = getModelRegistryRuntime(resolved.modelRegistry); return { - model: applySecretRefHeaderSentinels( - applyLocalNoAuthHeaderOverride(resolvedModel, resolvedAuth), - params.cfg, + model: bindModelLlmRuntime( + applySecretRefHeaderSentinels( + applyLocalNoAuthHeaderOverride(resolvedModel, resolvedAuth), + params.cfg, + ), + modelRuntime.llmRuntime, ), auth: resolvedAuth, ...(sourceAuthFingerprint ? { sourceAuthFingerprint } : {}), @@ -521,7 +528,17 @@ export async function completeWithPreparedSimpleCompletionModel(params: { cfg?: OpenClawConfig; options?: SimpleCompletionModelOptions; }): Promise { - const completionModel = prepareModelForSimpleCompletion({ model: params.model, cfg: params.cfg }); + const runtime = getModelLlmRuntime(params.model); + let completionModel = prepareModelForSimpleCompletion({ + // Direct SDK callers that did not use the preparation helper keep the shipped + // process-default behavior; all prepared host paths carry their lifecycle owner. + apiRegistry: runtime?.registry ?? defaultApiRegistry, + model: params.model, + cfg: params.cfg, + }); + if (runtime) { + completionModel = bindModelLlmRuntime(completionModel, runtime); + } const { reasoning: rawReasoning, ...options } = params.options ?? {}; const reasoning = normalizeSimpleCompletionReasoning(rawReasoning, completionModel); return await completeSimple(completionModel, params.context, { diff --git a/src/agents/simple-completion-transport.test.ts b/src/agents/simple-completion-transport.test.ts index 23273d336d7..9cef0d47870 100644 --- a/src/agents/simple-completion-transport.test.ts +++ b/src/agents/simple-completion-transport.test.ts @@ -1,8 +1,8 @@ -import { registerApiProvider, unregisterApiProviders } from "@openclaw/ai/internal/runtime"; +import { createApiRegistry, type ApiRegistry } from "@openclaw/ai"; // Simple completion transport tests cover provider-specific stream alias // selection before the generic completion helper invokes the LLM layer. import type { Model } from "openclaw/plugin-sdk/llm"; -import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; +import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import type { OpenClawConfig } from "../config/config.js"; import { createMoonshotThinkingWrapper } from "../llm/providers/stream-wrappers/moonshot-thinking.js"; import { mintSecretSentinel } from "../secrets/sentinel.js"; @@ -17,7 +17,7 @@ const createOpenClawTransportStreamFnForModel = vi.fn(); const createTransportAwareStreamFnForModel = vi.fn(); const prepareTransportAwareSimpleModel = vi.fn(); const resolveTransportAwareSimpleApi = vi.fn(); -const prepareGoogleSimpleCompletionModel = vi.fn((model: unknown) => model); +const prepareGoogleSimpleCompletionModel = vi.fn((_registry: unknown, model: unknown) => model); const pluginStreamFn = vi.fn(() => "plugin-stream-result" as never); vi.mock("./anthropic-vertex-stream.js", () => ({ @@ -51,17 +51,31 @@ vi.mock("../plugins/provider-runtime.js", async () => { }; }); -let prepareModelForSimpleCompletion: typeof import("./simple-completion-transport.js").prepareModelForSimpleCompletion; +let prepareModelForSimpleCompletionImpl: typeof import("./simple-completion-transport.js").prepareModelForSimpleCompletion; +let apiRegistry: ApiRegistry; const SIMPLE_COMPLETION_SOURCE_ID = "test:simple-completion-transport"; +function prepareModelForSimpleCompletion( + params: Omit< + Parameters< + typeof import("./simple-completion-transport.js").prepareModelForSimpleCompletion + >[0], + "apiRegistry" + >, +) { + return prepareModelForSimpleCompletionImpl({ ...params, apiRegistry }); +} + describe("prepareModelForSimpleCompletion", () => { beforeAll(async () => { // Dynamic import lets the mocked transport/provider modules settle before // the unit under test captures custom stream registration helpers. - ({ prepareModelForSimpleCompletion } = await import("./simple-completion-transport.js")); + ({ prepareModelForSimpleCompletion: prepareModelForSimpleCompletionImpl } = + await import("./simple-completion-transport.js")); }); beforeEach(() => { + apiRegistry = createApiRegistry(); createAnthropicVertexStreamFnForModel.mockReset(); ensureCustomApiRegistered.mockReset(); resolveProviderStreamFn.mockReset(); @@ -81,18 +95,14 @@ describe("prepareModelForSimpleCompletion", () => { createTransportAwareStreamFnForModel.mockReturnValue(undefined); prepareTransportAwareSimpleModel.mockImplementation((model) => model); resolveTransportAwareSimpleApi.mockReturnValue(undefined); - prepareGoogleSimpleCompletionModel.mockImplementation((model) => model); - }); - - afterEach(() => { - unregisterApiProviders(SIMPLE_COMPLETION_SOURCE_ID); + prepareGoogleSimpleCompletionModel.mockImplementation((_registry, model) => model); }); it("routes provider-owned simple-completion wrappers through an internal API alias", () => { const sourceApi = "moonshot-simple-source"; const sourceResult = { source: true }; let capturedApi: string | undefined; - registerApiProvider( + apiRegistry.registerApiProvider( { api: sourceApi, stream: () => sourceResult as never, @@ -137,7 +147,7 @@ describe("prepareModelForSimpleCompletion", () => { }), }), ); - const registeredStream = ensureCustomApiRegistered.mock.calls.at(-1)?.[1]; + const registeredStream = ensureCustomApiRegistered.mock.calls.at(-1)?.[2]; expect(registeredStream).toBeTypeOf("function"); const stream = registeredStream(result, { messages: [] }, {}); expect(stream).toBe(sourceResult); @@ -193,8 +203,12 @@ describe("prepareModelForSimpleCompletion", () => { ...model, headers: { Authorization: `Bearer ${secret}` }, }); - expect(ensureCustomApiRegistered).toHaveBeenCalledWith("ollama", expect.any(Function)); - const registeredStream = ensureCustomApiRegistered.mock.calls[0]?.[1] as StreamFn; + expect(ensureCustomApiRegistered).toHaveBeenCalledWith( + apiRegistry, + "ollama", + expect.any(Function), + ); + const registeredStream = ensureCustomApiRegistered.mock.calls[0]?.[2] as StreamFn; void registeredStream( { ...model, headers: { Authorization: `Bearer ${sentinel}` } } as never, {} as never, @@ -228,6 +242,7 @@ describe("prepareModelForSimpleCompletion", () => { expect(createAnthropicVertexStreamFnForModel).toHaveBeenCalledWith(model); expect(ensureCustomApiRegistered).toHaveBeenCalledWith( + apiRegistry, "openclaw-anthropic-vertex-simple:https%3A%2F%2Fus-central1-aiplatform.googleapis.com", "vertex-stream", ); @@ -263,6 +278,7 @@ describe("prepareModelForSimpleCompletion", () => { expect(prepareTransportAwareSimpleModel).toHaveBeenCalledWith(model, { cfg: undefined }); expect(buildTransportAwareSimpleStreamFn).toHaveBeenCalledWith(model, { cfg: undefined }); expect(ensureCustomApiRegistered).toHaveBeenCalledWith( + apiRegistry, "openclaw-openai-responses-transport", "transport-stream", ); @@ -286,7 +302,7 @@ describe("prepareModelForSimpleCompletion", () => { maxTokens: 8192, headers: {}, }; - prepareGoogleSimpleCompletionModel.mockImplementationOnce((m: unknown) => ({ + prepareGoogleSimpleCompletionModel.mockImplementationOnce((_registry: unknown, m: unknown) => ({ ...(m as Model<"google-generative-ai">), api: "openclaw-google-generative-ai-simple", })); @@ -295,7 +311,7 @@ describe("prepareModelForSimpleCompletion", () => { const result = prepareModelForSimpleCompletion({ model }); expect(prepareTransportAwareSimpleModel).toHaveBeenCalledWith(model, { cfg: undefined }); - expect(prepareGoogleSimpleCompletionModel).toHaveBeenCalledWith(model); + expect(prepareGoogleSimpleCompletionModel).toHaveBeenCalledWith(apiRegistry, model); expect(buildTransportAwareSimpleStreamFn).not.toHaveBeenCalled(); expect(result).toEqual({ ...model, @@ -330,6 +346,7 @@ describe("prepareModelForSimpleCompletion", () => { expect(buildTransportAwareSimpleStreamFn).toHaveBeenCalledWith(model, { cfg: undefined }); expect(ensureCustomApiRegistered).toHaveBeenCalledWith( + apiRegistry, "openclaw-google-generative-ai-transport", "google-transport-stream", ); @@ -380,6 +397,7 @@ describe("prepareModelForSimpleCompletion", () => { { cfg: undefined }, ); expect(ensureCustomApiRegistered).toHaveBeenCalledWith( + apiRegistry, "openclaw-openai-responses-transport", "codex-transport-stream", ); diff --git a/src/agents/simple-completion-transport.ts b/src/agents/simple-completion-transport.ts index 6ac27fbd94f..f36d0f3bd1d 100644 --- a/src/agents/simple-completion-transport.ts +++ b/src/agents/simple-completion-transport.ts @@ -1,4 +1,4 @@ -import { getApiProvider } from "@openclaw/ai/internal/runtime"; +import type { ApiRegistry } from "@openclaw/ai"; /** * Simple completion transport preparation. * @@ -65,11 +65,15 @@ function resolveProviderSimpleCompletionApi(model: Model): Api { .join(":")}`; } -function applyProviderSimpleCompletionWrapper(model: Model, cfg?: OpenClawConfig): Model { +function applyProviderSimpleCompletionWrapper( + registry: ApiRegistry, + model: Model, + cfg?: OpenClawConfig, +): Model { if (model.api.startsWith(PROVIDER_SIMPLE_COMPLETION_API_PREFIX)) { return model; } - const sourceProvider = getApiProvider(model.api); + const sourceProvider = registry.getApiProvider(model.api); if (!sourceProvider) { return model; } @@ -93,11 +97,12 @@ function applyProviderSimpleCompletionWrapper(model: Model, cfg?: OpenClawConfig } const api = resolveProviderSimpleCompletionApi(model); - ensureCustomApiRegistered(api, streamFn); + ensureCustomApiRegistered(registry, api, streamFn); return { ...model, api }; } function prepareCodexSimpleTransportModel( + registry: ApiRegistry, model: Model, cfg?: OpenClawConfig, ): Model | undefined { @@ -117,7 +122,7 @@ function prepareCodexSimpleTransportModel( return undefined; } - ensureCustomApiRegistered(api, streamFn); + ensureCustomApiRegistered(registry, api, streamFn); return { ...transportModel, api, @@ -125,38 +130,46 @@ function prepareCodexSimpleTransportModel( } export function prepareModelForSimpleCompletion(params: { + apiRegistry: ApiRegistry; model: Model; cfg?: OpenClawConfig; }): Model { - const { model, cfg } = params; + const { apiRegistry, model, cfg } = params; // Only provider-owned custom APIs need runtime stream registration here. - if (!getApiProvider(model.api) && registerProviderStreamForModel({ model, cfg })) { - return applyProviderSimpleCompletionWrapper(model, cfg); + if ( + !apiRegistry.getApiProvider(model.api) && + registerProviderStreamForModel({ model, cfg, apiRegistry }) + ) { + return applyProviderSimpleCompletionWrapper(apiRegistry, model, cfg); } - const codexTransportModel = prepareCodexSimpleTransportModel(model, cfg); + const codexTransportModel = prepareCodexSimpleTransportModel(apiRegistry, model, cfg); if (codexTransportModel) { - return applyProviderSimpleCompletionWrapper(codexTransportModel, cfg); + return applyProviderSimpleCompletionWrapper(apiRegistry, codexTransportModel, cfg); } const transportAwareModel = prepareTransportAwareSimpleModel(model, { cfg }); if (transportAwareModel !== model) { const streamFn = buildTransportAwareSimpleStreamFn(model, { cfg }); if (streamFn) { - ensureCustomApiRegistered(transportAwareModel.api, streamFn); - return applyProviderSimpleCompletionWrapper(transportAwareModel, cfg); + ensureCustomApiRegistered(apiRegistry, transportAwareModel.api, streamFn); + return applyProviderSimpleCompletionWrapper(apiRegistry, transportAwareModel, cfg); } } if (model.api === "google-generative-ai") { - return applyProviderSimpleCompletionWrapper(prepareGoogleSimpleCompletionModel(model), cfg); + return applyProviderSimpleCompletionWrapper( + apiRegistry, + prepareGoogleSimpleCompletionModel(apiRegistry, model), + cfg, + ); } if (model.provider === "anthropic-vertex") { const api = resolveAnthropicVertexSimpleApi(model.baseUrl); - ensureCustomApiRegistered(api, createAnthropicVertexStreamFnForModel(model)); - return applyProviderSimpleCompletionWrapper({ ...model, api }, cfg); + ensureCustomApiRegistered(apiRegistry, api, createAnthropicVertexStreamFnForModel(model)); + return applyProviderSimpleCompletionWrapper(apiRegistry, { ...model, api }, cfg); } - return applyProviderSimpleCompletionWrapper(model, cfg); + return applyProviderSimpleCompletionWrapper(apiRegistry, model, cfg); } diff --git a/src/agents/tools/pdf-tool.test.ts b/src/agents/tools/pdf-tool.test.ts index 329a3be4b9d..31b54e10c23 100644 --- a/src/agents/tools/pdf-tool.test.ts +++ b/src/agents/tools/pdf-tool.test.ts @@ -12,6 +12,10 @@ import * as modelDiscovery from "../agent-model-discovery.js"; import type { AuthProfileStore } from "../auth-profiles/types.js"; import * as modelAuth from "../model-auth.js"; import * as modelsConfig from "../models-config.js"; +import { + getModelRegistryRuntime, + initializeModelRegistryRuntime, +} from "../sessions/model-registry-runtime.js"; import * as pdfNativeProviders from "./pdf-native-providers.js"; import * as pdfModelConfigModule from "./pdf-tool.model-config.js"; import { resetPdfToolAuthEnv, withTempPdfAgentDir } from "./pdf-tool.test-support.js"; @@ -19,14 +23,6 @@ import { resetPdfToolAuthEnv, withTempPdfAgentDir } from "./pdf-tool.test-suppor const completeMock = vi.hoisted(() => vi.fn()); const registerProviderStreamForModelMock = vi.hoisted(() => vi.fn()); -vi.mock("../../llm/stream.js", async () => { - const actual = await vi.importActual("../../llm/stream.js"); - return { - ...actual, - complete: completeMock, - }; -}); - vi.mock("../provider-stream.js", () => ({ registerProviderStreamForModel: registerProviderStreamForModelMock, })); @@ -148,7 +144,10 @@ async function stubPdfToolInfra( maxTokens: 8192, input: params?.input ?? ["text", "document"], }) as never; - vi.spyOn(modelDiscovery, "discoverModels").mockReturnValue({ find } as never); + const modelRegistry = { find }; + initializeModelRegistryRuntime(modelRegistry); + getModelRegistryRuntime(modelRegistry).llmRuntime.complete = completeMock; + vi.spyOn(modelDiscovery, "discoverModels").mockReturnValue(modelRegistry as never); vi.spyOn(modelsConfig, "ensureOpenClawModelsJson").mockResolvedValue({ agentDir, @@ -744,6 +743,7 @@ describe("createPdfTool", () => { expect(modelAuth.requireApiKey).not.toHaveBeenCalled(); expect(setRuntimeApiKey).not.toHaveBeenCalled(); expect(registerProviderStreamForModelMock).toHaveBeenCalledWith({ + apiRegistry: expect.anything(), model: expect.objectContaining({ provider: "amazon-bedrock", api: "bedrock-converse-stream", diff --git a/src/agents/tools/pdf-tool.ts b/src/agents/tools/pdf-tool.ts index 98616648b9b..87d682c265f 100644 --- a/src/agents/tools/pdf-tool.ts +++ b/src/agents/tools/pdf-tool.ts @@ -9,7 +9,6 @@ import { } from "@openclaw/normalization-core/string-coerce"; import { Type } from "typebox"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; -import { complete } from "../../llm/stream.js"; import type { Context } from "../../llm/types.js"; import { classifyMediaReferenceSource, @@ -23,6 +22,7 @@ import { applySecretRefHeaderSentinels } from "../model-auth.js"; import { getModelProviderRequestTransport } from "../provider-request-config.js"; import { registerProviderStreamForModel } from "../provider-stream.js"; import { optionalFiniteNumberSchema } from "../schema/typebox.js"; +import { getModelRegistryRuntime } from "../sessions/model-registry-runtime.js"; import { readFiniteNumberParam, ToolInputError } from "./common.js"; import { coerceImageModelConfig, type ImageModelConfig } from "./image-tool.helpers.js"; import { @@ -245,10 +245,12 @@ async function runPdfPrompt(params: { // PDF-only model selections may not have loaded their provider plugin yet. // Register before complete() so plugin-owned APIs resolve on first use. + const modelRuntime = getModelRegistryRuntime(modelRegistry); registerProviderStreamForModel({ model, cfg: effectiveCfg, agentDir: params.agentDir, + apiRegistry: modelRuntime.apiRegistry, ...(params.workspaceDir ? { workspaceDir: params.workspaceDir } : {}), }); @@ -266,7 +268,7 @@ async function runPdfPrompt(params: { images: [], })); const context = buildPdfExtractionContext(params.prompt, textOnlyExtractions, model); - const message = await complete(model, context, { + const message = await modelRuntime.llmRuntime.complete(model, context, { apiKey, maxTokens: resolvePdfToolMaxTokens(model.maxTokens), }); @@ -275,7 +277,7 @@ async function runPdfPrompt(params: { } const context = buildPdfExtractionContext(params.prompt, extractions, model); - const message = await complete(model, context, { + const message = await modelRuntime.llmRuntime.complete(model, context, { apiKey, maxTokens: resolvePdfToolMaxTokens(model.maxTokens), }); diff --git a/src/gateway/worker-environments/inference-runtime.test.ts b/src/gateway/worker-environments/inference-runtime.test.ts index 583feed08a7..6255fd56f6a 100644 --- a/src/gateway/worker-environments/inference-runtime.test.ts +++ b/src/gateway/worker-environments/inference-runtime.test.ts @@ -13,6 +13,7 @@ import { resolveSimpleCompletionModelResolverWorkspace } from "../../agents/simp import type { SessionEntry } from "../../config/sessions.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; import { onTrustedInternalDiagnosticEvent } from "../../infra/diagnostic-events.js"; +import { bindModelLlmRuntime } from "../../llm/model-runtime-binding.js"; import type { AssistantMessage, Model, StreamFn, Usage } from "../../llm/types.js"; import { createAssistantMessageEventStream } from "../../llm/utils/event-stream.js"; import type { loadManifestMetadataSnapshot } from "../../plugins/manifest-contract-eligibility.js"; @@ -24,6 +25,10 @@ import { } from "./inference-runtime.js"; import { createWorkerToolCallStream } from "./inference-tool-call-stream.js"; +vi.mock("../../agents/sessions/model-registry-runtime.js", () => ({ + getModelRegistryRuntime: (owner: unknown) => owner, +})); + type Deps = { applyStreamPolicy: typeof applyExtraParamsToAgent; loadCatalog: typeof loadModelCatalog; @@ -172,7 +177,6 @@ function setup(entry: SessionEntry = sessionEntry) { authProfile?: string; catalogWorkspace?: string; prepareWorkspace?: string; - registerStream?: boolean; } = {}; const resolveModel = vi.fn( async (_provider, _model, _dir, _cfg, options) => { @@ -185,8 +189,12 @@ function setup(entry: SessionEntry = sessionEntry) { modelParams.modelResolver, ); await modelParams.modelResolver?.(PROVIDER, MODEL, modelParams.agentDir, modelParams.cfg, {}); + const apiRegistry = {}; return { - model: logicalModel, + model: bindModelLlmRuntime(logicalModel, { + registry: apiRegistry, + streamSimple: fallbackStream, + } as never), auth: { apiKey: AUTH_MARKER, profileId: PROFILE, @@ -201,8 +209,7 @@ function setup(entry: SessionEntry = sessionEntry) { const loadManifestSnapshot = vi.fn( () => ({ plugins: [] }) as unknown as ReturnType, ); - const resolveProviderStream = vi.fn((streamParams) => { - scope.registerStream = streamParams.registerStream; + const resolveProviderStream = vi.fn(() => { return stream; }); const resolveStream = vi.fn((streamParams) => { @@ -238,7 +245,6 @@ function setup(entry: SessionEntry = sessionEntry) { resolveProviderStream, resolveStream, applyStreamPolicy, - stream: fallbackStream, wrapStream: vi.fn((streamFn: StreamFn) => streamFn), createTrace: vi.fn(() => ({ traceId: "1".repeat(32), spanId: "2".repeat(16) })), }; @@ -364,7 +370,6 @@ describe("worker inference provider runtime", () => { authProfile: PROFILE, catalogWorkspace: WORKSPACE, prepareWorkspace: WORKSPACE, - registerStream: false, }); const [streamModel, streamContext, streamOptions] = runtime.stream.mock.calls[0] ?? []; expect(streamModel).toMatchObject({ baseUrl: ENDPOINT }); diff --git a/src/gateway/worker-environments/inference-runtime.ts b/src/gateway/worker-environments/inference-runtime.ts index a86884e9117..ff668bcc006 100644 --- a/src/gateway/worker-environments/inference-runtime.ts +++ b/src/gateway/worker-environments/inference-runtime.ts @@ -54,14 +54,13 @@ import { freezeDiagnosticTraceContext, type DiagnosticTraceContext, } from "../../infra/diagnostic-trace-context.js"; -import { streamSimple } from "../../llm/stream.js"; +import { getModelLlmRuntime } from "../../llm/model-runtime-binding.js"; import type { AssistantMessage, AssistantMessageEvent, Context, Model, SimpleStreamOptions, - StreamFn, Tool, Usage, } from "../../llm/types.js"; @@ -110,7 +109,6 @@ type WorkerInferenceRuntimeDependencies = { resolveProviderStream: typeof registerProviderStreamForModel; resolveStream: typeof resolveEmbeddedAgentStreamFn; applyStreamPolicy: typeof applyExtraParamsToAgent; - stream: StreamFn; wrapStream: typeof wrapStreamFnWithDiagnosticModelCallEvents; createTrace: typeof createDiagnosticTraceContextFromActiveScope; recordUsage: (params: WorkerInferenceUsageParams) => void; @@ -391,7 +389,6 @@ const DEFAULT_DEPENDENCIES: WorkerInferenceRuntimeDependencies = { resolveProviderStream: registerProviderStreamForModel, resolveStream: resolveEmbeddedAgentStreamFn, applyStreamPolicy: applyExtraParamsToAgent, - stream: streamSimple as StreamFn, wrapStream: wrapStreamFnWithDiagnosticModelCallEvents, createTrace: createDiagnosticTraceContextFromActiveScope, recordUsage: emitWorkerInferenceUsage, @@ -635,6 +632,10 @@ export function createWorkerInferenceExecutor( model: approved.model, }; const logicalModel = approved.prepared.model; + const llmRuntime = getModelLlmRuntime(logicalModel); + if (!llmRuntime) { + throw new Error("Prepared worker model has no lifecycle runtime owner"); + } const providerModel = logicalModel.provider === "openai" && logicalModel.api === "openai-chatgpt-responses" ? { @@ -647,12 +648,12 @@ export function createWorkerInferenceExecutor( cfg: config, agentDir: approved.agentDir, workspaceDir: approved.workspaceDir, - registerStream: false, }); const authValue = approved.prepared.auth.apiKey; const streamAgent = { streamFn: dependencies.resolveStream({ - currentStreamFn: dependencies.stream, + llmRuntime, + currentStreamFn: llmRuntime.streamSimple, ...(providerStream ? { providerStreamFn: providerStream } : {}), sessionId: request.sessionId, signal, diff --git a/src/llm/model-runtime-binding.ts b/src/llm/model-runtime-binding.ts new file mode 100644 index 00000000000..492cc2713e8 --- /dev/null +++ b/src/llm/model-runtime-binding.ts @@ -0,0 +1,32 @@ +import type { LlmRuntime } from "@openclaw/ai"; +import type { Model } from "./types.js"; + +const MODEL_LLM_RUNTIME = Symbol("openclaw.modelLlmRuntime"); +const streamLlmRuntimes = new WeakMap(); + +type RuntimeBoundModel = Model & { + [MODEL_LLM_RUNTIME]?: LlmRuntime; +}; + +/** Carries the prepared lifecycle runtime without changing the serialized model shape. */ +export function bindModelLlmRuntime(model: Model, runtime: LlmRuntime): Model { + const bound = { ...model } as RuntimeBoundModel; + Object.defineProperty(bound, MODEL_LLM_RUNTIME, { + value: runtime, + enumerable: false, + }); + return bound; +} + +export function getModelLlmRuntime(model: Model): LlmRuntime | undefined { + return (model as RuntimeBoundModel)[MODEL_LLM_RUNTIME]; +} + +/** Associates a prepared stream entry point with the runtime that owns it. */ +export function bindStreamLlmRuntime(streamFn: object, runtime: LlmRuntime): void { + streamLlmRuntimes.set(streamFn, runtime); +} + +export function getStreamLlmRuntime(streamFn: object | undefined): LlmRuntime | undefined { + return streamFn ? streamLlmRuntimes.get(streamFn) : undefined; +} diff --git a/src/llm/stream.test.ts b/src/llm/stream.test.ts new file mode 100644 index 00000000000..0f91dc38853 --- /dev/null +++ b/src/llm/stream.test.ts @@ -0,0 +1,37 @@ +import { createApiRegistry, createLlmRuntime } from "@openclaw/ai"; +import { describe, expect, it, vi } from "vitest"; +import { bindModelLlmRuntime } from "./model-runtime-binding.js"; +import { streamSimple } from "./stream.js"; +import { createAssistantMessageEventStream } from "./utils/event-stream.js"; + +describe("LLM stream facade", () => { + it("routes a prepared model through its lifecycle runtime", () => { + const registry = createApiRegistry(); + const runtime = createLlmRuntime(registry); + const expected = createAssistantMessageEventStream(); + const stream = vi.fn(() => expected); + registry.registerApiProvider({ + api: "test-lifecycle-api", + stream, + streamSimple: stream, + }); + const model = bindModelLlmRuntime( + { + api: "test-lifecycle-api", + provider: "test-provider", + id: "test-model", + name: "Test Model", + baseUrl: "https://example.test", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1024, + maxTokens: 512, + }, + runtime, + ); + + expect(streamSimple(model, { messages: [] })).toBe(expected); + expect(stream).toHaveBeenCalledOnce(); + }); +}); diff --git a/src/llm/stream.ts b/src/llm/stream.ts index bc4e8128ef2..377a0a37478 100644 --- a/src/llm/stream.ts +++ b/src/llm/stream.ts @@ -2,10 +2,54 @@ // This facade owns the process-default AI runtime wiring: it installs the // OpenClaw host policy ports and registers built-in providers exactly once, // before any caller imports the stream API. -import { defaultApiRegistry } from "@openclaw/ai/internal/runtime"; +import { defaultApiRegistry, defaultLlmRuntime } from "@openclaw/ai/internal/runtime"; import { registerBuiltInApiProviders } from "@openclaw/ai/providers"; +import { getModelLlmRuntime } from "./model-runtime-binding.js"; import "./ai-transport-host.js"; +import type { + Api, + AssistantMessage, + AssistantMessageEventStreamContract, + Context, + Model, + ProviderStreamOptions, + SimpleStreamOptions, +} from "./types.js"; registerBuiltInApiProviders(defaultApiRegistry); -export { complete, completeSimple, stream, streamSimple } from "@openclaw/ai/internal/runtime"; +function resolveRuntime(model: Model) { + return getModelLlmRuntime(model) ?? defaultLlmRuntime; +} + +export function stream( + model: Model, + context: Context, + options?: ProviderStreamOptions, +): AssistantMessageEventStreamContract { + return resolveRuntime(model).stream(model, context, options); +} + +export function complete( + model: Model, + context: Context, + options?: ProviderStreamOptions, +): Promise { + return resolveRuntime(model).complete(model, context, options); +} + +export function streamSimple( + model: Model, + context: Context, + options?: SimpleStreamOptions, +): AssistantMessageEventStreamContract { + return resolveRuntime(model).streamSimple(model, context, options); +} + +export function completeSimple( + model: Model, + context: Context, + options?: SimpleStreamOptions, +): Promise { + return resolveRuntime(model).completeSimple(model, context, options); +} diff --git a/src/media-understanding/image-model-runtime.ts b/src/media-understanding/image-model-runtime.ts index 35f8bb53e2d..02d9fffe733 100644 --- a/src/media-understanding/image-model-runtime.ts +++ b/src/media-understanding/image-model-runtime.ts @@ -11,6 +11,8 @@ import { ensureOpenClawModelsJson } from "../agents/models-config.js"; import { resolveProviderModelMaterializationAuthMode } from "../agents/provider-model-route-auth.js"; import { protectPreparedProviderRuntimeAuth } from "../agents/provider-secret-egress.js"; import { providerUsesCredentialScopedModelMetadata } from "../agents/runtime-plan/credential-scoped-model.js"; +import { getModelRegistryRuntime } from "../agents/sessions/model-registry-runtime.js"; +import { bindModelLlmRuntime } from "../llm/model-runtime-binding.js"; import type { Model } from "../llm/types.js"; import { prepareProviderRuntimeAuth } from "../plugins/provider-runtime.runtime.js"; import type { ImageDescriptionRequest } from "./types.js"; @@ -60,6 +62,7 @@ async function prepareResolvedImageRuntime( modelRegistry: Awaited>["modelRegistry"], ): Promise<{ apiKey: string; model: Model }> { let model = resolvedModel; + const modelRuntime = getModelRegistryRuntime(modelRegistry); const apiKeyInfo = await getApiKeyForModel({ model, cfg: params.cfg, @@ -113,7 +116,13 @@ async function prepareResolvedImageRuntime( apiKeyInfo.mode === "aws-sdk" && model.api === "bedrock-converse-stream" ) { - return { apiKey: "", model: applySecretRefHeaderSentinels(model, params.cfg) }; + return { + apiKey: "", + model: bindModelLlmRuntime( + applySecretRefHeaderSentinels(model, params.cfg), + modelRuntime.llmRuntime, + ), + }; } let apiKey = requireApiKey(apiKeyInfo, model.provider); const preparedAuth = protectPreparedProviderRuntimeAuth({ @@ -142,7 +151,13 @@ async function prepareResolvedImageRuntime( model = { ...model, baseUrl: runtimeBaseUrl }; } authStorage.setRuntimeApiKey(model.provider, apiKey); - return { apiKey, model: applySecretRefHeaderSentinels(model, params.cfg) }; + return { + apiKey, + model: bindModelLlmRuntime( + applySecretRefHeaderSentinels(model, params.cfg), + modelRuntime.llmRuntime, + ), + }; } export async function resolveImageRuntime( diff --git a/src/media-understanding/image.test.ts b/src/media-understanding/image.test.ts index 613ae2cf7cd..76c6de4fc8e 100644 --- a/src/media-understanding/image.test.ts +++ b/src/media-understanding/image.test.ts @@ -153,7 +153,23 @@ vi.mock("../plugins/provider-runtime.runtime.js", () => ({ })); vi.mock("../agents/embedded-agent-runner/model.js", () => ({ - resolveModelAsync: resolveModelAsyncMock, + resolveModelAsync: async (...args: unknown[]) => { + const result = await resolveModelAsyncMock(...args); + const modelRegistry = (result?.modelRegistry ?? {}) as Record; + return { + ...result, + modelRegistry: { + ...modelRegistry, + llmRuntime: modelRegistry.llmRuntime ?? { complete: completeMock }, + }, + }; + }, +})); + +vi.mock("../agents/sessions/model-registry-runtime.js", () => ({ + getModelRegistryRuntime: (owner: { llmRuntime?: unknown }) => ({ + llmRuntime: owner.llmRuntime ?? { complete: completeMock }, + }), })); vi.mock("../plugin-sdk/provider-auth.js", () => ({ @@ -1562,7 +1578,9 @@ describe("describeImageWithModel", () => { expect(resolveModelAsyncMock.mock.calls[2]?.[4]).toEqual( expect.objectContaining({ authStorage, - modelRegistry, + modelRegistry: expect.objectContaining({ + llmRuntime: expect.anything(), + }), authProfileId: "github-copilot:backup", }), ); diff --git a/src/media-understanding/image.ts b/src/media-understanding/image.ts index e7938692b60..65a30722ba8 100644 --- a/src/media-understanding/image.ts +++ b/src/media-understanding/image.ts @@ -435,6 +435,8 @@ async function describeImagesWithModelInternal( }); } + // Resolved models carry their lifecycle runtime, so registration targets that + // registry before the built-in fallback reaches complete(). const providerStreamFn = registerProviderStreamForModel({ model, cfg: params.cfg,