Agent tracing in Weave

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/<run_id>, bind-mounted to the host's ~/hackbot/artifacts.
       - ARTIFACTS_DIR=/artifacts

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).

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

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/<agent>/__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

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

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]