diff --git a/AGENTS.md b/AGENTS.md index d7bbf8654..8e2932d9b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -148,6 +148,12 @@ Desktop setup is configure-only and best-effort. Writes retain atomic replacemen failure rollback, and the Desktop-only five-second lock timeout. Desktop filters both static and discovered model IDs with case-insensitive `\bclaude-`; shared discovery and Claude Code lists are unchanged. If no IDs match, setup warns without writes. +When `claude_otel_tracing` is enabled, Desktop's reserved profile also owns +`otlpEndpoint`, `otlpProtocol`, `otlpHeadersHelper`, and `otlpTracesEnabled`. A generated +`.sh`/`.cmd` wrapper pins the configuring Python installation, workspace, and profile +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. See `docs/claude-desktop.md` for native validation limits and revert behavior.
diff --git a/README.md b/README.md index b801a9e34..ffd7fac7e 100644 --- a/README.md +++ b/README.md @@ -122,6 +122,8 @@ selects a dedicated Claude Desktop gateway profile. Its model list includes only Claude-named IDs; Claude Code's catalog is unchanged. Desktop setup errors warn without failing Claude Code configuration. See [Claude Desktop setup](docs/claude-desktop.md) for authentication, profile preservation, and native validation limits. +When managed Claude tracing is enabled, the same profile configures Desktop's native +OTLP trace export with a refreshing headers helper; no background service is needed. ## MCP Servers diff --git a/docs/claude-desktop.md b/docs/claude-desktop.md index de300cac9..5ad92a094 100644 --- a/docs/claude-desktop.md +++ b/docs/claude-desktop.md @@ -47,6 +47,60 @@ If authentication is missing, it opens browser login only when Desktop supplies a new Cowork task. Browser login has a bounded timeout and concurrent sign-ins are serialized. If the original message failed during login, resend it. +## OpenTelemetry tracing + +When the managed coding-agent configuration enables Claude tracing, `ug configure` +also enables Desktop's [native OTLP export](https://claude.com/docs/third-party/claude-desktop/configuration#otlp). +The profile sets: + +```json +{ + "otlpEndpoint": "https:///ai-gateway/otel", + "otlpProtocol": "http/protobuf", + "otlpHeadersHelper": "", + "otlpTracesEnabled": true +} +``` + +Desktop appends `/v1/traces` to the collector base URL. UG creates +`~/.ucode/claude-desktop-otel-headers.sh` on macOS/Linux or +`%USERPROFILE%\.ucode\claude-desktop-otel-headers.cmd` on Windows. This small wrapper +runs the configuring Python installation's `ug otel-headers` implementation with +the explicit workspace and CLI profile. It prints only a JSON `Authorization` +header, using the existing cached-token/silent-refresh mechanism. Neither the +profile nor wrapper stores a token. No proxy or background service is needed; +Codex's Linux-only OTEL service is unchanged. +The Windows wrapper selects UTF-8 and preserves literal arguments, including +non-ASCII profile names and trailing backslashes. + +Requirements: Desktop must support `otlpTracesEnabled` (added in 1.22209.0) and +`otlpHeadersHelper` (added in 1.30096.1), Claude tracing must be enabled in the managed +configuration, and the workspace must accept OTLP trace ingestion with the user's +credentials. Disable Claude tracing and rerun configure to remove the OTLP keys +and helper; `ug revert` also deletes the helper and UG profile. Writes and rollback +cover the wrapper alongside the profile and selection files. Other profiles remain +untouched. If Desktop's application-data directory has already been removed, +revert still deletes the helper and previous-selection state without recreating +Desktop directories. + +Restart Desktop and start a new Cowork task or Code session after configuring. +Code and Cowork sessions receive the headers helper and rerun it as tokens rotate, +rather than retaining an inference token copied at session start. Desktop caches +helper output for a few minutes, so rotation is not immediate. The telemetry helper +does not open browser login; if silent refresh fails, reauthenticate through Desktop's +inference sign-in or the Databricks CLI. Claude documents that helper failures can +send telemetry without authentication; check Desktop's logs. UG does not enable +raw-content capture. Desktop can also attempt logs and metrics exports to the collector base; +the gateway currently accepts only `/v1/traces`, not `/v1/logs` or `/v1/metrics`. + +Component tests verify the generated configuration and lifecycle, not native export +or ingestion. To validate natively, configure a tracing-enabled workspace, restart +Desktop, complete a small Cowork task, then confirm that its session spans reached +the workspace's configured trace table. Check Desktop logs for authentication or +trace-export errors, and verify further spans from the same open Code session after +token refresh and the helper cache interval. Do not share +unredacted credentials or diagnostic bundles. + ## Existing profiles and rollback Profiles use Anthropic's [documented locations](https://claude.com/docs/third-party/claude-desktop/data-storage#where-data-lives): @@ -119,7 +173,9 @@ 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, - connector/plugin compatibility, and native OTEL are not implemented here. + 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. - Observed inference applies to Cowork. Ordinary Chat is not claimed to use the same gateway path. diff --git a/src/ucode/agents/desktop/claude.py b/src/ucode/agents/desktop/claude.py index dfc54751e..f9e2a9c7b 100644 --- a/src/ucode/agents/desktop/claude.py +++ b/src/ucode/agents/desktop/claude.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import os import sys from collections.abc import Mapping, Sequence @@ -66,6 +67,7 @@ class _Paths: metadata: Path state: Path lock: Path + otel_helper: Path def claude_desktop_root_directory( @@ -104,11 +106,14 @@ def configure_claude_desktop( helper_args: Sequence[str], model_id_to_display_name: Mapping[str, str] | None = None, custom_headers: Mapping[str, str] | None = None, + otel_endpoint: str | None = None, + otel_headers_command: Sequence[str] = (), state_path: Path | str | None = None, profile_key: str | None = None, ) -> ConfigureResult: """Replace UG's reserved profile, register it, and select it.""" + paths = _resolve_paths(state_path=state_path) rendered = render_profile( model_ids=model_ids, model_id_to_display_name=model_id_to_display_name, @@ -116,8 +121,16 @@ def configure_claude_desktop( helper_command=helper_command, helper_args=helper_args, custom_headers=custom_headers, + otel_endpoint=otel_endpoint, + otel_headers_helper=paths.otel_helper, + ) + otel_helper = ( + desktop_config.render_otel_helper( + otel_headers_command, windows=sys.platform == "win32" + ).encode("utf-8") + if otel_endpoint + else None ) - paths = _resolve_paths(state_path=state_path) profile_path = paths.directory / f"{PROFILE_ID}.json" with desktop_config.locked(paths.lock): profile_before = desktop_config.read_optional_bytes(profile_path) @@ -136,13 +149,28 @@ def configure_claude_desktop( ) changed = desktop_config.write_configuration( [ - (profile_path, profile_before, rendered), + ( + paths.otel_helper, + desktop_config.read_optional_bytes(paths.otel_helper), + otel_helper, + 0o700, + ), + (profile_path, profile_before, _serialize_json(rendered, profile_before), 0o600), ( paths.metadata, metadata_before, - desktop_state.apply_metadata(metadata, profile_key=profile_key), + _serialize_json( + desktop_state.apply_metadata(metadata, profile_key=profile_key), + metadata_before, + ), + 0o600, + ), + ( + paths.state, + state_before, + _serialize_json({"prior_applied": prior_applied}, state_before), + 0o600, ), - (paths.state, state_before, {"prior_applied": prior_applied}), ] ) return ConfigureResult( @@ -157,9 +185,11 @@ def revert_claude_desktop(*, state_path: Path | str | None = None) -> RevertResu """Delete UG's reserved profile and registration, preserving unrelated profiles.""" directory = claude_desktop_directory() - if directory is None or not directory.exists(): + if directory is None: + return RevertResult(profile_id=None, changed=False) + paths = _resolve_paths(state_path=state_path, initialize=False) + if not any(path.exists() for path in (paths.directory, paths.state, paths.otel_helper)): return RevertResult(profile_id=None, changed=False) - paths = _resolve_paths(state_path=state_path) profile_path = paths.directory / f"{PROFILE_ID}.json" with desktop_config.locked(paths.lock): profile_before = desktop_config.read_optional_bytes(profile_path) @@ -180,21 +210,44 @@ def revert_claude_desktop(*, state_path: Path | str | None = None) -> RevertResu ) changed = desktop_config.write_configuration( [ - (profile_path, profile_before, None), - (paths.metadata, metadata_before, metadata_after), - (paths.state, state_before, None), + (profile_path, profile_before, None, 0o600), + ( + paths.metadata, + metadata_before, + _serialize_json(metadata_after, metadata_before), + 0o600, + ), + (paths.state, state_before, None, 0o600), + ( + paths.otel_helper, + desktop_config.read_optional_bytes(paths.otel_helper), + None, + 0o700, + ), ] ) return RevertResult(profile_id=PROFILE_ID if changed else None, changed=changed) -def _resolve_paths(*, state_path: Path | str | None) -> _Paths: +def _serialize_json(payload: Mapping[str, object] | None, before: bytes | None) -> bytes | None: + if payload is None: + return None + if before is not None: + try: + if json.loads(before) == payload: + return before + except (UnicodeError, json.JSONDecodeError): + pass + return json.dumps(payload, indent=2).encode("utf-8") + + +def _resolve_paths(*, state_path: Path | str | None, initialize: bool = True) -> _Paths: resolved_directory = claude_desktop_directory() if resolved_directory is None: raise ClaudeDesktopUnsupportedError( "Claude Desktop profile configuration is unavailable on this platform." ) - if not resolved_directory.exists(): + if initialize and not resolved_directory.exists(): can_initialize = resolved_directory.parent.is_dir() if sys.platform == "darwin": can_initialize = _claude_desktop_installed() @@ -210,7 +263,7 @@ def _resolve_paths(*, state_path: Path | str | None) -> _Paths: raise ClaudeDesktopUnavailableError( f"Could not initialize Claude Desktop profile directory at {resolved_directory}." ) from exc - if not resolved_directory.is_dir(): + if resolved_directory.exists() and not resolved_directory.is_dir(): raise ClaudeDesktopUnavailableError( f"Claude Desktop profile path is not a directory: {resolved_directory}." ) @@ -220,6 +273,9 @@ def _resolve_paths(*, state_path: Path | str | None) -> _Paths: metadata=resolved_directory / META_FILENAME, state=resolved_state, lock=resolved_state.with_name(f".{resolved_state.name}.lock"), + otel_helper=resolved_state.with_name( + "claude-desktop-otel-headers" + (".cmd" if sys.platform == "win32" else ".sh") + ), ) diff --git a/src/ucode/agents/desktop/claude_config.py b/src/ucode/agents/desktop/claude_config.py index 7d369bd4e..79a1516a9 100644 --- a/src/ucode/agents/desktop/claude_config.py +++ b/src/ucode/agents/desktop/claude_config.py @@ -4,6 +4,8 @@ import json import os +import shlex +import stat import tempfile import uuid from collections.abc import Iterator, Mapping, Sequence @@ -11,7 +13,6 @@ from pathlib import Path from typing import NotRequired, TypedDict, cast -from ucode.config_io import atomic_write_json from ucode.os_compatibility.file_lock_cross_os import ( acquire_exclusive_file_lock, release_file_lock, @@ -62,6 +63,8 @@ def render_profile( helper_args: Sequence[str], model_id_to_display_name: Mapping[str, str] | None = None, custom_headers: Mapping[str, str] | None = None, + otel_endpoint: str | None = None, + otel_headers_helper: Path | None = None, ) -> dict[str, object]: """Render resolved models, headers, and literal helper arguments in Desktop's schema.""" @@ -89,9 +92,34 @@ def render_profile( } if custom_headers is not None: profile["inferenceCustomHeaders"] = dict(custom_headers) + if otel_endpoint and otel_headers_helper is not None: + profile.update( + otlpEndpoint=otel_endpoint, + otlpProtocol="http/protobuf", + otlpHeadersHelper=os.fspath(otel_headers_helper), + otlpTracesEnabled=True, + ) return profile +def render_otel_helper(command: Sequence[str], *, windows: bool) -> str: + if not command: + raise ClaudeDesktopError("Claude Desktop OTEL headers command is required.") + if not windows: + return f"#!/bin/sh\nexec {shlex.join(command)}\n" + if any(character in argument for argument in command for character in '"\r\n\0'): + raise ClaudeDesktopError("Claude Desktop OTEL helper arguments contain invalid characters.") + # cmd.exe needs every arg quoted; the child parser needs trailing backslashes doubled. + quoted = " ".join( + '"' + + argument.replace("%", "%%") + + "\\" * (len(argument) - len(argument.rstrip("\\"))) + + '"' + for argument in command + ) + return f"@echo off\r\nsetlocal DisableDelayedExpansion\r\nchcp 65001 >nul\r\n{quoted}\r\n" + + @contextmanager def locked(lock_path: Path) -> Iterator[None]: try: @@ -176,35 +204,37 @@ def is_uuid(value: str) -> bool: def write_configuration( - changes: Sequence[tuple[Path, bytes | None, Mapping[str, object] | None]], + changes: Sequence[tuple[Path, bytes | None, bytes | None, int]], ) -> bool: - """Apply file changes, restoring exact original bytes if any write fails.""" + """Apply (path, before, after, mode) changes, restoring originals on failure.""" - for path, before, _after in changes: + for path, before, _after, _mode in changes: if read_optional_bytes(path) != before: raise ClaudeDesktopConflictError( "Claude Desktop configuration changed during setup. Quit Claude Desktop and retry." ) - originals = {path: before for path, before, _after in changes} - written: list[Path] = [] + original_modes = { + path: stat.S_IMODE(path.stat().st_mode) + for path, before, _after, _mode in changes + if before is not None + } + written: list[tuple[Path, bytes | None, int]] = [] try: - for path, before, after in changes: - contents = json.dumps(after, indent=2).encode("utf-8") if after is not None else None - if before == contents: + for path, before, after, mode in changes: + needs_exec_bit = ( + after is not None + and mode & stat.S_IXUSR + and not original_modes.get(path, mode) & stat.S_IXUSR + ) + if before == after and not needs_exec_bit: continue - if before is not None and after is not None: - try: - if json.loads(before) == after: - continue - except (UnicodeError, json.JSONDecodeError): - pass - written.append(path) - _write_or_remove(path, after) + written.append((path, before, original_modes.get(path, mode))) + _write_or_remove(path, after, mode=mode) except Exception as exc: # noqa: BLE001 - preserve the original failure through rollback rollback_errors: list[Exception] = [] - for path in reversed(written): + for path, before, mode in reversed(written): try: - _restore_original(path, originals[path]) + _write_or_remove(path, before, mode=mode) except Exception as rollback_exc: # noqa: BLE001 rollback_errors.append(rollback_exc) if rollback_errors: @@ -218,8 +248,8 @@ def write_configuration( return bool(written) -def _write_or_remove(path: Path, payload: Mapping[str, object] | None) -> None: - if payload is None: +def _write_or_remove(path: Path, contents: bytes | None, *, mode: int) -> None: + if contents is None: try: path.unlink() except FileNotFoundError: @@ -227,17 +257,12 @@ def _write_or_remove(path: Path, payload: Mapping[str, object] | None) -> None: except OSError as exc: raise ClaudeDesktopError(f"Could not remove Claude Desktop profile {path}.") from exc return - atomic_write_json(path, payload) - - -def _restore_original(path: Path, contents: bytes | None) -> None: - if contents is None: - _write_or_remove(path, None) - return + path.parent.mkdir(parents=True, exist_ok=True) fd, temporary = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp") try: with os.fdopen(fd, "wb") as handle: handle.write(contents) + os.chmod(temporary, mode) os.replace(temporary, path) finally: Path(temporary).unlink(missing_ok=True) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 1c57a1a2d..a1854ce15 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2463,7 +2463,8 @@ def otel_headers_cmd( ) raise typer.Exit(1) try: - token = get_databricks_token(workspace, profile, force_refresh=force_refresh) + with redirect_stdout(sys.stderr): + token = get_databricks_token(workspace, profile, force_refresh=force_refresh) except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None diff --git a/src/ucode/desktop_setup.py b/src/ucode/desktop_setup.py index 24979117d..d375d92db 100644 --- a/src/ucode/desktop_setup.py +++ b/src/ucode/desktop_setup.py @@ -101,6 +101,19 @@ def _configure_desktop(state: dict, *, parent_schema: str | None) -> None: profile, ], custom_headers=headers, + otel_endpoint=( + f"{workspace.rstrip('/')}/ai-gateway/otel" if state.get("claude_otel_tracing") else None + ), + otel_headers_command=[ + os.path.abspath(sys.executable), + "-m", + "ucode.cli", + "otel-headers", + "--host", + workspace, + "--profile", + profile, + ], profile_key=profile, ) print_success("Claude Desktop gateway profile configured. Quit and reopen Claude Desktop.") diff --git a/tests/README.md b/tests/README.md index 1af45c98c..619265af8 100644 --- a/tests/README.md +++ b/tests/README.md @@ -199,6 +199,22 @@ lists. Scoped discovery uses the same response for IDs and labels without extra Repeat configure refreshes or removes labels. Source lists/maps remain unchanged. These checks use temporary files on simulated platforms, not the native model picker. +Desktop OTEL component tests cover native trace settings gated by `claude_otel_tracing` +for static and discovered models on simulated macOS, Windows, and Linux. They verify +the collector base URL, protobuf protocol, refreshing `.sh`/`.cmd` headers helpers +without persisted tokens, literal argument quoting, repeated native helper execution +with rotated bearer overrides, Unicode and trailing-backslash arguments, and +JSON-only stdout/failure output. Windows wrappers select UTF-8 before the command. +Lifecycle checks cover repeat configure, workspace changes, disable/re-enable, rollback of helper bytes +and permissions at each write boundary, and revert preserving unrelated telemetry. +Revert also cleans up the helper and selection state after the Desktop library or +application-data root is removed, without recreating either directory; failed cleanup +restores helper bytes, permissions, and selection state for a later retry. +Byte-writer checks cover explicit permissions independent of contents, unchanged +JSON formatting, helper execute-permission repair, and failed atomic replacement +without leftover temporary files. +These checks do not prove native export, token refresh in open sessions, or ingestion. + `TestChildStdoutLaunch` in `test_cli.py` checks Claude `-p`/`--print` and Codex `exec` (plain and `--json`) and `app-server` status/error output on stderr, argument forwarding with and without ug's `--`, and an unchanged stdout descriptor. Detection checks exact tokens before diff --git a/tests/integration/README.md b/tests/integration/README.md index d5251e452..eda47a088 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -373,6 +373,22 @@ lists. Scoped discovery uses the same response for IDs and labels without extra Repeat configure refreshes or removes labels. Source lists/maps remain unchanged. These checks use temporary files on simulated platforms, not the native model picker. +Desktop OTEL coverage remains in `../test_claude_desktop.py` and +`../test_desktop_setup.py`: native settings follow Claude's tracing flag for both +model sources on simulated macOS, Windows, and Linux. Component checks cover the +collector base URL, protobuf protocol, refreshing `.sh`/`.cmd` helpers without +persisted tokens, quoting, repeated native helper execution with rotated bearer +overrides, Unicode and trailing-backslash arguments, Windows UTF-8 selection, +JSON-only stdout/failures, repeat configure, workspace changes, +disable/re-enable, helper-byte/permission rollback, and revert preserving unrelated +telemetry. Revert cleanup also covers a removed native library or application-data +root without recreating Desktop directories, plus rollback and retry after cleanup +failure. Byte-writer component cases check explicit permissions independent of +contents, unchanged JSON formatting, helper execute-permission repair, and +atomic-replacement failure cleanup. This suite does not launch +Desktop or establish native export, open-session token refresh, or trace ingestion. +See `../../docs/claude-desktop.md` for native validation steps and auth limits. + ## Test layout and format All user journeys are top-level tests. There is no separate regressions category: diff --git a/tests/test_claude_desktop.py b/tests/test_claude_desktop.py index acacf5b7c..244e86571 100644 --- a/tests/test_claude_desktop.py +++ b/tests/test_claude_desktop.py @@ -3,6 +3,11 @@ from __future__ import annotations import json +import os +import shlex +import stat +import subprocess +import sys import uuid from pathlib import Path from types import SimpleNamespace @@ -11,6 +16,7 @@ from ucode.agents.desktop import claude as desktop from ucode.agents.desktop import claude_config as desktop_config +from ucode.os_compatibility import subprocess_cross_os def _uuid() -> str: @@ -108,6 +114,168 @@ def test_render_profile_accepts_custom_headers(tmp_path: Path) -> None: assert profile["inferenceModels"] == [{"name": "system.ai.claude-opus-5-5"}] +def test_render_profile_uses_refreshing_otel_helper_without_storing_credentials(tmp_path): + profile = desktop.render_profile( + model_ids=("system.ai.claude-sonnet-5",), + base_url="https://workspace.example/ai-gateway/anthropic", + helper_command=tmp_path / "ug", + helper_args=("auth-helper",), + otel_endpoint="https://workspace.example/ai-gateway/otel", + otel_headers_helper=tmp_path / "otel-headers.sh", + ) + + assert {key: value for key, value in profile.items() if key.startswith("otlp")} == { + "otlpEndpoint": "https://workspace.example/ai-gateway/otel", + "otlpProtocol": "http/protobuf", + "otlpHeadersHelper": str(tmp_path / "otel-headers.sh"), + "otlpTracesEnabled": True, + } + assert "Authorization" not in json.dumps(profile) + + +def test_native_otel_helper_returns_current_headers_on_each_invocation(tmp_path): + helper = tmp_path / ( + "headers with spaces.cmd" if sys.platform == "win32" else "headers with spaces.sh" + ) + helper.write_bytes( + desktop_config.render_otel_helper( + [ + sys.executable, + "-m", + "ucode.cli", + "otel-headers", + "--host", + "https://workspace.example", + "--profile", + "selected & profile", + ], + windows=sys.platform == "win32", + ).encode("utf-8") + ) + helper.chmod(0o700) + env = { + **os.environ, + "HOME": str(tmp_path), + "USERPROFILE": str(tmp_path), + "PYTHONPATH": str(Path(__file__).resolve().parents[1] / "src"), + } + command = ( + [os.environ.get("COMSPEC", "cmd.exe"), "/d", "/c", str(helper)] + if sys.platform == "win32" + else [str(helper)] + ) + for token in ("first-component-token", "rotated-component-token"): + result = subprocess_cross_os.run( + command, + env={**env, "DATABRICKS_BEARER": token}, + capture_output=True, + text=True, + stdin=subprocess.DEVNULL, + timeout=30, + ) + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout) == {"Authorization": f"Bearer {token}"} + assert "component-token" not in helper.read_text() + + +def test_posix_otel_helper_preserves_literal_shell_arguments(tmp_path): + arguments = ["spaces & $HOME; 'quoted'", "line\nbreak", "100% !value!"] + command = [str(tmp_path / "python with spaces"), "otel-headers", *arguments] + helper = desktop_config.render_otel_helper(command, windows=False) + assert helper.startswith("#!/bin/sh\nexec ") + assert shlex.split(helper.removeprefix("#!/bin/sh\nexec ")) == command + + +def test_native_otel_helper_preserves_arguments_and_child_failure(tmp_path): + windows = sys.platform == "win32" + helper = tmp_path / ("headers with spaces.cmd" if windows else "headers with spaces.sh") + legacy_code_page = "@echo off\r\nchcp 437 >nul\r\n" if windows else "" + arguments = [ + "spaces in profile", + "profilé-工作区", + "a&b|ce^f(ghi)", + "%UG_WRAPPER_TEST%", + "bang!", + "'quoted'", + "C:\\dir\\", + "profile\\\\", + "after trailing backslashes", + ] + helper.write_bytes( + ( + legacy_code_page + + desktop_config.render_otel_helper( + [ + sys.executable, + "-c", + "import json,sys;print(json.dumps(sys.argv[1:]))", + *arguments, + ], + windows=windows, + ) + ).encode("utf-8") + ) + helper.chmod(0o700) + command = ( + [os.environ.get("COMSPEC", "cmd.exe"), "/d", "/c", str(helper)] + if windows + else [str(helper)] + ) + result = subprocess_cross_os.run( + command, + env={**os.environ, "UG_WRAPPER_TEST": "must not expand"}, + capture_output=True, + text=True, + stdin=subprocess.DEVNULL, + timeout=10, + ) + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout) == arguments + + helper.write_bytes( + desktop_config.render_otel_helper( + [sys.executable, "-c", "import sys;sys.stderr.write('Refresh failed');sys.exit(7)"], + windows=windows, + ).encode("utf-8") + ) + result = subprocess_cross_os.run( + command, + capture_output=True, + text=True, + stdin=subprocess.DEVNULL, + timeout=10, + ) + assert result.returncode == 7 + assert result.stdout == "" + assert result.stderr == "Refresh failed" + + +def test_windows_otel_helper_quotes_all_args_and_disables_expansion(): + assert desktop_config.render_otel_helper( + [r"C:\Program Files\Python!\python.exe", "--profile", "a&b|ce^f(ghi)%NAME%!"], + windows=True, + ) == ( + "@echo off\r\nsetlocal DisableDelayedExpansion\r\nchcp 65001 >nul\r\n" + '"C:\\Program Files\\Python!\\python.exe" "--profile" "a&b|ce^f(ghi)%%NAME%%!"\r\n' + ) + + +def test_windows_otel_helper_escapes_trailing_backslashes(): + helper = desktop_config.render_otel_helper( + ["python.exe", "--profile", "profile\\", "C:\\dir\\\\", "last"], windows=True + ) + + assert helper.splitlines()[-1] == ( + '"python.exe" "--profile" "profile\\\\" "C:\\dir\\\\\\\\" "last"' + ) + + +@pytest.mark.parametrize("argument", ['bad" & exit', "bad\rline", "bad\nline", "bad\0value"]) +def test_windows_otel_helper_rejects_unsafe_batch_arguments(argument): + with pytest.raises(desktop.ClaudeDesktopError, match="invalid characters"): + desktop_config.render_otel_helper(["python.exe", argument], windows=True) + + def test_render_profile_uses_catalog_labels_without_changing_routing_ids(tmp_path): model_ids = ["system.ai.claude-sonnet-5", "custom/claude-opus-5", "claude-haiku-4-5"] model_ids.append(model_ids[0]) @@ -249,6 +417,18 @@ def test_configure_leaves_unrelated_and_old_generated_profiles_untouched(tmp_pat assert (directory / name).read_bytes() == contents +def test_repeat_configure_preserves_equivalent_json_bytes(tmp_path, directory): + state_path = tmp_path / "state.json" + result = _configure(state_path) + paths = [result.profile_path, result.metadata_path, state_path] + for path in paths: + path.write_bytes((json.dumps(_read(path), sort_keys=True) + "\n").encode("utf-8")) + before = {path: path.read_bytes() for path in paths} + + assert not _configure(state_path).changed + assert {path: path.read_bytes() for path in paths} == before + + @pytest.mark.parametrize("profile_key", ["team profile", "研发", None]) def test_configure_replaces_renamed_entry_and_revert_removes_it(tmp_path, directory, profile_key): state_path = tmp_path / "state.json" @@ -432,12 +612,12 @@ def test_write_failure_rolls_back_exact_bytes( write = desktop_config._write_or_remove failed = False - def fail_once(path, payload): + def fail_once(path, contents, *, mode): nonlocal failed if path == targets[failed_file] and not failed: failed = True raise OSError("simulated write failure") - write(path, payload) + write(path, contents, mode=mode) monkeypatch.setattr(desktop_config, "_write_or_remove", fail_once) with pytest.raises(desktop.ClaudeDesktopError, match="rolled back"): @@ -491,10 +671,52 @@ def test_writer_rejects_concurrent_changes_without_overwrite(tmp_path): path = tmp_path / "profile.json" path.write_bytes(b"changed") with pytest.raises(desktop.ClaudeDesktopConflictError, match="changed during setup"): - desktop_config.write_configuration([(path, b"original", {"new": True})]) + desktop_config.write_configuration([(path, b"original", b'{"new": true}', 0o600)]) assert path.read_bytes() == b"changed" +@pytest.mark.parametrize("contents", [b'{"enabled": true}', b"#!/bin/sh\nexit 0\n", b"\xff\x00"]) +@pytest.mark.parametrize("mode", [0o600, 0o700]) +def test_writer_uses_explicit_permissions_independently_of_contents(tmp_path, contents, mode): + path = tmp_path / "nested" / "helper.cmd" + assert desktop_config.write_configuration([(path, None, contents, mode)]) + assert path.read_bytes() == contents + if os.name != "nt": + assert stat.S_IMODE(path.stat().st_mode) == mode + assert path.stat().st_mode & stat.S_IWUSR + assert not desktop_config.write_configuration([(path, contents, contents, mode)]) + + assert desktop_config.write_configuration([(path, contents, None, mode)]) + assert not path.exists() + + +def test_writer_failed_atomic_replace_restores_bytes_and_permissions(tmp_path, monkeypatch): + directory = tmp_path / "files" + directory.mkdir() + path = directory / "existing.cmd" + original = b"original\xff\x00" + path.write_bytes(original) + path.chmod(0o640) + original_mode = stat.S_IMODE(path.stat().st_mode) + new_path = directory / "new.json" + replace = os.replace + + def fail_new_file(source, destination): + if destination == new_path: + raise OSError("simulated atomic replacement failure") + replace(source, destination) + + monkeypatch.setattr(desktop_config.os, "replace", fail_new_file) + with pytest.raises(desktop.ClaudeDesktopError, match="rolled back"): + desktop_config.write_configuration( + [(path, original, b"replacement", 0o700), (new_path, None, b"{}", 0o600)] + ) + + assert path.read_bytes() == original + assert stat.S_IMODE(path.stat().st_mode) == original_mode + assert list(directory.iterdir()) == [path] + + @pytest.mark.parametrize("operation", ["configure", "revert"]) def test_profile_lock_timeout_preserves_files_and_closes_handle( tmp_path, directory, monkeypatch, operation diff --git a/tests/test_cli.py b/tests/test_cli.py index 45341ad23..23b231897 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2114,6 +2114,61 @@ def test_use_pat_honors_non_empty_bearer_env(self, monkeypatch): class TestOtelHeadersCommand: + def test_explicit_profile_refreshes_headers_and_keeps_diagnostics_off_stdout(self): + def token(workspace, profile, *, force_refresh): + assert (workspace, profile, force_refresh) == ( + "https://desktop.example", + "desktop-profile", + False, + ) + print("Authentication diagnostic") + return next(tokens) + + tokens = iter(["first-token", "refreshed-token"]) + with ( + patch( + "ucode.cli.load_state", + return_value={"workspace": "https://other.example", "profile": "other"}, + ), + patch("ucode.cli.get_databricks_token", side_effect=token), + ): + for expected in ("first-token", "refreshed-token"): + result = runner.invoke( + app, + [ + "otel-headers", + "--host", + "https://desktop.example", + "--profile", + "desktop-profile", + ], + ) + assert result.exit_code == 0 + assert json.loads(result.stdout) == {"Authorization": f"Bearer {expected}"} + assert "Authentication diagnostic" in result.stderr + + def test_failed_refresh_emits_no_header_or_token(self): + with ( + patch("ucode.cli.load_state", return_value={}), + patch( + "ucode.cli.get_databricks_token", + side_effect=RuntimeError("Reauthenticate with Databricks CLI"), + ), + ): + result = runner.invoke( + app, + [ + "otel-headers", + "--host", + "https://desktop.example", + "--profile", + "desktop-profile", + ], + ) + assert result.exit_code == 1 + assert result.stdout == "" + assert "Reauthenticate with Databricks CLI" in result.stderr + def test_prints_only_the_authorization_header_json(self): with ( patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), diff --git a/tests/test_desktop_setup.py b/tests/test_desktop_setup.py index 227448b31..8edfa8208 100644 --- a/tests/test_desktop_setup.py +++ b/tests/test_desktop_setup.py @@ -1,6 +1,9 @@ """Configure-only orchestration with real temporary Desktop profile files.""" import json +import os +import stat +from pathlib import Path from types import SimpleNamespace import pytest @@ -90,6 +93,240 @@ def test_configure_uses_claude_catalog_and_absolute_helper(desktop, tmp_path, ca assert "default-model" not in output +@pytest.mark.parametrize("source", ["discovery", "static"]) +@pytest.mark.parametrize("tracing", [None, False, True]) +def test_configure_otel_follows_claude_tracing_setting(desktop, source, tracing): + state = { + "workspace": "https://workspace.example/", + "profile": "chosen-profile", + "codex_otel_tracing": True, + } + if source == "static": + state["claude_static_models"] = ["system.ai.claude-sonnet-5"] + if tracing is not None: + state["claude_otel_tracing"] = tracing + original_state = dict(state) + + desktop_setup.configure_claude_desktop(state) + + profile = configured_profile(desktop) + assert {key: value for key, value in profile.items() if key.startswith("otlp")} == ( + { + "otlpEndpoint": "https://workspace.example/ai-gateway/otel", + "otlpProtocol": "http/protobuf", + "otlpHeadersHelper": str( + claude_desktop.APP_DIR + / ( + "claude-desktop-otel-headers.cmd" + if desktop_setup.sys.platform == "win32" + else "claude-desktop-otel-headers.sh" + ) + ), + "otlpTracesEnabled": True, + } + if tracing + else {} + ) + assert "component-token" not in json.dumps(profile) + helpers = list(claude_desktop.APP_DIR.glob("claude-desktop-otel-headers.*")) + assert len(helpers) == int(bool(tracing)) + if tracing: + helper = helpers[0].read_text() + assert "otel-headers" in helper + assert "https://workspace.example/" in helper + assert "chosen-profile" in helper + assert "component-token" not in helper + assert state == original_state + + +def test_otel_reconfigure_switch_disable_and_revert_preserve_other_profiles(desktop): + previous_id = "74074d15-5e9e-4465-a670-2d5bb1b83d00" + previous = {"otlpEndpoint": "https://other-collector.example", "otlpTracesEnabled": True} + previous_path = desktop / f"{previous_id}.json" + previous_path.write_text(json.dumps(previous)) + metadata_path = desktop / "_meta.json" + metadata = { + "entries": [{"id": previous_id, "name": "Other profile"}], + "appliedId": previous_id, + "unrelated": "keep", + } + metadata_path.write_text(json.dumps(metadata)) + state = { + "workspace": "https://workspace.example", + "profile": "chosen-profile", + "claude_otel_tracing": True, + } + desktop_setup.configure_claude_desktop(state) + helper_path = Path(configured_profile(desktop)["otlpHeadersHelper"]) + helper_before = helper_path.read_bytes() + before = {path.name: path.read_bytes() for path in desktop.glob("*.json")} + desktop_setup.configure_claude_desktop(state) + assert {path.name: path.read_bytes() for path in desktop.glob("*.json")} == before + assert helper_path.read_bytes() == helper_before + + desktop_setup.configure_claude_desktop({**state, "workspace": "https://next.example/"}) + profile = configured_profile(desktop) + assert profile["otlpEndpoint"] == "https://next.example/ai-gateway/otel" + assert profile["inferenceGatewayBaseUrl"] == "https://next.example/ai-gateway/anthropic" + assert "https://next.example/" in helper_path.read_text() + + desktop_setup.configure_claude_desktop({**state, "claude_otel_tracing": False}) + assert not any(key.startswith("otlp") for key in configured_profile(desktop)) + assert not helper_path.exists() + desktop_setup.configure_claude_desktop(state) + assert configured_profile(desktop)["otlpTracesEnabled"] is True + assert desktop_setup.revert_desktop() == "removed" + assert not (desktop / f"{claude_desktop.PROFILE_ID}.json").exists() + assert not helper_path.exists() + assert json.loads(previous_path.read_text()) == previous + assert json.loads(metadata_path.read_text()) == metadata + + +@pytest.mark.parametrize("removed_directory", ["library", "app-root"]) +def test_otel_revert_cleans_helper_after_desktop_removal(desktop, tmp_path, removed_directory): + desktop_setup.configure_claude_desktop( + { + "workspace": "https://workspace.example", + "profile": "chosen-profile", + "claude_otel_tracing": True, + } + ) + helper = Path(configured_profile(desktop)["otlpHeadersHelper"]) + state = claude_desktop.APP_DIR / claude_desktop.STATE_FILENAME + original_directory = desktop if removed_directory == "library" else desktop.parent + removed = original_directory.rename(tmp_path / "removed-desktop") + before = {path: path.read_bytes() for path in removed.rglob("*.json")} + + assert desktop_setup.revert_desktop() == "removed" + + assert not helper.exists() + assert not state.exists() + assert not original_directory.exists() + assert {path: path.read_bytes() for path in removed.rglob("*.json")} == before + assert desktop_setup.revert_desktop() == "unchanged" + assert not original_directory.exists() + + +@pytest.mark.parametrize("failed_file", ["helper", "state"]) +def test_otel_revert_after_desktop_removal_rolls_back(desktop, tmp_path, monkeypatch, failed_file): + desktop_setup.configure_claude_desktop( + { + "workspace": "https://workspace.example", + "profile": "chosen-profile", + "claude_otel_tracing": True, + } + ) + helper = Path(configured_profile(desktop)["otlpHeadersHelper"]) + state = claude_desktop.APP_DIR / claude_desktop.STATE_FILENAME + targets = {"helper": helper, "state": state} + before = {path: path.read_bytes() for path in targets.values()} + modes = {path: stat.S_IMODE(path.stat().st_mode) for path in targets.values()} + desktop.rename(tmp_path / "removed-desktop") + write = desktop_config._write_or_remove + + def fail_removal(path, contents, *, mode): + if path == targets[failed_file] and contents is None: + raise OSError("simulated removal failure") + write(path, contents, mode=mode) + + monkeypatch.setattr(desktop_config, "_write_or_remove", fail_removal) + + assert desktop_setup.revert_desktop() == "failed" + + assert {path: path.read_bytes() for path in targets.values()} == before + assert {path: stat.S_IMODE(path.stat().st_mode) for path in targets.values()} == modes + assert not desktop.exists() + monkeypatch.setattr(desktop_config, "_write_or_remove", write) + assert desktop_setup.revert_desktop() == "removed" + assert not helper.exists() + assert not state.exists() + assert not desktop.exists() + + +@pytest.mark.parametrize("failed_file", ["helper", "profile", "metadata", "state"]) +@pytest.mark.parametrize("operation", ["enable", "switch", "disable", "revert"]) +def test_otel_write_failure_rolls_back_existing_profile( + desktop, monkeypatch, failed_file, operation, capsys +): + state = { + "workspace": "https://workspace.example", + "profile": "chosen-profile", + "claude_otel_tracing": operation != "enable", + } + desktop_setup.configure_claude_desktop(state) + metadata_path = desktop / "_meta.json" + metadata = json.loads(metadata_path.read_text()) + previous_id = "74074d15-5e9e-4465-a670-2d5bb1b83d00" + (desktop / f"{previous_id}.json").write_text("{}") + metadata["entries"].append({"id": previous_id, "name": "Other profile"}) + metadata["appliedId"] = previous_id + metadata_path.write_text(json.dumps(metadata)) + selection_path = claude_desktop.APP_DIR / claude_desktop.STATE_FILENAME + targets = { + "helper": claude_desktop.APP_DIR + / ( + "claude-desktop-otel-headers.cmd" + if desktop_setup.sys.platform == "win32" + else "claude-desktop-otel-headers.sh" + ), + "profile": desktop / f"{claude_desktop.PROFILE_ID}.json", + "metadata": metadata_path, + "state": selection_path, + } + before = {path: path.read_bytes() if path.exists() else None for path in targets.values()} + if operation == "switch": + targets["helper"].chmod(0o600) + modes = {path: stat.S_IMODE(path.stat().st_mode) for path in targets.values() if path.exists()} + write = desktop_config._write_or_remove + failed = False + + def fail_once(path, contents, *, mode): + nonlocal failed + if path == targets[failed_file] and not failed: + failed = True + raise OSError("simulated write failure") + write(path, contents, mode=mode) + + monkeypatch.setattr(desktop_config, "_write_or_remove", fail_once) + if operation == "revert": + assert desktop_setup.revert_desktop() == "failed" + else: + desktop_setup.configure_claude_desktop( + { + **state, + "workspace": "https://next.example" + if operation == "switch" + else state["workspace"], + "claude_otel_tracing": operation != "disable", + } + ) + + assert failed + assert { + path: path.read_bytes() if path.exists() else None for path in targets.values() + } == before + assert {path: stat.S_IMODE(path.stat().st_mode) for path in modes} == modes + assert "rolled back" in capsys.readouterr().out + + +def test_reconfigure_repairs_otel_helper_execute_permission(desktop): + state = { + "workspace": "https://workspace.example", + "profile": "chosen-profile", + "claude_otel_tracing": True, + } + desktop_setup.configure_claude_desktop(state) + helper = Path(configured_profile(desktop)["otlpHeadersHelper"]) + original = helper.read_bytes() + helper.chmod(0o600) + desktop_setup.configure_claude_desktop(state) + + assert helper.read_bytes() == original + if os.name != "nt": + assert stat.S_IMODE(helper.stat().st_mode) == 0o700 + assert helper.stat().st_mode & stat.S_IWUSR + + @pytest.mark.parametrize("source", ["discovery", "static"]) def test_configure_filters_claude_names_without_changing_source(desktop, monkeypatch, source): claude_models = [