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] <open-swe@users.noreply.github.com>

* 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] <open-swe@users.noreply.github.com>

---------

Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>

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

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:
+        current_head = await get_pr_head_sha(owner, repo, pr_number)
-            review = await get_review(owner, repo, pr_number)
+        needs_seed = bool(current_head) and current_head != stored_head
-        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
 
     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)
+            # An existing chat can keep answering from its last seeded context, so
-            raise HTTPException(502, "could not load PR context") from exc
+            # a transient reseed failure shouldn't break the conversation. A fresh
-        if head_sha:
+            # chat (or one never seeded) has nothing to fall back to, so surface it.
-            configurable["chat_head_sha"] = head_sha
+            if created or not stored_head:
-            await langgraph_client().threads.update(
+                if isinstance(exc, HTTPException):
-                thread_id=thread_id, metadata={"chat_head_sha": head_sha}
+                    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")
+        else:
-        if not isinstance(run_input, dict):
+            if head_sha:
-            run_input = {}
+                configurable["chat_head_sha"] = head_sha
-        existing_files = run_input.get("files")
+                await langgraph_client().threads.update(
-        run_input["files"] = {
+                    thread_id=thread_id, metadata={"chat_head_sha": head_sha}
-            **(existing_files if isinstance(existing_files, dict) else {}),
+                )
-            **pr_files,
+            elif stored_head:
-        }
+                configurable["chat_head_sha"] = stored_head
-        params["input"] = run_input
+            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).
+    # below on the first `run.start` (with the caller as owner). Reuse the
-    await assert_chat_thread_access(thread_id, owner, repo, pr_number, login)
+    # 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)

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

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,

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)"}

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)