Enable 1h prompt caching and tool search for Claude client (#208)

* Enable 1h prompt caching and tool search for Claude client

Set ENABLE_PROMPT_CACHING_1H=1 and ENABLE_TOOL_SEARCH=true in the Claude
settings env block, matching Isaac's config.

The 1h cache depends on the extended-cache-ttl-2025-04-11 beta header,
which Claude Code only attaches when experimental betas are enabled. ucode
previously set CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS=1, which suppressed
that header, so the flag is removed.

Because deep_merge_dict preserves keys already in ucode-settings.json,
removing the key from the written env is not enough for existing installs.
Add a CLAUDE_REMOVED_ENV_KEYS list, stripped from the managed settings file
on every launch, so the stale key never lingers. This only touches ucode's
own managed file, so a user who wants experimental betas disabled can still
set the flag in their own ~/.claude/settings.json.

Keep build_agent_state in sync with the same flag change.

Co-authored-by: Isaac

* Set ENABLE_TOOL_SEARCH=1 and enable CLAUDE_CODE_USE_GATEWAY

Co-authored-by: Isaac

Author: lilly-luo

src/ucode/agents/claude.py

@@ -93,6 +93,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
@@ -166,8 +169,13 @@ 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": "1",
+        "CLAUDE_CODE_USE_GATEWAY": "1",
     }
     # 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) ✓")
@@ -343,6 +351,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,10 @@ 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": "1",
+                "CLAUDE_CODE_USE_GATEWAY": "1",
             },
         },
         "codex": {

tests/test_agent_claude.py

@@ -80,9 +80,22 @@ 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"] == "1"
+
+    def test_enables_use_gateway(self):
+        overlay, _ = claude.render_overlay(WS, "s4")
+        assert overlay["env"]["CLAUDE_CODE_USE_GATEWAY"] == "1"
 
     def test_sets_api_key_helper(self):
         overlay, _ = claude.render_overlay(WS, "s4")
@@ -366,6 +379,30 @@ 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"] == "1"
+        assert written[0]["env"]["CLAUDE_CODE_USE_GATEWAY"] == "1"
+
+
 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()