Skip to content

Commit dd1c690

Browse files
badlogicbadlogic
committed
fix(agent): add session context entry projection
1 parent 312bc71 commit dd1c690

4 files changed

Lines changed: 167 additions & 43 deletions

File tree

‎packages/agent/CHANGELOG.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@
22

33
## [Unreleased]
44

5+
### Added
6+
7+
- Added configurable harness session context entry transforms and custom-entry message projectors.
8+
59
### Fixed
610

711
- 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‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,8 @@ Stream options are shallow-copied when a snapshot is created. `headers` and `met
7979

8080
The session contains persisted entries only. Session reads return persisted state and do not include queued writes.
8181

82+
`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.
83+
8284
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.
8385

8486
### Pending session writes

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

Lines changed: 113 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,25 @@ import type {
1919
} from "../types.ts";
2020
import { SessionError } from "../types.ts";
2121

22-
export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionContext {
22+
export type ContextEntryTransform = (entries: readonly SessionTreeEntry[]) => readonly SessionTreeEntry[];
23+
24+
export type CustomEntryContextMessageProjector = (
25+
entry: CustomEntry,
26+
index: number,
27+
entries: readonly SessionTreeEntry[],
28+
) => readonly AgentMessage[] | undefined;
29+
30+
export interface SessionContextBuildOptions {
31+
/** Additional entry transforms applied after the default compaction transform. */
32+
entryTransforms?: readonly ContextEntryTransform[];
33+
/** Optional custom-entry projectors. Custom entries are omitted from model context by default. */
34+
entryProjectors?: Readonly<Record<string, CustomEntryContextMessageProjector>>;
35+
}
36+
37+
function deriveSessionContextState(pathEntries: readonly SessionTreeEntry[]): Omit<SessionContext, "messages"> {
2338
let thinkingLevel = "off";
2439
let model: { provider: string; modelId: string } | null = null;
2540
let activeToolNames: string[] | null = null;
26-
let compaction: CompactionEntry | null = null;
2741

2842
for (const entry of pathEntries) {
2943
if (entry.type === "thinking_level_change") {
@@ -34,56 +48,99 @@ export function buildSessionContext(pathEntries: SessionTreeEntry[]): SessionCon
3448
model = { provider: entry.message.provider, modelId: entry.message.model };
3549
} else if (entry.type === "active_tools_change") {
3650
activeToolNames = [...entry.activeToolNames];
37-
} else if (entry.type === "compaction") {
38-
compaction = entry;
3951
}
4052
}
4153

42-
const messages: AgentMessage[] = [];
43-
const appendMessage = (entry: SessionTreeEntry) => {
44-
if (entry.type === "message") {
45-
messages.push(entry.message as AgentMessage);
46-
} else if (entry.type === "custom_message") {
47-
messages.push(
48-
createCustomMessage(
49-
entry.customType,
50-
entry.content as string | (TextContent | ImageContent)[],
51-
entry.display,
52-
entry.details,
53-
entry.timestamp,
54-
),
55-
);
56-
} else if (entry.type === "branch_summary" && entry.summary) {
57-
messages.push(createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp));
58-
}
59-
};
60-
61-
if (compaction) {
62-
messages.push(createCompactionSummaryMessage(compaction.summary, compaction.tokensBefore, compaction.timestamp));
63-
const compactionIdx = pathEntries.findIndex((e) => e.type === "compaction" && e.id === compaction.id);
64-
let foundFirstKept = false;
65-
for (let i = 0; i < compactionIdx; i++) {
66-
const entry = pathEntries[i]!;
67-
if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true;
68-
if (foundFirstKept) appendMessage(entry);
69-
}
70-
for (let i = compactionIdx + 1; i < pathEntries.length; i++) {
71-
appendMessage(pathEntries[i]!);
72-
}
73-
} else {
74-
for (const entry of pathEntries) {
75-
appendMessage(entry);
54+
return { thinkingLevel, model, activeToolNames };
55+
}
56+
57+
export function defaultContextEntryTransform(pathEntries: readonly SessionTreeEntry[]): SessionTreeEntry[] {
58+
let compaction: CompactionEntry | null = null;
59+
for (const entry of pathEntries) {
60+
if (entry.type === "compaction") {
61+
compaction = entry;
7662
}
7763
}
64+
if (!compaction) {
65+
return [...pathEntries];
66+
}
7867

79-
return { messages, thinkingLevel, model, activeToolNames };
68+
const entries: SessionTreeEntry[] = [compaction];
69+
const compactionIdx = pathEntries.findIndex((entry) => entry.type === "compaction" && entry.id === compaction.id);
70+
let foundFirstKept = false;
71+
for (let i = 0; i < compactionIdx; i++) {
72+
const entry = pathEntries[i]!;
73+
if (entry.id === compaction.firstKeptEntryId) foundFirstKept = true;
74+
if (foundFirstKept) entries.push(entry);
75+
}
76+
for (let i = compactionIdx + 1; i < pathEntries.length; i++) {
77+
entries.push(pathEntries[i]!);
78+
}
79+
return entries;
80+
}
81+
82+
export function buildContextEntries(
83+
pathEntries: readonly SessionTreeEntry[],
84+
options: SessionContextBuildOptions = {},
85+
): SessionTreeEntry[] {
86+
let entries = defaultContextEntryTransform(pathEntries);
87+
for (const transform of options.entryTransforms ?? []) {
88+
entries = [...transform(entries)];
89+
}
90+
return entries;
91+
}
92+
93+
export function sessionEntryToContextMessages(
94+
entry: SessionTreeEntry,
95+
index: number,
96+
entries: readonly SessionTreeEntry[],
97+
options: SessionContextBuildOptions = {},
98+
): AgentMessage[] {
99+
if (entry.type === "message") {
100+
return [entry.message as AgentMessage];
101+
}
102+
if (entry.type === "custom_message") {
103+
return [
104+
createCustomMessage(
105+
entry.customType,
106+
entry.content as string | (TextContent | ImageContent)[],
107+
entry.display,
108+
entry.details,
109+
entry.timestamp,
110+
),
111+
];
112+
}
113+
if (entry.type === "compaction") {
114+
return [createCompactionSummaryMessage(entry.summary, entry.tokensBefore, entry.timestamp)];
115+
}
116+
if (entry.type === "branch_summary" && entry.summary) {
117+
return [createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp)];
118+
}
119+
if (entry.type === "custom") {
120+
return [...(options.entryProjectors?.[entry.customType]?.(entry, index, entries) ?? [])];
121+
}
122+
return [];
123+
}
124+
125+
export function buildSessionContext(
126+
pathEntries: readonly SessionTreeEntry[],
127+
options: SessionContextBuildOptions = {},
128+
): SessionContext {
129+
const state = deriveSessionContextState(pathEntries);
130+
const contextEntries = buildContextEntries(pathEntries, options);
131+
const messages = contextEntries.flatMap((entry, index) =>
132+
sessionEntryToContextMessages(entry, index, contextEntries, options),
133+
);
134+
return { ...state, messages };
80135
}
81136

82137
export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
83138
private storage: SessionStorage<TMetadata>;
139+
private contextBuildOptions: SessionContextBuildOptions;
84140

85-
constructor(storage: SessionStorage<TMetadata>) {
141+
constructor(storage: SessionStorage<TMetadata>, contextBuildOptions: SessionContextBuildOptions = {}) {
86142
this.storage = storage;
143+
this.contextBuildOptions = contextBuildOptions;
87144
}
88145

89146
getMetadata(): Promise<TMetadata> {
@@ -111,8 +168,22 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
111168
return this.storage.getPathToRoot(leafId);
112169
}
113170

114-
async buildContext(): Promise<SessionContext> {
115-
return buildSessionContext(await this.getBranch());
171+
async buildContextEntries(options: SessionContextBuildOptions = {}): Promise<SessionTreeEntry[]> {
172+
return buildContextEntries(await this.getBranch(), this.mergeContextBuildOptions(options));
173+
}
174+
175+
async buildContext(options: SessionContextBuildOptions = {}): Promise<SessionContext> {
176+
return buildSessionContext(await this.getBranch(), this.mergeContextBuildOptions(options));
177+
}
178+
179+
private mergeContextBuildOptions(options: SessionContextBuildOptions): SessionContextBuildOptions {
180+
return {
181+
entryTransforms: [...(this.contextBuildOptions.entryTransforms ?? []), ...(options.entryTransforms ?? [])],
182+
entryProjectors: {
183+
...(this.contextBuildOptions.entryProjectors ?? {}),
184+
...(options.entryProjectors ?? {}),
185+
},
186+
};
116187
}
117188

118189
getLabel(id: string): Promise<string | undefined> {

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

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,18 @@ import { describe, expect, it } from "vitest";
44
import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
55
import { JsonlSessionStorage } from "../../src/harness/session/jsonl-storage.ts";
66
import { InMemorySessionStorage } from "../../src/harness/session/memory-storage.ts";
7-
import { Session } from "../../src/harness/session/session.ts";
7+
import { type ContextEntryTransform, Session } from "../../src/harness/session/session.ts";
88
import type { SessionStorage } from "../../src/harness/types.ts";
99
import { createAssistantMessage, createTempDir, createUserMessage, getLatestTempDir } from "./session-test-utils.ts";
1010

11+
function getTextData(data: unknown): string {
12+
if (typeof data !== "object" || data === null || !("text" in data)) {
13+
return "";
14+
}
15+
const value = (data as { text?: unknown }).text;
16+
return typeof value === "string" ? value : "";
17+
}
18+
1119
async function runSessionSuite(
1220
name: string,
1321
createStorage: () => SessionStorage | Promise<SessionStorage>,
@@ -86,6 +94,45 @@ async function runSessionSuite(
8694
expect(context.messages[1]?.role).toBe("custom");
8795
});
8896

97+
it("keeps custom entries in context entries but omits them from messages by default", async () => {
98+
const session = new Session(await createStorage());
99+
await session.appendMessage(createUserMessage("one"));
100+
await session.appendCustomEntry("chat_message", { text: "hello" });
101+
const contextEntries = await session.buildContextEntries();
102+
const context = await session.buildContext();
103+
expect(contextEntries.map((entry) => entry.type)).toEqual(["message", "custom"]);
104+
expect(context.messages).toHaveLength(1);
105+
});
106+
107+
it("projects custom entries with configured custom-entry projectors", async () => {
108+
const session = new Session(await createStorage(), {
109+
entryProjectors: {
110+
chat_message: (entry) => [createUserMessage(`chat: ${getTextData(entry.data)}`)],
111+
},
112+
});
113+
await session.appendMessage(createUserMessage("one"));
114+
await session.appendCustomEntry("chat_message", { text: "hello" });
115+
const context = await session.buildContext();
116+
expect(context.messages.map((message) => message.role)).toEqual(["user", "user"]);
117+
expect(context.messages[1]).toMatchObject({ content: [{ type: "text", text: "chat: hello" }] });
118+
});
119+
120+
it("applies context entry transforms after default compaction selection", async () => {
121+
let observedFirstEntryType: string | undefined;
122+
const dropCompaction: ContextEntryTransform = (entries) => {
123+
observedFirstEntryType = entries[0]?.type;
124+
return entries.filter((entry) => entry.type !== "compaction");
125+
};
126+
const session = new Session(await createStorage(), { entryTransforms: [dropCompaction] });
127+
await session.appendMessage(createUserMessage("one"));
128+
const kept = await session.appendMessage(createUserMessage("two"));
129+
await session.appendCompaction("summary", kept, 1234);
130+
await session.appendMessage(createUserMessage("three"));
131+
const context = await session.buildContext();
132+
expect(observedFirstEntryType).toBe("compaction");
133+
expect(context.messages.map((message) => message.role)).toEqual(["user", "user"]);
134+
});
135+
89136
it("normalizes session names", async () => {
90137
const session = new Session(await createStorage());
91138
await session.appendSessionName(" hello\nworld\r\nagain ");

0 commit comments

Comments
 (0)