feat(stamphog): route PR review through the ai-gateway

.github/workflows/pr-approval-agent.yml

@@ -67,6 +67,10 @@ jobs:
                   ANTHROPIC_API_KEY: ${{ secrets.STAMPHOG_ANTHROPIC_API_KEY }}
                   POSTHOG_API_KEY: ${{ secrets.POSTHOG_API_TOKEN }}
                   GH_TOKEN: ${{ steps.app-token.outputs.token }}
+                  # Set both to route through the ai-gateway (gateway.py overrides
+                  # ANTHROPIC_* at runtime). Unset means direct Anthropic.
+                  AI_GATEWAY_URL: ${{ secrets.STAMPHOG_AI_GATEWAY_URL }}
+                  AI_GATEWAY_API_KEY: ${{ secrets.STAMPHOG_AI_GATEWAY_API_KEY }}
               run: |
                   uv run tools/pr-approval-agent/review_pr.py \
                     ${{ github.event.pull_request.number }} \

tools/pr-approval-agent/gateway.py

@@ -0,0 +1,76 @@
+# ruff: noqa: T201
+"""Route the Claude Agent SDK through PostHog's internal Go ai-gateway.
+
+Gated on AI_GATEWAY_URL + AI_GATEWAY_API_KEY; a bad/half-set config falls back to
+direct Anthropic instead of failing the review. The gateway is slugless, so the
+product rides on a header, not the path.
+"""
+
+import os
+import json
+from urllib.parse import urlparse
+
+# aio_ matches the other cutovers; no $ai_ prefix (gateway strips those).
+AI_PRODUCT = "aio_stamphog"
+
+
+def _misconfig(url: str, api_key: str) -> str | None:
+    if not (url and api_key):
+        return "AI_GATEWAY_URL and AI_GATEWAY_API_KEY must be set together"
+    parsed = urlparse(url)
+    # Require an absolute URL: a schemeless value parses all into .path and would
+    # pass the /v1 check, arming a broken base instead of falling back.
+    if parsed.scheme not in ("http", "https") or not parsed.netloc:
+        return "AI_GATEWAY_URL must be an absolute http(s) URL, e.g. https://<host>/v1"
+    # A query/fragment survives /v1 stripping and would corrupt the base URL.
+    if parsed.query or parsed.fragment:
+        return "AI_GATEWAY_URL must not contain a query string or fragment"
+    if not parsed.path.rstrip("/").endswith("/v1"):
+        return "AI_GATEWAY_URL must include the OpenAI base path, e.g. https://<host>/v1"
+    return None
+
+
+def resolve_gateway_config() -> tuple[str, str] | None:
+    """Validated (anthropic_base_url, phs_api_key), or None to use direct Anthropic.
+
+    Trailing /v1 is stripped; the Agent SDK re-appends /v1/messages.
+    """
+    url = os.environ.get("AI_GATEWAY_URL", "").strip()
+    api_key = os.environ.get("AI_GATEWAY_API_KEY", "").strip()
+    if not (url or api_key):
+        return None
+    reason = _misconfig(url, api_key)
+    if reason:
+        print(f"⚠️  ai-gateway misconfigured, falling back to direct Anthropic: {reason}")
+        return None
+    # Rebuild from parsed components (not raw-string slicing) so nothing past the
+    # path can leak into the base; _misconfig already rejected query/fragment.
+    parsed = urlparse(url)
+    path = parsed.path.rstrip("/")
+    if path.endswith("/v1"):
+        path = path[: -len("/v1")]
+    return f"{parsed.scheme}://{parsed.netloc}{path}", api_key
+
+
+def _properties_header(properties: dict[str, object]) -> str:
+    # Single X-PostHog-Properties JSON blob: the slugless Go gateway merges it onto
+    # $ai_generation and ignores per-property x-posthog-property-* headers. None
+    # dropped; newlines collapsed so a value can't break the header block.
+    clean: dict[str, object] = {}
+    for key, value in properties.items():
+        if value is None:
+            continue
+        clean[key] = value.replace("\r", " ").replace("\n", " ") if isinstance(value, str) else value
+    if not clean:
+        return ""
+    return f"X-PostHog-Properties: {json.dumps(clean, separators=(',', ':'))}"
+
+
+def gateway_env(base_url: str, api_key: str, properties: dict[str, object]) -> dict[str, str]:
+    # phs_ secret on both auth vars (SDK sends it as Bearer or x-api-key).
+    return {
+        "ANTHROPIC_BASE_URL": base_url,
+        "ANTHROPIC_AUTH_TOKEN": api_key,
+        "ANTHROPIC_API_KEY": api_key,
+        "ANTHROPIC_CUSTOM_HEADERS": _properties_header({"ai_product": AI_PRODUCT, **properties}),
+    }

