From bf2a8ecfdb89132de59f5b397867b1ea3beee1a3 Mon Sep 17 00:00:00 2001
From: lin-hongkuan <lin-hongkuan@users.noreply.github.com>
Date: Thu, 25 Jun 2026 22:33:35 +0800
Subject: [PATCH] fix(trajectory): preserve usage in truncated events

---
 src/trajectory/runtime.test.ts | 55 ++++++++++++++++++++++++++++---
 src/trajectory/runtime.ts      | 60 +++++++++++++++++++++++++++-------
 2 files changed, 98 insertions(+), 17 deletions(-)

diff --git a/src/trajectory/runtime.test.ts b/src/trajectory/runtime.test.ts
index 65690f05043..5a03f72caf0 100644
--- a/src/trajectory/runtime.test.ts
+++ b/src/trajectory/runtime.test.ts
@@ -9,10 +9,7 @@ import {
   resolveTrajectoryPointerFilePath,
   resolveTrajectoryPointerOpenFlags,
 } from "./paths.js";
-import {
-  createTrajectoryRuntimeRecorder,
-  toTrajectoryToolDefinitions,
-} from "./runtime.js";
+import { createTrajectoryRuntimeRecorder, toTrajectoryToolDefinitions } from "./runtime.js";
 
 type TrajectoryRuntimeRecorder = NonNullable<ReturnType<typeof createTrajectoryRuntimeRecorder>>;
 
@@ -57,7 +54,7 @@ describe("trajectory runtime", () => {
         env: { OPENCLAW_TRAJECTORY_DIR: "/tmp/traces" },
         sessionId: "../evil/session",
       }),
