diff --git a/agent/reviewer.py b/agent/reviewer.py index faca30ad..927d29d4 100644 --- a/agent/reviewer.py +++ b/agent/reviewer.py @@ -38,6 +38,7 @@ from .server import ( graph_loaded_for_execution, ) from .utils.auth import resolve_github_token +from .utils.github_token import get_github_token_from_thread from .utils.model import ModelKwargs, make_model from .utils.sandbox_paths import aresolve_sandbox_work_dir @@ -108,9 +109,14 @@ async def get_reviewer_agent(config: RunnableConfig) -> Pregel: return create_deep_agent(system_prompt="", tools=[]).with_config(config) if config["configurable"].get("source"): - _token, new_encrypted = await resolve_github_token(config, thread_id) - config["metadata"]["github_token_encrypted"] = new_encrypted - del _token + cached_token, cached_encrypted = await get_github_token_from_thread(thread_id) + if cached_token and cached_encrypted: + config["metadata"]["github_token_encrypted"] = cached_encrypted + del cached_token + else: + _token, new_encrypted = await resolve_github_token(config, thread_id) + config["metadata"]["github_token_encrypted"] = new_encrypted + del _token sandbox_backend = await ensure_sandbox_for_thread(thread_id) diff --git a/tests/test_reviewer.py b/tests/test_reviewer.py new file mode 100644 index 00000000..b713c213 --- /dev/null +++ b/tests/test_reviewer.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from langgraph.graph.state import RunnableConfig + +from agent import reviewer + + +class _DummyAgent: + def with_config(self, config: dict[str, object]) -> _DummyAgent: + self.config = config + return self + + +@pytest.mark.asyncio +async def test_reviewer_uses_cached_thread_token_for_slack_review_request() -> None: + config: RunnableConfig = { + "configurable": { + "__is_for_execution__": True, + "thread_id": "reviewer-thread-id", + "source": "slack", + "review_requested": True, + }, + "metadata": {}, + } + dummy_agent = _DummyAgent() + + with ( + patch( + "agent.reviewer.get_github_token_from_thread", + new_callable=AsyncMock, + return_value=("app-token", "encrypted-token"), + ) as mock_get_thread_token, + patch("agent.reviewer.resolve_github_token", new_callable=AsyncMock) as mock_resolve_token, + patch( + "agent.reviewer.ensure_sandbox_for_thread", + new_callable=AsyncMock, + return_value=MagicMock(), + ), + patch( + "agent.reviewer.aresolve_sandbox_work_dir", + new_callable=AsyncMock, + return_value="/workspace", + ), + patch("agent.reviewer.make_model", return_value=MagicMock()), + patch("agent.reviewer.create_deep_agent", return_value=dummy_agent), + ): + await reviewer.get_reviewer_agent(config) + + metadata = config["metadata"] + assert isinstance(metadata, dict) + assert metadata["github_token_encrypted"] == "encrypted-token" + mock_get_thread_token.assert_awaited_once_with("reviewer-thread-id") + mock_resolve_token.assert_not_called()