From dd1c690f36de449fac9b24c8b3d113837eaf7cfb Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Wed, 8 Jul 2026 14:58:51 +0200 Subject: [PATCH] fix(agent): add session context entry projection --- packages/agent/CHANGELOG.md | 4 + packages/agent/docs/agent-harness.md | 2 + packages/agent/src/harness/session/session.ts | 157 +++++++++++++----- packages/agent/test/harness/session.test.ts | 49 +++++- 4 files changed, 168 insertions(+), 44 deletions(-) diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 21ea92bd..9cff3432 100644 --- a/packages/agent/CHANGELOG.md +++ b/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)). diff --git a/packages/agent/docs/agent-harness.md b/packages/agent/docs/agent-harness.md index 64a8baec..16e0664c 100644 --- a/packages/agent/docs/agent-harness.md +++ b/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 diff --git a/packages/agent/src/harness/session/session.ts b/packages/agent/src/harness/session/session.ts index ce369dff..f5224c92 100644 --- a/packages/agent/src/harness/session/session.ts +++ b/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>; +} + +function deriveSessionContextState(pathEntries: readonly SessionTreeEntry[]): Omit { 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; } } - - const messages: AgentMessage[] = []; - 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); - } + if (!compaction) { + return [...pathEntries]; } - 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 { private storage: SessionStorage; + private contextBuildOptions: SessionContextBuildOptions; - constructor(storage: SessionStorage) { + constructor(storage: SessionStorage, contextBuildOptions: SessionContextBuildOptions = {}) { this.storage = storage; + this.contextBuildOptions = contextBuildOptions; } getMetadata(): Promise { @@ -111,8 +168,22 @@ export class Session { return this.storage.getPathToRoot(leafId); } - async buildContext(): Promise { - return buildSessionContext(await this.getBranch()); + async buildContextEntries(options: SessionContextBuildOptions = {}): Promise { + return buildContextEntries(await this.getBranch(), this.mergeContextBuildOptions(options)); + } + + async buildContext(options: SessionContextBuildOptions = {}): Promise { + 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 { diff --git a/packages/agent/test/harness/session.test.ts b/packages/agent/test/harness/session.test.ts index 39e285b5..c009123d 100644 --- a/packages/agent/test/harness/session.test.ts +++ b/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, @@ -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 ");