sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

refactor(agent): keep storage register keys internal

Browse source 
Represent snapshot registers as arrays and simplify memory session fork handling.
This commit is contained in:
David Brailovsky 2026-08-14 13:19:12 +02:00
parent 39b5f12a04
commit 3ad1d716ae
2 changed files with 32 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

50
packages/agent/src/harness/session/memory.ts
View file

@ -1,7 +1,7 @@
import { type Usage, uuidv7 } from "@earendil-works/pi-ai";
import type { AgentMessage } from "../../types.ts";
import { StorageBackedSession } from "./session.ts";
import { registerKey, StorageState, type StorageStateSnapshot } from "./storage-state.ts";
import { StorageState, type StorageStateSnapshot } from "./storage-state.ts";
import type {
BranchScan,
CommitResult,
@ -133,7 +133,7 @@ export class MemoryStorage implements Storage {
const snapshot = this.storageState.snapshot();
return {
entries: [...snapshot.entries.values()].sort((left, right) => left.seq - right.seq),
registers: [...snapshot.registers.values()],
registers: snapshot.registers,
};
});
this.commitQueue = result.then(
@ -443,23 +443,22 @@ export class MemorySessionRepo implements SessionRepo {
options: ForkOptions,
): MemoryStorage {
const sourceEntries = new Map(snapshot.entries.map((entry) => [entry.id, entry]));
const sourceRegisters = new Map(
snapshot.registers.map((register) => [registerKey(register.namespace, register.key), register]),
);
const sourceLeaves = snapshot.registers.filter((register) => isRegisterNamespace(register, "lane.leaf"));
this.validateForkSourceSnapshot(snapshot, sourceEntries, sourceRegisters, sourceLeaves);
this.validateForkSourceSnapshot(snapshot, sourceEntries, sourceLeaves);
const { copiedEntryIds, destinationLeaves } = this.selectForkContents(sourceEntries, sourceLeaves, options);
const { entryIds, laneToLeafId } = this.selectForkContents(sourceEntries, sourceLeaves, options);
const entries = new Map<string, Entry>();
for (const id of copiedEntryIds) entries.set(id, sourceEntries.get(id)!);
const registers = new Map<string, Register>();
for (const id of entryIds) entries.set(id, sourceEntries.get(id)!);
const registers: Register[] = [];
let nextSeq = Math.max(0, ...[...entries.values()].map((entry) => entry.seq)) + 1;
const setRegister = (namespace: RegisterNamespace, key: string, value: Register["value"]): void => {
registers.set(registerKey(namespace, key), { namespace, key, value, seq: nextSeq++ } as Register);
registers.push({ namespace, key, value, seq: nextSeq++ } as Register);
};
for (const [lane, leaf] of destinationLeaves) {
const configuration = sourceRegisters.get(registerKey("lane.config", lane));
for (const [lane, leaf] of laneToLeafId) {
const configuration = snapshot.registers.find(
(register) => register.namespace === "lane.config" && register.key === lane,
);
if (configuration !== undefined) setRegister("lane.config", lane, configuration.value);
setRegister("lane.leaf", lane, leaf);
setRegister("lane.state", lane, { currentOperationId: null, pendingNextRun: [] });
@ -468,7 +467,7 @@ export class MemorySessionRepo implements SessionRepo {
if (
register.namespace === "fact.name" ||
register.namespace === "fact.custom" ||
(register.namespace === "fact.label" && copiedEntryIds.has(register.key))
(register.namespace === "fact.label" && entryIds.has(register.key))
) {
setRegister(register.namespace, register.key, register.value);
}
@ -492,12 +491,12 @@ export class MemorySessionRepo implements SessionRepo {
sourceEntries: Map<string, Entry>,
sourceLeaves: Register<"lane.leaf">[],
options: ForkOptions,
): { copiedEntryIds: Set<string>; destinationLeaves: Map<string, string | null> } {
const copiedEntryIds = new Set<string>();
const destinationLeaves = new Map<string, string | null>();
): { entryIds: Set<string>; laneToLeafId: Map<string, string | null> } {
const entryIds = new Set<string>();
const laneToLeafId = new Map<string, string | null>();
if (options.scope === "tree") {
for (const id of sourceEntries.keys()) copiedEntryIds.add(id);
for (const register of sourceLeaves) destinationLeaves.set(register.key, register.value);
for (const id of sourceEntries.keys()) entryIds.add(id);
for (const register of sourceLeaves) laneToLeafId.set(register.key, register.value);
} else {
const mainLeaf = sourceLeaves.find((register) => register.key === "main");
if (mainLeaf === undefined) throw new Error("Source session is missing main lane");
@ -513,18 +512,17 @@ export class MemorySessionRepo implements SessionRepo {
while (entryId !== null) {
const entry = sourceEntries.get(entryId);
if (entry === undefined) throw new Error(`Corrupt source branch: missing parent ${entryId}`);
copiedEntryIds.add(entryId);
entryIds.add(entryId);
entryId = entry.parentId;
}
destinationLeaves.set("main", leaf);
laneToLeafId.set("main", leaf);
}
return { copiedEntryIds, destinationLeaves };
return { entryIds, laneToLeafId };
}
private validateForkSourceSnapshot(
snapshot: { entries: Entry[]; registers: Register[] },
sourceEntries: Map<string, Entry>,
sourceRegisters: Map<string, Register>,
sourceLeaves: Register<"lane.leaf">[],
): void {
const sourceLeafKeys = new Set(sourceLeaves.map((register) => register.key));
@ -542,10 +540,10 @@ export class MemorySessionRepo implements SessionRepo {
}
}
for (const leaf of sourceLeaves) {
if (!sourceRegisters.has(registerKey("lane.state", leaf.key))) {
if (!this.hasRegister(snapshot.registers, "lane.state", leaf.key)) {
throw new Error(`Source session lane ${JSON.stringify(leaf.key)} is missing lane.state`);
}
if (leaf.key !== "main" && !sourceRegisters.has(registerKey("lane.config", leaf.key))) {
if (leaf.key !== "main" && !this.hasRegister(snapshot.registers, "lane.config", leaf.key)) {
throw new Error(`Source session lane ${JSON.stringify(leaf.key)} is missing lane.config`);
}
if (leaf.value !== null && !sourceEntries.has(leaf.value)) {
@ -554,6 +552,10 @@ export class MemorySessionRepo implements SessionRepo {
}
}
private hasRegister(registers: Register[], namespace: RegisterNamespace, key: string): boolean {
return registers.some((register) => register.namespace === namespace && register.key === key);
}
private reserveId(id: string): void {
if (this.sessions.has(id) || this.pendingIds.has(id)) throw new Error(`Session already exists: ${id}`);
this.pendingIds.add(id);

10
packages/agent/src/harness/session/storage-state.ts
View file

@ -26,13 +26,13 @@ import type {
export interface StorageStateSnapshot {
entries: Map<string, Entry>;
registers: Map<string, Register>;
registers: Register[];
usage: Map<string, UsageRow>;
stats: SessionStats;
nextSeq: number;
}
export function registerKey(namespace: RegisterNamespace, key: string): string {
function registerKey(namespace: RegisterNamespace, key: string): string {
return `${namespace}\u0000${key}`;
}
@ -59,7 +59,9 @@ export class StorageState {
constructor(snapshot?: StorageStateSnapshot) {
this.entries = snapshot?.entries ?? new Map();
this.entriesBySeq = [...this.entries.values()].sort((left, right) => left.seq - right.seq);
this.registers = snapshot?.registers ?? new Map();
this.registers = new Map(
(snapshot?.registers ?? []).map((register) => [registerKey(register.namespace, register.key), register]),
);
this.usage = snapshot?.usage ?? new Map();
this.stats = snapshot?.stats ?? { messageCount: 0, usage: emptyUsage() };
this.nextSeq = snapshot?.nextSeq ?? 1;
@ -212,7 +214,7 @@ export class StorageState {
snapshot(): StorageStateSnapshot {
return {
entries: this.entries,
registers: this.registers,
registers: [...this.registers.values()],
usage: this.usage,
stats: this.stats,
nextSeq: this.nextSeq,