Add UNITY_GATEWAY_CUSTOM_CODING_AGENT header override

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

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:

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:

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 {}

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).

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 []

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

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()`)

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

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:

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

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/`.

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"

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"),
     ],