fix(agent): add session context entry projection

packages/agent/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Added
+
+- Added configurable harness session context entry transforms and custom-entry message projectors.
+
 ### Fixed
 
 - Fixed harness split-turn compaction to serialize summary requests so single-concurrency providers are not asked to run overlapping generations ([#5536](https://github.com/earendil-works/pi/issues/5536)).

packages/agent/docs/agent-harness.md

@@ -79,6 +79,8 @@ Stream options are shallow-copied when a snapshot is created. `headers` and `met
 
 The session contains persisted entries only. Session reads return persisted state and do not include queued writes.
 
+`Session.buildContextEntries()` returns the compaction-aware entry sequence used for model context construction. `Session.buildContext()` derives runtime state from the full active branch, then projects those context entries to `AgentMessage[]`. Custom entries are omitted from model context by default; applications can pass `entryProjectors` to the `Session` constructor or `buildContext()` to project selected custom entries into messages. Applications can also pass stacked `entryTransforms`, which run after the default compaction transform, to filter or reorder context entries before projection.
+
 Session storage implementations must persist leaf changes as `leaf` entries. `setLeafId()` is not an in-memory-only cursor update; it appends a durable entry whose `targetId` is the active tree leaf or `null` for root. Reopening storage must reconstruct the current leaf from the latest persisted leaf-affecting entry.
 
 ### Pending session writes

packages/agent/src/harness/session/session.ts

@@ -19,11 +19,25 @@ import type {
 } from "../types.ts";
 import { SessionError } from "../types.ts";
 
-export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionContext {
+export type ContextEntryTransform = (entries: readonly SessionTreeEntry[]) => readonly SessionTreeEntry[];
+
+export type CustomEntryContextMessageProjector = (
+	entry: CustomEntry,
+	index: number,
+	entries: readonly SessionTreeEntry[],
+) => readonly AgentMessage[] | undefined;
+
+export interface SessionContextBuildOptions {
+	/** Additional entry transforms applied after the default compaction transform. */
+	entryTransforms?: readonly ContextEntryTransform[];
+	/** Optional custom-entry projectors. Custom entries are omitted from model context by default. */
+	entryProjectors?: Readonly<Record<string, CustomEntryContextMessageProjector>>;
+}
+
+function deriveSessionContextState(pathEntries: readonly SessionTreeEntry[]): Omit<SessionContext, "messages"> {
 	let thinkingLevel = "off";
 	let model: { provider: string; modelId: string } | null = null;
 	let activeToolNames: string[] | null = null;
-	let compaction: CompactionEntry | null = null;
 
 	for (const entry of pathEntries) {
 		if (entry.type === "thinking_level_change") {
@@ -34,56 +48,99 @@ export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionCon
 			model = { provider: entry.message.provider, modelId: entry.message.model };
 		} else if (entry.type === "active_tools_change") {
 			activeToolNames = [...entry.activeToolNames];
-		} else if (entry.type === "compaction") {
+		}
+	}
+
+	return { thinkingLevel, model, activeToolNames };
+}
+
+export function defaultContextEntryTransform(pathEntries: readonly SessionTreeEntry[]): SessionTreeEntry[] {
+	let compaction: CompactionEntry | null = null;
+	for (const entry of pathEntries) {
+		if (entry.type === "compaction") {
 			compaction = entry;
 		}
 	}
-
+	if (!compaction) {
-	const messages: AgentMessage[] = [];
+		return [...pathEntries];
-	const appendMessage = (entry: SessionTreeEntry) => {
-		if (entry.type === "message") {
-			messages.push(entry.message as AgentMessage);
-		} else if (entry.type === "custom_message") {
-			messages.push(
-				createCustomMessage(
-					entry.customType,
-					entry.content as string | (TextContent | ImageContent)[],
-					entry.display,
-					entry.details,
-					entry.timestamp,
-				),
-			);
-		} else if (entry.type === "branch_summary" && entry.summary) {
-			messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp));
-		}
-	};
-
-	if (compaction) {
-		messages.push(createCompactionSummaryMessage(compaction.summary, compaction.tokensBefore, compaction.timestamp));
-		const compactionIdx = pathEntries.findIndex((e) => e.type === "compaction" && e.id === compaction.id);
-		let foundFirstKept = false;
-		for (let i = 0; i < compactionIdx; i++) {
-			const entry = pathEntries[i]!;
-			if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true;
-			if (foundFirstKept) appendMessage(entry);
-		}
-		for (let i = compactionIdx + 1; i < pathEntries.length; i++) {
-			appendMessage(pathEntries[i]!);
-		}
-	} else {
-		for (const entry of pathEntries) {
-			appendMessage(entry);
-		}
 	}
 
-	return { messages, thinkingLevel, model, activeToolNames };
+	const entries: SessionTreeEntry[] = [compaction];
+	const compactionIdx = pathEntries.findIndex((entry) => entry.type === "compaction" && entry.id === compaction.id);
+	let foundFirstKept = false;
+	for (let i = 0; i < compactionIdx; i++) {
+		const entry = pathEntries[i]!;
+		if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true;
+		if (foundFirstKept) entries.push(entry);
+	}
+	for (let i = compactionIdx + 1; i < pathEntries.length; i++) {
+		entries.push(pathEntries[i]!);
+	}
+	return entries;
+}
+
+export function buildContextEntries(
+	pathEntries: readonly SessionTreeEntry[],
+	options: SessionContextBuildOptions = {},
+): SessionTreeEntry[] {
+	let entries = defaultContextEntryTransform(pathEntries);
+	for (const transform of options.entryTransforms ?? []) {
+		entries = [...transform(entries)];
+	}
+	return entries;
+}
+
+export function sessionEntryToContextMessages(
+	entry: SessionTreeEntry,
+	index: number,
+	entries: readonly SessionTreeEntry[],
+	options: SessionContextBuildOptions = {},
+): AgentMessage[] {
+	if (entry.type === "message") {
+		return [entry.message as AgentMessage];
+	}
+	if (entry.type === "custom_message") {
+		return [
+			createCustomMessage(
+				entry.customType,
+				entry.content as string | (TextContent | ImageContent)[],
+				entry.display,
+				entry.details,
+				entry.timestamp,
+			),
+		];
+	}
+	if (entry.type === "compaction") {
+		return [createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp)];
+	}
+	if (entry.type === "branch_summary" && entry.summary) {
+		return [createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp)];
+	}
+	if (entry.type === "custom") {
+		return [...(options.entryProjectors?.[entry.customType]?.(entry, index, entries) ?? [])];
+	}
+	return [];
+}
+
+export function buildSessionContext(
+	pathEntries: readonly SessionTreeEntry[],
+	options: SessionContextBuildOptions = {},
+): SessionContext {
+	const state = deriveSessionContextState(pathEntries);
+	const contextEntries = buildContextEntries(pathEntries, options);
+	const messages = contextEntries.flatMap((entry, index) =>
+		sessionEntryToContextMessages(entry, index, contextEntries, options),
+	);
+	return { ...state, messages };
 }
 
 export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
 	private storage: SessionStorage<TMetadata>;
+	private contextBuildOptions: SessionContextBuildOptions;
 
-	constructor(storage: SessionStorage<TMetadata>) {
+	constructor(storage: SessionStorage<TMetadata>, contextBuildOptions: SessionContextBuildOptions = {}) {
 		this.storage = storage;
+		this.contextBuildOptions = contextBuildOptions;
 	}
 
 	getMetadata(): Promise<TMetadata> {
@@ -111,8 +168,22 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
 		return this.storage.getPathToRoot(leafId);
 	}
 
-	async buildContext(): Promise<SessionContext> {
+	async buildContextEntries(options: SessionContextBuildOptions = {}): Promise<SessionTreeEntry[]> {
-		return buildSessionContext(await this.getBranch());
+		return buildContextEntries(await this.getBranch(), this.mergeContextBuildOptions(options));
+	}
+
+	async buildContext(options: SessionContextBuildOptions = {}): Promise<SessionContext> {
+		return buildSessionContext(await this.getBranch(), this.mergeContextBuildOptions(options));
+	}
+
+	private mergeContextBuildOptions(options: SessionContextBuildOptions): SessionContextBuildOptions {
+		return {
+			entryTransforms: [...(this.contextBuildOptions.entryTransforms ?? []), ...(options.entryTransforms ?? [])],
+			entryProjectors: {
+				...(this.contextBuildOptions.entryProjectors ?? {}),
+				...(options.entryProjectors ?? {}),
+			},
+		};
 	}
 
 	getLabel(id: string): Promise<string | undefined> {

packages/agent/test/harness/session.test.ts

@@ -4,10 +4,18 @@ import { describe, expect, it } from "vitest";
 import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
 import { JsonlSessionStorage } from "../../src/harness/session/jsonl-storage.ts";
 import { InMemorySessionStorage } from "../../src/harness/session/memory-storage.ts";
-import { Session } from "../../src/harness/session/session.ts";
+import { type ContextEntryTransform, Session } from "../../src/harness/session/session.ts";
 import type { SessionStorage } from "../../src/harness/types.ts";
 import { createAssistantMessage, createTempDir, createUserMessage, getLatestTempDir } from "./session-test-utils.ts";
 
+function getTextData(data: unknown): string {
+	if (typeof data !== "object" || data === null || !("text" in data)) {
+		return "";
+	}
+	const value = (data as { text?: unknown }).text;
+	return typeof value === "string" ? value : "";
+}
+
 async function runSessionSuite(
 	name: string,
 	createStorage: () => SessionStorage | Promise<SessionStorage>,
@@ -86,6 +94,45 @@ async function runSessionSuite(
 			expect(context.messages[1]?.role).toBe("custom");
 		});
 
+		it("keeps custom entries in context entries but omits them from messages by default", async () => {
+			const session = new Session(await createStorage());
+			await session.appendMessage(createUserMessage("one"));
+			await session.appendCustomEntry("chat_message", { text: "hello" });
+			const contextEntries = await session.buildContextEntries();
+			const context = await session.buildContext();
+			expect(contextEntries.map((entry) => entry.type)).toEqual(["message", "custom"]);
+			expect(context.messages).toHaveLength(1);
+		});
+
+		it("projects custom entries with configured custom-entry projectors", async () => {
+			const session = new Session(await createStorage(), {
+				entryProjectors: {
+					chat_message: (entry) => [createUserMessage(`chat: ${getTextData(entry.data)}`)],
+				},
+			});
+			await session.appendMessage(createUserMessage("one"));
+			await session.appendCustomEntry("chat_message", { text: "hello" });
+			const context = await session.buildContext();
+			expect(context.messages.map((message) => message.role)).toEqual(["user", "user"]);
+			expect(context.messages[1]).toMatchObject({ content: [{ type: "text", text: "chat: hello" }] });
+		});
+
+		it("applies context entry transforms after default compaction selection", async () => {
+			let observedFirstEntryType: string | undefined;
+			const dropCompaction: ContextEntryTransform = (entries) => {
+				observedFirstEntryType = entries[0]?.type;
+				return entries.filter((entry) => entry.type !== "compaction");
+			};
+			const session = new Session(await createStorage(), { entryTransforms: [dropCompaction] });
+			await session.appendMessage(createUserMessage("one"));
+			const kept = await session.appendMessage(createUserMessage("two"));
+			await session.appendCompaction("summary", kept, 1234);
+			await session.appendMessage(createUserMessage("three"));
+			const context = await session.buildContext();
+			expect(observedFirstEntryType).toBe("compaction");
+			expect(context.messages.map((message) => message.role)).toEqual(["user", "user"]);
+		});
+
 		it("normalizes session names", async () => {
 			const session = new Session(await createStorage());
 			await session.appendSessionName(" hello\nworld\r\nagain ");