diff --git a/agents/build-repair/compose.yml b/agents/build-repair/compose.yml index ce91cfabe8..fcf5d93fb6 100644 --- a/agents/build-repair/compose.yml +++ b/agents/build-repair/compose.yml @@ -28,6 +28,8 @@ services: - BUGZILLA_MCP_URL=http://build-repair-broker:8765/mcp - SOURCE_REPO=/workspace/firefox - ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY:?error} + # Weave tracing locally: an API key enables it (deploys use W&B WIF instead). + - WANDB_API_KEY=${WANDB_API_KEY:-} # No uploader locally: summary/logs/artifacts are written under # /artifacts/, bind-mounted to the host's ~/hackbot/artifacts. - ARTIFACTS_DIR=/artifacts diff --git a/docs/hackbot/tracing.md b/docs/hackbot/tracing.md new file mode 100644 index 0000000000..e6f8610b02 --- /dev/null +++ b/docs/hackbot/tracing.md @@ -0,0 +1,33 @@ +## Tracing (Weave) + +Dashboards: [prod](https://wandb.ai/moz-bugbug/hackbot-prod/weave/agents), [dev](https://wandb.ai/moz-bugbug/hackbot-dev), [test](https://wandb.ai/moz-bugbug/hackbot-test) + +Tracing is handled once in the runtime (`hackbot_runtime.tracing`), so every agent +gets it for free — there's nothing to add per agent. On startup the runtime calls +`weave.init()`, which autopatches the Claude Agent SDK (or other supported frameworks) and captures each query, +model response, and tool call, labelled with the agent's name (so the Weave +**Agents** view shows `build-repair`, `bug-fix`, etc. instead of a generic +`claude_agent_sdk`). + +It is **opt-in**: the runtime only traces when it has W&B credentials, and never +fails a run if Weave can't start. `WEAVE_PROJECT` picks the destination project +(a bare `project` or `entity/project`; defaults to `hackbot-test`). + +**Locally**, add the key to your root `.env`; the agent's `compose.yml` should +pass `WANDB_API_KEY` through to the agent container. + +`.env`: + +```dotenv +WANDB_API_KEY=... +``` + +`compose.yml`: + +```yaml +environment: + - WANDB_API_KEY=${WANDB_API_KEY:-} +``` + +**In deployment**, the agent container holds no long-lived key: it authenticates +via W&B [Identity Federation](https://docs.wandb.ai/platform/hosting/iam/identity_federation). diff --git a/libs/hackbot-runtime/hackbot_runtime/runtime.py b/libs/hackbot-runtime/hackbot_runtime/runtime.py index de4eb6d7c8..4da970f380 100644 --- a/libs/hackbot-runtime/hackbot_runtime/runtime.py +++ b/libs/hackbot-runtime/hackbot_runtime/runtime.py @@ -9,7 +9,7 @@ from pydantic import ValidationError -from hackbot_runtime import anthropic_wif +from hackbot_runtime import anthropic_wif, tracing, wandb_wif from hackbot_runtime.config import HackbotConfig, load_config from hackbot_runtime.context import HackbotContext from hackbot_runtime.results import HackbotAgentResult @@ -34,13 +34,16 @@ def _configure_auth() -> None: - """Set up the model-provider credentials before the agent runs. + """Set up the model-provider and observability credentials before the run. - For now only Anthropic WIF is supported; this is where other providers will - be wired in when we start supporting them. + Both Anthropic (model access) and W&B (Weave tracing) support GCP Workload + Identity Federation, so in deployment neither needs a long-lived key in the + agent container; both fall back to their API-key env var locally. """ if anthropic_wif.configure(): log.info("Configured Anthropic WIF authentication") + if wandb_wif.configure(): + log.info("Configured W&B WIF authentication") def _configure_logging() -> None: @@ -189,7 +192,8 @@ def run(entrypoint: AgentMain, config: ConfigArg = None) -> NoReturn: try: _configure_auth() - outcome: object = entrypoint(ctx) + with tracing.trace_agent(entrypoint): + outcome: object = entrypoint(ctx) except Exception as exc: log.exception("Agent raised an exception") outcome = exc @@ -205,7 +209,8 @@ def run_async(entrypoint: AsyncAgentMain, config: ConfigArg = None) -> NoReturn: try: _configure_auth() - outcome: object = asyncio.run(entrypoint(ctx)) + with tracing.trace_agent(entrypoint): + outcome: object = asyncio.run(entrypoint(ctx)) except Exception as exc: log.exception("Agent raised an exception") outcome = exc diff --git a/libs/hackbot-runtime/hackbot_runtime/tracing.py b/libs/hackbot-runtime/hackbot_runtime/tracing.py new file mode 100644 index 0000000000..83ba0f71bd --- /dev/null +++ b/libs/hackbot-runtime/hackbot_runtime/tracing.py @@ -0,0 +1,82 @@ +"""Weights & Biases Weave tracing for all hackbot agents. + +``weave.init()`` autopatches the Claude Agent SDK, so calling it once in the +runtime before an agent's ``main()`` runs captures every query, model response, +and tool call as a trace with no per-agent instrumentation. It authenticates +either from ``WANDB_API_KEY`` (local/dev) or, in deployment, from the short-lived +identity token that :mod:`hackbot_runtime.wandb_wif` writes -- so the agent +container needs no long-lived W&B credential. Tracing is opt-in: it activates +only when one of those is present, and never fails the run if Weave can't start. + +The SDK integration otherwise labels every agent ``claude_agent_sdk``; +``agent_name_override`` relabels the spans with the running agent's name so the +dashboard can tell the agents apart. +""" + +import contextlib +import inspect +import logging +import os +from collections.abc import Callable, Iterator +from pathlib import Path + +from hackbot_runtime import wandb_wif + +log = logging.getLogger("hackbot_runtime") + +# Weave project traces land in when the deploy doesn't set WEAVE_PROJECT. Accepts +# either "project" or "entity/project". +DEFAULT_WEAVE_PROJECT = "hackbot-test" + + +def resolve_agent_name(entrypoint: Callable) -> str: + """The running agent's name, derived from its ``main()`` source file. + + Every agent's ``main()`` lives in ``hackbot_agents//__main__.py``, so + the directory holding it is the agent's package. We read the source file path + (not ``__module__``, which is just ``"__main__"`` under ``python -m``) and + take that directory name, normalized to the canonical hyphenated form + (``build_repair`` -> ``build-repair``). + """ + source = inspect.getfile(entrypoint) + return Path(source).parent.name.replace("_", "-") + + +def _init_weave() -> bool: + """Initialize Weave when credentials are available; return whether enabled. + + Credentials come from ``WANDB_API_KEY`` (local) or the identity token file set + by :mod:`wandb_wif` (deployment via federation). No-op when neither is present. + """ + if not ( + os.environ.get("WANDB_API_KEY") or os.environ.get(wandb_wif.TOKEN_FILE_ENV) + ): + return False + + project = os.environ.get("WEAVE_PROJECT", DEFAULT_WEAVE_PROJECT) + try: + import weave + + weave.init(project) + return True + except Exception: + log.exception("Failed to initialize Weave tracing; continuing without it") + return False + + +@contextlib.contextmanager +def trace_agent(entrypoint: Callable) -> Iterator[None]: + """Trace the agent run and label its Weave spans with the agent's name. + + A no-op when tracing isn't configured (no W&B credentials). + """ + if not _init_weave(): + yield + return + + from weave.conversation import agent_name_override + + agent = resolve_agent_name(entrypoint) + log.info("Enabled Weave tracing for agent %s", agent) + with agent_name_override(agent): + yield diff --git a/libs/hackbot-runtime/hackbot_runtime/wandb_wif.py b/libs/hackbot-runtime/hackbot_runtime/wandb_wif.py new file mode 100644 index 0000000000..91eceaa3ab --- /dev/null +++ b/libs/hackbot-runtime/hackbot_runtime/wandb_wif.py @@ -0,0 +1,205 @@ +"""GCP Workload Identity Federation for Weights & Biases (Weave) auth. + +W&B identity federation lets the wandb/weave SDK authenticate from a short-lived +JWT instead of a long-lived ``WANDB_API_KEY``: the SDK reads a JWT from the file +named by ``WANDB_IDENTITY_TOKEN_FILE`` and exchanges it for a W&B access token. On +GCP we mint that JWT as a Google-signed OIDC identity token (from the metadata +server, via ``google-auth``), so the agent container holds no W&B credential at +all -- the same shape as :mod:`hackbot_runtime.anthropic_wif`. + + ``WANDB_FEDERATION_AUDIENCE`` audience for the Google JWT ── set by deploy + ``WANDB_IDENTITY_TOKEN_FILE`` path to the Google JWT ── managed here + ``WANDB_IDENTITY_TOKEN_REFRESH_SECONDS`` refresh cadence (optional) ── set by deploy + +The deploy provisions the audience (which must match the W&B org's federation +trust config); this module owns the token file. It fetches the identity token, +writes it to a private file, points ``WANDB_IDENTITY_TOKEN_FILE`` at it, and keeps +it fresh in a background thread — Google identity tokens live ~1h and the SDK +re-reads the file when it re-exchanges, so a periodically-rewritten file keeps +auth alive for runs longer than a token's lifetime. + +When ``WANDB_FEDERATION_AUDIENCE`` is absent the runtime is in API-key mode and +this module is inert, so local/compose runs keep using ``WANDB_API_KEY``. + +See https://docs.wandb.ai/platform/hosting/iam/identity_federation +""" + +from __future__ import annotations + +import logging +import os +import tempfile +import threading +from pathlib import Path + +import google.auth.transport.requests +import google.oauth2.id_token +from google.auth.exceptions import GoogleAuthError + +from hackbot_runtime.providers import ProviderError + +log = logging.getLogger("hackbot_runtime.wandb_wif") + +# Presence of the audience is what flips the runtime from API-key mode to WIF +# mode — the deploy sets it only on federation-enabled workloads. Its value is +# the audience the Google identity token is minted for, which must match the +# W&B org's federation trust configuration. +AUDIENCE_ENV = "WANDB_FEDERATION_AUDIENCE" +TOKEN_FILE_ENV = "WANDB_IDENTITY_TOKEN_FILE" +REFRESH_INTERVAL_ENV = "WANDB_IDENTITY_TOKEN_REFRESH_SECONDS" + +# Google identity tokens expire after ~1h; refresh well inside that window so +# the file always carries plenty of remaining lifetime for the SDK to exchange. +DEFAULT_REFRESH_INTERVAL = 1800 + +# Set once configure() succeeds; keeps the daemon refresher alive for the +# process lifetime and makes a second configure() call a no-op. +_refresher: _TokenFileRefresher | None = None + + +def is_enabled() -> bool: + """True when the deploy provisioned a federation audience (WIF mode).""" + return bool(os.environ.get(AUDIENCE_ENV)) + + +def fetch_gcp_identity_token(audience: str) -> str: + """Fetch a Google-signed OIDC identity token for ``audience``. + + Delegates to ``google.oauth2.id_token``, which sources the token from the + GCE/Cloud Run/GKE metadata server (or a ``GOOGLE_APPLICATION_CREDENTIALS`` + service-account file for local testing). + """ + request = google.auth.transport.requests.Request() + try: + token = google.oauth2.id_token.fetch_id_token(request, audience) + except GoogleAuthError as exc: + raise ProviderError( + f"Failed to fetch a Google identity token for W&B federation: {exc}" + ) from exc + if not token: + raise ProviderError( + "Google returned an empty identity token; the workload has no usable " + "service account." + ) + return token + + +def _refresh_interval() -> float: + raw = os.environ.get(REFRESH_INTERVAL_ENV) + if not raw: + return DEFAULT_REFRESH_INTERVAL + try: + value = float(raw) + except ValueError: + log.warning( + "%s=%r is not a number; falling back to %ss", + REFRESH_INTERVAL_ENV, + raw, + DEFAULT_REFRESH_INTERVAL, + ) + return DEFAULT_REFRESH_INTERVAL + if value <= 0: + log.warning( + "%s=%s is not positive; falling back to %ss", + REFRESH_INTERVAL_ENV, + raw, + DEFAULT_REFRESH_INTERVAL, + ) + return DEFAULT_REFRESH_INTERVAL + return value + + +def _default_token_path() -> Path: + """A private, process-owned path for the identity token file. + + ``mkdtemp`` gives a 0700 directory so the bearer token isn't world-readable. + """ + return Path(tempfile.mkdtemp(prefix="wandb-wif-")) / "identity-token" + + +class _TokenFileRefresher: + """Keeps ``token_file`` populated with a fresh Google identity token. + + The first write happens synchronously in :meth:`start` so a metadata-server + failure surfaces immediately (fail fast) rather than as an opaque auth error + once the agent is mid-run. Subsequent writes run on a daemon thread. + """ + + def __init__(self, token_file: Path, interval: float, audience: str): + self._token_file = token_file + self._interval = interval + self._audience = audience + self._stop = threading.Event() + self._thread = threading.Thread( + target=self._loop, name="wandb-wif-refresh", daemon=True + ) + + def _write_token(self) -> None: + """Atomically replace the token file so the SDK never reads a partial write.""" + token = fetch_gcp_identity_token(self._audience) + self._token_file.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(dir=self._token_file.parent, prefix=".identity-") + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + handle.write(token) + os.replace(tmp, self._token_file) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise + + def _loop(self) -> None: + while not self._stop.wait(self._interval): + try: + self._write_token() + except (ProviderError, OSError) as exc: + log.warning( + "Failed to refresh W&B WIF identity token; " + "serving the previously written token: %s", + exc, + ) + + def start(self) -> None: + self._write_token() + self._thread.start() + + def stop(self) -> None: + self._stop.set() + + +def configure() -> bool: + """Wire up W&B WIF auth for the weave SDK when the deploy enables it. + + Returns ``True`` when WIF was configured, ``False`` in API-key mode. Idempotent: + a second call while the refresher is already running is a no-op. + """ + global _refresher + if not is_enabled(): + return False + if _refresher is not None: + return True + + if os.environ.get("WANDB_API_KEY"): + log.error( + "WANDB_API_KEY is set while Workload Identity Federation is configured; " + "the API key takes precedence and shadows WIF. Unset it if you intend " + "to authenticate via federation." + ) + return False + + audience = os.environ[AUDIENCE_ENV] + token_file = Path(os.environ.get(TOKEN_FILE_ENV) or _default_token_path()) + interval = _refresh_interval() + refresher = _TokenFileRefresher(token_file, interval, audience) + refresher.start() + os.environ[TOKEN_FILE_ENV] = str(token_file) + _refresher = refresher + log.info( + "W&B auth: GCP Workload Identity Federation " + "(identity token file %s, refresh every %ss)", + token_file, + interval, + ) + return True diff --git a/libs/hackbot-runtime/pyproject.toml b/libs/hackbot-runtime/pyproject.toml index 8186fc62e3..1b4c24324c 100644 --- a/libs/hackbot-runtime/pyproject.toml +++ b/libs/hackbot-runtime/pyproject.toml @@ -8,6 +8,7 @@ dependencies = [ "pydantic-settings>=2.1.0", "google-auth>=2.0.0", "agent-tools", + "weave>=0.53.1" ] [project.optional-dependencies] diff --git a/libs/hackbot-runtime/tests/test_tracing.py b/libs/hackbot-runtime/tests/test_tracing.py new file mode 100644 index 0000000000..4a1bee7135 --- /dev/null +++ b/libs/hackbot-runtime/tests/test_tracing.py @@ -0,0 +1,82 @@ +"""Tests for Weave tracing setup and agent-name derivation.""" + +import pytest +from hackbot_runtime import tracing, wandb_wif +from weave.conversation.agent_context import resolve_agent_name as weave_agent_name + + +def _entrypoint_at(path: str): + """A function whose source file is ``path`` (mimics an agent's main()).""" + ns: dict = {} + exec(compile("def main(): pass", path, "exec"), ns) + return ns["main"] + + +@pytest.fixture(autouse=True) +def _clear_env(monkeypatch): + for var in ("WANDB_API_KEY", wandb_wif.TOKEN_FILE_ENV, "WEAVE_PROJECT"): + monkeypatch.delenv(var, raising=False) + + +@pytest.mark.parametrize( + "path, expected", + [ + ("/app/hackbot_agents/build_repair/__main__.py", "build-repair"), + ("/x/hackbot_agents/test_plan_generator/__main__.py", "test-plan-generator"), + ("/venv/site-packages/hackbot_agents/bug_fix/__main__.py", "bug-fix"), + ], +) +def test_resolve_agent_name_from_source_path(path, expected): + assert tracing.resolve_agent_name(_entrypoint_at(path)) == expected + + +def test_init_weave_is_inert_without_credentials(monkeypatch): + import weave + + monkeypatch.setattr( + weave, "init", lambda *a, **k: pytest.fail("weave.init should not be called") + ) + assert tracing._init_weave() is False + + +def test_init_weave_enabled_by_api_key(monkeypatch): + import weave + + calls = {} + monkeypatch.setattr( + weave, "init", lambda project, **kw: calls.update(project=project) + ) + monkeypatch.setenv("WANDB_API_KEY", "dummy") + monkeypatch.setenv("WEAVE_PROJECT", "team/prod") + + assert tracing._init_weave() is True + assert calls == {"project": "team/prod"} + + +def test_init_weave_enabled_by_wif_token_file(monkeypatch): + """Federation leaves no API key -- the identity token file enables tracing.""" + import weave + + calls = {} + monkeypatch.setattr( + weave, "init", lambda project, **kw: calls.update(project=project) + ) + monkeypatch.setenv(wandb_wif.TOKEN_FILE_ENV, "/run/wandb-identity-token") + + assert tracing._init_weave() is True + assert calls == {"project": tracing.DEFAULT_WEAVE_PROJECT} + + +def test_trace_agent_is_noop_without_credentials(): + entry = _entrypoint_at("/app/hackbot_agents/build_repair/__main__.py") + with tracing.trace_agent(entry): + assert weave_agent_name("claude_agent_sdk") == "claude_agent_sdk" + + +def test_trace_agent_labels_spans_when_enabled(monkeypatch): + monkeypatch.setattr(tracing, "_init_weave", lambda: True) + entry = _entrypoint_at("/app/hackbot_agents/build_repair/__main__.py") + + with tracing.trace_agent(entry): + assert weave_agent_name("claude_agent_sdk") == "build-repair" + assert weave_agent_name("claude_agent_sdk") == "claude_agent_sdk" diff --git a/libs/hackbot-runtime/tests/test_wandb_wif.py b/libs/hackbot-runtime/tests/test_wandb_wif.py new file mode 100644 index 0000000000..9d9dab5b93 --- /dev/null +++ b/libs/hackbot-runtime/tests/test_wandb_wif.py @@ -0,0 +1,218 @@ +"""Tests for W&B GCP Workload Identity Federation auth setup.""" + +import logging +import os + +import pytest +from google.auth.exceptions import GoogleAuthError +from hackbot_runtime import wandb_wif +from hackbot_runtime.providers import ProviderError + +_AUDIENCE = "https://example.wandb.example/federation" + + +@pytest.fixture(autouse=True) +def _reset_module_state(): + yield + if wandb_wif._refresher is not None: + wandb_wif._refresher.stop() + wandb_wif._refresher = None + + +@pytest.fixture +def _no_federation_env(monkeypatch): + for var in ( + wandb_wif.AUDIENCE_ENV, + wandb_wif.TOKEN_FILE_ENV, + wandb_wif.REFRESH_INTERVAL_ENV, + ): + monkeypatch.delenv(var, raising=False) + + +def test_is_enabled_follows_audience(monkeypatch): + monkeypatch.delenv(wandb_wif.AUDIENCE_ENV, raising=False) + assert wandb_wif.is_enabled() is False + monkeypatch.setenv(wandb_wif.AUDIENCE_ENV, _AUDIENCE) + assert wandb_wif.is_enabled() is True + + +def test_configure_is_inert_without_federation(_no_federation_env, monkeypatch): + called = False + + def _should_not_fetch(*_a, **_k): + nonlocal called + called = True + return "tok" + + monkeypatch.setattr(wandb_wif, "fetch_gcp_identity_token", _should_not_fetch) + + assert wandb_wif.configure() is False + assert called is False + assert wandb_wif.TOKEN_FILE_ENV not in os.environ + + +def test_configure_writes_token_and_sets_env(_no_federation_env, monkeypatch, tmp_path): + token_file = tmp_path / "identity-token" + monkeypatch.setenv(wandb_wif.AUDIENCE_ENV, _AUDIENCE) + monkeypatch.setenv(wandb_wif.TOKEN_FILE_ENV, str(token_file)) + monkeypatch.delenv("WANDB_API_KEY", raising=False) + monkeypatch.setattr( + wandb_wif, "fetch_gcp_identity_token", lambda *a, **k: "google.jwt.token" + ) + + assert wandb_wif.configure() is True + + assert token_file.read_text() == "google.jwt.token" + assert os.environ[wandb_wif.TOKEN_FILE_ENV] == str(token_file) + + +def test_configure_refuses_when_api_key_set( + _no_federation_env, monkeypatch, tmp_path, caplog +): + monkeypatch.setenv(wandb_wif.AUDIENCE_ENV, _AUDIENCE) + monkeypatch.setenv(wandb_wif.TOKEN_FILE_ENV, str(tmp_path / "identity-token")) + monkeypatch.setenv("WANDB_API_KEY", "leftover") + monkeypatch.setattr(wandb_wif, "fetch_gcp_identity_token", lambda *a, **k: "tok") + + with caplog.at_level(logging.ERROR, logger=wandb_wif.log.name): + assert wandb_wif.configure() is False + + assert any( + rec.levelno == logging.ERROR and "WANDB_API_KEY" in rec.message + for rec in caplog.records + ) + assert os.environ["WANDB_API_KEY"] == "leftover" + assert wandb_wif._refresher is None + + +def test_configure_is_idempotent(_no_federation_env, monkeypatch, tmp_path): + monkeypatch.setenv(wandb_wif.AUDIENCE_ENV, _AUDIENCE) + monkeypatch.setenv(wandb_wif.TOKEN_FILE_ENV, str(tmp_path / "tok")) + monkeypatch.delenv("WANDB_API_KEY", raising=False) + + calls = 0 + + def _fetch(*_a, **_k): + nonlocal calls + calls += 1 + return "tok" + + monkeypatch.setattr(wandb_wif, "fetch_gcp_identity_token", _fetch) + + assert wandb_wif.configure() is True + assert wandb_wif.configure() is True + assert calls == 1 + + +def test_configure_defaults_token_path_when_unset(_no_federation_env, monkeypatch): + monkeypatch.setenv(wandb_wif.AUDIENCE_ENV, _AUDIENCE) + monkeypatch.delenv("WANDB_API_KEY", raising=False) + monkeypatch.setattr(wandb_wif, "fetch_gcp_identity_token", lambda *a, **k: "tok") + + assert wandb_wif.configure() is True + + path = os.environ[wandb_wif.TOKEN_FILE_ENV] + assert "wandb-wif-" in path + assert os.path.exists(path) + + +def test_configure_passes_audience_to_fetch(_no_federation_env, monkeypatch, tmp_path): + monkeypatch.setenv(wandb_wif.AUDIENCE_ENV, _AUDIENCE) + monkeypatch.setenv(wandb_wif.TOKEN_FILE_ENV, str(tmp_path / "tok")) + monkeypatch.delenv("WANDB_API_KEY", raising=False) + + seen = {} + + def _fetch(audience): + seen["audience"] = audience + return "tok" + + monkeypatch.setattr(wandb_wif, "fetch_gcp_identity_token", _fetch) + + assert wandb_wif.configure() is True + assert seen["audience"] == _AUDIENCE + + +def test_fetch_uses_google_auth_with_audience(monkeypatch): + captured = {} + + def _fake_fetch(request, audience): + captured["request"] = request + captured["audience"] = audience + return "signed.jwt.value" + + monkeypatch.setattr(wandb_wif.google.oauth2.id_token, "fetch_id_token", _fake_fetch) + + token = wandb_wif.fetch_gcp_identity_token(_AUDIENCE) + + assert token == "signed.jwt.value" + assert captured["audience"] == _AUDIENCE + assert captured["request"] is not None + + +def test_fetch_wraps_google_auth_error(monkeypatch): + def _raise(*_a, **_k): + raise GoogleAuthError("metadata server unreachable") + + monkeypatch.setattr(wandb_wif.google.oauth2.id_token, "fetch_id_token", _raise) + with pytest.raises(ProviderError, match="Failed to fetch a Google identity token"): + wandb_wif.fetch_gcp_identity_token(_AUDIENCE) + + +def test_fetch_rejects_empty(monkeypatch): + monkeypatch.setattr( + wandb_wif.google.oauth2.id_token, "fetch_id_token", lambda *a, **k: "" + ) + with pytest.raises(ProviderError, match="empty identity token"): + wandb_wif.fetch_gcp_identity_token(_AUDIENCE) + + +@pytest.mark.parametrize( + "raw, expected", + [ + (None, wandb_wif.DEFAULT_REFRESH_INTERVAL), + ("900", 900), + ("not-a-number", wandb_wif.DEFAULT_REFRESH_INTERVAL), + ("0", wandb_wif.DEFAULT_REFRESH_INTERVAL), + ("-5", wandb_wif.DEFAULT_REFRESH_INTERVAL), + ], +) +def test_refresh_interval_parsing(monkeypatch, raw, expected): + if raw is None: + monkeypatch.delenv(wandb_wif.REFRESH_INTERVAL_ENV, raising=False) + else: + monkeypatch.setenv(wandb_wif.REFRESH_INTERVAL_ENV, raw) + assert wandb_wif._refresh_interval() == expected + + +def test_token_file_write_is_atomic_and_clean(monkeypatch, tmp_path): + token_file = tmp_path / "identity-token" + monkeypatch.setattr( + wandb_wif, "fetch_gcp_identity_token", lambda *a, **k: "the.jwt" + ) + refresher = wandb_wif._TokenFileRefresher( + token_file, interval=1800, audience=_AUDIENCE + ) + + refresher._write_token() + + assert token_file.read_text() == "the.jwt" + assert list(tmp_path.iterdir()) == [token_file] + + +def test_refresh_loop_survives_provider_error(monkeypatch, tmp_path): + refresher = wandb_wif._TokenFileRefresher( + tmp_path / "tok", interval=0, audience=_AUDIENCE + ) + calls = [] + + def _boom(*_a, **_k): + calls.append(1) + refresher.stop() + raise ProviderError("transient") + + monkeypatch.setattr(wandb_wif, "fetch_gcp_identity_token", _boom) + + refresher._loop() + + assert calls == [1] diff --git a/uv.lock b/uv.lock index 0de875449c..45ed2b712c 100644 --- a/uv.lock +++ b/uv.lock @@ -2593,6 +2593,7 @@ dependencies = [ { name = "google-auth" }, { name = "pydantic-settings" }, { name = "requests" }, + { name = "weave" }, ] [package.optional-dependencies] @@ -2609,6 +2610,7 @@ requires-dist = [ { name = "google-auth", specifier = ">=2.0.0" }, { name = "pydantic-settings", specifier = ">=2.1.0" }, { name = "requests", specifier = ">=2.32.0" }, + { name = "weave", specifier = ">=0.53.1" }, ] provides-extras = ["claude-sdk"] @@ -7424,7 +7426,7 @@ wheels = [ [[package]] name = "weave" -version = "0.52.43" +version = "0.53.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cachetools" }, @@ -7435,7 +7437,6 @@ dependencies = [ { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp-proto-http" }, { name = "opentelemetry-sdk" }, - { name = "opentelemetry-semantic-conventions" }, { name = "packaging" }, { name = "polyfile-weave" }, { name = "pydantic" }, @@ -7443,9 +7444,9 @@ dependencies = [ { name = "tenacity" }, { name = "tzdata", marker = "sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9d/df/dcb32bf366eca5e11d345a8e15238256396a9ff7346a29b3290a358b9cd8/weave-0.52.43.tar.gz", hash = "sha256:3c4abcb4da33e0db8a03bd52739503d99c90bdb1922b704ca2b2df8a0f6d561e", size = 977397, upload-time = "2026-06-15T21:29:23.062Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a3/34/c9ddb250e27708b581f52fd75f94ca04edbe7ed287560fe13995e042521d/weave-0.53.1.tar.gz", hash = "sha256:623c1918394e8493a3344fb89ffc71b3e1e96c3bd247a77309aa2b6e78475cb9", size = 1113042, upload-time = "2026-07-02T21:17:45.087Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c1/a0/e823f714e337c6342a831307d21d2a12c87bc6a6015fc0ac62be9dd8e01f/weave-0.52.43-py3-none-any.whl", hash = "sha256:d7fb10e893c9da57c7663c8c51a1a67eb59216f9da0e4056a244a41de63faf19", size = 1199217, upload-time = "2026-06-15T21:29:21.267Z" }, + { url = "https://files.pythonhosted.org/packages/bf/87/99d77ddb9f1008a92e36e6d72b0ff63e17fbe22abe6b383bbfb80001d70b/weave-0.53.1-py3-none-any.whl", hash = "sha256:95c45e22fe38cf241f4ff26d120c86bcff3652e4c813e0bcf5950ec26e7f5db8", size = 1349841, upload-time = "2026-07-02T21:17:42.969Z" }, ] [[package]]