feat: optimize ui

src/index.ts

@@ -15,6 +15,7 @@ import { CONFIG_FILE } from "./constants";
 import createWriteStream from "pino-rotating-file-stream";
 import { HOME_DIR } from "./constants";
 import { configureLogging } from "./utils/log";
+import { sessionUsageCache } from "./utils/cache";
 
 async function initializeClaudeConfig() {
   const homeDir = homedir();
@@ -52,10 +53,10 @@ async function run(options: RunOptions = {}) {
   // Clean up old log files, keeping only the 10 most recent ones
   await cleanupLogFiles();
   const config = await initConfig();
-  
+
   // Configure logging based on config
   configureLogging(config);
-  
+
   let HOST = config.HOST;
 
   if (config.HOST && !config.APIKEY) {
@@ -88,15 +89,18 @@ async function run(options: RunOptions = {}) {
     : port;
 
   // Configure logger based on config settings
-  const loggerConfig = config.LOG !== false ? {
-    level: config.LOG_LEVEL || "debug",
-    stream: createWriteStream({
-      path: HOME_DIR,
-      filename: config.LOGNAME || `./logs/ccr-${+new Date()}.log`,
-      maxFiles: 3,
-      interval: "1d",
-    }),
-  } : false;
+  const loggerConfig =
+    config.LOG !== false
+      ? {
+          level: config.LOG_LEVEL || "debug",
+          stream: createWriteStream({
+            path: HOME_DIR,
+            filename: config.LOGNAME || `./logs/ccr-${+new Date()}.log`,
+            maxFiles: 3,
+            interval: "1d",
+          }),
+        }
+      : false;
 
   const server = createServer({
     jsonPath: CONFIG_FILE,
@@ -129,6 +133,33 @@ async function run(options: RunOptions = {}) {
       router(req, reply, config);
     }
   });
+  server.addHook("onSend", async (req, reply, payload) => {
+    if (req.sessionId && req.url.startsWith("/v1/messages")) {
+      if (payload instanceof ReadableStream) {
+        const [originalStream, clonedStream] = payload.tee();
+        const reader1 = clonedStream.getReader();
+        while (true) {
+          const { done, value } = await reader1.read();
+          if (done) break;
+          // Process the value if needed
+          const dataStr = new TextDecoder().decode(value);
+          if (!dataStr.startsWith("event: message_delta")) {
+            continue;
+          }
+          const str = dataStr.slice(27);
+          try {
+            const message = JSON.parse(str);
+            sessionUsageCache.put(req.sessionId, message.usage);
+          } catch {}
+        }
+
+        return originalStream;
+      } else {
+        sessionUsageCache.put(req.sessionId, payload.usage);
+      }
+    }
+    return payload;
+  });
   server.start();
 }
 

src/utils/cache.ts

@@ -0,0 +1,47 @@
+// LRU cache for session usage
+
+export interface Usage {
+  input_tokens: number;
+  output_tokens: number;
+}
+
+class LRUCache<K, V> {
+  private capacity: number;
+  private cache: Map<K, V>;
+
+  constructor(capacity: number) {
+    this.capacity = capacity;
+    this.cache = new Map<K, V>();
+  }
+
+  get(key: K): V | undefined {
+    if (!this.cache.has(key)) {
+      return undefined;
+    }
+    const value = this.cache.get(key) as V;
+    // Move to end to mark as recently used
+    this.cache.delete(key);
+    this.cache.set(key, value);
+    return value;
+  }
+
+  put(key: K, value: V): void {
+    if (this.cache.has(key)) {
+      // If key exists, delete it to update its position
+      this.cache.delete(key);
+    } else if (this.cache.size >= this.capacity) {
+      // If cache is full, delete the least recently used item
+      const leastRecentlyUsedKey = this.cache.keys().next().value;
+      if (leastRecentlyUsedKey !== undefined) {
+        this.cache.delete(leastRecentlyUsedKey);
+      }
+    }
+    this.cache.set(key, value);
+  }
+
+  values(): V[] {
+    return Array.from(this.cache.values());
+  }
+}
+
+export const sessionUsageCache = new LRUCache<string, Usage>(100);

src/utils/codeCommand.ts

@@ -11,12 +11,14 @@ export async function executeCodeCommand(args: string[] = []) {
   const config = await readConfigFile();
   const env: Record<string, string> = {
     ...process.env,
-    ANTHROPIC_AUTH_TOKEN: "test",
+    ANTHROPIC_AUTH_TOKEN: config?.APIKEY || "test",
     ANTHROPIC_BASE_URL: `http://127.0.0.1:${config.PORT || 3456}`,
     API_TIMEOUT_MS: String(config.API_TIMEOUT_MS ?? 600000), // Default to 10 minutes if not set
   };
 
-  const settingsFlag: Record<string, any> = {};
+  const settingsFlag: Record<string, any> = {
+    env,
+  };
   if (config?.StatusLine?.enabled) {
     settingsFlag.statusLine = {
       type: "command",
@@ -41,10 +43,10 @@ export async function executeCodeCommand(args: string[] = []) {
     env.ANTHROPIC_SMALL_FAST_MODEL = config.ANTHROPIC_SMALL_FAST_MODEL;
   }
 
-  if (config?.APIKEY) {
-    env.ANTHROPIC_API_KEY = config.APIKEY;
-    delete env.ANTHROPIC_AUTH_TOKEN;
-  }
+  // if (config?.APIKEY) {
+  //   env.ANTHROPIC_API_KEY = config.APIKEY;
+  //   delete env.ANTHROPIC_AUTH_TOKEN;
+  // }
 
   // Increment reference count when command starts
   incrementReferenceCount();

src/utils/log.ts

@@ -16,7 +16,7 @@ let logLevel: string = "info";
 // Function to configure logging
 export function configureLogging(config: { LOG?: boolean; LOG_LEVEL?: string }) {
   isLogEnabled = config.LOG !== false; // Default to true if not explicitly set to false
-  logLevel = config.LOG_LEVEL || "info";
+  logLevel = config.LOG_LEVEL || "debug";
 }
 
 export function log(...args: any[]) {

src/utils/router.ts

@@ -5,6 +5,7 @@ import {
 } from "@anthropic-ai/sdk/resources/messages";
 import { get_encoding } from "tiktoken";
 import { log } from "./log";
+import { sessionUsageCache, Usage } from "./cache";
 
 const enc = get_encoding("cl100k_base");
 
@@ -62,7 +63,12 @@ const calculateTokenCount = (
   return tokenCount;
 };
 
-const getUseModel = async (req: any, tokenCount: number, config: any) => {
+const getUseModel = async (
+  req: any,
+  tokenCount: number,
+  config: any,
+  lastUsage?: Usage | undefined
+) => {
   if (req.body.model.includes(",")) {
     const [provider, model] = req.body.model.split(",");
     const finalProvider = config.Providers.find(
@@ -78,7 +84,15 @@ const getUseModel = async (req: any, tokenCount: number, config: any) => {
   }
   // if tokenCount is greater than the configured threshold, use the long context model
   const longContextThreshold = config.Router.longContextThreshold || 60000;
-  if (tokenCount > longContextThreshold && config.Router.longContext) {
+  const lastUsageThreshold =
+    lastUsage &&
+    lastUsage.input_tokens > longContextThreshold &&
+    tokenCount > 20000;
+  const tokenCountThreshold = tokenCount > longContextThreshold;
+  if (
+    (lastUsageThreshold || tokenCountThreshold) &&
+    config.Router.longContext
+  ) {
     log(
       "Using long context model due to token count:",
       tokenCount,
@@ -125,14 +139,15 @@ const getUseModel = async (req: any, tokenCount: number, config: any) => {
   return config.Router!.default;
 };
 
-export const router = async (req: any, _res: any, config: any) => { 
+export const router = async (req: any, _res: any, config: any) => {
   // Parse sessionId from metadata.user_id
   if (req.body.metadata?.user_id) {
-    const parts = req.body.metadata.user_id.split('_session_');
+    const parts = req.body.metadata.user_id.split("_session_");
     if (parts.length > 1) {
       req.sessionId = parts[1];
     }
   }
+  const lastMessageUsage = sessionUsageCache.get(req.sessionId);
   const { messages, system = [], tools }: MessageCreateParamsBase = req.body;
   try {
     const tokenCount = calculateTokenCount(
@@ -152,7 +167,7 @@ export const router = async (req: any, _res: any, config: any) => {
       }
     }
     if (!model) {
-      model = await getUseModel(req, tokenCount, config);
+      model = await getUseModel(req, tokenCount, config, lastMessageUsage);
     }
     req.body.model = model;
   } catch (error: any) {