From db517cff0e16e6d3b371c6b506c9fe61bbbb44ab Mon Sep 17 00:00:00 2001
From: Ari Webb <ari@letta.com>
Date: Tue, 23 Dec 2025 11:35:05 -0800
Subject: [PATCH] fix: cache context window from models.list, use for byok
 (#364)

---
 src/agent/available-models.ts | 18 +++++++++++++++++-
 src/cli/App.tsx               |  9 +++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/agent/available-models.ts b/src/agent/available-models.ts
index 653107b..52fdf51 100644
--- a/src/agent/available-models.ts
+++ b/src/agent/available-models.ts
@@ -4,6 +4,7 @@ const CACHE_TTL_MS = 5 * 60 * 1000; // 5 minutes
 
 type CacheEntry = {
   handles: Set<string>;
+  contextWindows: Map<string, number>; // handle -> max_context_window
   fetchedAt: number;
 };
 
@@ -47,7 +48,14 @@ async function fetchFromNetwork(): Promise<CacheEntry> {
   const handles = new Set(
     modelsList.map((m) => m.handle).filter((h): h is string => !!h),
   );
-  return { handles, fetchedAt: Date.now() };
+  // Build context window map from API response
+  const contextWindows = new Map<string, number>();
+  for (const model of modelsList) {
+    if (model.handle && model.max_context_window) {
+      contextWindows.set(model.handle, model.max_context_window);
+    }
+  }
+  return { handles, contextWindows, fetchedAt: Date.now() };
 }
 
 export async function getAvailableModelHandles(options?: {
@@ -99,3 +107,11 @@ export function prefetchAvailableModelHandles(): void {
     // Ignore failures; UI will handle errors on-demand.
   });
 }
+
+/**
+ * Get the max_context_window for a model handle from the cached API response.
+ * Returns undefined if not cached or handle not found.
+ */
+export function getModelContextWindow(handle: string): number | undefined {
+  return cache?.contextWindows.get(handle);
+}
diff --git a/src/cli/App.tsx b/src/cli/App.tsx
index 2b22dad..244373d 100644
--- a/src/cli/App.tsx
+++ b/src/cli/App.tsx
@@ -3972,11 +3972,20 @@ DO NOT respond to these messages or otherwise consider them in your response unl
           // If not found in static list, it might be a BYOK model where id === handle
           if (!selectedModel && modelId.includes("/")) {
             // Treat it as a BYOK model - the modelId is actually the handle
+            // Look up the context window from the API-cached model info
+            const { getModelContextWindow } = await import(
+              "../agent/available-models"
+            );
+            const apiContextWindow = getModelContextWindow(modelId);
+
             selectedModel = {
               id: modelId,
               handle: modelId,
               label: modelId.split("/").pop() ?? modelId,
               description: "Custom model",
+              updateArgs: apiContextWindow
+                ? { context_window: apiContextWindow }
+                : undefined,
             } as unknown as (typeof models)[number];
           }