From da81f65595935cba69bd2a9eb97932808100577c Mon Sep 17 00:00:00 2001 From: Pooja Date: Sat, 17 Oct 2026 19:37:11 +0000 Subject: [PATCH] Support custom coding-agent attribution via Unity Gateway --- README.md | 21 +++ src/ucode/agents/claude.py | 6 +- src/ucode/agents/codex.py | 2 + src/ucode/agents/coding_agent.py | 26 ++++ src/ucode/agents/gemini.py | 11 ++ src/ucode/agents/kilo.py | 2 + src/ucode/agents/opencode.py | 2 + src/ucode/agents/pi.py | 9 +- src/ucode/constants.py | 3 + src/ucode/desktop_setup.py | 2 + tests/README.md | 7 + tests/integration/README.md | 6 + tests/test_agent_codex.py | 16 +++ tests/test_claude_managed_headers.py | 1 + tests/test_coding_agent.py | 204 +++++++++++++++++++++++++++ tests/test_desktop_setup.py | 32 +++++ 16 files changed, 348 insertions(+), 2 deletions(-) create mode 100644 src/ucode/agents/coding_agent.py create mode 100644 tests/test_coding_agent.py diff --git a/README.md b/README.md index e21fb4983..c1c5307f5 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,27 @@ ug removes its shared reference on discovery/validation failure, reconfiguration revert, or before installing/updating Codex. After an update, run `ug codex` to refresh discovery or `ug configure` for a managed static list, then restart the app server. +## Custom Coding-Agent Attribution + +A wrapper or custom agent can identify itself to AI Gateway: + +```bash +UNITY_GATEWAY_CUSTOM_CODING_AGENT=my-agent ug claude +``` + +This sends `Databricks-Custom-Coding-Agent: my-agent` without changing authentication +or `User-Agent`. The gateway uses this identity instead of the native agent name +and does not record the native agent's version for the override. Leave the variable +unset or blank to retain native agent/version detection. This does not add surface headers. +Names containing control characters are warned about and ignored. + +Supported adapters: Claude Code, Codex, Gemini CLI, OpenCode, Kilo, and Pi. +Gemini CLI additionally requires names without commas. Pi reads the identity from +the environment rather than interpreting it as a literal config value. +Admin-managed headers take precedence. Claude Desktop picks up the override during +`ug configure`; rerun configure without the variable to remove it from that profile. +Copilot CLI and MCP-only agents do not support this override. + ## Configure ```bash diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 251d92fb7..a7bedf781 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -27,6 +27,7 @@ ) from ucode.constants import ( AI_GATEWAY_REQUEST_TAGS_HEADER, + CUSTOM_CODING_AGENT_HEADER, LOOPBACK_HOST, MCP_CLEANUP_SCOPES, MCP_USER_SCOPE, @@ -103,6 +104,7 @@ from ucode.ui import print_note, print_success, print_warning from .args import LaunchOptions, has_explicit_model_arg +from .coding_agent import custom_coding_agent_headers from .request_tags import request_tags_header_value GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" @@ -242,6 +244,7 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: { "x-databricks-use-coding-agent-mode", "user-agent", + CUSTOM_CODING_AGENT_HEADER.casefold(), MODEL_PROVIDER_SERVICE_HEADER.casefold(), MODEL_SERVICE_PARENT_SCHEMA_HEADER.casefold(), SMART_ROUTER_RECIPE_HEADER.casefold(), # Remove recipe headers written by older ug versions. @@ -252,7 +255,7 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: CLAUDE_OPTIONAL_CUSTOM_HEADER_NAMES = frozenset({AI_GATEWAY_REQUEST_TAGS_HEADER.casefold()}) # These attribute inference traffic; neither selects the gateway or provider. CLAUDE_OPTIONAL_ATTRIBUTION_HEADER_NAMES = frozenset( - {"user-agent", SMART_ROUTER_RECIPE_HEADER.casefold()} + {"user-agent", SMART_ROUTER_RECIPE_HEADER.casefold(), CUSTOM_CODING_AGENT_HEADER.casefold()} ) # Relayed drops the user scope to deliberately omit the stale apiKeyHelper. Only applied to relayed # launches — normal launches keep loading user settings (hooks/permissions) as before. @@ -508,6 +511,7 @@ def render_overlay( "x-databricks-use-coding-agent-mode: true", f"User-Agent: ucode/{ug_version()} claude/{agent_version('claude')}", ] + header_lines.extend(f"{name}: {value}" for name, value in custom_coding_agent_headers().items()) if provider: header_lines.append(f"{MODEL_PROVIDER_SERVICE_HEADER}: {provider}") elif parent_schema: diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 3408c9365..c39fd0d73 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -94,6 +94,7 @@ from .args import LaunchOptions from .codex_catalog import prepare_codex_catalog, validate_codex_catalog +from .coding_agent import custom_coding_agent_headers from .request_tags import request_tags_header_value CODEX_CONFIG_DIR = Path.home() / ".codex" @@ -218,6 +219,7 @@ def _provider_block( http_headers = { "User-Agent": f"ucode/{ug_version()} codex/{agent_version('codex')}", } + http_headers.update(custom_coding_agent_headers()) if provider: http_headers[MODEL_PROVIDER_SERVICE_HEADER] = provider elif parent_schema: diff --git a/src/ucode/agents/coding_agent.py b/src/ucode/agents/coding_agent.py new file mode 100644 index 000000000..3d21fdee5 --- /dev/null +++ b/src/ucode/agents/coding_agent.py @@ -0,0 +1,26 @@ +"""Custom coding-agent attribution shared by agent adapters.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping + +from ucode.constants import ( + CUSTOM_CODING_AGENT_HEADER, + UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, +) +from ucode.ui import print_warning_err + + +def custom_coding_agent_headers(env: Mapping[str, str] | None = None) -> dict[str, str]: + """Return the explicitly configured identity override, or no additional headers.""" + source = os.environ if env is None else env + value = source.get(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR) or "" + if any(ord(character) < 32 or ord(character) == 127 for character in value): + print_warning_err( + f"Ignoring {UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR}: " + "value contains control characters. Provide a single-line agent name." + ) + return {} + value = value.strip() + return {CUSTOM_CODING_AGENT_HEADER: value} if value else {} diff --git a/src/ucode/agents/gemini.py b/src/ucode/agents/gemini.py index 6d49ea7ed..e589685e8 100644 --- a/src/ucode/agents/gemini.py +++ b/src/ucode/agents/gemini.py @@ -20,6 +20,7 @@ write_dotenv, write_json_file, ) +from ucode.constants import UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR from ucode.databricks import ( TOKEN_REFRESH_INTERVAL_SECONDS, build_tool_base_url, @@ -33,8 +34,10 @@ set_provider_service, ) from ucode.telemetry import agent_version, ug_version +from ucode.ui import print_warning_err from .args import LaunchOptions +from .coding_agent import custom_coding_agent_headers GEMINI_CONFIG_DIR = Path.home() / ".gemini" GEMINI_ENV_PATH = GEMINI_CONFIG_DIR / "ucode.env" @@ -142,6 +145,14 @@ def render_env_overlay( # so a key named `User-Agent` overrides the default. Resolved via # upstream issue google-gemini/gemini-cli#10088. custom_headers = f"User-Agent:ucode/{ug_version()} gemini/{agent_version('gemini')}" + for name, value in custom_coding_agent_headers().items(): + if "," in value: + print_warning_err( + f"Ignoring {UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR}: " + "Gemini CLI cannot send a custom agent name containing commas." + ) + else: + custom_headers += f",{name}:{value}" if provider: # A Model Provider Service routes by this header; the request still names # the service's target model in `GEMINI_MODEL` (pinned by the launch path). diff --git a/src/ucode/agents/kilo.py b/src/ucode/agents/kilo.py index 053554ac5..1fecf9a78 100644 --- a/src/ucode/agents/kilo.py +++ b/src/ucode/agents/kilo.py @@ -34,6 +34,7 @@ from ucode.telemetry import agent_version, ug_version from .args import LaunchOptions, explicit_model_arg_value, has_explicit_model_arg +from .coding_agent import custom_coding_agent_headers KILO_XDG_CONFIG_HOME = APP_DIR / "kilo-xdg" KILO_CONFIG_DIR = KILO_XDG_CONFIG_HOME / "kilo" @@ -258,6 +259,7 @@ def render_overlay( ua_header = { "User-Agent": f"ucode/{ug_version()} kilo/{agent_version('kilo')}", } + ua_header.update(custom_coding_agent_headers()) anthropic_models = opencode_models.get("anthropic") or [] gemini_models = opencode_models.get("gemini") or [] diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index b9eb1f655..d634f7866 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -29,6 +29,7 @@ from ucode.telemetry import agent_version, ug_version from .args import LaunchOptions, explicit_model_arg_value, has_explicit_model_arg +from .coding_agent import custom_coding_agent_headers from .request_tags import request_tags_header_value OPENCODE_XDG_CONFIG_HOME = APP_DIR / "opencode-xdg" @@ -306,6 +307,7 @@ def render_overlay( ua_header = { "User-Agent": f"ucode/{ug_version()} opencode/{agent_version('opencode')}", } + ua_header.update(custom_coding_agent_headers()) request_tags = request_tags_header_value() if request_tags: ua_header[AI_GATEWAY_REQUEST_TAGS_HEADER] = request_tags diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index 0dde09687..1e4f9d912 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -42,7 +42,11 @@ read_json_safe, write_json_file, ) -from ucode.constants import AI_GATEWAY_REQUEST_TAGS_HEADER +from ucode.constants import ( + AI_GATEWAY_REQUEST_TAGS_HEADER, + CUSTOM_CODING_AGENT_HEADER, + UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, +) from ucode.databricks import ( ANTHROPIC_FAMILIES, build_auth_token_argv, @@ -55,6 +59,7 @@ from ucode.telemetry import agent_version, ug_version from .args import LaunchOptions +from .coding_agent import custom_coding_agent_headers from .request_tags import request_tags_header_value PI_UCODE_HOME = APP_DIR / "pi-home" @@ -121,6 +126,8 @@ def render_overlay( # Pi expands header values that match an env var name. Our UA contains # `/` and a space so it can never collide — safe to pass as a literal. ua_headers = {"User-Agent": f"ucode/{ug_version()} pi/{agent_version('pi')}"} + if custom_coding_agent_headers(): + ua_headers[CUSTOM_CODING_AGENT_HEADER] = UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR request_tags = request_tags_header_value() if request_tags: # Similarly to UA, request_tags is a JSON object (validated in `request_tags_header_value()`) diff --git a/src/ucode/constants.py b/src/ucode/constants.py index 3ef84ee25..46454e51f 100644 --- a/src/ucode/constants.py +++ b/src/ucode/constants.py @@ -28,6 +28,9 @@ AI_GATEWAY_REQUEST_TAGS_HEADER = "Databricks-Ai-Gateway-Request-Tags" AI_GATEWAY_REQUEST_TAGS_ENV_VAR = "AI_GATEWAY_REQUEST_TAGS" +CUSTOM_CODING_AGENT_HEADER = "Databricks-Custom-Coding-Agent" +UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR = "UNITY_GATEWAY_CUSTOM_CODING_AGENT" + # MCP server registration scopes. Claude Code supports local/project/user; the # other CLIs only take the user-scope name. Kept here (a leaf module) so both diff --git a/src/ucode/desktop_setup.py b/src/ucode/desktop_setup.py index 24979117d..47d8d8a63 100644 --- a/src/ucode/desktop_setup.py +++ b/src/ucode/desktop_setup.py @@ -6,6 +6,7 @@ import re import sys +from ucode.agents.coding_agent import custom_coding_agent_headers from ucode.agents.desktop.claude import ( claude_desktop_directory, revert_claude_desktop, @@ -76,6 +77,7 @@ def _configure_desktop(state: dict, *, parent_schema: str | None) -> None: ) headers = {"x-databricks-use-coding-agent-mode": "true"} + headers.update(custom_coding_agent_headers()) if provider: headers[MODEL_PROVIDER_SERVICE_HEADER] = provider elif parent_schema: diff --git a/tests/README.md b/tests/README.md index 17c67b403..b31b3c076 100644 --- a/tests/README.md +++ b/tests/README.md @@ -12,6 +12,13 @@ mocks, monkeypatching, fake binaries/services, or fabricated ug state. | Integration CUJs | `integration/test_*.py` | Public configure, TUI, script, command, protocol, and lifecycle journeys | | Installation | `integration/test_installation.py` | Fresh installed package, CLI, and local helpers without credentials on Linux and advisory native Windows | +`test_coding_agent.py` covers `UNITY_GATEWAY_CUSTOM_CODING_AGENT` validation and +header rendering for Claude Code, modern/legacy Codex, Gemini, OpenCode, Kilo, and Pi. +It checks unchanged native User-Agent values, admin precedence, repeat-render cleanup, +Claude header merging, Gemini delimiter rejection, and Pi's environment-reference transport. +Codex config-write and Desktop setup tests also check removal on reconfiguration. +These are component checks; no live request or backend attribution is asserted. + Claude mod entry-point coverage lives in `test_claude_mod.py` (generated plugin contents), `test_claude_smart_routing_v2.py` (launch wiring), and the explicitly selected `native_claude/test_mod.py` suite. Claude 2.1.290 validates the module and runs its diff --git a/tests/integration/README.md b/tests/integration/README.md index e9f5f2f75..a6ab2853e 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -1,5 +1,11 @@ # Integration tests +Custom coding-agent attribution is covered by component tests in +`../test_coding_agent.py`, `../test_agent_codex.py`, `../test_claude_managed_headers.py`, +and `../test_desktop_setup.py`: validation, adapter headers, admin precedence, +and removal on later renders/reconfiguration. This suite does not yet assert +`UNITY_GATEWAY_CUSTOM_CODING_AGENT` on live inference or in gateway usage logs. + The separate [dedicated-workspace CUJ](../e2e_cuj/AGENTS.md) lives in `tests/e2e_cuj/` and runs directly with pytest. It does not use this suite's runner or config fixtures. It reuses the session, terminal, file-task, and transcript helpers in `utils/`. diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index f887f4d20..0c8e2681f 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -234,6 +234,22 @@ def test_managed_keys_include_http_headers(self): class TestCodexWriteConfig: + def test_rewrite_removes_custom_agent_override(self, tmp_path, monkeypatch): + config_path = tmp_path / "ucode.config.toml" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + monkeypatch.setattr(codex, "agent_version", lambda _binary: "0.145.0") + monkeypatch.setenv("UNITY_GATEWAY_CUSTOM_CODING_AGENT", "custom-agent") + state = {"workspace": WS, "codex_models": []} + codex.write_tool_config(state) + headers = read_toml_safe(config_path)["model_providers"]["Databricks"]["http_headers"] + assert headers["Databricks-Custom-Coding-Agent"] == "custom-agent" + + monkeypatch.delenv("UNITY_GATEWAY_CUSTOM_CODING_AGENT") + codex.write_tool_config(state) + headers = read_toml_safe(config_path)["model_providers"]["Databricks"]["http_headers"] + assert "Databricks-Custom-Coding-Agent" not in headers + def test_writes_ucode_profile_config_file(self, tmp_path, monkeypatch): config_path = tmp_path / ".codex" / "ucode.config.toml" backup_path = tmp_path / "codex-ucode-config.backup.toml" diff --git a/tests/test_claude_managed_headers.py b/tests/test_claude_managed_headers.py index 8d6885427..70e9046cd 100644 --- a/tests/test_claude_managed_headers.py +++ b/tests/test_claude_managed_headers.py @@ -11,6 +11,7 @@ ("x-mode: true", "x-mode: true\nUser-Agent: ug/new"), ("x-mode: true\nUser-Agent: isaac", "User-Agent: ug/new\nx-mode: true"), ("x-mode: true", "x-mode: true\nDatabricks-Smart-Router-Recipe: recipe"), + ("x-mode: true", "x-mode: true\nDatabricks-Custom-Coding-Agent: custom"), ("X-Mode: true\nX-Tags: source=isaac", "x-tags: source=isaac\nx-mode:\ttrue "), ("X-Mode: true\r\n", "x-mode: true\n\n"), ], diff --git a/tests/test_coding_agent.py b/tests/test_coding_agent.py new file mode 100644 index 000000000..f3bdede8e --- /dev/null +++ b/tests/test_coding_agent.py @@ -0,0 +1,204 @@ +"""Component coverage for explicit custom coding-agent attribution.""" + +from __future__ import annotations + +import pytest + +from ucode.agents import claude, codex, gemini, kilo, opencode, pi +from ucode.agents.coding_agent import custom_coding_agent_headers +from ucode.constants import ( + CUSTOM_CODING_AGENT_HEADER, + UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, +) + +WORKSPACE = "https://example.databricks.com" + + +@pytest.mark.parametrize( + ("env", "expected"), + [ + ({}, {}), + ({UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR: ""}, {}), + ({UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR: " "}, {}), + ( + {UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR: "my-agent"}, + {CUSTOM_CODING_AGENT_HEADER: "my-agent"}, + ), + ( + {UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR: " My Agent/1:custom "}, + {CUSTOM_CODING_AGENT_HEADER: "My Agent/1:custom"}, + ), + ({"AI_GATEWAY_CUSTOM_CODING_AGENT": "old-name"}, {}), + ], +) +def test_custom_coding_agent_headers(env, expected): + assert custom_coding_agent_headers(env) == expected + + +@pytest.mark.parametrize( + "value", ["agent\nInjected: value", "agent\r", "\nagent", "agent\x00", "agent\t", "agent\x7f"] +) +def test_control_characters_warn_on_stderr_and_omit_header(value, capsys): + assert custom_coding_agent_headers({UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR: value}) == {} + output = capsys.readouterr() + assert not output.out + assert UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR in output.err + assert "control characters" in output.err + assert "Injected" not in output.err + + +def test_helper_reads_process_environment(monkeypatch): + monkeypatch.setenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, "wrapper") + assert custom_coding_agent_headers() == {CUSTOM_CODING_AGENT_HEADER: "wrapper"} + monkeypatch.delenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR) + assert custom_coding_agent_headers() == {} + + +def _parse_header_lines(value: str, separator: str) -> dict[str, str]: + headers = {} + for line in value.split(separator): + name, _, header_value = line.partition(":") + headers[name.strip()] = header_value.strip() + return headers + + +@pytest.fixture(params=["claude", "codex", "codex-legacy", "gemini", "opencode", "kilo", "pi"]) +def rendered_headers(request, monkeypatch): + agent = request.param + module = { + "claude": claude, + "codex": codex, + "codex-legacy": codex, + "gemini": gemini, + "opencode": opencode, + "kilo": kilo, + "pi": pi, + }[agent] + monkeypatch.setattr(module, "ug_version", lambda: "1.2.3") + monkeypatch.setattr(module, "agent_version", lambda _binary: "4.5.6") + monkeypatch.delenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, raising=False) + monkeypatch.delenv("AI_GATEWAY_REQUEST_TAGS", raising=False) + + def render(): + if agent == "claude": + overlay, _ = claude.render_overlay(WORKSPACE, "claude-model") + return [_parse_header_lines(overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"], "\n")] + if agent in ("codex", "codex-legacy"): + renderer = codex.render_overlay if agent == "codex" else codex.render_legacy_overlay + overlay = renderer(WORKSPACE) + return [overlay["model_providers"]["Databricks"]["http_headers"]] + if agent == "gemini": + overlay = gemini.render_env_overlay(WORKSPACE, "gemini-model", "token") + return [_parse_header_lines(overlay["GEMINI_CLI_CUSTOM_HEADERS"], ",")] + if agent == "pi": + overlay, _ = pi.render_overlay( + "claude-model", + "token", + {"claude": WORKSPACE, "openai": WORKSPACE, "gemini": WORKSPACE}, + {"sonnet": "claude-model"}, + ["gpt-model"], + ["gemini-model"], + ) + return [provider["headers"] for provider in overlay["providers"].values()] + overlay, _ = module.render_overlay( + "claude-model", + "token", + {"anthropic": WORKSPACE, "gemini": WORKSPACE, "oss": WORKSPACE}, + {"anthropic": ["claude-model"], "gemini": ["gemini-model"], "oss": ["oss-model"]}, + ) + return [ + model["headers"] + for provider in overlay["provider"].values() + for model in provider["models"].values() + ] + + return agent, render + + +def test_adapters_emit_override_without_changing_user_agent(rendered_headers, monkeypatch): + agent, render = rendered_headers + baseline = render() + monkeypatch.setenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, "custom-agent") + overridden = render() + expected = UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR if agent == "pi" else "custom-agent" + assert overridden + for original, headers in zip(baseline, overridden, strict=True): + assert headers[CUSTOM_CODING_AGENT_HEADER] == expected + assert { + name: value for name, value in headers.items() if name != CUSTOM_CODING_AGENT_HEADER + } == original + assert "Databricks-Custom-Client-Surfaces" not in headers + + +@pytest.mark.parametrize("value", ["", " ", "bad\nInjected: header"]) +def test_adapters_omit_empty_or_invalid_override(rendered_headers, monkeypatch, value): + _, render = rendered_headers + baseline = render() + monkeypatch.setenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, value) + assert render() == baseline + + +def test_adapters_remove_override_on_next_render(rendered_headers, monkeypatch): + _, render = rendered_headers + baseline = render() + monkeypatch.setenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, "first-agent") + assert all(CUSTOM_CODING_AGENT_HEADER in headers for headers in render()) + monkeypatch.delenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR) + assert render() == baseline + + +@pytest.mark.parametrize("agent", [claude, codex]) +def test_admin_header_overrides_environment_case_insensitively(agent, monkeypatch): + monkeypatch.setenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, "session-agent") + managed_headers = {CUSTOM_CODING_AGENT_HEADER.lower(): "admin-agent"} + if agent is claude: + overlay, _ = claude.render_overlay( + WORKSPACE, "claude-model", managed_http_headers=managed_headers + ) + headers = _parse_header_lines(overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"], "\n") + else: + overlay = codex.render_overlay(WORKSPACE, managed_http_headers=managed_headers) + headers = overlay["model_providers"]["Databricks"]["http_headers"] + assert [ + value + for name, value in headers.items() + if name.casefold() == CUSTOM_CODING_AGENT_HEADER.casefold() + ] == ["admin-agent"] + + +def test_claude_merge_replaces_and_removes_override_but_preserves_other_headers(monkeypatch): + existing = "X-Caller: keep\ndatabricks-custom-coding-agent: old" + monkeypatch.setenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, "new-agent") + overlay, _ = claude.render_overlay(WORKSPACE, "claude-model") + merged = claude._merge_anthropic_custom_headers( + existing, overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"] + ) + assert _parse_header_lines(merged, "\n")[CUSTOM_CODING_AGENT_HEADER] == "new-agent" + monkeypatch.delenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR) + overlay, _ = claude.render_overlay(WORKSPACE, "claude-model") + merged = claude._merge_anthropic_custom_headers( + merged, overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"] + ) + headers = _parse_header_lines(merged, "\n") + assert headers["X-Caller"] == "keep" + assert not any(name.casefold() == CUSTOM_CODING_AGENT_HEADER.casefold() for name in headers) + + +def test_gemini_warns_and_omits_comma_delimited_override(monkeypatch, capsys): + monkeypatch.setenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, "agent,Injected:value") + overlay = gemini.render_env_overlay(WORKSPACE, "gemini-model", "token") + assert CUSTOM_CODING_AGENT_HEADER not in overlay["GEMINI_CLI_CUSTOM_HEADERS"] + assert "Injected" not in overlay["GEMINI_CLI_CUSTOM_HEADERS"] + assert "Gemini CLI" in capsys.readouterr().err + + +@pytest.mark.parametrize("identity", ["COLLIDING_ENV_VAR", "!do-not-execute"]) +def test_pi_uses_environment_reference_not_a_dynamic_literal(identity, monkeypatch): + monkeypatch.setenv(UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR, identity) + monkeypatch.setenv("COLLIDING_ENV_VAR", "unrelated-value") + overlay, _ = pi.render_overlay( + "claude-model", "token", {"claude": WORKSPACE}, {"sonnet": "claude-model"}, [], [] + ) + header = overlay["providers"]["databricks-claude"]["headers"][CUSTOM_CODING_AGENT_HEADER] + assert header == UNITY_GATEWAY_CUSTOM_CODING_AGENT_ENV_VAR + assert pi.build_runtime_env("token")[header] == identity diff --git a/tests/test_desktop_setup.py b/tests/test_desktop_setup.py index 227448b31..7b763efd8 100644 --- a/tests/test_desktop_setup.py +++ b/tests/test_desktop_setup.py @@ -61,6 +61,38 @@ def configured_profile(directory): return json.loads((directory / f"{metadata['appliedId']}.json").read_text()) +def test_custom_agent_header_is_removed_on_reconfigure(desktop, monkeypatch): + state = {"workspace": "https://workspace.example", "profile": "chosen-profile"} + monkeypatch.setenv("UNITY_GATEWAY_CUSTOM_CODING_AGENT", "custom-agent") + desktop_setup.configure_claude_desktop(state) + assert ( + configured_profile(desktop)["inferenceCustomHeaders"]["Databricks-Custom-Coding-Agent"] + == "custom-agent" + ) + + monkeypatch.delenv("UNITY_GATEWAY_CUSTOM_CODING_AGENT") + desktop_setup.configure_claude_desktop(state) + assert ( + "Databricks-Custom-Coding-Agent" + not in configured_profile(desktop)["inferenceCustomHeaders"] + ) + + +def test_admin_custom_agent_header_overrides_environment(desktop, monkeypatch): + monkeypatch.setenv("UNITY_GATEWAY_CUSTOM_CODING_AGENT", "session-agent") + desktop_setup.configure_claude_desktop( + { + "workspace": "https://workspace.example", + "profile": "chosen-profile", + "claude_http_headers": {"databricks-custom-coding-agent": "admin-agent"}, + } + ) + headers = configured_profile(desktop)["inferenceCustomHeaders"] + assert [ + value for name, value in headers.items() if name.lower() == "databricks-custom-coding-agent" + ] == ["admin-agent"] + + def test_configure_uses_claude_catalog_and_absolute_helper(desktop, tmp_path, capsys): desktop_setup.configure_claude_desktop( {"workspace": "https://workspace.example/", "profile": "chosen-profile"}