From 7ef472f265cdfe8502c25035ccf41a75aa28f01e Mon Sep 17 00:00:00 2001 From: dawn <90008@klbr.net> Date: Wed, 16 Sep 2026 07:54:54 +0300 Subject: [PATCH] feat(agent): resolve the loop model before each provider request AgentLoopConfig.getModel lets a caller retarget a loop that is already running. Agent now passes its live state model, so an extension that changes the model during a turn affects the next request instead of waiting for the next run. --- .../.changes/per-turn-model-resolution.md | 1 + packages/agent/src/agent-loop.ts | 30 ++++++++----- packages/agent/src/agent.ts | 1 + packages/agent/src/types.ts | 8 ++++ packages/agent/test/agent-loop.test.ts | 44 +++++++++++++++++++ 5 files changed, 72 insertions(+), 12 deletions(-) create mode 100644 packages/agent/.changes/per-turn-model-resolution.md diff --git a/packages/agent/.changes/per-turn-model-resolution.md b/packages/agent/.changes/per-turn-model-resolution.md new file mode 100644 index 000000000..6046d75ee --- /dev/null +++ b/packages/agent/.changes/per-turn-model-resolution.md @@ -0,0 +1 @@ +- Added `AgentLoopConfig.getModel` so a loop resolves its model before each provider request, which lets an extension retarget a running loop between turns. diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index e11bbd663..990397376 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -8,6 +8,7 @@ import { type AssistantMessageEvent, type Context, EventStream, + type Model, streamSimple, type ToolResultMessage, validateToolArguments, @@ -129,16 +130,21 @@ function cloneUsage(usage: AssistantMessage["usage"]): AssistantMessage["usage"] return { ...usage, cost: { ...usage.cost } }; } -function createAbortedAssistantMessage( - config: AgentLoopConfig, - partialMessage: AssistantMessage | null, -): AssistantMessage { +/** + * The model for this request. A loop that resolves its model per call can be + * retargeted between turns; the snapshot stays the fallback. + */ +function resolveRequestModel(config: AgentLoopConfig): Model { + return config.getModel?.() ?? config.model; +} + +function createAbortedAssistantMessage(model: Model, partialMessage: AssistantMessage | null): AssistantMessage { return { role: "assistant", content: partialMessage ? cloneAssistantContent(partialMessage.content) : [{ type: "text", text: "" }], - api: partialMessage?.api ?? config.model.api, - provider: partialMessage?.provider ?? config.model.provider, - model: partialMessage?.model ?? config.model.id, + api: partialMessage?.api ?? model.api, + provider: partialMessage?.provider ?? model.provider, + model: partialMessage?.model ?? model.id, usage: cloneUsage(partialMessage?.usage ?? EMPTY_USAGE), stopReason: "aborted", errorMessage: ABORT_ERROR_MESSAGE, @@ -455,10 +461,11 @@ async function streamAssistantResponse( emit: AgentEventSink, streamFn?: StreamFn, ): Promise { + const model = resolveRequestModel(config); let partialMessage: AssistantMessage | null = null; let addedPartial = false; const finishAbortedMessage = async () => { - const finalMessage = createAbortedAssistantMessage(config, partialMessage); + const finalMessage = createAbortedAssistantMessage(model, partialMessage); if (addedPartial) { context.messages[context.messages.length - 1] = finalMessage; } else { @@ -481,9 +488,8 @@ async function streamAssistantResponse( const streamFunction = streamFn || streamSimple; const resolvedApiKey = - (config.getApiKey - ? await maybePromiseWithAbort(config.getApiKey(config.model.provider), signal) - : undefined) || config.apiKey; + (config.getApiKey ? await maybePromiseWithAbort(config.getApiKey(model.provider), signal) : undefined) || + config.apiKey; const llmContext: Context = { systemPrompt: config.getSystemPrompt?.() ?? context.systemPrompt, @@ -492,7 +498,7 @@ async function streamAssistantResponse( }; const response = await maybePromiseWithAbort( - streamFunction(config.model, llmContext, { + streamFunction(model, llmContext, { ...config, apiKey: resolvedApiKey, signal, diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index d13c9e16c..e9790927b 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -466,6 +466,7 @@ export class Agent { let skipInitialSteeringPoll = options.skipInitialSteeringPoll === true; return { model: this._state.model, + getModel: () => this._state.model, reasoning: this._state.thinkingLevel, serviceTier: this._state.serviceTier, sessionId: this.sessionId, diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 1d74a46cf..4a0ab73e9 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -171,6 +171,14 @@ export interface AgentLoopConfig extends SimpleStreamOptions { /** Resolves the system prompt immediately before each LLM call. */ getSystemPrompt?: () => string; + /** + * Resolves the model immediately before each LLM call. + * + * Lets a caller retarget a loop that is already running, for example when an + * extension routes every turn by policy. Defaults to `model`. + */ + getModel?: () => Model; + /** * Resolves an API key dynamically for each LLM call. * diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index a4f0648ff..e792c8ed8 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -2121,3 +2121,47 @@ describe("agentLoopContinue with AgentMessage", () => { ]); }); }); + +describe("AgentLoopConfig.getModel", () => { + it("retargets the model between turns of one loop", async () => { + const first = createModel(); + const second: Model<"openai-responses"> = { ...createModel(), id: "mock-2", provider: "other" }; + let selected = first; + const models: string[] = []; + const tool: AgentTool = { + name: "wait", + label: "Wait", + description: "Wait", + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text" as const, text: "waited" }], details: {} }), + }; + const context: AgentContext = { + systemPrompt: "You are helpful.", + messages: [], + tools: [tool], + }; + const config: AgentLoopConfig = { + model: first, + getModel: () => selected, + convertToLlm: identityConverter, + }; + let turn = 0; + const streamFn = (model: Model) => { + models.push(`${model.provider}/${model.id}`); + const stream = new MockAssistantStream(); + const message = + turn++ === 0 + ? createAssistantMessage([{ type: "toolCall", id: "tool_1", name: "wait", arguments: {} }], "toolUse") + : createAssistantMessage([{ type: "text", text: "done" }]); + if (turn === 1) selected = second; + queueMicrotask(() => { + stream.push({ type: "done", reason: message.stopReason === "toolUse" ? "toolUse" : "stop", message }); + }); + return stream; + }; + + await runAgentLoop([createUserMessage("Hello")], context, config, () => {}, undefined, streamFn as never); + + expect(models).toEqual(["openai/mock", "other/mock-2"]); + }); +}); -- 2.51.2