mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-21 16:01:05 +00:00
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
This commit is contained in:
parent
1a97331af1
commit
c3f63dd243
2 changed files with 157 additions and 2 deletions
|
|
@ -1,12 +1,12 @@
|
||||||
import { Agent, type Attachment, ProviderTransport, type ThinkingLevel } from "@mariozechner/pi-agent-core";
|
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 { ProcessTerminal, TUI } from "@mariozechner/pi-tui";
|
||||||
import chalk from "chalk";
|
import chalk from "chalk";
|
||||||
import { existsSync, readFileSync, statSync } from "fs";
|
import { existsSync, readFileSync, statSync } from "fs";
|
||||||
import { homedir } from "os";
|
import { homedir } from "os";
|
||||||
import { extname, join, resolve } from "path";
|
import { extname, join, resolve } from "path";
|
||||||
import { getChangelogPath, getNewEntries, parseChangelog } from "./changelog.js";
|
import { getChangelogPath, getNewEntries, parseChangelog } from "./changelog.js";
|
||||||
import { compact } from "./compaction.js";
|
import { calculateContextTokens, compact, shouldCompact } from "./compaction.js";
|
||||||
import {
|
import {
|
||||||
APP_NAME,
|
APP_NAME,
|
||||||
CONFIG_DIR_NAME,
|
CONFIG_DIR_NAME,
|
||||||
|
|
@ -820,6 +820,61 @@ async function runRpcMode(
|
||||||
sessionManager: SessionManager,
|
sessionManager: SessionManager,
|
||||||
settingsManager: SettingsManager,
|
settingsManager: SettingsManager,
|
||||||
): Promise<void> {
|
): Promise<void> {
|
||||||
|
// 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)
|
// Subscribe to all events and output as JSON (same pattern as tui-renderer)
|
||||||
agent.subscribe(async (event) => {
|
agent.subscribe(async (event) => {
|
||||||
console.log(JSON.stringify(event));
|
console.log(JSON.stringify(event));
|
||||||
|
|
@ -836,6 +891,11 @@ async function runRpcMode(
|
||||||
if (sessionManager.shouldInitializeSession(agent.state.messages)) {
|
if (sessionManager.shouldInitializeSession(agent.state.messages)) {
|
||||||
sessionManager.startSession(agent.state);
|
sessionManager.startSession(agent.state);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Check for auto-compaction after assistant messages
|
||||||
|
if (event.message.role === "assistant") {
|
||||||
|
await checkAutoCompaction();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -6,6 +6,7 @@ import * as readline from "node:readline";
|
||||||
import { fileURLToPath } from "node:url";
|
import { fileURLToPath } from "node:url";
|
||||||
import type { AgentEvent } from "@mariozechner/pi-agent-core";
|
import type { AgentEvent } from "@mariozechner/pi-agent-core";
|
||||||
import { afterEach, beforeEach, describe, expect, test } from "vitest";
|
import { afterEach, beforeEach, describe, expect, test } from "vitest";
|
||||||
|
import type { CompactionEntry } from "../src/session-manager.js";
|
||||||
|
|
||||||
const __dirname = dirname(fileURLToPath(import.meta.url));
|
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("user");
|
||||||
expect(roles).toContain("assistant");
|
expect(roles).toContain("assistant");
|
||||||
}, 90000);
|
}, 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<void>((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);
|
||||||
});
|
});
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue