Add /context-size command to set context window for local providers

Author: Gustav-Simonsson

src/provider/context-window.ts

@@ -23,10 +23,26 @@ export function contextTokensFromUsage(usage: TokenUsage | undefined): number {
 // Exact model-id match wins over the family heuristics below.
 let contextWindowRegistry: Record<string, number> = {};
 
+// Provider-level overrides for local models set via /context-size.
+// Takes precedence over the model registry and heuristics.
+let providerContextOverrides: Record<string, number | undefined> = {};
+
 export function setModelContextWindows(windows: Record<string, number> | undefined): void {
   contextWindowRegistry = windows ?? {};
 }
 
+export function setProviderContextOverrides(overrides: Record<string, number> | undefined): void {
+  providerContextOverrides = overrides ?? {};
+}
+
+export function setProviderContextWindow(provider: string, tokens?: number): void {
+  if (tokens === undefined) {
+    providerContextOverrides[provider] = undefined;
+  } else {
+    providerContextOverrides[provider] = tokens;
+  }
+}
+
 function heuristicWindow(model: string): number {
   const m = model.toLowerCase();
   if (m.includes("gpt-5") || m.includes("codex")) return 400_000;
@@ -45,6 +61,12 @@ function heuristicWindow(model: string): number {
 // full identity as given, the bare model id, and `canonicalProvider/model` —
 // so a custom-named provider still exact-matches the registry instead of
 // silently missing and falling through to the heuristic.
+function providerFromModel(model: string): string | undefined {
+  const colonIndex = model.indexOf(":");
+  if (colonIndex === -1) return undefined;
+  return model.slice(0, colonIndex);
+}
+
 function lookupCandidates(model: string): string[] {
   const colonIndex = model.indexOf(":");
   if (colonIndex === -1) return [model];
@@ -59,12 +81,20 @@ function lookupCandidates(model: string): string[] {
 /** True when the registry has an entry for `model` under any known form, so a
  * caller can distinguish a confident lookup from the heuristic fallback. */
 export function hasContextWindowFor(model: string): boolean {
+  const provider = providerFromModel(model);
+  if (provider !== undefined && providerContextOverrides[provider] !== undefined) {
+    return true;
+  }
   return lookupCandidates(model).some(
     (candidate) => contextWindowRegistry[candidate] !== undefined,
   );
 }
 
 export function contextWindowFor(model: string): number {
+  const provider = providerFromModel(model);
+  if (provider !== undefined && providerContextOverrides[provider] !== undefined) {
+    return providerContextOverrides[provider];
+  }
   for (const candidate of lookupCandidates(model)) {
     const exact = contextWindowRegistry[candidate];
     if (exact !== undefined) return exact;

src/tui/commands/built-in.ts

@@ -226,4 +226,21 @@ export function registerBuiltInCommands(): void {
       };
     },
   });
+
+  registerCommand({
+    name: "context-size",
+    description: "Set context window size for the current local model (tokens)",
+    argumentHint: "<tokens>",
+    handler: (args, ctx) => {
+      const tokens = Number(args.trim());
+      if (!Number.isFinite(tokens) || tokens <= 0) {
+        return { type: "message", text: "Usage: /context-size <positive-integer>" };
+      }
+      if (ctx.setContextWindow === undefined) {
+        return { type: "message", text: "Context size is not available in this mode." };
+      }
+      const msg = ctx.setContextWindow(tokens);
+      return { type: "message", text: msg };
+    },
+  });
 }

src/tui/commands/registry.ts

@@ -27,6 +27,12 @@ export interface CommandContext {
   getSkipPermissions?: () => boolean;
   /** Live-flip skip-permissions and persist `/yolo` as the user-global default. */
   setSkipPermissions?: (value: boolean) => void;
+  /**
+   * Set the context window for the current provider (local models only).
+   * Returns a user-facing status message. The implementation persists the
+   * change to the global settings file.
+   */
+  setContextWindow?: (tokens: number) => string;
 }
 
 export type CommandResult =

src/tui/model-catalog.ts

@@ -235,10 +235,11 @@ function pricingImpact(pricing: PricingCache | null, model: string): string {
   return `${formatPrice(price.inputPricePerToken)} / ${formatPrice(price.outputPricePerToken)} per Mtok${ratioText}.`;
 }
 
-function whatLine(model: string): string {
-  const reasoning = modelReasoningCapability(model);
-  const context = contextWindowFor(model);
-  const confident = hasContextWindowFor(model);
+function whatLine(modelId: string): string {
+  const bare = modelId.includes(":") ? modelId.slice(modelId.indexOf(":") + 1) : modelId;
+  const reasoning = modelReasoningCapability(bare);
+  const context = contextWindowFor(modelId);
+  const confident = hasContextWindowFor(modelId);
   const contextText =
     context > 0
       ? `${Math.round(context / 1000)}k context${confident ? "" : " (estimated)"}`
@@ -268,15 +269,15 @@ export function describeModelCatalogOption(
 
   if (option.warning !== undefined) {
     return {
-      what: whatLine(model),
+      what: whatLine(option.id),
       impact:
         "A Go model reached over the Zen path. Billed as Zen credits, not your Go subscription.",
       tone: "consequence",
     };
   }
 
   return {
-    what: whatLine(model),
+    what: whatLine(option.id),
     impact: pricingImpact(pricing, model),
     tone: "plain",
   };

src/tui/runner.ts

@@ -133,7 +133,11 @@ import {
   maskContextMeterWhenNoTurns,
   type CostSummary,
 } from "../cost/cost-summary.js";
-import { contextTokensFromUsage } from "../provider/context-window.js";
+import {
+  contextTokensFromUsage,
+  setProviderContextWindow,
+  setProviderContextOverrides,
+} from "../provider/context-window.js";
 import {
   advertisedToolNamesForSessionMode,
   advertisedTools,
@@ -587,6 +591,17 @@ export function setUpCommandRegistry(
 
 export async function runTUI(initialConfig: Config): Promise<number> {
   let config = initialConfig;
+  // Seed provider context-window overrides from persisted settings so the
+  // model picker shows the user-set value after restart.
+  if (config.settings?.providers) {
+    const overrides: Record<string, number> = {};
+    for (const [providerName, p] of Object.entries(config.settings.providers)) {
+      if (typeof p.contextWindow === "number" && p.contextWindow > 0) {
+        overrides[providerName] = p.contextWindow;
+      }
+    }
+    setProviderContextOverrides(overrides);
+  }
   const inferenceDeps = await createInferenceDependencies();
 
   // Auto-discover plugins from the repo's plugins/ directory and user plugin
@@ -2113,6 +2128,23 @@ export async function runTUI(initialConfig: Config): Promise<number> {
       beginFeedbackCapture: () => {
         armFeedbackCapture();
       },
+      setContextWindow: (tokens: number) => {
+        const providerName = config.providerName;
+        const providerSettings = config.settings?.providers?.[providerName];
+        const isLocal = config.keyless || providerSettings?.keyless === true;
+        if (!isLocal) {
+          return "Context size can only be set for local models.";
+        }
+        void persistGlobalSettings("contextWindow", (base) => {
+          const prov = base.providers[providerName];
+          if (prov) {
+            prov.contextWindow = tokens;
+          }
+          return base;
+        });
+        setProviderContextWindow(providerName, tokens);
+        return `Context window set to ${tokens} tokens for ${providerName}.`;
+      },
     };
 
     // Routed through the shell's notice path rather than straight into the