diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 9fed042ab..f1d1b9aab 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/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`. diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index ce13db622..4eed3ca71 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/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. diff --git a/packages/coding-agent/src/core/model-registry.ts b/packages/coding-agent/src/core/model-registry.ts index cfb8dbf4e..bdd2d9938 100644 --- a/packages/coding-agent/src/core/model-registry.ts +++ b/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( + model: Model, + context: Context, + options?: ModelsApiStreamOptions, + ): AssistantMessageEventStream { + return this.runtime.stream(model, context, options); + } + + /** Stream with provider-neutral options and request-time authentication. */ + streamSimple(model: Model, context: Context, options?: ModelsSimpleStreamOptions): AssistantMessageEventStream { + return this.runtime.streamSimple(model, context, options); + } + complete( model: Model, context: Context, diff --git a/packages/coding-agent/test/suite/regressions/8964-extension-provider-streaming.test.ts b/packages/coding-agent/test/suite/regressions/8964-extension-provider-streaming.test.ts new file mode 100644 index 000000000..30ad50dae --- /dev/null +++ b/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(); + } + }, +);