285 lines
14 KiB
Python
285 lines
14 KiB
Python
"""Lazy creation and history isolation through real HTTP, SQL and graph paths."""
|
|
|
|
import asyncio
|
|
import json
|
|
from contextlib import asynccontextmanager
|
|
from types import SimpleNamespace
|
|
from typing import Annotated, TypedDict
|
|
from unittest.mock import AsyncMock
|
|
|
|
import httpx
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
from langgraph.checkpoint.base import empty_checkpoint
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
|
from langgraph.graph import END, START, StateGraph
|
|
from langgraph.graph.message import add_messages
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from app.gateway.authz import AuthContext
|
|
from app.gateway.routers import threads
|
|
from deerflow.persistence.task_deeplink_threads.model import TaskDeeplinkThreadRow
|
|
from deerflow.persistence.thread_meta.model import ThreadMetaRow
|
|
from deerflow.persistence.thread_meta.sql import ThreadMetaRepository
|
|
from deerflow.runtime.user_context import reset_current_user, set_current_user
|
|
|
|
|
|
@asynccontextmanager
|
|
async def client_for(saver):
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(ThreadMetaRow.__table__.create)
|
|
await conn.run_sync(TaskDeeplinkThreadRow.__table__.create)
|
|
app = FastAPI()
|
|
app.state.thread_store = ThreadMetaRepository(async_sessionmaker(engine, expire_on_commit=False))
|
|
app.state.checkpointer = saver
|
|
app.include_router(threads.router)
|
|
|
|
@app.middleware("http")
|
|
async def authenticated(request, call_next):
|
|
name = request.headers.get("X-User", "owner")
|
|
user = SimpleNamespace(id=name) if name != "anonymous" else None
|
|
request.state.auth = AuthContext(user, ["threads:read", "threads:write", "runs:read", "runs:create", "runs:cancel"] if user else [])
|
|
token = set_current_user(user) if user else None
|
|
try:
|
|
return await call_next(request)
|
|
finally:
|
|
if token:
|
|
reset_current_user(token)
|
|
|
|
try:
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
|
yield client, app
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
def test_create_list_and_metadata_do_not_touch_checkpoint(monkeypatch):
|
|
async def run():
|
|
async with client_for(InMemorySaver()) as (client, app):
|
|
def unavailable(_):
|
|
raise AssertionError("Metadata operation accessed the checkpoint backend")
|
|
|
|
monkeypatch.setattr(threads, "get_checkpointer", unavailable)
|
|
metadata = {"taskId": "task", "actionId": "action", "agent_id": "agent", "owner_id": "forged"}
|
|
created = await client.post("/api/threads", json={"thread_id": "lazy", "metadata": metadata})
|
|
assert created.status_code == 200, created.text
|
|
assert created.json()["metadata"] == {k: v for k, v in metadata.items() if k != "owner_id"}
|
|
assert "checkpoint_create" not in created.headers["server-timing"]
|
|
duplicate = await client.post("/api/threads", json={"thread_id": "lazy", "metadata": {"changed": True}})
|
|
assert duplicate.json()["metadata"] == created.json()["metadata"]
|
|
record = await app.state.thread_store.get("lazy", user_id=None)
|
|
assert record["user_id"] == "owner"
|
|
listing = await client.post("/api/threads/search", json={"metadata": {"taskId": "task", "actionId": "action", "agent_id": "agent"}}, headers={"X-User": "other"})
|
|
assert [row["thread_id"] for row in listing.json()] == ["lazy"]
|
|
info = await client.get("/api/threads/lazy", headers={"X-User": "other"})
|
|
assert info.status_code == 200
|
|
assert info.json()["metadata"] == created.json()["metadata"]
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_roundtable_init_stream_creates_only_metadata_and_keeps_seat_mapping(monkeypatch):
|
|
from app.gateway.routers import multi_agent
|
|
|
|
monkeypatch.setattr(multi_agent, "ensure_roundtable_functional_agents", lambda: None)
|
|
monkeypatch.setattr(multi_agent, "_get_agents_batch_direct", AsyncMock(return_value={"seat-a": {"name": "A"}, "seat-b": {"name": "B"}}))
|
|
monkeypatch.setattr(multi_agent, "_make_loopback_client", lambda *_: pytest.fail("init must not seed checkpoint via HTTP"))
|
|
async with client_for(InMemorySaver()) as (client, app):
|
|
del app.state.checkpointer
|
|
app.include_router(multi_agent.router)
|
|
response = await client.post("/api/multi-agent/init", json={"agent_names": ["seat-a", "seat-b"]})
|
|
assert response.status_code == 200
|
|
frames = [json.loads(line.removeprefix("data: ")) for line in response.text.splitlines() if line.startswith("data: ")]
|
|
assert [frame["event"] for frame in frames] == ["seat_ready", "seat_ready", "coordinator_ready", "init_done"]
|
|
mapping = frames[-1]["thread_ids"]
|
|
assert set(mapping) == {"seat-a", "seat-b", multi_agent.COORDINATOR_AGENT_ID}
|
|
for thread_id in mapping.values():
|
|
record = await app.state.thread_store.get(thread_id, user_id=None)
|
|
assert record["metadata"]["thread_type"] == "roundtable"
|
|
|
|
|
|
def test_lazy_create_completes_while_real_saver_history_lock_is_held():
|
|
async def run():
|
|
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
|
|
await saver.aput({"configurable": {"thread_id": "old", "checkpoint_ns": ""}}, empty_checkpoint(), {}, {})
|
|
history = saver.alist({"configurable": {"thread_id": "old"}})
|
|
await anext(history)
|
|
try:
|
|
assert saver.lock.locked()
|
|
async with client_for(saver) as (client, _):
|
|
results = await asyncio.wait_for(asyncio.gather(*(client.post("/api/threads", json={"thread_id": f"new-{i}"}) for i in range(20))), timeout=3)
|
|
assert all(response.status_code == 200 for response in results)
|
|
assert saver.lock.locked()
|
|
finally:
|
|
await history.aclose()
|
|
assert await saver.aget_tuple({"configurable": {"thread_id": "new-0"}}) is None
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_first_run_initializes_checkpoint_and_second_run_keeps_conversation():
|
|
class State(TypedDict):
|
|
messages: Annotated[list, add_messages]
|
|
|
|
async def run():
|
|
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver, client_for(saver) as (client, _):
|
|
response = await client.post("/api/threads", json={"thread_id": "conversation"})
|
|
assert response.status_code == 200
|
|
config = {"configurable": {"thread_id": "conversation", "checkpoint_ns": ""}}
|
|
assert await saver.aget_tuple(config) is None
|
|
graph = StateGraph(State)
|
|
graph.add_node("answer", lambda state: {"messages": [AIMessage(content=f"reply {len(state['messages'])}")]})
|
|
graph.add_edge(START, "answer")
|
|
graph.add_edge("answer", END)
|
|
compiled = graph.compile(checkpointer=saver)
|
|
first = await compiled.ainvoke({"messages": [HumanMessage(content="one")]}, config)
|
|
assert len(first["messages"]) == 2
|
|
second = await compiled.ainvoke({"messages": [HumanMessage(content="two")]}, config)
|
|
assert [message.content for message in second["messages"]] == ["one", "reply 1", "two", "reply 3"]
|
|
history = await client.post("/api/threads/conversation/history", json={"limit": 1})
|
|
assert [message["content"] for message in history.json()[0]["values"]["messages"]] == ["one", "reply 1", "two", "reply 3"]
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.parametrize("shared", [False, True])
|
|
def test_metadata_only_threads_keep_private_and_task_sharing_rules(shared):
|
|
async def run():
|
|
async with client_for(InMemorySaver()) as (client, _):
|
|
metadata = {"taskId": "task", "actionId": "action", "agent_id": "agent"} if shared else {"agent_id": "agent"}
|
|
assert (await client.post("/api/threads", json={"thread_id": "lazy", "metadata": metadata})).status_code == 200
|
|
for user, expected in [("owner", 200), ("other", 200 if shared else 404), ("anonymous", 401)]:
|
|
headers = {"X-User": user}
|
|
info = await client.get("/api/threads/lazy", headers=headers)
|
|
state = await client.get("/api/threads/lazy/state", headers=headers)
|
|
history = await client.post("/api/threads/lazy/history", json={"limit": 1}, headers=headers)
|
|
assert info.status_code == state.status_code == history.status_code == expected
|
|
if expected == 200:
|
|
assert state.json()["values"] == {}
|
|
assert state.json()["metadata"] == metadata
|
|
assert history.json() == []
|
|
assert (await client.get("/api/threads/missing/state")).status_code == 404
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_limit_one_history_uses_latest_tuple_not_the_iterator(monkeypatch):
|
|
class State(TypedDict):
|
|
messages: Annotated[list, add_messages]
|
|
|
|
async def run():
|
|
async with client_for(InMemorySaver()) as (client, app):
|
|
config = {"configurable": {"thread_id": "hello", "checkpoint_ns": ""}}
|
|
graph = StateGraph(State)
|
|
graph.add_node("answer", lambda state: {"messages": [AIMessage(content="你好")]})
|
|
graph.add_edge(START, "answer")
|
|
graph.add_edge("answer", END)
|
|
compiled = graph.compile(checkpointer=app.state.checkpointer)
|
|
await compiled.ainvoke({"messages": [HumanMessage(content="hi")]}, config)
|
|
|
|
async def forbidden(*_args, **_kwargs):
|
|
raise AssertionError("limit=1 history must not walk alist")
|
|
yield None
|
|
|
|
monkeypatch.setattr(app.state.checkpointer, "alist", forbidden)
|
|
result = await client.post("/api/threads/hello/history", json={"limit": 1})
|
|
assert result.status_code == 200, result.text
|
|
contents = [message["content"] for message in result.json()[0]["values"]["messages"]]
|
|
assert contents == ["hi", "你好"]
|
|
assert "checkpoint_load" in result.headers["server-timing"]
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_history_releases_saver_lock_before_serialization_and_ignores_subgraphs(monkeypatch):
|
|
async def run():
|
|
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver, client_for(saver) as (client, _):
|
|
for namespace, title in [("", "root"), ("agent:subgraph", "child")]:
|
|
checkpoint = empty_checkpoint()
|
|
checkpoint["channel_values"] = {"title": title, "messages": [AIMessage(content=title)], "todos": [title]}
|
|
await saver.aput({"configurable": {"thread_id": "thread", "checkpoint_ns": namespace}}, checkpoint, {}, {})
|
|
serialize = threads.serialize_channel_values
|
|
|
|
def after_release(values, **kwargs):
|
|
assert not saver.lock.locked(), "Serialization held the shared saver lock"
|
|
return serialize(values, **kwargs)
|
|
|
|
monkeypatch.setattr(threads, "serialize_channel_values", after_release)
|
|
result = await client.post("/api/threads/thread/history", json={"limit": 10})
|
|
assert result.status_code == 200, result.text
|
|
assert len(result.json()) == 1
|
|
assert result.json()[0]["values"]["todos"] == ["root"]
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_cancelled_history_closes_iterator():
|
|
async def run():
|
|
closed = asyncio.Event()
|
|
entered = asyncio.Event()
|
|
|
|
class Saver:
|
|
async def alist(self, *args, **kwargs):
|
|
try:
|
|
entered.set()
|
|
await asyncio.Event().wait()
|
|
yield None
|
|
finally:
|
|
closed.set()
|
|
|
|
async with client_for(Saver()) as (client, _):
|
|
task = asyncio.create_task(client.post("/api/threads/thread/history", json={"limit": 1}))
|
|
await entered.wait()
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert closed.is_set()
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.parametrize("kind", ["foreground", "background"])
|
|
def test_roundtable_internal_creation_defers_checkpoint_and_can_run_twice(kind):
|
|
from starlette.requests import Request
|
|
|
|
from app.gateway import roundtable_inprocess_gateway
|
|
from app.gateway.routers import multi_agent
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, add_messages]
|
|
|
|
async def run():
|
|
saver = InMemorySaver()
|
|
async with client_for(saver) as (client, app):
|
|
token = set_current_user(SimpleNamespace(id="owner"))
|
|
try:
|
|
# Both helpers must work without even obtaining a checkpointer.
|
|
del app.state.checkpointer
|
|
if kind == "foreground":
|
|
tid = await multi_agent._create_thread_direct(Request({"type": "http", "app": app}), metadata={"thread_type": "roundtable", "taskId": "shared", "agent_id": "agent"})
|
|
else:
|
|
tid = await roundtable_inprocess_gateway._create_thread(app)
|
|
record = await app.state.thread_store.get(tid)
|
|
assert record["metadata"]["thread_type"] == "roundtable"
|
|
assert record["user_id"] == "owner"
|
|
config = {"configurable": {"thread_id": tid, "checkpoint_ns": ""}}
|
|
assert await saver.aget_tuple(config) is None
|
|
app.state.checkpointer = saver
|
|
assert (await client.get(f"/api/threads/{tid}", headers={"X-User": "other"})).status_code == 200
|
|
graph = StateGraph(State)
|
|
graph.add_node("answer", lambda state: {"messages": [AIMessage(content="reply")]})
|
|
graph.add_edge(START, "answer")
|
|
graph.add_edge("answer", END)
|
|
compiled = graph.compile(checkpointer=saver)
|
|
await compiled.ainvoke({"messages": [HumanMessage(content="first")]}, config)
|
|
result = await compiled.ainvoke({"messages": [HumanMessage(content="second")]}, config)
|
|
assert [message.content for message in result["messages"]] == ["first", "reply", "second", "reply"]
|
|
finally:
|
|
reset_current_user(token)
|
|
|
|
asyncio.run(run())
|