tools/pr-approval-agent/reviewer.py

@@ -5,6 +5,7 @@
 and reach a verdict on whether a PR is safe to auto-approve.
 """
 
+import os
 import re
 import json
 import asyncio
@@ -14,18 +15,20 @@
 
 from claude_agent_sdk import ClaudeAgentOptions, ResultMessage, query
 from claude_agent_sdk.types import AssistantMessage, ToolUseBlock
+from gateway import gateway_env, resolve_gateway_config
 from github import PRData
 
+# Traced wrapper, bound only with a PostHog key. Gateway mode uses the plain
+# `query` so the gateway's own $ai_generation isn't double-counted.
+_traced_query = None
 try:
-    import os
-
     import posthoganalytics
 
     posthoganalytics.api_key = os.environ.get("POSTHOG_API_KEY", "")  # ty: ignore[invalid-assignment]
     posthoganalytics.host = os.environ.get("POSTHOG_HOST", "https://us.i.posthog.com")  # ty: ignore[invalid-assignment]
 
     if posthoganalytics.api_key:
-        from posthoganalytics.ai.claude_agent_sdk import query  # type: ignore[no-redef]  # noqa: F811
+        from posthoganalytics.ai.claude_agent_sdk import query as _traced_query  # type: ignore[no-redef]
 
         _POSTHOG_AI_AVAILABLE = True
     else:
@@ -36,6 +39,15 @@
 MODEL = "claude-sonnet-4-6"
 
 
+def _apply_gateway_route(gateway: tuple[str, str] | None, attribution: dict[str, object]):
+    """Apply the gateway env and return the plain SDK ``query`` when configured, else None."""
+    if gateway is None:
+        return None
+    base_url, api_key = gateway
+    os.environ.update(gateway_env(base_url, api_key, attribution))
+    return query
+
+
 _CONTROL_CHARS_RE = re.compile(r"[^\x20-\x7E\n\t]")
 
 
@@ -233,11 +245,24 @@ async def _review(self, pr: PRData, classification: dict, gate_context: dict) ->
             extra_args={"no-session-persistence": None},
         )
 
+        attribution = {
+            "stamphog_pr_number": pr.number,
+            "stamphog_repo": pr.repo,
+            "stamphog_author": pr.author,
+            "stamphog_tier": classification.get("tier", ""),
+            "stamphog_t1_subclass": classification.get("t1_subclass", ""),
+            "stamphog_breadth": classification.get("breadth", ""),
+            "stamphog_commit_type": classification.get("commit_type") or "",
+            "stamphog_gate_verdict": gate_context.get("gate_verdict", ""),
+            "stamphog_files_changed": len(pr.files),
+            "stamphog_lines_total": pr.lines_total,
+        }
+
+        active_query = _apply_gateway_route(resolve_gateway_config(), attribution)
         posthog_kwargs: dict = {}
-        if _POSTHOG_AI_AVAILABLE:
-            # Unique reviewer usernames, sanitized — labels and title are
-            # author-controlled so we sanitize them too (cheap insurance
-            # against weird unicode landing in analytics).
+        props: dict = {}
+        if active_query is None and _POSTHOG_AI_AVAILABLE:
+            active_query = _traced_query
             reviewers = sorted({_sanitize_untrusted(r["user"], max_len=50) for r in pr.reviews if r.get("user")})
             safe_labels = [_sanitize_untrusted(label, max_len=100) for label in pr.labels]
             trace_name = f"stamphog PR #{pr.number}: {_sanitize_untrusted(pr.title, max_len=100)}"
@@ -273,14 +298,15 @@ async def _review(self, pr: PRData, classification: dict, gate_context: dict) ->
                     "stamphog_llm_verdict": "",
                 },
             }
+            # Live ref: the verdict update below reaches the trace ($ai_trace is
+            # sent after the generator ends).
+            props = posthog_kwargs["posthog_properties"]
 
-        # Keep a reference so we can mutate it when the verdict arrives —
-        # the SDK sends the $ai_trace event after the generator completes,
-        # so updates here propagate to the trace.
-        props = posthog_kwargs.get("posthog_properties", {})
+        if active_query is None:
+            active_query = query
 
         structured_output = None
-        async for message in query(prompt=prompt, options=options, **posthog_kwargs):
+        async for message in active_query(prompt=prompt, options=options, **posthog_kwargs):
             if self.verbose:
                 print(f"\033[2m    [{type(message).__name__}]\033[0m", flush=True)
             if isinstance(message, ResultMessage):

tools/pr-approval-agent/test_gateway.py

@@ -0,0 +1,105 @@
+"""Tests for ai-gateway routing in gateway.py."""
+
+import pytest
+
+from gateway import AI_PRODUCT, gateway_env, resolve_gateway_config
+
+
+@pytest.fixture(autouse=True)
+def _clear_gateway_env(monkeypatch):
+    monkeypatch.delenv("AI_GATEWAY_URL", raising=False)
+    monkeypatch.delenv("AI_GATEWAY_API_KEY", raising=False)
+
+
+def test_unset_returns_none_direct_path(monkeypatch):
+    assert resolve_gateway_config() is None
+
+
+@pytest.mark.parametrize(
+    "url,api_key",
+    [
+        pytest.param("https://gateway.us.posthog.com/v1", "", id="url-only"),
+        pytest.param("", "phs_secret", id="key-only"),
+    ],
+)
+def test_half_applied_config_falls_back(monkeypatch, url, api_key):
+    monkeypatch.setenv("AI_GATEWAY_URL", url)
+    monkeypatch.setenv("AI_GATEWAY_API_KEY", api_key)
+    assert resolve_gateway_config() is None
+
+
+def test_url_without_v1_falls_back(monkeypatch):
+    monkeypatch.setenv("AI_GATEWAY_URL", "https://gateway.us.posthog.com")
+    monkeypatch.setenv("AI_GATEWAY_API_KEY", "phs_secret")
+    assert resolve_gateway_config() is None
+
+
+def test_schemeless_url_falls_back(monkeypatch):
+    # Schemeless parses all into .path, so the /v1 check alone would pass.
+    monkeypatch.setenv("AI_GATEWAY_URL", "gateway.us.posthog.com/v1")
+    monkeypatch.setenv("AI_GATEWAY_API_KEY", "phs_secret")
+    assert resolve_gateway_config() is None
+
+
+@pytest.mark.parametrize(
+    "bad_url",
+    [
+        pytest.param("https://gateway.us.posthog.com/v1?x=y", id="query"),
+        pytest.param("https://gateway.us.posthog.com/v1#frag", id="fragment"),
+    ],
+)
+def test_url_with_query_or_fragment_falls_back(monkeypatch, bad_url):
+    # A query/fragment survives /v1 stripping and would corrupt the base URL.
+    monkeypatch.setenv("AI_GATEWAY_URL", bad_url)
+    monkeypatch.setenv("AI_GATEWAY_API_KEY", "phs_secret")
+    assert resolve_gateway_config() is None
+
+
+@pytest.mark.parametrize(
+    "configured_url,expected_base",
+    [
+        pytest.param("https://gateway.us.posthog.com/v1", "https://gateway.us.posthog.com", id="strips-v1"),
+        pytest.param(
+            "https://gateway.us.posthog.com/v1/", "https://gateway.us.posthog.com", id="strips-v1-trailing-slash"
+        ),
+    ],
+)
+def test_resolves_and_strips_v1_for_anthropic_base(monkeypatch, configured_url, expected_base):
+    # SDK re-appends /v1/messages, so resolve() returns the bare host.
+    monkeypatch.setenv("AI_GATEWAY_URL", configured_url)
+    monkeypatch.setenv("AI_GATEWAY_API_KEY", "phs_secret")
+    assert resolve_gateway_config() == (expected_base, "phs_secret")
+
+
+def test_gateway_env_points_sdk_at_gateway():
+    env = gateway_env("https://gateway.us.posthog.com", "phs_secret", {"stamphog_pr_number": 123})
+    assert env["ANTHROPIC_BASE_URL"] == "https://gateway.us.posthog.com"
+    assert env["ANTHROPIC_AUTH_TOKEN"] == "phs_secret"
+    assert env["ANTHROPIC_API_KEY"] == "phs_secret"
+
+
+def test_gateway_env_tags_ai_product_and_attribution():
+    # Go gateway reads one X-PostHog-Properties JSON blob, not x-posthog-property-*.
+    env = gateway_env("https://host", "phs_secret", {"stamphog_pr_number": 123, "stamphog_repo": "PostHog/posthog"})
+    headers = env["ANTHROPIC_CUSTOM_HEADERS"]
+    assert headers == (
+        'X-PostHog-Properties: {"ai_product":"aio_stamphog","stamphog_pr_number":123,"stamphog_repo":"PostHog/posthog"}'
+    )
+    assert "x-posthog-property-" not in headers
+
+
+def test_ai_product_uses_aio_prefix_no_reserved_prefix():
+    assert AI_PRODUCT == "aio_stamphog"
+    assert not AI_PRODUCT.startswith("$")
+
+
+def test_header_values_are_single_line():
+    env = gateway_env("https://host", "phs_secret", {"stamphog_pr_title": "line one\nline two"})
+    header = env["ANTHROPIC_CUSTOM_HEADERS"]
+    assert "\n" not in header
+    assert '"stamphog_pr_title":"line one line two"' in header
+
+
+def test_none_attribution_values_dropped():
+    env = gateway_env("https://host", "phs_secret", {"stamphog_commit_type": None})
+    assert "stamphog_commit_type" not in env["ANTHROPIC_CUSTOM_HEADERS"]

tools/pr-approval-agent/test_reviewer_gateway.py

@@ -0,0 +1,47 @@
+"""Routing-pin tests for the ai-gateway switch in reviewer.py."""
+
+import os
+import sys
+
+import pytest
+from unittest.mock import MagicMock
+
+# reviewer.py's claude_agent_sdk dep is installed by `uv run`, not the test venv.
+sys.modules.setdefault("claude_agent_sdk", MagicMock())
+sys.modules.setdefault("claude_agent_sdk.types", MagicMock())
+
+import reviewer  # noqa: E402
+from reviewer import _apply_gateway_route  # noqa: E402
+
+_ANTHROPIC_ENV_KEYS = ("ANTHROPIC_BASE_URL", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY", "ANTHROPIC_CUSTOM_HEADERS")
+
+
+@pytest.fixture(autouse=True)
+def _restore_anthropic_env():
+    saved = {k: os.environ.get(k) for k in _ANTHROPIC_ENV_KEYS}
+    yield
+    for k, v in saved.items():
+        if v is None:
+            os.environ.pop(k, None)
+        else:
+            os.environ[k] = v
+
+
+def test_no_gateway_returns_none_and_leaves_env_untouched():
+    before = {k: os.environ.get(k) for k in _ANTHROPIC_ENV_KEYS}
+    assert _apply_gateway_route(None, {"stamphog_pr_number": 1}) is None
+    assert {k: os.environ.get(k) for k in _ANTHROPIC_ENV_KEYS} == before
+
+
+def test_gateway_mode_uses_plain_query_and_applies_env():
+    active = _apply_gateway_route(("https://gateway.us.posthog.com", "phs_secret"), {"stamphog_pr_number": 7})
+    # Plain query, not the traced wrapper (else the gateway's $ai_generation double-counts).
+    assert active is reviewer.query
+    assert os.environ["ANTHROPIC_BASE_URL"] == "https://gateway.us.posthog.com"
+    assert os.environ["ANTHROPIC_AUTH_TOKEN"] == "phs_secret"
+    assert os.environ["ANTHROPIC_API_KEY"] == "phs_secret"
+    headers = os.environ["ANTHROPIC_CUSTOM_HEADERS"]
+    # Single X-PostHog-Properties JSON blob; the slugless gateway ignores per-property headers.
+    assert '"ai_product":"aio_stamphog"' in headers
+    assert '"stamphog_pr_number":7' in headers
+    assert "x-posthog-property-" not in headers