nibzard · caydyan · Jun 14, 2026
diff --git a/helpers/context_window_helpers.py b/helpers/context_window_helpers.py
@@ -0,0 +1,144 @@
+from dataclasses import dataclass
+import re
+
+
+CONTEXT_SECTION_PATTERN = re.compile(
+    r"<<SECTION: (?P<title>.*?) >>\n(?P<body>.*?)(?:\n)?<<END_SECTION: (?P=title) >>",
+    re.DOTALL,
+)
+
+HIGH_PRIORITY_FILES = {
+    "package.json",
+    "requirements.txt",
+    "pyproject.toml",
+    "Cargo.toml",
+    "Cargo.lock",
+    "go.mod",
+    "go.sum",
+    "pom.xml",
+    "build.gradle",
+    "Dockerfile",
+    "docker-compose.yml",
+    "compose.yml",
+    "Gemfile",
+    "Pipfile",
+    "Pipfile.lock",
+    "setup.py",
+    "Makefile",
+    "CMakeLists.txt",
+}
+
+TRUNCATION_NOTICE = "\n\n<<TRUNCATED: Context window budget reached >>"
+
+
+@dataclass
+class ContextSection:
+    title: str
+    text: str
+    order: int
+    priority: int
+
+
+def _encoding_for_context():
+    import tiktoken
+
+    return tiktoken.encoding_for_model("gpt-4o-mini")
+
+
+def count_context_tokens(text, encoding=None):
+    selected_encoding = encoding or _encoding_for_context()
+    return len(selected_encoding.encode(text))
+
+
+def _section_priority(title):
+    normalized_title = title.strip()
+    if normalized_title == "Repository Structure":
+        return 0
+    if normalized_title == "Repository Languages":
+        return 1
+    if normalized_title == "Existing devcontainer.json":
+        return 2
+    if normalized_title.startswith("Content of "):
+        file_name = normalized_title.removeprefix("Content of ").strip()
+        if file_name in HIGH_PRIORITY_FILES:
+            return 3
+        if file_name.lower() == "readme.md":
+            return 4
+    return 5
+
+
+def _parse_context_sections(context):
+    sections = []
+    for order, match in enumerate(CONTEXT_SECTION_PATTERN.finditer(context)):
+        title = match.group("title")
+        sections.append(
+            ContextSection(
+                title=title,
+                text=match.group(0),
+                order=order,
+                priority=_section_priority(title),
+            )
+        )
+    return sections
+
+
+def _truncate_text_to_budget(text, max_tokens, encoding):
+    if max_tokens <= 0:
+        return ""
+
+    tokens = encoding.encode(text)
+    if len(tokens) <= max_tokens:
+        return text
+    return encoding.decode(tokens[:max_tokens])
+
+
+def _truncate_section_to_budget(section, max_tokens, encoding):
+    notice_tokens = count_context_tokens(TRUNCATION_NOTICE, encoding)
+    if max_tokens <= notice_tokens:
+        return _truncate_text_to_budget(section.text, max_tokens, encoding)
+
+    truncated = _truncate_text_to_budget(section.text, max_tokens - notice_tokens, encoding)
+    if not truncated:
+        return ""
+    return truncated.rstrip() + TRUNCATION_NOTICE
+
+
+def optimize_context_window(context, max_tokens=120000, encoding=None):
+    selected_encoding = encoding or _encoding_for_context()
+    initial_tokens = count_context_tokens(context, selected_encoding)
+
+    if initial_tokens <= max_tokens:
+        return context
+
+    sections = _parse_context_sections(context)
+    if not sections:
+        return _truncate_text_to_budget(context, max_tokens, selected_encoding)
+
+    selected_sections = []
+    used_tokens = 0
+
+    for section in sorted(sections, key=lambda item: (item.priority, item.order)):
+        separator = "\n\n" if selected_sections else ""
+        separator_tokens = count_context_tokens(separator, selected_encoding)
+        section_tokens = count_context_tokens(section.text, selected_encoding)
+        available_tokens = max_tokens - used_tokens - separator_tokens
+
+        if available_tokens <= 0:
+            break
+
+        if section_tokens <= available_tokens:
+            selected_sections.append(section.text)
+            used_tokens += separator_tokens + section_tokens
+            continue
+
+        truncated_section = _truncate_section_to_budget(section, available_tokens, selected_encoding)
+        if truncated_section:
+            selected_sections.append(truncated_section)
+        break
+
+    optimized_context = "\n\n".join(selected_sections)
+    optimized_tokens = count_context_tokens(optimized_context, selected_encoding)
+    if optimized_tokens > max_tokens:
+        return _truncate_text_to_budget(optimized_context, max_tokens, selected_encoding)
+
+    return optimized_context
diff --git a/helpers/devcontainer_helpers.py b/helpers/devcontainer_helpers.py
@@ -4,61 +4,24 @@
 import logging
 import os
 import jsonschema
-import tiktoken
+from helpers.context_window_helpers import count_context_tokens, optimize_context_window
 from helpers.jinja_helper import process_template
 from schemas import DevContainerModel
 from supabase_client import supabase
 from models import DevContainer
 
 
-import logging
-import tiktoken
-
 def truncate_context(context, max_tokens=120000):
     logging.info(f"Starting truncate_context with max_tokens={max_tokens}")
     logging.debug(f"Initial context length: {len(context)} characters")
 
