diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 01b3a1f9..2da60c36 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -29,6 +29,7 @@ build_auth_shell_command, build_tool_base_url, get_databricks_token, + state_oauth_client_id, ) from ucode.launcher import exec_or_spawn from ucode.managed_files import ( @@ -360,6 +361,7 @@ def render_overlay( relayed_base_url: str | None = None, route_root_model: str | None = None, custom_model: str | None = None, + oauth_client_id: str | None = None, ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for Claude settings.json. @@ -477,7 +479,9 @@ def render_overlay( if relayed: keys = [["env", k] for k in env] else: - overlay["apiKeyHelper"] = build_auth_shell_command(workspace, profile, use_pat=use_pat) + overlay["apiKeyHelper"] = build_auth_shell_command( + workspace, profile, use_pat=use_pat, oauth_client_id=oauth_client_id + ) keys = [["apiKeyHelper"]] + [["env", k] for k in env] # Disable Claude Code's built-in WebSearch: it declares Anthropic's hosted @@ -641,6 +645,7 @@ def write_tool_config( relayed_base_url=relayed_base_url, route_root_model=route_root_model, custom_model=custom_model, + oauth_client_id=state_oauth_client_id(state), ) tracing_env_vars = tracing_env(state, "claude") stop_hook_command = claude_tracing_stop_hook_command() if tracing_env_vars else None diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 625329f3..e4df8c17 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -27,6 +27,7 @@ build_auth_token_argv, build_tool_base_url, get_databricks_token, + state_oauth_client_id, ) from ucode.launcher import exec_or_spawn from ucode.managed_files import ( @@ -141,8 +142,11 @@ def _provider_block( databricks_profile: str | None, use_pat: bool = False, provider: str | None = None, + oauth_client_id: str | None = None, ) -> dict: - auth_argv = build_auth_token_argv(workspace, databricks_profile, use_pat=use_pat) + auth_argv = build_auth_token_argv( + workspace, databricks_profile, use_pat=use_pat, oauth_client_id=oauth_client_id + ) base_url = build_tool_base_url("codex", workspace) http_headers = { "User-Agent": f"ucode/{ucode_version()} codex/{agent_version('codex')}", @@ -173,13 +177,14 @@ def render_overlay( databricks_profile: str | None = None, use_pat: bool = False, provider: str | None = None, + oauth_client_id: str | None = None, ) -> dict: overlay: dict = {"model_provider": CODEX_MODEL_PROVIDER_NAME} if model: overlay["model"] = model overlay["model_providers"] = { CODEX_MODEL_PROVIDER_NAME: _provider_block( - workspace, databricks_profile, use_pat, provider + workspace, databricks_profile, use_pat, provider, oauth_client_id ), } return overlay @@ -191,6 +196,7 @@ def render_legacy_overlay( databricks_profile: str | None = None, use_pat: bool = False, provider: str | None = None, + oauth_client_id: str | None = None, ) -> dict: """Overlay for Codex CLI < 0.134.0, which only reads `~/.codex/config.toml`. @@ -205,7 +211,7 @@ def render_legacy_overlay( "profiles": {CODEX_PROFILE_NAME: profile_block}, "model_providers": { CODEX_MODEL_PROVIDER_NAME: _provider_block( - workspace, databricks_profile, use_pat, provider + workspace, databricks_profile, use_pat, provider, oauth_client_id ), }, } @@ -320,6 +326,7 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non databricks_profile, use_pat=bool(state.get("use_pat")), provider=provider, + oauth_client_id=state_oauth_client_id(state), ) doc = read_toml_safe(LEGACY_CODEX_CONFIG_PATH) deep_merge_dict(doc, overlay) @@ -345,6 +352,7 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non databricks_profile, use_pat=bool(state.get("use_pat")), provider=provider, + oauth_client_id=state_oauth_client_id(state), ) def compose(base: dict) -> dict: diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index b07aef56..b38735ed 100644 --- a/src/ucode/agents/opencode.py +++ b/src/ucode/agents/opencode.py @@ -23,6 +23,7 @@ build_opencode_base_urls, get_databricks_token, model_token_limits, + state_oauth_client_id, ) from ucode.state import mark_tool_managed, save_state from ucode.telemetry import agent_version, ucode_version @@ -174,6 +175,7 @@ def render_auth_plugin(state: dict) -> str: state["workspace"], state.get("profile"), use_pat=bool(state.get("use_pat")), + oauth_client_id=state_oauth_client_id(state), ) # A 401 must not return the same still-unexpired cached credential. argv.append("--force-refresh") diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6650974d..05108f85 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -42,6 +42,7 @@ from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH from ucode.config_io import is_dry_run, restore_file, set_dry_run from ucode.databricks import ( + OAUTH_CLIENT_ID_STATE_KEY, apply_pat_environment, build_shared_base_urls, discover_claude_models, @@ -520,6 +521,7 @@ def configure_shared_state( skip_preflight: bool = False, fable_enabled: bool | None = None, databricks_ai_tools_enabled: bool | None = None, + oauth_client_id: str | None = None, ) -> dict: """Log into Databricks, verify AI Gateway, fetch model lists, persist state. @@ -543,12 +545,19 @@ def configure_shared_state( ``ANTHROPIC_DEFAULT_FABLE_MODEL`` pin (default off). ``None`` means "inherit": a launch re-run keeps whatever the workspace was configured with; ``True``/ ``False`` come from an explicit ``configure --enable-fable``/``--disable-fable``. + ``oauth_client_id`` authenticates through a custom OAuth app instead of the + built-in ``databricks-cli`` one, whose refresh token expires after 7 days. + ``None`` means "inherit"; an empty string clears a prior one. It is persisted + before the login below runs, and baked into every generated agent config, so + the agent's own token helper mints from the same app. """ workspace = normalize_workspace_url(workspace) prior_state = load_state() previous_workspace = prior_state.get("workspace") if use_pat is None: use_pat = bool(prior_state.get("use_pat")) and previous_workspace == workspace + if oauth_client_id is None and previous_workspace == workspace: + oauth_client_id = prior_state.get(OAUTH_CLIENT_ID_STATE_KEY) if fable_enabled is None: fable_enabled = bool(prior_state.get("fable_enabled")) and previous_workspace == workspace if databricks_ai_tools_enabled is None: @@ -582,6 +591,13 @@ def configure_shared_state( state["use_pat"] = True else: state.pop("use_pat", None) + # Persist the custom OAuth app before the login below: `run_databricks_login` + # and every later token mint resolve the client id from this entry, and the + # generated agent configs pin it from here. + if oauth_client_id: + state[OAUTH_CLIENT_ID_STATE_KEY] = oauth_client_id + else: + state.pop(OAUTH_CLIENT_ID_STATE_KEY, None) # Persist the Fable opt-in so launches keep pinning the family; an explicit # `configure --disable-fable` (fable_enabled=False) clears it. if fable_enabled: @@ -627,11 +643,11 @@ def configure_shared_state( # empty one as absent, so it never shadows the PAT. Pass the validated # token to avoid re-reading ~/.databrickscfg. ensure_pat_bearer(profile, pat) - ensure_databricks_auth(workspace, profile) + ensure_databricks_auth(workspace, profile, oauth_client_id=oauth_client_id) elif force_login: - run_databricks_login(workspace, profile) + run_databricks_login(workspace, profile, oauth_client_id=oauth_client_id) else: - ensure_databricks_auth(workspace, profile) + ensure_databricks_auth(workspace, profile, oauth_client_id=oauth_client_id) # After login the profile exists in ~/.databrickscfg, so a host->profile # lookup is reliable even when it returned nothing above. if profile is None: @@ -639,7 +655,7 @@ def configure_shared_state( if profile: state["profile"] = profile with spinner("Verifying Unity AI Gateway..."): - token = get_databricks_token(workspace, profile) + token = get_databricks_token(workspace, profile, oauth_client_id=oauth_client_id) model_service_probe = probe_unity_gateway_capabilities(workspace, token) if model_service_probe.resource_available: print_success("Unity AI Gateway connected") @@ -753,6 +769,7 @@ def _configure_shared_workspace_states( use_pat: bool = False, fable_enabled: bool | None = None, databricks_ai_tools_enabled: bool | None = None, + oauth_client_id: str | None = None, ) -> list[dict]: if not workspaces: raise RuntimeError("At least one workspace must be provided.") @@ -767,6 +784,7 @@ def _configure_shared_workspace_states( use_pat=use_pat, fable_enabled=fable_enabled, databricks_ai_tools_enabled=databricks_ai_tools_enabled, + oauth_client_id=oauth_client_id, ) ) return states @@ -849,6 +867,7 @@ def configure_workspace_command( skip_unavailable: bool = False, fable_enabled: bool | None = None, databricks_ai_tools_enabled: bool | None = None, + oauth_client_id: str | None = None, offer_optional_setup: bool = False, ) -> int: if tool is not None and selected_tools is not None: @@ -869,6 +888,7 @@ def configure_workspace_command( use_pat=use_pat, fable_enabled=fable_enabled, databricks_ai_tools_enabled=databricks_ai_tools_enabled, + oauth_client_id=oauth_client_id, ) state = states[0] state = configure_single_tool(tool, state) @@ -908,6 +928,7 @@ def configure_workspace_command( use_pat=use_pat, fable_enabled=fable_enabled, databricks_ai_tools_enabled=databricks_ai_tools_enabled, + oauth_client_id=oauth_client_id, ) state = states[0] save_state(state) @@ -1743,14 +1764,16 @@ def claude_router_hook_cmd( sys.stdout.write(json.dumps(output)) -def _auto_configure_tool(tool: str) -> None: +def _auto_configure_tool(tool: str, oauth_client_id: str | None = None) -> None: """First-time setup for a single tool — mirrors configure_workspace_command.""" existing = load_state() workspace = existing.get("workspace") profile = existing.get("profile") if not workspace: workspace, profile = _prompt_for_configuration(tool) - state = configure_shared_state(workspace, profile=profile, tools=[tool]) + state = configure_shared_state( + workspace, profile=profile, tools=[tool], oauth_client_id=oauth_client_id + ) state = configure_single_tool(tool, state) @@ -2011,6 +2034,7 @@ def _can_launch_from_cached_config( model: str | None, explicit_provider: str | None, workspace_url: str | None, + oauth_client_id: str | None = None, ) -> bool: """Return whether a normal Claude/Codex launch can use its cached config.""" if tool not in CAN_USE_CACHED_CONFIG_AGENTS: @@ -2019,6 +2043,13 @@ def _can_launch_from_cached_config( if refresh or model or explicit_provider is not None: return False + # A `--oauth-client-id` that disagrees with what the cached config was written against would + # launch the agent with a token helper pointing at the *other* OAuth app. Reconfigure instead. + if oauth_client_id is not None and oauth_client_id != ( + state.get(OAUTH_CLIENT_ID_STATE_KEY) or "" + ): + return False + if tool == "codex" and smart_routing_v2.enabled(): if not state.get("codex_models") or not state.get("oss_models"): return False @@ -2054,6 +2085,7 @@ def _launch_tool( managed: dict | None = None, recommendation: dict | None = None, model: str | None = None, + oauth_client_id: str | None = None, ) -> None: try: tool = normalize_tool(tool_name) @@ -2081,7 +2113,7 @@ def _launch_tool( ) ensure_bootstrap_dependencies(tool, update_existing=needs_auto_configure) if needs_auto_configure: - _auto_configure_tool(tool) + _auto_configure_tool(tool, oauth_client_id=oauth_client_id) state = ensure_provider_state(tool) # Remembered before the fallback below collapses the two cases: a managed config may not # silently override a provider the user typed on the command line (it errors instead). @@ -2097,6 +2129,7 @@ def _launch_tool( model=model, explicit_provider=explicit_provider, workspace_url=workspace_url, + oauth_client_id=oauth_client_id, ): print_section(_launch_title(tool)) if forwarded_model: @@ -2125,6 +2158,7 @@ def _launch_tool( tools=[tool], skip_model_discovery=bool(provider) or managed_models_known, skip_preflight=skip_preflight, + oauth_client_id=oauth_client_id, ) # An admin-published managed config wins over the developer's own settings. Layered on after # `configure_shared_state`, whose returned state it overrides, and before the provider and @@ -2378,6 +2412,19 @@ def _disable_managed_config_if_requested(skip_managed_config: bool) -> None: ), ] +# Sign in through a custom OAuth app rather than the built-in `databricks-cli` one, whose refresh +# token lasts 7 days. Persisted by the launch's configure pass and baked into the agent's own token +# helper, so the agent keeps minting from the same app. +OauthClientIdOption = Annotated[ + str | None, + typer.Option( + "--oauth-client-id", + help="Authenticate with this custom OAuth app instead of the built-in `databricks-cli` " + "app, and remember it for this workspace. Pass an empty string to go back to the " + "built-in app.", + ), +] + @app.callback(invoke_without_command=True) def default( @@ -2403,6 +2450,7 @@ def default( skip_preflight: SkipPreflightOption = False, skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, + oauth_client_id: OauthClientIdOption = None, ) -> None: """Configure and launch coding agents through Databricks AI Gateway. @@ -2416,7 +2464,11 @@ def default( _disable_managed_config_if_requested(skip_managed_config) try: _launch_managed_default( - ctx, dry_run=dry_run, skip_preflight=skip_preflight, workspace=workspace + ctx, + dry_run=dry_run, + skip_preflight=skip_preflight, + workspace=workspace, + oauth_client_id=oauth_client_id, ) except typer.Exit: # `typer.Exit` subclasses RuntimeError, so it has to be re-raised ahead of the handler @@ -2434,6 +2486,7 @@ def _launch_managed_default( dry_run: bool, skip_preflight: bool, workspace: str | None, + oauth_client_id: str | None = None, ) -> None: """Route bare ``ucode`` by whether the workspace publishes a managed config.""" if not managed_agent_config_enabled(): @@ -2476,6 +2529,7 @@ def _launch_managed_default( workspace_url=workspace, managed=managed, recommendation=recommendation, + oauth_client_id=oauth_client_id, ) @@ -2519,6 +2573,7 @@ def codex_cmd( skip_preflight: SkipPreflightOption = False, skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, + oauth_client_id: OauthClientIdOption = None, enable_smart_routing_flag: Annotated[ bool, typer.Option( @@ -2552,6 +2607,7 @@ def codex_cmd( refresh=refresh, skip_preflight=skip_preflight, workspace_url=workspace, + oauth_client_id=oauth_client_id, ) @@ -2588,6 +2644,7 @@ def claude_cmd( skip_preflight: SkipPreflightOption = False, skip_managed_config: SkipManagedConfigOption = False, workspace: WorkspaceOption = None, + oauth_client_id: OauthClientIdOption = None, enable_model_discovery: Annotated[ bool, typer.Option( @@ -2632,6 +2689,7 @@ def claude_cmd( refresh=refresh, skip_preflight=skip_preflight, workspace_url=workspace, + oauth_client_id=oauth_client_id, ) @@ -2778,6 +2836,16 @@ def configure( "CI / headless environments.", ), ] = False, + oauth_client_id: Annotated[ + str | None, + typer.Option( + "--oauth-client-id", + help="Authenticate with this custom OAuth app instead of the built-in " + "`databricks-cli` app, whose refresh token expires after 7 days. The id is " + "remembered for the workspace(s) and baked into the configured agents' token " + "helpers. Pass an empty string to go back to the built-in app.", + ), + ] = None, skip_validate: Annotated[ bool, typer.Option( @@ -2904,6 +2972,11 @@ def configure( skip_kwargs["use_pat"] = True if skip_validate: skip_kwargs["skip_validate"] = True + # Only when the flag was passed: `None` lets configure_shared_state inherit the + # workspace's saved client id rather than clearing it. An explicit `--oauth-client-id ""` + # is forwarded and clears it, going back to the built-in `databricks-cli` app. + if oauth_client_id is not None: + skip_kwargs["oauth_client_id"] = oauth_client_id # Only forward the Fable opt-in when the user passed the flag; `None` # (neither flag given) lets configure_shared_state inherit the prior # workspace setting instead of clobbering it. @@ -2972,6 +3045,7 @@ def configure( tools=[], force_login=not use_pat, use_pat=use_pat, + oauth_client_id=oauth_client_id, ) else: # Neither model agents nor cursor -> empty/invalid --agents list. @@ -2988,6 +3062,7 @@ def configure( tools=[], force_login=not use_pat, use_pat=use_pat, + oauth_client_id=oauth_client_id, ) else: # Tool binaries are installed after the user picks which agents diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index ea090cfb..1ff2bae2 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -942,6 +942,17 @@ def resolve_oauth_client_id(workspace: str, explicit: str | None = None) -> str return value.strip() if isinstance(value, str) and value.strip() else None +def state_oauth_client_id(state: dict) -> str | None: + """Return the custom OAuth client id recorded in a hydrated ``state`` dict, if any. + + For the config writers, which already hold the workspace's state and must *pin* the id into + the file they render (an agent runs its token helper as a bare command, so the app has to be + named in the command itself). ``resolve_oauth_client_id`` stays the entry point for code that + only has a host.""" + value = state.get(OAUTH_CLIENT_ID_STATE_KEY) + return value.strip() if isinstance(value, str) and value.strip() else None + + def resolve_oauth_account_id(workspace: str, explicit: str | None = None) -> str | None: """Resolve the account id for an account-level custom OAuth app, if any. @@ -956,7 +967,9 @@ def resolve_oauth_account_id(workspace: str, explicit: str | None = None) -> str return value.strip() if isinstance(value, str) and value.strip() else None -def has_valid_databricks_auth(workspace: str, profile: str | None = None) -> bool: +def has_valid_databricks_auth( + workspace: str, profile: str | None = None, *, oauth_client_id: str | None = None +) -> bool: # Honor the CI short-circuit (see ``get_databricks_token``): if a # pre-fetched bearer is available, treat auth as valid and skip the # `databricks auth token` shell-out (which only knows user-OAuth). @@ -965,7 +978,7 @@ def has_valid_databricks_auth(workspace: str, profile: str | None = None) -> boo # A workspace configured with a custom OAuth app keeps its own token cache; # `databricks auth token` knows nothing about it, and would answer for the # built-in client's session instead. - client_id = resolve_oauth_client_id(workspace) + client_id = resolve_oauth_client_id(workspace, oauth_client_id) if client_id: from ucode import oauth @@ -1158,7 +1171,9 @@ def apply_pat_environment(state: dict) -> None: ensure_pat_bearer(state.get("profile")) -def run_databricks_login(workspace: str, profile: str | None = None) -> None: +def run_databricks_login( + workspace: str, profile: str | None = None, *, oauth_client_id: str | None = None +) -> None: """Run databricks auth login unconditionally. When ``profile`` is provided, it is passed via ``--profile``. Otherwise we @@ -1166,8 +1181,10 @@ def run_databricks_login(workspace: str, profile: str | None = None) -> None: refreshed in place rather than overwriting another profile's tokens. A workspace configured with a custom OAuth app takes the PKCE flow instead: - `databricks auth login` has no way to select a client id.""" - client_id = resolve_oauth_client_id(workspace) + `databricks auth login` has no way to select a client id. ``oauth_client_id`` + is the id being configured *now* — it is passed explicitly because `ug + configure` has not persisted it yet at this point.""" + client_id = resolve_oauth_client_id(workspace, oauth_client_id) if client_id: from ucode import oauth @@ -1195,21 +1212,31 @@ def run_databricks_login(workspace: str, profile: str | None = None) -> None: def ensure_databricks_auth( - workspace: str, profile: str | None = None, *, quiet: bool = False + workspace: str, + profile: str | None = None, + *, + quiet: bool = False, + oauth_client_id: str | None = None, ) -> None: """Check auth and login only if needed (used by launch path). ``quiet`` suppresses the "already available" line for a caller that only needs a token before some later step re-authenticates and reports it — otherwise the same success prints twice. A login that actually runs is never silent. + + ``oauth_client_id`` pins a custom OAuth app for both halves: the check has to + ask about the *same* app the login would sign into, or a workspace with a live + built-in session would report "already available" and never sign in. """ with spinner("Checking Databricks auth..."): - auth_is_valid = has_valid_databricks_auth(workspace, profile) + auth_is_valid = has_valid_databricks_auth( + workspace, profile, oauth_client_id=oauth_client_id + ) if auth_is_valid: if not quiet: print_success(f"Databricks auth already available for {workspace}") return - run_databricks_login(workspace, profile) + run_databricks_login(workspace, profile, oauth_client_id=oauth_client_id) def get_databricks_token( diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index f855d086..d4ce680a 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -42,6 +42,7 @@ list_mcp_services, list_uc_functions_catalog_schemas, list_vector_search_catalog_schemas, + state_oauth_client_id, workspace_hostname, ) from ucode.state import load_full_state, load_state, save_state @@ -300,13 +301,16 @@ def configure_client_mcp_server( *, use_pat: bool = False, always_load: bool = False, + oauth_client_id: str | None = None, ) -> list[str]: # Every client registers the same `ucode mcp-proxy ...` stdio command; the # proxy forwards to `url` and refreshes the Databricks token itself. Only the # per-client registration syntax differs. `always_load` (skills registry) is # a Claude-only hint to load the server's tools at session start; other # clients don't support it and ignore it. - argv = build_mcp_proxy_argv(url, workspace, profile, use_pat=use_pat) + argv = build_mcp_proxy_argv( + url, workspace, profile, use_pat=use_pat, oauth_client_id=oauth_client_id + ) if client == "claude": removed_scopes = [ scope for scope in MCP_CLEANUP_SCOPES if remove_claude_mcp_server(name, scope) @@ -1168,7 +1172,15 @@ def apply_managed_mcp_servers( for server in (state.get("managed_mcp_servers") or []) if isinstance(server, dict) and tool in (server.get("clients") or []) ] - apply_mcp_server_changes(previous, working, [tool], workspace, profile, use_pat=use_pat) + apply_mcp_server_changes( + previous, + working, + [tool], + workspace, + profile, + use_pat=use_pat, + oauth_client_id=state_oauth_client_id(state), + ) return working @@ -1219,7 +1231,13 @@ def apply_managed_skills( original = list(state.get("mcp_servers") or []) working = _resolve_skills_mcp_servers(workspace, [tool], new_locations, original) changed = apply_mcp_server_changes( - original, working, [tool], workspace, profile, use_pat=use_pat + original, + working, + [tool], + workspace, + profile, + use_pat=use_pat, + oauth_client_id=state_oauth_client_id(state), ) if not (changed or original != working or prev_managed != desired): return [] @@ -1458,6 +1476,7 @@ def apply_mcp_server_changes( profile: str | None = None, *, use_pat: bool = False, + oauth_client_id: str | None = None, ) -> bool: original_by_name = _servers_by_name(original_servers) working_by_name = _servers_by_name(working_servers) @@ -1492,7 +1511,14 @@ def apply_mcp_server_changes( for client in clients: work[client].append( lambda c=client, n=name, u=url, al=always_load: configure_client_mcp_server( - c, n, u, workspace, profile, use_pat=use_pat, always_load=al + c, + n, + u, + workspace, + profile, + use_pat=use_pat, + always_load=al, + oauth_client_id=oauth_client_id, ) ) changed = True @@ -1869,6 +1895,7 @@ def configure_mcp_command( workspace, profile, use_pat=bool(state.get("use_pat")), + oauth_client_id=state_oauth_client_id(state), ) if changed or original_mcp_servers_for_location != working_mcp_servers: state["mcp_servers"] = working_mcp_servers @@ -1966,6 +1993,7 @@ def configure_mcp_command( workspace, profile, use_pat=bool(state.get("use_pat")), + oauth_client_id=state_oauth_client_id(state), ) if changed or original_mcp_servers != working_mcp_servers: state["mcp_servers"] = working_mcp_servers @@ -2074,7 +2102,13 @@ def remove_mcp_command(agents: set[str] | None = None) -> int: if targets: removal_view.append({**server, "clients": targets}) changed = apply_mcp_server_changes( - removal_view, [], clients, workspace, profile, use_pat=bool(state.get("use_pat")) + removal_view, + [], + clients, + workspace, + profile, + use_pat=bool(state.get("use_pat")), + oauth_client_id=state_oauth_client_id(state), ) # Update saved state: drop a fully-removed server, or keep it with the named @@ -2178,7 +2212,9 @@ def _update_skills_mcp( """Rebuild the single skills connection for ``locations`` and persist it.""" original = list(state.get("mcp_servers") or []) working = _resolve_skills_mcp_servers(workspace, clients, locations, original) - changed = apply_mcp_server_changes(original, working, clients, workspace, profile) + changed = apply_mcp_server_changes( + original, working, clients, workspace, profile, oauth_client_id=state_oauth_client_id(state) + ) if changed or original != working: state["mcp_servers"] = working save_state(state) diff --git a/src/ucode/smart_routing/claude_hooks.py b/src/ucode/smart_routing/claude_hooks.py index 837326c0..7a2adda8 100644 --- a/src/ucode/smart_routing/claude_hooks.py +++ b/src/ucode/smart_routing/claude_hooks.py @@ -11,7 +11,7 @@ import shlex -from ucode.databricks import build_auth_token_argv +from ucode.databricks import build_auth_token_argv, state_oauth_client_id from ucode.smart_routing import hooks ROUTING_HOOK_COMMAND_MARKER = "claude-router-hook" @@ -90,6 +90,11 @@ def _routing_hook_argv(state: dict, event: str) -> list[str]: argv += ["--profile", profile] if state.get("use_pat"): argv.append("--use-pat") + # The hook re-mints a token in its own process, so it needs the workspace's custom OAuth app + # named on the command line just like the agents' token helpers do. + oauth_client_id = state_oauth_client_id(state) + if oauth_client_id: + argv += ["--oauth-client-id", oauth_client_id] # The route-subagent hook resolves the router's chosen arm back to a routable # workspace id, so it needs the discovered claude model ids. claude_models = state.get("claude_models") diff --git a/src/ucode/smart_routing/codex_hooks.py b/src/ucode/smart_routing/codex_hooks.py index 421cabda..af5237ac 100644 --- a/src/ucode/smart_routing/codex_hooks.py +++ b/src/ucode/smart_routing/codex_hooks.py @@ -6,7 +6,7 @@ import shlex import subprocess -from ucode.databricks import build_auth_token_argv +from ucode.databricks import build_auth_token_argv, state_oauth_client_id from ucode.smart_routing import hooks ROUTING_HOOK_COMMAND_MARKER = "codex-router-hook" @@ -96,6 +96,11 @@ def _routing_hook_argv( argv += ["--profile", profile] if state.get("use_pat"): argv.append("--use-pat") + # The hook re-mints a token in its own process, so it needs the workspace's custom OAuth app + # named on the command line just like the agents' token helpers do. + oauth_client_id = state_oauth_client_id(state) + if oauth_client_id: + argv += ["--oauth-client-id", oauth_client_id] models = available_models if available_models is not None else routing_models(state) for model in models: if isinstance(model, str) and model: diff --git a/src/ucode/state.py b/src/ucode/state.py index 6344a7e5..97b62af5 100644 --- a/src/ucode/state.py +++ b/src/ucode/state.py @@ -9,6 +9,7 @@ build_auth_shell_command, build_auth_token_argv, build_shared_base_urls, + state_oauth_client_id, ) STATE_PATH = APP_DIR / "state.json" @@ -161,8 +162,15 @@ def build_agent_state(state: dict) -> dict[str, dict]: base_urls_value = state.get("base_urls") base_urls = base_urls_value if isinstance(base_urls_value, dict) else {} use_pat = bool(state.get("use_pat")) - auth_command = build_auth_shell_command(workspace, profile, use_pat=use_pat) - auth_argv = build_auth_token_argv(workspace, profile, use_pat=use_pat) + # Pinned into the auth command rather than left to run-time resolution: the agent runs it as a + # bare shell command, so the app it mints from has to be named in the command itself. + oauth_client_id = state_oauth_client_id(state) + auth_command = build_auth_shell_command( + workspace, profile, use_pat=use_pat, oauth_client_id=oauth_client_id + ) + auth_argv = build_auth_token_argv( + workspace, profile, use_pat=use_pat, oauth_client_id=oauth_client_id + ) 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") diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 61f07438..36e4a0da 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -236,6 +236,14 @@ def test_sets_api_key_helper(self): assert "apiKeyHelper" in overlay assert WS in overlay["apiKeyHelper"] + def test_api_key_helper_pins_a_custom_oauth_app(self): + overlay, _ = claude.render_overlay(WS, "s4", oauth_client_id="custom-app-id") + assert "--oauth-client-id custom-app-id" in overlay["apiKeyHelper"] + + def test_api_key_helper_omits_the_flag_without_a_custom_app(self): + overlay, _ = claude.render_overlay(WS, "s4") + assert "--oauth-client-id" not in overlay["apiKeyHelper"] + def test_relayed_omits_api_key_helper(self): # Claude Code's own subscription OAuth must own Authorization; an # apiKeyHelper would outrank it. @@ -762,6 +770,17 @@ def test_managed_file_preserves_other_keys(self, monkeypatch): assert written["env"]["ANTHROPIC_BASE_URL"] assert written["apiKeyHelper"] + def test_written_settings_pin_the_workspace_custom_oauth_app(self, monkeypatch): + # write_tool_config reads the id out of state, so a `ug configure --oauth-client-id` + # lands in settings.json without the launch having to pass it again. + private_writes: list = [] + managed_writes: list = [] + self._patch(monkeypatch, private_writes, managed_writes, {}) + state = {"workspace": WS, "codex_models": [], "oauth_client_id": "custom-app-id"} + claude.write_tool_config(state, "databricks-claude-sonnet-4") + _, payload = private_writes[0] + assert "--oauth-client-id custom-app-id" in payload["apiKeyHelper"] + def test_managed_file_strips_stale_gateway_model_discovery(self, monkeypatch): private_writes: list = [] managed_writes: list = [] diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 2cd11dc3..782cdb11 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -86,6 +86,21 @@ def test_auth_contains_workspace(self): auth = overlay["model_providers"]["ucode-databricks"]["auth"] assert any(WS in arg for arg in auth["args"]) + def test_auth_pins_a_custom_oauth_app(self): + overlay = codex.render_overlay(WS, oauth_client_id="custom-app-id") + auth = overlay["model_providers"]["ucode-databricks"]["auth"] + assert auth["args"][-2:] == ["--oauth-client-id", "custom-app-id"] + + def test_auth_omits_the_flag_without_a_custom_app(self): + overlay = codex.render_overlay(WS) + auth = overlay["model_providers"]["ucode-databricks"]["auth"] + assert "--oauth-client-id" not in auth["args"] + + def test_legacy_overlay_pins_a_custom_oauth_app(self): + overlay = codex.render_legacy_overlay(WS, oauth_client_id="custom-app-id") + auth = overlay["model_providers"]["ucode-databricks"]["auth"] + assert auth["args"][-2:] == ["--oauth-client-id", "custom-app-id"] + def test_auth_refresh_interval(self): overlay = codex.render_overlay(WS) auth = overlay["model_providers"]["ucode-databricks"]["auth"] diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 96c162dc..a5866aa5 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -72,7 +72,7 @@ def test_calls_cross_platform_auth_token_helper_only_when_refreshing(self, monke monkeypatch.setattr( opencode, "build_auth_token_argv", - lambda workspace, profile, use_pat=False: [ + lambda workspace, profile, use_pat=False, **_kwargs: [ "/opt/ucode", "auth-token", "--host", @@ -91,6 +91,16 @@ def test_calls_cross_platform_auth_token_helper_only_when_refreshing(self, monke assert "run(AUTH_COMMAND[0], AUTH_COMMAND.slice(1)" in plugin assert '"chat.headers"' not in plugin + def test_pins_the_workspace_custom_oauth_app(self): + plugin = opencode.render_auth_plugin({"workspace": WS, "oauth_client_id": "custom-app-id"}) + + assert '"--oauth-client-id", "custom-app-id"' in plugin + + def test_omits_the_flag_without_a_custom_oauth_app(self): + plugin = opencode.render_auth_plugin({"workspace": WS}) + + assert "--oauth-client-id" not in plugin + def test_installs_cached_refreshing_fetch_on_databricks_providers(self): plugin = opencode.render_auth_plugin({"workspace": WS}) diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index ad0429d9..d22c74fb 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -555,3 +555,26 @@ def read_until(suffix): "replayed": "\x1b[200~fix\nthe parser\x1b[201~\r", "restored_before_replay": True, } + + +class TestRoutingHookCustomOauthApp: + """The route-subagent hook re-mints a token in its own process, so the workspace's + custom OAuth app has to be named on its command line — the same way the agents' + token helpers name it.""" + + STATE = { + "workspace": "https://example.databricks.com", + "claude_models": {"sonnet": "databricks-claude-sonnet-4"}, + } + + def _route_command(self, state: dict) -> str: + doc: dict = {} + claude_hooks.sync_smart_routing_hooks(doc, state, enabled=True) + return doc["hooks"]["PreToolUse"][0]["hooks"][0]["command"] + + def test_flag_is_appended_when_the_workspace_has_a_custom_app(self): + command = self._route_command({**self.STATE, "oauth_client_id": "custom-app-id"}) + assert "--oauth-client-id custom-app-id" in command + + def test_flag_is_absent_without_a_custom_app(self): + assert "--oauth-client-id" not in self._route_command(dict(self.STATE)) diff --git a/tests/test_cli.py b/tests/test_cli.py index bdaa1638..2c81ec66 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -435,6 +435,36 @@ def test_no_workspace_flag_leaves_current_workspace(self): assert result.exit_code == 0, result.output mock_set.assert_not_called() + @pytest.mark.parametrize("tool", ["claude", "codex"]) + @pytest.mark.parametrize( + "args_extra,expected", + [ + (["--oauth-client-id", "custom-app-id"], "custom-app-id"), + # No flag forwards None, which configure_shared_state reads as "inherit" — it must + # not look like an explicit request for the built-in app. + ([], None), + ], + ) + def test_oauth_client_id_flag_reaches_the_configure_pass(self, tool, args_extra, expected): + """`ug --oauth-client-id X` configures and launches against app X.""" + patches = _patch_launch(tool) + with ( + patches[0], + patches[1], + patches[2], + patch("ucode.cli.configure_shared_state", return_value=MINIMAL_STATE) as mock_shared, + patches[4], + patches[5], + patches[6], + patches[7], + # Forced off so the configure pass runs for both agents regardless of what this + # machine has cached (or of ENABLE_SMART_ROUTING_V2 in the environment). + patch("ucode.cli._can_launch_from_cached_config", return_value=False), + ): + result = runner.invoke(app, [tool, *args_extra]) + assert result.exit_code == 0, result.output + assert mock_shared.call_args.kwargs["oauth_client_id"] == expected + def test_codex_enable_smart_routing_is_consumed_by_ucode(self): enabled_during_launch = [] with patch( @@ -1749,7 +1779,7 @@ def test_triggers_when_no_workspace(self): result = runner.invoke(app, ["claude"]) assert result.exit_code == 0, result.output mock_bootstrap.assert_called_once_with("claude", update_existing=True) - mock_auto.assert_called_once_with("claude") + mock_auto.assert_called_once_with("claude", oauth_client_id=None) def test_triggers_when_tool_not_in_available_tools(self): """Auto-configure runs when workspace exists but the tool wasn't configured.""" @@ -1774,7 +1804,7 @@ def test_triggers_when_tool_not_in_available_tools(self): result = runner.invoke(app, ["claude"]) assert result.exit_code == 0, result.output mock_bootstrap.assert_called_once_with("claude", update_existing=True) - mock_auto.assert_called_once_with("claude") + mock_auto.assert_called_once_with("claude", oauth_client_id=None) def test_skipped_when_already_configured(self): """Auto-configure is skipped when workspace and tool are already set up.""" @@ -1954,6 +1984,42 @@ def test_rejects_dynamic_launch_overrides(self, override): is False ) + def test_rejects_an_oauth_client_id_that_differs_from_the_cached_one(self): + # The cached config's token helper names the app it was written against; launching it + # under a different --oauth-client-id would mint from the wrong app. + import ucode.cli as cli_mod + + state = {**MINIMAL_STATE, "oauth_client_id": "configured-app"} + with ( + patch("ucode.cli.managed_agent_config_enabled", return_value=False), + patch("ucode.cli.smart_routing_v2.enabled", return_value=False), + patch("ucode.cli.codex_agent.has_ucode_config", return_value=True), + patch("ucode.cli.codex_agent.managed_config_is_current", return_value=True), + ): + assert ( + cli_mod._can_launch_from_cached_config( + "codex", state, **self._kwargs(oauth_client_id="other-app") + ) + is False + ) + # Repeating the id the config already carries is not a change. + assert ( + cli_mod._can_launch_from_cached_config( + "codex", state, **self._kwargs(oauth_client_id="configured-app") + ) + is True + ) + # No flag at all keeps inheriting the cached config. + assert cli_mod._can_launch_from_cached_config("codex", state, **self._kwargs()) is True + # `--oauth-client-id ""` asks to go back to the built-in app, which the cached + # config cannot serve either. + assert ( + cli_mod._can_launch_from_cached_config( + "codex", state, **self._kwargs(oauth_client_id="") + ) + is False + ) + class TestPassthroughArgs: @pytest.mark.parametrize( @@ -2578,6 +2644,7 @@ def fake_configure_shared_state( use_pat=False, fable_enabled=None, databricks_ai_tools_enabled=None, + oauth_client_id=None, ): captured["workspace"] = workspace captured["profile"] = profile @@ -2616,6 +2683,7 @@ def fake_configure_shared_state( use_pat=False, fable_enabled=None, databricks_ai_tools_enabled=None, + oauth_client_id=None, ): configured_shared.append( (workspace, profile, tuple(tools) if tools is not None else None, force_login) @@ -2878,13 +2946,15 @@ def _stub_deps(monkeypatch, *, pat_token, existing_state=None): saved: list[dict] = [] monkeypatch.setattr(cli_mod, "load_state", lambda: dict(existing_state or {})) monkeypatch.setattr(cli_mod, "save_state", lambda s: saved.append(dict(s))) - monkeypatch.setattr(cli_mod, "run_databricks_login", lambda w, p: logins.append((w, p))) monkeypatch.setattr( - cli_mod, "ensure_databricks_auth", lambda w, p=None: ensures.append((w, p)) + cli_mod, "run_databricks_login", lambda w, p, **_k: logins.append((w, p)) + ) + monkeypatch.setattr( + cli_mod, "ensure_databricks_auth", lambda w, p=None, **_k: ensures.append((w, p)) ) monkeypatch.setattr(cli_mod, "resolve_pat_token", lambda p: pat_token) monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) - monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") + monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p, **_k: "token") monkeypatch.setattr( cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE ) @@ -3296,6 +3366,116 @@ def test_skip_validate_skips_single_tool_validation(self, monkeypatch): assert installed == [["claude"]] +class TestConfigureSharedStateOauthClientId: + """`--oauth-client-id` pins a custom OAuth app: it is persisted for the workspace, + passed to the login/auth check being performed *now* (state isn't saved yet when the + login runs), inherited by later launches, and cleared by an explicit empty string.""" + + WS = "https://example.databricks.com" + CLIENT_ID = "0844d280-b84d-45f2-b675-5d2c8fdc3825" + + @staticmethod + def _stub_deps(monkeypatch, *, existing_state=None): + import ucode.cli as cli_mod + + calls: dict[str, list] = {"ensure": [], "login": [], "token": []} + saved: list[dict] = [] + monkeypatch.setattr(cli_mod, "load_state", lambda: dict(existing_state or {})) + monkeypatch.setattr(cli_mod, "save_state", lambda s: saved.append(dict(s))) + monkeypatch.setattr( + cli_mod, + "ensure_databricks_auth", + lambda w, p=None, **kwargs: calls["ensure"].append(kwargs.get("oauth_client_id")), + ) + monkeypatch.setattr( + cli_mod, + "run_databricks_login", + lambda w, p=None, **kwargs: calls["login"].append(kwargs.get("oauth_client_id")), + ) + monkeypatch.setattr( + cli_mod, + "get_databricks_token", + lambda w, p=None, **kwargs: ( + calls["token"].append(kwargs.get("oauth_client_id")) or "token" + ), + ) + monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) + monkeypatch.setattr( + cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE + ) + monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], None)) + monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) + monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) + monkeypatch.setattr(cli_mod, "discover_codex_models", lambda w, t: ([], None)) + monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) + return cli_mod, calls, saved + + def test_persists_and_reaches_the_auth_check_and_the_token_mint(self, monkeypatch): + cli_mod, calls, _saved = self._stub_deps(monkeypatch) + + state = cli_mod.configure_shared_state(self.WS, oauth_client_id=self.CLIENT_ID) + + assert state["oauth_client_id"] == self.CLIENT_ID + # Passed explicitly rather than left to host resolution: the state entry that would + # answer that lookup is not on disk yet when this auth check runs. + assert calls["ensure"] == [self.CLIENT_ID] + assert calls["token"] == [self.CLIENT_ID] + + def test_force_login_signs_into_the_app_being_configured(self, monkeypatch): + cli_mod, calls, _saved = self._stub_deps(monkeypatch) + + cli_mod.configure_shared_state(self.WS, force_login=True, oauth_client_id=self.CLIENT_ID) + + assert calls["login"] == [self.CLIENT_ID] + assert calls["ensure"] == [] + + def test_launch_inherits_the_persisted_client_id(self, monkeypatch): + # A launch re-run passes None, which must mean "inherit", not "clear". + cli_mod, calls, _saved = self._stub_deps( + monkeypatch, + existing_state={"workspace": self.WS, "oauth_client_id": self.CLIENT_ID}, + ) + + state = cli_mod.configure_shared_state(self.WS) + + assert state["oauth_client_id"] == self.CLIENT_ID + assert calls["ensure"] == [self.CLIENT_ID] + + def test_empty_string_clears_the_persisted_client_id(self, monkeypatch): + cli_mod, calls, _saved = self._stub_deps( + monkeypatch, + existing_state={"workspace": self.WS, "oauth_client_id": self.CLIENT_ID}, + ) + + state = cli_mod.configure_shared_state(self.WS, oauth_client_id="") + + assert "oauth_client_id" not in state + assert calls["ensure"] == [""] + + def test_not_inherited_across_workspaces(self, monkeypatch): + cli_mod, _calls, _saved = self._stub_deps( + monkeypatch, + existing_state={"workspace": "https://other.databricks.com", "oauth_client_id": "old"}, + ) + + state = cli_mod.configure_shared_state(self.WS) + + assert "oauth_client_id" not in state + + def test_skip_preflight_still_persists_it(self, monkeypatch): + # --skip-preflight runs no login, but the agent configs written afterwards read the + # client id back out of state, so it still has to land there. + cli_mod, calls, saved = self._stub_deps(monkeypatch) + + state = cli_mod.configure_shared_state( + self.WS, skip_preflight=True, oauth_client_id=self.CLIENT_ID + ) + + assert state["oauth_client_id"] == self.CLIENT_ID + assert saved[-1]["oauth_client_id"] == self.CLIENT_ID + assert calls["ensure"] == [] and calls["login"] == [] + + class TestConfigureSharedStateMcpCleanup: """A workspace switch should scrub the previous workspace's MCP entries from installed client configs. Switching to the same workspace must not.""" @@ -3305,10 +3485,10 @@ def _stub_external_deps(monkeypatch): import ucode.cli as cli_mod monkeypatch.setattr(cli_mod, "normalize_workspace_url", lambda w: w) - monkeypatch.setattr(cli_mod, "run_databricks_login", lambda w, p: None) - monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda w, p=None: None) + monkeypatch.setattr(cli_mod, "run_databricks_login", lambda w, p, **_k: None) + monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda w, p=None, **_k: None) monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) - monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") + monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p, **_k: "token") monkeypatch.setattr( cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE ) @@ -3367,10 +3547,10 @@ def _stub(monkeypatch): import ucode.cli as cli_mod monkeypatch.setattr(cli_mod, "normalize_workspace_url", lambda w: w) - monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda w, p=None: None) - monkeypatch.setattr(cli_mod, "run_databricks_login", lambda w, p: None) + monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda w, p=None, **_k: None) + monkeypatch.setattr(cli_mod, "run_databricks_login", lambda w, p, **_k: None) monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) - monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") + monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p, **_k: "token") monkeypatch.setattr( cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE ) diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index 097d8969..fcdb3e67 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -5,7 +5,7 @@ import json import urllib.error -from ucode.smart_routing import codex_routing +from ucode.smart_routing import codex_hooks, codex_routing from ucode.smart_routing.codex_hooks import routing_models WS = "https://example.databricks.com" @@ -432,3 +432,22 @@ def test_decision_record_persists_rationale(tmp_path, monkeypatch): # returned none" from a display-placement bug. record = json.loads(decisions.read_text().strip()) assert record["rationale"] == "Cross-cutting refactor needs the strongest model." + + +class TestRoutingHookCustomOauthApp: + """Codex's route-subagent hook mints its own token, so it needs the workspace's + custom OAuth app named on its command line (mirrors claude_hooks).""" + + STATE = {"workspace": WS, "codex_models": ["databricks-gpt-5"]} + + def _route_command(self, state: dict) -> str: + doc: dict = {} + codex_hooks.sync_smart_routing_hooks(doc, state, enabled=True) + return doc["hooks"]["PreToolUse"][0]["hooks"][0]["command"] + + def test_flag_is_appended_when_the_workspace_has_a_custom_app(self): + command = self._route_command({**self.STATE, "oauth_client_id": "custom-app-id"}) + assert "--oauth-client-id custom-app-id" in command + + def test_flag_is_absent_without_a_custom_app(self): + assert "--oauth-client-id" not in self._route_command(dict(self.STATE)) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index e9ed1358..4813c1d8 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -281,6 +281,22 @@ def test_configure_dispatches_proxy_argv_to_cursor_writer(self, monkeypatch): assert removed_scopes == [] assert calls == [("github", _proxy_argv())] + def test_configure_pins_a_custom_oauth_app_in_the_proxy_argv(self, monkeypatch): + # The stdio bridge mints its own token, so the app has to be in the argv the + # client config records. + calls: list[tuple[str, list[str]]] = [] + monkeypatch.setattr( + mcp.cursor, + "write_mcp_server_config", + lambda name, argv: calls.append((name, argv)) or False, + ) + + mcp.configure_client_mcp_server( + "cursor", "github", GH_URL, WS, "p", oauth_client_id="custom-app-id" + ) + + assert calls[0][1][-2:] == ["--oauth-client-id", "custom-app-id"] + def test_configure_reports_user_scope_on_replace(self, monkeypatch): monkeypatch.setattr(mcp.cursor, "write_mcp_server_config", lambda name, argv: True) assert mcp.configure_client_mcp_server("cursor", "github", GH_URL, WS, "p") == [ diff --git a/tests/test_state.py b/tests/test_state.py index 36c8ce4f..668c6339 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -267,6 +267,43 @@ def test_use_pat_state_builds_pat_auth_command(self): assert "--use-pat" in result[agent]["auth_command"] assert "--profile DEFAULT" in result[agent]["auth_command"] + def test_custom_oauth_app_is_pinned_into_the_auth_command(self): + # Agents run `auth_command` as a bare command line, so the app it mints from has to + # be named in the command itself rather than resolved later. + result = build_agent_state( + { + "workspace": "https://example.databricks.com", + "profile": "DEFAULT", + "oauth_client_id": "custom-app-id", + "base_urls": FAKE_URLS, + } + ) + for agent in ("claude", "codex", "pi"): + assert "--oauth-client-id custom-app-id" in result[agent]["auth_command"] + # Codex reads argv, not a shell string. + codex_auth = result["codex"]["auth"] + assert codex_auth["args"][-2:] == ["--oauth-client-id", "custom-app-id"] + + def test_no_custom_oauth_app_leaves_the_command_unchanged(self): + result = build_agent_state( + { + "workspace": "https://example.databricks.com", + "base_urls": FAKE_URLS, + } + ) + for agent in ("claude", "codex", "pi"): + assert "--oauth-client-id" not in result[agent]["auth_command"] + + def test_blank_client_id_is_treated_as_absent(self): + result = build_agent_state( + { + "workspace": "https://example.databricks.com", + "oauth_client_id": " ", + "base_urls": FAKE_URLS, + } + ) + assert "--oauth-client-id" not in result["claude"]["auth_command"] + # --------------------------------------------------------------------------- # mark_tool_managed