diff --git a/.github/workflows/pr-approval-agent.yml b/.github/workflows/pr-approval-agent.yml index 9ca86fd1c1..a12473fc90 100644 --- a/.github/workflows/pr-approval-agent.yml +++ b/.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 }} \ diff --git a/tools/pr-approval-agent/gateway.py b/tools/pr-approval-agent/gateway.py new file mode 100644 index 0000000000..1522118dde --- /dev/null +++ b/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:///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:///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}), + } diff --git a/tools/pr-approval-agent/reviewer.py b/tools/pr-approval-agent/reviewer.py index 140a64821e..ea935011a6 100644 --- a/tools/pr-approval-agent/reviewer.py +++ b/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): diff --git a/tools/pr-approval-agent/test_gateway.py b/tools/pr-approval-agent/test_gateway.py new file mode 100644 index 0000000000..dff67f119c --- /dev/null +++ b/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"] diff --git a/tools/pr-approval-agent/test_reviewer_gateway.py b/tools/pr-approval-agent/test_reviewer_gateway.py new file mode 100644 index 0000000000..9a4a476ddb --- /dev/null +++ b/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