ton/pi_harness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(agent): add session context entry projection

Browse Source
This commit is contained in:
Mario Zechner
2026-07-08 14:58:51 +02:00
parent 312bc713bb
commit dd1c690f36
4 changed files with 168 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/agent/CHANGELOG.md
+4
View File
@@ -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
+2
View File
@@ -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
+114 -43
View File
@@ -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;
}
}
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<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> {
return buildSessionContext(await this.getBranch());
async buildContextEntries(options: SessionContextBuildOptions = {}): Promise<SessionTreeEntry[]> {
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
+48 -1
View File
@@ -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 ");