248 lines
11 KiB
Python
248 lines
11 KiB
Python
"""No-model concurrency diagnostics using actual HTTP routes and SQLite saver."""
|
|
|
|
import asyncio
|
|
import logging
|
|
from contextlib import asynccontextmanager
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
import httpx
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from langgraph.checkpoint.base import empty_checkpoint
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
from starlette.exceptions import HTTPException
|
|
from starlette.requests import Request
|
|
|
|
from app.gateway.operation_timing import OperationTiming
|
|
from app.gateway.routers.threads import ThreadCreateRequest, ThreadSearchRequest, create_thread, router, search_threads
|
|
|
|
|
|
class ThreadStore:
|
|
def __init__(self):
|
|
self.rows = {}
|
|
|
|
async def get(self, thread_id):
|
|
return self.rows.get(thread_id)
|
|
|
|
async def create(self, thread_id, **kwargs):
|
|
self.rows[thread_id] = {"thread_id": thread_id, "metadata": kwargs.get("metadata", {})}
|
|
return self.rows[thread_id]
|
|
|
|
async def search(self, **kwargs):
|
|
return list(self.rows.values())
|
|
|
|
|
|
def client_for(saver, store=None):
|
|
app = FastAPI()
|
|
app.state.checkpointer = saver
|
|
app.state.thread_store = store or ThreadStore()
|
|
app.include_router(router)
|
|
return httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test")
|
|
|
|
|
|
def phase_ms(response):
|
|
return {name: float(value) for name, value in (item.strip().split(";dur=") for item in response.headers["server-timing"].split(","))}
|
|
|
|
|
|
def test_twenty_creates_preserve_immediate_checkpoints_and_do_not_log_content(caplog):
|
|
async def run():
|
|
saver = InMemorySaver()
|
|
async with client_for(saver) as client:
|
|
responses = await asyncio.gather(*(client.post("/api/threads", json={"initialize_checkpoint": True, "metadata": {"private_content": "do-not-log-me"}}) for _ in range(20)))
|
|
ids = set()
|
|
for response in responses:
|
|
assert response.status_code == 200
|
|
thread_id = response.json()["thread_id"]
|
|
ids.add(thread_id)
|
|
assert set(phase_ms(response)) == {"metadata_lookup", "metadata_create", "checkpoint_create", "handler_total"}
|
|
checkpoint = await saver.aget_tuple({"configurable": {"thread_id": thread_id}})
|
|
assert checkpoint is not None
|
|
assert checkpoint.checkpoint["channel_values"] == {}
|
|
assert len(ids) == 20
|
|
|
|
with caplog.at_level(logging.INFO):
|
|
asyncio.run(run())
|
|
assert "GATEWAY_OPERATION_TIMING operation=thread_create outcome=ok" in caplog.text
|
|
assert "checkpointer=InMemorySaver" in caplog.text
|
|
assert "do-not-log-me" not in caplog.text
|
|
|
|
|
|
def test_direct_internal_callers_do_not_need_a_response_argument():
|
|
async def run():
|
|
app = FastAPI()
|
|
app.state.checkpointer = InMemorySaver()
|
|
app.state.thread_store = ThreadStore()
|
|
request = Request({"type": "http", "app": app})
|
|
created = await create_thread(body=ThreadCreateRequest(), request=request)
|
|
found = await search_threads(body=ThreadSearchRequest(), request=request)
|
|
assert [item.thread_id for item in found] == [created.thread_id]
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_analysis_project_calls_real_create_thread_without_response(monkeypatch):
|
|
from app.gateway.routers import analysis_projects
|
|
|
|
async def run():
|
|
project = {"id": "project", "user_id": "owner", "mode": "qa", "name": "Test", "agent_id": None}
|
|
project_store = SimpleNamespace(bind_thread=AsyncMock())
|
|
monkeypatch.setattr(analysis_projects, "owned_project", AsyncMock(return_value=project))
|
|
monkeypatch.setattr(analysis_projects, "project_store", lambda request: project_store)
|
|
monkeypatch.setattr(analysis_projects, "validate_project_agent", AsyncMock())
|
|
app = FastAPI()
|
|
app.state.checkpointer = InMemorySaver()
|
|
app.state.thread_store = ThreadStore()
|
|
request = Request({"type": "http", "app": app})
|
|
created = await analysis_projects.new_project_thread("project", request)
|
|
assert created.metadata["analysis_project_id"] == "project"
|
|
assert created.thread_id in app.state.thread_store.rows
|
|
project_store.bind_thread.assert_awaited_once()
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.parametrize("kind", ["capture", "ask"])
|
|
def test_browser_context_calls_real_create_thread_without_response(monkeypatch, tmp_path, kind):
|
|
from app.gateway.routers import browser_context
|
|
from deerflow.config.browser_context_config import BrowserContextConfig
|
|
|
|
async def run():
|
|
app = FastAPI()
|
|
app.state.checkpointer = InMemorySaver()
|
|
app.state.thread_store = ThreadStore()
|
|
request = Request({"type": "http", "app": app})
|
|
config = SimpleNamespace(browser_context=BrowserContextConfig(enabled=True))
|
|
if kind == "capture":
|
|
monkeypatch.setattr(browser_context, "ensure_uploads_dir", lambda thread_id: tmp_path)
|
|
monkeypatch.setattr(browser_context, "_sync_to_sandbox_if_needed", AsyncMock())
|
|
body = browser_context.BrowserContextCaptureRequest(content="Page body", title="Test")
|
|
result = await browser_context.capture_browser_context.__wrapped__(body, request, config)
|
|
else:
|
|
monkeypatch.setattr(browser_context, "start_run", AsyncMock(return_value=SimpleNamespace(task=None, run_id="test-run", error=None)))
|
|
body = browser_context.BrowserContextAskRequest(content="Page body", question="Summarize")
|
|
result = await browser_context.ask_browser_context.__wrapped__(body, request, config)
|
|
assert result.success
|
|
assert result.thread_id in app.state.thread_store.rows
|
|
assert await app.state.checkpointer.aget_tuple({"configurable": {"thread_id": result.thread_id}}) is None
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_timing_handler_failure_does_not_replace_request_error():
|
|
log = Mock()
|
|
log.warning.side_effect = RuntimeError("broken log destination")
|
|
failure = HTTPException(409, "original error", headers={"X-Test": "retained"})
|
|
with pytest.raises(HTTPException) as captured:
|
|
with OperationTiming(None, log, "test") as timing, timing.phase("metadata_create"):
|
|
raise failure
|
|
assert captured.value is failure
|
|
assert failure.headers["X-Test"] == "retained"
|
|
assert "Server-Timing" in failure.headers
|
|
|
|
|
|
def test_sqlite_history_lock_blocks_create_but_not_metadata_search_or_idempotency(caplog):
|
|
async def run():
|
|
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
|
|
await saver.aput({"configurable": {"thread_id": "existing", "checkpoint_ns": ""}}, empty_checkpoint(), {}, {})
|
|
store = ThreadStore()
|
|
await store.create("existing")
|
|
# Installed SQLite saver holds its lock while alist is suspended
|
|
# at yield, as can happen during checkpoint-history consumption.
|
|
history = saver.alist({"configurable": {"thread_id": "existing"}})
|
|
await anext(history)
|
|
assert saver.lock.locked()
|
|
async with client_for(saver, store) as client:
|
|
create = asyncio.create_task(client.post("/api/threads", json={"thread_id": "new", "initialize_checkpoint": True}))
|
|
try:
|
|
search = await asyncio.wait_for(client.post("/api/threads/search", json={}), timeout=2)
|
|
existing = await asyncio.wait_for(client.post("/api/threads", json={"thread_id": "existing"}), timeout=2)
|
|
assert search.status_code == existing.status_code == 200
|
|
assert set(phase_ms(search)) == {"metadata_search", "handler_total"}
|
|
assert set(phase_ms(existing)) == {"metadata_lookup", "handler_total"}
|
|
await asyncio.sleep(0.06)
|
|
assert not create.done()
|
|
assert "new" in store.rows # Metadata creation already finished.
|
|
finally:
|
|
await history.aclose()
|
|
response = await asyncio.wait_for(create, timeout=2)
|
|
phases = phase_ms(response)
|
|
assert response.status_code == 200
|
|
assert phases["checkpoint_create"] >= 40
|
|
assert phases["checkpoint_create"] > phases["metadata_create"]
|
|
assert await saver.aget_tuple({"configurable": {"thread_id": "new"}}) is not None
|
|
|
|
with caplog.at_level(logging.INFO):
|
|
asyncio.run(run())
|
|
assert "checkpointer=AsyncSqliteSaver" in caplog.text
|
|
|
|
|
|
@pytest.mark.parametrize("stage", ["metadata_create", "checkpoint_create"])
|
|
def test_failures_keep_timing_headers_and_existing_error_contract(stage, caplog):
|
|
class FailingStore(ThreadStore):
|
|
async def create(self, thread_id, **kwargs):
|
|
raise RuntimeError("synthetic metadata failure")
|
|
|
|
class FailingSaver:
|
|
async def aput(self, *args):
|
|
raise RuntimeError("synthetic checkpoint failure")
|
|
|
|
async def run():
|
|
saver = InMemorySaver() if stage == "metadata_create" else FailingSaver()
|
|
store = FailingStore() if stage == "metadata_create" else ThreadStore()
|
|
async with client_for(saver, store) as client:
|
|
response = await client.post("/api/threads", json={"initialize_checkpoint": True})
|
|
assert response.status_code == 500
|
|
assert response.json() == {"detail": "Failed to create thread"}
|
|
assert stage in phase_ms(response)
|
|
|
|
asyncio.run(run())
|
|
assert "outcome=error" in caplog.text
|
|
|
|
|
|
def test_installed_postgres_saver_serializes_cursor_checkout_even_with_pool(monkeypatch):
|
|
"""Characterize the installed saver, without contacting PostgreSQL."""
|
|
from langgraph.checkpoint.postgres import aio
|
|
|
|
async def run():
|
|
pool = object()
|
|
entered = asyncio.Event()
|
|
release = asyncio.Event()
|
|
checkouts = 0
|
|
|
|
@asynccontextmanager
|
|
async def cursor(**kwargs):
|
|
yield object()
|
|
|
|
@asynccontextmanager
|
|
async def get_connection(conn):
|
|
nonlocal checkouts
|
|
assert conn is pool
|
|
checkouts += 1
|
|
yield SimpleNamespace(cursor=cursor)
|
|
|
|
monkeypatch.setattr(aio._ainternal, "get_connection", get_connection)
|
|
saver = aio.AsyncPostgresSaver(pool)
|
|
|
|
async def first():
|
|
async with saver._cursor():
|
|
entered.set()
|
|
await release.wait()
|
|
|
|
async def second():
|
|
async with saver._cursor():
|
|
pass
|
|
|
|
first_task = asyncio.create_task(first())
|
|
await entered.wait()
|
|
second_task = asyncio.create_task(second())
|
|
await asyncio.sleep(0.02)
|
|
assert checkouts == 1
|
|
assert not second_task.done()
|
|
release.set()
|
|
await asyncio.gather(first_task, second_task)
|
|
assert checkouts == 2
|
|
|
|
asyncio.run(run())
|