perf: warm Code Mode sandbox and cache repository context for batch refactors

.github/workflows/pytest.yml

@@ -0,0 +1,32 @@
+name: PyLattice tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Install package and tests
+        run: python -m pip install -e . "pytest>=8,<10" "pytest-asyncio>=0.25,<2" "ruff>=0.11,<1"
+      - name: Compile Python sources
+        run: python -m compileall -q src tests
+      - name: Lint Code Mode changes
+        run: |
+          python -m ruff check src/agent_tui/agent.py src/agent_tui/codemode.py \
+            src/agent_tui/context.py src/agent_tui/defaults.py src/agent_tui/loop.py \
+            src/agent_tui/policy.py src/agent_tui/prompts.py src/agent_tui/tools.py \
+            tests/test_agent.py tests/test_codemode.py tests/test_repository_context.py \
+            tests/test_tools.py benchmarks/benchmark_repository_context.py
+      - name: Run regression suite
+        run: python -m pytest -q
+      - name: Benchmark cold vs indexed repository scans
+        run: python benchmarks/benchmark_repository_context.py

README.md

@@ -49,17 +49,52 @@ Each request uses a 64,000-character context budget, task-relevant history, and
 
 Enabled by default via `--code-mode` flag or `AGENT_TUI_CODE_MODE=true`. The decision model (`--router-model` or `AGENT_TUI_ROUTER_MODEL`) chooses `execute_code` or `finish` once per turn using the latest result. The executor then writes the selected program or final answer, with a catalog of workspace, skill, memory, delegation, and MCP tools. On a routing fallback, the executor chooses between both tools. A program can make dependent tool calls without requiring another model request between them.
 
-**Example:**
+**Fast repository inspection:**
 
 ```python
-listing = await call_tool("list_files", {"path": "src"})
-results = []
-for path in listing["output"]["files"][:3]:
-    result = await call_tool("read_file", {"path": path})
-    results.append(result["output"])
-results
+scan = await call_tool("context_collect", {"query": "router", "max_files": 6})
+{"files": scan["output"]["files"], "truncated": scan["output"]["truncated"]}
 ```
 
+This performs one bounded filesystem scan and remembers the **full verified source** of
+selected files in Context Manager. Later Code Mode programs automatically receive
+relevant excerpts, instead of re-reading the same files through many tool calls.
+Use `{"paths": ["src/first.py", "src/second.py"]}` to target exact paths.
+The result includes a SHA-256 for optimistic concurrency control. A bounded, metadata-
+validated search index (8 MiB) lets repeated searches skip unchanged disk reads.
+Responses report `index_hits` and `disk_reads` for performance diagnosis. To reproduce
+a cold-versus-warm scan comparison, run
+`python benchmarks/benchmark_repository_context.py` after installing the package.
+
+**Batched refactoring and file creation:**
+
+```python
+scan = await call_tool("context_collect", {"paths": ["src/app.py"]})
+revision = scan["output"]["files"][0]["sha256"]
+change = await call_tool("apply_patchset", {
+    "changes": [
+        {
+            "action": "edit", "path": "src/app.py",
+            "old_text": "old_function()", "new_text": "new_function()",
+            "expected_sha256": revision,
+        },
+        {
+            "action": "create", "path": "src/new_module.py",
+            "content": "def new_function():\n    return 42\n",
+        },
+    ]
+})
+change["output"]  # Read/test the modified files before finishing
+```
+
+`apply_patchset` preflights every change and rejects stale context, duplicate paths,
+ambiguous edits, and existing-file creations. It stages writes before commit and uses
+one normal approval with a diff for each file. If a commit fails partway, it attempts
+a safe rollback; filesystem-level atomicity across multiple paths is not guaranteed.
+
+The Monty process pool is reused within a task, but **variables do not persist**
+between programs: source state lives in Context Manager, not in the sandbox.
+
 **Behavior:**
 
 - Send one `execute_code` tool call with a raw Python `code` string, without Markdown fences or prose. Use Python `True`, `False`, and `None` inside that string.

