Skip to content

Commit 7dc5b94

Browse files
steipetesteipete
committed
refactor: parse session reads without manager
1 parent c76ee64 commit 7dc5b94

2 files changed

Lines changed: 108 additions & 202 deletions

File tree

‎src/gateway/session-utils.fs.test.ts‎

Lines changed: 66 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -667,6 +667,51 @@ describe("readSessionMessages", () => {
667667
}
668668
});
669669

670+
test("honors byte caps for sync recent tree-message reads", () => {
671+
const sessionId = "test-session-recent-tree-byte-cap";
672+
const transcriptPath = path.join(tmpDir, `${sessionId}.jsonl`);
673+
const hugeContent = "huge ".repeat(4096);
674+
const lines = [
675+
JSON.stringify({ type: "session", version: 3, id: sessionId }),
676+
JSON.stringify({
677+
type: "message",
678+
id: "root",
679+
parentId: null,
680+
message: { role: "user", content: "root" },
681+
}),
682+
JSON.stringify({
683+
type: "message",
684+
id: "huge",
685+
parentId: "root",
686+
message: { role: "assistant", content: hugeContent },
687+
}),
688+
JSON.stringify({
689+
type: "message",
690+
id: "tail",
691+
parentId: "huge",
692+
message: { role: "assistant", content: "tail" },
693+
}),
694+
];
695+
fs.writeFileSync(transcriptPath, `${lines.join("\n")}\n`, "utf-8");
696+
const readFileSpy = vi.spyOn(fs, "readFileSync");
697+
const sessionManagerOpenSpy = vi.spyOn(SessionManager, "open");
698+
699+
try {
700+
const out = readRecentSessionMessages(sessionId, storePath, undefined, {
701+
maxMessages: 2,
702+
maxBytes: 2048,
703+
});
704+
705+
expect(out).toEqual([expect.objectContaining({ role: "assistant", content: "tail" })]);
706+
expect(JSON.stringify(out)).not.toContain("huge");
707+
expect(readFileSpy).not.toHaveBeenCalled();
708+
expect(sessionManagerOpenSpy).not.toHaveBeenCalled();
709+
} finally {
710+
readFileSpy.mockRestore();
711+
sessionManagerOpenSpy.mockRestore();
712+
}
713+
});
714+
670715
test("counts transcript messages without loading the whole file", () => {
671716
const sessionId = "test-session-count-large";
672717
const transcriptPath = path.join(tmpDir, `${sessionId}.jsonl`);
@@ -857,22 +902,28 @@ describe("readSessionMessages", () => {
857902
const rawTranscript = fs.readFileSync(sessionFile, "utf-8");
858903
expect(rawTranscript).toContain("original wrapped prompt");
859904
expect(rawTranscript).toContain("clean prompt");
905+
const sessionManagerOpenSpy = vi.spyOn(SessionManager, "open");
860906

861-
const out = readSessionMessages(sessionId, storePath, sessionFile);
862-
expect(out).toHaveLength(2);
863-
expect(out).toEqual([
864-
expect.objectContaining({
865-
role: "user",
866-
content: "clean prompt",
867-
__openclaw: expect.objectContaining({ seq: 1 }),
868-
}),
869-
expect.objectContaining({
870-
role: "assistant",
871-
content: [{ type: "text", text: "clean answer" }],
872-
__openclaw: expect.objectContaining({ seq: 2 }),
873-
}),
874-
]);
875-
expect(JSON.stringify(out)).not.toContain("original wrapped prompt");
907+
try {
908+
const out = readSessionMessages(sessionId, storePath, sessionFile);
909+
expect(out).toHaveLength(2);
910+
expect(out).toEqual([
911+
expect.objectContaining({
912+
role: "user",
913+
content: "clean prompt",
914+
__openclaw: expect.objectContaining({ seq: 1 }),
915+
}),
916+
expect.objectContaining({
917+
role: "assistant",
918+
content: [{ type: "text", text: "clean answer" }],
919+
__openclaw: expect.objectContaining({ seq: 2 }),
920+
}),
921+
]);
922+
expect(JSON.stringify(out)).not.toContain("original wrapped prompt");
923+
expect(sessionManagerOpenSpy).not.toHaveBeenCalled();
924+
} finally {
925+
sessionManagerOpenSpy.mockRestore();
926+
}
876927
});
877928

878929
test.each([

‎src/gateway/session-utils.fs.ts‎

Lines changed: 42 additions & 187 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
import fs from "node:fs";
22
import { StringDecoder } from "node:string_decoder";
3-
import { SessionManager, type SessionEntry } from "@mariozechner/pi-coding-agent";
43
import { deriveSessionTotalTokens, hasNonzeroUsage, normalizeUsage } from "../agents/usage.js";
54
import { jsonUtf8Bytes } from "../infra/json-utf8-bytes.js";
65
import { hasInterSessionUserProvenance } from "../sessions/input-provenance.js";
@@ -151,69 +150,7 @@ export function readSessionMessages(
151150
return [];
152151
}
153152

154-
const lines = fs.readFileSync(filePath, "utf-8").split(/\r?\n/);
155-
const hasTreeEntries = lines.some(hasSessionTreeEntry);
156-
let branchEntries: SessionEntry[] | null = null;
157-
if (hasTreeEntries) {
158-
try {
159-
branchEntries = SessionManager.open(filePath).getBranch();
160-
} catch {
161-
branchEntries = null;
162-
}
163-
}
164-
165-
if (branchEntries) {
166-
const messages: unknown[] = [];
167-
let messageSeq = 0;
168-
for (const entry of branchEntries) {
169-
if (entry.type === "message" && entry.message) {
170-
messageSeq += 1;
171-
messages.push(
172-
attachOpenClawTranscriptMeta(entry.message, {
173-
...(typeof entry.id === "string" ? { id: entry.id } : {}),
174-
seq: messageSeq,
175-
}),
176-
);
177-
continue;
178-
}
179-
180-
if (entry.type === "compaction") {
181-
const ts = typeof entry.timestamp === "string" ? Date.parse(entry.timestamp) : Number.NaN;
182-
const timestamp = Number.isFinite(ts) ? ts : Date.now();
183-
messageSeq += 1;
184-
messages.push({
185-
role: "system",
186-
content: [{ type: "text", text: "Compaction" }],
187-
timestamp,
188-
__openclaw: {
189-
kind: "compaction",
190-
id: typeof entry.id === "string" ? entry.id : undefined,
191-
seq: messageSeq,
192-
},
193-
});
194-
}
195-
}
196-
return messages;
197-
}
198-
199-
const messages: unknown[] = [];
200-
let messageSeq = 0;
201-
for (const line of lines) {
202-
if (!line.trim()) {
203-
continue;
204-
}
205-
try {
206-
const parsed = JSON.parse(line);
207-
const message = parsedSessionEntryToMessage(parsed, messageSeq + 1);
208-
if (message) {
209-
messageSeq += 1;
210-
messages.push(message);
211-
}
212-
} catch {
213-
// ignore bad lines
214-
}
215-
}
216-
return messages;
153+
return transcriptRecordsToMessages(readSelectedTranscriptRecords(filePath));
217154
}
218155

219156
type ReadRecentSessionMessagesOptions = {
@@ -283,25 +220,7 @@ export function readRecentSessionMessages(
283220
.filter((line) => line.trim().length > 0)
284221
.slice(-maxLines);
285222

286-
if (lines.some(hasSessionTreeEntry)) {
287-
return readSessionMessages(sessionId, storePath, sessionFile).slice(-maxMessages);
288-
}
289-
290-
const messages: unknown[] = [];
291-
let messageSeq = 0;
292-
for (const line of lines) {
293-
try {
294-
const parsed = JSON.parse(line);
295-
const message = parsedSessionEntryToMessage(parsed, messageSeq + 1);
296-
if (message) {
297-
messageSeq += 1;
298-
messages.push(message);
299-
}
300-
} catch {
301-
// ignore bad tail lines
302-
}
303-
}
304-
return messages.slice(-maxMessages);
223+
return parseRecentTranscriptTailMessages(lines, maxMessages);
305224
}) ?? []
306225
);
307226
}
@@ -401,24 +320,51 @@ function selectBoundedActiveTailRecords(entries: TailTranscriptRecord[]): TailTr
401320
return selected.toReversed();
402321
}
403322

