Add agent harness state (#79)

Author: AarushiShah-db

src/ucode/databricks.py

@@ -9,6 +9,7 @@
 import os
 import platform
 import re
+import shlex
 import shutil
 import subprocess
 from typing import cast
@@ -527,9 +528,16 @@ def list_databricks_apps(workspace: str) -> list[dict]:
 
 
 def build_auth_shell_command(workspace: str) -> str:
+    workspace_arg = shlex.quote(workspace.rstrip("/"))
+    cli_command = (
+        "env -u DATABRICKS_CONFIG_PROFILE "
+        f"databricks auth token --host {workspace_arg} --force-refresh --output json "
+        "| jq -r '.access_token'"
+    )
     return (
-        f"databricks auth token --host {workspace} --force-refresh --output json "
-        f"| jq -r '.access_token'"
+        'if [ -n "${DATABRICKS_BEARER:-}" ]; then '
+        'printf "%s\\n" "$DATABRICKS_BEARER"; '
+        f"else {cli_command}; fi"
     )
 
 

src/ucode/state.py

@@ -5,10 +5,12 @@
 import json
 
 from ucode.config_io import APP_DIR, is_dry_run
-from ucode.databricks import build_shared_base_urls
+from ucode.databricks import build_auth_shell_command, build_shared_base_urls
 
 STATE_PATH = APP_DIR / "state.json"
 STATE_VERSION = 3
+AUTH_COMMAND_TIMEOUT_MS = 5000
+AUTH_REFRESH_INTERVAL_MS = 900_000
 
 
 def load_full_state() -> dict:
@@ -52,6 +54,12 @@ def save_state(state: dict) -> None:
 
 
 def hydrate_state(state: dict) -> dict:
+    """Normalize a workspace state entry and add derived harness config.
+
+    :param state: Raw workspace state entry from ``state.json``.
+    :returns: Hydrated workspace state with stable ``managed_configs``,
+        ``base_urls``, and per-agent ``agents`` entries.
+    """
     if not isinstance(state, dict):
         return {}
 
@@ -71,12 +79,81 @@ def hydrate_state(state: dict) -> dict:
     workspace = hydrated.get("workspace")
     if workspace:
         hydrated["base_urls"] = build_shared_base_urls(workspace)
+        hydrated["agents"] = build_agent_state(hydrated)
     else:
         hydrated["base_urls"] = {}
+        hydrated["agents"] = {}
 
     return hydrated
 
 
+def build_agent_state(state: dict) -> dict[str, dict]:
+    """Build per-agent harness configuration for a workspace.
+
+    The returned shape is intended for downstream tools that want to reuse
+    ucode's configured gateway URLs and auth command without duplicating
+    endpoint construction logic.
+
+    :param state: Hydrated workspace state containing ``workspace``,
+        ``base_urls``, and discovered model lists.
+    :returns: Mapping from agent name to its reusable configuration.
+    """
+    workspace = state.get("workspace")
+    if not isinstance(workspace, str) or not workspace:
+        return {}
+
+    base_urls_value = state.get("base_urls")
+    base_urls = base_urls_value if isinstance(base_urls_value, dict) else {}
+    auth_command = build_auth_shell_command(workspace)
+    claude_models_value = state.get("claude_models")
+    claude_models: dict = claude_models_value if isinstance(claude_models_value, dict) else {}
+    codex_models_value = state.get("codex_models")
+    codex_models = codex_models_value if isinstance(codex_models_value, list) else []
+    gemini_models_value = state.get("gemini_models")
+    gemini_models = gemini_models_value if isinstance(gemini_models_value, list) else []
+
+    claude_model = (
+        claude_models.get("opus") or claude_models.get("sonnet") or claude_models.get("haiku")
+    )
+    codex_model = codex_models[0] if codex_models else None
+    pi_model = claude_model or codex_model or (gemini_models[0] if gemini_models else None)
+
+    agents: dict[str, dict] = {
+        "claude": {
+            "model": claude_model,
+            "base_url": base_urls.get("claude"),
+            "auth_command": auth_command,
+            "auth_refresh_interval_ms": AUTH_REFRESH_INTERVAL_MS,
+            "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",
+            },
+        },
+        "codex": {
+            "model": codex_model,
+            "base_url": base_urls.get("codex"),
+            "auth_command": auth_command,
+            "auth": {
+                "command": "sh",
+                "args": ["-c", auth_command],
+                "timeout_ms": AUTH_COMMAND_TIMEOUT_MS,
+                "refresh_interval_ms": AUTH_REFRESH_INTERVAL_MS,
+            },
+        },
+        "pi": {
+            "model": pi_model,
+            "base_urls": base_urls.get("pi") if isinstance(base_urls.get("pi"), dict) else {},
+            "auth_command": auth_command,
+            "auth_refresh_interval_ms": AUTH_REFRESH_INTERVAL_MS,
+        },
+    }
+    return {
+        name: {key: value for key, value in config.items() if value is not None}
+        for name, config in agents.items()
+    }
+
+
 def clear_state() -> None:
     """Remove the current workspace entry from state."""
     full = load_full_state()

