From c3f63dd243db3bac0e246b075c4c49ca52eea3c1 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Thu, 4 Dec 2025 02:49:22 +0100 Subject: [PATCH] feat(coding-agent): add auto-compaction to RPC mode, add RPC compaction test - RPC mode now auto-compacts when context exceeds threshold (same as TUI) - Add RPC test for manual compaction via compact command - Auto-compaction emits compaction event with auto: true flag --- packages/coding-agent/src/main.ts | 64 ++++++++++++++++- packages/coding-agent/test/rpc.test.ts | 95 ++++++++++++++++++++++++++ 2 files changed, 157 insertions(+), 2 deletions(-) diff --git a/packages/coding-agent/src/main.ts b/packages/coding-agent/src/main.ts index ef5002a7..e6e7431f 100644 --- a/packages/coding-agent/src/main.ts +++ b/packages/coding-agent/src/main.ts @@ -1,12 +1,12 @@ import { Agent, type Attachment, ProviderTransport, type ThinkingLevel } from "@mariozechner/pi-agent-core"; -import type { Api, KnownProvider, Model } from "@mariozechner/pi-ai"; +import type { Api, AssistantMessage, KnownProvider, Model } from "@mariozechner/pi-ai"; import { ProcessTerminal, TUI } from "@mariozechner/pi-tui"; import chalk from "chalk"; import { existsSync, readFileSync, statSync } from "fs"; import { homedir } from "os"; import { extname, join, resolve } from "path"; import { getChangelogPath, getNewEntries, parseChangelog } from "./changelog.js"; -import { compact } from "./compaction.js"; +import { calculateContextTokens, compact, shouldCompact } from "./compaction.js"; import { APP_NAME, CONFIG_DIR_NAME, @@ -820,6 +820,61 @@ async function runRpcMode( sessionManager: SessionManager, settingsManager: SettingsManager, ): Promise { + // Track if auto-compaction is in progress + let autoCompactionInProgress = false; + + // Auto-compaction helper + const checkAutoCompaction = async () => { + if (autoCompactionInProgress) return; + + const settings = settingsManager.getCompactionSettings(); + if (!settings.enabled) return; + + // Get last non-aborted assistant message + const messages = agent.state.messages; + let lastAssistant: AssistantMessage | null = null; + for (let i = messages.length - 1; i >= 0; i--) { + const msg = messages[i]; + if (msg.role === "assistant") { + const assistantMsg = msg as AssistantMessage; + if (assistantMsg.stopReason !== "aborted") { + lastAssistant = assistantMsg; + break; + } + } + } + if (!lastAssistant) return; + + const contextTokens = calculateContextTokens(lastAssistant.usage); + const contextWindow = agent.state.model.contextWindow; + + if (!shouldCompact(contextTokens, contextWindow, settings)) return; + + // Trigger auto-compaction + autoCompactionInProgress = true; + try { + const apiKey = await getApiKeyForModel(agent.state.model); + if (!apiKey) { + throw new Error(`No API key for ${agent.state.model.provider}`); + } + + const entries = sessionManager.loadEntries(); + const compactionEntry = await compact(entries, agent.state.model, settings, apiKey); + + sessionManager.saveCompaction(compactionEntry); + const loaded = loadSessionFromEntries(sessionManager.loadEntries()); + agent.replaceMessages(loaded.messages); + + // Emit auto-compaction event + console.log(JSON.stringify({ ...compactionEntry, auto: true })); + } catch (error: unknown) { + const message = error instanceof Error ? error.message : String(error); + console.log(JSON.stringify({ type: "error", error: `Auto-compaction failed: ${message}` })); + } finally { + autoCompactionInProgress = false; + } + }; + // Subscribe to all events and output as JSON (same pattern as tui-renderer) agent.subscribe(async (event) => { console.log(JSON.stringify(event)); @@ -836,6 +891,11 @@ async function runRpcMode( if (sessionManager.shouldInitializeSession(agent.state.messages)) { sessionManager.startSession(agent.state); } + + // Check for auto-compaction after assistant messages + if (event.message.role === "assistant") { + await checkAutoCompaction(); + } } }); diff --git a/packages/coding-agent/test/rpc.test.ts b/packages/coding-agent/test/rpc.test.ts index cd13a469..377695e6 100644 --- a/packages/coding-agent/test/rpc.test.ts +++ b/packages/coding-agent/test/rpc.test.ts @@ -6,6 +6,7 @@ import * as readline from "node:readline"; import { fileURLToPath } from "node:url"; import type { AgentEvent } from "@mariozechner/pi-agent-core"; import { afterEach, beforeEach, describe, expect, test } from "vitest"; +import type { CompactionEntry } from "../src/session-manager.js"; const __dirname = dirname(fileURLToPath(import.meta.url)); @@ -135,4 +136,98 @@ describe.skipIf(!process.env.ANTHROPIC_API_KEY && !process.env.ANTHROPIC_OAUTH_T expect(roles).toContain("user"); expect(roles).toContain("assistant"); }, 90000); + + test("should handle manual compaction", async () => { + // Spawn agent in RPC mode + agent = spawn( + "node", + ["dist/cli.js", "--mode", "rpc", "--provider", "anthropic", "--model", "claude-sonnet-4-5"], + { + cwd: join(__dirname, ".."), + env: { + ...process.env, + PI_CODING_AGENT_DIR: sessionDir, + }, + }, + ); + + const events: (AgentEvent | CompactionEntry | { type: "error"; error: string })[] = []; + + const rl = readline.createInterface({ input: agent.stdout!, terminal: false }); + + let stderr = ""; + agent.stderr?.on("data", (data) => { + stderr += data.toString(); + }); + + // Helper to wait for a specific event type + const waitForEvent = (eventType: string, timeout = 60000) => + new Promise((resolve, reject) => { + const timer = setTimeout(() => reject(new Error(`Timeout waiting for ${eventType}`)), timeout); + + const checkExisting = () => { + if (events.some((e) => e.type === eventType)) { + clearTimeout(timer); + resolve(); + return true; + } + return false; + }; + + if (checkExisting()) return; + + const handler = (line: string) => { + try { + const event = JSON.parse(line); + events.push(event); + if (event.type === eventType) { + clearTimeout(timer); + rl.off("line", handler); + resolve(); + } + } catch { + // Ignore non-JSON + } + }; + rl.on("line", handler); + }); + + // First, send a prompt to have some messages to compact + agent.stdin!.write(JSON.stringify({ type: "prompt", message: "Say hello" }) + "\n"); + await waitForEvent("agent_end"); + + // Clear events to focus on compaction + events.length = 0; + + // Send compact command + agent.stdin!.write(JSON.stringify({ type: "compact" }) + "\n"); + await waitForEvent("compaction"); + + // Verify compaction event + const compactionEvent = events.find((e) => e.type === "compaction") as CompactionEntry | undefined; + expect(compactionEvent).toBeDefined(); + expect(compactionEvent!.summary).toBeDefined(); + expect(compactionEvent!.tokensBefore).toBeGreaterThan(0); + + // Wait for file writes + await new Promise((resolve) => setTimeout(resolve, 200)); + + agent.kill("SIGTERM"); + + // Verify compaction was saved to session file + const sessionsPath = join(sessionDir, "sessions"); + const sessionDirs = readdirSync(sessionsPath); + const cwdSessionDir = join(sessionsPath, sessionDirs[0]); + const sessionFiles = readdirSync(cwdSessionDir).filter((f) => f.endsWith(".jsonl")); + const sessionContent = readFileSync(join(cwdSessionDir, sessionFiles[0]), "utf8"); + const entries = sessionContent + .trim() + .split("\n") + .map((line) => JSON.parse(line)); + + // Should have a compaction entry + const compactionEntries = entries.filter((e: { type: string }) => e.type === "compaction"); + expect(compactionEntries.length).toBe(1); + expect(compactionEntries[0].summary).toBeDefined(); + }, 120000); });