-    encoding = tiktoken.encoding_for_model("gpt-4o-mini")
-    tokens = encoding.encode(context)
-
-    logging.info(f"Initial token count: {len(tokens)}")
-
-    if len(tokens) <= max_tokens:
-        logging.info("Context is already within token limit. No truncation needed.")
-        return context
-
-    logging.info(f"Context size is {len(tokens)} tokens. Truncation needed.")
-
-    # Prioritize keeping the repository structure and languages
-    structure_end = context.find("<<END_SECTION: Repository Structure >>")
-    languages_end = context.find("<<END_SECTION: Repository Languages >>")
-
-    logging.debug(f"Structure end position: {structure_end}")
-    logging.debug(f"Languages end position: {languages_end}")
-
-    important_content = context[:languages_end] + "<<END_SECTION: Repository Languages >>\n\n"
-    remaining_content = context[languages_end + len("<<END_SECTION: Repository Languages >>\n\n"):]
-
-    important_tokens = encoding.encode(important_content)
-    logging.debug(f"Important content token count: {len(important_tokens)}")
-
-    if len(important_tokens) > max_tokens:
-        logging.warning("Important content alone exceeds max_tokens. Truncating important content.")
-        important_content = encoding.decode(important_tokens[:max_tokens])
-        return important_content
-
-    remaining_tokens = max_tokens - len(important_tokens)
-    logging.info(f"Tokens available for remaining content: {remaining_tokens}")
-
-    truncated_remaining = encoding.decode(encoding.encode(remaining_content)[:remaining_tokens])
-
-    final_context = important_content + truncated_remaining
-    final_tokens = encoding.encode(final_context)
+    optimized_context = optimize_context_window(context, max_tokens=max_tokens)
+    final_tokens = count_context_tokens(optimized_context)
 
-    logging.info(f"Final token count: {len(final_tokens)}")
-    logging.debug(f"Final context length: {len(final_context)} characters")
+    logging.info(f"Final token count: {final_tokens}")
+    logging.debug(f"Final context length: {len(optimized_context)} characters")
 
-    return final_context
+    return optimized_context
 
 def generate_devcontainer_json(instructor_client, repo_url, repo_context, devcontainer_url=None, max_retries=2, regenerate=False):
     existing_devcontainer = None
@@ -136,4 +99,4 @@ def save_devcontainer(new_devcontainer):
         return result.data[0] if result.data else None
     except Exception as e:
         logging.error(f"Error saving devcontainer to Supabase: {str(e)}")
-        raise
+        raise
diff --git a/tests/test_context_window_optimization.py b/tests/test_context_window_optimization.py
@@ -0,0 +1,72 @@
+import unittest
+
+from helpers.context_window_helpers import optimize_context_window
+
+
+class FakeEncoding:
+    def encode(self, text):
+        return [ord(character) for character in text]
+
+    def decode(self, tokens):
+        return "".join(chr(token) for token in tokens)
+
+
+ENCODING = FakeEncoding()
+
+
+def section(title, body):
+    return f"<<SECTION: {title} >>\n{body}\n<<END_SECTION: {title} >>"
+
+
+def count_tokens(text):
+    return len(ENCODING.encode(text))
+
+
+class ContextWindowOptimizationTest(unittest.TestCase):
+    def test_keeps_high_priority_sections_before_large_lower_priority_content(self):
+        structure = section("Repository Structure", "package.json\nREADME.md\nsrc/app.py")
+        languages = section("Repository Languages", "TypeScript: 120 lines\nCSS: 10 lines")
+        package_json = section("Content of package.json", '{"scripts":{"dev":"vite"}}')
+        readme = section("Content of README.md", "README_START " + ("details " * 300) + "README_TAIL")
+        changelog = section("Content of CHANGELOG.md", "CHANGELOG_START " + ("noise " * 300) + "CHANGELOG_TAIL")
+        context = "\n\n".join([readme, changelog, structure, package_json, languages])
+        budget = count_tokens("\n\n".join([structure, languages, package_json])) + 10
+
+        result = optimize_context_window(context, max_tokens=budget, encoding=ENCODING)
+
+        self.assertLessEqual(count_tokens(result), budget)
+        self.assertIn("Repository Structure", result)
+        self.assertIn("Repository Languages", result)
+        self.assertIn('"scripts":{"dev":"vite"}', result)
+        self.assertNotIn("README_TAIL", result)
+        self.assertNotIn("CHANGELOG_TAIL", result)
+
+    def test_existing_devcontainer_is_preserved_before_readme_content(self):
+        structure = section("Repository Structure", ".devcontainer/devcontainer.json\nREADME.md")
+        languages = section("Repository Languages", "Python: 20 lines")
+        existing_devcontainer = section(
+            "Existing devcontainer.json",
+            '{"image":"mcr.microsoft.com/devcontainers/python:3.12"}',
+        )
+        readme = section("Content of README.md", "README_START " + ("details " * 300) + "README_TAIL")
+        context = "\n\n".join([readme, structure, existing_devcontainer, languages])
+        budget = count_tokens("\n\n".join([structure, languages, existing_devcontainer])) + 10
+
+        result = optimize_context_window(context, max_tokens=budget, encoding=ENCODING)
+
+        self.assertLessEqual(count_tokens(result), budget)
+        self.assertIn("mcr.microsoft.com/devcontainers/python:3.12", result)
+        self.assertNotIn("README_TAIL", result)
+
+    def test_plain_context_falls_back_to_token_prefix(self):
+        context = "alpha " * 300
+        budget = 50
+
+        result = optimize_context_window(context, max_tokens=budget, encoding=ENCODING)
+
+        self.assertLessEqual(count_tokens(result), budget)
+        self.assertTrue(result.startswith("alpha"))
+
+
+if __name__ == "__main__":
+    unittest.main()