import assert from "node:assert/strict"; import { mkdirSync, writeFileSync } from "node:fs"; import { dirname } from "node:path"; import { mock, runTest } from "./test-helpers.test.ts"; mock.module("@earendil-works/pi-coding-agent", { namedExports: { getAgentDir: () => "/tmp/.pi/agent", getSettingsListTheme: () => ({}), isToolCallEventType: (toolName: string, event: Record) => event.toolName === toolName, }, }); mock.module("@earendil-works/pi-tui", { namedExports: { Box: class {}, Container: class { addChild(): void {} render(): string[] { return []; } invalidate(): void {} }, SettingsList: class { handleInput(): void {} updateValue(): void {} }, Spacer: class {}, Text: class {}, truncateToWidth: (text: string) => text, visibleWidth: (text: string) => text.length, }, }); const indexModule = await import("./index.ts"); const { createBoundedNoticeTracker, shouldInjectSourceFilterTroubleshootingNote, injectGuidelineIntoPrompt } = indexModule; const rtkIntegrationExtension = indexModule.default; const { DEFAULT_RTK_INTEGRATION_CONFIG } = await import("./types.ts"); const { CONFIG_PATH } = await import("./constants.ts"); function writeTestConfig(commandRewritingEnabled: boolean): void { mkdirSync(dirname(CONFIG_PATH), { recursive: true }); writeFileSync( CONFIG_PATH, `${JSON.stringify({ ...DEFAULT_RTK_INTEGRATION_CONFIG, commandRewritingEnabled }, null, 2)}\n`, "utf-8", ); } type Notification = { message: string; level: "info" | "warning" | "error" }; type ExtensionHandler = (event: Record, ctx: Record) => Promise | void>; function createNotificationContext(notifications: Notification[]): Record { return { hasUI: true, ui: { notify(message: string, level: "info" | "warning" | "error") { notifications.push({ message, level }); }, }, }; } function firstText(content: unknown): string { if (!Array.isArray(content) || content.length === 0) { return ""; } const block = content[0] as { type?: string; text?: string }; return block.type === "text" && typeof block.text === "string" ? block.text : ""; } function configWith(overrides: { enabled?: boolean; compactionEnabled?: boolean; readCompactionEnabled?: boolean; sourceFilteringEnabled?: boolean; sourceFilteringLevel?: "none" | "minimal" | "aggressive"; smartTruncateEnabled?: boolean; truncateEnabled?: boolean; }): typeof DEFAULT_RTK_INTEGRATION_CONFIG { const base = DEFAULT_RTK_INTEGRATION_CONFIG; return { ...base, enabled: overrides.enabled ?? base.enabled, outputCompaction: { ...base.outputCompaction, enabled: overrides.compactionEnabled ?? base.outputCompaction.enabled, readCompaction: { ...base.outputCompaction.readCompaction, enabled: overrides.readCompactionEnabled ?? base.outputCompaction.readCompaction.enabled, }, sourceCodeFilteringEnabled: overrides.sourceFilteringEnabled ?? base.outputCompaction.sourceCodeFilteringEnabled, sourceCodeFiltering: overrides.sourceFilteringLevel ?? base.outputCompaction.sourceCodeFiltering, smartTruncate: { ...base.outputCompaction.smartTruncate, enabled: overrides.smartTruncateEnabled ?? base.outputCompaction.smartTruncate.enabled, }, truncate: { ...base.outputCompaction.truncate, enabled: overrides.truncateEnabled ?? base.outputCompaction.truncate.enabled, }, }, }; } runTest("bounded notice tracker evicts old entries and supports reset", () => { const tracker = createBoundedNoticeTracker(2); assert.equal(tracker.remember("first"), true); assert.equal(tracker.remember("second"), true); assert.equal(tracker.remember("first"), false); assert.equal(tracker.remember("third"), true); assert.equal(tracker.remember("second"), false); assert.equal(tracker.remember("first"), true); tracker.reset(); assert.equal(tracker.remember("third"), true); }); runTest("bounded notice tracker coerces invalid limits to a safe minimum", () => { const tracker = createBoundedNoticeTracker(0); assert.equal(tracker.remember("alpha"), true); assert.equal(tracker.remember("beta"), true); assert.equal(tracker.remember("alpha"), true); }); runTest("source-filter note injected when source filtering is active", () => { assert.equal( shouldInjectSourceFilterTroubleshootingNote( configWith({ readCompactionEnabled: true, sourceFilteringEnabled: true, sourceFilteringLevel: "minimal", smartTruncateEnabled: true, }), ), true, ); assert.equal( shouldInjectSourceFilterTroubleshootingNote( configWith({ readCompactionEnabled: true, sourceFilteringEnabled: true, sourceFilteringLevel: "aggressive", smartTruncateEnabled: true, }), ), true, ); }); runTest("source-filter note skipped when extension is disabled", () => { assert.equal(shouldInjectSourceFilterTroubleshootingNote(configWith({ enabled: false })), false); }); runTest("source-filter note skipped when compaction is disabled", () => { assert.equal(shouldInjectSourceFilterTroubleshootingNote(configWith({ compactionEnabled: false })), false); }); runTest("source-filter note skipped when read compaction is disabled", () => { assert.equal( shouldInjectSourceFilterTroubleshootingNote( configWith({ readCompactionEnabled: false, sourceFilteringEnabled: true, sourceFilteringLevel: "minimal", smartTruncateEnabled: true, }), ), false, ); }); runTest("source-filter note skipped when source filtering flag is off", () => { assert.equal( shouldInjectSourceFilterTroubleshootingNote(configWith({ sourceFilteringEnabled: false })), false, ); }); runTest("source-filter note skipped when filtering level is 'none'", () => { assert.equal( shouldInjectSourceFilterTroubleshootingNote( configWith({ sourceFilteringEnabled: true, sourceFilteringLevel: "none" }), ), false, ); }); runTest("source-filter note skipped when all read filtering safeguards are disabled", () => { assert.equal( shouldInjectSourceFilterTroubleshootingNote( configWith({ smartTruncateEnabled: false, truncateEnabled: false }), ), false, ); }); runTest("injectGuidelineIntoPrompt inserts bullet inside Guidelines section when header is at index 0", () => { const guideline = "Test guideline for RTK troubleshooting."; const prompt = [ "Guidelines:", "- Use mcp for MCP discovery first", "- Be concise in your responses", "", "", "Project-specific instructions", ].join("\n"); const result = injectGuidelineIntoPrompt(prompt, guideline); assert.ok(result.includes(guideline), "guideline must be present"); const bullet = `- ${guideline}`; assert.ok(result.includes(bullet), "guideline must be a bullet line"); const guidelinesStart = result.indexOf("Guidelines:\n"); const bulletIndex = result.indexOf(bullet); const projectContextIndex = result.indexOf("\n"); assert.equal(guidelinesStart, 0, "Guidelines header must be at index 0"); assert.ok(bulletIndex > guidelinesStart, "bullet must be after Guidelines header"); assert.ok(projectContextIndex !== -1, "project_context section must exist"); assert.ok(bulletIndex < projectContextIndex, "bullet must be before project_context section"); assert.equal(injectGuidelineIntoPrompt(result, guideline), result, "should be idempotent"); }); runTest("injectGuidelineIntoPrompt inserts bullet inside Guidelines section when header is mid-prompt", () => { const guideline = "Test guideline for RTK troubleshooting."; const prompt = [ "You are an expert coding assistant.", "", "Guidelines:", "- Be concise in your responses", "", "Pi documentation:", "- Main documentation: /path/to/readme", ].join("\n"); const result = injectGuidelineIntoPrompt(prompt, guideline); const bullet = `- ${guideline}`; assert.ok(result.includes(bullet), "guideline must be a bullet line"); const guidelinesStart = result.indexOf("\nGuidelines:\n"); const bulletIndex = result.indexOf(bullet); const piDocsIndex = result.indexOf("\nPi documentation:"); assert.ok(guidelinesStart !== -1, "Guidelines header must exist"); assert.ok(bulletIndex > guidelinesStart, "bullet must be after Guidelines header"); assert.ok(bulletIndex < piDocsIndex, "bullet must be before Pi documentation section"); }); runTest("injectGuidelineIntoPrompt falls back to appending when no Guidelines section exists", () => { const guideline = "Test guideline for RTK troubleshooting."; const prompt = "You are a coding assistant with no guidelines section."; const result = injectGuidelineIntoPrompt(prompt, guideline); assert.ok(result.includes(guideline), "guideline must be present"); assert.ok(result.endsWith(guideline), "guideline must be appended at the end"); }); await runTest("session_start refreshes RTK provenance and runtime guard skips missing rewrites", async () => { writeTestConfig(true); const handlers: Record = {}; const notifications: Notification[] = []; const execCommands: string[] = []; let rtkAvailable = false; let rewriteCalls = 0; rtkIntegrationExtension({ exec: async (command: string, args: string[]) => { execCommands.push(command); if (command === "which" || command === "where") { return { code: 0, stdout: "/opt/rtk/bin/rtk\n", stderr: "" }; } if (args[0] === "--version") { return rtkAvailable ? { code: 0, stdout: "rtk 1.0.0", stderr: "" } : { code: 1, stdout: "", stderr: "missing rtk" }; } if (args[0] === "rewrite") { rewriteCalls += 1; return { code: 3, stdout: "rtk git status", stderr: "" }; } return { code: 1, stdout: "", stderr: "unexpected" }; }, on(eventName: string, handler: ExtensionHandler) { handlers[eventName] = handler; }, registerCommand() {}, } as never); const sessionStartHandler = handlers.session_start; const toolCallHandler = handlers.tool_call; assert.ok(sessionStartHandler); assert.ok(toolCallHandler); await sessionStartHandler({}, createNotificationContext(notifications)); const skippedEvent = { toolName: "bash", input: { command: "git status" } }; await toolCallHandler(skippedEvent, createNotificationContext(notifications)); assert.equal((skippedEvent.input as { command: string }).command, "git status"); assert.equal(rewriteCalls, 0); assert.ok(notifications.some((notice) => notice.message.includes("rtk binary unavailable"))); rtkAvailable = true; await sessionStartHandler({}, createNotificationContext(notifications)); const rewrittenEvent = { toolName: "bash", input: { command: "git status" } }; await toolCallHandler(rewrittenEvent, createNotificationContext(notifications)); assert.equal(rewriteCalls, 1); assert.ok((rewrittenEvent.input as { command: string }).command.includes("rtk git status")); const remoteEvent = { toolName: "ssh_bash", input: { command: "git status" } }; await toolCallHandler(remoteEvent, createNotificationContext(notifications)); assert.equal(remoteEvent.input.command, "git status"); assert.equal(rewriteCalls, 1, "ssh_bash output support must not enable remote command rewriting"); assert.ok(execCommands.includes("/opt/rtk/bin/rtk")); writeTestConfig(false); }); await runTest("tool execution lifecycle sanitizes streamed bash output", async () => { const handlers: Record = {}; rtkIntegrationExtension({ exec: async () => ({ code: 0, stdout: "rtk 1.0.0", stderr: "" }), on(eventName: string, handler: ExtensionHandler) { handlers[eventName] = handler; }, registerCommand() {}, } as never); const startHandler = handlers.tool_execution_start; const updateHandler = handlers.tool_execution_update; const endHandler = handlers.tool_execution_end; assert.ok(startHandler); assert.ok(updateHandler); assert.ok(endHandler); await startHandler( { toolName: "bash", toolCallId: "bash-1", args: { command: "rtk git status" } }, {}, ); const updateEvent = { toolName: "bash", toolCallId: "bash-1", args: { command: "rtk git status" }, partialResult: { content: [ { type: "text", text: "\x1B[32mworking tree clean\x1B[0m\n", }, ], }, }; await updateHandler(updateEvent, {}); assert.equal(firstText(updateEvent.partialResult.content), "working tree clean\n"); const endEvent = { toolName: "bash", toolCallId: "bash-1", result: { content: [{ type: "text", text: "\x1B[31merror: build failed\x1B[0m\n" }] }, }; await endHandler(endEvent, {}); assert.equal(firstText(endEvent.result.content), "error: build failed\n"); await startHandler( { toolName: "ssh_bash", toolCallId: "ssh-bash-1", args: { command: "npm test" } }, {}, ); const remoteUpdate = { toolName: "ssh_bash", toolCallId: "ssh-bash-1", args: { command: "npm test" }, partialResult: { content: [{ type: "text", text: "\x1B[32mremote test\x1B[0m\n" }] }, }; await updateHandler(remoteUpdate, {}); assert.equal(firstText(remoteUpdate.partialResult.content), "remote test\n"); const remoteEnd = { toolName: "ssh_bash", toolCallId: "ssh-bash-1", result: { content: [{ type: "text", text: "\x1B[31mremote failed\x1B[0m\n" }] }, }; await endHandler(remoteEnd, {}); assert.equal(firstText(remoteEnd.result.content), "remote failed\n"); }); await runTest("tool_result lifecycle merges compaction metadata with existing details", async () => { const handlers: Record = {}; const notifications: Notification[] = []; rtkIntegrationExtension({ exec: async () => ({ code: 0, stdout: "rtk 1.0.0", stderr: "" }), on(eventName: string, handler: ExtensionHandler) { handlers[eventName] = handler; }, registerCommand() {}, } as never); const toolResultHandler = handlers.tool_result; assert.ok(toolResultHandler); const result = await toolResultHandler( { toolName: "bash", input: { command: "printf TODO" }, content: [{ type: "text", text: "\x1B[31msrc/a.ts\n 1: TODO\x1B[0m\n" }], details: { metadata: { requestId: "abc" }, traceId: "trace-1" }, }, createNotificationContext(notifications), ); assert.ok(result); assert.equal(firstText(result.content), "src/a.ts\n 1: TODO\n"); assert.equal((result.details as { traceId?: string }).traceId, "trace-1"); const details = result.details as { rtkCompaction?: { applied: boolean }; metadata?: Record }; assert.equal(details.rtkCompaction?.applied, true); assert.deepEqual(details.metadata?.requestId, "abc"); assert.equal((details.metadata?.rtkCompaction as { applied?: boolean } | undefined)?.applied, true); assert.equal(notifications.length, 0); }); await runTest("tool_call surfaces RTK rewrite errors through existing UI warning path", async () => { writeTestConfig(true); const handlers: Record, ctx: Record) => Promise | void>> = {}; const notifications: Notification[] = []; rtkIntegrationExtension({ exec: async (_command: string, args: string[]) => { if (args[0] === "--version") { return { code: 0, stdout: "rtk 1.0.0", stderr: "" }; } return { code: 2, stdout: "", stderr: "denied unsafe rewrite" }; }, on(eventName: string, handler: (event: Record, ctx: Record) => Promise | void>) { handlers[eventName] = handler; }, registerCommand() {}, } as never); const toolCallHandler = handlers.tool_call; assert.ok(toolCallHandler); const event = { toolName: "bash", input: { command: "git status" } }; await toolCallHandler(event, { hasUI: true, ui: { notify(message: string, level: "info" | "warning" | "error") { notifications.push({ message, level }); }, }, }); assert.equal((event.input as { command: string }).command, "git status"); assert.equal(notifications.length, 1); assert.equal(notifications[0]?.level, "warning"); assert.ok(notifications[0]?.message.includes("rtk rewrite skipped")); assert.ok(notifications[0]?.message.includes("denied unsafe rewrite")); writeTestConfig(false); }); console.log("All index tests passed.");