From d56842ab7a5ac98860451188bf3ef99f975d3960 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vojt=C4=9Bch=20=C5=A0pl=C3=ADchal?= Date: Wed, 25 Mar 2026 22:32:00 +0100 Subject: [PATCH] fix(model-registry): scope custom provider stream handlers to prevent clobbering built-in API handlers When a custom provider (e.g. claude-code-cli) registers a streamSimple handler with the same api type as a built-in (e.g. 'anthropic-messages'), the global API provider registry was overwritten, routing ALL models of that api type through the custom handler. This caused anthropic/claude-opus-4-6 requests to be dispatched through the Claude Code SDK subprocess instead of the Anthropic API, resulting in 'Tool not found' errors for Glob, Read, Edit, Bash (SDK tool names not present in pi's tool registry). Fix: wrap the registered handler with a model.provider guard so it only fires for models from the registering provider, delegating to the previous handler for all other providers. Closes #2536 --- .../src/core/model-registry-auth-mode.test.ts | 70 +++++++++++++++++++ .../src/core/model-registry.ts | 31 +++++++- 2 files changed, 99 insertions(+), 2 deletions(-) diff --git a/packages/pi-coding-agent/src/core/model-registry-auth-mode.test.ts b/packages/pi-coding-agent/src/core/model-registry-auth-mode.test.ts index 66f88fa86..be27f6c60 100644 --- a/packages/pi-coding-agent/src/core/model-registry-auth-mode.test.ts +++ b/packages/pi-coding-agent/src/core/model-registry-auth-mode.test.ts @@ -572,3 +572,73 @@ describe("ModelRegistry authMode — streamSimple apiKey boundary", () => { assert.equal((captured as Record).reasoning, "high", "reasoning must pass through"); }); }); + +// ─── Provider-scoped stream routing (#2533) ─────────────────────────────────── + +describe("ModelRegistry authMode — provider-scoped stream routing", () => { + it("does not clobber built-in stream handler when custom provider uses same api", () => { + const registry = createRegistry(() => true); + const customSpy = createStreamSpy(); + + // Register a custom provider with the same API type as a built-in (anthropic-messages). + // This simulates the claude-code-cli extension registering with api: "anthropic-messages". + registry.registerProvider("custom-cli", { + authMode: "externalCli", + baseUrl: "local://custom", + api: "anthropic-messages", + streamSimple: customSpy.streamSimple, + models: [createProviderModel("custom-model", "anthropic-messages")], + }); + + // The built-in anthropic-messages provider should still be accessible + // when calling streamSimple with a model from the built-in provider. + const provider = getApiProvider("anthropic-messages" as Api); + assert.ok(provider, "anthropic-messages provider must still be registered"); + + // Call with a built-in anthropic model — should NOT hit the custom spy. + // The built-in handler will throw (no API key), which proves the routing + // correctly delegates to the built-in instead of the custom handler. + assert.throws( + () => provider.streamSimple( + makeModel("anthropic", "claude-sonnet-4-6", "anthropic-messages"), + makeContext(), + { maxTokens: 4096 } as SimpleStreamOptions, + ), + (err: Error) => err.message.includes("API key"), + "built-in Anthropic handler must be invoked (throws because no API key in tests)", + ); + + assert.equal( + customSpy.getCapturedOptions(), + undefined, + "custom provider's streamSimple must NOT be called for anthropic provider models", + ); + }); + + it("routes to custom provider when model.provider matches", () => { + const registry = createRegistry(() => true); + const customSpy = createStreamSpy(); + + registry.registerProvider("custom-cli", { + authMode: "externalCli", + baseUrl: "local://custom", + api: "anthropic-messages", + streamSimple: customSpy.streamSimple, + models: [createProviderModel("custom-model", "anthropic-messages")], + }); + + const provider = getApiProvider("anthropic-messages" as Api); + assert.ok(provider); + + // Call with the custom provider's model — should hit the custom spy + provider.streamSimple( + makeModel("custom-cli", "custom-model", "anthropic-messages"), + makeContext(), + { maxTokens: 2048 } as SimpleStreamOptions, + ); + + const captured = customSpy.getCapturedOptions(); + assert.ok(captured, "custom provider's streamSimple must be called for its own models"); + assert.equal(captured.maxTokens, 2048); + }); +}); diff --git a/packages/pi-coding-agent/src/core/model-registry.ts b/packages/pi-coding-agent/src/core/model-registry.ts index d68778a0e..9a92cd1b7 100644 --- a/packages/pi-coding-agent/src/core/model-registry.ts +++ b/packages/pi-coding-agent/src/core/model-registry.ts @@ -6,6 +6,7 @@ import { type Api, type AssistantMessageEventStream, type Context, + getApiProvider, getModels, getProviders, type KnownProvider, @@ -635,11 +636,37 @@ export class ModelRegistry { }) : rawStreamSimple; + // Guard: if there's already a handler registered for this API, wrap + // the new one so it only fires for models from this provider and + // delegates to the previous handler for all other providers. Without + // this, a custom provider using api:"anthropic-messages" would clobber + // the built-in Anthropic stream handler (#2536). + const existingProvider = getApiProvider(config.api as Api); + const scopedStream = existingProvider + ? (model: Model, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream => { + if (model.provider === providerName) { + return streamSimple(model, context, options); + } + return existingProvider.streamSimple(model, context, options); + } + : streamSimple; + + const newFullStream = (model: Model, context: Context, options?: SimpleStreamOptions) => + scopedStream(model, context, options as SimpleStreamOptions); + const scopedFullStream = existingProvider + ? (model: Model, context: Context, options?: Record) => { + if (model.provider === providerName) { + return newFullStream(model, context, options as SimpleStreamOptions); + } + return existingProvider.stream(model, context, options); + } + : newFullStream; + registerApiProvider( { api: config.api, - stream: (model, context, options) => streamSimple(model, context, options as SimpleStreamOptions), - streamSimple, + stream: scopedFullStream as any, + streamSimple: scopedStream, }, `provider:${providerName}`, );