From edaaec398cc42c767355fddac71900614b7bcafa Mon Sep 17 00:00:00 2001 From: roboomp Date: Sun, 28 Jun 2026 20:57:13 +0000 Subject: [PATCH] fix(task): routed side requests through provider cap Passed the provider-capped stream wrapper into AgentSession side-channel requests so /btw, /omfg, IRC auto-replies, and handoff generation share the same per-provider concurrency limit as normal turns. Added focused coverage for runEphemeralTurn and handoff generation using the configured side stream function. --- packages/agent/src/compaction/compaction.ts | 9 ++- packages/coding-agent/src/sdk.ts | 15 ++-- .../coding-agent/src/session/agent-session.ts | 15 +++- .../test/agent-session-handoff.test.ts | 73 ++++++++++++++++++- .../agent-session-message-pipeline.test.ts | 57 ++++++++++++++- 5 files changed, 158 insertions(+), 11 deletions(-) diff --git a/packages/agent/src/compaction/compaction.ts b/packages/agent/src/compaction/compaction.ts index 866ef3b6e..4bac74dfb 100644 --- a/packages/agent/src/compaction/compaction.ts +++ b/packages/agent/src/compaction/compaction.ts @@ -6,6 +6,7 @@ */ import { + type Api, type ApiKey, type AssistantMessage, type Context, @@ -828,6 +829,12 @@ export interface HandoffFromContextOptions { * anything provided here. */ streamOptions: SimpleStreamOptions; + /** Optional completion transport override for host-level request wrappers. */ + completeImpl?: ( + model: Model, + ctx: Context, + options: SimpleStreamOptions, + ) => Promise; /** See {@link HandoffOptions.telemetry}. */ telemetry?: AgentTelemetry; /** See {@link HandoffOptions.thinkingLevel}. */ @@ -859,7 +866,7 @@ export async function generateHandoffFromContext( reasoning: resolveCompactionEffort(model, options.thinkingLevel), toolChoice: "none", }, - { telemetry: options.telemetry, oneshotKind: "handoff" }, + { telemetry: options.telemetry, oneshotKind: "handoff", completeImpl: options.completeImpl }, ); if (response.stopReason === "error") { diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index c3972bd2b..5789a45dd 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -2536,13 +2536,13 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} // One-shot launch-latency marker: fired the first time the loop dispatches // a chat request to the provider transport. See onFirstChatDispatch. let notifyFirstChatDispatch = options.onFirstChatDispatch; - // Shared, settings-aware stream wrapper used by both the main agent and - // the advisor (via AgentSessionConfig.streamFn). Keeps OpenRouter - // sticky-routing variants, antigravity endpoint routing, in-flight caps, - // and the loop guard consistent across every agent the session drives. - // Wrapped in a per-provider concurrency limiter so each LLM HTTP request - // — not the whole subagent lifecycle — holds the slot, preventing the - // nested-spawn deadlock from issue #3749. + // Shared, settings-aware stream wrapper used by the main agent, advisor, + // and side-channel requests (`/btw`, `/omfg`, IRC auto-replies, handoff). + // Keeps OpenRouter sticky-routing variants, antigravity endpoint routing, + // in-flight caps, and the loop guard consistent across every provider call + // the session drives. Wrapped in a per-provider concurrency limiter so + // each LLM HTTP request — not the whole subagent lifecycle — holds the + // slot, preventing the nested-spawn deadlock from issue #3749. const settingsAwareStreamFn = wrapStreamFnWithProviderConcurrency( settings, createSettingsAwareStreamFn(settings), @@ -2716,6 +2716,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} transformProviderContext, onPayload, onResponse, + sideStreamFn: settingsAwareStreamFn, advisorStreamFn: settingsAwareStreamFn, convertToLlm: convertToLlmFinal, rebuildSystemPrompt, diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index a40ce2414..da99737a9 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -534,6 +534,13 @@ export interface AgentSessionConfig { * inherits this so its requests undergo the same shaping as the main turn. */ transformProviderContext?: (context: Context, model: Model) => Context | Promise; + /** + * Stream wrapper passed to side-channel requests (`/btw`, `/omfg`, IRC + * auto-replies, and handoff generation) so they apply the same provider + * shaping and host-level request wrappers as normal agent turns. Defaults + * to plain `streamSimple` when omitted. + */ + sideStreamFn?: StreamFn; /** * Stream wrapper passed to the advisor agent so its requests apply the * session's `providers.openrouterVariant`, `providers.antigravityEndpoint`, @@ -1459,6 +1466,7 @@ export class AgentSession { #onResponse: SimpleStreamOptions["onResponse"] | undefined; #onSseEvent: SimpleStreamOptions["onSseEvent"] | undefined; #transformProviderContext: ((context: Context, model: Model) => Context | Promise) | undefined; + #sideStreamFn: StreamFn; #advisorStreamFn: StreamFn | undefined; #convertToLlm: (messages: AgentMessage[]) => Message[] | Promise; #rebuildSystemPrompt: @@ -1788,6 +1796,7 @@ export class AgentSession { this.#requestedToolNames = config.requestedToolNames; this.#transformContext = config.transformContext ?? (messages => messages); this.#transformProviderContext = config.transformProviderContext; + this.#sideStreamFn = config.sideStreamFn ?? streamSimple; this.#advisorStreamFn = config.advisorStreamFn; this.#onPayload = config.onPayload; this.rawSseDebugBuffer = config.rawSseDebugBuffer ?? new RawSseDebugBuffer(); @@ -9195,6 +9204,10 @@ export class AgentSession { model, { streamOptions: handoffStreamOptions, + completeImpl: async (requestModel, requestContext, requestOptions) => { + const stream = await this.#sideStreamFn(requestModel, requestContext, requestOptions); + return stream.result(); + }, telemetry: resolveTelemetry(this.agent.telemetry, this.sessionId), // Honor the user's /model thinking selection on the handoff path. // Clamped per-model inside generateHandoffFromContext via @@ -12923,7 +12936,7 @@ export class AgentSession { let providerReplyText = ""; let emittedReplyText = ""; let assistantMessage: AssistantMessage | undefined; - const stream = streamSimple(model, obfuscateProviderContext(this.#obfuscator, context), options); + const stream = await this.#sideStreamFn(model, obfuscateProviderContext(this.#obfuscator, context), options); for await (const event of stream) { if (event.type === "text_delta") { providerReplyText += event.delta; diff --git a/packages/coding-agent/test/agent-session-handoff.test.ts b/packages/coding-agent/test/agent-session-handoff.test.ts index 8e06d90d0..52251ec70 100644 --- a/packages/coding-agent/test/agent-session-handoff.test.ts +++ b/packages/coding-agent/test/agent-session-handoff.test.ts @@ -1,9 +1,10 @@ import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "bun:test"; import * as path from "node:path"; -import { Agent, type AgentMessage } from "@oh-my-pi/pi-agent-core"; +import { Agent, type AgentMessage, type StreamFn } from "@oh-my-pi/pi-agent-core"; import * as compactionModule from "@oh-my-pi/pi-agent-core/compaction"; import type { AssistantMessage, Model, ToolCall } from "@oh-my-pi/pi-ai"; import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock"; +import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream"; import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; @@ -156,6 +157,76 @@ describe("AgentSession handoff", () => { expect(sessionManager.getEntries().filter(entry => entry.type === "compaction")).toHaveLength(0); }); + it("runs handoff generation through the configured side stream function", async () => { + const handoffText = "## Goal\nContinue via side stream"; + let sideStreamCalls = 0; + let capturedSideSessionId: string | undefined; + const sideStreamFn: StreamFn = (requestModel, _context, options) => { + sideStreamCalls++; + capturedSideSessionId = options?.sessionId; + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + const message: AssistantMessage = { + role: "assistant", + content: [{ type: "text", text: handoffText }], + api: requestModel.api, + provider: requestModel.provider, + model: requestModel.id, + stopReason: "stop", + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + timestamp: Date.now(), + }; + stream.push({ type: "done", reason: "stop", message }); + }); + return stream; + }; + await session.dispose(); + session = new AgentSession({ + agent: new Agent({ + initialState: { + model, + systemPrompt: ["Test"], + tools: [], + messages: [], + }, + }), + sessionManager, + settings: Settings.isolated({ + "compaction.enabled": true, + "compaction.autoContinue": false, + }), + modelRegistry, + obfuscator, + sideStreamFn, + }); + const preHandoffSessionId = session.sessionId; + + const generateHandoffSpy = vi + .spyOn(compactionModule, "generateHandoffFromContext") + .mockImplementation(async (context, requestModel, options) => { + expect(options.completeImpl).toBeDefined(); + const message = await options.completeImpl!(requestModel, context, options.streamOptions); + return message.content + .filter(block => block.type === "text") + .map(block => block.text) + .join("\n"); + }); + + const result = await session.handoff(); + + expect(generateHandoffSpy).toHaveBeenCalledTimes(1); + expect(result?.document).toBe(handoffText); + expect(sideStreamCalls).toBe(1); + expect(capturedSideSessionId).toStartWith(`${preHandoffSessionId}:side:`); + }); + it("preserves queued steering and follow-up messages across the handoff reset", async () => { // Defect 2: handoff() calls agent.reset(), which clears the core steering/follow-up // queues. Steers/follow-ups already queued (the mis-routed first compaction message, diff --git a/packages/coding-agent/test/agent-session-message-pipeline.test.ts b/packages/coding-agent/test/agent-session-message-pipeline.test.ts index 4c1d4db3a..5c47c74b4 100644 --- a/packages/coding-agent/test/agent-session-message-pipeline.test.ts +++ b/packages/coding-agent/test/agent-session-message-pipeline.test.ts @@ -1,5 +1,11 @@ import { afterEach, describe, expect, it, vi } from "bun:test"; -import { Agent, type AgentMessage, type AgentTool, AppendOnlyContextManager } from "@oh-my-pi/pi-agent-core"; +import { + Agent, + type AgentMessage, + type AgentTool, + AppendOnlyContextManager, + type StreamFn, +} from "@oh-my-pi/pi-agent-core"; import { type Api, type Context, @@ -278,6 +284,55 @@ describe("AgentSession message pipeline", () => { expect(capturedOptions?.preferWebsockets).toBe(false); }); + it("runs ephemeral side-channel requests through the configured side stream function", async () => { + const model = buildModel({ + id: "side-stream-model", + name: "Side Stream Model", + api: "anthropic", + provider: "test-provider", + baseUrl: "", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 4096, + maxTokens: 1024, + } as ModelSpec) as Model; + let capturedOptions: SimpleStreamOptions | undefined; + let capturedContext: Context | undefined; + const sideStreamFn: StreamFn = (_model, context, options) => { + capturedContext = context; + capturedOptions = options; + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + const message = createAssistantMessage("Side answer"); + stream.push({ type: "text_delta", contentIndex: 0, delta: "Side answer", partial: message }); + stream.push({ type: "done", reason: "stop", message }); + }); + return stream; + }; + const session = new AgentSession({ + agent: new Agent({ + initialState: { + model, + systemPrompt: ["system prompt"], + messages: [], + tools: [], + }, + }), + sessionManager: SessionManager.inMemory(), + settings: Settings.isolated({ "compaction.enabled": false }), + modelRegistry: createModelRegistryStub() as never, + sideStreamFn, + }); + sessions.push(session); + + const result = await session.runEphemeralTurn({ promptText: "Question?" }); + + expect(result.replyText).toBe("Side answer"); + expect(capturedContext?.messages.at(-1)?.content).toEqual([{ type: "text", text: "Question?" }]); + expect(capturedOptions?.sessionId).toStartWith(`${session.sessionId}:side:`); + }); + it("rotates ephemeral side-channel credentials on Google Resource exhausted", async () => { const api = "test-ephemeral-google-resource-exhausted"; const googleErrorMessage = "Google API error (429): Resource exhausted. Please try again later.";