maximhar
952db9fd93
* feat: add ephemeral /btw side-question panel * fix: preserve /btw live context and payload hooks * fix: send /btw question through session pipeline
292 lines
11 KiB
TypeScript
292 lines
11 KiB
TypeScript
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<string, unknown>,
|
|
];
|
|
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);
|
|
});
|
|
});
|