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
This commit is contained in:
parent
55c8988900
commit
d56842ab7a
2 changed files with 99 additions and 2 deletions
|
|
@ -572,3 +572,73 @@ describe("ModelRegistry authMode — streamSimple apiKey boundary", () => {
|
|||
assert.equal((captured as Record<string, unknown>).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);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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<Api>, 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<Api>, context: Context, options?: SimpleStreamOptions) =>
|
||||
scopedStream(model, context, options as SimpleStreamOptions);
|
||||
const scopedFullStream = existingProvider
|
||||
? (model: Model<Api>, context: Context, options?: Record<string, unknown>) => {
|
||||
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}`,
|
||||
);
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue