From 070938b0c6d7ab7cc129636ddba19c5916e6b272 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Mon, 13 Jul 2026 15:26:25 +0000 Subject: [PATCH] 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 --- src/ucode/agents/claude.py | 14 +++++++++++++- src/ucode/state.py | 4 +++- tests/test_agent_claude.py | 36 ++++++++++++++++++++++++++++++++++-- tests/test_e2e.py | 2 -- 4 files changed, 50 insertions(+), 6 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 422e763..5c5c6c3 100644 --- a/src/ucode/agents/claude.py +++ b/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: diff --git a/src/ucode/state.py b/src/ucode/state.py index 372af7f..a36a281 100644 --- a/src/ucode/state.py +++ b/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": { diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index b2d3a27..ef377d1 100644 --- a/tests/test_agent_claude.py +++ b/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 diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 604452a..ff9e2b9 100644 --- a/tests/test_e2e.py +++ b/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()