fix(cli): record actual token usage for CLI backend responses
Previously, CLI backend responses wrote hardcoded zeros for token usage in session transcripts (input: 0, output: 0, totalTokens: 0). This caused the UI to show incorrect token counts and status to fall back to stale accumulated values. Changes: - Add usage parameter to appendMessageToTranscript and related functions in session-utils.fs.ts to accept NormalizedUsage from CLI backends - Pass result.meta.agentMeta?.usage when persisting assistant messages in agent-runner-execution.ts - Create CliSessionManager class with SDK-aligned API for future use: static factories (open/create), accessor methods, write locking - Add comprehensive tests for both session-utils.fs usage parameter and the new CliSessionManager class (17 + 2 new tests) Transcript entries now include actual input/output/cache token counts from CLI backends like claude-cli and opus.
This commit is contained in:
Ross
parent
4cb8f37202
commit
7fc6476688
@ -219,6 +219,7 @@ export async function runAgentTurnWithFallback(params: {
|
||||
createIfMissing: true,
|
||||
provider,
|
||||
model,
|
||||
usage: result.meta.agentMeta?.usage,
|
||||
});
|
||||
if (!assistantPersistResult.ok) {
|
||||
logVerbose(
|
||||
|
||||
347
src/gateway/cli-session-manager.test.ts
Normal file
347
src/gateway/cli-session-manager.test.ts
Normal file
@ -0,0 +1,347 @@
|
||||
import fs from "node:fs";
|
||||
import os from "node:os";
|
||||
import path from "node:path";
|
||||
import { afterEach, beforeEach, describe, expect, test } from "vitest";
|
||||
import { CliSessionManager } from "./cli-session-manager.js";
|
||||
|
||||
describe("CliSessionManager", () => {
|
||||
let tmpDir: string;
|
||||
|
||||
beforeEach(() => {
|
||||
tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "clawdbot-cli-session-manager-test-"));
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
fs.rmSync(tmpDir, { recursive: true, force: true });
|
||||
});
|
||||
|
||||
describe("static factories", () => {
|
||||
test("create() creates a new session manager with transcript file", () => {
|
||||
const sessionId = "test-session-create";
|
||||
const sessionFile = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
|
||||
const mgr = CliSessionManager.create({
|
||||
sessionId,
|
||||
sessionFile,
|
||||
});
|
||||
|
||||
expect(mgr.getSessionId()).toBe(sessionId);
|
||||
expect(mgr.getSessionFile()).toBe(sessionFile);
|
||||
expect(fs.existsSync(sessionFile)).toBe(true);
|
||||
|
||||
// Verify header was written
|
||||
const content = fs.readFileSync(sessionFile, "utf-8");
|
||||
const header = JSON.parse(content.trim());
|
||||
expect(header.type).toBe("session");
|
||||
expect(header.id).toBe(sessionId);
|
||||
});
|
||||
|
||||
test("open() opens existing session and loads entries", () => {
|
||||
const sessionId = "test-session-open";
|
||||
const sessionFile = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
|
||||
// Write some existing entries
|
||||
const lines = [
|
||||
JSON.stringify({ type: "session", version: 1, id: sessionId }),
|
||||
JSON.stringify({
|
||||
type: "message",
|
||||
id: "msg1",
|
||||
timestamp: new Date().toISOString(),
|
||||
message: {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Hello" }],
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
}),
|
||||
JSON.stringify({
|
||||
type: "message",
|
||||
id: "msg2",
|
||||
timestamp: new Date().toISOString(),
|
||||
message: {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "Hi" }],
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
}),
|
||||
];
|
||||
fs.writeFileSync(sessionFile, lines.join("\n") + "\n", "utf-8");
|
||||
|
||||
const mgr = CliSessionManager.open(sessionFile, { sessionId });
|
||||
|
||||
expect(mgr.getSessionId()).toBe(sessionId);
|
||||
expect(mgr.getEntries()).toHaveLength(2);
|
||||
expect(mgr.getLeafId()).toBe("msg2");
|
||||
});
|
||||
|
||||
test("open() creates file if it does not exist", () => {
|
||||
const sessionId = "test-session-open-new";
|
||||
const sessionFile = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
|
||||
expect(fs.existsSync(sessionFile)).toBe(false);
|
||||
|
||||
const mgr = CliSessionManager.open(sessionFile, { sessionId });
|
||||
|
||||
expect(fs.existsSync(sessionFile)).toBe(true);
|
||||
expect(mgr.getEntries()).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe("accessors", () => {
|
||||
test("getSessionId() returns session ID", () => {
|
||||
const sessionId = "accessor-test-session";
|
||||
const mgr = CliSessionManager.create({
|
||||
sessionId,
|
||||
sessionFile: path.join(tmpDir, `${sessionId}.jsonl`),
|
||||
});
|
||||
|
||||
expect(mgr.getSessionId()).toBe(sessionId);
|
||||
});
|
||||
|
||||
test("getLeafId() returns null initially, then entry ID after append", () => {
|
||||
const sessionId = "leaf-test-session";
|
||||
const mgr = CliSessionManager.create({
|
||||
sessionId,
|
||||
sessionFile: path.join(tmpDir, `${sessionId}.jsonl`),
|
||||
});
|
||||
|
||||
expect(mgr.getLeafId()).toBeNull();
|
||||
|
||||
const entryId = mgr.appendMessage({ role: "user", content: "Hello" });
|
||||
|
||||
expect(mgr.getLeafId()).toBe(entryId);
|
||||
});
|
||||
|
||||
test("getLeafEntry() returns the last appended entry", () => {
|
||||
const sessionId = "leaf-entry-test";
|
||||
const mgr = CliSessionManager.create({
|
||||
sessionId,
|
||||
sessionFile: path.join(tmpDir, `${sessionId}.jsonl`),
|
||||
});
|
||||
|
||||
expect(mgr.getLeafEntry()).toBeUndefined();
|
||||
|
||||
mgr.appendMessage({ role: "user", content: "Hello" });
|
||||
const leafEntry = mgr.getLeafEntry();
|
||||
|
||||
expect(leafEntry).toBeDefined();
|
||||
expect(leafEntry?.type).toBe("message");
|
||||
expect(leafEntry?.message?.role).toBe("user");
|
||||
});
|
||||
|
||||
test("isPersisted() returns true when sessionFile is set", () => {
|
||||
const mgr = CliSessionManager.create({
|
||||
sessionId: "persisted-test",
|
||||
sessionFile: path.join(tmpDir, "persisted.jsonl"),
|
||||
});
|
||||
|
||||
expect(mgr.isPersisted()).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe("appendMessage", () => {
|
||||
test("appends user message and returns entry ID", () => {
|
||||
const sessionId = "append-user-test";
|
||||
const sessionFile = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
const mgr = CliSessionManager.create({ sessionId, sessionFile });
|
||||
|
||||
const entryId = mgr.appendMessage({ role: "user", content: "Hello world" });
|
||||
|
||||
expect(entryId).toBeDefined();
|
||||
expect(typeof entryId).toBe("string");
|
||||
expect(entryId.length).toBe(8);
|
||||
|
||||
const entries = mgr.getEntries();
|
||||
expect(entries).toHaveLength(1);
|
||||
expect(entries[0]?.message?.role).toBe("user");
|
||||
});
|
||||
|
||||
test("appends assistant message with usage data", () => {
|
||||
const sessionId = "append-assistant-usage-test";
|
||||
const sessionFile = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
const mgr = CliSessionManager.create({
|
||||
sessionId,
|
||||
sessionFile,
|
||||
provider: "claude-cli",
|
||||
model: "opus",
|
||||
});
|
||||
|
||||
const entryId = mgr.appendMessage({
|
||||
role: "assistant",
|
||||
content: "Hello!",
|
||||
usage: {
|
||||
input: 100,
|
||||
output: 50,
|
||||
cacheRead: 10,
|
||||
cacheWrite: 5,
|
||||
total: 165,
|
||||
},
|
||||
});
|
||||
|
||||
expect(entryId).toBeDefined();
|
||||
|
||||
// Verify written to file
|
||||
const content = fs.readFileSync(sessionFile, "utf-8");
|
||||
const lines = content.trim().split("\n");
|
||||
// Skip header, get the message entry
|
||||
const messageEntry = JSON.parse(lines[lines.length - 1]);
|
||||
|
||||
expect(messageEntry.message.role).toBe("assistant");
|
||||
expect(messageEntry.message.usage).toEqual({
|
||||
input: 100,
|
||||
output: 50,
|
||||
cacheRead: 10,
|
||||
cacheWrite: 5,
|
||||
totalTokens: 165,
|
||||
});
|
||||
expect(messageEntry.message.provider).toBe("claude-cli");
|
||||
expect(messageEntry.message.model).toBe("opus");
|
||||
expect(messageEntry.message.stopReason).toBe("cli_backend");
|
||||
});
|
||||
|
||||
test("usage defaults to zeros when not provided", () => {
|
||||
const sessionId = "append-no-usage-test";
|
||||
const sessionFile = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
const mgr = CliSessionManager.create({ sessionId, sessionFile });
|
||||
|
||||
mgr.appendMessage({ role: "assistant", content: "Response" });
|
||||
|
||||
const content = fs.readFileSync(sessionFile, "utf-8");
|
||||
const lines = content.trim().split("\n");
|
||||
const messageEntry = JSON.parse(lines[lines.length - 1]);
|
||||
|
||||
expect(messageEntry.message.usage).toEqual({
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: undefined,
|
||||
cacheWrite: undefined,
|
||||
totalTokens: 0,
|
||||
});
|
||||
});
|
||||
|
||||
test("handles string content", () => {
|
||||
const sessionId = "string-content-test";
|
||||
const sessionFile = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
const mgr = CliSessionManager.create({ sessionId, sessionFile });
|
||||
|
||||
mgr.appendMessage({ role: "user", content: "Plain string message" });
|
||||
|
||||
const entry = mgr.getLeafEntry();
|
||||
expect(entry?.message?.content).toEqual([{ type: "text", text: "Plain string message" }]);
|
||||
});
|
||||
|
||||
test("handles array content", () => {
|
||||
const sessionId = "array-content-test";
|
||||
const sessionFile = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
const mgr = CliSessionManager.create({ sessionId, sessionFile });
|
||||
|
||||
mgr.appendMessage({
|
||||
role: "user",
|
||||
content: [
|
||||
{ type: "text", text: "Part 1" },
|
||||
{ type: "text", text: "Part 2" },
|
||||
],
|
||||
});
|
||||
|
||||
const entry = mgr.getLeafEntry();
|
||||
expect(entry?.message?.content).toEqual([
|
||||
{ type: "text", text: "Part 1" },
|
||||
{ type: "text", text: "Part 2" },
|
||||
]);
|
||||
});
|
||||
|
||||
test("sets parentId to previous leaf", () => {
|
||||
const sessionId = "parent-id-test";
|
||||
const sessionFile = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
const mgr = CliSessionManager.create({ sessionId, sessionFile });
|
||||
|
||||
const id1 = mgr.appendMessage({ role: "user", content: "First" });
|
||||
const id2 = mgr.appendMessage({ role: "assistant", content: "Second" });
|
||||
|
||||
const entries = mgr.getEntries();
|
||||
expect(entries[0]?.parentId).toBeUndefined();
|
||||
expect(entries[1]?.parentId).toBe(id1);
|
||||
expect(mgr.getLeafId()).toBe(id2);
|
||||
});
|
||||
});
|
||||
|
||||
describe("appendCustomEntry", () => {
|
||||
test("appends custom entry and returns entry ID", () => {
|
||||
const sessionId = "custom-entry-test";
|
||||
const sessionFile = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
const mgr = CliSessionManager.create({ sessionId, sessionFile });
|
||||
|
||||
const entryId = mgr.appendCustomEntry("tool_result", { result: "success" });
|
||||
|
||||
expect(entryId).toBeDefined();
|
||||
expect(mgr.getLeafId()).toBe(entryId);
|
||||
|
||||
const content = fs.readFileSync(sessionFile, "utf-8");
|
||||
const lines = content.trim().split("\n");
|
||||
const customEntry = JSON.parse(lines[lines.length - 1]);
|
||||
|
||||
expect(customEntry.type).toBe("custom");
|
||||
expect(customEntry.customType).toBe("tool_result");
|
||||
expect(customEntry.data).toEqual({ result: "success" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("CLI backend session ID", () => {
|
||||
test("setCliBackendSessionId and getCliBackendSessionId", () => {
|
||||
const mgr = CliSessionManager.create({
|
||||
sessionId: "cli-backend-id-test",
|
||||
sessionFile: path.join(tmpDir, "cli-backend.jsonl"),
|
||||
});
|
||||
|
||||
expect(mgr.getCliBackendSessionId()).toBeUndefined();
|
||||
|
||||
mgr.setCliBackendSessionId("claude-cli-session-123");
|
||||
|
||||
expect(mgr.getCliBackendSessionId()).toBe("claude-cli-session-123");
|
||||
});
|
||||
});
|
||||
|
||||
describe("archive", () => {
|
||||
test("archives transcript file with reason suffix", () => {
|
||||
const sessionId = "archive-test";
|
||||
const sessionFile = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
const mgr = CliSessionManager.create({ sessionId, sessionFile });
|
||||
|
||||
mgr.appendMessage({ role: "user", content: "Message to archive" });
|
||||
|
||||
const archivedPath = mgr.archive("context-overflow");
|
||||
|
||||
expect(archivedPath).toContain("context-overflow");
|
||||
expect(fs.existsSync(archivedPath)).toBe(true);
|
||||
expect(fs.existsSync(sessionFile)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("concurrent write safety", () => {
|
||||
test("appendMessageAsync serializes writes", async () => {
|
||||
const sessionId = "concurrent-test";
|
||||
const sessionFile = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
const mgr = CliSessionManager.create({ sessionId, sessionFile });
|
||||
|
||||
// Fire multiple async appends concurrently
|
||||
const promises = [
|
||||
mgr.appendMessageAsync({ role: "user", content: "Message 1" }),
|
||||
mgr.appendMessageAsync({ role: "assistant", content: "Response 1" }),
|
||||
mgr.appendMessageAsync({ role: "user", content: "Message 2" }),
|
||||
mgr.appendMessageAsync({ role: "assistant", content: "Response 2" }),
|
||||
];
|
||||
|
||||
const ids = await Promise.all(promises);
|
||||
|
||||
expect(ids).toHaveLength(4);
|
||||
expect(new Set(ids).size).toBe(4); // All unique IDs
|
||||
|
||||
const entries = mgr.getEntries();
|
||||
expect(entries).toHaveLength(4);
|
||||
|
||||
// Verify parent chain is correct
|
||||
for (let i = 1; i < entries.length; i++) {
|
||||
expect(entries[i]?.parentId).toBe(entries[i - 1]?.id);
|
||||
}
|
||||
});
|
||||
});
|
||||
});
|
||||
359
src/gateway/cli-session-manager.ts
Normal file
359
src/gateway/cli-session-manager.ts
Normal file
@ -0,0 +1,359 @@
|
||||
import crypto from "node:crypto";
|
||||
import fs from "node:fs";
|
||||
import path from "node:path";
|
||||
|
||||
import type { NormalizedUsage } from "../agents/usage.js";
|
||||
import { ensureTranscriptFile } from "./session-utils.fs.js";
|
||||
|
||||
/**
|
||||
* Content block type for CLI messages (matches SDK format).
|
||||
*/
|
||||
export type CliContentBlock =
|
||||
| { type: "text"; text: string }
|
||||
| { type: "input_text"; text: string }
|
||||
| { type: "output_text"; text: string };
|
||||
|
||||
/**
|
||||
* Message type for CLI session manager with usage support.
|
||||
*/
|
||||
export type CliMessage = {
|
||||
role: "user" | "assistant";
|
||||
content: string | CliContentBlock[];
|
||||
provider?: string;
|
||||
model?: string;
|
||||
/** Token usage for assistant messages. */
|
||||
usage?: NormalizedUsage;
|
||||
stopReason?: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* Entry in a CLI session transcript (matches SDK SessionEntry format).
|
||||
*/
|
||||
export type CliSessionEntry = {
|
||||
type: "message" | "custom" | "session";
|
||||
id: string;
|
||||
parentId?: string;
|
||||
timestamp: string;
|
||||
message?: {
|
||||
role: string;
|
||||
content: CliContentBlock[];
|
||||
timestamp: number;
|
||||
stopReason?: string;
|
||||
usage?: {
|
||||
input: number;
|
||||
output: number;
|
||||
cacheRead?: number;
|
||||
cacheWrite?: number;
|
||||
totalTokens: number;
|
||||
};
|
||||
provider?: string;
|
||||
model?: string;
|
||||
};
|
||||
customType?: string;
|
||||
data?: unknown;
|
||||
};
|
||||
|
||||
/**
|
||||
* Parameters for creating a new CLI session manager.
|
||||
*/
|
||||
export type CliSessionManagerParams = {
|
||||
sessionId: string;
|
||||
sessionFile?: string;
|
||||
storePath?: string;
|
||||
provider?: string;
|
||||
model?: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* In-process lock for session transcript writes.
|
||||
* Ensures atomic user+assistant message pairs within the same Node process.
|
||||
*/
|
||||
const sessionWriteLocks = new Map<string, Promise<void>>();
|
||||
|
||||
async function withSessionLock<T>(sessionId: string, fn: () => T | Promise<T>): Promise<T> {
|
||||
const pending = sessionWriteLocks.get(sessionId);
|
||||
if (pending) {
|
||||
await pending.catch(() => {});
|
||||
}
|
||||
|
||||
let resolve: () => void;
|
||||
const lock = new Promise<void>((r) => {
|
||||
resolve = r;
|
||||
});
|
||||
sessionWriteLocks.set(sessionId, lock);
|
||||
|
||||
try {
|
||||
return await fn();
|
||||
} finally {
|
||||
resolve!();
|
||||
if (sessionWriteLocks.get(sessionId) === lock) {
|
||||
sessionWriteLocks.delete(sessionId);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* CLI Session Manager - manages session transcripts for CLI backends.
|
||||
*
|
||||
* API aligned with SDK's SessionManager where appropriate:
|
||||
* - Static factories (open/create) instead of public constructor
|
||||
* - Accessor methods (getSessionId, getSessionFile, etc.)
|
||||
* - appendMessage returns entry ID (string)
|
||||
* - Write lock for concurrency safety
|
||||
*
|
||||
* Key difference from SDK: CLI doesn't need branching/navigation (linear transcript only).
|
||||
*/
|
||||
export class CliSessionManager {
|
||||
private readonly sessionId: string;
|
||||
private readonly sessionFile: string | undefined;
|
||||
private readonly storePath: string | undefined;
|
||||
private readonly provider: string | undefined;
|
||||
private readonly model: string | undefined;
|
||||
private readonly entries: CliSessionEntry[] = [];
|
||||
private leafId: string | null = null;
|
||||
private cliBackendSessionId: string | undefined;
|
||||
|
||||
private constructor(params: CliSessionManagerParams) {
|
||||
this.sessionId = params.sessionId;
|
||||
this.sessionFile = params.sessionFile;
|
||||
this.storePath = params.storePath;
|
||||
this.provider = params.provider;
|
||||
this.model = params.model;
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a new CLI session manager (SDK-style factory).
|
||||
* Creates the transcript file if it doesn't exist.
|
||||
*/
|
||||
static create(params: CliSessionManagerParams): CliSessionManager {
|
||||
const manager = new CliSessionManager(params);
|
||||
const transcriptPath = manager.resolveTranscriptPath();
|
||||
if (transcriptPath) {
|
||||
ensureTranscriptFile({ transcriptPath, sessionId: params.sessionId });
|
||||
}
|
||||
return manager;
|
||||
}
|
||||
|
||||
/**
|
||||
* Opens an existing session from a transcript file (SDK-style factory).
|
||||
* Creates the transcript file if it doesn't exist.
|
||||
*/
|
||||
static open(sessionFile: string, params?: Partial<CliSessionManagerParams>): CliSessionManager {
|
||||
const sessionId = params?.sessionId ?? crypto.randomUUID();
|
||||
const manager = new CliSessionManager({
|
||||
sessionId,
|
||||
sessionFile,
|
||||
storePath: params?.storePath,
|
||||
provider: params?.provider,
|
||||
model: params?.model,
|
||||
});
|
||||
|
||||
// Ensure transcript file exists
|
||||
ensureTranscriptFile({ transcriptPath: sessionFile, sessionId });
|
||||
|
||||
// Load existing entries from transcript
|
||||
manager.loadEntries(sessionFile);
|
||||
|
||||
return manager;
|
||||
}
|
||||
|
||||
// --- Accessors (matching SDK surface) ---
|
||||
|
||||
getSessionId(): string {
|
||||
return this.sessionId;
|
||||
}
|
||||
|
||||
getSessionFile(): string | undefined {
|
||||
return this.sessionFile ?? this.resolveTranscriptPath() ?? undefined;
|
||||
}
|
||||
|
||||
getLeafId(): string | null {
|
||||
return this.leafId;
|
||||
}
|
||||
|
||||
getLeafEntry(): CliSessionEntry | undefined {
|
||||
if (!this.leafId) return undefined;
|
||||
return this.entries.find((e) => e.id === this.leafId);
|
||||
}
|
||||
|
||||
getEntries(): CliSessionEntry[] {
|
||||
return [...this.entries];
|
||||
}
|
||||
|
||||
isPersisted(): boolean {
|
||||
return Boolean(this.sessionFile || this.storePath);
|
||||
}
|
||||
|
||||
// --- CLI-specific extensions ---
|
||||
|
||||
setCliBackendSessionId(cliSessionId: string): void {
|
||||
this.cliBackendSessionId = cliSessionId;
|
||||
}
|
||||
|
||||
getCliBackendSessionId(): string | undefined {
|
||||
return this.cliBackendSessionId;
|
||||
}
|
||||
|
||||
// --- Append Methods (matching SDK pattern, returns entry ID) ---
|
||||
|
||||
/**
|
||||
* Appends a message to the session transcript.
|
||||
* Returns the entry ID (string) like SDK SessionManager.
|
||||
*/
|
||||
appendMessage(message: CliMessage): string {
|
||||
const transcriptPath = this.resolveTranscriptPath();
|
||||
if (!transcriptPath) {
|
||||
throw new Error("Cannot append message: no transcript path available");
|
||||
}
|
||||
|
||||
const now = Date.now();
|
||||
const entryId = crypto.randomUUID().slice(0, 8);
|
||||
const content = this.normalizeContent(message.content);
|
||||
|
||||
const messageBody: CliSessionEntry["message"] = {
|
||||
role: message.role,
|
||||
content,
|
||||
timestamp: now,
|
||||
};
|
||||
|
||||
if (message.role === "assistant") {
|
||||
messageBody.stopReason = message.stopReason ?? "cli_backend";
|
||||
const u = message.usage;
|
||||
messageBody.usage = {
|
||||
input: u?.input ?? 0,
|
||||
output: u?.output ?? 0,
|
||||
cacheRead: u?.cacheRead,
|
||||
cacheWrite: u?.cacheWrite,
|
||||
totalTokens: u?.total ?? (u?.input ?? 0) + (u?.output ?? 0),
|
||||
};
|
||||
if (message.provider ?? this.provider) {
|
||||
messageBody.provider = message.provider ?? this.provider;
|
||||
}
|
||||
if (message.model ?? this.model) {
|
||||
messageBody.model = message.model ?? this.model;
|
||||
}
|
||||
}
|
||||
|
||||
const entry: CliSessionEntry = {
|
||||
type: "message",
|
||||
id: entryId,
|
||||
parentId: this.leafId ?? undefined,
|
||||
timestamp: new Date(now).toISOString(),
|
||||
message: messageBody,
|
||||
};
|
||||
|
||||
// Append to transcript file
|
||||
try {
|
||||
fs.appendFileSync(transcriptPath, `${JSON.stringify(entry)}\n`, "utf-8");
|
||||
} catch (err) {
|
||||
throw new Error(
|
||||
`Failed to write to transcript: ${err instanceof Error ? err.message : String(err)}`,
|
||||
);
|
||||
}
|
||||
|
||||
// Update in-memory state
|
||||
this.entries.push(entry);
|
||||
this.leafId = entryId;
|
||||
|
||||
return entryId;
|
||||
}
|
||||
|
||||
/**
|
||||
* Appends a message asynchronously with session-level locking.
|
||||
* Returns the entry ID (string) like SDK SessionManager.
|
||||
*/
|
||||
async appendMessageAsync(message: CliMessage): Promise<string> {
|
||||
return withSessionLock(this.sessionId, () => this.appendMessage(message));
|
||||
}
|
||||
|
||||
/**
|
||||
* Appends a custom entry to the session transcript.
|
||||
* Returns the entry ID (string) like SDK SessionManager.
|
||||
*/
|
||||
appendCustomEntry(customType: string, data?: unknown): string {
|
||||
const transcriptPath = this.resolveTranscriptPath();
|
||||
if (!transcriptPath) {
|
||||
throw new Error("Cannot append custom entry: no transcript path available");
|
||||
}
|
||||
|
||||
const entryId = crypto.randomUUID().slice(0, 8);
|
||||
const entry: CliSessionEntry = {
|
||||
type: "custom",
|
||||
id: entryId,
|
||||
parentId: this.leafId ?? undefined,
|
||||
timestamp: new Date().toISOString(),
|
||||
customType,
|
||||
data,
|
||||
};
|
||||
|
||||
try {
|
||||
fs.appendFileSync(transcriptPath, `${JSON.stringify(entry)}\n`, "utf-8");
|
||||
} catch (err) {
|
||||
throw new Error(
|
||||
`Failed to write custom entry: ${err instanceof Error ? err.message : String(err)}`,
|
||||
);
|
||||
}
|
||||
|
||||
this.entries.push(entry);
|
||||
this.leafId = entryId;
|
||||
|
||||
return entryId;
|
||||
}
|
||||
|
||||
/**
|
||||
* Archives the session transcript with a reason suffix.
|
||||
* Returns the path to the archived file.
|
||||
*/
|
||||
archive(reason: string): string {
|
||||
const transcriptPath = this.resolveTranscriptPath();
|
||||
if (!transcriptPath || !fs.existsSync(transcriptPath)) {
|
||||
throw new Error("Cannot archive: transcript file not found");
|
||||
}
|
||||
|
||||
const ts = new Date().toISOString().replaceAll(":", "-");
|
||||
const archived = `${transcriptPath}.${reason}.${ts}`;
|
||||
fs.renameSync(transcriptPath, archived);
|
||||
return archived;
|
||||
}
|
||||
|
||||
// --- Private helpers ---
|
||||
|
||||
private resolveTranscriptPath(): string | null {
|
||||
if (this.sessionFile) return this.sessionFile;
|
||||
if (this.storePath) {
|
||||
return path.join(path.dirname(this.storePath), `${this.sessionId}.jsonl`);
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
private normalizeContent(content: string | CliContentBlock[]): CliContentBlock[] {
|
||||
if (typeof content === "string") {
|
||||
return [{ type: "text", text: content }];
|
||||
}
|
||||
return content;
|
||||
}
|
||||
|
||||
private loadEntries(filePath: string): void {
|
||||
if (!fs.existsSync(filePath)) return;
|
||||
|
||||
try {
|
||||
const data = fs.readFileSync(filePath, "utf-8");
|
||||
const lines = data.split(/\r?\n/);
|
||||
|
||||
for (const line of lines) {
|
||||
if (!line.trim()) continue;
|
||||
try {
|
||||
const parsed = JSON.parse(line) as CliSessionEntry;
|
||||
if (parsed.type === "message" || parsed.type === "custom") {
|
||||
this.entries.push(parsed);
|
||||
this.leafId = parsed.id;
|
||||
}
|
||||
} catch {
|
||||
// Skip malformed lines
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
// File read error - start fresh
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -573,6 +573,63 @@ describe("appendMessageToTranscript", () => {
|
||||
expect(msg.message.provider).toBeUndefined();
|
||||
expect(msg.message.model).toBeUndefined();
|
||||
});
|
||||
|
||||
test("includes usage data when provided for assistant message", () => {
|
||||
const sessionId = "test-append-usage";
|
||||
const transcriptPath = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
const header = JSON.stringify({ type: "session", version: 1, id: sessionId });
|
||||
fs.writeFileSync(transcriptPath, header + "\n", "utf-8");
|
||||
|
||||
const result = appendMessageToTranscript({
|
||||
message: "Response with usage",
|
||||
role: "assistant",
|
||||
sessionId,
|
||||
storePath,
|
||||
usage: {
|
||||
input: 100,
|
||||
output: 50,
|
||||
cacheRead: 10,
|
||||
cacheWrite: 5,
|
||||
total: 165,
|
||||
},
|
||||
});
|
||||
|
||||
expect(result.ok).toBe(true);
|
||||
|
||||
const content = fs.readFileSync(transcriptPath, "utf-8");
|
||||
const lines = content.trim().split("\n");
|
||||
const msg = JSON.parse(lines[1]);
|
||||
expect(msg.message.usage).toEqual({
|
||||
input: 100,
|
||||
output: 50,
|
||||
cacheRead: 10,
|
||||
cacheWrite: 5,
|
||||
totalTokens: 165,
|
||||
});
|
||||
});
|
||||
|
||||
test("defaults usage to zeros when not provided for assistant message", () => {
|
||||
const sessionId = "test-append-no-usage";
|
||||
const transcriptPath = path.join(tmpDir, `${sessionId}.jsonl`);
|
||||
const header = JSON.stringify({ type: "session", version: 1, id: sessionId });
|
||||
fs.writeFileSync(transcriptPath, header + "\n", "utf-8");
|
||||
|
||||
const result = appendMessageToTranscript({
|
||||
message: "Response without usage",
|
||||
role: "assistant",
|
||||
sessionId,
|
||||
storePath,
|
||||
});
|
||||
|
||||
expect(result.ok).toBe(true);
|
||||
|
||||
const content = fs.readFileSync(transcriptPath, "utf-8");
|
||||
const lines = content.trim().split("\n");
|
||||
const msg = JSON.parse(lines[1]);
|
||||
expect(msg.message.usage.input).toBe(0);
|
||||
expect(msg.message.usage.output).toBe(0);
|
||||
expect(msg.message.usage.totalTokens).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe("ensureTranscriptFile", () => {
|
||||
|
||||
@ -4,6 +4,7 @@ import os from "node:os";
|
||||
import path from "node:path";
|
||||
|
||||
import { CURRENT_SESSION_VERSION } from "@mariozechner/pi-coding-agent";
|
||||
import type { NormalizedUsage } from "../agents/usage.js";
|
||||
import { resolveSessionTranscriptPath } from "../config/sessions.js";
|
||||
import { stripEnvelope } from "./chat-sanitize.js";
|
||||
import type { SessionPreviewItem } from "./session-utils.types.js";
|
||||
@ -112,6 +113,8 @@ function appendMessageToTranscriptImpl(
|
||||
stopReason?: string;
|
||||
provider?: string;
|
||||
model?: string;
|
||||
/** Token usage for assistant messages (CLI backends). */
|
||||
usage?: NormalizedUsage;
|
||||
},
|
||||
transcriptPath: string,
|
||||
): TranscriptAppendResult {
|
||||
@ -137,7 +140,14 @@ function appendMessageToTranscriptImpl(
|
||||
};
|
||||
if (params.role === "assistant") {
|
||||
messageBody.stopReason = params.stopReason ?? "cli_backend";
|
||||
messageBody.usage = { input: 0, output: 0, totalTokens: 0 };
|
||||
const u = params.usage;
|
||||
messageBody.usage = {
|
||||
input: u?.input ?? 0,
|
||||
output: u?.output ?? 0,
|
||||
cacheRead: u?.cacheRead,
|
||||
cacheWrite: u?.cacheWrite,
|
||||
totalTokens: u?.total ?? (u?.input ?? 0) + (u?.output ?? 0),
|
||||
};
|
||||
if (params.provider) messageBody.provider = params.provider;
|
||||
if (params.model) messageBody.model = params.model;
|
||||
}
|
||||
@ -176,6 +186,8 @@ export function appendMessageToTranscript(params: {
|
||||
provider?: string;
|
||||
/** Model that generated the response (for assistant messages). */
|
||||
model?: string;
|
||||
/** Token usage for assistant messages (CLI backends). */
|
||||
usage?: NormalizedUsage;
|
||||
}): TranscriptAppendResult {
|
||||
const transcriptPath = resolveTranscriptPathFromParams(params);
|
||||
if (!transcriptPath) {
|
||||
@ -199,6 +211,8 @@ export async function appendMessageToTranscriptAsync(params: {
|
||||
stopReason?: string;
|
||||
provider?: string;
|
||||
model?: string;
|
||||
/** Token usage for assistant messages (CLI backends). */
|
||||
usage?: NormalizedUsage;
|
||||
}): Promise<TranscriptAppendResult> {
|
||||
const transcriptPath = resolveTranscriptPathFromParams(params);
|
||||
if (!transcriptPath) {
|
||||
@ -223,6 +237,8 @@ export function appendAssistantMessageToTranscript(params: {
|
||||
provider?: string;
|
||||
/** Model that generated the response. */
|
||||
model?: string;
|
||||
/** Token usage for assistant messages (CLI backends). */
|
||||
usage?: NormalizedUsage;
|
||||
}): TranscriptAppendResult {
|
||||
return appendMessageToTranscript({
|
||||
...params,
|
||||
@ -241,6 +257,8 @@ export async function appendAssistantMessageToTranscriptAsync(params: {
|
||||
createIfMissing?: boolean;
|
||||
provider?: string;
|
||||
model?: string;
|
||||
/** Token usage for assistant messages (CLI backends). */
|
||||
usage?: NormalizedUsage;
|
||||
}): Promise<TranscriptAppendResult> {
|
||||
return appendMessageToTranscriptAsync({
|
||||
...params,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user