153 lines
6.1 KiB
Python
153 lines
6.1 KiB
Python
"""Actual async entry points must not block another request's MCP discovery.
|
|
|
|
Run each scenario in a child process: the regression blocks the event loop, so
|
|
an asyncio timeout alone cannot fail the test or clean up the waiting threads.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextvars
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.parametrize("entrypoint", ["open_chat", "open_chat_stream", "workflow"])
|
|
def test_async_entrypoint_coexists_with_cold_mcp_discovery(entrypoint):
|
|
backend_root = Path(__file__).resolve().parents[1]
|
|
env = dict(os.environ)
|
|
env["PYTHONPATH"] = os.pathsep.join(filter(None, (str(backend_root), env.get("PYTHONPATH"))))
|
|
result = subprocess.run(
|
|
[sys.executable, str(Path(__file__).resolve()), entrypoint],
|
|
cwd=backend_root,
|
|
env=env,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=40,
|
|
check=False,
|
|
)
|
|
assert result.returncode == 0, result.stdout + result.stderr
|
|
assert "concurrent MCP initialization completed" in result.stdout
|
|
|
|
|
|
async def _exercise_entrypoint(entrypoint: str) -> None:
|
|
import httpx
|
|
from fastapi import FastAPI
|
|
from langchain_core.messages import AIMessage
|
|
|
|
from app.gateway.routers import open_chat
|
|
from app.gateway.workflow_agent_runner import WorkflowAgentRunner
|
|
from deerflow.agents.lead_agent import agent as lead_agent
|
|
from deerflow.config import app_config
|
|
from deerflow.mcp import cache, tools
|
|
from deerflow.runtime.runs.worker import build_agent_with_timing
|
|
|
|
owner_loop = asyncio.get_running_loop()
|
|
discovery_started = asyncio.Event()
|
|
release_discovery = asyncio.Event()
|
|
built_thread = contextvars.ContextVar("entrypoint_built_thread", default=None)
|
|
discovery_calls = 0
|
|
|
|
async def discover():
|
|
nonlocal discovery_calls
|
|
assert asyncio.get_running_loop() is owner_loop
|
|
discovery_calls += 1
|
|
discovery_started.set()
|
|
await release_discovery.wait()
|
|
return ["discovered tool"]
|
|
|
|
class Agent:
|
|
def __init__(self, thread_id):
|
|
self.thread_id = thread_id
|
|
|
|
async def ainvoke(self, state, *, config):
|
|
assert built_thread.get() == self.thread_id
|
|
return {"messages": [AIMessage(content="answer")]}
|
|
|
|
async def astream(self, state, *, config, stream_mode):
|
|
assert built_thread.get() == self.thread_id
|
|
yield "values", {"messages": [AIMessage(content="answer")]}
|
|
|
|
def factory(config):
|
|
assert cache.get_cached_mcp_tools() == ["discovered tool"]
|
|
thread_id = config["configurable"]["thread_id"]
|
|
built_thread.set(thread_id)
|
|
return Agent(thread_id)
|
|
|
|
factory._deerflow_background_build = True
|
|
app = FastAPI()
|
|
app.include_router(open_chat.router)
|
|
|
|
@app.get("/probe")
|
|
async def probe():
|
|
return {"ok": True}
|
|
|
|
with pytest.MonkeyPatch.context() as patch:
|
|
patch.setattr(cache, "_mcp_tools_cache", None)
|
|
patch.setattr(cache, "_cache_initialized", False)
|
|
patch.setattr(cache, "_config_mtime", None)
|
|
patch.setattr(cache, "_initialization_lock", asyncio.Lock())
|
|
patch.setattr(cache, "_get_config_mtime", lambda: 1.0)
|
|
patch.setattr(tools, "get_mcp_tools", discover)
|
|
patch.setattr(open_chat, "make_lead_agent", factory)
|
|
patch.setattr(lead_agent, "make_lead_agent", factory)
|
|
patch.setattr(
|
|
app_config,
|
|
"get_app_config",
|
|
lambda: SimpleNamespace(workflows=SimpleNamespace(agent_recursion_limit=100)),
|
|
)
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
|
# This is the same background construction path used by run_agent.
|
|
first = asyncio.create_task(
|
|
build_agent_with_timing(factory, {"configurable": {"thread_id": "normal-chat"}}, None, "normal-chat")
|
|
)
|
|
second = None
|
|
try:
|
|
await asyncio.wait_for(discovery_started.wait(), timeout=5)
|
|
if entrypoint == "workflow":
|
|
second = asyncio.create_task(
|
|
WorkflowAgentRunner(app).run_agent(
|
|
run_id="test-run", node_id="test-node", owner_id="test-owner", agent_id="test-agent", prompt="hello"
|
|
)
|
|
)
|
|
else:
|
|
path = "/api/open/chat/stream" if entrypoint == "open_chat_stream" else "/api/open/chat"
|
|
second = asyncio.create_task(client.post(path, json={"message": "hello", "include_files": False}))
|
|
|
|
# Let the real HTTP handler / workflow runner enter construction
|
|
# while discovery holds its lock and is awaiting network work.
|
|
await asyncio.sleep(0.05)
|
|
assert not first.done() and not second.done()
|
|
for _ in range(5):
|
|
response = await asyncio.wait_for(client.get("/probe"), timeout=1)
|
|
assert response.json() == {"ok": True}
|
|
await asyncio.sleep(0)
|
|
finally:
|
|
release_discovery.set()
|
|
results = await asyncio.wait_for(asyncio.gather(first, *([second] if second else [])), timeout=5)
|
|
|
|
assert discovery_calls == 1
|
|
assert built_thread.get() is None # Other request contexts remain isolated.
|
|
result = results[1]
|
|
if entrypoint == "workflow":
|
|
assert result["text"] == "answer"
|
|
assert result["thread_id"] == "wf-test-run-test-node"
|
|
elif entrypoint == "open_chat_stream":
|
|
assert result.status_code == 200
|
|
assert "event: result" in result.text
|
|
assert '"answer": "answer"' in result.text
|
|
assert "event: error" not in result.text
|
|
else:
|
|
assert result.status_code == 200
|
|
assert result.json()["answer"] == "answer"
|
|
print("concurrent MCP initialization completed", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(_exercise_entrypoint(sys.argv[1]))
|