Merge pull request #2543 from splichy/fix/provider-scoped-stream-routing

fix(model-registry): scope custom provider stream handlers to prevent clobbering built-ins

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<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);
+	});
+});

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<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}`,
 			);