improved token-caching

Author: bhouston

packages/agent/src/core/tokens.ts

@@ -95,6 +95,6 @@ export class TokenTracker {
   }
 
   toString() {
-    return `${this.name}: ${this.tokenUsage.toString()}`;
+    return `${this.name}: ${this.getTotalUsage().toString()}`;
   }
 }

packages/agent/src/core/toolAgent.ts

@@ -199,7 +199,15 @@ function addCacheControlToContentBlocks(
 ): ContentBlockParam[] {
   return content.map((c, i) => {
     if (i === content.length - 1) {
-      if (c.type === 'text' || c.type === 'document' || c.type === 'image') {
+      if (
+        c.type === 'text' ||
+        c.type === 'document' ||
+        c.type === 'image' ||
+        c.type === 'tool_use' ||
+        c.type === 'tool_result' ||
+        c.type === 'thinking' ||
+        c.type === 'redacted_thinking'
+      ) {
         return { ...c, cache_control: { type: 'ephemeral' } };
       }
     }
@@ -209,7 +217,7 @@ function addCacheControlToContentBlocks(
 function addCacheControlToMessages(
   messages: Anthropic.Messages.MessageParam[],
 ): Anthropic.Messages.MessageParam[] {
-  return messages.map((m) => {
+  return messages.map((m, i) => {
     if (typeof m.content === 'string') {
       return {
         ...m,
@@ -224,7 +232,10 @@ function addCacheControlToMessages(
     }
     return {
       ...m,
-      content: addCacheControlToContentBlocks(m.content),
+      content:
+        i >= messages.length - 2
+          ? addCacheControlToContentBlocks(m.content)
+          : m.content,
     };
   });
 }

packages/agent/src/tools/system/shellStart.ts

@@ -189,7 +189,7 @@ export const shellStartTool: Tool<Parameters, ReturnType> = {
     { logger },
   ) => {
     logger.info(
-      `🖥️ Starting "${command}", ${description} (timeout: ${timeout}ms)`,
+      `Starting "${command}", ${description} (timeout: ${timeout}ms)`,
     );
   },
   logReturns: (output, { logger }) => {

packages/cli/src/commands/$default.ts

@@ -1,6 +1,7 @@
 import * as fs from 'fs/promises';
 import { createInterface } from 'readline/promises';
 
+import chalk from 'chalk';
 import {
   toolAgent,
   Logger,
@@ -107,6 +108,15 @@ export const command: CommandModule<object, DefaultArgs> = {
 
       const tools = getTools();
 
+      // Error handling
+      process.on('SIGINT', () => {
+        logger.log(
+          tokenTracker.logLevel,
+          chalk.blueBright(`[Token Usage Total] ${tokenTracker.toString()}`),
+        );
+        process.exit(0);
+      });
+
       const result = await toolAgent(prompt, tools, undefined, {
         logger,
         headless: argv.headless ?? true,
@@ -123,6 +133,9 @@ export const command: CommandModule<object, DefaultArgs> = {
       logger.error('An error occurred:', error);
     }
 
-    logger.info(`Token Usage: ${tokenTracker.toString()}`);
+    logger.log(
+      tokenTracker.logLevel,
+      chalk.blueBright(`[Token Usage Total] ${tokenTracker.toString()}`),
+    );
   },
 };

packages/cli/src/index.ts

@@ -35,12 +35,6 @@ const main = async () => {
     console.log();
   }
 
-  // Error handling
-  process.on('SIGINT', () => {
-    logger.warn('\nGracefully shutting down...');
-    process.exit(0);
-  });
-
   process.on('uncaughtException', (error) => {
     logger.error(
       'Fatal error:',