benchmarks/benchmark_repository_context.py

@@ -0,0 +1,70 @@
+"""Reproducible cold-versus-warm Code Mode repository scan microbenchmark.
+
+Run after installing PyLattice: python benchmarks/benchmark_repository_context.py.
+This measures filesystem context collection only; not LLM or router latency.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import statistics
+import tempfile
+import time
+from pathlib import Path
+
+from agent_tui.config import Settings
+from agent_tui.context import ContextManager
+from agent_tui.tools import ToolRegistry
+
+
+async def main() -> None:
+    with tempfile.TemporaryDirectory(prefix="pylattice-context-bench-") as directory:
+        root = Path(directory)
+        source_dir = root / "src"
+        source_dir.mkdir()
+        for number in range(240):
+            marker = "needle" if number == 175 else "other_symbol"
+            body = (f"def helper_{number}():\n    return '{marker}'\n" * 200)
+            (source_dir / f"module_{number:03d}.py").write_text(body, encoding="utf-8")
+
+        registry = ToolRegistry(Settings(workspace=root, max_output_chars=16000))
+        registry.bind_context(ContextManager(registry.settings))
+        query = {"query": "needle", "path": "src", "max_files": 4}
+
+        started = time.perf_counter()
+        first = await registry.execute("context_collect", query)
+        cold_ms = 1000 * (time.perf_counter() - started)
+        if not first.ok:
+            raise RuntimeError(first.content)
+        cold = json.loads(first.content)
+
+        elapsed: list[float] = []
+        warm = {}
+        for _ in range(7):
+            started = time.perf_counter()
+            result = await registry.execute("context_collect", query)
+            elapsed.append(1000 * (time.perf_counter() - started))
+            if not result.ok:
+                raise RuntimeError(result.content)
+            warm = json.loads(result.content)
+
+        print(
+            json.dumps(
+                {
+                    "files_in_repo": 240,
+                    "source_bytes": sum(p.stat().st_size for p in source_dir.iterdir()),
+                    "cold_ms": round(cold_ms, 3),
+                    "warm_median_ms": round(statistics.median(elapsed), 3),
+                    "cold_disk_reads": cold["disk_reads"],
+                    "warm_disk_reads": warm["disk_reads"],
+                    "warm_index_hits": warm["index_hits"],
+                    "matches": [entry["path"] for entry in warm["files"]],
+                },
+                indent=2,
+            )
+        )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

src/agent_tui/agent.py

@@ -72,6 +72,8 @@ def __init__(
         self.planner = runtime.get("planner")
         self.subagents = runtime.get("subagents") if runtime.allow_delegation else None
         self.tool_bindings.register(self.registry, self.skills, self.memory, self.subagents)
+        if hasattr(self.registry, "bind_context") and hasattr(self.context, "cache_file"):
+            self.registry.bind_context(self.context)
 
     @property
     def requires_api_key(self) -> bool:
@@ -113,6 +115,8 @@ def clear(self) -> None:
         if self.running:
             raise RuntimeError("Stop the current run before clearing the conversation")
         self.context.clear()
+        if hasattr(self.registry, "clear_index"):
+            self.registry.clear_index()
         self.skills.clear()
 
     def _context(

src/agent_tui/codemode.py

@@ -5,6 +5,7 @@
 import asyncio
 import json
 from collections.abc import Awaitable, Callable
+from contextlib import nullcontext
 from typing import Any
 
 from pydantic_monty import AsyncMonty, CollectString, MontyError, MontySyntaxError, MontyTypingError
@@ -109,6 +110,10 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> ToolReply: ...
    content includes display prefixes like "12: text"; those prefixes are not file content.
    search_files: {"matches": [{"path": str, "line": int, "text": str}],
                   "truncated": bool, "skipped": int}; query is literal text, not regex.
+   context_collect: {"files": [{"path": str, "sha256": str, "content": str,
+                      "start_line": int, "total_lines": int, "excerpt": bool}],
+                      "disk_reads": int, "index_hits": int, "truncated": bool}.
+   apply_patchset: {"changes": [{"path": str, "bytes": int}], "count": int}.
    edit_file/write_file: {"path": str, "bytes": int}.
    run_command: {"exit_code": int, "output": str, "truncated": bool, "timed_out": bool}.
 Use result["output"]["content"] for file text and result["output"]["output"] for command
@@ -119,11 +124,20 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> ToolReply: ...
 or print short evidence. Do not end with only assignments, an unawaited coroutine, a set,
 bytes or a custom object. Do not use a top-level return. Intermediate results are not shown
 to the next model turn unless included in the final expression or printed.
-8. For changes, inspect source first, apply exact edit_file patches without line-number
-prefixes, and verify with a read or appropriate command. write_file is for new files or
-explicitly requested full replacements. Never reconstruct an entire file from a partial
-read. Batch a change only when it is clear; otherwise return the source and reason in the
-next turn. Continue using execute_code while requested edits or verification remain.
+8. For repository work, prefer context_collect over repeated list_files/search_files/
+read_file calls: use {"query": "term", "max_files": 6} or {"paths": ["src/a.py",
+"src/b.py"]}. Files are cached in Context Manager across independent Code Mode programs
+and relevant verified excerpts are added to later model contexts. Returned content has
+no line-number prefixes; use exact old_text from it. Excerpts are not whole files.
+For refactors or new files, use apply_patchset with a list of actions. An edit is
+{"action": "edit", "path": "src/a.py", "old_text": "exact old block",
+"new_text": "replacement", "expected_sha256": "digest from context_collect"}.
+A create is {"action": "create", "path": "src/new.py", "content": "source"}.
+The patchset validates all files before writing and requires the usual approval.
+Do not overwrite existing files with create. Check read-only tool availability.
+For single edits, edit_file still works; write_file is for new files or requested
+full replacements. Never reconstruct a full file from a partial excerpt. Verify
+changed files with context_collect or a test command before finishing.
 9. Keep output and loops bounded. A program has at most 32 tool calls, 5 seconds of
 computation (excluding tools/approvals), 64 MiB memory and 24,000 source characters.
 On failure, read error, phase, tools and retry_hint. Correct the cause in a NEW program.
@@ -133,13 +147,9 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> ToolReply: ...
 and blockers honestly. A direct final answer is allowed only when finish is selected or
 routing falls back and no work remains; never answer directly when execute_code is selected.
 
-Example Python source for bounded inspection:
-listing = await call_tool("list_files", {"path": "src"})
-results = []
-for path in listing["output"]["files"][:3]:
-    result = await call_tool("read_file", {"path": path})
-    results.append(result["output"])
-{"files": results, "truncated": listing["output"]["truncated"]}
+Example Python source for bounded inspection in one host tool call:
+result = await call_tool("context_collect", {"query": "example", "max_files": 4})
+{"files": result["output"]["files"], "truncated": result["output"]["truncated"]}
 """
 
 
@@ -214,6 +224,7 @@ async def execute_code(
     run_tool: ToolRunner,
     *,
     format_result: Callable[[ToolResult], dict[str, Any]] = ToolResult.as_dict,
+    pool: AsyncMonty | None = None,
 ) -> ToolResult:
     """Expose only a JSON tool bridge; never share host objects, mounts or OS handlers."""
     if not code.strip() or len(code) > MAX_CODE_CHARS:
@@ -274,8 +285,12 @@ async def call_tool(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
 
     try:
         async with (
-            AsyncMonty(max_processes=1, request_timeout=EXECUTION_SECONDS + 2) as pool,
-            pool.checkout(
+            (
+                nullcontext(pool)
+                if pool is not None
+                else AsyncMonty(max_processes=1, request_timeout=EXECUTION_SECONDS + 2)
+            ) as runtime_pool,
+            runtime_pool.checkout(
                 type_check=True,
                 type_check_stubs=TOOL_BRIDGE_STUBS,
                 type_check_format="concise",