From 35f807cfafbaa63a914e3d9bdbd19dadd5abd1f1 Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Sun, 17 May 2026 01:57:03 +0200 Subject: [PATCH] fix(coding-agent): route compaction through streamFn closes #4484 --- packages/coding-agent/CHANGELOG.md | 6 ++ .../coding-agent/src/core/agent-session.ts | 46 ++++++++---- .../src/core/compaction/compaction.ts | 59 ++++++++++++---- .../suite/agent-session-compaction.test.ts | 70 ++++++++++++++++++- 4 files changed, 154 insertions(+), 27 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 97de3a6e..0f86cf10 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## [Unreleased] + +### Fixed + +- Fixed compaction summary calls to use custom agent stream functions, preserving proxy-backed LLM routing ([#4484](https://github.com/earendil-works/pi/issues/4484)). + ## [0.74.1] - 2026-05-16 ### New Features diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 9868e8a6..1d51163b 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -31,6 +31,7 @@ import { isContextOverflow, modelsAreEqual, resetApiProviders, + streamSimple, } from "@earendil-works/pi-ai"; import { theme } from "../modes/interactive/theme/theme.js"; import { stripFrontmatter } from "../utils/frontmatter.js"; @@ -367,6 +368,18 @@ export class AgentSession { throw new Error(formatNoApiKeyFoundMessage(model.provider)); } + private async _getCompactionRequestAuth(model: Model): Promise<{ + apiKey?: string; + headers?: Record; + }> { + if (this.agent.streamFn === streamSimple) { + return this._getRequiredRequestAuth(model); + } + + const result = await this._modelRegistry.getApiKeyAndHeaders(model); + return result.ok ? { apiKey: result.apiKey, headers: result.headers } : {}; + } + /** * Install tool hooks once on the Agent instance. * @@ -1618,7 +1631,7 @@ export class AgentSession { throw new Error(formatNoModelSelectedMessage()); } - const { apiKey, headers } = await this._getRequiredRequestAuth(this.model); + const { apiKey, headers } = await this._getCompactionRequestAuth(this.model); const pathEntries = this.sessionManager.getBranch(); const settings = this.settingsManager.getCompactionSettings(); @@ -1676,6 +1689,7 @@ export class AgentSession { customInstructions, this._compactionAbortController.signal, this.thinkingLevel, + this.agent.streamFn, ); summary = result.summary; firstKeptEntryId = result.firstKeptEntryId; @@ -1864,18 +1878,25 @@ export class AgentSession { return; } - const authResult = await this._modelRegistry.getApiKeyAndHeaders(this.model); - if (!authResult.ok || !authResult.apiKey) { - this._emit({ - type: "compaction_end", - reason, - result: undefined, - aborted: false, - willRetry: false, - }); - return; + let apiKey: string | undefined; + let headers: Record | undefined; + if (this.agent.streamFn === streamSimple) { + const authResult = await this._modelRegistry.getApiKeyAndHeaders(this.model); + if (!authResult.ok || !authResult.apiKey) { + this._emit({ + type: "compaction_end", + reason, + result: undefined, + aborted: false, + willRetry: false, + }); + return; + } + apiKey = authResult.apiKey; + headers = authResult.headers; + } else { + ({ apiKey, headers } = await this._getCompactionRequestAuth(this.model)); } - const { apiKey, headers } = authResult; const pathEntries = this.sessionManager.getBranch(); @@ -1941,6 +1962,7 @@ export class AgentSession { undefined, this._autoCompactionAbortController.signal, this.thinkingLevel, + this.agent.streamFn, ); summary = compactResult.summary; firstKeptEntryId = compactResult.firstKeptEntryId; diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 5075e124..c12be0ee 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -5,8 +5,8 @@ * and after compaction the session is reloaded. */ -import type { AgentMessage, ThinkingLevel } from "@earendil-works/pi-agent-core"; -import type { AssistantMessage, Model, Usage } from "@earendil-works/pi-ai"; +import type { AgentMessage, StreamFn, ThinkingLevel } from "@earendil-works/pi-agent-core"; +import type { AssistantMessage, Context, Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai"; import { completeSimple } from "@earendil-works/pi-ai"; import { convertToLlm, @@ -523,6 +523,34 @@ Use this EXACT format: Keep each section concise. Preserve exact file paths, function names, and error messages.`; +function createSummarizationOptions( + model: Model, + maxTokens: number, + apiKey: string | undefined, + headers: Record | undefined, + signal: AbortSignal | undefined, + thinkingLevel: ThinkingLevel | undefined, +): SimpleStreamOptions { + const options: SimpleStreamOptions = { maxTokens, signal, apiKey, headers }; + if (model.reasoning && thinkingLevel && thinkingLevel !== "off") { + options.reasoning = thinkingLevel; + } + return options; +} + +async function completeSummarization( + model: Model, + context: Context, + options: SimpleStreamOptions, + streamFn?: StreamFn, +): Promise { + if (!streamFn) { + return completeSimple(model, context, options); + } + const stream = await streamFn(model, context, options); + return stream.result(); +} + /** * Generate a summary of the conversation using the LLM. * If previousSummary is provided, uses the update prompt to merge. @@ -531,12 +559,13 @@ export async function generateSummary( currentMessages: AgentMessage[], model: Model, reserveTokens: number, - apiKey: string, + apiKey: string | undefined, headers?: Record, signal?: AbortSignal, customInstructions?: string, previousSummary?: string, thinkingLevel?: ThinkingLevel, + streamFn?: StreamFn, ): Promise { const maxTokens = Math.min( Math.floor(0.8 * reserveTokens), @@ -569,15 +598,13 @@ export async function generateSummary( }, ]; - const completionOptions = - model.reasoning && thinkingLevel && thinkingLevel !== "off" - ? { maxTokens, signal, apiKey, headers, reasoning: thinkingLevel } - : { maxTokens, signal, apiKey, headers }; + const completionOptions = createSummarizationOptions(model, maxTokens, apiKey, headers, signal, thinkingLevel); - const response = await completeSimple( + const response = await completeSummarization( model, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, completionOptions, + streamFn, ); if (response.stopReason === "error") { @@ -720,11 +747,12 @@ Be concise. Focus on what's needed to understand the kept suffix.`; export async function compact( preparation: CompactionPreparation, model: Model, - apiKey: string, + apiKey: string | undefined, headers?: Record, customInstructions?: string, signal?: AbortSignal, thinkingLevel?: ThinkingLevel, + streamFn?: StreamFn, ): Promise { const { firstKeptEntryId, @@ -754,6 +782,7 @@ export async function compact( customInstructions, previousSummary, thinkingLevel, + streamFn, ) : Promise.resolve("No prior history."), generateTurnPrefixSummary( @@ -764,6 +793,7 @@ export async function compact( headers, signal, thinkingLevel, + streamFn, ), ]); // Merge into single summary @@ -780,6 +810,7 @@ export async function compact( customInstructions, previousSummary, thinkingLevel, + streamFn, ); } @@ -806,10 +837,11 @@ async function generateTurnPrefixSummary( messages: AgentMessage[], model: Model, reserveTokens: number, - apiKey: string, + apiKey: string | undefined, headers?: Record, signal?: AbortSignal, thinkingLevel?: ThinkingLevel, + streamFn?: StreamFn, ): Promise { const maxTokens = Math.min( Math.floor(0.5 * reserveTokens), @@ -826,12 +858,11 @@ async function generateTurnPrefixSummary( }, ]; - const response = await completeSimple( + const response = await completeSummarization( model, { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages }, - model.reasoning && thinkingLevel && thinkingLevel !== "off" - ? { maxTokens, signal, apiKey, headers, reasoning: thinkingLevel } - : { maxTokens, signal, apiKey, headers }, + createSummarizationOptions(model, maxTokens, apiKey, headers, signal, thinkingLevel), + streamFn, ); if (response.stopReason === "error") { diff --git a/packages/coding-agent/test/suite/agent-session-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-compaction.test.ts index cb9f1808..e2d28911 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -1,4 +1,9 @@ -import { type AssistantMessage, fauxAssistantMessage, type Model } from "@earendil-works/pi-ai"; +import { + type AssistantMessage, + createAssistantMessageEventStream, + fauxAssistantMessage, + type Model, +} from "@earendil-works/pi-ai"; import { afterEach, describe, expect, it, vi } from "vitest"; import { createHarness, type Harness } from "./harness.js"; @@ -41,6 +46,43 @@ function createAssistant( }; } +function useSummaryStreamFn(harness: Harness, summary: string): () => number { + let callCount = 0; + harness.session.agent.streamFn = (model) => { + callCount++; + const stream = createAssistantMessageEventStream(); + queueMicrotask(() => { + const message: AssistantMessage = { + ...fauxAssistantMessage(summary), + api: model.api, + provider: model.provider, + model: model.id, + usage: createUsage(10), + }; + stream.push({ type: "done", reason: "stop", message }); + }); + return stream; + }; + return () => callCount; +} + +function seedCompactableSession(harness: Harness): void { + const now = Date.now(); + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "message to compact" }], + timestamp: now - 1000, + }); + harness.sessionManager.appendMessage( + createAssistant(harness, { + stopReason: "stop", + totalTokens: 100, + timestamp: now - 500, + }), + ); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; +} + describe("AgentSession compaction characterization", () => { const harnesses: Harness[] = []; @@ -95,6 +137,32 @@ describe("AgentSession compaction characterization", () => { await expect(harness.session.compact()).rejects.toThrow(`No API key found for ${harness.getModel().provider}.`); }); + it("manually compacts with a custom streamFn when registry auth is absent", async () => { + const harness = await createHarness({ withConfiguredAuth: false }); + harnesses.push(harness); + seedCompactableSession(harness); + const getStreamCallCount = useSummaryStreamFn(harness, "summary from custom stream"); + + const result = await harness.session.compact(); + + expect(result.summary).toBe("summary from custom stream"); + expect(getStreamCallCount()).toBe(1); + }); + + it("auto-compacts with a custom streamFn when registry auth is absent", async () => { + const harness = await createHarness({ withConfiguredAuth: false }); + harnesses.push(harness); + seedCompactableSession(harness); + const getStreamCallCount = useSummaryStreamFn(harness, "auto summary from custom stream"); + const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; + + await sessionInternals._runAutoCompaction("threshold", false); + + const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); + expect(compactionEntries).toHaveLength(1); + expect(getStreamCallCount()).toBe(1); + }); + it("cancels in-progress manual compaction when abortCompaction is called", async () => { const harness = await createHarness({ extensionFactories: [