feat(agent): add initial harness foundation
This commit is contained in:
Mario Zechner
310
packages/agent/test/harness/compaction.test.ts
Normal file
310
packages/agent/test/harness/compaction.test.ts
Normal file
@@ -0,0 +1,310 @@
|
||||
import type { AgentMessage } from "@mariozechner/pi-agent-core";
|
||||
import {
|
||||
type AssistantMessage,
|
||||
type FauxProviderRegistration,
|
||||
fauxAssistantMessage,
|
||||
type Model,
|
||||
registerFauxProvider,
|
||||
type Usage,
|
||||
} from "@mariozechner/pi-ai";
|
||||
import { afterEach, beforeEach, describe, expect, it } from "vitest";
|
||||
import {
|
||||
calculateContextTokens,
|
||||
compact,
|
||||
DEFAULT_COMPACTION_SETTINGS,
|
||||
estimateContextTokens,
|
||||
findCutPoint,
|
||||
generateSummary,
|
||||
prepareCompaction,
|
||||
serializeConversation,
|
||||
shouldCompact,
|
||||
} from "../../src/harness/compaction.js";
|
||||
import { buildSessionContext } from "../../src/harness/session-tree.js";
|
||||
import type {
|
||||
CompactionEntry,
|
||||
CompactionSettings,
|
||||
MessageEntry,
|
||||
ModelChangeEntry,
|
||||
SessionTreeEntry,
|
||||
ThinkingLevelChangeEntry,
|
||||
} from "../../src/harness/types.js";
|
||||
|
||||
let nextId = 0;
|
||||
function createId(): string {
|
||||
return `entry-${nextId++}`;
|
||||
}
|
||||
|
||||
function createMockUsage(input: number, output: number, cacheRead = 0, cacheWrite = 0): Usage {
|
||||
return {
|
||||
input,
|
||||
output,
|
||||
cacheRead,
|
||||
cacheWrite,
|
||||
totalTokens: input + output + cacheRead + cacheWrite,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
}
|
||||
|
||||
function createUserMessage(text: string): AgentMessage {
|
||||
return {
|
||||
role: "user",
|
||||
content: [{ type: "text", text }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function createAssistantMessage(text: string, usage = createMockUsage(100, 50)): 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 createMessageEntry(message: AgentMessage, parentId: string | null = null): MessageEntry {
|
||||
return {
|
||||
type: "message",
|
||||
id: createId(),
|
||||
parentId,
|
||||
timestamp: new Date().toISOString(),
|
||||
message,
|
||||
};
|
||||
}
|
||||
|
||||
function createCompactionEntry(
|
||||
summary: string,
|
||||
firstKeptEntryId: string,
|
||||
parentId: string | null = null,
|
||||
): CompactionEntry {
|
||||
return {
|
||||
type: "compaction",
|
||||
id: createId(),
|
||||
parentId,
|
||||
timestamp: new Date().toISOString(),
|
||||
summary,
|
||||
firstKeptEntryId,
|
||||
tokensBefore: 1234,
|
||||
};
|
||||
}
|
||||
|
||||
function createThinkingLevelEntry(level: string, parentId: string | null = null): ThinkingLevelChangeEntry {
|
||||
return {
|
||||
type: "thinking_level_change",
|
||||
id: createId(),
|
||||
parentId,
|
||||
timestamp: new Date().toISOString(),
|
||||
thinkingLevel: level,
|
||||
};
|
||||
}
|
||||
|
||||
function createModelChangeEntry(provider: string, modelId: string, parentId: string | null = null): ModelChangeEntry {
|
||||
return {
|
||||
type: "model_change",
|
||||
id: createId(),
|
||||
parentId,
|
||||
timestamp: new Date().toISOString(),
|
||||
provider,
|
||||
modelId,
|
||||
};
|
||||
}
|
||||
|
||||
function createFauxModel(reasoning: boolean): { faux: FauxProviderRegistration; model: Model<string> } {
|
||||
const faux = registerFauxProvider({
|
||||
models: [
|
||||
{
|
||||
id: reasoning ? "reasoning-model" : "non-reasoning-model",
|
||||
reasoning,
|
||||
contextWindow: 200000,
|
||||
maxTokens: 8192,
|
||||
},
|
||||
],
|
||||
});
|
||||
fauxRegistrations.push(faux);
|
||||
return { faux, model: faux.getModel() };
|
||||
}
|
||||
|
||||
const fauxRegistrations: FauxProviderRegistration[] = [];
|
||||
|
||||
afterEach(() => {
|
||||
while (fauxRegistrations.length > 0) {
|
||||
fauxRegistrations.pop()?.unregister();
|
||||
}
|
||||
});
|
||||
|
||||
describe("harness compaction", () => {
|
||||
beforeEach(() => {
|
||||
nextId = 0;
|
||||
});
|
||||
|
||||
it("calculates total context tokens from usage", () => {
|
||||
expect(calculateContextTokens(createMockUsage(1000, 500, 200, 100))).toBe(1800);
|
||||
expect(calculateContextTokens(createMockUsage(0, 0, 0, 0))).toBe(0);
|
||||
});
|
||||
|
||||
it("checks compaction threshold", () => {
|
||||
const settings: CompactionSettings = {
|
||||
enabled: true,
|
||||
reserveTokens: 10000,
|
||||
keepRecentTokens: 20000,
|
||||
};
|
||||
expect(shouldCompact(95000, 100000, settings)).toBe(true);
|
||||
expect(shouldCompact(89000, 100000, settings)).toBe(false);
|
||||
expect(shouldCompact(95000, 100000, { ...settings, enabled: false })).toBe(false);
|
||||
});
|
||||
|
||||
it("finds a cut point based on token differences", () => {
|
||||
const entries: SessionTreeEntry[] = [];
|
||||
let parentId: string | null = null;
|
||||
for (let i = 0; i < 10; i++) {
|
||||
const user = createMessageEntry(createUserMessage(`User ${i}`), parentId);
|
||||
entries.push(user);
|
||||
const assistant = createMessageEntry(
|
||||
createAssistantMessage(`Assistant ${i}`, createMockUsage(0, 100, (i + 1) * 1000, 0)),
|
||||
user.id,
|
||||
);
|
||||
entries.push(assistant);
|
||||
parentId = assistant.id;
|
||||
}
|
||||
|
||||
const result = findCutPoint(entries, 0, entries.length, 2500);
|
||||
expect(entries[result.firstKeptEntryIndex]?.type).toBe("message");
|
||||
});
|
||||
|
||||
it("builds session context with a compaction entry", () => {
|
||||
const u1 = createMessageEntry(createUserMessage("1"));
|
||||
const a1 = createMessageEntry(createAssistantMessage("a"), u1.id);
|
||||
const u2 = createMessageEntry(createUserMessage("2"), a1.id);
|
||||
const a2 = createMessageEntry(createAssistantMessage("b"), u2.id);
|
||||
const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id, a2.id);
|
||||
const u3 = createMessageEntry(createUserMessage("3"), compaction.id);
|
||||
const a3 = createMessageEntry(createAssistantMessage("c"), u3.id);
|
||||
const loaded = buildSessionContext([u1, a1, u2, a2, compaction, u3, a3]);
|
||||
expect(loaded.messages).toHaveLength(5);
|
||||
expect(loaded.messages[0]?.role).toBe("compactionSummary");
|
||||
});
|
||||
|
||||
it("tracks model and thinking level changes in built context", () => {
|
||||
const user = createMessageEntry(createUserMessage("1"));
|
||||
const modelChange = createModelChangeEntry("openai", "gpt-4", user.id);
|
||||
const assistant = createMessageEntry(createAssistantMessage("a"), modelChange.id);
|
||||
const thinkingChange = createThinkingLevelEntry("high", assistant.id);
|
||||
const loaded = buildSessionContext([user, modelChange, assistant, thinkingChange]);
|
||||
expect(loaded.model).toEqual({ provider: "anthropic", modelId: "claude-sonnet-4-5" });
|
||||
expect(loaded.thinkingLevel).toBe("high");
|
||||
});
|
||||
|
||||
it("prepares compaction using the latest compaction summary as previousSummary", () => {
|
||||
const u1 = createMessageEntry(createUserMessage("user msg 1"));
|
||||
const a1 = createMessageEntry(createAssistantMessage("assistant msg 1"), u1.id);
|
||||
const u2 = createMessageEntry(createUserMessage("user msg 2"), a1.id);
|
||||
const a2 = createMessageEntry(createAssistantMessage("assistant msg 2", createMockUsage(5000, 1000)), u2.id);
|
||||
const compaction1 = createCompactionEntry("First summary", u2.id, a2.id);
|
||||
const u3 = createMessageEntry(createUserMessage("user msg 3"), compaction1.id);
|
||||
const a3 = createMessageEntry(createAssistantMessage("assistant msg 3", createMockUsage(8000, 2000)), u3.id);
|
||||
const pathEntries = [u1, a1, u2, a2, compaction1, u3, a3];
|
||||
const preparation = prepareCompaction(pathEntries, DEFAULT_COMPACTION_SETTINGS);
|
||||
expect(preparation).toBeDefined();
|
||||
expect(preparation?.previousSummary).toBe("First summary");
|
||||
expect(preparation?.firstKeptEntryId).toBeTruthy();
|
||||
expect(preparation?.tokensBefore).toBe(estimateContextTokens(buildSessionContext(pathEntries).messages).tokens);
|
||||
});
|
||||
|
||||
it("serializes conversation with truncated tool results", () => {
|
||||
const longContent = "x".repeat(5000);
|
||||
const messages = convertMessages([
|
||||
{
|
||||
role: "toolResult",
|
||||
toolCallId: "tc1",
|
||||
toolName: "read",
|
||||
content: [{ type: "text", text: longContent }],
|
||||
isError: false,
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
]);
|
||||
const result = serializeConversation(messages);
|
||||
expect(result).toContain("[Tool result]:");
|
||||
expect(result).toContain("[... 3000 more characters truncated]");
|
||||
});
|
||||
|
||||
it("passes reasoning through generateSummary only for reasoning models with thinking enabled", async () => {
|
||||
const messages: AgentMessage[] = [createUserMessage("Summarize this.")];
|
||||
const seenOptions: Array<Record<string, unknown> | undefined> = [];
|
||||
const { faux: fauxReasoning, model: reasoningModel } = createFauxModel(true);
|
||||
fauxReasoning.setResponses([
|
||||
(_context, options) => {
|
||||
seenOptions.push(options as Record<string, unknown> | undefined);
|
||||
return fauxAssistantMessage("## Goal\nTest summary");
|
||||
},
|
||||
]);
|
||||
await generateSummary(
|
||||
messages,
|
||||
reasoningModel,
|
||||
2000,
|
||||
"test-key",
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
"medium",
|
||||
);
|
||||
expect(seenOptions[0]).toMatchObject({ reasoning: "medium", apiKey: "test-key" });
|
||||
|
||||
const { faux: fauxOff, model: offModel } = createFauxModel(true);
|
||||
fauxOff.setResponses([
|
||||
(_context, options) => {
|
||||
seenOptions.push(options as Record<string, unknown> | undefined);
|
||||
return fauxAssistantMessage("## Goal\nTest summary");
|
||||
},
|
||||
]);
|
||||
await generateSummary(messages, offModel, 2000, "test-key", undefined, undefined, undefined, undefined, "off");
|
||||
expect(seenOptions[1]).not.toHaveProperty("reasoning");
|
||||
|
||||
const { faux: fauxNonReasoning, model: nonReasoningModel } = createFauxModel(false);
|
||||
fauxNonReasoning.setResponses([
|
||||
(_context, options) => {
|
||||
seenOptions.push(options as Record<string, unknown> | undefined);
|
||||
return fauxAssistantMessage("## Goal\nTest summary");
|
||||
},
|
||||
]);
|
||||
await generateSummary(
|
||||
messages,
|
||||
nonReasoningModel,
|
||||
2000,
|
||||
"test-key",
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
"medium",
|
||||
);
|
||||
expect(seenOptions[2]).not.toHaveProperty("reasoning");
|
||||
});
|
||||
|
||||
it("returns a compaction result with file details", async () => {
|
||||
const u1 = createMessageEntry(createUserMessage("read a file"));
|
||||
const assistantMessage: AssistantMessage = {
|
||||
...createAssistantMessage("calling tool", createMockUsage(1000, 200)),
|
||||
content: [{ type: "toolCall", id: "tool-1", name: "read", arguments: { path: "src/index.ts" } }],
|
||||
};
|
||||
const a1 = createMessageEntry(assistantMessage, u1.id);
|
||||
const u2 = createMessageEntry(createUserMessage("continue"), a1.id);
|
||||
const a2 = createMessageEntry(createAssistantMessage("done", createMockUsage(4000, 500)), u2.id);
|
||||
const preparation = prepareCompaction([u1, a1, u2, a2], DEFAULT_COMPACTION_SETTINGS);
|
||||
expect(preparation).toBeDefined();
|
||||
const { faux, model } = createFauxModel(false);
|
||||
faux.setResponses([fauxAssistantMessage("## Goal\nTest summary")]);
|
||||
const result = await compact(preparation!, model, "test-key");
|
||||
expect(result.summary.length).toBeGreaterThan(0);
|
||||
expect(result.firstKeptEntryId).toBeTruthy();
|
||||
expect(result.details).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
function convertMessages(messages: any[]): any[] {
|
||||
return messages;
|
||||
}
|
||||
174
packages/agent/test/harness/session-tree.test.ts
Normal file
174
packages/agent/test/harness/session-tree.test.ts
Normal file
@@ -0,0 +1,174 @@
|
||||
import { existsSync, mkdirSync, readFileSync, rmSync } from "node:fs";
|
||||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import type { AgentMessage } from "@mariozechner/pi-agent-core";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import {
|
||||
DefaultSessionTree,
|
||||
InMemorySessionTreeStorage,
|
||||
JsonlSessionTreeStorage,
|
||||
} from "../../src/harness/session-tree.js";
|
||||
import type { SessionTreeStorage } from "../../src/harness/types.js";
|
||||
|
||||
function createUserMessage(text: string): AgentMessage {
|
||||
return {
|
||||
role: "user",
|
||||
content: [{ type: "text", text }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function createAssistantMessage(text: string): AgentMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text }],
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
model: "claude-sonnet-4-5",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
const tempDirs: string[] = [];
|
||||
|
||||
afterEach(() => {
|
||||
while (tempDirs.length > 0) {
|
||||
const dir = tempDirs.pop()!;
|
||||
if (existsSync(dir)) {
|
||||
rmSync(dir, { recursive: true, force: true });
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
async function runSessionTreeSuite(name: string, createStorage: () => SessionTreeStorage, inspect?: () => void) {
|
||||
describe(name, () => {
|
||||
it("appends messages and builds context in order", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendMessage(createAssistantMessage("two"));
|
||||
const context = await tree.buildContext();
|
||||
expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
|
||||
});
|
||||
|
||||
it("tracks model and thinking level changes", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendModelChange("openai", "gpt-4.1");
|
||||
await tree.appendThinkingLevelChange("high");
|
||||
const context = await tree.buildContext();
|
||||
expect(context.thinkingLevel).toBe("high");
|
||||
expect(context.model).toEqual({ provider: "openai", modelId: "gpt-4.1" });
|
||||
});
|
||||
|
||||
it("supports branching by moving the leaf and appending a new branch", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
const user1 = await tree.appendMessage(createUserMessage("one"));
|
||||
const assistant1 = await tree.appendMessage(createAssistantMessage("two"));
|
||||
await tree.appendMessage(createUserMessage("three"));
|
||||
await tree.moveTo(user1);
|
||||
await tree.appendMessage(createAssistantMessage("branched"));
|
||||
const branch = await tree.getBranch();
|
||||
expect(branch.map((entry) => entry.id)).toContain(user1);
|
||||
expect(branch.map((entry) => entry.id)).not.toContain(assistant1);
|
||||
const context = await tree.buildContext();
|
||||
expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
|
||||
});
|
||||
|
||||
it("supports moving the leaf to root", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.moveTo(null);
|
||||
expect(await tree.getLeafId()).toBeNull();
|
||||
expect((await tree.buildContext()).messages).toEqual([]);
|
||||
});
|
||||
|
||||
it("reconstructs compaction summaries in context", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendMessage(createAssistantMessage("two"));
|
||||
const user2 = await tree.appendMessage(createUserMessage("three"));
|
||||
await tree.appendMessage(createAssistantMessage("four"));
|
||||
await tree.appendCompaction("summary", user2, 1234);
|
||||
await tree.appendMessage(createUserMessage("five"));
|
||||
const context = await tree.buildContext();
|
||||
expect(context.messages[0]?.role).toBe("compactionSummary");
|
||||
expect(context.messages).toHaveLength(4);
|
||||
});
|
||||
|
||||
it("supports branch summary entries in context", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
const user1 = await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendBranchSummary(user1, "summary text");
|
||||
const context = await tree.buildContext();
|
||||
expect(context.messages[1]?.role).toBe("branchSummary");
|
||||
});
|
||||
|
||||
it("supports custom message entries in context", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendCustomMessageEntry("custom", "hello", true, { ok: true });
|
||||
const context = await tree.buildContext();
|
||||
expect(context.messages[1]?.role).toBe("custom");
|
||||
});
|
||||
|
||||
it("supports labels and session info entries without affecting context", async () => {
|
||||
const tree = new DefaultSessionTree(createStorage());
|
||||
const user1 = await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendLabelChange(user1, "checkpoint");
|
||||
await tree.appendSessionInfo("name");
|
||||
const entries = await tree.getEntries();
|
||||
expect(entries.some((entry) => entry.type === "label")).toBe(true);
|
||||
expect(entries.some((entry) => entry.type === "session_info")).toBe(true);
|
||||
expect((await tree.buildContext()).messages).toHaveLength(1);
|
||||
});
|
||||
|
||||
it("persists leaf changes and appended entries via storage", async () => {
|
||||
const storage = createStorage();
|
||||
const tree = new DefaultSessionTree(storage);
|
||||
const user1 = await tree.appendMessage(createUserMessage("one"));
|
||||
await tree.appendMessage(createAssistantMessage("two"));
|
||||
await tree.moveTo(user1);
|
||||
await tree.appendMessage(createAssistantMessage("branched"));
|
||||
const tree2 = new DefaultSessionTree(storage);
|
||||
const context = await tree2.buildContext();
|
||||
expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
|
||||
inspect?.();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
runSessionTreeSuite("SessionTree with in-memory storage", () => new InMemorySessionTreeStorage());
|
||||
|
||||
runSessionTreeSuite(
|
||||
"SessionTree with JSONL storage",
|
||||
() => {
|
||||
const dir = join(tmpdir(), `pi-agent-session-tree-${Date.now()}-${Math.random().toString(36).slice(2)}`);
|
||||
mkdirSync(dir, { recursive: true });
|
||||
tempDirs.push(dir);
|
||||
return new JsonlSessionTreeStorage(join(dir, "session.jsonl"), { cwd: dir });
|
||||
},
|
||||
() => {
|
||||
const dir = tempDirs[tempDirs.length - 1]!;
|
||||
const filePath = join(dir, "session.jsonl");
|
||||
const lines = readFileSync(filePath, "utf8").trim().split("\n");
|
||||
expect(lines.length).toBeGreaterThan(1);
|
||||
const header = JSON.parse(lines[0]!);
|
||||
expect(header.type).toBe("session");
|
||||
expect(header.version).toBe(3);
|
||||
for (const line of lines.slice(1)) {
|
||||
const entry = JSON.parse(line);
|
||||
expect(entry.type).not.toBe("entry");
|
||||
expect(entry.type).not.toBe("leaf");
|
||||
expect(typeof entry.id).toBe("string");
|
||||
}
|
||||
},
|
||||
);
|
||||
Reference in New Issue
Block a user