Enable 1h prompt caching and tool search for Claude client

src/ucode/agents/claude.py

@@ -91,6 +91,9 @@ def _resolve_web_search_model(state: dict) -> str | None:
     "ANTHROPIC_DEFAULT_HAIKU_MODEL",
     "ANTHROPIC_DEFAULT_HAIKU_MODEL_NAME",
 )
+# Env keys ucode used to write but no longer does; stripped from the managed
+# settings file on every launch so stale values never linger.
+CLAUDE_REMOVED_ENV_KEYS = ("CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS",)
 CLAUDE_TRACING_STOP_HOOK_SUFFIX = " autolog claude stop-hook"
 # Tracing is driven by an `mlflow autolog claude stop-hook` Stop hook, run by
 # the `mlflow` CLI on each session end. Pin to 3.11.x: 3.12 dropped the Unity
@@ -163,8 +166,12 @@ def render_overlay(
     env: dict[str, str] = {
         "ANTHROPIC_BASE_URL": base_url,
         "ANTHROPIC_CUSTOM_HEADERS": custom_headers,
-        "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1",
         "CLAUDE_CODE_API_KEY_HELPER_TTL_MS": "900000",
+        # 1h prompt caching needs the extended-cache-ttl beta header, which
+        # Claude Code only sends when experimental betas are enabled — so we must
+        # not set CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS (see CLAUDE_REMOVED_ENV_KEYS).
+        "ENABLE_PROMPT_CACHING_1H": "1",
+        "ENABLE_TOOL_SEARCH": "true",
     }
     # Intentionally NOT setting ANTHROPIC_MODEL. Setting it produces a duplicate
     # catalog row in Claude Code's /model picker (e.g. "Opus 4.8 (1M context) ✓")
@@ -331,6 +338,11 @@ def write_tool_config(
         for key in CLAUDE_MANAGED_MODEL_ENV_KEYS:
             if key not in overlay_env:
                 merged_env.pop(key, None)
+    # deep_merge_dict keeps keys already in the file, so drop the ones ucode no
+    # longer writes.
+    if isinstance(merged_env, dict):
+        for key in CLAUDE_REMOVED_ENV_KEYS:
+            merged_env.pop(key, None)
     write_json_file(CLAUDE_SETTINGS_PATH, merged)
 
     if web_search_model:

src/ucode/state.py

@@ -153,7 +153,9 @@ def build_agent_state(state: dict) -> dict[str, dict]:
             "env": {
                 "ANTHROPIC_BASE_URL": base_urls.get("claude"),
                 "CLAUDE_CODE_API_KEY_HELPER_TTL_MS": str(AUTH_REFRESH_INTERVAL_MS),
-                "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1",
+                # Kept in sync with agents/claude.py render_overlay.
+                "ENABLE_PROMPT_CACHING_1H": "1",
+                "ENABLE_TOOL_SEARCH": "true",
             },
         },
         "codex": {

tests/test_agent_claude.py

@@ -80,9 +80,18 @@ def test_sets_custom_headers(self):
         overlay, _ = claude.render_overlay(WS, "s4")
         assert "x-databricks-use-coding-agent-mode" in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"]
 
-    def test_disables_experimental_betas(self):
+    def test_does_not_disable_experimental_betas(self):
+        # Would suppress the beta header 1h prompt caching needs.
         overlay, _ = claude.render_overlay(WS, "s4")
-        assert overlay["env"]["CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS"] == "1"
+        assert "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS" not in overlay["env"]
+
+    def test_enables_prompt_caching_1h(self):
+        overlay, _ = claude.render_overlay(WS, "s4")
+        assert overlay["env"]["ENABLE_PROMPT_CACHING_1H"] == "1"
+
+    def test_enables_tool_search(self):
+        overlay, _ = claude.render_overlay(WS, "s4")
+        assert overlay["env"]["ENABLE_TOOL_SEARCH"] == "true"
 
     def test_sets_api_key_helper(self):
         overlay, _ = claude.render_overlay(WS, "s4")
@@ -335,6 +344,29 @@ def test_explicit_override_used_over_codex_models(self, monkeypatch):
         assert calls == [("register", WS, "explicit-model")]
 
 
+class TestWriteToolConfigStripsRemovedEnvKeys:
+    """Stale keys ucode no longer writes are dropped from the merged settings."""
+
+    def _patch(self, monkeypatch, existing, written):
+        monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True)
+        monkeypatch.setattr(claude, "read_json_safe", lambda path: existing)
+        monkeypatch.setattr(
+            claude, "write_json_file", lambda path, payload: written.append(payload)
+        )
+        monkeypatch.setattr(claude, "save_state", lambda state: None)
+        monkeypatch.setattr(claude, "_register_web_search_mcp", lambda *a, **kw: True)
+
+    def test_strips_stale_disable_experimental_betas(self, monkeypatch):
+        existing = {"env": {"CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1"}}
+        written: list = []
+        self._patch(monkeypatch, existing, written)
+        state = {"workspace": WS, "codex_models": []}
+        claude.write_tool_config(state, "databricks-claude-sonnet-4")
+        assert "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS" not in written[0]["env"]
+        assert written[0]["env"]["ENABLE_PROMPT_CACHING_1H"] == "1"
+        assert written[0]["env"]["ENABLE_TOOL_SEARCH"] == "true"
+
+
 class TestRegisterWebSearchMcp:
     def test_clears_existing_then_adds(self, monkeypatch):
         import ucode.mcp as mcp_mod

tests/test_e2e.py

@@ -504,7 +504,6 @@ def test_launch_claude_per_model(
                 "ANTHROPIC_MODEL": model_id,
                 "ANTHROPIC_BASE_URL": base_url,
                 "ANTHROPIC_API_KEY": e2e_token,
-                "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1",
             }
             cmd = claude.validate_cmd("claude")
             result = _run_agent(cmd, env=env, timeout=90)
@@ -571,7 +570,6 @@ def test_launch_claude_through_provider(
             "CLAUDE_CONFIG_DIR": str(config_dir),
             "ANTHROPIC_BASE_URL": build_tool_base_url("claude", e2e_workspace),
             "ANTHROPIC_API_KEY": e2e_token,
-            "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1",
         }
         result = _run_agent(claude.validate_cmd("claude"), env=env, timeout=90)
         combined = (result.stdout + result.stderr).strip()