404-
function parseRecentTranscriptTailMessages(lines: string[], maxMessages: number): unknown[] {
405-
const entries = lines.flatMap((line) => {
406-
const entry = parseTailTranscriptRecord(line);
407-
return entry ? [entry] : [];
323+
function readTranscriptRecords(filePath: string): TailTranscriptRecord[] {
324+
const records: TailTranscriptRecord[] = [];
325+
visitTranscriptLines(filePath, (line) => {
326+
if (!line.trim()) {
327+
return;
328+
}
329+
const record = parseTailTranscriptRecord(line);
330+
if (record && record.record.type !== "session") {
331+
records.push(record);
332+
}
408333
});
409-
const selected = entries.some(tailRecordHasTreeLink)
410-
? selectBoundedActiveTailRecords(entries)
411-
: entries;
334+
return records;
335+
}
336+
337+
function selectActiveTranscriptRecords(records: TailTranscriptRecord[]): TailTranscriptRecord[] {
338+
return records.some(tailRecordHasTreeLink) ? selectBoundedActiveTailRecords(records) : records;
339+
}
340+
341+
function readSelectedTranscriptRecords(filePath: string): TailTranscriptRecord[] {
342+
try {
343+
return selectActiveTranscriptRecords(readTranscriptRecords(filePath));
344+
} catch {
345+
return [];
346+
}
347+
}
348+
349+
function transcriptRecordsToMessages(records: TailTranscriptRecord[]): unknown[] {
412350
const messages: unknown[] = [];
413351
let messageSeq = 0;
414-
for (const entry of selected) {
352+
for (const entry of records) {
415353
const message = parsedSessionEntryToMessage(entry.record, messageSeq + 1);
416354
if (message) {
417355
messageSeq += 1;
418356
messages.push(message);
419357
}
420358
}
421-
return messages.slice(-maxMessages);
359+
return messages;
360+
}
361+
362+
function parseRecentTranscriptTailMessages(lines: string[], maxMessages: number): unknown[] {
363+
const entries = lines.flatMap((line) => {
364+
const entry = parseTailTranscriptRecord(line);
365+
return entry ? [entry] : [];
366+
});
367+
return transcriptRecordsToMessages(selectActiveTranscriptRecords(entries)).slice(-maxMessages);
422368
}
423369

424370
function visitTranscriptLines(filePath: string, visit: (line: string) => void): void {
@@ -479,61 +425,6 @@ async function visitTranscriptLinesAsync(
479425
}
480426
}
481427

482-
function transcriptHasTreeEntries(filePath: string): boolean {
483-
let hasTreeEntries = false;
484-
try {
485-
visitTranscriptLines(filePath, (line) => {
486-
if (!hasTreeEntries && hasSessionTreeEntry(line)) {
487-
hasTreeEntries = true;
488-
}
489-
});
490-
} catch {
491-
return false;
492-
}
493-
return hasTreeEntries;
494-
}
495-
496-
function visitSessionManagerBranchMessages(
497-
filePath: string,
498-
visit: (message: unknown, seq: number) => void,
499-
): number {
500-
const branchEntries = SessionManager.open(filePath).getBranch();
501-
let messageSeq = 0;
502-
for (const entry of branchEntries) {
503-
if (entry.type === "message" && entry.message) {
504-
messageSeq += 1;
505-
visit(
506-
attachOpenClawTranscriptMeta(entry.message, {
507-
...(typeof entry.id === "string" ? { id: entry.id } : {}),
508-
seq: messageSeq,
509-
}),
510-
messageSeq,
511-
);
512-
continue;
513-
}
514-
515-
if (entry.type === "compaction") {
516-
const ts = typeof entry.timestamp === "string" ? Date.parse(entry.timestamp) : Number.NaN;
517-
const timestamp = Number.isFinite(ts) ? ts : Date.now();
518-
messageSeq += 1;
519-
visit(
520-
{
521-
role: "system",
522-
content: [{ type: "text", text: "Compaction" }],
523-
timestamp,
524-
__openclaw: {
525-
kind: "compaction",
526-
id: typeof entry.id === "string" ? entry.id : undefined,
527-
seq: messageSeq,
528-
},
529-
},
530-
messageSeq,
531-
);
532-
}
533-
}
534-
return messageSeq;
535-
}
536-
537428
export function visitSessionMessages(
538429
sessionId: string,
539430
storePath: string | undefined,
@@ -545,35 +436,11 @@ export function visitSessionMessages(
545436
return 0;
546437
}
547438

548-
if (transcriptHasTreeEntries(filePath)) {
549-
try {
550-
return visitSessionManagerBranchMessages(filePath, visit);
551-
} catch {
552-
return 0;
553-
}
439+
const messages = transcriptRecordsToMessages(readSelectedTranscriptRecords(filePath));
440+
for (const [index, message] of messages.entries()) {
441+
visit(message, index + 1);
554442
}
555-
556-
let messageSeq = 0;
557-
try {
558-
visitTranscriptLines(filePath, (line) => {
559-
if (!line.trim()) {
560-
return;
561-
}
562-
try {
563-
const parsed = JSON.parse(line);
564-
const message = parsedSessionEntryToMessage(parsed, messageSeq + 1);
565-
if (message) {
566-
messageSeq += 1;
567-
visit(message, messageSeq);
568-
}
569-
} catch {
570-
// ignore bad lines
571-
}
572-
});
573-
} catch {
574-
return 0;
575-
}
576-
return messageSeq;
443+
return messages.length;
577444
}
578445

579446
export function readSessionMessageCount(
@@ -763,18 +630,6 @@ export function readRecentSessionTranscriptLines(params: {
763630
return { lines, totalLines };
764631
}
765632

766-
function hasSessionTreeEntry(line: string): boolean {
767-
if (!line.trim()) {
768-
return false;
769-
}
770-
try {
771-
const parsed = JSON.parse(line) as { type?: unknown; id?: unknown; parentId?: unknown };
772-
return parsed.type !== "session" && typeof parsed.id === "string" && "parentId" in parsed;
773-
} catch {
774-
return false;
775-
}
776-
}
777-
778633
function parsedSessionEntryToMessage(parsed: unknown, seq: number): unknown {
779634
if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
780635
return null;

0 commit comments

Comments
 (0)