From 4c7d45273f6565ecd893faa84f086f4f50f76426 Mon Sep 17 00:00:00 2001 From: Johannes du Plessis Date: Tue, 23 Jun 2026 12:32:10 -0700 Subject: [PATCH] perf: cut review-chat time-to-first-token (#1598) * perf: cut review-chat time-to-first-token The sandbox-less PR review chat paid several blocking network round-trips before the first token on every message. Cache GitHub App installation tokens in-process (per scope, until ~10m before expiry, above the proxy's 5m refresh window) so the chat graph factory and proxy stop re-minting one each turn. Also drop the duplicate thread-metadata read in the commands proxy and replace the heavy per-message get_review staleness check with a single lightweight PR head-SHA lookup. Co-authored-by: open-swe[bot] * fix: keep review chat alive when reseed fails Address review: a moved PR head now triggers _build_pr_context (and thus get_review). For an existing chat, fall back to the last seeded context on HTTPException instead of failing the command; fresh chats still surface the error since they have no prior context. Co-authored-by: open-swe[bot] --------- Co-authored-by: open-swe[bot] --- agent/dashboard/review_api.py | 16 +++++ agent/dashboard/review_chat_api.py | 93 +++++++++++++++++++----------- agent/utils/github_app.py | 65 ++++++++++++++++++++- tests/test_github_app.py | 92 +++++++++++++++++++++++++++++ tests/test_review_api.py | 29 ++++++++++ tests/test_review_chat.py | 90 ++++++++++++++++++++++++++++- 6 files changed, 350 insertions(+), 35 deletions(-) diff --git a/agent/dashboard/review_api.py b/agent/dashboard/review_api.py index 8aa8a0bf..3351c884 100644 --- a/agent/dashboard/review_api.py +++ b/agent/dashboard/review_api.py @@ -348,6 +348,22 @@ async def _fetch_check_runs(owner: str, repo: str, sha: str, token: str) -> list return out +async def get_pr_head_sha(owner: str, repo: str, pr_number: int) -> str: + """Return the PR's current head SHA from GitHub, or "" if unavailable. + + A lightweight alternative to :func:`get_review` for callers that only need to + detect whether the PR head has moved (e.g. the chat staleness check). + """ + try: + token = await _require_app_token() + payload = await _github_get(f"/repos/{owner}/{repo}/pulls/{pr_number}", token) + except HTTPException: + return "" + head = payload.get("head") if isinstance(payload, dict) else None + sha = head.get("sha") if isinstance(head, dict) else None + return sha if isinstance(sha, str) else "" + + async def get_review(owner: str, repo: str, pr_number: int) -> dict[str, Any]: thread_id = reviewer_thread_id(owner, repo, pr_number) client = langgraph_client() diff --git a/agent/dashboard/review_chat_api.py b/agent/dashboard/review_chat_api.py index 3c74175e..a918870a 100644 --- a/agent/dashboard/review_chat_api.py +++ b/agent/dashboard/review_chat_api.py @@ -25,7 +25,7 @@ from ..reviewer_findings import REVIEWER_THREAD_KIND from ..utils.github_app import get_github_app_installation_token from ..utils.thread_ops import langgraph_client, langgraph_url from .options import SUPPORTED_MODEL_IDS, model_supports_effort -from .review_api import classify_finding, get_review, reviewer_thread_id +from .review_api import classify_finding, get_pr_head_sha, get_review, reviewer_thread_id from .thread_api import ( _DASHBOARD_STREAM_MODES, _langgraph_proxy_headers, @@ -37,6 +37,8 @@ logger = logging.getLogger(__name__) _CHAT_ASSISTANT_ID = "chat" _CHAT_SOURCE = "review_chat" +# Sentinel: caller did not pre-fetch the thread metadata, so fetch it here. +_UNFETCHED = object() _PROXY_REQUEST_TIMEOUT = httpx.Timeout(30.0, connect=5.0) _MAX_DIFF_CHARS = 400_000 @@ -314,6 +316,7 @@ async def _enrich_chat_command( pr_number: int, login: str, thread_id: str, + thread_metadata: Any = _UNFETCHED, ) -> dict[str, Any]: if command.get("method") != "run.start": return command @@ -323,7 +326,11 @@ async def _enrich_chat_command( params = {} command["params"] = params - metadata = await _get_chat_thread_metadata(thread_id) + # Reuse the metadata the access check already fetched to avoid a duplicate + # thread read on the hot path; fall back to fetching it when not supplied. + if thread_metadata is _UNFETCHED: + thread_metadata = await _get_chat_thread_metadata(thread_id) + metadata = thread_metadata created = metadata is None if created: await _create_chat_thread( @@ -358,46 +365,59 @@ async def _enrich_chat_command( stored_head = metadata.get("chat_head_sha") if isinstance(metadata, dict) else None stored_head = stored_head if isinstance(stored_head, str) else "" + # Detect head movement with a single lightweight PR lookup instead of a full + # ``get_review`` (which also fetches check runs + the reviewer thread). The + # full review is only fetched below, by ``_build_pr_context``, when we + # actually need to reseed. On lookup failure, keep the existing context. review: dict[str, Any] | None = None needs_seed = created if not created: - try: - review = await get_review(owner, repo, pr_number) - except HTTPException: - review = None # transient/missing review: keep the existing context - if review is not None: - current_head = _review_head_sha(review) - needs_seed = bool(current_head) and current_head != stored_head + current_head = await get_pr_head_sha(owner, repo, pr_number) + needs_seed = bool(current_head) and current_head != stored_head if needs_seed: - token = await get_github_app_installation_token(repositories=[repo]) - if not token: - raise HTTPException(503, "GitHub App token unavailable") try: + token = await get_github_app_installation_token(repositories=[repo]) + if not token: + raise HTTPException(503, "GitHub App token unavailable") pr_files, head_sha = await _build_pr_context( owner, repo, pr_number, token, review=review ) - except HTTPException: - raise except Exception as exc: # noqa: BLE001 - logger.warning("Failed to seed PR chat context for %s/%s#%s", owner, repo, pr_number) - raise HTTPException(502, "could not load PR context") from exc - if head_sha: - configurable["chat_head_sha"] = head_sha - await langgraph_client().threads.update( - thread_id=thread_id, metadata={"chat_head_sha": head_sha} + # An existing chat can keep answering from its last seeded context, so + # a transient reseed failure shouldn't break the conversation. A fresh + # chat (or one never seeded) has nothing to fall back to, so surface it. + if created or not stored_head: + if isinstance(exc, HTTPException): + raise + logger.warning( + "Failed to seed PR chat context for %s/%s#%s", owner, repo, pr_number + ) + raise HTTPException(502, "could not load PR context") from exc + logger.warning( + "Failed to reseed PR chat context for %s/%s#%s; keeping last seeded context", + owner, + repo, + pr_number, ) - elif stored_head: configurable["chat_head_sha"] = stored_head - run_input = params.get("input") - if not isinstance(run_input, dict): - run_input = {} - existing_files = run_input.get("files") - run_input["files"] = { - **(existing_files if isinstance(existing_files, dict) else {}), - **pr_files, - } - params["input"] = run_input + else: + if head_sha: + configurable["chat_head_sha"] = head_sha + await langgraph_client().threads.update( + thread_id=thread_id, metadata={"chat_head_sha": head_sha} + ) + elif stored_head: + configurable["chat_head_sha"] = stored_head + run_input = params.get("input") + if not isinstance(run_input, dict): + run_input = {} + existing_files = run_input.get("files") + run_input["files"] = { + **(existing_files if isinstance(existing_files, dict) else {}), + **pr_files, + } + params["input"] = run_input elif stored_head: configurable["chat_head_sha"] = stored_head @@ -420,8 +440,9 @@ async def proxy_review_chat_commands( content_type: str = "application/json", ) -> tuple[int, bytes, str | None]: # Reject threads the caller doesn't own; a missing thread is created lazily - # below on the first `run.start` (with the caller as owner). - await assert_chat_thread_access(thread_id, owner, repo, pr_number, login) + # below on the first `run.start` (with the caller as owner). Reuse the + # fetched metadata to avoid a second thread read during enrichment. + metadata = await assert_chat_thread_access(thread_id, owner, repo, pr_number, login) _require_json_content_type(content_type) try: parsed = json.loads(body) @@ -431,7 +452,13 @@ async def proxy_review_chat_commands( raise HTTPException(400, "command body must be a JSON object") enriched = await _enrich_chat_command( - parsed, owner=owner, repo=repo, pr_number=pr_number, login=login, thread_id=thread_id + parsed, + owner=owner, + repo=repo, + pr_number=pr_number, + login=login, + thread_id=thread_id, + thread_metadata=metadata, ) url = f"{langgraph_url().rstrip('/')}/threads/{thread_id}/commands" headers = _langgraph_proxy_headers(content_type=content_type) diff --git a/agent/utils/github_app.py b/agent/utils/github_app.py index 4d45a360..f69041bd 100644 --- a/agent/utils/github_app.py +++ b/agent/utils/github_app.py @@ -6,6 +6,7 @@ import logging import os import time from collections.abc import Sequence +from datetime import UTC, datetime, timedelta from typing import Any import httpx @@ -17,6 +18,58 @@ GITHUB_APP_ID = os.environ.get("GITHUB_APP_ID", "") GITHUB_APP_PRIVATE_KEY = os.environ.get("GITHUB_APP_PRIVATE_KEY", "") GITHUB_APP_INSTALLATION_ID = os.environ.get("GITHUB_APP_INSTALLATION_ID", "") +# Installation tokens are valid for 1 hour. Reuse a minted token until it is +# within this window of expiring so chat/review requests don't pay a fresh +# JWT-sign + GitHub round-trip every message. The margin stays above the proxy's +# 5-minute refresh window (``github_proxy.PROXY_TOKEN_REFRESH_WINDOW``) so a +# near-expiry proxy refresh still mints a genuinely fresh token. +_TOKEN_CACHE_MARGIN = timedelta(minutes=10) +# scope key -> (token, expires_at, good_until). In-process only; never persisted. +_TOKEN_CACHE: dict[tuple[tuple[int, ...], tuple[str, ...]], tuple[str, str | None, datetime]] = {} + + +def _scope_key( + repository_ids: Sequence[int] | None, repositories: Sequence[str] | None +) -> tuple[tuple[int, ...], tuple[str, ...]]: + """Cache key segregating repo-scoped tokens from installation-wide ones.""" + ids = tuple(sorted(int(i) for i in repository_ids)) if repository_ids else () + names = tuple(sorted(str(r) for r in repositories)) if repositories else () + return ids, names + + +def _parse_expiry(expires_at: Any) -> datetime | None: + """Best-effort parse of a GitHub ``expires_at`` ISO timestamp to a UTC datetime.""" + if not isinstance(expires_at, str): + return None + raw = expires_at.strip() + if not raw: + return None + if raw.endswith("Z"): + raw = raw[:-1] + "+00:00" + try: + parsed = datetime.fromisoformat(raw) + except ValueError: + return None + return parsed if parsed.tzinfo else parsed.replace(tzinfo=UTC) + + +def _cached_token( + key: tuple[tuple[int, ...], tuple[str, ...]], *, now: datetime +) -> tuple[str, str | None] | None: + cached = _TOKEN_CACHE.get(key) + if cached is None: + return None + token, expires_at, good_until = cached + if now < good_until: + return token, expires_at + _TOKEN_CACHE.pop(key, None) + return None + + +def clear_app_token_cache() -> None: + """Drop all cached installation tokens (test/maintenance hook).""" + _TOKEN_CACHE.clear() + def _generate_app_jwt() -> str: """Generate a short-lived JWT signed with the GitHub App private key.""" @@ -53,6 +106,12 @@ async def get_github_app_installation_token_with_expiry( logger.debug("GitHub App env vars not fully configured, skipping app token") return None, None + key = _scope_key(repository_ids, repositories) + now = datetime.now(UTC) + cached = _cached_token(key, now=now) + if cached is not None: + return cached + body: dict[str, Any] = {} if repository_ids: body["repository_ids"] = list(repository_ids) @@ -73,7 +132,11 @@ async def get_github_app_installation_token_with_expiry( ) response.raise_for_status() data = response.json() - return data.get("token"), data.get("expires_at") + token, expires_at = data.get("token"), data.get("expires_at") + parsed = _parse_expiry(expires_at) + if isinstance(token, str) and token and parsed is not None: + _TOKEN_CACHE[key] = (token, expires_at, parsed - _TOKEN_CACHE_MARGIN) + return token, expires_at except Exception: logger.exception("Failed to get GitHub App installation token") return None, None diff --git a/tests/test_github_app.py b/tests/test_github_app.py index 795d8798..38c37ef8 100644 --- a/tests/test_github_app.py +++ b/tests/test_github_app.py @@ -1,5 +1,6 @@ from __future__ import annotations +from datetime import UTC, datetime, timedelta from typing import Any import pytest @@ -7,6 +8,13 @@ import pytest from agent.utils import github_app +@pytest.fixture(autouse=True) +def _clear_token_cache() -> Any: + github_app.clear_app_token_cache() + yield + github_app.clear_app_token_cache() + + class _FakeResponse: def raise_for_status(self) -> None: pass @@ -29,6 +37,90 @@ class _FakeAsyncClient: return _FakeResponse() +def _configure(monkeypatch: pytest.MonkeyPatch, client_cls: type) -> None: + monkeypatch.setattr(github_app, "GITHUB_APP_ID", "1") + monkeypatch.setattr(github_app, "GITHUB_APP_PRIVATE_KEY", "key") + monkeypatch.setattr(github_app, "GITHUB_APP_INSTALLATION_ID", "2") + monkeypatch.setattr(github_app, "_generate_app_jwt", lambda: "jwt") + monkeypatch.setattr(github_app.httpx, "AsyncClient", client_cls) + + +class _CountingResponse: + def __init__(self, expires_at: str) -> None: + self._expires_at = expires_at + + def raise_for_status(self) -> None: + pass + + def json(self) -> dict[str, str]: + return {"token": "tok-123", "expires_at": self._expires_at} + + +class _CountingClient: + posts = 0 + expires_at = "2099-01-01T00:00:00Z" + + async def __aenter__(self) -> _CountingClient: + return self + + async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None: + return None + + async def post(self, url: str, **kwargs: Any) -> _CountingResponse: + type(self).posts += 1 + return _CountingResponse(type(self).expires_at) + + +@pytest.mark.asyncio +async def test_token_is_cached_until_near_expiry(monkeypatch: pytest.MonkeyPatch) -> None: + future = (datetime.now(UTC) + timedelta(hours=1)).isoformat() + + class Client(_CountingClient): + posts = 0 + expires_at = future + + _configure(monkeypatch, Client) + + t1, _ = await github_app.get_github_app_installation_token_with_expiry() + t2, _ = await github_app.get_github_app_installation_token_with_expiry() + + assert t1 == t2 == "tok-123" + assert Client.posts == 1 # second call served from the in-process cache + + +@pytest.mark.asyncio +async def test_cache_is_scoped_per_repository_set(monkeypatch: pytest.MonkeyPatch) -> None: + future = (datetime.now(UTC) + timedelta(hours=1)).isoformat() + + class Client(_CountingClient): + posts = 0 + expires_at = future + + _configure(monkeypatch, Client) + + await github_app.get_github_app_installation_token_with_expiry(repositories=["a"]) + await github_app.get_github_app_installation_token_with_expiry(repositories=["b"]) + await github_app.get_github_app_installation_token_with_expiry(repositories=["a"]) + + assert Client.posts == 2 # distinct scopes mint separately; the repeat is cached + + +@pytest.mark.asyncio +async def test_near_expiry_token_is_not_cached(monkeypatch: pytest.MonkeyPatch) -> None: + soon = (datetime.now(UTC) + timedelta(minutes=2)).isoformat() + + class Client(_CountingClient): + posts = 0 + expires_at = soon + + _configure(monkeypatch, Client) + + await github_app.get_github_app_installation_token_with_expiry() + await github_app.get_github_app_installation_token_with_expiry() + + assert Client.posts == 2 # within the safety margin -> re-minted every call + + @pytest.mark.asyncio async def test_installation_token_can_be_scoped_to_repository_ids( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/test_review_api.py b/tests/test_review_api.py index 257541a3..2e66fb59 100644 --- a/tests/test_review_api.py +++ b/tests/test_review_api.py @@ -1,6 +1,7 @@ import pytest from fastapi import HTTPException +from agent.dashboard import review_api from agent.dashboard.review_api import ( _ALLOWED_IMAGE_CONTENT_TYPES, _finding_counts, @@ -10,6 +11,7 @@ from agent.dashboard.review_api import ( _serialize_finding, _thread_review_summary, classify_finding, + get_pr_head_sha, reviewer_thread_id, ) from agent.webapp import generate_reviewer_thread_id @@ -101,6 +103,33 @@ def test_image_content_type_allowlist_excludes_svg(): assert "image/png" in _ALLOWED_IMAGE_CONTENT_TYPES +@pytest.mark.asyncio +async def test_get_pr_head_sha_returns_head(monkeypatch): + async def fake_token(): + return "tok" + + async def fake_get(path, token, **kwargs): + assert path == "/repos/acme/repo/pulls/7" + return {"head": {"sha": "abc123"}} + + monkeypatch.setattr(review_api, "_require_app_token", fake_token) + monkeypatch.setattr(review_api, "_github_get", fake_get) + assert await get_pr_head_sha("acme", "repo", 7) == "abc123" + + +@pytest.mark.asyncio +async def test_get_pr_head_sha_empty_on_failure(monkeypatch): + async def fake_token(): + return "tok" + + async def fake_get(path, token, **kwargs): + raise HTTPException(404, "not found") + + monkeypatch.setattr(review_api, "_require_app_token", fake_token) + monkeypatch.setattr(review_api, "_github_get", fake_get) + assert await get_pr_head_sha("acme", "repo", 7) == "" + + async def test_require_image_in_pr_rejects_unreferenced_url(monkeypatch): async def fake_github_get(path, token, **kwargs): return {"body": "see ![diagram](https://x.githubusercontent.com/a.png)"} diff --git a/tests/test_review_chat.py b/tests/test_review_chat.py index bedfca54..853295ad 100644 --- a/tests/test_review_chat.py +++ b/tests/test_review_chat.py @@ -6,6 +6,7 @@ from types import SimpleNamespace from typing import Any import pytest +from fastapi import HTTPException from agent.dashboard import review_chat_api @@ -376,7 +377,9 @@ def _client_for_enrich(existing_metadata: dict[str, Any] | None) -> tuple[Any, d return client, captured -def _patch_enrich_deps(monkeypatch, *, metadata: dict[str, Any] | None) -> dict[str, Any]: +def _patch_enrich_deps( + monkeypatch, *, metadata: dict[str, Any] | None, current_head: str = "abc123def456" +) -> dict[str, Any]: client, captured = _client_for_enrich(metadata) monkeypatch.setattr(review_chat_api, "langgraph_client", lambda: client) @@ -389,9 +392,14 @@ def _patch_enrich_deps(monkeypatch, *, metadata: dict[str, Any] | None) -> dict[ async def fake_token(repositories=None): return "app-token" + async def fake_head(owner, repo, pr_number): + captured["head_calls"] = captured.get("head_calls", 0) + 1 + return current_head + monkeypatch.setattr(review_chat_api, "get_review", fake_get_review) monkeypatch.setattr(review_chat_api, "fetch_pr_diff", fake_diff) monkeypatch.setattr(review_chat_api, "get_github_app_installation_token", fake_token) + monkeypatch.setattr(review_chat_api, "get_pr_head_sha", fake_head) return captured @@ -457,6 +465,86 @@ async def test_enrich_chat_command_reseeds_on_head_change(monkeypatch) -> None: assert {"chat_head_sha": "abc123def456"} in captured["updated"] +@pytest.mark.asyncio +async def test_enrich_chat_command_keeps_context_when_reseed_fails(monkeypatch) -> None: + # Existing chat whose head moved, but loading the fresh context fails: the + # command must keep answering from the last seeded context instead of erroring. + captured = _patch_enrich_deps( + monkeypatch, + metadata={"kind": "review_chat", "chat_head_sha": "old-stale-sha"}, + current_head="new-head-sha", + ) + + async def failing_build(*args, **kwargs): + raise HTTPException(404, "review not found") + + monkeypatch.setattr(review_chat_api, "_build_pr_context", failing_build) + command = {"method": "run.start", "params": {"input": {"messages": []}}} + + enriched = await review_chat_api._enrich_chat_command( + command, + owner="acme", + repo="repo", + pr_number=7, + login="octocat", + thread_id="ct-1", + thread_metadata={"kind": "review_chat", "chat_head_sha": "old-stale-sha"}, + ) + + params = enriched["params"] + assert "files" not in params["input"] # no reseed + assert params["config"]["configurable"]["chat_head_sha"] == "old-stale-sha" + assert params["assistant_id"] == "chat" + assert captured["updated"] == [] # head metadata not advanced on failure + + +@pytest.mark.asyncio +async def test_enrich_chat_command_surfaces_reseed_failure_on_create(monkeypatch) -> None: + # A brand-new chat has no prior context to fall back to, so a seeding failure + # must surface rather than silently produce an empty conversation. + _patch_enrich_deps(monkeypatch, metadata=None) + + async def failing_build(*args, **kwargs): + raise HTTPException(404, "review not found") + + monkeypatch.setattr(review_chat_api, "_build_pr_context", failing_build) + command = {"method": "run.start", "params": {"input": {"messages": []}}} + + with pytest.raises(HTTPException): + await review_chat_api._enrich_chat_command( + command, owner="acme", repo="repo", pr_number=7, login="octocat", thread_id="ct-1" + ) + + +@pytest.mark.asyncio +async def test_enrich_chat_command_uses_passed_metadata_without_refetch(monkeypatch) -> None: + # When the caller supplies the already-fetched metadata, enrichment must not + # issue a second thread read on the hot path. + captured = _patch_enrich_deps( + monkeypatch, metadata={"kind": "review_chat", "chat_head_sha": "abc123def456"} + ) + + async def boom(thread_id: str) -> None: + raise AssertionError("metadata was supplied; must not refetch the thread") + + monkeypatch.setattr(review_chat_api, "_get_chat_thread_metadata", boom) + command = {"method": "run.start", "params": {"input": {"messages": []}}} + + enriched = await review_chat_api._enrich_chat_command( + command, + owner="acme", + repo="repo", + pr_number=7, + login="octocat", + thread_id="ct-1", + thread_metadata={"kind": "review_chat", "chat_head_sha": "abc123def456"}, + ) + + assert enriched["params"]["config"]["configurable"]["chat_head_sha"] == "abc123def456" + assert "files" not in enriched["params"]["input"] + assert captured.get("head_calls") == 1 # lightweight head lookup, no full get_review + + @pytest.mark.asyncio async def test_enrich_chat_command_ignores_non_run_start(monkeypatch) -> None: _patch_enrich_deps(monkeypatch, metadata=None)