Add context window management: implement context size configuration for models and apply token limit trimming in message history

habedi · habedi · commit 4dbc2581de45 · 2025-08-19T09:22:56.000+02:00
diff --git a/package.json b/package.json
@@ -31,20 +31,22 @@
     "@types/marked": "^5.0.2",
     "ai": "^5.0.15",
     "dotenv": "^16.4.5",
+    "gpt-tokenizer": "^3.0.1",
     "html-to-text": "^9.0.5",
     "ink": "^6.2.0",
     "ink-spinner": "^5.0.0",
     "ink-text-input": "^6.0.0",
     "json5": "^2.2.3",
     "marked": "^16.2.0",
     "ollama-ai-provider-v2": "^1.1.1",
+    "os-locale": "^6.0.2",
     "react": "^19.1.1",
     "simple-git": "^3.28.0",
+    "typescript": "^5.5.3",
     "winston": "^3.17.0",
     "zod": "^4.0.17",
-    "zustand": "^5.0.7",
-    "typescript": "^5.5.3",
-    "zod-to-ts": "^1.2.0"
+    "zod-to-ts": "^1.2.0",
+    "zustand": "^5.0.7"
   },
   "devDependencies": {
     "@types/node": "^20.14.9",
diff --git a/src/agent/context-window.ts b/src/agent/context-window.ts
@@ -0,0 +1,70 @@
+import { encode } from "gpt-tokenizer";
+import type { ModelMessage } from "ai";
+import type { ModelConfig } from "@/config.js";
+import logger from "@/logger.js";
+
+function getTokenCount(text: string): number {
+    return encode(text).length;
+}
+
+export function applyContextWindow(
+    history: ModelMessage[],
+    modelConfig: ModelConfig,
+): ModelMessage[] {
+    const { context: contextLimit } = modelConfig;
+    const safeContextLimit = contextLimit * 0.8;
+
+    let totalTokens = 0;
+    for (const message of history) {
+        if (typeof message.content === "string") {
+            totalTokens += getTokenCount(message.content);
+        } else {
+            // Handle content arrays
+            for (const part of message.content) {
+                if ("text" in part) {
+                    totalTokens += getTokenCount(part.text);
+                }
+            }
+        }
+    }
+
+    if (totalTokens <= safeContextLimit) {
+        logger.info(
+            `Token count (${totalTokens}) is within the safe limit of ${safeContextLimit}. No trimming needed.`,
+        );
+        return history;
+    }
+
+    logger.warn(
+        `Token count (${totalTokens}) exceeds the safe limit of ${safeContextLimit}. Trimming history...`,
+    );
+
+    const trimmedHistory = [...history];
+
+    // Always preserve the system prompt (if it's the first message)
+    const hasSystemPrompt = trimmedHistory[0]?.role === "system";
+    const startIndex = hasSystemPrompt ? 1 : 0;
+
+    while (totalTokens > safeContextLimit && trimmedHistory.length > startIndex + 1) {
+        const removedMessage = trimmedHistory.splice(startIndex, 1)[0];
+        if (removedMessage) {
+            let removedTokens = 0;
+            if (typeof removedMessage.content === "string") {
+                removedTokens = getTokenCount(removedMessage.content);
+            } else {
+                for (const part of removedMessage.content) {
+                    if ("text" in part) {
+                        removedTokens += getTokenCount(part.text);
+                    }
+                }
+            }
+            totalTokens -= removedTokens;
+            logger.info(
+                `Removed message at index ${startIndex} to save ${removedTokens} tokens. New total: ${totalTokens}`,
+            );
+        }
+    }
+
+    logger.info(`History trimmed. Final token count: ${totalTokens}.`);
+    return trimmedHistory;
+}
diff --git a/src/agent/state.ts b/src/agent/state.ts
@@ -13,6 +13,7 @@ import path from "path";
 import simpleGit from "simple-git";
 import { HistoryItem, ToolRequestItem } from "./history.js";
 import type { ModelMessage } from "ai";
+import { applyContextWindow } from "./context-window.js";
 
 function loadCommandHistory(): string[] {
     try {
@@ -269,7 +270,7 @@ export const useStore = create<AppState & AppActions>((set, get) => ({
                 const { history, config } = get();
                 if (!config) throw new FatalError("Configuration not loaded.");
 
-                const sdkCompliantHistory = history
+                let sdkCompliantHistory = history
                     .map((item): ModelMessage | null => {
                         switch (item.role) {
                             case "user":
@@ -306,6 +307,15 @@ export const useStore = create<AppState & AppActions>((set, get) => ({
                     })
                     .filter(Boolean) as ModelMessage[];
 
+                const modelConfig = config.models.find((m) => m.name === config.defaultModel);
+                if (!modelConfig) {
+                    throw new FatalError(
+                        `Model ${config.defaultModel} not found in configuration.`,
+                    );
+                }
+
+                sdkCompliantHistory = applyContextWindow(sdkCompliantHistory, modelConfig);
+
                 const systemPrompt = await generateSystemPrompt(config);
 
                 const { textStream, toolCalls: toolCallPartsPromise } =
diff --git a/src/agent/system-prompt.ts b/src/agent/system-prompt.ts
@@ -5,6 +5,39 @@ import { zodToTs, printNode } from "zod-to-ts";
 import fs from "fs/promises";
 import path from "path";
 import os from "os";
+import { osLocale } from "os-locale"; // NEW: Import the library for locale detection.
+
+/**
+ * Dynamically gets the user's system locale (language and region).
+ * Handles cross-platform differences and provides a safe fallback.
+ * @returns A promise that resolves to the user's locale string (e.g., "en-US").
+ */
+async function getUserLocale(): Promise<string> {
+    try {
+        // `os-locale` is the standard way to solve this problem in Node.js.
+        // It correctly checks LANG, LC_ALL, etc., on Linux/macOS and uses OS APIs on Windows.
+        const locale = await osLocale();
+        // The library might return 'en_US'. We convert it to the IETF BCP 47 standard 'en-US'.
+        return locale.replace("_", "-");
+    } catch (e) {
+        // If detection fails for any reason, fall back to a sensible default.
+        return "en-US";
+    }
+}
+
+/**
+ * Dynamically gets the user's system timezone.
+ * @returns The user's timezone string (e.g., "Europe/Oslo").
+ */
+function getUserTimezone(): string {
+    try {
+        // The Intl API is the standard, modern way to get the system timezone in JavaScript.
+        return Intl.DateTimeFormat().resolvedOptions().timeZone;
+    } catch (e) {
+        // Fallback in case the environment is unusual and the API fails.
+        return "Europe/Oslo";
+    }
+}
 
 /**
  * Recursively searches for instruction files (TOBI.md, AGENTS.md) upwards from the current directory.
@@ -14,21 +47,17 @@ import os from "os";
 async function findInstructionFile(currentDir: string): Promise<string | null> {
     const instructionFiles = ["TOBI.md", "AGENTS.md"];
     const homeDir = os.homedir();
-
     let dir = currentDir;
-    // Stop if we reach the root directory or the user's home directory
     while (dir !== path.dirname(dir) && dir.startsWith(homeDir)) {
         for (const fileName of instructionFiles) {
             const filePath = path.join(dir, fileName);
             try {
-                // Check if the file exists and we can read it
                 await fs.access(filePath, fs.constants.R_OK);
                 return await fs.readFile(filePath, "utf-8");
             } catch {
-                // File does not exist or is not readable, continue searching
+                // Continue searching
             }
         }
-        // Move to the parent directory
         dir = path.dirname(dir);
     }
     return null;
@@ -39,24 +68,28 @@ async function findInstructionFile(currentDir: string): Promise<string | null> {
  * @param config The application configuration.
  * @returns A promise that resolves to the generated system prompt string.
  */
-export async function generateSystemPrompt(_config: Config): Promise<string> {
+export async function generateSystemPrompt(config: Config): Promise<string> {
     const cwd = process.cwd();
 
-    // 1. Get workspace files
+    // 1. Gather environmental context, now including locale and timezone
     const dirents = await fs.readdir(cwd, { withFileTypes: true });
     const filesAndDirs = dirents.map((d) => (d.isDirectory() ? `${d.name}/` : d.name)).join("\n");
+    const osPlatform = os.platform();
+
+    // NEW: Call the dynamic helper functions.
+    const userLocale = await getUserLocale();
+    const userTimezone = getUserTimezone();
+    // NEW: Use the detected locale and timezone to format the date correctly for the user.
+    const currentDate = new Date().toLocaleString(userLocale, { timeZone: userTimezone });
 
-    // 2. Find instruction file
+    // 2. Find project-specific instructions
     const instructionContent = await findInstructionFile(cwd);
 
     // 3. Generate tool definitions as TypeScript types
     const toolDefinitions = Object.values(toolModules)
         .map((module) => {
-            // We pass the zod schema for the arguments to zodToTs
             const { node } = zodToTs(module.schema.shape.arguments);
-            // Then we print the resulting TypeScript AST node to a string.
             const argumentsString = printNode(node);
-
             return (
                 `// ${module.description}\n` +
                 `type ${module.schema.shape.name.value} = {\n` +
@@ -67,37 +100,63 @@ export async function generateSystemPrompt(_config: Config): Promise<string> {
         })
         .join("\n\n");
 
-    // 4. Assemble the prompt
+    // 4. Assemble the prompt, now with the new context
     const promptParts: string[] = [
-        `You are a helpful AI assistant named Tobi. You can use tools to help the user with coding and file system tasks.`,
-    ];
+        `You are Tobi, an autonomous AI software engineer. Your role is to assist the user, named "${config.defaultModel}", by executing tasks with the tools provided. You operate with maximum efficiency and precision.`,
 
-    promptParts.push("### Current Workspace");
-    promptParts.push("Here is a list of files and directories in the current working directory:");
-    promptParts.push("```");
-    promptParts.push(filesAndDirs);
-    promptParts.push("```");
+        // UPDATED: The Environment section now includes the new dynamic information.
+        `### Environment\n` +
+            `* **Operating System:** ${osPlatform}\n` +
+            `* **User Locale:** ${userLocale}\n` +
+            `* **User Timezone:** ${userTimezone}\n` +
+            `* **Current Date & Time:** ${currentDate}\n` +
+            `* **Working Directory:** ${cwd}\n` +
+            "* **Directory Contents:**\n" +
+            "```\n" +
+            `${filesAndDirs || "(empty)"}\n` +
+            "```",
+
+        "### Rules of Engagement\n" +
+            "1.  **Think Step-by-Step:** Before acting, briefly state your plan to achieve the user's goal.\n" +
+            "2.  **Execute Autonomously:** You are autonomous. Use your tools to execute your plan without asking for permission. The user will intervene if your plan is incorrect.\n" +
+            "3.  **One Tool at a Time:** You can only call one tool per turn. Decompose complex tasks into a sequence of single tool calls.\n" +
+            "4.  **Stay on Task:** Your responses should consist of your thought process followed by a tool call. Avoid conversational filler or apologies.\n" +
+            "5.  **Code Concisely:** Do not add comments to code unless explicitly requested by the user.",
+    ];
 
     if (instructionContent) {
-        promptParts.push("### User-Provided Instructions");
         promptParts.push(
-            "The user has provided the following instructions in a `TOBI.md` or `AGENTS.md` file. Follow them carefully.",
+            "### User-Provided Instructions\n" +
+                "The user has provided the following project-specific instructions. Adhere to them strictly.\n" +
+                "```markdown\n" +
+                instructionContent +
+                "\n```",
         );
-        promptParts.push("```markdown");
-        promptParts.push(instructionContent);
-        promptParts.push("```");
     }
 
-    promptParts.push("### Available Tools");
-    promptParts.push(
-        "You have the following tools available. To use a tool, respond with a JSON object that strictly adheres to the TypeScript type definition of the tool.",
-    );
     promptParts.push(
-        "The following are TypeScript type definitions for the tools. The `name` property is the tool to call, and you must provide the corresponding `arguments` object.",
+        "### Tool Reference\n" +
+            "To use a tool, you must respond with a single JSON object containing the `tool_calls` property. This object must conform to the following TypeScript definitions.\n\n" +
+            "**Example:** To list files, you would respond with:\n" +
+            "```json\n" +
+            JSON.stringify(
+                {
+                    tool_calls: [
+                        {
+                            name: "list",
+                            arguments: { path: "." },
+                        },
+                    ],
+                },
+                null,
+                2,
+            ) +
+            "\n```\n\n" +
+            "**Tool Definitions:**\n" +
+            "```typescript\n" +
+            toolDefinitions +
+            "\n```",
     );
-    promptParts.push("```typescript");
-    promptParts.push(toolDefinitions);
-    promptParts.push("```");
 
     return promptParts.join("\n\n");
 }
diff --git a/src/config.ts b/src/config.ts
@@ -12,6 +12,7 @@ const modelSchema = z.object({
         .enum(["openai", "google", "anthropic", "ollama"])
         .describe("The provider of the model."),
     modelId: z.string().describe("The actual model ID used by the API."),
+    context: z.number().describe("The context window size for the model."),
     baseUrl: z.url().optional().describe("Optional base URL for providers like Ollama."),
 });
 export type ModelConfig = z.infer<typeof modelSchema>;
@@ -57,17 +58,38 @@ const defaultConfig: Config = {
     systemPrompt: `You are a helpful AI assistant named Tobi. You can use tools to help the user with coding and file system tasks.`,
     defaultModel: "gpt-4.1-mini",
     models: [
-        { name: "gpt-4.1-mini", provider: "openai", modelId: "gpt-4.1-mini" },
-        { name: "gpt-4.1", provider: "openai", modelId: "gpt-4.1" },
-        { name: "gpt-4o", provider: "openai", modelId: "gpt-4o" },
-        { name: "claude-3.5-sonnet", provider: "anthropic", modelId: "claude-3-5-sonnet" },
-        { name: "claude-3.5-haiku", provider: "anthropic", modelId: "claude-3-5-haiku" },
-        { name: "gemini-2.5-pro", provider: "google", modelId: "models/gemini-2.5-pro" },
-        { name: "gemini-2.5-flash", provider: "google", modelId: "models/gemini-2.5-flash" },
+        { name: "gpt-4.1-mini", provider: "openai", modelId: "gpt-4.1-mini", context: 128000 },
+        { name: "gpt-4.1", provider: "openai", modelId: "gpt-4.1", context: 128000 },
+        { name: "gpt-4o", provider: "openai", modelId: "gpt-4o", context: 128000 },
+        {
+            name: "claude-3.5-sonnet",
+            provider: "anthropic",
+            modelId: "claude-3-5-sonnet",
+            context: 200000,
+        },
+        {
+            name: "claude-3.5-haiku",
+            provider: "anthropic",
+            modelId: "claude-3-5-haiku",
+            context: 200000,
+        },
+        {
+            name: "gemini-2.5-pro",
+            provider: "google",
+            modelId: "models/gemini-2.5-pro",
+            context: 1000000,
+        },
+        {
+            name: "gemini-2.5-flash",
+            provider: "google",
+            modelId: "models/gemini-2.5-flash",
+            context: 1000000,
+        },
         {
             name: "qwen3",
             provider: "ollama",
             modelId: "qwen3:8b",
+            context: 32768,
             baseUrl: "http://localhost:11434/v1",
         },
     ],
