feat: add usage tracking, output mode, and cli exit mode (#7)

src/cli/helpers/accumulator.ts

@@ -55,6 +55,14 @@ export type Buffers = {
   toolCallIdToLineId: Map<string, string>;
   lastOtid: string | null; // Track the last otid to detect transitions
   pendingRefresh?: boolean; // Track throttled refresh state
+  usage: {
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+    cachedTokens: number;
+    reasoningTokens: number;
+    stepCount: number;
+  };
 };
 
 export function createBuffers(): Buffers {
@@ -65,6 +73,14 @@ export function createBuffers(): Buffers {
     pendingToolByRun: new Map(),
     toolCallIdToLineId: new Map(),
     lastOtid: null,
+    usage: {
+      promptTokens: 0,
+      completionTokens: 0,
+      totalTokens: 0,
+      cachedTokens: 0,
+      reasoningTokens: 0,
+      stepCount: 0,
+    },
   };
 }
 
@@ -339,8 +355,26 @@ export function onChunk(
       break;
     }
 
+    case "usage_statistics": {
+      // Accumulate usage statistics from the stream
+      // These messages arrive after stop_reason in the stream
+      if (chunk.promptTokens !== undefined) {
+        b.usage.promptTokens += chunk.promptTokens;
+      }
+      if (chunk.completionTokens !== undefined) {
+        b.usage.completionTokens += chunk.completionTokens;
+      }
+      if (chunk.totalTokens !== undefined) {
+        b.usage.totalTokens += chunk.totalTokens;
+      }
+      if (chunk.stepCount !== undefined) {
+        b.usage.stepCount += chunk.stepCount;
+      }
+      break;
+    }
+
     default:
-      break; // ignore ping/usage/etc
+      break; // ignore ping/etc
   }
 }
 

src/cli/helpers/stream.ts

@@ -16,6 +16,7 @@ type DrainResult = {
   lastRunId?: string | null;
   lastSeqId?: number | null;
   approval?: ApprovalRequest | null; // present only if we ended due to approval
+  apiDurationMs: number; // time spent in API call
 };
 
 export async function drainStream(
@@ -23,6 +24,8 @@ export async function drainStream(
   buffers: ReturnType<typeof createBuffers>,
   refresh: () => void,
 ): Promise<DrainResult> {
+  const startTime = performance.now();
+
   let approvalRequestId: string | null = null;
   let toolCallId: string | null = null;
   let toolName: string | null = null;
@@ -78,10 +81,15 @@ export async function drainStream(
 
     if (chunk.messageType === "stop_reason") {
       stopReason = chunk.stopReason;
-      break; // end of turn
+      // Continue reading stream to get usage_statistics that may come after
     }
   }
 
+  // Stream has ended, check if we captured a stop reason
+  if (!stopReason) {
+    stopReason = Letta.StopReasonType.Error;
+  }
+
   // Mark the final line as finished now that stream has ended
   markCurrentLineAsFinished(buffers);
   queueMicrotask(refresh);
@@ -96,9 +104,7 @@ export async function drainStream(
         }
       : null;
 
-  if (!stopReason) {
-    stopReason = Letta.StopReasonType.Error;
-  }
+  const apiDurationMs = performance.now() - startTime;
 
-  return { stopReason, approval, lastRunId, lastSeqId };
+  return { stopReason, approval, lastRunId, lastSeqId, apiDurationMs };
 }