refactor(agent): expose concrete harness session

This commit is contained in:
Mario Zechner
2026-05-06 12:57:19 +02:00
parent 0a032ae1b3
commit 530f14c018
7 changed files with 25 additions and 64 deletions

View File

@@ -1,7 +1,7 @@
import { randomUUID } from "node:crypto"; import { randomUUID } from "node:crypto";
import { readFileSync } from "node:fs"; import { readFileSync } from "node:fs";
import type { ImageContent, Model } from "@mariozechner/pi-ai"; import type { ImageContent, Model } from "@mariozechner/pi-ai";
import type { Agent } from "../agent.js"; import { Agent } from "../agent.js";
import type { AgentEvent, AgentMessage, AgentTool, ThinkingLevel } from "../types.js"; import type { AgentEvent, AgentMessage, AgentTool, ThinkingLevel } from "../types.js";
import { collectEntriesForBranchSummary, generateBranchSummary } from "./compaction/branch-summarization.js"; import { collectEntriesForBranchSummary, generateBranchSummary } from "./compaction/branch-summarization.js";
import { compact, DEFAULT_COMPACTION_SETTINGS, prepareCompaction } from "./compaction/compaction.js"; import { compact, DEFAULT_COMPACTION_SETTINGS, prepareCompaction } from "./compaction/compaction.js";
@@ -68,7 +68,7 @@ export class DefaultAgentHarness implements AgentHarness {
>(); >();
constructor(options: AgentHarnessOptions) { constructor(options: AgentHarnessOptions) {
this.agent = options.agent; this.agent = new Agent({});
this.env = options.env; this.env = options.env;
this.session = options.session; this.session = options.session;
this.promptTemplates = options.promptTemplates ?? []; this.promptTemplates = options.promptTemplates ?? [];

View File

@@ -1,11 +1,11 @@
import { DefaultAgentHarness } from "./agent-harness.js"; import { DefaultAgentHarness } from "./agent-harness.js";
import { DefaultSession } from "./session/session.js"; import { Session } from "./session/session.js";
import type { AgentHarness, AgentHarnessOptions, Session, SessionMetadata, SessionStorage } from "./types.js"; import type { AgentHarness, AgentHarnessOptions, SessionMetadata, SessionStorage } from "./types.js";
export function createSession<TMetadata extends SessionMetadata>( export function createSession<TMetadata extends SessionMetadata>(
storage: SessionStorage<TMetadata>, storage: SessionStorage<TMetadata>,
): Session<TMetadata> { ): Session<TMetadata> {
return new DefaultSession(storage); return new Session(storage);
} }
export function createAgentHarness(options: AgentHarnessOptions): AgentHarness { export function createAgentHarness(options: AgentHarnessOptions): AgentHarness {

View File

@@ -1,6 +1,6 @@
import { v7 as uuidv7 } from "uuid"; import { v7 as uuidv7 } from "uuid";
import type { Session, SessionMetadata, SessionStorage, SessionTreeEntry } from "../../types.js"; import type { SessionMetadata, SessionStorage, SessionTreeEntry } from "../../types.js";
import { DefaultSession } from "../session.js"; import { Session } from "../session.js";
export function createSessionId(): string { export function createSessionId(): string {
return uuidv7(); return uuidv7();
@@ -11,7 +11,7 @@ export function createTimestamp(): string {
} }
export function toSession<TMetadata extends SessionMetadata>(storage: SessionStorage<TMetadata>): Session<TMetadata> { export function toSession<TMetadata extends SessionMetadata>(storage: SessionStorage<TMetadata>): Session<TMetadata> {
return new DefaultSession(storage); return new Session(storage);
} }
export async function getEntriesToFork( export async function getEntriesToFork(

View File

@@ -9,7 +9,6 @@ import type {
LabelEntry, LabelEntry,
MessageEntry, MessageEntry,
ModelChangeEntry, ModelChangeEntry,
Session,
SessionContext, SessionContext,
SessionInfoEntry, SessionInfoEntry,
SessionMetadata, SessionMetadata,
@@ -75,7 +74,7 @@ export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionCon
return { messages, thinkingLevel, model }; return { messages, thinkingLevel, model };
} }
export class DefaultSession<TMetadata extends SessionMetadata = SessionMetadata> implements Session<TMetadata> { export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
private storage: SessionStorage<TMetadata>; private storage: SessionStorage<TMetadata>;
constructor(storage: SessionStorage<TMetadata>) { constructor(storage: SessionStorage<TMetadata>) {

View File

@@ -1,5 +1,6 @@
import type { ImageContent, Model, TextContent } from "@mariozechner/pi-ai"; import type { ImageContent, Model, TextContent } from "@mariozechner/pi-ai";
import type { Agent, AgentEvent, AgentMessage, ThinkingLevel } from "../index.js"; import type { Agent, AgentEvent, AgentMessage, ThinkingLevel } from "../index.js";
import type { Session } from "./session/session.js";
export type SourceScope = "user" | "project" | "temporary"; export type SourceScope = "user" | "project" | "temporary";
export type SourceOrigin = "package" | "top-level"; export type SourceOrigin = "package" | "top-level";
@@ -183,43 +184,7 @@ export interface SessionStorage<TMetadata extends SessionMetadata = SessionMetad
getEntries(): Promise<SessionTreeEntry[]>; getEntries(): Promise<SessionTreeEntry[]>;
} }
export interface Session<TMetadata extends SessionMetadata = SessionMetadata> { export type { Session } from "./session/session.js";
getMetadata(): Promise<TMetadata>;
getStorage(): SessionStorage<TMetadata>;
getLeafId(): Promise<string | null>;
getEntry(id: string): Promise<SessionTreeEntry | undefined>;
getEntries(): Promise<SessionTreeEntry[]>;
getBranch(fromId?: string): Promise<SessionTreeEntry[]>;
buildContext(): Promise<SessionContext>;
getLabel(id: string): Promise<string | undefined>;
getSessionName(): Promise<string | undefined>;
appendMessage(message: AgentMessage): Promise<string>;
appendThinkingLevelChange(thinkingLevel: string): Promise<string>;
appendModelChange(provider: string, modelId: string): Promise<string>;
appendCompaction<T = unknown>(
summary: string,
firstKeptEntryId: string,
tokensBefore: number,
details?: T,
fromHook?: boolean,
): Promise<string>;
appendCustomEntry(customType: string, data?: unknown): Promise<string>;
appendCustomMessageEntry<T = unknown>(
customType: string,
content: string | (TextContent | ImageContent)[],
display: boolean,
details?: T,
): Promise<string>;
appendLabel(targetId: string, label: string | undefined): Promise<string>;
appendSessionName(name: string): Promise<string>;
moveTo(
entryId: string | null,
summary?: { summary: string; details?: unknown; fromHook?: boolean },
): Promise<string | undefined>;
}
export interface SessionCreateOptions { export interface SessionCreateOptions {
id?: string; id?: string;
@@ -554,7 +519,6 @@ export interface BranchSummaryResult {
} }
export interface AgentHarnessOptions { export interface AgentHarnessOptions {
agent: Agent;
env: ExecutionEnv; env: ExecutionEnv;
session: Session; session: Session;
promptTemplates?: PromptTemplate[]; promptTemplates?: PromptTemplate[];

View File

@@ -1,5 +1,4 @@
import { describe, expect, it } from "vitest"; import { describe, expect, it } from "vitest";
import { Agent } from "../../src/agent.js";
import { NodeExecutionEnv } from "../../src/harness/execution-env.js"; import { NodeExecutionEnv } from "../../src/harness/execution-env.js";
import { createAgentHarness, createSession } from "../../src/harness/factory.js"; import { createAgentHarness, createSession } from "../../src/harness/factory.js";
import { InMemorySessionStorage } from "../../src/harness/session/storage/memory.js"; import { InMemorySessionStorage } from "../../src/harness/session/storage/memory.js";
@@ -17,9 +16,8 @@ describe("harness factories", () => {
it("creates agent harnesses", () => { it("creates agent harnesses", () => {
const session = createSession(new InMemorySessionStorage()); const session = createSession(new InMemorySessionStorage());
const env = new NodeExecutionEnv({ cwd: process.cwd() }); const env = new NodeExecutionEnv({ cwd: process.cwd() });
const agent = new Agent(); const harness = createAgentHarness({ env, session });
const harness = createAgentHarness({ agent, env, session }); expect(harness.env).toBe(env);
expect(harness.agent).toBe(agent);
expect(harness.conversation.session).toBe(session); expect(harness.conversation.session).toBe(session);
}); });
}); });

View File

@@ -1,7 +1,7 @@
import { readFileSync } from "node:fs"; import { readFileSync } from "node:fs";
import { join } from "node:path"; import { join } from "node:path";
import { describe, expect, it } from "vitest"; import { describe, expect, it } from "vitest";
import { DefaultSession } from "../../src/harness/session/session.js"; import { Session } from "../../src/harness/session/session.js";
import { JsonlSessionStorage } from "../../src/harness/session/storage/jsonl.js"; import { JsonlSessionStorage } from "../../src/harness/session/storage/jsonl.js";
import { InMemorySessionStorage } from "../../src/harness/session/storage/memory.js"; import { InMemorySessionStorage } from "../../src/harness/session/storage/memory.js";
import type { SessionStorage } from "../../src/harness/types.js"; import type { SessionStorage } from "../../src/harness/types.js";
@@ -14,7 +14,7 @@ async function runSessionSuite(
) { ) {
describe(name, () => { describe(name, () => {
it("appends messages and builds context in order", async () => { it("appends messages and builds context in order", async () => {
const session = new DefaultSession(await createStorage()); const session = new Session(await createStorage());
await session.appendMessage(createUserMessage("one")); await session.appendMessage(createUserMessage("one"));
await session.appendMessage(createAssistantMessage("two")); await session.appendMessage(createAssistantMessage("two"));
const context = await session.buildContext(); const context = await session.buildContext();
@@ -22,7 +22,7 @@ async function runSessionSuite(
}); });
it("tracks model and thinking level changes", async () => { it("tracks model and thinking level changes", async () => {
const session = new DefaultSession(await createStorage()); const session = new Session(await createStorage());
await session.appendMessage(createUserMessage("one")); await session.appendMessage(createUserMessage("one"));
await session.appendModelChange("openai", "gpt-4.1"); await session.appendModelChange("openai", "gpt-4.1");
await session.appendThinkingLevelChange("high"); await session.appendThinkingLevelChange("high");
@@ -32,7 +32,7 @@ async function runSessionSuite(
}); });
it("supports branching by moving the leaf and appending a new branch", async () => { it("supports branching by moving the leaf and appending a new branch", async () => {
const session = new DefaultSession(await createStorage()); const session = new Session(await createStorage());
const user1 = await session.appendMessage(createUserMessage("one")); const user1 = await session.appendMessage(createUserMessage("one"));
const assistant1 = await session.appendMessage(createAssistantMessage("two")); const assistant1 = await session.appendMessage(createAssistantMessage("two"));
await session.appendMessage(createUserMessage("three")); await session.appendMessage(createUserMessage("three"));
@@ -46,7 +46,7 @@ async function runSessionSuite(
}); });
it("supports moving the leaf to root", async () => { it("supports moving the leaf to root", async () => {
const session = new DefaultSession(await createStorage()); const session = new Session(await createStorage());
await session.appendMessage(createUserMessage("one")); await session.appendMessage(createUserMessage("one"));
await session.moveTo(null); await session.moveTo(null);
expect(await session.getLeafId()).toBeNull(); expect(await session.getLeafId()).toBeNull();
@@ -54,7 +54,7 @@ async function runSessionSuite(
}); });
it("reconstructs compaction summaries in context", async () => { it("reconstructs compaction summaries in context", async () => {
const session = new DefaultSession(await createStorage()); const session = new Session(await createStorage());
await session.appendMessage(createUserMessage("one")); await session.appendMessage(createUserMessage("one"));
await session.appendMessage(createAssistantMessage("two")); await session.appendMessage(createAssistantMessage("two"));
const user2 = await session.appendMessage(createUserMessage("three")); const user2 = await session.appendMessage(createUserMessage("three"));
@@ -67,7 +67,7 @@ async function runSessionSuite(
}); });
it("supports moving with branch summary entries in context", async () => { it("supports moving with branch summary entries in context", async () => {
const session = new DefaultSession(await createStorage()); const session = new Session(await createStorage());
const user1 = await session.appendMessage(createUserMessage("one")); const user1 = await session.appendMessage(createUserMessage("one"));
const summaryId = await session.moveTo(user1, { summary: "summary text" }); const summaryId = await session.moveTo(user1, { summary: "summary text" });
expect(summaryId).toBeTruthy(); expect(summaryId).toBeTruthy();
@@ -78,7 +78,7 @@ async function runSessionSuite(
}); });
it("supports custom message entries in context", async () => { it("supports custom message entries in context", async () => {
const session = new DefaultSession(await createStorage()); const session = new Session(await createStorage());
await session.appendMessage(createUserMessage("one")); await session.appendMessage(createUserMessage("one"));
await session.appendCustomMessageEntry("custom", "hello", true, { ok: true }); await session.appendCustomMessageEntry("custom", "hello", true, { ok: true });
const context = await session.buildContext(); const context = await session.buildContext();
@@ -86,7 +86,7 @@ async function runSessionSuite(
}); });
it("supports labels and session info entries without affecting context", async () => { it("supports labels and session info entries without affecting context", async () => {
const session = new DefaultSession(await createStorage()); const session = new Session(await createStorage());
const user1 = await session.appendMessage(createUserMessage("one")); const user1 = await session.appendMessage(createUserMessage("one"));
await session.appendLabel(user1, "checkpoint"); await session.appendLabel(user1, "checkpoint");
await session.appendSessionName("name"); await session.appendSessionName("name");
@@ -99,20 +99,20 @@ async function runSessionSuite(
}); });
it("rejects labels for missing entries", async () => { it("rejects labels for missing entries", async () => {
const session = new DefaultSession(await createStorage()); const session = new Session(await createStorage());
await expect(session.appendLabel("missing", "checkpoint")).rejects.toThrow("Entry missing not found"); await expect(session.appendLabel("missing", "checkpoint")).rejects.toThrow("Entry missing not found");
}); });
it("persists leaf changes and appended entries via storage", async () => { it("persists leaf changes and appended entries via storage", async () => {
const storage = await createStorage(); const storage = await createStorage();
const session = new DefaultSession(storage); const session = new Session(storage);
const user1 = await session.appendMessage(createUserMessage("one")); const user1 = await session.appendMessage(createUserMessage("one"));
await session.appendMessage(createAssistantMessage("two")); await session.appendMessage(createAssistantMessage("two"));
await session.appendLabel(user1, "checkpoint"); await session.appendLabel(user1, "checkpoint");
await session.appendSessionName("name"); await session.appendSessionName("name");
await session.moveTo(user1); await session.moveTo(user1);
await session.appendMessage(createAssistantMessage("branched")); await session.appendMessage(createAssistantMessage("branched"));
const session2 = new DefaultSession(storage); const session2 = new Session(storage);
const context = await session2.buildContext(); const context = await session2.buildContext();
expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]); expect(context.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
expect(await session2.getLabel(user1)).toBe("checkpoint"); expect(await session2.getLabel(user1)).toBe("checkpoint");