tests/test_databricks.py

@@ -105,6 +105,8 @@ def test_parses_access_token(self):
         assert "jq" in cmd
         assert ".access_token" in cmd
         assert "--force-refresh" in cmd
+        assert "DATABRICKS_BEARER" in cmd
+        assert "DATABRICKS_CONFIG_PROFILE" in cmd
 
     def test_returns_token_when_auth_succeeds(self, tmp_path):
         # Fake databricks binary that always returns a valid token JSON.
@@ -118,10 +120,31 @@ def test_returns_token_when_auth_succeeds(self, tmp_path):
             ["sh", "-c", cmd],
             capture_output=True,
             text=True,
-            env={**os.environ, "PATH": f"{tmp_path}:{os.environ['PATH']}"},
+            env={
+                **os.environ,
+                "PATH": f"{tmp_path}:{os.environ['PATH']}",
+                "DATABRICKS_BEARER": "",
+            },
         )
         assert result.stdout.strip() == "good-token"
 
+    def test_prefers_databricks_bearer(self, tmp_path):
+        fake = tmp_path / "databricks"
+        fake.write_text("#!/bin/sh\nexit 1\n")
+        fake.chmod(0o755)
+        cmd = build_auth_shell_command(WS)
+        result = subprocess.run(
+            ["sh", "-c", cmd],
+            capture_output=True,
+            text=True,
+            env={
+                **os.environ,
+                "PATH": f"{tmp_path}:{os.environ['PATH']}",
+                "DATABRICKS_BEARER": "bearer-token",
+            },
+        )
+        assert result.stdout.strip() == "bearer-token"
+
 
 class TestGetDatabricksToken:
     def _fake_databricks(self, tmp_path, script: str) -> dict:

tests/test_state.py

@@ -10,6 +10,7 @@
 import ucode.state as state_mod
 from ucode.state import (
     STATE_VERSION,
+    build_agent_state,
     clear_state,
     hydrate_state,
     load_full_state,
@@ -27,6 +28,12 @@
         "anthropic": f"{FAKE_WS}/ai-gateway/anthropic/v1",
         "gemini": f"{FAKE_WS}/ai-gateway/gemini/v1beta",
     },
+    "copilot": f"{FAKE_WS}/ai-gateway/mlflow/v1",
+    "pi": {
+        "claude": f"{FAKE_WS}/ai-gateway/anthropic",
+        "openai": f"{FAKE_WS}/ai-gateway/codex/v1",
+        "gemini": f"{FAKE_WS}/ai-gateway/gemini/v1beta",
+    },
 }
 
 
@@ -140,7 +147,7 @@ def test_clear_when_no_state_is_noop(self):
 class TestHydrateState:
     def test_empty_input_returns_empty(self):
         result = hydrate_state({})
-        assert result == {"managed_configs": {}, "base_urls": {}}
+        assert result == {"managed_configs": {}, "base_urls": {}, "agents": {}}
 
     def test_non_dict_returns_empty(self):
         assert hydrate_state(None) == {}  # type: ignore[arg-type]
@@ -153,6 +160,28 @@ def test_populates_base_urls_when_workspace_present(self):
     def test_no_base_urls_when_no_workspace(self):
         result = hydrate_state({"claude_models": {}})
         assert result["base_urls"] == {}
+        assert result["agents"] == {}
+
+    def test_populates_agent_state_when_workspace_present(self):
+        result = hydrate_state(
+            {
+                "workspace": FAKE_WS,
+                "claude_models": {"opus": "claude-opus"},
+                "codex_models": ["gpt-5"],
+            }
+        )
+
+        assert result["agents"]["claude"]["model"] == "claude-opus"
+        assert result["agents"]["claude"]["base_url"] == FAKE_URLS["claude"]
+        assert result["agents"]["claude"]["auth_command"].startswith("if [ -n")
+        assert result["agents"]["codex"]["model"] == "gpt-5"
+        assert result["agents"]["codex"]["base_url"] == FAKE_URLS["codex"]
+        assert (
+            result["agents"]["codex"]["auth"]["args"][1]
+            == result["agents"]["codex"]["auth_command"]
+        )
+        assert result["agents"]["pi"]["model"] == "claude-opus"
+        assert result["agents"]["pi"]["base_urls"] == FAKE_URLS["pi"]
 
     def test_normalizes_managed_configs_dict_entry(self):
         state = {"managed_configs": {"claude": {"keys": [["env", "X"]]}}}
@@ -171,6 +200,12 @@ def test_drops_falsy_managed_configs(self):
         assert "claude" not in result["managed_configs"]
 
 
+class TestBuildAgentState:
+    def test_returns_empty_without_workspace(self):
+        result = build_agent_state({"base_urls": FAKE_URLS})
+        assert result == {}
+
+
 # ---------------------------------------------------------------------------
 # mark_tool_managed
 # ---------------------------------------------------------------------------