diff --git a/packages/coding-agent/src/modes/components/btw-panel.ts b/packages/coding-agent/src/modes/components/btw-panel.ts new file mode 100644 index 000000000..362857953 --- /dev/null +++ b/packages/coding-agent/src/modes/components/btw-panel.ts @@ -0,0 +1,104 @@ +import { type Component, Container, Markdown, Spacer, Text, type TUI } from "@oh-my-pi/pi-tui"; +import { replaceTabs } from "../../tools/render-utils"; +import { getMarkdownTheme, theme } from "../theme/theme"; +import { DynamicBorder } from "./dynamic-border"; + +type BtwPanelState = "running" | "complete" | "aborted" | "error"; + +interface BtwPanelComponentOptions { + question: string; + tui: TUI; +} + +export class BtwPanelComponent extends Container { + #question: string; + #tui: TUI; + #state: BtwPanelState = "running"; + #answer = ""; + #errorMessage: string | undefined; + #closed = false; + + constructor(options: BtwPanelComponentOptions) { + super(); + this.#question = options.question; + this.#tui = options.tui; + this.#rebuild(); + } + + appendText(delta: string): void { + if (!delta || this.#closed) return; + this.#answer += delta; + this.#rebuild(); + } + + setAnswer(text: string): void { + if (this.#closed) return; + this.#answer = text; + this.#rebuild(); + } + + markComplete(): void { + if (this.#closed) return; + this.#state = "complete"; + this.#errorMessage = undefined; + this.#rebuild(); + } + + markAborted(): void { + if (this.#closed) return; + this.#state = "aborted"; + this.#errorMessage = undefined; + this.#rebuild(); + } + + markError(message: string): void { + if (this.#closed) return; + this.#state = "error"; + this.#errorMessage = message; + this.#rebuild(); + } + + close(): void { + this.#closed = true; + } + + #rebuild(): void { + this.clear(); + this.addChild(new DynamicBorder(str => theme.fg("dim", str))); + this.addChild(new Spacer(1)); + this.addChild(new Text(theme.fg("accent", replaceTabs(this.#question)), 1, 0)); + this.addChild(new Spacer(1)); + this.addChild(this.#contentComponent()); + this.addChild(new Spacer(1)); + this.addChild(new Text(this.#footerLine(), 1, 0)); + this.addChild(new Spacer(1)); + this.addChild(new DynamicBorder(str => theme.fg("dim", str))); + this.#tui.requestRender(); + } + + #footerLine(): string { + switch (this.#state) { + case "running": + return theme.fg("muted", "Esc cancel /btw"); + case "complete": + return theme.fg("muted", "Esc dismiss"); + case "aborted": + return theme.fg("warning", `${theme.status.warning} Cancelled · Esc dismiss`); + case "error": + return theme.fg("error", `${theme.status.error} Error · Esc dismiss`); + } + } + + #contentComponent(): Component { + if (this.#state === "error") { + return new Text(theme.fg("error", replaceTabs(this.#errorMessage ?? "Unknown error")), 1, 0); + } + const text = replaceTabs(this.#answer).trim(); + if (!text) { + const waiting = + this.#state === "running" ? `${theme.status.pending} Waiting for response…` : "No text returned."; + return new Text(theme.fg("dim", waiting), 1, 0); + } + return new Markdown(text, 1, 0, getMarkdownTheme()); + } +} diff --git a/packages/coding-agent/src/modes/controllers/btw-controller.ts b/packages/coding-agent/src/modes/controllers/btw-controller.ts new file mode 100644 index 000000000..1fdb4501e --- /dev/null +++ b/packages/coding-agent/src/modes/controllers/btw-controller.ts @@ -0,0 +1,193 @@ +import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; +import { type AssistantMessage, type Context, streamSimple } from "@oh-my-pi/pi-ai"; +import { renderPromptTemplate } from "../../config/prompt-templates"; +import btwUserPrompt from "../../prompts/system/btw-user.md" with { type: "text" }; +import { toReasoningEffort } from "../../thinking"; +import { BtwPanelComponent } from "../components/btw-panel"; +import type { InteractiveModeContext } from "../types"; + +interface BtwRequest { + component: BtwPanelComponent; + abortController: AbortController; + question: string; +} + +export class BtwController { + #activeRequest: BtwRequest | undefined; + readonly #streamFn: typeof streamSimple; + + constructor( + private readonly ctx: InteractiveModeContext, + options?: { streamFn?: typeof streamSimple }, + ) { + this.#streamFn = options?.streamFn ?? streamSimple; + } + + hasActiveRequest(): boolean { + return this.#activeRequest !== undefined; + } + + handleEscape(): boolean { + if (!this.#activeRequest) return false; + this.#closeActiveRequest({ abort: this.#activeRequest.abortController.signal.aborted === false }); + return true; + } + + dispose(): void { + this.#closeActiveRequest({ abort: true }); + } + + async start(question: string): Promise { + const trimmedQuestion = question.trim(); + if (!trimmedQuestion) { + this.ctx.showStatus("Usage: /btw "); + return; + } + + const model = this.ctx.session.model; + if (!model) { + this.ctx.showError("No active model available for /btw."); + return; + } + + this.#closeActiveRequest({ abort: true }); + + const request: BtwRequest = { + component: new BtwPanelComponent({ question: trimmedQuestion, tui: this.ctx.ui }), + abortController: new AbortController(), + question: trimmedQuestion, + }; + this.ctx.btwContainer.clear(); + this.ctx.btwContainer.addChild(request.component); + this.ctx.ui.requestRender(); + this.#activeRequest = request; + void this.#runRequest(request, model); + } + + async #runRequest( + request: BtwRequest, + model: NonNullable, + ): Promise { + try { + const apiKey = await this.ctx.session.modelRegistry.getApiKey(model, this.ctx.session.sessionId); + if (!apiKey) { + throw new Error(`No API key for provider: ${model.provider}`); + } + + const llmMessages = await this.ctx.session.convertMessagesToLlm( + [...this.#buildMessageSnapshot(), this.#buildQuestionMessage(request.question)], + request.abortController.signal, + ); + const context: Context = { + systemPrompt: this.ctx.session.systemPrompt, + messages: llmMessages, + }; + const options = this.ctx.session.prepareSimpleStreamOptions({ + apiKey, + sessionId: this.ctx.session.sessionId, + reasoning: toReasoningEffort(this.ctx.session.thinkingLevel), + serviceTier: this.ctx.session.serviceTier, + signal: request.abortController.signal, + toolChoice: "none", + }); + const stream = this.#streamFn(model, context, options); + + for await (const event of stream) { + if (!this.#isActiveRequest(request)) { + return; + } + if (event.type === "text_delta") { + request.component.appendText(event.delta); + continue; + } + if (event.type === "done") { + const finalText = this.#assistantText(event.message); + if (finalText) { + request.component.setAnswer(finalText); + } + request.component.markComplete(); + return; + } + if (event.type === "error") { + if (event.reason === "aborted" || request.abortController.signal.aborted) { + request.component.markAborted(); + } else { + request.component.markError( + this.#assistantText(event.error) || event.error.errorMessage || "BTW request failed.", + ); + } + return; + } + } + } catch (error) { + if (!this.#isActiveRequest(request)) { + return; + } + if (request.abortController.signal.aborted) { + request.component.markAborted(); + return; + } + request.component.markError(error instanceof Error ? error.message : String(error)); + } + } + + #buildQuestionMessage(question: string): AgentMessage { + return { + role: "user", + content: [ + { + type: "text", + text: renderPromptTemplate(btwUserPrompt, { question }), + }, + ], + attribution: "user", + timestamp: Date.now(), + }; + } + + #buildMessageSnapshot(): AgentMessage[] { + const messages = this.ctx.session.messages.slice(); + if (!this.ctx.session.isStreaming || !this.ctx.streamingMessage) { + return messages; + } + const streamingText = this.ctx.extractAssistantText(this.ctx.streamingMessage); + const lastMessage = messages.at(-1); + if (!streamingText) { + return lastMessage?.role === "assistant" ? messages.slice(0, -1) : messages; + } + const normalizedStreamingMessage: AssistantMessage = { + ...this.ctx.streamingMessage, + content: [{ type: "text", text: streamingText }], + }; + if (lastMessage?.role === "assistant") { + return [...messages.slice(0, -1), normalizedStreamingMessage]; + } + return [...messages, normalizedStreamingMessage]; + } + + #assistantText(message: AssistantMessage): string { + let text = ""; + for (const content of message.content) { + if (content.type === "text") { + text += content.text; + } + } + return text.trim(); + } + + #closeActiveRequest(options: { abort: boolean }): void { + const request = this.#activeRequest; + if (!request) return; + this.#activeRequest = undefined; + if (options.abort) { + request.abortController.abort(); + } + request.component.close(); + this.ctx.btwContainer.clear(); + this.ctx.ui.requestRender(); + } + + #isActiveRequest(request: BtwRequest): boolean { + return this.#activeRequest === request; + } +} diff --git a/packages/coding-agent/src/modes/controllers/input-controller.ts b/packages/coding-agent/src/modes/controllers/input-controller.ts index b6b4f97d2..ac08a8fc6 100644 --- a/packages/coding-agent/src/modes/controllers/input-controller.ts +++ b/packages/coding-agent/src/modes/controllers/input-controller.ts @@ -29,6 +29,7 @@ export class InputController { this.ctx.editor.shouldBypassAutocompleteOnEscape = () => Boolean( this.ctx.loadingAnimation || + this.ctx.hasActiveBtw() || this.ctx.session.isStreaming || this.ctx.session.isCompacting || this.ctx.session.isGeneratingHandoff || @@ -40,6 +41,9 @@ export class InputController { this.ctx.retryEscapeHandler, ); this.ctx.editor.onEscape = () => { + if (this.ctx.hasActiveBtw() && this.ctx.handleBtwEscape()) { + return; + } if (this.ctx.loadingAnimation) { if (this.ctx.cancelPendingSubmission()) { return; @@ -435,6 +439,9 @@ export class InputController { this.ctx.showWarning("Agent is idle; nothing to background"); return; } + if (this.ctx.hasActiveBtw()) { + this.ctx.handleBtwEscape(); + } this.ctx.isBackgrounded = true; const backgroundUiContext = this.ctx.createBackgroundUiContext(); diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index faec12381..22f66b8fb 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -36,6 +36,7 @@ import type { PythonExecutionComponent } from "./components/python-execution"; import { StatusLineComponent } from "./components/status-line"; import type { ToolExecutionHandle } from "./components/tool-execution"; import { WelcomeComponent } from "./components/welcome"; +import { BtwController } from "./controllers/btw-controller"; import { CommandController } from "./controllers/command-controller"; import { EventController } from "./controllers/event-controller"; import { ExtensionUiController } from "./controllers/extension-ui-controller"; @@ -89,6 +90,7 @@ export class InteractiveMode implements InteractiveModeContext { pendingMessagesContainer: Container; statusContainer: Container; todoContainer: Container; + btwContainer: Container; editor: CustomEditor; editorContainer: Container; statusLine: StatusLineComponent; @@ -151,6 +153,7 @@ export class InteractiveMode implements InteractiveModeContext { mcpManager?: import("../mcp").MCPManager; readonly #toolUiContextSetter: (uiContext: ExtensionUIContext, hasUI: boolean) => void; + readonly #btwController: BtwController; readonly #commandController: CommandController; readonly #eventController: EventController; readonly #extensionUiController: ExtensionUiController; @@ -192,6 +195,7 @@ export class InteractiveMode implements InteractiveModeContext { this.pendingMessagesContainer = new Container(); this.statusContainer = new Container(); this.todoContainer = new Container(); + this.btwContainer = new Container(); this.editor = new CustomEditor(getEditorTheme()); this.editor.setUseTerminalCursor(this.ui.getShowHardwareCursor()); this.editor.setAutocompleteMaxVisible(settings.get("autocompleteMaxVisible")); @@ -248,6 +252,7 @@ export class InteractiveMode implements InteractiveModeContext { this.#pendingSlashCommands = [...BUILTIN_SLASH_COMMANDS, ...hookCommands, ...customCommands, ...skillCommandList]; this.#uiHelpers = new UiHelpers(this); + this.#btwController = new BtwController(this); this.#extensionUiController = new ExtensionUiController(this); this.#eventController = new EventController(this); this.#commandController = new CommandController(this); @@ -328,6 +333,7 @@ export class InteractiveMode implements InteractiveModeContext { this.ui.addChild(this.pendingMessagesContainer); this.ui.addChild(this.statusContainer); this.ui.addChild(this.todoContainer); + this.ui.addChild(this.btwContainer); this.ui.addChild(this.statusLine); // Only renders hook statuses (main status in editor border) this.ui.addChild(new Spacer(1)); this.ui.addChild(this.editorContainer); @@ -861,6 +867,7 @@ export class InteractiveMode implements InteractiveModeContext { // Flush pending session writes before shutdown await this.sessionManager.flush(); + this.#btwController.dispose(); // Emit shutdown event to hooks await this.session.dispose(); @@ -1073,11 +1080,13 @@ export class InteractiveMode implements InteractiveModeContext { } handleClearCommand(): Promise { + this.#btwController.dispose(); this.#extensionUiController.clearExtensionTerminalInputListeners(); return this.#commandController.handleClearCommand(); } handleForkCommand(): Promise { + this.#btwController.dispose(); return this.#commandController.handleForkCommand(); } @@ -1233,6 +1242,7 @@ export class InteractiveMode implements InteractiveModeContext { } handleResumeSession(sessionPath: string): Promise { + this.#btwController.dispose(); return this.#selectorController.handleResumeSession(sessionPath); } @@ -1269,6 +1279,18 @@ export class InteractiveMode implements InteractiveModeContext { return this.#inputController.handleImagePaste(); } + handleBtwCommand(question: string): Promise { + return this.#btwController.start(question); + } + + hasActiveBtw(): boolean { + return this.#btwController.hasActiveRequest(); + } + + handleBtwEscape(): boolean { + return this.#btwController.handleEscape(); + } + cycleThinkingLevel(): void { this.#inputController.cycleThinkingLevel(); } diff --git a/packages/coding-agent/src/modes/types.ts b/packages/coding-agent/src/modes/types.ts index b6c4fc7d0..b4eb2cf3a 100644 --- a/packages/coding-agent/src/modes/types.ts +++ b/packages/coding-agent/src/modes/types.ts @@ -56,6 +56,7 @@ export interface InteractiveModeContext { pendingMessagesContainer: Container; statusContainer: Container; todoContainer: Container; + btwContainer: Container; editor: CustomEditor; editorContainer: Container; statusLine: StatusLineComponent; @@ -205,6 +206,9 @@ export interface InteractiveModeContext { handleDequeue(): void; handleBackgroundCommand(): void; handleImagePaste(): Promise; + handleBtwCommand(question: string): Promise; + hasActiveBtw(): boolean; + handleBtwEscape(): boolean; cycleThinkingLevel(): void; cycleRoleModel(options?: { temporary?: boolean }): Promise; toggleToolOutputExpansion(): void; diff --git a/packages/coding-agent/src/prompts/system/btw-user.md b/packages/coding-agent/src/prompts/system/btw-user.md new file mode 100644 index 000000000..857614841 --- /dev/null +++ b/packages/coding-agent/src/prompts/system/btw-user.md @@ -0,0 +1,8 @@ + +This is an ephemeral side question for the current interactive session. +Answer briefly and directly using the conversation context already provided. +Do not use tools. +Do not ask follow-up questions. +Question: +{{question}} + diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index 71cac1761..0d6c22485 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -1340,6 +1340,16 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} if (!obfuscator?.hasSecrets()) return converted; return obfuscateMessages(obfuscator, converted); }; + const transformContext = extensionRunner + ? async (messages: AgentMessage[], _signal?: AbortSignal) => { + return await extensionRunner.emitContext(messages); + } + : undefined; + const onPayload = extensionRunner + ? async (payload: unknown, _model?: Model) => { + return await extensionRunner.emitBeforeProviderRequest(payload); + } + : undefined; const setToolUIContext = (uiContext: ExtensionUIContext, hasUI: boolean) => { toolContextStore.setUIContext(uiContext, hasUI); @@ -1362,17 +1372,9 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} tools: initialTools, }, convertToLlm: convertToLlmFinal, - onPayload: extensionRunner - ? async (payload, _model) => { - return extensionRunner.emitBeforeProviderRequest(payload); - } - : undefined, + onPayload, sessionId: sessionManager.getSessionId(), - transformContext: extensionRunner - ? async messages => { - return extensionRunner.emitContext(messages); - } - : undefined, + transformContext, steeringMode: settings.get("steeringMode") ?? "one-at-a-time", followUpMode: settings.get("followUpMode") ?? "one-at-a-time", interruptMode: settings.get("interruptMode") ?? "immediate", @@ -1447,6 +1449,9 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} skillsSettings: settings.getGroup("skills") as Required, modelRegistry, toolRegistry, + transformContext, + onPayload, + convertToLlm: convertToLlmFinal, rebuildSystemPrompt, ttsrManager, forceCopilotAgentInitiator, diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 434f16871..54bcb0185 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -34,6 +34,7 @@ import type { Model, ProviderSessionState, ServiceTier, + SimpleStreamOptions, TextContent, ToolCall, ToolChoice, @@ -126,6 +127,7 @@ import { bashExecutionToText, type CompactionSummaryMessage, type CustomMessage, + convertToLlm, type FileMentionMessage, type HookMessage, type PythonExecutionMessage, @@ -191,6 +193,12 @@ export interface AgentSessionConfig { modelRegistry: ModelRegistry; /** Tool registry for LSP and settings */ toolRegistry?: Map; + /** Current session pre-LLM message transform pipeline */ + transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => AgentMessage[] | Promise; + /** Provider payload hook used by the active session request path */ + onPayload?: SimpleStreamOptions["onPayload"]; + /** Current session message-to-LLM conversion pipeline */ + convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise; /** System prompt builder that can consider tool availability */ rebuildSystemPrompt?: (toolNames: string[], tools: Map) => Promise; /** TTSR manager for time-traveling stream rules */ @@ -381,6 +389,9 @@ export class AgentSession { // Tool registry and prompt builder for extensions #toolRegistry: Map; + #transformContext: (messages: AgentMessage[], signal?: AbortSignal) => AgentMessage[] | Promise; + #onPayload: SimpleStreamOptions["onPayload"] | undefined; + #convertToLlm: (messages: AgentMessage[]) => Message[] | Promise; #rebuildSystemPrompt: ((toolNames: string[], tools: Map) => Promise) | undefined; #baseSystemPrompt: string; #forceCopilotAgentInitiator = false; @@ -425,6 +436,9 @@ export class AgentSession { this.#skillsSettings = config.skillsSettings; this.#modelRegistry = config.modelRegistry; this.#toolRegistry = config.toolRegistry ?? new Map(); + this.#transformContext = config.transformContext ?? (messages => messages); + this.#onPayload = config.onPayload; + this.#convertToLlm = config.convertToLlm ?? convertToLlm; this.#rebuildSystemPrompt = config.rebuildSystemPrompt; this.#baseSystemPrompt = this.agent.state.systemPrompt; this.#ttsrManager = config.ttsrManager; @@ -1714,6 +1728,31 @@ export class AgentSession { return this.agent.state.messages; } + /** Convert session messages using the same pre-LLM pipeline as the active session. */ + async convertMessagesToLlm(messages: AgentMessage[], signal?: AbortSignal): Promise { + const transformedMessages = await this.#transformContext(messages, signal); + return await this.#convertToLlm(transformedMessages); + } + + /** Apply session-level stream hooks to a direct side request. */ + prepareSimpleStreamOptions(options: SimpleStreamOptions): SimpleStreamOptions { + if (!this.#onPayload) return options; + if (!options.onPayload) { + return { ...options, onPayload: this.#onPayload }; + } + const sessionOnPayload = this.#onPayload; + const requestOnPayload = options.onPayload; + return { + ...options, + onPayload: async (payload, model) => { + const sessionPayload = await sessionOnPayload(payload, model); + const sessionResolvedPayload = sessionPayload ?? payload; + const requestPayload = await requestOnPayload(sessionResolvedPayload, model); + return requestPayload ?? sessionResolvedPayload; + }, + }; + } + /** Current steering mode */ get steeringMode(): "all" | "one-at-a-time" { return this.agent.getSteeringMode(); diff --git a/packages/coding-agent/src/slash-commands/builtin-registry.ts b/packages/coding-agent/src/slash-commands/builtin-registry.ts index 89615b98e..927fe72f7 100644 --- a/packages/coding-agent/src/slash-commands/builtin-registry.ts +++ b/packages/coding-agent/src/slash-commands/builtin-registry.ts @@ -465,6 +465,17 @@ const BUILTIN_SLASH_COMMAND_REGISTRY: ReadonlyArray = [ runtime.ctx.editor.setText(""); }, }, + { + name: "btw", + description: "Ask an ephemeral side question using the current session context", + inlineHint: "", + allowArgs: true, + handle: async (command, runtime) => { + const question = command.text.slice(`/${command.name}`.length).trim(); + runtime.ctx.editor.setText(""); + await runtime.ctx.handleBtwCommand(question); + }, + }, { name: "background", aliases: ["bg"], diff --git a/packages/coding-agent/test/agent-session-message-pipeline.test.ts b/packages/coding-agent/test/agent-session-message-pipeline.test.ts new file mode 100644 index 000000000..07248758d --- /dev/null +++ b/packages/coding-agent/test/agent-session-message-pipeline.test.ts @@ -0,0 +1,93 @@ +import { afterEach, describe, expect, it, vi } from "bun:test"; +import { Agent, type AgentMessage } from "@oh-my-pi/pi-agent-core"; +import type { Message, SimpleStreamOptions } from "@oh-my-pi/pi-ai"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; + +function createAgent(): Agent { + return new Agent({ + initialState: { + systemPrompt: "system prompt", + messages: [], + tools: [], + }, + }); +} + +describe("AgentSession message pipeline", () => { + const sessions: AgentSession[] = []; + + afterEach(async () => { + vi.restoreAllMocks(); + for (const session of sessions.splice(0)) { + await session.dispose(); + } + }); + + it("applies transformContext before convertToLlm", async () => { + const inputMessages: AgentMessage[] = [{ role: "user", content: "hello", timestamp: Date.now() }]; + const transformedMessages: AgentMessage[] = [ + ...inputMessages, + { role: "user", content: "injected context", timestamp: Date.now() }, + ]; + const convertedMessages: Message[] = [ + { + role: "user", + content: [{ type: "text", text: "converted" }], + attribution: "user", + timestamp: Date.now(), + }, + ]; + const transformContext = vi.fn(async (messages: AgentMessage[], signal?: AbortSignal) => { + expect(signal).toBe(abortController.signal); + return [...messages, ...transformedMessages.slice(messages.length)]; + }); + const convertToLlm = vi.fn(async (_messages: AgentMessage[]) => { + return convertedMessages; + }); + const abortController = new AbortController(); + const session = new AgentSession({ + agent: createAgent(), + sessionManager: SessionManager.inMemory(), + settings: Settings.isolated({ "compaction.enabled": false }), + modelRegistry: {} as never, + transformContext, + convertToLlm, + }); + sessions.push(session); + + const result = await session.convertMessagesToLlm(inputMessages, abortController.signal); + + expect(transformContext).toHaveBeenCalledWith(inputMessages, abortController.signal); + expect(convertToLlm).toHaveBeenCalledWith(transformedMessages); + expect(result).toEqual(convertedMessages); + }); + + it("composes session payload hooks into direct side-request options", async () => { + const sessionOnPayload = vi.fn(async (payload: unknown) => ({ + ...(payload as Record), + session: true, + })); + const requestOnPayload = vi.fn(async () => undefined); + const session = new AgentSession({ + agent: createAgent(), + sessionManager: SessionManager.inMemory(), + settings: Settings.isolated({ "compaction.enabled": false }), + modelRegistry: {} as never, + onPayload: sessionOnPayload, + }); + sessions.push(session); + const options: SimpleStreamOptions = { + apiKey: "key", + onPayload: requestOnPayload, + }; + + const prepared = session.prepareSimpleStreamOptions(options); + const result = await prepared.onPayload?.({ original: true }); + + expect(sessionOnPayload).toHaveBeenCalledWith({ original: true }, undefined); + expect(requestOnPayload).toHaveBeenCalledWith({ original: true, session: true }, undefined); + expect(result).toEqual({ original: true, session: true }); + }); +}); diff --git a/packages/coding-agent/test/input-controller-escape.test.ts b/packages/coding-agent/test/input-controller-escape.test.ts index eaac71cc5..22a3de918 100644 --- a/packages/coding-agent/test/input-controller-escape.test.ts +++ b/packages/coding-agent/test/input-controller-escape.test.ts @@ -49,7 +49,11 @@ function createContext(): { cancelPendingSubmission: ReturnType; clearQueue: ReturnType; ensureLoadingAnimation: ReturnType; + handleBtwCommand: ReturnType; + handleBtwEscape: ReturnType; + hasActiveBtw: ReturnType; onInputCallback: ReturnType; + prompt: ReturnType; requestRender: ReturnType; startPendingSubmission: ReturnType; }; @@ -62,7 +66,11 @@ function createContext(): { const cancelPendingSubmission = vi.fn(() => false); const clearQueue = vi.fn(() => ({ steering: [], followUp: [] })); const onInputCallback = vi.fn(); + const prompt = vi.fn(); const requestRender = vi.fn(); + const handleBtwCommand = vi.fn(async () => {}); + const handleBtwEscape = vi.fn(() => true); + const hasActiveBtw = vi.fn(() => false); const startPendingSubmission = vi.fn((input: { text: string; images?: InteractiveModeContext["pendingImages"] }) => { ensureLoadingAnimation(); return createSubmission(input); @@ -104,6 +112,7 @@ function createContext(): { abortBash, abortPython, clearQueue, + prompt, } as unknown as InteractiveModeContext["session"], sessionManager: { getSessionName: () => "existing session", @@ -129,6 +138,9 @@ function createContext(): { toggleTodoExpansion: vi.fn(), handleHotkeysCommand: vi.fn(), handleSTTToggle: vi.fn(), + handleBtwEscape, + handleBtwCommand, + hasActiveBtw, showTreeSelector: vi.fn(), showUserMessageSelector: vi.fn(), showSessionSelector: vi.fn(), @@ -145,7 +157,11 @@ function createContext(): { cancelPendingSubmission, clearQueue, ensureLoadingAnimation, + handleBtwCommand, + handleBtwEscape, + hasActiveBtw, onInputCallback, + prompt, requestRender, startPendingSubmission, }, @@ -175,6 +191,21 @@ describe("InputController escape behavior", () => { expect(spies.abort).not.toHaveBeenCalled(); }); + it("runs /btw as a builtin side request instead of steering the active stream", async () => { + const { ctx, editor, spies } = createContext(); + (ctx.session as { isStreaming: boolean }).isStreaming = true; + const controller = new InputController(ctx); + + controller.setupEditorSubmitHandler(); + editor.setText("/btw why is it doing that?"); + await editor.onSubmit?.("/btw why is it doing that?"); + + expect(spies.handleBtwCommand).toHaveBeenCalledWith("why is it doing that?"); + expect(spies.prompt).not.toHaveBeenCalled(); + expect(editor.addToHistory).not.toHaveBeenCalled(); + expect(editor.getText()).toBe(""); + }); + it("falls back to aborting the active session when no pending optimistic submission exists", () => { const { ctx, editor, spies } = createContext(); ctx.loadingAnimation = {} as InteractiveModeContext["loadingAnimation"]; @@ -214,6 +245,51 @@ describe("InputController escape behavior", () => { expect(spies.abort).not.toHaveBeenCalled(); }); + it("dismisses an active /btw panel before aborting the main stream", () => { + const { ctx, editor, spies } = createContext(); + (ctx.session as { isStreaming: boolean }).isStreaming = true; + spies.hasActiveBtw.mockReturnValue(true); + const controller = new InputController(ctx); + + controller.setupKeyHandlers(); + expect(editor.shouldBypassAutocompleteOnEscape?.()).toBe(true); + editor.onEscape?.(); + + expect(spies.handleBtwEscape).toHaveBeenCalledTimes(1); + expect(spies.abort).not.toHaveBeenCalled(); + }); + + it("dismisses an active /btw panel before canceling a pending optimistic submission", () => { + const { ctx, editor, spies } = createContext(); + ctx.loadingAnimation = {} as InteractiveModeContext["loadingAnimation"]; + spies.hasActiveBtw.mockReturnValue(true); + const controller = new InputController(ctx); + + controller.setupKeyHandlers(); + expect(editor.shouldBypassAutocompleteOnEscape?.()).toBe(true); + editor.onEscape?.(); + + expect(spies.handleBtwEscape).toHaveBeenCalledTimes(1); + expect(spies.cancelPendingSubmission).not.toHaveBeenCalled(); + expect(spies.clearQueue).not.toHaveBeenCalled(); + expect(spies.abort).not.toHaveBeenCalled(); + }); + + it("dismisses an active /btw panel before aborting bash", () => { + const { ctx, editor, spies } = createContext(); + (ctx.session as { isBashRunning: boolean }).isBashRunning = true; + spies.hasActiveBtw.mockReturnValue(true); + const controller = new InputController(ctx); + + controller.setupKeyHandlers(); + expect(editor.shouldBypassAutocompleteOnEscape?.()).toBe(true); + editor.onEscape?.(); + + expect(spies.handleBtwEscape).toHaveBeenCalledTimes(1); + expect(spies.abortBash).not.toHaveBeenCalled(); + expect(spies.abort).not.toHaveBeenCalled(); + }); + it("aborts streaming even when the working loader is no longer present", () => { const { ctx, editor, spies } = createContext(); (ctx.session as { isStreaming: boolean }).isStreaming = true; diff --git a/packages/coding-agent/test/modes/controllers/btw-controller.test.ts b/packages/coding-agent/test/modes/controllers/btw-controller.test.ts new file mode 100644 index 000000000..a2723344b --- /dev/null +++ b/packages/coding-agent/test/modes/controllers/btw-controller.test.ts @@ -0,0 +1,291 @@ +import { beforeAll, describe, expect, it, vi } from "bun:test"; +import { type AgentMessage, ThinkingLevel } from "@oh-my-pi/pi-agent-core"; +import type { AssistantMessage, Message, Usage } from "@oh-my-pi/pi-ai"; +import { getBundledModel } from "@oh-my-pi/pi-ai"; +import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream"; +import { BtwController } from "@oh-my-pi/pi-coding-agent/modes/controllers/btw-controller"; +import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; +import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types"; +import { Container, type TUI } from "@oh-my-pi/pi-tui"; + +const usage: Usage = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, +}; + +function createAssistantMessage(text: string): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4-5", + usage, + stopReason: "stop", + timestamp: Date.now(), + }; +} + +function createUserMessage(text: string): Message { + return { + role: "user", + content: [{ type: "text", text }], + attribution: "user", + timestamp: Date.now(), + }; +} + +beforeAll(async () => { + await initTheme(); +}); + +describe("BtwController", () => { + it("builds a tool-less side request from the current session prefix and preserves payload hooks", async () => { + const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; + const convertMessagesToLlm = vi.fn(async (messages: AgentMessage[]) => { + const question = messages.at(-1); + if (!question || !("content" in question)) { + throw new Error("Expected the /btw question to be present in the conversion pipeline"); + } + const questionText = + typeof question.content === "string" + ? question.content + : question.content + .filter(content => content.type === "text") + .map(content => content.text) + .join(""); + return [createUserMessage(sessionMessages[0].content), createUserMessage(questionText)]; + }); + const getApiKey = vi.fn(async () => "key"); + const onPayload = vi.fn(async payload => payload); + const prepareSimpleStreamOptions = vi.fn(options => ({ ...options, onPayload })); + const prompt = vi.fn(); + const requestRender = vi.fn(); + const btwContainer = new Container(); + const sessionMessages = [{ role: "user" as const, content: "hello", timestamp: Date.now() }]; + const streamFn = vi.fn((_model, _context, _options) => { + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Answer") }); + }); + return stream; + }); + const ctx = { + ui: { requestRender } as unknown as TUI, + btwContainer, + session: { + model, + messages: sessionMessages, + isStreaming: false, + sessionId: "session-1", + serviceTier: "priority", + thinkingLevel: ThinkingLevel.High, + systemPrompt: "system prompt", + modelRegistry: { getApiKey } as unknown as InteractiveModeContext["session"]["modelRegistry"], + convertMessagesToLlm, + prepareSimpleStreamOptions, + prompt, + } as unknown as InteractiveModeContext["session"], + streamingMessage: undefined, + extractAssistantText: (message: AssistantMessage) => + message.content + .filter(content => content.type === "text") + .map(content => content.text) + .join(""), + showStatus: vi.fn(), + showError: vi.fn(), + } as unknown as InteractiveModeContext; + const controller = new BtwController(ctx, { streamFn }); + + await controller.start("What changed?"); + await Bun.sleep(0); + + const convertCall = convertMessagesToLlm.mock.calls[0] as unknown as [AgentMessage[], AbortSignal]; + expect(convertCall).toBeDefined(); + expect(convertCall[0]).toHaveLength(2); + expect(convertCall[0][0]).toEqual(ctx.session.messages[0]); + expect(convertCall[0][1]?.role).toBe("user"); + expect(convertCall[1]).toBeInstanceOf(AbortSignal); + const appendedQuestion = convertCall[0][1] as { + role: "user"; + content: Array<{ type: string; text?: string }>; + }; + expect(appendedQuestion.content[0]?.type).toBe("text"); + expect(appendedQuestion.content[0]?.text).toContain("What changed?"); + + expect(getApiKey).toHaveBeenCalledWith(model, "session-1"); + expect(prepareSimpleStreamOptions).toHaveBeenCalledTimes(1); + expect(streamFn).toHaveBeenCalledTimes(1); + const [, context, options] = streamFn.mock.calls[0] as [ + unknown, + { systemPrompt?: string; messages: Message[] }, + Record, + ]; + expect(context.systemPrompt).toBe("system prompt"); + expect(context.messages).toHaveLength(2); + expect((context.messages[1]?.content as Array<{ type: string; text?: string }>)[0]?.text).toContain( + "What changed?", + ); + expect(options.apiKey).toBe("key"); + expect(options.sessionId).toBe("session-1"); + expect(options.serviceTier).toBe("priority"); + expect(options.reasoning).toBe(ThinkingLevel.High); + expect(options.toolChoice).toBe("none"); + expect(options.onPayload).toBe(onPayload); + expect("providerSessionState" in options).toBe(false); + expect(prompt).not.toHaveBeenCalled(); + expect(ctx.session.messages).toEqual(sessionMessages); + expect(btwContainer.children).toHaveLength(1); + expect(controller.hasActiveRequest()).toBe(true); + }); + + it("appends the active streaming assistant snapshot when session history still ends with the user message", async () => { + const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; + const convertMessagesToLlm = vi.fn(async () => []); + const streamFn = vi.fn((_model, _context, _options) => { + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") }); + }); + return stream; + }); + const streamingMessage = { + ...createAssistantMessage("partial answer"), + content: [ + { type: "text", text: "partial answer" }, + { type: "toolCall", id: "tool-1", name: "read", arguments: {} }, + ], + } as AssistantMessage; + const ctx = { + ui: { requestRender: vi.fn() } as unknown as TUI, + btwContainer: new Container(), + session: { + model, + messages: [createUserMessage("latest user")], + isStreaming: true, + sessionId: "session-1", + serviceTier: undefined, + thinkingLevel: ThinkingLevel.Off, + systemPrompt: "system prompt", + modelRegistry: { + getApiKey: async () => "key", + } as unknown as InteractiveModeContext["session"]["modelRegistry"], + convertMessagesToLlm, + prepareSimpleStreamOptions: (options => + options) as InteractiveModeContext["session"]["prepareSimpleStreamOptions"], + } as unknown as InteractiveModeContext["session"], + streamingMessage, + extractAssistantText: () => "partial answer", + showStatus: vi.fn(), + showError: vi.fn(), + } as unknown as InteractiveModeContext; + const controller = new BtwController(ctx, { streamFn }); + + await controller.start("Why?"); + await Bun.sleep(0); + + const firstCall = convertMessagesToLlm.mock.calls[0] as unknown as [AgentMessage[], AbortSignal]; + expect(firstCall).toBeDefined(); + expect(firstCall[1]).toBeInstanceOf(AbortSignal); + const snapshot = firstCall[0]; + expect(snapshot).toHaveLength(3); + expect(snapshot[0]?.role).toBe("user"); + const normalizedAssistant = snapshot[1] as AssistantMessage; + expect(normalizedAssistant.role).toBe("assistant"); + expect(normalizedAssistant.content).toEqual([{ type: "text", text: "partial answer" }]); + const appendedQuestion = snapshot[2] as { role: "user"; content: Array<{ type: string; text?: string }> }; + expect(appendedQuestion.role).toBe("user"); + expect(appendedQuestion.content[0]?.text).toContain("Why?"); + }); + + it("replaces an existing request by aborting the previous btw stream and keeping one panel", async () => { + const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; + const firstStream = new AssistantMessageEventStream(); + const secondStream = new AssistantMessageEventStream(); + const btwContainer = new Container(); + const streamFn = vi + .fn() + .mockImplementationOnce((_model, _context, _options) => firstStream) + .mockImplementationOnce((_model, _context, _options) => { + queueMicrotask(() => { + secondStream.push({ type: "done", reason: "stop", message: createAssistantMessage("Second") }); + }); + return secondStream; + }); + const ctx = { + ui: { requestRender: vi.fn() } as unknown as TUI, + btwContainer, + session: { + model, + messages: [], + isStreaming: false, + sessionId: "session-1", + serviceTier: undefined, + thinkingLevel: ThinkingLevel.Off, + systemPrompt: "system prompt", + modelRegistry: { + getApiKey: async () => "key", + } as unknown as InteractiveModeContext["session"]["modelRegistry"], + convertMessagesToLlm: async () => [], + prepareSimpleStreamOptions: (options => + options) as InteractiveModeContext["session"]["prepareSimpleStreamOptions"], + } as unknown as InteractiveModeContext["session"], + streamingMessage: undefined, + extractAssistantText: vi.fn(), + showStatus: vi.fn(), + showError: vi.fn(), + } as unknown as InteractiveModeContext; + const controller = new BtwController(ctx, { streamFn }); + + await controller.start("First?"); + await controller.start("Second?"); + await Bun.sleep(0); + + const firstOptions = streamFn.mock.calls[0]?.[2] as { signal: AbortSignal }; + expect(firstOptions.signal.aborted).toBe(true); + expect(streamFn).toHaveBeenCalledTimes(2); + expect(btwContainer.children).toHaveLength(1); + expect(controller.hasActiveRequest()).toBe(true); + }); + + it("clears the btw panel when the active request is dismissed", async () => { + const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; + const btwContainer = new Container(); + const streamFn = vi.fn((_model, _context, _options) => new AssistantMessageEventStream()); + const ctx = { + ui: { requestRender: vi.fn() } as unknown as TUI, + btwContainer, + session: { + model, + messages: [], + isStreaming: false, + sessionId: "session-1", + serviceTier: undefined, + thinkingLevel: ThinkingLevel.Off, + systemPrompt: "system prompt", + modelRegistry: { + getApiKey: async () => "key", + } as unknown as InteractiveModeContext["session"]["modelRegistry"], + convertMessagesToLlm: async () => [], + prepareSimpleStreamOptions: (options => + options) as InteractiveModeContext["session"]["prepareSimpleStreamOptions"], + } as unknown as InteractiveModeContext["session"], + streamingMessage: undefined, + extractAssistantText: vi.fn(), + showStatus: vi.fn(), + showError: vi.fn(), + } as unknown as InteractiveModeContext; + const controller = new BtwController(ctx, { streamFn }); + + await controller.start("Question?"); + + expect(btwContainer.children).toHaveLength(1); + expect(controller.handleEscape()).toBe(true); + expect(btwContainer.children).toHaveLength(0); + expect(controller.hasActiveRequest()).toBe(false); + }); +}); diff --git a/packages/coding-agent/test/slash-commands/btw.test.ts b/packages/coding-agent/test/slash-commands/btw.test.ts new file mode 100644 index 000000000..1638229df --- /dev/null +++ b/packages/coding-agent/test/slash-commands/btw.test.ts @@ -0,0 +1,43 @@ +import { describe, expect, it, vi } from "bun:test"; +import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types"; +import { executeBuiltinSlashCommand } from "@oh-my-pi/pi-coding-agent/slash-commands/builtin-registry"; + +function createRuntime() { + const handleBtwCommand = vi.fn(async () => {}); + const setText = vi.fn(); + return { + handleBtwCommand, + setText, + runtime: { + ctx: { + editor: { setText } as unknown as InteractiveModeContext["editor"], + handleBtwCommand, + } as unknown as InteractiveModeContext, + handleBackgroundCommand: () => {}, + }, + }; +} + +describe("/btw slash command", () => { + it("routes the full question through the interactive btw handler", async () => { + const harness = createRuntime(); + + const handled = await executeBuiltinSlashCommand("/btw why is it doing that?", harness.runtime); + + expect(handled).toBe(true); + expect(harness.setText).toHaveBeenCalledWith(""); + expect(harness.handleBtwCommand).toHaveBeenCalledWith("why is it doing that?"); + }); + + it("preserves the raw multi-word suffix after /btw", async () => { + const harness = createRuntime(); + + const handled = await executeBuiltinSlashCommand( + "/btw explain why the cache reuse matters here", + harness.runtime, + ); + + expect(handled).toBe(true); + expect(harness.handleBtwCommand).toHaveBeenCalledWith("explain why the cache reuse matters here"); + }); +});