tests/test_graph.py

from pathlib import Path
from typing import Any

from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.tools import tool
from langgraph.checkpoint.memory import InMemorySaver

import app.graph as graph_module
from app.auth import AuthenticatedCaller
from app.config import Settings
from app.graph import PreparedRun, ResearchAgent
from app.pricing import PricingRegistry
from app.reports import ReportService
from app.schemas import ChatRequest, ProviderCredentials


class ScriptedModel(BaseChatModel):
    responses: list[Any]

    @property
    def _llm_type(self) -> str:
        return "scripted-test"

    def bind_tools(self, tools: Any, **kwargs: Any) -> "ScriptedModel":
        return self

    def _generate(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        response = self.responses.pop(0)
        if isinstance(response, Exception):
            raise response
        return ChatResult(generations=[ChatGeneration(message=response)])


class ContextAwareModel(BaseChatModel):
    @property
    def _llm_type(self) -> str:
        return "context-aware-test"

    def bind_tools(self, tools: Any, **kwargs: Any) -> "ContextAwareModel":
        return self

    def _generate(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        human_count = sum(isinstance(message, HumanMessage) for message in messages)
        message = AIMessage(
            content=f"human_messages={human_count}",
            usage_metadata={"input_tokens": 10, "output_tokens": 2, "total_tokens": 12},
        )
        return ChatResult(generations=[ChatGeneration(message=message)])


async def test_graph_loops_through_tool_and_emits_metered_final(
    monkeypatch: Any, tmp_path: Path
) -> None:
    scripted = ScriptedModel(
        responses=[
            AIMessage(
                content="",
                tool_calls=[{"name": "test_lookup", "args": {"query": "wildfire"}, "id": "call-1"}],
                usage_metadata={"input_tokens": 100, "output_tokens": 10, "total_tokens": 110},
            ),
            AIMessage(
                content="Final researched answer.",
                usage_metadata={
                    "input_tokens": 200,
                    "output_tokens": 20,
                    "total_tokens": 220,
                    "input_token_details": {"cache_read": 50},
                },
            ),
        ]
    )

    @tool("test_lookup")
    async def test_lookup(query: str) -> str:
        """Return deterministic evidence for a test query."""

        return '{"kind":"search_results","results":[]}'

    monkeypatch.setattr(graph_module, "build_chat_model", lambda *args, **kwargs: scripted)
    monkeypatch.setattr(graph_module, "build_research_tools", lambda *args, **kwargs: [test_lookup])

    settings = Settings(
        environment="test",
        service_token="test-service-token",
        checkpoint_backend="memory",
        artifact_dir=tmp_path,
        run_timeout_seconds=30,
    )
    registry = PricingRegistry.load()
    reports = ReportService(tmp_path)
    agent = ResearchAgent(
        settings=settings,
        registry=registry,
        checkpointer=InMemorySaver(),
        reports=reports,
    )
    request = ChatRequest(
        thread_id="thread-1",
        message="Research wildfire prevention",
        model="openai/gpt-5.6-luna",
        credentials=ProviderCredentials(api_key="sk-test-provider"),
    )
    prepared = PreparedRun(
        request=request,
        caller=AuthenticatedCaller("user-1"),
        definition=registry.get(request.model),
        base_url="https://api.openai.com/v1",
        max_iterations=1,
    )

    events = [event async for event in agent.stream(prepared)]
    names = [event.event for event in events]
    usage = next(event.data for event in events if event.event == "usage")
    final = next(event.data for event in events if event.event == "final")

    assert "tool.started" in names
    assert "tool.completed" in names
    assert final["content"] == "Final researched answer."
    assert usage["input_tokens"] == 250
    assert usage["cache_read_tokens"] == 50
    assert usage["output_tokens"] == 30
    assert names[-1] == "done"


async def test_checkpoint_retains_context_only_for_same_user_thread(
    monkeypatch: Any, tmp_path: Path
) -> None:
    monkeypatch.setattr(
        graph_module, "build_chat_model", lambda *args, **kwargs: ContextAwareModel()
    )
    settings = Settings(
        environment="test",
        service_token="test-service-token",
        checkpoint_backend="memory",
        artifact_dir=tmp_path,
        run_timeout_seconds=30,
    )
    registry = PricingRegistry.load()
    agent = ResearchAgent(
        settings=settings,
        registry=registry,
        checkpointer=InMemorySaver(),
        reports=ReportService(tmp_path),
    )

    async def run(user_id: str, thread_id: str, message: str) -> str:
        request = ChatRequest(
            thread_id=thread_id,
            message=message,
            model="openai/gpt-5.6-luna",
            credentials=ProviderCredentials(api_key="sk-test-provider"),
        )
        prepared = PreparedRun(
            request=request,
            caller=AuthenticatedCaller(user_id),
            definition=registry.get(request.model),
            base_url="https://api.openai.com/v1",
            max_iterations=1,
        )
        events = [event async for event in agent.stream(prepared)]
        return next(event.data["content"] for event in events if event.event == "final")

    assert await run("user-1", "thread-a", "first") == "human_messages=1"
    assert await run("user-1", "thread-a", "second") == "human_messages=2"
    assert await run("user-2", "thread-a", "other user") == "human_messages=1"
    assert await run("user-1", "thread-b", "other thread") == "human_messages=1"


async def test_failed_later_model_call_still_emits_partial_billable_usage(
    monkeypatch: Any, tmp_path: Path
) -> None:
    scripted = ScriptedModel(
        responses=[
            AIMessage(
                content="",
                tool_calls=[{"name": "test_lookup", "args": {"query": "x"}, "id": "call-1"}],
                usage_metadata={"input_tokens": 80, "output_tokens": 10, "total_tokens": 90},
            ),
            RuntimeError("provider failed after first billed call"),
        ]
    )

    @tool("test_lookup")
    async def test_lookup(query: str) -> str:
        """Return deterministic test evidence."""

        return '{"kind":"search_results","results":[]}'

    monkeypatch.setattr(graph_module, "build_chat_model", lambda *args, **kwargs: scripted)
    monkeypatch.setattr(graph_module, "build_research_tools", lambda *args, **kwargs: [test_lookup])
    settings = Settings(
        environment="test",
        service_token="test-service-token",
        checkpoint_backend="memory",
        artifact_dir=tmp_path,
        run_timeout_seconds=30,
    )
    registry = PricingRegistry.load()
    agent = ResearchAgent(
        settings=settings,
        registry=registry,
        checkpointer=InMemorySaver(),
        reports=ReportService(tmp_path),
    )
    request = ChatRequest(
        thread_id="thread-failure",
        message="research this",
        model="openai/gpt-5.6-luna",
        credentials=ProviderCredentials(api_key="sk-test-provider"),
    )
    prepared = PreparedRun(
        request=request,
        caller=AuthenticatedCaller("user-1"),
        definition=registry.get(request.model),
        base_url="https://api.openai.com/v1",
        max_iterations=3,
    )

    events = [event async for event in agent.stream(prepared)]
    usage = next(event.data for event in events if event.event == "usage")

    assert usage["partial"] is True
    assert usage["input_tokens"] == 80
    assert usage["output_tokens"] == 10
    assert [event.event for event in events][-2:] == ["error", "done"]