-    ).toBe("/tmp/traces/___evil_session.jsonl");
+    ).toBe(path.join(path.resolve("/tmp/traces"), "___evil_session.jsonl"));
   });
 
   it("records sanitized runtime events by default", () => {
@@ -136,6 +133,54 @@ describe("trajectory runtime", () => {
     );
   });
 
+  it("preserves usage when truncating oversized runtime events", () => {
+    const writes: string[] = [];
+    const usage = {
+      input: 384_954,
+      output: 5_624,
+      cacheRead: 333_824,
+      reasoningTokens: 2_038,
+      total: 724_402,
+    };
+    const promptCache = { readTokens: 333_824, writeTokens: 51_130 };
+    const recorder = createTrajectoryRuntimeRecorder({
+      sessionId: "session-1",
+      sessionFile: "/tmp/session.jsonl",
+      writer: {
+        filePath: "/tmp/session.trajectory.jsonl",
+        write: (line) => {
+          writes.push(line);
+        },
+        flush: async () => undefined,
+      },
+    });
+
+    const runtimeRecorder = expectTrajectoryRuntimeRecorder(recorder);
+    runtimeRecorder.recordEvent("model.completed", {
+      usage,
+      promptCache,
+      messagesSnapshot: Array.from({ length: 12 }, (_value, index) => ({
+        role: index % 2 === 0 ? "user" : "assistant",
+        content: `message-${index} ${"x".repeat(32_000)}`,
+      })),
+    });
+
+    expect(writes).toHaveLength(1);
+    const parsed = JSON.parse(writes[0]);
+    expect(parsed.type).toBe("model.completed");
+    expect(parsed.data).toMatchObject({
+      truncated: true,
+      reason: "trajectory-event-size-limit",
+      usage,
+      promptCache,
+    });
+    expect(parsed.data.messagesSnapshot).toBeUndefined();
+    expect(parsed.data.droppedFields).toContain("messagesSnapshot");
+    expect(Buffer.byteLength(writes[0], "utf8")).toBeLessThanOrEqual(
+      TRAJECTORY_RUNTIME_EVENT_MAX_BYTES + 1,
+    );
+  });
+
   it("rotates runtime capture at the file budget and keeps newer events", async () => {
     const tmpDir = makeTempDir();
     const sessionFile = path.join(tmpDir, "session.jsonl");
diff --git a/src/trajectory/runtime.ts b/src/trajectory/runtime.ts
index 414d7b33be5..4ce5e6613f2 100644
--- a/src/trajectory/runtime.ts
+++ b/src/trajectory/runtime.ts
@@ -52,6 +52,7 @@ const TRAJECTORY_RUNTIME_DATA_STRING_MAX_CHARS = 32_768;
 const TRAJECTORY_RUNTIME_DATA_ARRAY_MAX_ITEMS = 64;
 const TRAJECTORY_RUNTIME_DATA_OBJECT_MAX_KEYS = 64;
 const TRAJECTORY_RUNTIME_DATA_MAX_DEPTH = 6;
+const TRAJECTORY_RUNTIME_OVERSIZE_PRESERVED_DATA_KEYS = ["usage", "promptCache"] as const;
 
 type TrajectoryRuntimeWriterDiagnostics = Omit<QueuedFileWriterDiagnostics, "activeOperation"> & {
   activeOperation: QueuedFileWriterDiagnostics["activeOperation"] | "file-replace";
@@ -128,19 +129,54 @@ function truncateOversizedTrajectoryEvent(
   if (bytes <= TRAJECTORY_RUNTIME_EVENT_MAX_BYTES) {
     return line;
   }
-  const truncated = safeJsonStringify({
-    ...event,
-    data: {
-      truncated: true,
-      originalBytes: bytes,
-      limitBytes: TRAJECTORY_RUNTIME_EVENT_MAX_BYTES,
-      reason: "trajectory-event-size-limit",
-    },
-  });
-  if (truncated && Buffer.byteLength(truncated, "utf8") <= TRAJECTORY_RUNTIME_EVENT_MAX_BYTES) {
-    return truncated;
+
+  const originalData = event.data ?? {};
+  const originalDataKeys = Object.keys(originalData);
+  const preservedDataKeys = new Set<string>();
+  const baseData = {
+    truncated: true,
+    originalBytes: bytes,
+    limitBytes: TRAJECTORY_RUNTIME_EVENT_MAX_BYTES,
+    reason: "trajectory-event-size-limit",
+  };
+  const buildTruncatedEventLine = (includeDroppedFields: boolean): string | undefined => {
+    const data: Record<string, unknown> = { ...baseData };
+    for (const key of TRAJECTORY_RUNTIME_OVERSIZE_PRESERVED_DATA_KEYS) {
+      if (preservedDataKeys.has(key)) {
+        data[key] = originalData[key];
+      }
+    }
+    if (includeDroppedFields) {
+      const droppedFields = originalDataKeys.filter((key) => !preservedDataKeys.has(key));
+      if (droppedFields.length > 0) {
+        data.droppedFields = droppedFields;
+      }
+    }
+    const truncated = safeJsonStringify({ ...event, data });
+    if (truncated && Buffer.byteLength(truncated, "utf8") <= TRAJECTORY_RUNTIME_EVENT_MAX_BYTES) {
+      return truncated;
+    }
+    return undefined;
+  };
+
+  let best = buildTruncatedEventLine(true) ?? buildTruncatedEventLine(false);
+  if (!best) {
+    return undefined;
   }
-  return undefined;
+
+  for (const key of TRAJECTORY_RUNTIME_OVERSIZE_PRESERVED_DATA_KEYS) {
+    if (!Object.hasOwn(originalData, key)) {
+      continue;
+    }
+    preservedDataKeys.add(key);
+    const next = buildTruncatedEventLine(true) ?? buildTruncatedEventLine(false);
+    if (next) {
+      best = next;
+      continue;
+    }
+    preservedDataKeys.delete(key);
+  }
+  return best;
 }
 
 function truncatedTrajectoryValue(reason: string, details: Record<string, unknown> = {}): unknown {