fix(coding-agent): allow extensions to stream from custom providers (#9272)

- fixes #8964

packages/coding-agent/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+- Added `ctx.modelRegistry.stream()` and `streamSimple()` for extension model calls through configured providers with resolved authentication ([#8964](https://github.com/earendil-works/pi/issues/8964)).
+
 ### Changed
 
 - Enabled strict-prefer JSON-schema sampling by default for built-in `read`, `bash`, `powershell`, `edit`, and `write` tools, without requiring `PI_EXPERIMENTAL`. Extensions can re-register tool definitions with `constrainedSampling: false`.

packages/coding-agent/docs/extensions.md

@@ -1016,6 +1016,12 @@ Access to models, providers, and resolved authentication. `ctx.modelRegistry.get
 
 `ctx.scopedModels` is the read-only list of models scoped to the current session — the same set the `/scoped-models` command shows. It is resolved at session start from the `--models` CLI flag and the `enabledModels` setting (matched against the available catalogue with minimatch on `provider/modelId` or a bare `modelId`). It is empty when no scoping is configured, meaning every available model is usable. Each entry is `{ model, thinkingLevel? }`, where `thinkingLevel` is set only when a pattern pinned it (e.g. `anthropic/*:high`). Use it to populate a model picker that mirrors the built-in one instead of enumerating the whole catalogue via `ctx.modelRegistry.getAvailable()`.
 
+#### Streaming model calls
+
+Use `ctx.modelRegistry.streamSimple(model, context, options)` for provider-neutral options such as `reasoning`, or `stream()` for API-specific options. Both use configured providers and resolve authentication, including for providers registered with `pi.registerProvider()`. Use these instead of `pi-ai/compat` streaming functions, which cannot see extension provider registrations.
+
+Both return an `AssistantMessageEventStream`. Iterate it for response events and await `.result()` for the final message. Setup failures produce error events and error results.
+
 ### ctx.signal
 
 The current agent abort signal, or `undefined` when no agent turn is active.

packages/coding-agent/src/core/model-registry.ts

@@ -1,12 +1,14 @@
 import type {
 	Api,
 	AssistantMessage,
+	AssistantMessageEventStream,
 	AuthResult,
 	Context,
 	Model,
 	ModelsApiStreamOptions,
 	ModelsRefreshOptions,
 	ModelsRefreshResult,
+	ModelsSimpleStreamOptions,
 	Provider,
 	ProviderHeaders,
 } from "@earendil-works/pi-ai";
@@ -100,6 +102,20 @@ export class ModelRegistry {
 		return this.runtime.getProvider(provider);
 	}
 
+	/** Stream through the configured provider with request-time authentication. */
+	stream<TApi extends Api>(
+		model: Model<TApi>,
+		context: Context,
+		options?: ModelsApiStreamOptions<TApi>,
+	): AssistantMessageEventStream {
+		return this.runtime.stream(model, context, options);
+	}
+
+	/** Stream with provider-neutral options and request-time authentication. */
+	streamSimple(model: Model<Api>, context: Context, options?: ModelsSimpleStreamOptions): AssistantMessageEventStream {
+		return this.runtime.streamSimple(model, context, options);
+	}
+
 	complete<TApi extends Api>(
 		model: Model<TApi>,
 		context: Context,

packages/coding-agent/test/suite/regressions/8964-extension-provider-streaming.test.ts

@@ -0,0 +1,63 @@
+import { type AssistantMessage, fauxAssistantMessage, fauxProvider } from "@earendil-works/pi-ai";
+import { getApiProvider } from "@earendil-works/pi-ai/compat";
+import { expect, it } from "vitest";
+import { createHarness } from "../harness.ts";
+
+// Regression for #8964: extensions can stream responses from providers registered with pi.registerProvider().
+it.each(["stream", "streamSimple"] as const)(
+	"allows an extension command to use ctx.modelRegistry.%s",
+	async (method) => {
+		const faux = fauxProvider({ provider: "extension-provider", api: "issue-8964-extension-api" });
+		let receivedApiKey: string | undefined;
+		faux.setResponses([
+			(_context, options) => {
+				receivedApiKey = options?.apiKey;
+				return fauxAssistantMessage("custom provider response");
+			},
+		]);
+		let streamedText = "";
+		let result: AssistantMessage | undefined;
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.registerProvider(faux.provider.id, {
+						api: faux.api,
+						baseUrl: faux.getModel().baseUrl,
+						apiKey: "extension-key",
+						models: faux.models,
+						streamSimple: faux.provider.streamSimple,
+					});
+				},
+				(pi) => {
+					pi.registerCommand("stream-custom", {
+						description: "Stream a response from the custom provider",
+						handler: async (_args, ctx) => {
+							const model = ctx.modelRegistry.find(faux.provider.id, faux.getModel().id)!;
+							const stream = ctx.modelRegistry[method](model, {
+								messages: [{ role: "user", content: "Hello", timestamp: Date.now() }],
+							});
+							for await (const event of stream) {
+								if (event.type === "text_delta") streamedText += event.delta;
+							}
+							result = await stream.result();
+						},
+					});
+				},
+			],
+		});
+		try {
+			expect(getApiProvider(faux.api)).toBeUndefined();
+			await harness.session.prompt("/stream-custom");
+
+			expect(receivedApiKey).toBe("extension-key");
+			expect(streamedText).toBe("custom provider response");
+			expect(result).toMatchObject({
+				stopReason: "stop",
+				content: [{ type: "text", text: "custom provider response" }],
+			});
+			expect(getApiProvider(faux.api)).toBeUndefined();
+		} finally {
+			harness.cleanup();
+		}
+	},
+);