diff --git a/AGENTS.md b/AGENTS.md index a46cbbdac..069543de4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -21,7 +21,7 @@ Tests live in `tests/`. - Use Python 3.12+. - Keep changes scoped to the requested behavior. -- Follow the existing module boundaries: CLI orchestration in `cli.py`, CLI-agent behavior in `agents/.py`, desktop-app adapters in `agents/desktop/`, shared agent dispatch in `agents/__init__.py`, Databricks calls in `databricks.py`, skill download (UC fetch client + on-disk writer + download orchestration) in `skills_download.py`, MCP-connection state glue in `mcp.py`, and presentation helpers in `ui.py`. Skill download persists no disk state — it writes files to `--path` (or the home dir) and registers only the schema-less skills MCP connection. `ug configure skills` with no `--location` (or `--mcp` with no `--location`) registers that schema-less connection without downloading anything. +- Follow the existing module boundaries: CLI orchestration in `cli.py`, CLI-agent behavior in `agents/.py`, desktop-app adapters in `agents/desktop/` with CLI setup/cleanup orchestration in `agents/desktop/setup.py`, shared agent dispatch in `agents/__init__.py`, Databricks calls in `databricks.py`, skill download (UC fetch client + on-disk writer + download orchestration) in `skills_download.py`, MCP-connection state glue in `mcp.py`, and presentation helpers in `ui.py`. Skill download persists no disk state — it writes files to `--path` (or the home dir) and registers only the schema-less skills MCP connection. `ug configure skills` with no `--location` (or `--mcp` with no `--location`) registers that schema-less connection without downloading anything. - Prefer existing helpers for config file writes, state persistence, UI messages, and Databricks authentication. - Use `ucode.os_compatibility.subprocess_cross_os.run` / `ucode.os_compatibility.subprocess_cross_os.popen` for subprocesses; they resolve Windows npm wrappers automatically and default text streams to UTF-8 with replacement for invalid bytes. Keep commands as argument lists. Do not repeat these encoding defaults at callsites; explicit encoding/error overrides and binary mode remain supported. Ruff rejects direct `subprocess.run` / `subprocess.Popen` calls outside `os_compatibility/subprocess_cross_os.py` and tests. - Use `acquire_exclusive_file_lock` / `release_file_lock` from `ucode.os_compatibility.file_lock_cross_os` instead of calling `fcntl.flock`, `fcntl.lockf`, or `msvcrt.locking` directly. Ruff enforces this outside the compatibility module and tests. @@ -154,6 +154,11 @@ When `claude_otel_tracing` is enabled, Desktop's reserved profile also owns for the existing `otel-headers` command. Native sessions rerun it as tokens rotate; no bearer is persisted. Disable/revert removes the wrapper, with atomic writes and rollback shared with the profile. Disabled tracing omits these fields on replacement. +Managed Claude setup also installs a generated Cowork plugin containing only the current +workspace's managed skills and Claude-targeted MCP servers. The source lives at +`~/.ucode/claude-cowork-marketplace`, and `~/.ucode/claude-cowork.json` records the selected +account/org root for cleanup. Cowork must have started once. Setup and cleanup use scoped +`claude plugin ... --cowork --json` commands with `CLAUDE_CONFIG_DIR`; failures stay best-effort. See `docs/claude-desktop.md` for native validation limits and revert behavior.
diff --git a/docs/claude-desktop.md b/docs/claude-desktop.md index 8d291d8a5..371adbea7 100644 --- a/docs/claude-desktop.md +++ b/docs/claude-desktop.md @@ -6,14 +6,60 @@ Windows, and Linux. This uses the same workspace, Databricks CLI OAuth profile, model source, and managed HTTP headers as Claude Code. There is no separate Desktop checkbox or managed field. +## Requirements and setup + +- Claude Desktop installed on macOS, Windows, or Linux. On Windows/Linux, open it + once to initialize application data. +- A named Databricks CLI OAuth profile with gateway/model access. A bearer token + alone is not enough for Desktop's credential helper. +- At least one Claude-named model in the configured static list or gateway catalog. + +Quit Claude Desktop, configure Claude Code, then reopen Desktop normally: + ```bash ug configure --profile --agents claude ``` -Quit Claude Desktop before configuring, then reopen it normally from your app launcher. -No terminal process needs to stay running. Claude Desktop must already be installed; -UG does not install it or manage the customer's MDM process. On Windows and Linux, -open Desktop once before configuring so it initializes its application-data directory. +No terminal process needs to stay running. UG does not install Desktop or manage +the customer's MDM process. Explicit-agent, selected-agent, and managed configure +flows all set up Desktop after successful Claude Code configuration. Ordinary agent +launches and `ug agents add` do not configure Desktop. + +## Cowork workspace skills and MCP servers + +Whenever Claude is configured, `ug configure` also installs a generated +`unity-gateway@unity-gateway-cowork` plugin into that account's Cowork scope. This applies +to managed and unmanaged workspaces. Start one Cowork session before configuring so Desktop +creates its active account/organization directory. If several Cowork profiles exist and UG +cannot identify one active profile, setup warns without changing any of them. + +The plugin contains the current workspace's Claude-targeted MCP servers. In a managed +workspace, that includes both user-selected and managed MCP servers, plus managed skills. +In an unmanaged workspace, it includes user-selected MCP servers. User, project, unmanaged +Databricks AI Tools, and other-workspace skills are not copied. `ug configure --mcp`, +`ug mcp add`, and `ug mcp remove` resync the plugin. MCP entries run the absolute +`ug mcp-proxy` executable with the selected workspace and CLI profile; no bearer token is +stored in the plugin. Saved PAT mode is supported. Custom OAuth is not yet supported for +Cowork MCP servers. + +UG owns these local files: + +- Generated marketplace: `~/.ucode/claude-cowork-marketplace` +- Selected Cowork root record: `~/.ucode/claude-cowork.json` + +Start a new Cowork session after configure. Plugin MCP servers connect asynchronously and +may take a few seconds to appear, so a tool can be absent on the first model turn. Deferred +tool search or a retry after another tool call should find it. The resulting tool name follows +Claude's `mcp__plugin____` pattern. + +Cowork plugin support can still be disabled by Anthropic account feature flags or managed +Desktop policy. UG treats installation as best-effort: failure warns but does not fail Claude +Code setup. If the current workspace has no Claude-targeted MCP servers or managed skills, UG +keeps the plugin installed but replaces its contents with an empty plugin so a prior workspace's +tools and skills do not remain active. If Cowork has never been started and UG has no prior plugin +state, setup skips silently. Only `ug revert` removes the plugin registration and generated +marketplace. The native Desktop gateway profile is separate and remains under the cleanup rules +below. ## Models and authentication @@ -133,23 +179,32 @@ otherwise it clears the selection. A later selection of another profile is prese Only this previous-selection ID is saved in `~/.ucode/claude-desktop.json`. Repeated configure retains it while UG remains selected. +Revert also uninstalls `unity-gateway@unity-gateway-cowork`, removes its marketplace +registration, and deletes UG's generated marketplace. This removes the MCP connectors and +skills delivered by that plugin. It also cleans up an exact UG registration left in another +Cowork profile by an older failed revert. Unrelated Desktop plugins, connectors, and skills are +left untouched. + +The revert summary reports `Claude Desktop profile: removed`, `unchanged`, or `failed`. + Old random-ID and workspace-derived profiles from earlier builds are left untouched; there is no automatic migration or cleanup. Workspace-specific UUID normalization is no longer needed; normal CLI URL handling is unchanged. Files are replaced atomically one at a time with best-effort rollback of their exact previous bytes on write failure. This is not a crash-atomic transaction across files. -Workspace state is cleared even if Desktop revert fails. Desktop's separate -previous-selection file is retained on failure, so cleanup can still be retried -with `ug revert`. +Workspace and Cowork cleanup state is retained if Desktop revert fails, so cleanup can be +retried with `ug revert`. The generated Cowork marketplace is also retained until its plugin +and marketplace registrations are removed successfully. A Desktop cleanup failure produces +one warning and exit code 1; the summary reports the failed cleanup. Profile configure/revert waits at most five seconds for its file lock. A timeout does not interrupt the holder or change profile, metadata, or selection-state files. -Desktop failures are warnings: successful Claude Code setup still returns exit 0. -`--dry-run` performs no Desktop discovery or writes. Removing Claude from managed -configuration does not revoke a previously written Desktop profile. Desktop does -not run UG's launch-time managed-agent checks; gateway authorization remains the -enforcement boundary. +Desktop setup failures only warn: successful Claude Code configuration still returns exit 0. +`--dry-run` performs no Desktop or Cowork discovery, writes, or plugin commands. Removing Claude +from managed configuration preserves both the Cowork plugin and a previously written native +Desktop profile until `ug revert`. Desktop does not run UG's launch-time managed-agent checks; +gateway authorization remains the enforcement boundary. ## Validation limits @@ -174,9 +229,10 @@ orchestration; they do not prove GUI inference on every app version. models without renaming them; that name filter does not reproduce every native validation rule. - Anthropic subscription relay, custom OAuth/PAT profiles, Desktop Smart Routing, - and connector/plugin compatibility are not implemented here. Native OTEL settings - are generated, but native export and ingestion remain unverified; session auth - limits are described above. + and third-party connectors outside UG's generated Cowork plugin are not implemented here. + Cowork's MCP proxy supports saved PAT mode but not custom OAuth. Native OTEL settings are + generated, but native export and ingestion remain unverified; session auth limits are + described above. - Observed inference applies to Cowork. Ordinary Chat is not claimed to use the same gateway path. diff --git a/src/ucode/agents/desktop/cowork.py b/src/ucode/agents/desktop/cowork.py new file mode 100644 index 000000000..e06c5f707 --- /dev/null +++ b/src/ucode/agents/desktop/cowork.py @@ -0,0 +1,456 @@ +"""Install Unity Gateway workspace skills and MCP servers into Claude Cowork.""" + +from __future__ import annotations + +import json +import os +import shutil +import tempfile +from pathlib import Path + +from ucode import config_io +from ucode.agents.desktop.claude import claude_desktop_root_directory +from ucode.databricks import build_mcp_proxy_argv +from ucode.os_compatibility import subprocess_cross_os +from ucode.skills_state import records_for_scope +from ucode.ui import print_warning + +MARKETPLACE_NAME = "unity-gateway-cowork" +PLUGIN_NAME = "unity-gateway" +PLUGIN_ID = f"{PLUGIN_NAME}@{MARKETPLACE_NAME}" +STATE_FILENAME = "claude-cowork.json" +MARKETPLACE_DIRECTORY_NAME = "claude-cowork-marketplace" + +_COWORK_MARKERS = ( + "cowork_settings.json", + "cowork_plugins", + "cowork_plugins/known_marketplaces.json", + "cowork_plugins/installed_plugins.json", +) + + +def configure_cowork(state: dict, mcp_servers: list[dict]) -> bool: + """Generate and install UG's Cowork plugin for the active workspace.""" + + workspace = state.get("workspace") + if not isinstance(workspace, str) or not workspace: + raise RuntimeError("No Databricks workspace is configured for Cowork.") + + previous_state = config_io.read_json_safe(_state_path()) + skill_records = _managed_skill_records(workspace) + if not previous_state and not skill_records and not mcp_servers: + return False + if not previous_state and not _cowork_sessions_root_exists(): + return False + + profile = state.get("profile") + if not isinstance(profile, str) or not profile: + raise RuntimeError("Select an explicit Databricks CLI profile for Cowork.") + + mcp_config = _cowork_mcp_config( + mcp_servers, + workspace, + profile, + use_pat=bool(state.get("use_pat")), + custom_oauth=bool(state.get("custom_oauth")), + ) + if not previous_state and not skill_records and not mcp_config["mcpServers"]: + return False + + cowork_root = discover_cowork_root() + previous_root = previous_state.get("cowork_root") + if isinstance(previous_root, str) and previous_root and Path(previous_root) != cowork_root: + _remove_plugin_registration(Path(previous_root)) + marketplace_root = _marketplace_root() + plugin_root = marketplace_root / "plugins" / PLUGIN_NAME + _write_marketplace(marketplace_root, plugin_root, skill_records, mcp_config) + config_io.atomic_write_json( + _state_path(), + {"version": 1, "cowork_root": str(cowork_root)}, + ) + _install_plugin(cowork_root, marketplace_root) + return True + + +def revert_cowork() -> bool: + """Remove only UG's Cowork plugin, marketplace registration, and generated source.""" + + state_path = _state_path() + state = config_io.read_json_safe(state_path) + raw_root = state.get("cowork_root") + roots: list[Path] = [] + if isinstance(raw_root, str) and raw_root: + roots.append(Path(raw_root)) + for root in _cowork_roots(): + if root not in roots and _has_ug_registration(root): + roots.append(root) + + changed = False + for root in roots: + changed = _remove_plugin_registration(root) or changed + + marketplace_root = _marketplace_root() + if marketplace_root.exists() or marketplace_root.is_symlink(): + if marketplace_root.is_symlink(): + marketplace_root.unlink() + else: + shutil.rmtree(marketplace_root) + changed = True + if state_path.exists(): + state_path.unlink() + changed = True + return changed + + +def _cowork_sessions_root_exists() -> bool: + root = claude_desktop_root_directory() + return root is not None and (root / "local-agent-mode-sessions").is_dir() + + +def _remove_plugin_registration(cowork_root: Path) -> bool: + """Remove UG's plugin and marketplace registration from one Cowork root.""" + + claude = _claude_binary() + env = os.environ.copy() + env["CLAUDE_CONFIG_DIR"] = str(cowork_root) + plugin_result = _run_claude_json( + claude, + ["plugin", "uninstall", PLUGIN_ID, "--cowork", "--json"], + env, + accepted_failures=("not installed", "not_installed", "not found"), + ) + marketplace_result = _run_claude_json( + claude, + [ + "plugin", + "marketplace", + "remove", + MARKETPLACE_NAME, + "--cowork", + "--json", + ], + env, + accepted_failures=("not found", "not_found", "not registered"), + ) + return bool(plugin_result or marketplace_result) + + +def discover_cowork_root(*, desktop_root: Path | None = None) -> Path: + """Find the active account/org Cowork configuration root without guessing IDs.""" + + root = desktop_root or claude_desktop_root_directory() + if root is None: + raise RuntimeError("Claude Cowork configuration is unavailable on this platform.") + candidates = _cowork_roots(desktop_root=root) + if not candidates: + raise RuntimeError("Open and start Claude Cowork once, then run `ug configure` again.") + if len(candidates) == 1: + return candidates[0] + + mtimes = {candidate: _cowork_mtime(candidate) for candidate in candidates} + newest = max(mtimes.values()) + winners = [candidate for candidate, mtime in mtimes.items() if mtime == newest] + if len(winners) == 1: + return winners[0] + marked_winners = [candidate for candidate in winners if _has_cowork_marker(candidate)] + if len(marked_winners) == 1: + return marked_winners[0] + paths = ", ".join(str(candidate) for candidate in candidates) + raise RuntimeError(f"Could not determine the active Cowork profile; candidates: {paths}.") + + +def _cowork_roots(*, desktop_root: Path | None = None) -> list[Path]: + root = desktop_root or claude_desktop_root_directory() + if root is None: + return [] + sessions_root = root / "local-agent-mode-sessions" + return sorted( + grandchild for child in _directories(sessions_root) for grandchild in _directories(child) + ) + + +def _has_ug_registration(cowork_root: Path) -> bool: + for relative_path in ( + "cowork_plugins/installed_plugins.json", + "cowork_plugins/known_marketplaces.json", + ): + try: + content = (cowork_root / relative_path).read_text(encoding="utf-8") + except OSError: + continue + if PLUGIN_ID in content or MARKETPLACE_NAME in content: + return True + return False + + +def _directories(path: Path) -> list[Path]: + try: + return [entry for entry in path.iterdir() if entry.is_dir()] + except OSError: + return [] + + +def _has_cowork_marker(candidate: Path) -> bool: + return any((candidate / marker).exists() for marker in _COWORK_MARKERS) + + +def _cowork_mtime(candidate: Path) -> int: + mtimes: list[int] = [] + for path in (candidate, *(candidate / marker for marker in _COWORK_MARKERS)): + try: + mtimes.append(path.stat().st_mtime_ns) + except OSError: + continue + return max(mtimes, default=0) + + +def _managed_skill_records(workspace: str) -> list[dict]: + normalized_workspace = workspace.rstrip("/") + return [ + record + for record in records_for_scope("managed", str(Path.home())) + if str(record.get("workspace") or "").rstrip("/") == normalized_workspace + ] + + +def _cowork_mcp_config( + servers: list[dict], + workspace: str, + profile: str, + *, + use_pat: bool = False, + custom_oauth: bool = False, +) -> dict: + entries: dict[str, dict] = {} + for server in servers: + name = server.get("name") + url = server.get("url") + clients = server.get("clients") or [] + if not isinstance(name, str) or not name or not isinstance(url, str) or not url: + continue + if "claude" not in clients: + continue + if custom_oauth: + print_warning(f"Skipping Cowork MCP server `{name}`: custom OAuth is not supported.") + continue + argv = build_mcp_proxy_argv( + url, + workspace, + profile, + use_pat=use_pat, + ) + executable = Path(argv[0]) + if not executable.is_absolute(): + raise RuntimeError( + "Could not resolve the absolute `ug` executable required by Claude Cowork." + ) + entries[name] = {"type": "stdio", "command": str(executable), "args": argv[1:]} + return {"mcpServers": entries} + + +def _write_marketplace( + marketplace_root: Path, + plugin_root: Path, + skill_records: list[dict], + mcp_config: dict, +) -> None: + marketplace_root.parent.mkdir(parents=True, exist_ok=True) + staging = Path( + tempfile.mkdtemp(prefix=f".{marketplace_root.name}-", dir=marketplace_root.parent) + ) + backup = marketplace_root.with_name(f".{marketplace_root.name}.previous") + try: + staging_plugin = staging / plugin_root.relative_to(marketplace_root) + (staging / ".claude-plugin").mkdir(parents=True) + (staging_plugin / ".claude-plugin").mkdir(parents=True) + marketplace: dict[str, object] = { + "name": MARKETPLACE_NAME, + "owner": {"name": "Databricks"}, + "plugins": [ + { + "name": PLUGIN_NAME, + "description": "Unity Gateway workspace skills and MCP servers for Cowork", + "source": f"./plugins/{PLUGIN_NAME}", + } + ], + } + plugin: dict[str, object] = { + "name": PLUGIN_NAME, + "description": "Unity Gateway workspace skills and MCP servers for Cowork", + "author": {"name": "Databricks"}, + } + if mcp_config["mcpServers"]: + plugin["mcpServers"] = "./.mcp.json" + _write_json(staging_plugin / ".mcp.json", mcp_config) + _write_json(staging / ".claude-plugin" / "marketplace.json", marketplace) + _write_json(staging_plugin / ".claude-plugin" / "plugin.json", plugin) + _copy_skills(staging_plugin / "skills", skill_records) + + if backup.exists(): + shutil.rmtree(backup) + if marketplace_root.exists(): + os.replace(marketplace_root, backup) + try: + os.replace(staging, marketplace_root) + except OSError: + if backup.exists() and not marketplace_root.exists(): + os.replace(backup, marketplace_root) + raise + if backup.exists(): + shutil.rmtree(backup) + finally: + if staging.exists(): + shutil.rmtree(staging) + + +def _copy_skills(destination: Path, records: list[dict]) -> None: + copied: set[str] = set() + for record in records: + bundle_name = record.get("bundle_name") + if not isinstance(bundle_name, str) or not bundle_name or bundle_name in copied: + continue + source = _claude_skill_directory(record) + if source is None: + continue + destination.mkdir(parents=True, exist_ok=True) + shutil.copytree(source, destination / bundle_name) + copied.add(bundle_name) + + +def _claude_skill_directory(record: dict) -> Path | None: + existing = [Path(raw) for raw in (record.get("dirs") or []) if Path(raw).is_dir()] + return next( + ( + path + for path in existing + if path.parent.name == "skills" and path.parent.parent.name == ".claude" + ), + existing[0] if existing else None, + ) + + +def _install_plugin(cowork_root: Path, marketplace_root: Path) -> None: + claude = _claude_binary() + env = os.environ.copy() + env["CLAUDE_CONFIG_DIR"] = str(cowork_root) + _run_claude( + claude, + ["plugin", "validate", str(marketplace_root)], + env, + ) + _run_claude_json( + claude, + ["plugin", "marketplace", "add", str(marketplace_root), "--cowork", "--json"], + env, + accepted_failures=("already exists", "already_exists", "already registered"), + ) + _run_claude_json( + claude, + [ + "plugin", + "marketplace", + "update", + MARKETPLACE_NAME, + "--cowork", + "--json", + ], + env, + ) + installed = _run_claude_json( + claude, + ["plugin", "list", "--cowork", "--json"], + env, + ) + action = "update" if _contains_plugin_id(installed) else "install" + _run_claude_json( + claude, + ["plugin", action, PLUGIN_ID, "--cowork", "--json"], + env, + ) + + +def _claude_binary() -> str: + executable = shutil.which("claude") + if not executable or not Path(executable).is_absolute(): + raise RuntimeError("Claude CLI is required to configure Cowork plugins.") + return executable + + +def _run_claude(claude: str, args: list[str], env: dict[str, str]) -> str: + result = subprocess_cross_os.run( + [claude, *args], + check=False, + capture_output=True, + text=True, + timeout=120, + env=env, + ) + if result.returncode != 0: + detail = (result.stderr or result.stdout or "").strip() + raise RuntimeError( + f"`claude {' '.join(args[:3])}` failed{f': {detail}' if detail else ''}." + ) + return result.stdout or "" + + +def _run_claude_json( + claude: str, + args: list[str], + env: dict[str, str], + *, + accepted_failures: tuple[str, ...] = (), +) -> object: + result = subprocess_cross_os.run( + [claude, *args], + check=False, + capture_output=True, + text=True, + timeout=120, + env=env, + ) + output = f"{result.stdout or ''}\n{result.stderr or ''}".strip() + lowered = output.lower() + if result.returncode != 0: + if any(marker in lowered for marker in accepted_failures): + return {} + raise RuntimeError( + f"`claude {' '.join(args[:3])}` failed{f': {output}' if output else ''}." + ) + try: + payload = json.loads(result.stdout or "null") + except json.JSONDecodeError as exc: + raise RuntimeError(f"`claude {' '.join(args[:3])}` returned invalid JSON.") from exc + if isinstance(payload, dict) and payload.get("outcome") not in (None, "ok", "success"): + message = str(payload.get("message") or "Claude plugin command failed.") + if any(marker in message.lower() for marker in accepted_failures): + return {} + raise RuntimeError(message) + return payload + + +def _contains_plugin_id(payload: object) -> bool: + if isinstance(payload, str): + return payload == PLUGIN_ID + if isinstance(payload, list): + return any(_contains_plugin_id(value) for value in payload) + if isinstance(payload, dict): + plugin_id = payload.get("pluginId") or payload.get("id") + if plugin_id == PLUGIN_ID: + return True + if payload.get("name") == PLUGIN_NAME and payload.get("marketplace") == MARKETPLACE_NAME: + return True + return any(_contains_plugin_id(value) for value in payload.values()) + return False + + +def _write_json(path: Path, payload: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + + +def _marketplace_root() -> Path: + return config_io.APP_DIR / MARKETPLACE_DIRECTORY_NAME + + +def _state_path() -> Path: + return config_io.APP_DIR / STATE_FILENAME diff --git a/src/ucode/agents/desktop/setup.py b/src/ucode/agents/desktop/setup.py index d2a58ab65..b063d0b89 100644 --- a/src/ucode/agents/desktop/setup.py +++ b/src/ucode/agents/desktop/setup.py @@ -1,4 +1,4 @@ -"""Best-effort Desktop setup after successful Claude Code configuration.""" +"""Best-effort Desktop configuration and cleanup for CLI commands.""" from __future__ import annotations @@ -13,6 +13,7 @@ from ucode.agents.desktop.claude import ( configure_claude_desktop as write_claude_desktop_profile, ) +from ucode.agents.desktop.cowork import configure_cowork, revert_cowork from ucode.config_io import is_dry_run from ucode.constants import MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER from ucode.databricks import build_tool_base_url, get_databricks_token, list_anthropic_model_catalog @@ -34,6 +35,31 @@ def configure_claude_desktop(state: dict, *, parent_schema: str | None = None) - print_warning(f"Could not configure Claude Desktop: {exc}. Claude Code setup is unchanged.") +def configure_claude_cowork(state: dict, mcp_servers: list[dict]) -> None: + if sys.platform not in ("darwin", "win32", "linux"): + return + if is_dry_run(): + print_note("Would configure Claude Cowork workspace skills and MCP servers.") + return + try: + configured = configure_cowork(state, mcp_servers) + except Exception as exc: # noqa: BLE001 -- Cowork must not block Claude Code setup + print_warning(f"Could not configure Claude Cowork: {exc}. Claude Code setup is unchanged.") + return + if configured: + print_success("Claude Cowork workspace plugin configured. Start a new Cowork session.") + + +def revert_claude_cowork() -> str: + if sys.platform not in ("darwin", "win32", "linux") or is_dry_run(): + return "unchanged" + try: + return "removed" if revert_cowork() else "unchanged" + except Exception as exc: # noqa: BLE001 -- Preserve other cleanup operations on failure + print_warning(f"Could not revert Claude Cowork: {exc}. Cowork setup was retained.") + return "failed" + + def _configure_desktop(state: dict, *, parent_schema: str | None) -> None: directory = claude_desktop_directory() if directory is None: @@ -122,11 +148,17 @@ def _configure_desktop(state: dict, *, parent_schema: str | None) -> None: def revert_desktop() -> str: if sys.platform not in ("darwin", "win32", "linux") or is_dry_run(): return "unchanged" + cowork_result = revert_claude_cowork() if claude_desktop_directory() is None: - return "unchanged" + return cowork_result try: result = revert_claude_desktop() except Exception as exc: # noqa: BLE001 -- Preserve other revert operations on failure print_warning(f"Could not revert Claude Desktop: {exc}. Desktop setup was retained.") return "failed" - return "removed" if result.changed else "unchanged" + native_result = "removed" if result.changed else "unchanged" + if cowork_result == "failed": + return "failed" + if "removed" in (native_result, cowork_result): + return "removed" + return "unchanged" diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 724bd6592..576b4b0e6 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -135,6 +135,7 @@ remove_skills_locations_command, revert_mcp_configs, ) +from ucode.mcp import managed_mcp_servers as load_managed_mcp_servers from ucode.os_compatibility import subprocess_cross_os from ucode.skills_download import ( configure_location_skills_download_command, @@ -814,6 +815,11 @@ def _setup_single_agent( from ucode.agents.desktop.setup import configure_claude_desktop configure_claude_desktop(state, parent_schema=parent_schema) + if not is_dry_run(): + resolved_mcps = _configure_managed_mcp_servers(managed) + skills_configured = _configure_managed_skills(managed) + if resolved_mcps is not None and skills_configured: + _sync_claude_cowork_mcp_servers(state, managed_mcp_servers=resolved_mcps) install_databricks_ai_tools_for_agents( [tool], state, force_refresh=tool not in ("claude", "codex") ) @@ -1062,8 +1068,13 @@ def _configure_workspace_command( ) registered_mcps: list[str] = [] if not is_dry_run(): - registered_mcps = _configure_managed_mcp_servers(managed) - _configure_managed_skills(managed) + resolved_mcps = _configure_managed_mcp_servers(managed) + skills_configured = _configure_managed_skills(managed) + registered_mcps = [ + str(server["name"]) for server in (resolved_mcps or []) if server.get("name") + ] + if "claude" in configured_tools and resolved_mcps is not None and skills_configured: + _sync_claude_cowork_mcp_servers(state, managed_mcp_servers=resolved_mcps) _summarize_managed_config(managed, configured_tools, registered_mcps) return 0 @@ -1115,7 +1126,8 @@ def _configure_workspace_command( state = configure_selected_tools(state, picked) # No managed config here: undo what a prior managed workspace left behind (its MCP servers and - # skills), so switching workspaces doesn't strand the old registry and skills. + # skills), so switching workspaces doesn't strand the old registry and skills. Cowork is synced + # below from the current workspace's user-selected MCP servers. if not is_dry_run(): _configure_managed_mcp_servers(None) _configure_managed_skills(None) @@ -1134,6 +1146,8 @@ def _configure_workspace_command( from ucode.agents.desktop.setup import configure_claude_desktop configure_claude_desktop(state) + if not is_dry_run(): + _sync_claude_cowork_mcp_servers(state) summary_lines = [f"[bold]Workspace:[/bold] [cyan]{state['workspace']}[/cyan]"] for tool_name in picked: spec = TOOL_SPECS[tool_name] @@ -1423,8 +1437,6 @@ def revert() -> int: # Older Codex (< 0.134.0) had ucode edit the shared ~/.codex/config.toml in # place; restoring the per-profile file above does not undo that. legacy_codex_stripped = revert_legacy_shared_config() - clear_state() - print_heading("Revert") print_kv("Workspace", state.get("workspace") or "none") for tool, spec in TOOL_SPECS.items(): @@ -1441,8 +1453,9 @@ def revert() -> int: "restored" if mcp_results.get(client) else "unchanged", ) if desktop_result == "failed": - print_warning("ug state cleared. Fix the Claude Desktop error and run ug revert again.") + print_warning("ug state retained. Fix the Claude Desktop error and run ug revert again.") return 1 + clear_state() print_success("ug state cleared") return 0 @@ -1785,6 +1798,28 @@ def _not_installed_message(agents: list[str]) -> str: return f"Not installed or not set up: {names}. Install the CLI on PATH, then rerun." +def _sync_claude_cowork_mcp_servers( + state: dict | None = None, *, managed_mcp_servers: list[dict] | None = None +) -> None: + """Keep Claude Cowork's generated plugin aligned with Claude-targeted MCP state.""" + current_state = state if state is not None else load_state() + configured_tools = set(current_state.get("available_tools") or []) | set( + current_state.get("last_configured_tools") or [] + ) + if "claude" not in configured_tools: + return + + user_mcp_servers = list(current_state.get("mcp_servers") or []) + if managed_mcp_servers is None: + managed_mcp_servers = load_managed_mcp_servers(current_state, {"claude"}) + from ucode.agents.desktop.setup import configure_claude_cowork + + configure_claude_cowork( + current_state, + [*user_mcp_servers, *managed_mcp_servers], + ) + + def _configure_optional_setup(state: dict, tools: list[str]) -> None: enabled = prompt_yes_no("Configure MCP servers, skills, and plugins?") state["databricks_ai_tools_enabled"] = enabled @@ -1794,6 +1829,8 @@ def _configure_optional_setup(state: dict, tools: list[str]) -> None: install_databricks_ai_tools_for_agents(tools, state) configure_mcp_command() + if "claude" in tools: + _sync_claude_cowork_mcp_servers() @mcp_app.command("add") @@ -1891,6 +1928,7 @@ def mcp_add( else None ) add_mcp_command(location=location, services=selected, agents=scope) + _sync_claude_cowork_mcp_servers() except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None @@ -1924,6 +1962,7 @@ def mcp_remove( ) try: remove_mcp_command(agents=requested_agents) + _sync_claude_cowork_mcp_servers() except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None @@ -2933,16 +2972,16 @@ def _print_budget_panel(recommendation: dict, tool: str, managed: dict | None = console.print(panel) -def _configure_managed_mcp_servers(managed: dict | None) -> list[str]: +def _configure_managed_mcp_servers(managed: dict | None) -> list[dict] | None: """Register the managed config's MCP servers for every enabled MCP-client agent. Runs during ``ug configure`` after the enabled agents are configured, so a workspace-published server reaches each agent's `/mcp` list without the developer re-adding it. ``managed`` is None when the (now-current) workspace has no managed config: the reconcile then unregisters any servers a prior managed workspace registered, so switching workspaces resets the MCP registry. - Best-effort: a failure warns and leaves the rest of configure intact. Returns the names of the - servers registered this run (the completion summary counts them; state alone can't, since #717 - writes them to agents' OS-managed files rather than ``managed_mcp_servers``). + Best-effort: a failure warns and leaves the rest of configure intact. Returns the resolved + servers, or ``None`` when reconciliation failed. The full entries also feed Claude Cowork's + generated plugin; state alone can't, since #717 writes some servers to OS-managed files. """ managed = managed or {} agents = {tool for tool in managed_enabled_tools(managed) if tool in MCP_CLIENTS} @@ -2956,17 +2995,17 @@ def _configure_managed_mcp_servers(managed: dict | None) -> list[str]: "Skipped workspace MCP setup this run — MCP service discovery was rate-limited " "(HTTP 429). Existing MCP servers are unchanged; run `ug configure` again to retry." ) - return [] + return None except RuntimeError as exc: print_warning(f"Could not register your workspace's MCP servers: {exc}") - return [] + return None names = [str(server["name"]) for server in registered if server.get("name")] if names: print_note(f"Registered workspace MCP server(s): {', '.join(names)}") - return names + return registered -def _configure_managed_skills(managed: dict | None) -> None: +def _configure_managed_skills(managed: dict | None) -> bool: """Download and reconcile the managed config's skills for every agent's ``/skills`` picker. Mirrors :func:`_configure_managed_mcp_servers`: runs during ``ug configure`` after the enabled @@ -2979,11 +3018,12 @@ def _configure_managed_skills(managed: dict | None) -> None: written, removed = reconcile_managed_skills(managed or {}) except (RuntimeError, OSError) as exc: print_warning(f"Could not sync your workspace's skills: {exc}") - return + return False if written: print_note(f"Downloaded workspace skill(s): {', '.join(written)}") if removed: print_note(f"Removed workspace skill(s) no longer configured: {', '.join(removed)}") + return True def _child_owns_stdout(tool: str, tool_args: list[str]) -> bool: @@ -4358,6 +4398,7 @@ def configure( "interactive picker." ) configure_mcp_command(services=services) + _sync_claude_cowork_mcp_servers() if ( fully_interactive and not combined_optional_setup @@ -4365,6 +4406,7 @@ def configure( and prompt_yes_no("Configure MCP servers now?") ): configure_mcp_command() + _sync_claude_cowork_mcp_servers() save_config_scope(scope) except typer.Exit: # `typer.Exit` subclasses RuntimeError, so it has to be re-raised ahead of the handler @@ -4422,10 +4464,11 @@ def status_cmd() -> None: def revert_cmd() -> None: """Clear ug state and restore backed-up agent config files.""" try: - revert() + result = revert() except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None + raise typer.Exit(result) @app.command("doctor", rich_help_panel="Manage") diff --git a/src/ucode/mcp_connection_login.py b/src/ucode/mcp_connection_login.py index b4dc44a18..17eb0414d 100644 --- a/src/ucode/mcp_connection_login.py +++ b/src/ucode/mcp_connection_login.py @@ -1,12 +1,12 @@ -"""Per-connection login for AI Gateway MCP services, driven by the proxy on a 401. +"""Per-connection login for AI Gateway MCP services, driven by the proxy on demand. A connection-backed AI Gateway MCP service (e.g. ``system.ai.github``) needs a per-user connection credential before its tools can be used. Until the user has -logged in to the underlying SaaS, AI Gateway answers requests with an HTTP 401 -(RFC 9728 ``WWW-Authenticate``). +logged in to the underlying SaaS, AI Gateway reports missing connection authentication +as an HTTP 401 or an MCP JSON-RPC authentication error. -The ``ug mcp-proxy`` bridge (see ``mcp_proxy``) sees that 401 in its httpx auth -flow and, for a connection-backed service, runs the Databricks CLI U2M login +The ``ug mcp-proxy`` bridge (see ``mcp_proxy``) sees that failure and, for a +connection-backed service, runs the Databricks CLI U2M login with an RFC 8707 ``resource`` indicator naming the service, then retries the request. A resource-aware ``/oidc`` drives the connection's own SaaS login before minting the token, so the credential exists on retry — transparently to @@ -38,7 +38,7 @@ def connection_from_url(url: str) -> str | None: ``https://ws/ai-gateway/mcp-services/system.ai.github`` -> ``system.ai.github``. A URL that is not an mcp-services endpoint (or names no service) returns - ``None`` — only connection-backed services get the login-on-401 treatment. + ``None`` — only connection-backed services get automatic connection login. """ marker = url.find(AIGW_MCP_SERVICES_SEGMENT) if marker == -1: diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index e2cf2fc22..f870cd2e7 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -35,6 +35,7 @@ import sys from collections.abc import AsyncIterator +from threading import Lock from types import ModuleType, TracebackType from typing import Protocol, Self @@ -42,6 +43,8 @@ from anyio import to_thread from mcp.client.streamable_http import streamable_http_client from mcp.server.stdio import stdio_server +from mcp.shared.message import SessionMessage +from mcp.types import JSONRPCError, JSONRPCNotification, JSONRPCRequest, JSONRPCResponse from ucode.databricks import ( ensure_pat_bearer, @@ -53,6 +56,9 @@ # Exit code used when the proxy cannot continue. MCP clients surface a non-zero # exit far more usefully than a timeout, so bail out instead of hanging. AUTH_FAILURE_EXIT_CODE = 2 +CONNECTION_AUTH_REQUIRED_ERROR_CODE = -32006 +AUTHENTICATE_TOOL_NAME = "authenticate" +TOOLS_REFRESH_TIMEOUT_SECONDS = 5.0 def _httpx() -> ModuleType: @@ -117,7 +123,64 @@ def _fail_fast(message: str) -> None: raise SystemExit(AUTH_FAILURE_EXIT_CODE) -def _build_token_auth(url: str, workspace: str, profile: str | None, *, use_pat: bool = False): +class _ConnectionLogin: + """Run automatic login once per proxy session; explicit login may retry.""" + + def __init__( + self, + url: str, + workspace: str, + profile: str | None, + *, + use_pat: bool = False, + ) -> None: + self.url = url + self.workspace = workspace + self.profile = profile + self.connection = None if use_pat else connection_from_url(url) + self._login_attempted = False + self._completed = False + self._lock = Lock() + + def run_if_needed(self) -> bool: + """Run automatic connection login once, returning whether the connection is ready.""" + return self._run_if_needed(check_service=True, retry=False) + + def authenticate_or_confirm_ready(self) -> bool: + """Authenticate explicitly, allowing retry after a cancelled or failed login.""" + return self._run_if_needed(check_service=False, retry=True) + + def _run_if_needed(self, *, check_service: bool, retry: bool) -> bool: + with self._lock: + if self.connection is None: + return False + if self._completed: + return True + if check_service: + try: + token = get_databricks_token(self.workspace, self.profile) + except RuntimeError: + return False + if not mcp_service_needs_connection_login(self.workspace, token, self.connection): + return False + if self._login_attempted and not retry: + return False + self._login_attempted = True + ok, detail = run_connection_login(self.url, self.workspace, profile=self.profile) + if not ok: + raise ProxyAuthError(f"connection login for '{self.connection}' failed: {detail}") + self._completed = True + return True + + +def _build_token_auth( + url: str, + workspace: str, + profile: str | None, + *, + use_pat: bool = False, + connection_login: _ConnectionLogin | None = None, +): """Build an httpx ``Auth`` that injects a fresh bearer and drives the per-connection login **lazily, on a 401** — not eagerly at startup. @@ -126,8 +189,10 @@ def _build_token_auth(url: str, workspace: str, profile: str | None, *, use_pat: the gateway answers with an RFC 9728 ``401`` until the user holds the per-user connection credential. Rather than block startup logging in every configured server (which pops N browsers and stalls the agent's MCP startup), we let the - bridge come up immediately — ``tools/list`` needs no credential — and only - run ``databricks auth login --resource`` when a request actually gets a 401. + bridge come up immediately and only run + ``databricks auth login --resource`` when a request actually reports missing + connection authentication. JSON-RPC authentication errors are handled by the + bridge below because httpx auth only sees HTTP responses. So a browser opens only for a service you actually *use*, one at a time, and an already-signed-in service never prompts. The login runs at most once per session; PAT profiles have no connection OAuth to drive, so they never do it. @@ -137,8 +202,9 @@ def _build_token_auth(url: str, workspace: str, profile: str | None, *, use_pat: ``async_auth_flow``, so the blocking login runs off the event loop.""" httpx = _httpx() # PAT auth has no interactive OAuth to drive, so never treat it as connection-backed. - connection = None if use_pat else connection_from_url(url) - login = {"attempted": False} + connection_login = connection_login or _ConnectionLogin( + url, workspace, profile, use_pat=use_pat + ) def _mint(request): # get_databricks_token honors the DATABRICKS_BEARER short-circuit and PAT @@ -156,30 +222,7 @@ def _needs_connection_login(response) -> bool: # A 401 from a connection-backed service means the per-user credential is # missing. Only drive the login once per session — if it doesn't resolve # the 401, re-running it on every subsequent request would loop browsers. - return connection is not None and response.status_code == 401 and not login["attempted"] - - def _connection_login_or_fail() -> bool: - """Drive the one-time connection sign-in for a 401. Returns whether a login was actually - performed, so the caller retries the request only then. - - A 401 is only fixable by a connection sign-in when the service is actually connection-backed - (its UC connection uses per-user OAuth-U2M). A no-login service (e.g. system.ai.web_search) - that 401s has a bad Databricks token, not a missing connection credential — so don't open a - doomed connection-login browser it can't complete (AIGTWY-4856), and don't bother retrying. - Classified once, on the first 401; a token/lookup failure also skips the login.""" - login["attempted"] = True - if connection is None: # unreachable via _needs_connection_login; narrows for the checker - return False - try: - token = get_databricks_token(workspace, profile) - except RuntimeError: - return False - if not mcp_service_needs_connection_login(workspace, token, connection): - return False - ok, detail = run_connection_login(url, workspace, profile=profile) - if not ok: - raise ProxyAuthError(f"connection login for '{connection}' failed: {detail}") - return True + return response is not None and response.status_code == 401 class _DatabricksTokenAuth(httpx.Auth): def sync_auth_flow(self, request): @@ -197,7 +240,7 @@ async def async_auth_flow(self, request): return # Run the blocking browser login off the event loop so the bridge's # other pumps aren't starved while the user completes the sign-in. - if not await to_thread.run_sync(_connection_login_or_fail): + if not await to_thread.run_sync(connection_login.run_if_needed): return _mint(request) yield request @@ -205,36 +248,274 @@ async def async_auth_flow(self, request): return _DatabricksTokenAuth() -async def _pump[T]( - source: _ReceiveStream[T], - dest: _SendStream[T], -) -> None: - """Forward every message from ``source`` to ``dest``. +class _McpBridge: + """Bridge MCP messages and recover connection authentication failures.""" + + def __init__(self, connection_login: _ConnectionLogin) -> None: + self.connection_login = connection_login + self.pending_requests: dict[int | str, SessionMessage] = {} + self.retried_requests: set[int | str] = set() + self.authentication_required = False + self.authentication_completed = False + self._tools_refresh_waiters: list[anyio.Event] = [] + + def _authentication_tool(self) -> dict: + connection = self.connection_login.connection or "this service" + return { + "name": AUTHENTICATE_TOOL_NAME, + "description": ( + f"Sign in to the Unity Gateway {connection} connection. Call this when the user " + "requests this service but its normal tools are unavailable. After sign-in, " + "immediately retry the user's original request in the same turn using the newly " + "available tools; do not ask the user to send another message." + ), + "inputSchema": { + "type": "object", + "properties": {}, + "additionalProperties": False, + }, + } + + def _authentication_tool_list_response(self, request_id: int | str) -> SessionMessage: + return SessionMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=request_id, + result={"tools": [self._authentication_tool()]}, + ) + ) + + @staticmethod + def _with_tool_list_changes( + message: SessionMessage, + original: SessionMessage | None, + ) -> SessionMessage: + payload = message.message + request = original.message if original is not None else None + if ( + not isinstance(payload, JSONRPCResponse) + or not isinstance(request, JSONRPCRequest) + or request.method != "initialize" + ): + return message + + result = dict(payload.result) + capabilities = dict(result.get("capabilities") or {}) + tools = dict(capabilities.get("tools") or {}) + tools["listChanged"] = True + capabilities["tools"] = tools + result["capabilities"] = capabilities + return SessionMessage( + JSONRPCResponse(jsonrpc="2.0", id=payload.id, result=result), + metadata=message.metadata, + ) - The proxy is transport-level: it never inspects or rewrites MCP method - payloads, so new methods and capabilities pass through untouched.""" - async with source, dest: - async for message in source: - await dest.send(message) + def _mark_tools_refreshed(self) -> None: + waiters, self._tools_refresh_waiters = self._tools_refresh_waiters, [] + for waiter in waiters: + waiter.set() + async def _handle_authentication_tool( + self, + request_id: int | str, + client: _SendStream[SessionMessage], + ) -> None: + try: + authenticated = await to_thread.run_sync( + self.connection_login.authenticate_or_confirm_ready + ) + except ProxyAuthError as exc: + result = { + "content": [{"type": "text", "text": str(exc)}], + "isError": True, + } + else: + if not authenticated: + result = { + "content": [ + { + "type": "text", + "text": "Authentication did not complete. Retry this tool after sign-in.", + } + ], + "isError": True, + } + else: + self.authentication_required = False + self.authentication_completed = True + refreshed = anyio.Event() + self._tools_refresh_waiters.append(refreshed) + await client.send( + SessionMessage( + JSONRPCNotification( + jsonrpc="2.0", + method="notifications/tools/list_changed", + ) + ) + ) + with anyio.move_on_after(TOOLS_REFRESH_TIMEOUT_SECONDS) as scope: + await refreshed.wait() + if scope.cancel_called: + if refreshed in self._tools_refresh_waiters: + self._tools_refresh_waiters.remove(refreshed) + text = ( + "Authentication complete, but the tool list did not refresh before the " + "timeout. Retry the user's original request now. If its tools are still " + "unavailable, tell the user to start a new session." + ) + else: + text = ( + "Authentication complete. Retry the user's original request now in this " + "same turn using the newly available tools. Do not ask the user to send " + "another message." + ) + result = { + "content": [{"type": "text", "text": text}], + "isError": False, + } + + await client.send( + SessionMessage(JSONRPCResponse(jsonrpc="2.0", id=request_id, result=result)) + ) -async def _pump_upstream[T]( - source: _ReceiveStream[T | Exception], - dest: _SendStream[T], -) -> None: - """Forward upstream messages, failing if the transport closes first.""" - async with source, dest: - async for message in source: - if isinstance(message, Exception): - detail = " ".join(str(message).split()) or type(message).__name__ - raise ProxyTransportError(f"upstream MCP transport failed: {detail}") from message - await dest.send(message) - raise ProxyTransportError("upstream MCP transport closed unexpectedly") + async def _authenticate_and_retry( + self, + request_id: int | str, + original: SessionMessage, + auth_error: SessionMessage, + client: _SendStream[SessionMessage], + upstream: _SendStream[SessionMessage], + ) -> None: + try: + authenticated = await to_thread.run_sync(self.connection_login.run_if_needed) + except ProxyAuthError: + authenticated = False + if authenticated: + await upstream.send(original) + return + self.pending_requests.pop(request_id, None) + self.retried_requests.discard(request_id) + await client.send(auth_error) + + async def pump_client( + self, + source: _ReceiveStream[SessionMessage | Exception], + dest: _SendStream[SessionMessage], + client: _SendStream[SessionMessage] | None = None, + ) -> None: + async with source, dest: + async with anyio.create_task_group() as tg: + async for message in source: + if isinstance(message, Exception): + detail = " ".join(str(message).split()) or type(message).__name__ + raise ProxyTransportError( + f"client MCP transport failed: {detail}" + ) from message + payload = message.message + if isinstance(payload, JSONRPCRequest): + if ( + client is not None + and self.authentication_required + and payload.method == "tools/call" + and isinstance(payload.params, dict) + and payload.params.get("name") == AUTHENTICATE_TOOL_NAME + ): + tg.start_soon( + self._handle_authentication_tool, + payload.id, + client, + ) + continue + self.pending_requests[payload.id] = message + await dest.send(message) + + async def pump_upstream( + self, + source: _ReceiveStream[SessionMessage | Exception], + dest: _SendStream[SessionMessage], + upstream: _SendStream[SessionMessage], + ) -> None: + async with source, dest: + async with anyio.create_task_group() as tg: + async for message in source: + if isinstance(message, Exception): + detail = " ".join(str(message).split()) or type(message).__name__ + raise ProxyTransportError( + f"upstream MCP transport failed: {detail}" + ) from message + + payload = message.message + original = None + if isinstance(payload, (JSONRPCResponse, JSONRPCError)): + original = self.pending_requests.get(payload.id) + if self.connection_login.connection is not None: + message = self._with_tool_list_changes(message, original) + payload = message.message + if isinstance(payload, JSONRPCError): + request_id = payload.id + request = original.message if original is not None else None + if ( + isinstance(request, JSONRPCRequest) + and request_id is not None + and request.method == "tools/list" + and payload.error.code == CONNECTION_AUTH_REQUIRED_ERROR_CODE + and self.connection_login.connection is not None + and not self.authentication_completed + ): + self.authentication_required = True + self.pending_requests.pop(request_id, None) + self.retried_requests.discard(request_id) + await dest.send(self._authentication_tool_list_response(request_id)) + continue + if ( + isinstance(request, JSONRPCRequest) + and request_id is not None + and request.method == "tools/call" + and request_id not in self.retried_requests + and payload.error.code == CONNECTION_AUTH_REQUIRED_ERROR_CODE + and self.connection_login.connection is not None + ): + self.retried_requests.add(request_id) + tg.start_soon( + self._authenticate_and_retry, + request_id, + original, + message, + dest, + upstream, + ) + continue + + tools_refreshed = False + if isinstance(payload, JSONRPCResponse) and original is not None: + request = original.message + tools_refreshed = ( + isinstance(request, JSONRPCRequest) + and request.method == "tools/list" + and self.authentication_completed + ) + if ( + isinstance(payload, (JSONRPCResponse, JSONRPCError)) + and payload.id is not None + ): + self.pending_requests.pop(payload.id, None) + self.retried_requests.discard(payload.id) + await dest.send(message) + if tools_refreshed: + self._mark_tools_refreshed() + raise ProxyTransportError("upstream MCP transport closed unexpectedly") async def _run(url: str, workspace: str, profile: str | None, use_pat: bool = False) -> None: httpx = _httpx() - auth = _build_token_auth(url, workspace, profile, use_pat=use_pat) + connection_login = _ConnectionLogin(url, workspace, profile, use_pat=use_pat) + auth = _build_token_auth( + url, + workspace, + profile, + use_pat=use_pat, + connection_login=connection_login, + ) # 2.x-native shape: hand the transport a pre-built AsyncClient carrying our # per-request auth. Works on mcp 1.28+ and 2.x; `streamable_http_client` # yields a (read, write) pair in both. @@ -254,9 +535,10 @@ async def _run(url: str, workspace: str, profile: str | None, use_pat: bool = Fa http_read, http_write = streams[0], streams[1] async with stdio_server() as (stdio_read, stdio_write): # Bidirectional bridge: client stdin -> Databricks, Databricks -> client stdout. + bridge = _McpBridge(connection_login) async with anyio.create_task_group() as tg: - tg.start_soon(_pump_upstream, http_read, stdio_write) - await _pump(stdio_read, http_write) + tg.start_soon(bridge.pump_upstream, http_read, stdio_write, http_write) + await bridge.pump_client(stdio_read, http_write, stdio_write) tg.cancel_scope.cancel() @@ -303,12 +585,11 @@ def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool "Set DATABRICKS_BEARER, or reconfigure the profile." ) - # The per-connection sign-in is NOT driven here. Doing it eagerly at startup - # blocks the agent's MCP startup and, with N configured connection-backed - # servers, pops N browsers at once and stalls until each times out. Instead - # the bridge opens immediately (tools/list needs no credential) and the login - # is driven lazily, on a 401, from the httpx auth flow (see _build_token_auth) - # — so a browser opens only for a service whose tool you actually call. + # The per-connection sign-in is NOT driven before opening the bridge. The + # bridge starts first and drives login only when this service reports missing + # connection auth, either as HTTP 401 or MCP error -32006. Services whose + # connections are already authenticated, and services such as web search that + # need no connection, never open a browser. # Pre-flight the token before opening the bridge. Without this, the first # token failure surfaces from inside the transport's task group, where it can diff --git a/tests/README.md b/tests/README.md index c14c2e96f..3f33b6a64 100644 --- a/tests/README.md +++ b/tests/README.md @@ -168,6 +168,9 @@ cover Desktop token-only stdout, explicit workspace/profile arguments, interacti browser recovery, bounded login/lock waits, cancellation, and lock contention. Recovery regressions check token recheck before browser login, lock release after success/failure, one shared deadline, and network/CLI failures at each recovery stage. +Proxy bridge coverage also checks that unauthenticated tool discovery exposes a temporary +authentication tool, advertises dynamic tool-list changes, refreshes the real tools after login, +and tells the model to retry the original request in the same turn. Invalid token output must not appear in helper errors or stdout/stderr. Shared CLI discovery selects the newest absolute binary for cached tokens and every recovery subprocess; a missing CLI fails before token fetch or browser login. @@ -222,8 +225,19 @@ changing the UUID, and configure/revert warnings on profile-lock timeout without changing profile or selection-state files. An existing reserved UUID without records is replaced on configure and deleted on revert. Workspace changes reuse the same file. CLI revert clears workspace state even when Desktop cleanup fails. A lock-failure -regression verifies that Desktop files and previous-selection state survive, then -retry deletes the UG profile and restores the previous selection without workspace state. +regression invokes the Typer CLI and checks exit code 1, one warning, and a failed +Desktop summary while profile and previous-selection files survive. Retry returns +exit code 0, deletes the UG profile, and restores the previous selection without +workspace state; a repeated revert also returns 0. + +`test_cowork_plugin.py` covers active Cowork-root discovery, ambiguous-profile refusal, +generated marketplace/plugin/MCP manifests, current-workspace managed-skill selection, +Claude-only MCP filtering, absolute proxy executables, saved PAT mode, custom-OAuth refusal, +install/update selection, idempotent cleanup, and saved-root removal. CLI component tests +cover managed setup after successful skill/MCP reconciliation, preservation after transient +reconciliation failures, and cleanup for unmanaged or non-Claude configurations. These tests +mock the Claude plugin CLI and do not establish native Desktop policy, account-feature, +late MCP connection, tool discovery, or inference behavior. Model-source tests require Desktop-only, case-insensitive `\bclaude-` filtering instead of copying every catalog entry. Mixed static/discovered lists exclude GPT/OSS @@ -294,7 +308,8 @@ Fake CLI subprocesses check that re-auth runs only when a successful help probe advertises `--no-browser`. Search component tests verify token reuse and transient timeout retries. `test_mcp_connection_login.py` verifies token/failure invalidation after successful and failed login attempts; `test_mcp_proxy.py` checks fresh tokens -on async post-login retries. These checks do not establish live login or +on async post-login retries and recovery when a connection-backed service returns +the MCP authentication-required error. These checks do not establish live login or launch-time savings. `test_mcp_web_search_concurrency.py` drives the real stdio dispatcher with controlled diff --git a/tests/e2e_cuj/test_cuj4_smart_routing.py b/tests/e2e_cuj/test_cuj4_smart_routing.py index 8ba46efc0..42ec2daa8 100644 --- a/tests/e2e_cuj/test_cuj4_smart_routing.py +++ b/tests/e2e_cuj/test_cuj4_smart_routing.py @@ -417,15 +417,27 @@ def test_smart_router_config_version(self, cuj, SMART_ROUTER_CONFIG_VERSION): for path, records in agent_sessions(session, agent).items() if path not in existing_sessions and is_child_session(agent, path, records) } - assert len(children) == 1, (agent, children.keys()) - assert any( - child_task.value in answer - for records in children.values() - for answer in assistant_answers(agent, records) - ), agent - decisions = read_jsonl(decisions_path)[decision_count:] + completed_children = { + path: records + for path, records in children.items() + if any( + child_task.value in answer + for answer in assistant_answers(agent, records) + ) + } + assert len(completed_children) == 1, (agent, children.keys()) + decisions = [ + decision + for decision in read_jsonl(decisions_path)[decision_count:] + if child_task.filename in (decision.get("task_name") or "") + ] requests = recorder.requests_after(checkpoint) - routes = [request for request in requests if request.path == ROUTING_PATH] + routes = [ + request + for request in requests + if request.path == ROUTING_PATH + and child_task.filename in request.payload["task"]["prompt"] + ] assert len(routes) == int(subagent_routed), (agent, routes) if subagent_routed: assert_subagent_routed( diff --git a/tests/integration/README.md b/tests/integration/README.md index 4e188c967..8b03c9f7e 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -241,6 +241,8 @@ covered by the hermetic proxy tests in `tests/test_gateway_proxy_integration.py` `utils/agents/claude.py` and `utils/agents/codex.py` own the existing native answer parsing, session paths, and child-session detection. `utils/evidence.py` dispatches through them without changing existing journey commands or completion assertions. +Claude child answers include sidechain `SubagentHandback` messages, which newer +Claude versions use instead of final text blocks; parent handbacks remain excluded. ## Run a specific combination @@ -402,9 +404,16 @@ configure/revert warnings that preserve profile and selection-state files on loc An existing reserved UUID without records is replaced on configure and deleted on revert; these remain filesystem component checks. CLI revert clears workspace state even when Desktop cleanup fails. A component -regression verifies that a lock failure preserves Desktop files and selection state, -then retry restores the previous selection without workspace state. This does not -establish native Desktop coverage. +regression invokes the Typer CLI and checks that a lock failure returns exit code 1, +emits one warning and a failed Desktop summary, and preserves profile and selection +files. Retry restores the previous selection without workspace state and returns 0; +a repeated revert also returns 0. This does not establish native Desktop coverage. + +`../test_cowork_plugin.py` covers Cowork-root discovery, generated plugin contents, +scoped Claude plugin commands, install/update selection, PAT/custom-OAuth behavior, and +owned-file cleanup. `../test_desktop_configure_cli.py` covers managed orchestration and +cleanup. Both are mocked component coverage; the integration suite does not launch Cowork, +connect a plugin MCP server, discover its deferred tool, or run a task through it. Model-source tests require Desktop-only, case-insensitive `\bclaude-` filtering instead of copying every catalog entry. Mixed static/discovered lists exclude GPT/OSS diff --git a/tests/integration/utils/agents/claude.py b/tests/integration/utils/agents/claude.py index 32bf4e4f1..52f2be47b 100644 --- a/tests/integration/utils/agents/claude.py +++ b/tests/integration/utils/agents/claude.py @@ -23,7 +23,6 @@ def is_child_session(path: str, records: list[dict]) -> bool: def assistant_answers(records: list[dict]) -> list[str]: answers = [] - handbacks = {} for record in records: message = record.get("message", {}) if record.get("type") == "assistant" and message.get("role") == "assistant": @@ -31,15 +30,7 @@ def assistant_answers(records: list[dict]) -> list[str]: if part.get("type") == "text" and isinstance(part.get("text"), str): answers.append(part["text"]) elif _is_subagent_handback(record, part): - # Claude 2.1.290 can deliver a child's final answer through this - # tool. The call alone is not evidence that delivery succeeded. - handbacks[part["id"]] = part["input"]["message"] - elif _is_successful_child_turn_completion(record): - for part in message.get("content", []): - if part.get("type") == "tool_result" and not part.get("is_error"): - answer = handbacks.pop(part.get("tool_use_id"), None) - if answer is not None: - answers.append(answer) + answers.append(part["input"]["message"]) return answers @@ -53,16 +44,6 @@ def _is_subagent_handback(record: dict, part: dict) -> bool: ) -def _is_successful_child_turn_completion(record: dict) -> bool: - return ( - record.get("type") == "user" - and record.get("message", {}).get("role") == "user" - and record.get("isSidechain") is True - and record.get("toolEndsTurn") is True - and record.get("toolUseResult", {}).get("success") is True - ) - - def completed_task_models(records: list[dict], answer_value: str) -> set[str]: """Parse Claude's response-reported model on the assistant's matching answer.""" assert isinstance(answer_value, str) and answer_value.strip(), ( diff --git a/tests/test_cli.py b/tests/test_cli.py index deaac9f8e..46353c2da 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2035,6 +2035,28 @@ def test_mcp_list_forwards_agents_option(self, monkeypatch): assert result.exit_code == 0, result.output assert calls == [{"claude", "codex"}] + def test_mcp_add_syncs_claude_cowork(self): + with ( + patch("ucode.cli.add_mcp_command") as mock_add, + patch("ucode.cli._sync_claude_cowork_mcp_servers") as mock_sync, + ): + result = runner.invoke(app, ["mcp", "add", "--names", "system.ai.github"]) + + assert result.exit_code == 0, result.output + mock_add.assert_called_once_with(location=None, services={"system.ai.github"}, agents=None) + mock_sync.assert_called_once_with() + + def test_mcp_remove_syncs_claude_cowork(self): + with ( + patch("ucode.cli.remove_mcp_command") as mock_remove, + patch("ucode.cli._sync_claude_cowork_mcp_servers") as mock_sync, + ): + result = runner.invoke(app, ["mcp", "remove"]) + + assert result.exit_code == 0, result.output + mock_remove.assert_called_once_with(agents=None) + mock_sync.assert_called_once_with() + class TestAuthTokenCommand: """`ucode auth-token` is the cross-platform apiKeyHelper (#116).""" @@ -3411,6 +3433,24 @@ def test_reverts_mcp_configs_before_clearing_state(self): assert cleared == [True] assert "Claude Code MCP config: restored" in result.output + def test_desktop_revert_failure_is_actionable(self): + with ( + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.agents.desktop.setup.revert_desktop", return_value="failed"), + patch("ucode.cli.revert_mcp_configs", return_value={}), + patch("ucode.cli.claude_agent.revert_managed_settings", return_value="unchanged"), + patch("ucode.cli.codex_agent.revert_managed_config", return_value="unchanged"), + patch("ucode.cli.restore_file", return_value=False), + patch("ucode.cli.revert_legacy_shared_config", return_value=False), + patch("ucode.cli.clear_state") as clear_state, + ): + result = runner.invoke(app, ["revert"]) + + assert result.exit_code == 1 + clear_state.assert_not_called() + assert "ug state retained" in result.output + assert "Fix the Claude Desktop error and run ug revert again" in result.output + @pytest.mark.parametrize("args", [[], ["claude"], ["codex"]]) def test_unknown_routing_config_does_not_block_launch(monkeypatch, args): @@ -3787,11 +3827,24 @@ def test_optional_setup_installs_ai_tools_and_configures_mcp(self): import ucode.cli as cli_mod state = {"available_tools": ["claude", "codex"]} + configured_state = { + **state, + "mcp_servers": [ + { + "name": "system-ai-github", + "url": "https://workspace.example/api/2.0/mcp/system.ai.github", + "clients": ["claude"], + } + ], + } with ( patch("ucode.cli.prompt_yes_no", return_value=True) as mock_prompt, patch("ucode.cli.save_state") as mock_save, patch("ucode.cli.install_databricks_ai_tools_for_agents") as mock_install, patch("ucode.cli.configure_mcp_command") as mock_mcp, + patch("ucode.cli.load_state", return_value=configured_state), + patch("ucode.cli.load_managed_mcp_servers", return_value=[]), + patch("ucode.agents.desktop.setup.configure_claude_cowork") as mock_configure_cowork, ): cli_mod._configure_optional_setup(state, ["claude", "codex"]) @@ -3800,6 +3853,91 @@ def test_optional_setup_installs_ai_tools_and_configures_mcp(self): mock_save.assert_called_once_with(state) mock_install.assert_called_once_with(["claude", "codex"], state) mock_mcp.assert_called_once_with() + mock_configure_cowork.assert_called_once_with( + configured_state, configured_state["mcp_servers"] + ) + + def test_cowork_sync_merges_user_and_managed_mcp_servers(self): + import ucode.cli as cli_mod + + state = { + **MINIMAL_STATE, + "available_tools": ["claude"], + "mcp_servers": [ + { + "name": "user-server", + "url": "https://workspace.example/user", + "clients": ["claude"], + } + ], + } + managed = [ + { + "name": "managed-server", + "url": "https://workspace.example/managed", + "clients": ["claude"], + } + ] + with ( + patch("ucode.cli.load_state", return_value=state), + patch("ucode.agents.desktop.setup.configure_claude_cowork") as mock_configure, + ): + cli_mod._sync_claude_cowork_mcp_servers(state, managed_mcp_servers=managed) + + mock_configure.assert_called_once_with(state, [*state["mcp_servers"], *managed]) + + def test_cowork_sync_loads_managed_mcp_servers_when_not_supplied(self): + import ucode.cli as cli_mod + + state = { + **MINIMAL_STATE, + "available_tools": ["claude"], + "mcp_servers": [], + } + managed = [ + { + "name": "managed-server", + "url": "https://workspace.example/managed", + "clients": ["claude"], + } + ] + with ( + patch("ucode.cli.load_state", side_effect=AssertionError("must not reload state")), + patch("ucode.cli.load_managed_mcp_servers", return_value=managed) as mock_managed, + patch("ucode.agents.desktop.setup.configure_claude_cowork") as mock_configure, + ): + cli_mod._sync_claude_cowork_mcp_servers(state) + + mock_managed.assert_called_once_with(state, {"claude"}) + mock_configure.assert_called_once_with(state, managed) + + def test_cowork_sync_preserves_explicit_empty_managed_mcp_state(self): + import ucode.cli as cli_mod + + state = { + **MINIMAL_STATE, + "available_tools": ["claude"], + "mcp_servers": [], + } + with ( + patch( + "ucode.cli.load_managed_mcp_servers", + side_effect=AssertionError("explicit managed state must be preserved"), + ), + patch("ucode.agents.desktop.setup.configure_claude_cowork") as mock_configure, + ): + cli_mod._sync_claude_cowork_mcp_servers(state, managed_mcp_servers=[]) + + mock_configure.assert_called_once_with(state, []) + + def test_cowork_sync_skips_when_claude_is_not_configured(self): + import ucode.cli as cli_mod + + state = {**MINIMAL_STATE, "available_tools": ["codex"]} + with patch("ucode.agents.desktop.setup.configure_claude_cowork") as mock_configure: + cli_mod._sync_claude_cowork_mcp_servers(state) + + mock_configure.assert_not_called() def test_optional_setup_decline_does_nothing(self): import ucode.cli as cli_mod @@ -4096,6 +4234,7 @@ def test_mcp_with_agents_configures_then_registers_services(self): patch("ucode.cli.install_tool_binary"), patch("ucode.cli.configure_workspace_command") as mock_cfg, patch("ucode.cli.configure_mcp_command") as mock_mcp, + patch("ucode.cli._sync_claude_cowork_mcp_servers") as mock_sync, ): result = runner.invoke( app, @@ -4106,6 +4245,7 @@ def test_mcp_with_agents_configures_then_registers_services(self): selected_tools=["claude"], ) mock_mcp.assert_called_once_with(services={"system.ai.slack", "system.ai.github"}) + mock_sync.assert_called_once_with() def test_mcp_only_configures_workspace_without_agent_picker(self): # `--mcp` with no --agents (e.g. Cursor): configure the workspace directly, @@ -4115,6 +4255,7 @@ def test_mcp_only_configures_workspace_without_agent_picker(self): patch("ucode.cli.configure_workspace_command") as mock_cfg, patch("ucode.cli._configure_shared_workspace_states") as mock_shared, patch("ucode.cli.configure_mcp_command") as mock_mcp, + patch("ucode.cli._sync_claude_cowork_mcp_servers") as mock_sync, ): result = runner.invoke( app, @@ -4135,6 +4276,7 @@ def test_mcp_only_configures_workspace_without_agent_picker(self): mock_shared.call_args.kwargs.get("tools") == [] or mock_shared.call_args.args[1] == [] ) mock_mcp.assert_called_once_with(services={"system.ai.slack"}) + mock_sync.assert_called_once_with() def test_mcp_rejects_bare_short_name(self): with ( @@ -4735,7 +4877,7 @@ def test_configure_managed_mcp_servers_skips_gracefully_on_rate_limit(self, monk result = cli_mod._configure_managed_mcp_servers({"enabled_agents": {"claude": {}}}) - assert result == [] + assert result is None assert warned == [] # not surfaced as a failure assert notes and "rate-limited" in notes[0].lower() and "429" in notes[0] @@ -4757,6 +4899,8 @@ def test_unmanaged_workspace_reconciles_managed_mcp_servers(self, monkeypatch): ) calls: list = [] monkeypatch.setattr(cli_mod, "_configure_managed_mcp_servers", lambda m: calls.append(m)) + cowork_sync = MagicMock() + monkeypatch.setattr(cli_mod, "_sync_claude_cowork_mcp_servers", cowork_sync) assert ( cli_mod.configure_workspace_command( @@ -4765,6 +4909,7 @@ def test_unmanaged_workspace_reconciles_managed_mcp_servers(self, monkeypatch): == 0 ) assert calls == [None] + cowork_sync.assert_called_once() def test_configures_available_subset_by_default(self, monkeypatch): """A workspace with no OpenAI models still configures claude and pi.""" diff --git a/tests/test_cowork_plugin.py b/tests/test_cowork_plugin.py new file mode 100644 index 000000000..a0b170963 --- /dev/null +++ b/tests/test_cowork_plugin.py @@ -0,0 +1,449 @@ +"""Claude Cowork plugin generation and scoped CLI installation.""" + +import json +import os +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from ucode.agents.desktop import cowork + + +def _cowork_root(desktop_root: Path, account: str, org: str) -> Path: + root = desktop_root / "local-agent-mode-sessions" / account / org + root.mkdir(parents=True) + return root + + +def test_discover_cowork_root_requires_started_session(tmp_path): + with pytest.raises(RuntimeError, match="start Claude Cowork once"): + cowork.discover_cowork_root(desktop_root=tmp_path / "Claude-3p") + + +def test_discover_cowork_root_uses_only_candidate(tmp_path): + desktop_root = tmp_path / "Claude-3p" + expected = _cowork_root(desktop_root, "account", "org") + + assert cowork.discover_cowork_root(desktop_root=desktop_root) == expected + + +def test_discover_cowork_root_prefers_marked_candidate(tmp_path): + desktop_root = tmp_path / "Claude-3p" + _cowork_root(desktop_root, "account-a", "org-a") + expected = _cowork_root(desktop_root, "account-b", "org-b") + (expected / "cowork_plugins").mkdir() + + assert cowork.discover_cowork_root(desktop_root=desktop_root) == expected + + +def test_discover_cowork_root_prefers_new_active_candidate_over_old_marked_candidate(tmp_path): + desktop_root = tmp_path / "Claude-3p" + old = _cowork_root(desktop_root, "account-a", "org-a") + old_marker = old / "cowork_plugins" + old_marker.mkdir() + new = _cowork_root(desktop_root, "account-b", "org-b") + os.utime(old, ns=(100, 100)) + os.utime(old_marker, ns=(100, 100)) + os.utime(new, ns=(200, 200)) + + assert cowork.discover_cowork_root(desktop_root=desktop_root) == new + + +def test_discover_cowork_root_uses_unique_newest_marker(tmp_path): + desktop_root = tmp_path / "Claude-3p" + old = _cowork_root(desktop_root, "account-a", "org-a") + new = _cowork_root(desktop_root, "account-b", "org-b") + for index, root in enumerate((old, new), start=1): + marker = root / "cowork_settings.json" + marker.write_text("{}") + timestamp = 1_900_000_000_000_000_000 + index * 1_000_000_000 + os.utime(marker, ns=(timestamp, timestamp)) + os.utime(root, ns=(timestamp, timestamp)) + + assert cowork.discover_cowork_root(desktop_root=desktop_root) == new + + +def test_discover_cowork_root_rejects_tied_candidates(tmp_path): + desktop_root = tmp_path / "Claude-3p" + candidates = [ + _cowork_root(desktop_root, "account-a", "org-a"), + _cowork_root(desktop_root, "account-b", "org-b"), + ] + for root in candidates: + marker = root / "cowork_settings.json" + marker.write_text("{}") + os.utime(marker, ns=(100, 100)) + os.utime(root, ns=(100, 100)) + + with pytest.raises(RuntimeError, match="Could not determine the active Cowork profile"): + cowork.discover_cowork_root(desktop_root=desktop_root) + + +def test_configure_generates_plugin_and_installs_in_cowork_scope(tmp_path, monkeypatch): + app_dir = tmp_path / ".ucode" + cowork_root = tmp_path / "Claude-3p" / "local-agent-mode-sessions" / "acct" / "org" + cowork_root.mkdir(parents=True) + claude_skill = tmp_path / "home" / ".claude" / "skills" / "review" + agent_skill = tmp_path / "home" / ".agents" / "skills" / "review" + claude_skill.mkdir(parents=True) + agent_skill.mkdir(parents=True) + (claude_skill / "SKILL.md").write_text("managed skill") + (agent_skill / "SKILL.md").write_text("other copy") + ug_binary = tmp_path / "bin" / "ug" + monkeypatch.setattr(cowork.config_io, "APP_DIR", app_dir) + monkeypatch.setattr(cowork, "discover_cowork_root", lambda: cowork_root) + monkeypatch.setattr(cowork, "_cowork_sessions_root_exists", lambda: True) + monkeypatch.setattr(cowork.shutil, "which", lambda name: str(ug_binary.parent / name)) + monkeypatch.setattr( + cowork, + "records_for_scope", + lambda scope, base: [ + { + "bundle_name": "review", + "workspace": "https://workspace.example/", + "dirs": [str(agent_skill), str(claude_skill)], + }, + { + "bundle_name": "foreign", + "workspace": "https://other.example", + "dirs": [str(claude_skill)], + }, + ], + ) + calls = [] + + def run(argv, **kwargs): + calls.append((argv, kwargs)) + stdout = "[]" if argv[1:4] == ["plugin", "list", "--cowork"] else '{"outcome":"ok"}' + if argv[1:3] == ["plugin", "validate"]: + stdout = "valid" + return SimpleNamespace(returncode=0, stdout=stdout, stderr="") + + monkeypatch.setattr(cowork.subprocess_cross_os, "run", run) + + assert cowork.configure_cowork( + {"workspace": "https://workspace.example", "profile": "chosen"}, + [ + { + "name": "github", + "url": "https://workspace.example/api/2.0/mcp/github", + "clients": ["claude", "codex"], + }, + { + "name": "codex-only", + "url": "https://workspace.example/api/2.0/mcp/codex", + "clients": ["codex"], + }, + ], + ) + + plugin_root = app_dir / cowork.MARKETPLACE_DIRECTORY_NAME / "plugins" / cowork.PLUGIN_NAME + plugin = json.loads((plugin_root / ".claude-plugin" / "plugin.json").read_text()) + assert plugin["mcpServers"] == "./.mcp.json" + mcp = json.loads((plugin_root / ".mcp.json").read_text()) + assert mcp == { + "mcpServers": { + "github": { + "type": "stdio", + "command": str(ug_binary), + "args": [ + "mcp-proxy", + "--url", + "https://workspace.example/api/2.0/mcp/github", + "--host", + "https://workspace.example", + "--profile", + "chosen", + ], + } + } + } + assert (plugin_root / "skills" / "review" / "SKILL.md").read_text() == "managed skill" + assert not (plugin_root / "skills" / "foreign").exists() + assert [call[0][1:3] for call in calls] == [ + ["plugin", "validate"], + ["plugin", "marketplace"], + ["plugin", "marketplace"], + ["plugin", "list"], + ["plugin", "install"], + ] + assert all(call[1]["env"]["CLAUDE_CONFIG_DIR"] == str(cowork_root) for call in calls) + assert json.loads((app_dir / cowork.STATE_FILENAME).read_text())["cowork_root"] == str( + cowork_root + ) + + +def test_configure_updates_installed_plugin(tmp_path, monkeypatch): + cowork_root = tmp_path / "cowork" + cowork_root.mkdir() + monkeypatch.setattr(cowork.config_io, "APP_DIR", tmp_path / ".ucode") + monkeypatch.setattr(cowork, "discover_cowork_root", lambda: cowork_root) + monkeypatch.setattr(cowork, "_cowork_sessions_root_exists", lambda: True) + bin_dir = tmp_path / "bin" + monkeypatch.setattr(cowork.shutil, "which", lambda name: str(bin_dir / name)) + monkeypatch.setattr(cowork, "records_for_scope", lambda *args: []) + calls = [] + + def run(argv, **kwargs): + calls.append(argv) + stdout = ( + json.dumps([{"pluginId": cowork.PLUGIN_ID}]) + if argv[1:4] == ["plugin", "list", "--cowork"] + else '{"outcome":"ok"}' + ) + if argv[1:3] == ["plugin", "validate"]: + stdout = "valid" + return SimpleNamespace(returncode=0, stdout=stdout, stderr="") + + monkeypatch.setattr(cowork.subprocess_cross_os, "run", run) + + assert cowork.configure_cowork( + {"workspace": "https://workspace.example", "profile": "chosen"}, + [{"name": "github", "url": "https://mcp.example", "clients": ["claude"]}], + ) + + assert calls[-1][1:3] == ["plugin", "update"] + + +def test_configure_records_root_before_plugin_command_failure(tmp_path, monkeypatch): + app_dir = tmp_path / ".ucode" + cowork_root = tmp_path / "cowork" + cowork_root.mkdir() + monkeypatch.setattr(cowork.config_io, "APP_DIR", app_dir) + monkeypatch.setattr(cowork, "discover_cowork_root", lambda: cowork_root) + monkeypatch.setattr(cowork, "_cowork_sessions_root_exists", lambda: True) + monkeypatch.setattr(cowork, "records_for_scope", lambda *args: []) + monkeypatch.setattr( + cowork, + "_cowork_mcp_config", + lambda *args, **kwargs: {"mcpServers": {"github": {"type": "stdio"}}}, + ) + + def fail_install(*args): + raise RuntimeError("plugin command failed") + + monkeypatch.setattr(cowork, "_install_plugin", fail_install) + + with pytest.raises(RuntimeError, match="plugin command failed"): + cowork.configure_cowork( + {"workspace": "https://workspace.example", "profile": "chosen"}, + [{"name": "github"}], + ) + + assert json.loads((app_dir / cowork.STATE_FILENAME).read_text()) == { + "version": 1, + "cowork_root": str(cowork_root), + } + + +def test_configure_removes_registration_from_previous_cowork_root(tmp_path, monkeypatch): + app_dir = tmp_path / ".ucode" + app_dir.mkdir(exist_ok=True) + old_root = tmp_path / "old-cowork" + new_root = tmp_path / "new-cowork" + (app_dir / cowork.STATE_FILENAME).write_text( + json.dumps({"version": 1, "cowork_root": str(old_root)}) + ) + removed_roots = [] + monkeypatch.setattr(cowork.config_io, "APP_DIR", app_dir) + monkeypatch.setattr(cowork, "discover_cowork_root", lambda: new_root) + monkeypatch.setattr(cowork, "records_for_scope", lambda *args: []) + monkeypatch.setattr( + cowork, + "_cowork_mcp_config", + lambda *args, **kwargs: {"mcpServers": {"github": {"type": "stdio"}}}, + ) + monkeypatch.setattr(cowork, "_remove_plugin_registration", removed_roots.append) + monkeypatch.setattr(cowork, "_write_marketplace", lambda *args: None) + monkeypatch.setattr(cowork, "_install_plugin", lambda *args: None) + + assert cowork.configure_cowork( + {"workspace": "https://workspace.example", "profile": "chosen"}, + [{"name": "github"}], + ) + + assert removed_roots == [old_root] + assert json.loads((app_dir / cowork.STATE_FILENAME).read_text())["cowork_root"] == str(new_root) + + +def test_configure_requires_absolute_ug_binary(monkeypatch): + monkeypatch.setattr(cowork, "build_mcp_proxy_argv", lambda *args, **kwargs: ["ug", "proxy"]) + + with pytest.raises(RuntimeError, match="absolute `ug` executable"): + cowork._cowork_mcp_config( + [{"name": "github", "url": "https://mcp.example", "clients": ["claude"]}], + "https://workspace.example", + "chosen", + ) + + +def test_cowork_mcp_config_uses_saved_pat_mode(tmp_path, monkeypatch): + ug_binary = tmp_path / "bin" / "ug" + monkeypatch.setattr( + cowork, + "build_mcp_proxy_argv", + lambda *args, **kwargs: ( + [str(ug_binary), "mcp-proxy", "--use-pat"] + if kwargs["use_pat"] + else [str(ug_binary), "mcp-proxy"] + ), + ) + + config = cowork._cowork_mcp_config( + [{"name": "github", "url": "https://mcp.example", "clients": ["claude"]}], + "https://workspace.example", + "chosen", + use_pat=True, + ) + + assert config["mcpServers"]["github"]["args"][-1] == "--use-pat" + + +def test_cowork_mcp_config_skips_custom_oauth(capsys): + config = cowork._cowork_mcp_config( + [{"name": "github", "url": "https://mcp.example", "clients": ["claude"]}], + "https://workspace.example", + "chosen", + custom_oauth=True, + ) + + assert config == {"mcpServers": {}} + assert "Skipping Cowork MCP server `github`" in capsys.readouterr().out + + +def test_empty_content_without_prior_plugin_or_cowork_sessions_skips_silently( + tmp_path, monkeypatch +): + app_dir = tmp_path / ".ucode" + monkeypatch.setattr(cowork.config_io, "APP_DIR", app_dir) + monkeypatch.setattr(cowork, "records_for_scope", lambda *args: []) + monkeypatch.setattr(cowork, "claude_desktop_root_directory", lambda: tmp_path / "Claude-3p") + monkeypatch.setattr( + cowork, + "discover_cowork_root", + lambda: pytest.fail("Cowork discovery should not run for empty content"), + ) + + assert not cowork.configure_cowork({"workspace": "https://workspace.example"}, []) + + +def test_empty_content_replaces_prior_plugin(tmp_path, monkeypatch): + app_dir = tmp_path / ".ucode" + app_dir.mkdir(exist_ok=True) + cowork_root = tmp_path / "cowork" + state_path = app_dir / cowork.STATE_FILENAME + state_path.write_text(json.dumps({"version": 1, "cowork_root": str(cowork_root)})) + monkeypatch.setattr(cowork.config_io, "APP_DIR", app_dir) + monkeypatch.setattr(cowork, "records_for_scope", lambda *args: []) + monkeypatch.setattr(cowork, "discover_cowork_root", lambda: cowork_root) + writes = [] + installs = [] + monkeypatch.setattr(cowork, "_write_marketplace", lambda *args: writes.append(args)) + monkeypatch.setattr(cowork, "_install_plugin", lambda *args: installs.append(args)) + + assert cowork.configure_cowork( + {"workspace": "https://workspace.example", "profile": "chosen"}, [] + ) + + assert writes[0][3] == {"mcpServers": {}} + assert installs[0][0] == cowork_root + + +def test_revert_uses_saved_root_and_removes_owned_files(tmp_path, monkeypatch): + app_dir = tmp_path / ".ucode" + app_dir.mkdir(exist_ok=True) + cowork_root = tmp_path / "cowork" + marketplace = app_dir / cowork.MARKETPLACE_DIRECTORY_NAME + marketplace.mkdir() + (app_dir / cowork.STATE_FILENAME).write_text( + json.dumps({"version": 1, "cowork_root": str(cowork_root)}) + ) + monkeypatch.setattr(cowork.config_io, "APP_DIR", app_dir) + bin_dir = tmp_path / "bin" + monkeypatch.setattr(cowork.shutil, "which", lambda name: str(bin_dir / name)) + calls = [] + + def run(argv, **kwargs): + calls.append((argv, kwargs)) + return SimpleNamespace(returncode=0, stdout='{"outcome":"ok"}', stderr="") + + monkeypatch.setattr(cowork.subprocess_cross_os, "run", run) + + assert cowork.revert_cowork() + assert [call[0][1:3] for call in calls] == [ + ["plugin", "uninstall"], + ["plugin", "marketplace"], + ] + assert all(call[1]["env"]["CLAUDE_CONFIG_DIR"] == str(cowork_root) for call in calls) + assert not marketplace.exists() + assert not (app_dir / cowork.STATE_FILENAME).exists() + + +def test_revert_accepts_idempotent_json_failure(tmp_path, monkeypatch): + app_dir = tmp_path / ".ucode" + app_dir.mkdir(exist_ok=True) + (app_dir / cowork.STATE_FILENAME).write_text( + json.dumps({"version": 1, "cowork_root": str(tmp_path / "cowork")}) + ) + monkeypatch.setattr(cowork.config_io, "APP_DIR", app_dir) + bin_dir = tmp_path / "bin" + monkeypatch.setattr(cowork.shutil, "which", lambda name: str(bin_dir / name)) + + def run(argv, **kwargs): + message = "Plugin not installed" if argv[1:3] == ["plugin", "uninstall"] else "Not found" + return SimpleNamespace( + returncode=0, + stdout=json.dumps({"outcome": "failed", "message": message}), + stderr="", + ) + + monkeypatch.setattr(cowork.subprocess_cross_os, "run", run) + + assert cowork.revert_cowork() + assert not (app_dir / cowork.STATE_FILENAME).exists() + + +def test_revert_removes_orphaned_registration_without_saved_state(tmp_path, monkeypatch): + app_dir = tmp_path / ".ucode" + desktop_root = tmp_path / "Claude-3p" + unrelated_root = _cowork_root(desktop_root, "account-a", "org-a") + orphaned_root = _cowork_root(desktop_root, "account-b", "org-b") + plugins_dir = orphaned_root / "cowork_plugins" + plugins_dir.mkdir() + (plugins_dir / "installed_plugins.json").write_text( + json.dumps({cowork.PLUGIN_ID: {"installPath": "/missing/source"}}) + ) + monkeypatch.setattr(cowork.config_io, "APP_DIR", app_dir) + monkeypatch.setattr(cowork, "claude_desktop_root_directory", lambda: desktop_root) + bin_dir = tmp_path / "bin" + monkeypatch.setattr(cowork.shutil, "which", lambda name: str(bin_dir / name)) + calls = [] + + def run(argv, **kwargs): + calls.append((argv, kwargs)) + return SimpleNamespace(returncode=0, stdout='{"outcome":"ok"}', stderr="") + + monkeypatch.setattr(cowork.subprocess_cross_os, "run", run) + + assert cowork.revert_cowork() + assert len(calls) == 2 + assert all(call[1]["env"]["CLAUDE_CONFIG_DIR"] == str(orphaned_root) for call in calls) + assert all(call[1]["env"]["CLAUDE_CONFIG_DIR"] != str(unrelated_root) for call in calls) + + +def test_revert_retains_owned_files_when_claude_cli_is_missing(tmp_path, monkeypatch): + app_dir = tmp_path / ".ucode" + app_dir.mkdir(exist_ok=True) + marketplace = app_dir / cowork.MARKETPLACE_DIRECTORY_NAME + marketplace.mkdir() + state_path = app_dir / cowork.STATE_FILENAME + state_path.write_text(json.dumps({"version": 1, "cowork_root": str(tmp_path / "cowork")})) + monkeypatch.setattr(cowork.config_io, "APP_DIR", app_dir) + monkeypatch.setattr(cowork.shutil, "which", lambda name: None) + + with pytest.raises(RuntimeError, match="Claude CLI is required"): + cowork.revert_cowork() + + assert marketplace.exists() + assert state_path.exists() diff --git a/tests/test_desktop_configure_cli.py b/tests/test_desktop_configure_cli.py index 63f75946d..cd09cc671 100644 --- a/tests/test_desktop_configure_cli.py +++ b/tests/test_desktop_configure_cli.py @@ -23,11 +23,13 @@ def configuration(monkeypatch): monkeypatch.setattr(cli, "install_databricks_ai_tools_for_agents", lambda *a, **k: None) monkeypatch.setattr(cli, "check_gateway_endpoint", lambda *a, **k: True) monkeypatch.setattr(cli, "_configure_managed_mcp_servers", lambda *a, **k: []) - monkeypatch.setattr(cli, "_configure_managed_skills", lambda *a, **k: None) + monkeypatch.setattr(cli, "_configure_managed_skills", lambda *a, **k: True) monkeypatch.setattr(cli, "_summarize_managed_config", lambda *a, **k: None) monkeypatch.setattr(cli, "refresh_managed_config", lambda *a, **k: (None, False)) monkeypatch.setattr(cli, "configure_single_tool", lambda *a, **k: state) monkeypatch.setattr(cli, "configure_selected_tools", lambda *a, **k: state) + monkeypatch.setattr(desktop_setup, "configure_claude_cowork", lambda *a, **k: None) + monkeypatch.setattr(desktop_setup, "revert_claude_cowork", lambda *a, **k: "unchanged") return state @@ -73,7 +75,9 @@ def test_successful_claude_setup_configures_desktop_once(configuration, monkeypa def test_failed_claude_is_not_setup_from_prior_available_tools(configuration, monkeypatch): configuration["last_configured_tools"] = ["codex"] configure = Mock() + revert_cowork = Mock(return_value="removed") monkeypatch.setattr(desktop_setup, "configure_claude_desktop", configure) + monkeypatch.setattr(desktop_setup, "revert_claude_cowork", revert_cowork) assert ( cli._configure_workspace_command( selected_tools=["claude", "codex"], @@ -82,11 +86,14 @@ def test_failed_claude_is_not_setup_from_prior_available_tools(configuration, mo == 0 ) configure.assert_not_called() + revert_cowork.assert_not_called() def test_codex_only_does_not_configure_desktop(configuration, monkeypatch): configure = Mock() + revert_cowork = Mock(return_value="removed") monkeypatch.setattr(desktop_setup, "configure_claude_desktop", configure) + monkeypatch.setattr(desktop_setup, "revert_claude_cowork", revert_cowork) assert ( cli._configure_workspace_command( tool="codex", workspaces=[("https://workspace.example", "chosen-profile")] @@ -94,6 +101,7 @@ def test_codex_only_does_not_configure_desktop(configuration, monkeypatch): == 0 ) configure.assert_not_called() + revert_cowork.assert_not_called() def test_self_managed_agent_add_does_not_configure_desktop(configuration, monkeypatch): @@ -129,6 +137,41 @@ def test_single_managed_setup_passes_catalog_scope_to_desktop(configuration, mon configure.assert_called_once_with(configuration, parent_schema="main.models") +def test_single_managed_claude_configures_cowork(configuration, monkeypatch): + managed = {"enabled_agents": {"claude": {}}} + servers = [{"name": "github", "clients": ["claude"]}] + configure_cowork = Mock() + configure_mcps = Mock(return_value=servers) + monkeypatch.setattr(desktop_setup, "configure_claude_desktop", lambda *a, **k: None) + monkeypatch.setattr(desktop_setup, "configure_claude_cowork", configure_cowork) + monkeypatch.setattr(cli, "_configure_managed_mcp_servers", configure_mcps) + monkeypatch.setattr(cli, "_configure_managed_skills", lambda configured: True) + monkeypatch.setattr(cli, "refresh_managed_config", lambda *a, **k: (managed, False)) + + assert ( + cli._configure_workspace_command( + tool="claude", workspaces=[("https://workspace.example", "chosen-profile")] + ) + == 0 + ) + configure_mcps.assert_called_once_with(managed) + configure_cowork.assert_called_once_with(configuration, servers) + + +def test_single_unmanaged_claude_reconciles_empty_cowork(configuration, monkeypatch): + configure_cowork = Mock() + monkeypatch.setattr(desktop_setup, "configure_claude_desktop", lambda *a, **k: None) + monkeypatch.setattr(desktop_setup, "configure_claude_cowork", configure_cowork) + + assert ( + cli._configure_workspace_command( + tool="claude", workspaces=[("https://workspace.example", "chosen-profile")] + ) + == 0 + ) + configure_cowork.assert_called_once_with(configuration, []) + + @pytest.mark.parametrize("flow", ["single", "managed", "selected"]) def test_desktop_failure_preserves_successful_configure_exit( configuration, monkeypatch, capsys, flow @@ -155,3 +198,100 @@ def fail(*a, **k): == 0 ) assert "Desktop profile is not writable" in capsys.readouterr().out + + +def test_managed_claude_configures_cowork_after_mcp_and_skill_sync(configuration, monkeypatch): + servers = [ + { + "name": "github", + "url": "https://workspace.example/api/2.0/mcp/github", + "clients": ["claude"], + } + ] + configure_cowork = Mock() + monkeypatch.setattr(desktop_setup, "configure_claude_desktop", lambda *a, **k: None) + monkeypatch.setattr(desktop_setup, "configure_claude_cowork", configure_cowork) + monkeypatch.setattr(cli, "_configure_managed_mcp_servers", lambda managed: servers) + monkeypatch.setattr(cli, "_configure_managed_skills", lambda managed: True) + monkeypatch.setattr( + cli, + "refresh_managed_config", + lambda *a, **k: ({"enabled_agents": {"claude": {}}}, False), + ) + + assert ( + cli._configure_workspace_command( + workspaces=[("https://workspace.example", "chosen-profile")] + ) + == 0 + ) + configure_cowork.assert_called_once_with(configuration, servers) + + +@pytest.mark.parametrize("failed_sync", ["mcp", "skills"]) +def test_managed_sync_failure_preserves_existing_cowork(configuration, monkeypatch, failed_sync): + configure_cowork = Mock() + monkeypatch.setattr(desktop_setup, "configure_claude_desktop", lambda *a, **k: None) + monkeypatch.setattr(desktop_setup, "configure_claude_cowork", configure_cowork) + monkeypatch.setattr( + cli, + "_configure_managed_mcp_servers", + lambda managed: None if failed_sync == "mcp" else [], + ) + monkeypatch.setattr(cli, "_configure_managed_skills", lambda managed: failed_sync != "skills") + monkeypatch.setattr( + cli, + "refresh_managed_config", + lambda *a, **k: ({"enabled_agents": {"claude": {}}}, False), + ) + + assert ( + cli._configure_workspace_command( + workspaces=[("https://workspace.example", "chosen-profile")] + ) + == 0 + ) + configure_cowork.assert_not_called() + + +def test_unmanaged_configure_syncs_user_mcp_servers_to_cowork(configuration, monkeypatch): + user_servers = [ + { + "name": "system-ai-github", + "url": "https://workspace.example/api/2.0/mcp/system.ai.github", + "clients": ["claude"], + } + ] + configuration["mcp_servers"] = user_servers + configure_cowork = Mock() + monkeypatch.setattr(desktop_setup, "configure_claude_desktop", lambda *a, **k: None) + monkeypatch.setattr(desktop_setup, "configure_claude_cowork", configure_cowork) + monkeypatch.setattr(cli, "load_state", lambda: configuration) + + assert ( + cli._configure_workspace_command( + selected_tools=["claude"], + workspaces=[("https://workspace.example", "chosen-profile")], + ) + == 0 + ) + configure_cowork.assert_called_once_with(configuration, user_servers) + + +def test_managed_config_without_claude_preserves_cowork(configuration, monkeypatch): + configuration["last_configured_tools"] = ["codex"] + revert_cowork = Mock(return_value="removed") + monkeypatch.setattr(desktop_setup, "revert_claude_cowork", revert_cowork) + monkeypatch.setattr( + cli, + "refresh_managed_config", + lambda *a, **k: ({"enabled_agents": {"codex": {}}}, False), + ) + + assert ( + cli._configure_workspace_command( + workspaces=[("https://workspace.example", "chosen-profile")] + ) + == 0 + ) + revert_cowork.assert_not_called() diff --git a/tests/test_desktop_setup.py b/tests/test_desktop_setup.py index 27dd90aec..564a0a26b 100644 --- a/tests/test_desktop_setup.py +++ b/tests/test_desktop_setup.py @@ -7,6 +7,7 @@ from types import SimpleNamespace import pytest +from typer.testing import CliRunner from ucode import cli from ucode import state as state_store @@ -682,7 +683,7 @@ def time_out(handle, *, timeout): assert "retry" in output -def test_revert_failure_clears_workspace_state_and_allows_retry(desktop, monkeypatch, capsys): +def test_revert_failure_retains_workspace_state_and_allows_retry(desktop, monkeypatch): previous_id = "12345678-1234-4abc-8def-123456789abc" (desktop / f"{previous_id}.json").write_text("{}") metadata_path = desktop / "_meta.json" @@ -707,17 +708,28 @@ def test_revert_failure_clears_workspace_state_and_allows_retry(desktop, monkeyp def lock_timeout(handle, *, timeout): raise TimeoutError("Lock is held by another operation") + runner = CliRunner() with monkeypatch.context() as blocked: blocked.setattr(desktop_config, "acquire_exclusive_file_lock", lock_timeout) - assert cli.revert() == 1 + failed = runner.invoke(cli.app, ["revert"]) + + assert failed.exit_code == 1, failed.output + warnings = [line for line in failed.output.splitlines() if line.startswith("! ")] + assert len(warnings) == 2, failed.output + assert "Could not revert Claude Desktop" in warnings[0] + assert ( + warnings[1] == "! ug state retained. Fix the Claude Desktop error and run ug revert again." + ) + assert "Claude Desktop profile: failed" in failed.output - assert state_store.load_state() == {} - assert state_store.load_full_state()["workspaces"] == {} + assert state_store.load_state()["workspace"] == state["workspace"] + assert state["workspace"] in state_store.load_full_state()["workspaces"] assert selection_path.read_bytes() == selection_before assert {path.name: path.read_bytes() for path in desktop.glob("*.json")} == profiles_before - assert "ug state cleared" in capsys.readouterr().out - assert cli.revert() == 0 + retried = runner.invoke(cli.app, ["revert"]) + assert retried.exit_code == 0, retried.output + assert state_store.load_state() == {} assert not (desktop / f"{claude_desktop.PROFILE_ID}.json").exists() assert not selection_path.exists() assert json.loads(metadata_path.read_text()) == { @@ -725,7 +737,8 @@ def lock_timeout(handle, *, timeout): "appliedId": previous_id, } assert (desktop / f"{previous_id}.json").read_text() == "{}" - assert cli.revert() == 0 + repeated = runner.invoke(cli.app, ["revert"]) + assert repeated.exit_code == 0, repeated.output def test_dry_run_does_not_discover_or_create_profiles(desktop, monkeypatch): @@ -782,6 +795,13 @@ def test_revert_without_desktop_ownership_leaves_existing_settings_untouched(des assert {path.name: path.read_bytes() for path in desktop.iterdir()} == before +def test_revert_cleans_cowork_when_native_desktop_is_unavailable(monkeypatch): + monkeypatch.setattr(desktop_setup, "claude_desktop_directory", lambda: None) + monkeypatch.setattr(desktop_setup, "revert_claude_cowork", lambda: "removed") + + assert desktop_setup.revert_desktop() == "removed" + + def test_reconfigure_recreates_deleted_profile_with_fixed_id(desktop, capsys): state = {"workspace": "https://workspace.example", "profile": "chosen-profile"} desktop_setup.configure_claude_desktop(state) diff --git a/tests/test_integration_evidence.py b/tests/test_integration_evidence.py index e0c22846d..6d97e80b1 100644 --- a/tests/test_integration_evidence.py +++ b/tests/test_integration_evidence.py @@ -176,6 +176,50 @@ def test_claude_model_identity_uses_assistant_answer_not_tool_output(): assert claude.completed_task_models(records, "value") == {"catalog.models.claude_sonnet"} +def test_claude_assistant_answers_include_subagent_handback(): + records = [ + { + "type": "assistant", + "isSidechain": True, + "message": { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "tool-use", + "name": "SubagentHandback", + "input": {"message": "child answer"}, + } + ], + }, + } + ] + + assert claude.assistant_answers(records) == ["child answer"] + + +def test_claude_assistant_answers_ignore_parent_handback(): + records = [ + { + "type": "assistant", + "isSidechain": False, + "message": { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "tool-use", + "name": "SubagentHandback", + "input": {"message": "not a child answer"}, + } + ], + }, + } + ] + + assert claude.assistant_answers(records) == [] + + def completed_records(agent, model, answer="value", turn_id="matching"): if agent == "claude": return [ diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index 8dba0aaa8..238f9ec18 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -6,6 +6,7 @@ import tomllib from contextlib import asynccontextmanager from pathlib import Path +from threading import Event import anyio import httpx @@ -280,6 +281,42 @@ def test_login_failure_becomes_a_proxy_auth_error(self, monkeypatch): with pytest.raises(mcp_proxy.ProxyAuthError, match="user cancelled"): self._drive(auth, httpx.Request("POST", CONN_URL), httpx.Response(401)) + def test_login_failure_requires_explicit_retry(self, monkeypatch): + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + monkeypatch.setattr(mcp_proxy, "mcp_service_needs_connection_login", lambda *a, **k: True) + outcomes = iter([(False, "user cancelled"), (True, "signed in")]) + monkeypatch.setattr(mcp_proxy, "run_connection_login", lambda *a, **k: next(outcomes)) + login = mcp_proxy._ConnectionLogin(CONN_URL, WS, "p") + + with pytest.raises(mcp_proxy.ProxyAuthError, match="user cancelled"): + login.run_if_needed() + + assert not login.run_if_needed() + assert login.authenticate_or_confirm_ready() + + def test_token_lookup_failure_can_be_retried(self, monkeypatch): + tokens = iter([RuntimeError("token unavailable"), "tok"]) + + def get_token(*args): + token = next(tokens) + if isinstance(token, Exception): + raise token + return token + + monkeypatch.setattr(mcp_proxy, "get_databricks_token", get_token) + monkeypatch.setattr(mcp_proxy, "mcp_service_needs_connection_login", lambda *a, **k: True) + logins = [] + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, **k: logins.append(1) or (True, "signed in"), + ) + login = mcp_proxy._ConnectionLogin(CONN_URL, WS, "p") + + assert not login.run_if_needed() + assert login.run_if_needed() + assert logins == [1] + def test_async_auth_flow_drives_login_on_401(self, monkeypatch): # The proxy uses the async client, so async_auth_flow is the real path; the # blocking login runs off the event loop via anyio.to_thread. @@ -306,81 +343,548 @@ def test_sync_auth_flow_is_rejected(self, monkeypatch): CONN_URL = f"{WS}/ai-gateway/mcp-services/system.ai.github" -class TestPump: - def test_forwards_all_messages_in_order(self): - async def scenario() -> list[str]: - src_send, src_recv = anyio.create_memory_object_stream(10) - dst_send, dst_recv = anyio.create_memory_object_stream(10) - # Preload the source, then close its send end so _pump's `async for` - # terminates once drained. - for msg in ["a", "b", "c"]: - await src_send.send(msg) - await src_send.aclose() - - await mcp_proxy._pump(src_recv, dst_send) - - received: list[str] = [] - # _pump closed dst_send on exit, so this drains then stops. - async with dst_recv: - async for msg in dst_recv: - received.append(msg) - return received - - assert anyio.run(scenario) == ["a", "b", "c"] - - def test_closes_destination_when_source_exhausts(self): - # A closed dest send-stream is what lets the *other* pump's reader - # terminate, so the bridge tears down cleanly when one side hangs up. - async def scenario() -> bool: - src_send, src_recv = anyio.create_memory_object_stream(1) - dst_send, dst_recv = anyio.create_memory_object_stream(1) - await src_send.aclose() - - await mcp_proxy._pump(src_recv, dst_send) - - with pytest.raises(anyio.EndOfStream): - dst_recv.receive_nowait() - return True - - assert anyio.run(scenario) is True - - def test_client_errors_are_forwarded(self): - async def scenario() -> Exception: - src_send, src_recv = anyio.create_memory_object_stream(1) - dst_send, dst_recv = anyio.create_memory_object_stream(1) - error = ValueError("malformed client message") - await src_send.send(error) - await src_send.aclose() - - await mcp_proxy._pump(src_recv, dst_send) - return await dst_recv.receive() - - error = anyio.run(scenario) - assert isinstance(error, ValueError) - assert str(error) == "malformed client message" - - def test_upstream_errors_are_raised(self): - async def scenario() -> None: - src_send, src_recv = anyio.create_memory_object_stream(1) - dst_send, _ = anyio.create_memory_object_stream(1) - await src_send.send(httpx.ReadTimeout("upstream timed out")) - await src_send.aclose() - - with pytest.raises(mcp_proxy.ProxyTransportError, match="upstream timed out"): - await mcp_proxy._pump_upstream(src_recv, dst_send) +class TestMcpBridgeConnectionLogin: + @staticmethod + def _request(request_id=1): + return mcp_proxy.SessionMessage( + mcp_proxy.JSONRPCRequest( + jsonrpc="2.0", + id=request_id, + method="tools/list", + ) + ) + + @staticmethod + def _tool_call_request(request_id=1, name="search_pull_requests"): + return mcp_proxy.SessionMessage( + mcp_proxy.JSONRPCRequest( + jsonrpc="2.0", + id=request_id, + method="tools/call", + params={"name": name, "arguments": {}}, + ) + ) + + @staticmethod + def _error(code=mcp_proxy.CONNECTION_AUTH_REQUIRED_ERROR_CODE, request_id=1): + from mcp.types import ErrorData + + return mcp_proxy.SessionMessage( + mcp_proxy.JSONRPCError( + jsonrpc="2.0", + id=request_id, + error=ErrorData(code=code, message="authentication required"), + ) + ) + + @staticmethod + def _response(request_id=1): + return mcp_proxy.SessionMessage( + mcp_proxy.JSONRPCResponse( + jsonrpc="2.0", + id=request_id, + result={"tools": []}, + ) + ) + + def test_auth_error_runs_login_replays_request_and_suppresses_first_error(self, monkeypatch): + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + monkeypatch.setattr(mcp_proxy, "mcp_service_needs_connection_login", lambda *a: True) + logins: list[tuple[str, str | None]] = [] + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda url, ws, **kwargs: ( + logins.append((url, kwargs.get("profile"))) or (True, "signed in") + ), + ) + + async def scenario(): + client_send, client_read = anyio.create_memory_object_stream(10) + to_client, client_receive = anyio.create_memory_object_stream(10) + from_upstream, upstream_read = anyio.create_memory_object_stream(10) + to_upstream, upstream_receive = anyio.create_memory_object_stream(10) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + request = self._tool_call_request() + response = self._response() + + async with anyio.create_task_group() as tg: + tg.start_soon(bridge.pump_client, client_read, to_upstream) + tg.start_soon(bridge.pump_upstream, upstream_read, to_client, to_upstream) + with anyio.fail_after(1): + await client_send.send(request) + assert await upstream_receive.receive() is request + await from_upstream.send(self._error()) + assert await upstream_receive.receive() is request + with pytest.raises(anyio.WouldBlock): + client_receive.receive_nowait() + await from_upstream.send(response) + assert await client_receive.receive() is response + tg.cancel_scope.cancel() + + anyio.run(scenario) + assert logins == [(CONN_URL, "p")] + + def test_non_tool_auth_error_is_forwarded_without_login(self, monkeypatch): + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, **k: pytest.fail("non-tool requests must not open a browser"), + ) + + async def scenario(): + from_upstream, upstream_read = anyio.create_memory_object_stream(1) + to_client, client_receive = anyio.create_memory_object_stream(1) + to_upstream, upstream_receive = anyio.create_memory_object_stream(1) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + bridge.pending_requests[1] = mcp_proxy.SessionMessage( + mcp_proxy.JSONRPCRequest( + jsonrpc="2.0", + id=1, + method="resources/list", + ) + ) + error = self._error() + + async with anyio.create_task_group() as tg: + tg.start_soon(bridge.pump_upstream, upstream_read, to_client, to_upstream) + with anyio.fail_after(1): + await from_upstream.send(error) + assert await client_receive.receive() is error + with pytest.raises(anyio.WouldBlock): + upstream_receive.receive_nowait() + tg.cancel_scope.cancel() + + anyio.run(scenario) + + def test_tools_list_auth_error_returns_authentication_tool_without_login(self, monkeypatch): + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, **k: pytest.fail("listing tools must not open a browser"), + ) + + async def scenario(): + from_upstream, upstream_read = anyio.create_memory_object_stream(1) + to_client, client_receive = anyio.create_memory_object_stream(1) + to_upstream, upstream_receive = anyio.create_memory_object_stream(1) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + bridge.pending_requests[1] = self._request() + + async with anyio.create_task_group() as tg: + tg.start_soon(bridge.pump_upstream, upstream_read, to_client, to_upstream) + with anyio.fail_after(1): + await from_upstream.send(self._error()) + response = (await client_receive.receive()).message + assert isinstance(response, mcp_proxy.JSONRPCResponse) + assert response.result["tools"][0]["name"] == "authenticate" + assert "same turn" in response.result["tools"][0]["description"] + with pytest.raises(anyio.WouldBlock): + upstream_receive.receive_nowait() + tg.cancel_scope.cancel() + + anyio.run(scenario) + + def test_initialize_advertises_tool_list_changes(self): + async def scenario(): + from_upstream, upstream_read = anyio.create_memory_object_stream(1) + to_client, client_receive = anyio.create_memory_object_stream(1) + to_upstream, _ = anyio.create_memory_object_stream(1) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + bridge.pending_requests[1] = mcp_proxy.SessionMessage( + mcp_proxy.JSONRPCRequest( + jsonrpc="2.0", + id=1, + method="initialize", + params={}, + ) + ) + response = mcp_proxy.SessionMessage( + mcp_proxy.JSONRPCResponse( + jsonrpc="2.0", + id=1, + result={"capabilities": {"tools": {}}}, + ) + ) + + async with anyio.create_task_group() as tg: + tg.start_soon(bridge.pump_upstream, upstream_read, to_client, to_upstream) + with anyio.fail_after(1): + await from_upstream.send(response) + received = (await client_receive.receive()).message + assert isinstance(received, mcp_proxy.JSONRPCResponse) + assert received.result["capabilities"]["tools"]["listChanged"] is True + tg.cancel_scope.cancel() + + anyio.run(scenario) + + def test_initialize_without_connection_preserves_capabilities(self): + async def scenario(): + to_client, client_receive = anyio.create_memory_object_stream(1) + from_upstream, upstream_read = anyio.create_memory_object_stream(1) + to_upstream, _ = anyio.create_memory_object_stream(1) + bridge = mcp_proxy._McpBridge( + mcp_proxy._ConnectionLogin(CONN_URL, WS, "p", use_pat=True) + ) + bridge.pending_requests[1] = mcp_proxy.SessionMessage( + mcp_proxy.JSONRPCRequest( + jsonrpc="2.0", + id=1, + method="initialize", + params={}, + ) + ) + response = mcp_proxy.SessionMessage( + mcp_proxy.JSONRPCResponse( + jsonrpc="2.0", + id=1, + result={"capabilities": {"tools": {}}}, + ) + ) + + async with anyio.create_task_group() as tg: + tg.start_soon(bridge.pump_upstream, upstream_read, to_client, to_upstream) + with anyio.fail_after(1): + await from_upstream.send(response) + assert await client_receive.receive() is response + tg.cancel_scope.cancel() + + anyio.run(scenario) + + def test_authentication_tool_runs_login_after_upstream_auth_error(self, monkeypatch): + monkeypatch.setattr( + mcp_proxy, + "mcp_service_needs_connection_login", + lambda *a: pytest.fail("explicit authentication must not trust classification"), + ) + logins = [] + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, **k: logins.append(1) or (True, "signed in"), + ) + + async def scenario(): + client_send, client_read = anyio.create_memory_object_stream(10) + to_client, client_receive = anyio.create_memory_object_stream(10) + from_upstream, upstream_read = anyio.create_memory_object_stream(10) + to_upstream, upstream_receive = anyio.create_memory_object_stream(10) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + bridge.authentication_required = True + + async with anyio.create_task_group() as tg: + tg.start_soon(bridge.pump_client, client_read, to_upstream, to_client) + tg.start_soon(bridge.pump_upstream, upstream_read, to_client, to_upstream) + with anyio.fail_after(1): + await client_send.send( + self._tool_call_request(2, mcp_proxy.AUTHENTICATE_TOOL_NAME) + ) + notification = (await client_receive.receive()).message + assert isinstance(notification, mcp_proxy.JSONRPCNotification) + assert notification.method == "notifications/tools/list_changed" + + refresh_request = self._request(3) + await client_send.send(refresh_request) + assert await upstream_receive.receive() is refresh_request + refreshed = self._response(3) + await from_upstream.send(refreshed) + assert await client_receive.receive() is refreshed + + authentication = (await client_receive.receive()).message + assert isinstance(authentication, mcp_proxy.JSONRPCResponse) + assert authentication.result["isError"] is False + tg.cancel_scope.cancel() + + anyio.run(scenario) + assert logins == [1] + + def test_authentication_tool_refreshes_tools_before_returning(self, monkeypatch): + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + monkeypatch.setattr(mcp_proxy, "mcp_service_needs_connection_login", lambda *a: True) + logins: list[int] = [] + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, **k: logins.append(1) or (True, "signed in"), + ) + + async def scenario(): + client_send, client_read = anyio.create_memory_object_stream(10) + to_client, client_receive = anyio.create_memory_object_stream(10) + from_upstream, upstream_read = anyio.create_memory_object_stream(10) + to_upstream, upstream_receive = anyio.create_memory_object_stream(10) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + bridge.authentication_required = True + + async with anyio.create_task_group() as tg: + tg.start_soon(bridge.pump_client, client_read, to_upstream, to_client) + tg.start_soon(bridge.pump_upstream, upstream_read, to_client, to_upstream) + with anyio.fail_after(1): + await client_send.send( + self._tool_call_request(2, mcp_proxy.AUTHENTICATE_TOOL_NAME) + ) + notification = (await client_receive.receive()).message + assert isinstance(notification, mcp_proxy.JSONRPCNotification) + assert notification.method == "notifications/tools/list_changed" + + refresh_request = self._request(3) + await client_send.send(refresh_request) + assert await upstream_receive.receive() is refresh_request + refreshed = mcp_proxy.SessionMessage( + mcp_proxy.JSONRPCResponse( + jsonrpc="2.0", + id=3, + result={"tools": [{"name": "search_pull_requests"}]}, + ) + ) + await from_upstream.send(refreshed) + assert await client_receive.receive() is refreshed + + authentication = (await client_receive.receive()).message + assert isinstance(authentication, mcp_proxy.JSONRPCResponse) + assert authentication.id == 2 + assert authentication.result["isError"] is False + assert "same turn" in authentication.result["content"][0]["text"] + tg.cancel_scope.cancel() + + anyio.run(scenario) + assert logins == [1] + + def test_authentication_tool_reports_tool_refresh_timeout(self, monkeypatch): + monkeypatch.setattr(mcp_proxy, "TOOLS_REFRESH_TIMEOUT_SECONDS", 0.01) + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, **k: (True, "signed in"), + ) + + async def scenario(): + to_client, client_receive = anyio.create_memory_object_stream(10) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + bridge.authentication_required = True + + await bridge._handle_authentication_tool(2, to_client) + notification = (await client_receive.receive()).message + response = (await client_receive.receive()).message + + assert isinstance(notification, mcp_proxy.JSONRPCNotification) + assert isinstance(response, mcp_proxy.JSONRPCResponse) + assert "did not refresh before the timeout" in response.result["content"][0]["text"] + assert "retry authentication" not in response.result["content"][0]["text"] + assert "start a new session" in response.result["content"][0]["text"] + assert bridge._tools_refresh_waiters == [] + + anyio.run(scenario) + + def test_second_auth_error_is_forwarded_without_another_login(self, monkeypatch): + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + monkeypatch.setattr(mcp_proxy, "mcp_service_needs_connection_login", lambda *a: True) + logins: list[int] = [] + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, **k: logins.append(1) or (True, "signed in"), + ) + + async def scenario(): + to_client, client_receive = anyio.create_memory_object_stream(10) + from_upstream, upstream_read = anyio.create_memory_object_stream(10) + to_upstream, upstream_receive = anyio.create_memory_object_stream(10) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + request = self._tool_call_request() + bridge.pending_requests[1] = request + + async with anyio.create_task_group() as tg: + tg.start_soon(bridge.pump_upstream, upstream_read, to_client, to_upstream) + with anyio.fail_after(1): + await from_upstream.send(self._error()) + assert await upstream_receive.receive() is request + second_error = self._error() + await from_upstream.send(second_error) + assert await client_receive.receive() is second_error + tg.cancel_scope.cancel() + + anyio.run(scenario) + assert logins == [1] + + def test_automatic_login_does_not_block_other_upstream_messages(self, monkeypatch): + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + monkeypatch.setattr(mcp_proxy, "mcp_service_needs_connection_login", lambda *a: True) + login_started = Event() + release_login = Event() + + def login(*args, **kwargs): + login_started.set() + assert release_login.wait(2) + return True, "signed in" + + monkeypatch.setattr(mcp_proxy, "run_connection_login", login) + + async def scenario(): + source_send, source_read = anyio.create_memory_object_stream(2) + dest_send, dest_read = anyio.create_memory_object_stream(2) + retry_send, retry_read = anyio.create_memory_object_stream(2) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + first_request = self._tool_call_request(1) + bridge.pending_requests = { + 1: first_request, + 2: self._tool_call_request(2), + } + later_response = self._response(2) + + async with anyio.create_task_group() as tg: + tg.start_soon(bridge.pump_upstream, source_read, dest_send, retry_send) + with anyio.fail_after(1): + await source_send.send(self._error(request_id=1)) + assert await anyio.to_thread.run_sync(login_started.wait, 1) + await source_send.send(later_response) + assert await dest_read.receive() is later_response + release_login.set() + assert await retry_read.receive() is first_request + tg.cancel_scope.cancel() anyio.run(scenario) - def test_upstream_eof_is_an_error(self): - async def scenario() -> None: - src_send, src_recv = anyio.create_memory_object_stream(1) - dst_send, _ = anyio.create_memory_object_stream(1) - await src_send.aclose() + def test_distinct_auth_errors_retry_after_one_login(self, monkeypatch): + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + monkeypatch.setattr(mcp_proxy, "mcp_service_needs_connection_login", lambda *a: True) + logins: list[int] = [] + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, **k: logins.append(1) or (True, "signed in"), + ) + + async def scenario(): + to_client, client_receive = anyio.create_memory_object_stream(10) + from_upstream, upstream_read = anyio.create_memory_object_stream(10) + to_upstream, upstream_receive = anyio.create_memory_object_stream(10) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + first_request = self._tool_call_request(1) + second_request = self._tool_call_request(2) + bridge.pending_requests = {1: first_request, 2: second_request} + + async with anyio.create_task_group() as tg: + tg.start_soon(bridge.pump_upstream, upstream_read, to_client, to_upstream) + with anyio.fail_after(1): + await from_upstream.send(self._error(request_id=1)) + assert await upstream_receive.receive() is first_request + await from_upstream.send(self._error(request_id=2)) + assert await upstream_receive.receive() is second_request + with pytest.raises(anyio.WouldBlock): + client_receive.receive_nowait() + tg.cancel_scope.cancel() + + anyio.run(scenario) + assert logins == [1] + + @pytest.mark.parametrize( + ("url", "use_pat"), + [ + (URL, False), + (CONN_URL, True), + ], + ) + def test_non_connection_and_pat_auth_errors_pass_through_without_login( + self, monkeypatch, url, use_pat + ): + logins: list[int] = [] + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, **k: logins.append(1) or (True, "signed in"), + ) + + async def scenario(): + source_send, source_read = anyio.create_memory_object_stream(1) + dest_send, dest_read = anyio.create_memory_object_stream(1) + retry_send, retry_read = anyio.create_memory_object_stream(1) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(url, WS, "p", use_pat=use_pat)) + bridge.pending_requests[1] = self._request() + error = self._error() + await source_send.send(error) + await source_send.aclose() + + with pytest.raises(mcp_proxy.ProxyTransportError, match="closed unexpectedly"): + await bridge.pump_upstream(source_read, dest_send, retry_send) + assert await dest_read.receive() is error + with pytest.raises(anyio.WouldBlock): + retry_read.receive_nowait() + + anyio.run(scenario) + assert logins == [] + + def test_other_jsonrpc_error_passes_through(self, monkeypatch): + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, **k: pytest.fail("login must not run"), + ) + + async def scenario(): + source_send, source_read = anyio.create_memory_object_stream(1) + dest_send, dest_read = anyio.create_memory_object_stream(1) + retry_send, _ = anyio.create_memory_object_stream(1) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + bridge.pending_requests[1] = self._request() + error = self._error(code=-32601) + await source_send.send(error) + await source_send.aclose() + + with pytest.raises(mcp_proxy.ProxyTransportError, match="closed unexpectedly"): + await bridge.pump_upstream(source_read, dest_send, retry_send) + assert await dest_read.receive() is error + + anyio.run(scenario) + + def test_automatic_login_failure_forwards_error_and_keeps_bridge_alive(self, monkeypatch): + monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + monkeypatch.setattr(mcp_proxy, "mcp_service_needs_connection_login", lambda *a: True) + monkeypatch.setattr( + mcp_proxy, + "run_connection_login", + lambda *a, **k: (False, "user cancelled"), + ) + + async def scenario(): + source_send, source_read = anyio.create_memory_object_stream(2) + dest_send, dest_read = anyio.create_memory_object_stream(2) + retry_send, _ = anyio.create_memory_object_stream(2) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + bridge.pending_requests = { + 1: self._tool_call_request(1), + 2: self._tool_call_request(2), + } + auth_error = self._error(request_id=1) + later_response = self._response(2) + + async with anyio.create_task_group() as tg: + tg.start_soon(bridge.pump_upstream, source_read, dest_send, retry_send) + with anyio.fail_after(1): + await source_send.send(auth_error) + assert await dest_read.receive() is auth_error + await source_send.send(later_response) + assert await dest_read.receive() is later_response + tg.cancel_scope.cancel() + + anyio.run(scenario) + + def test_notifications_are_forwarded_without_becoming_pending_requests(self): + from mcp.types import JSONRPCNotification + + async def scenario(): + source_send, source_read = anyio.create_memory_object_stream(1) + dest_send, dest_read = anyio.create_memory_object_stream(1) + bridge = mcp_proxy._McpBridge(mcp_proxy._ConnectionLogin(CONN_URL, WS, "p")) + notification = mcp_proxy.SessionMessage( + JSONRPCNotification(jsonrpc="2.0", method="notifications/initialized") + ) + await source_send.send(notification) + await source_send.aclose() - with pytest.raises( - mcp_proxy.ProxyTransportError, match="upstream MCP transport closed" - ): - await mcp_proxy._pump_upstream(src_recv, dst_send) + await bridge.pump_client(source_read, dest_send) + assert await dest_read.receive() is notification + assert bridge.pending_requests == {} anyio.run(scenario)