mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-30 07:00:37 +00:00
fix(threads): preserve concurrent create idempotency in memory mode (#7691)
* fix(threads): preserve concurrent create idempotency in memory mode * fix(threads): serialize detached run metadata creation * docs(gateway): clarify metadata admission timeout budget --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
f7da737118
commit
b86be89f09
@ -3504,6 +3504,10 @@ Thread IDs may be supplied by callers and do not have to be UUIDs. Explicit
|
||||
IDs must contain 1–64 ASCII letters, digits, hyphens, or underscores
|
||||
(`^[A-Za-z0-9_-]{1,64}$`). DeerFlow generates a UUID only when `thread_id` is
|
||||
omitted or `None`; an explicitly supplied empty string is invalid.
|
||||
Concurrent `POST /api/threads` requests and automatic run metadata creation share
|
||||
same-ID admission in one Gateway, preserving the first thread's metadata and status.
|
||||
HTTP retries also preserve its current checkpoint. SQL-backed Gateways retain
|
||||
duplicate-insert recovery across workers.
|
||||
Existing route-addressable threads created under older, looser rules remain
|
||||
readable and deletable, but cannot start new runs or create new filesystem or
|
||||
sandbox state. Legacy deletion skips local path cleanup when the ID is not
|
||||
|
||||
@ -4,6 +4,11 @@ Reject external `sandbox`, `thread_data`, `viewed_images`, `goal_outcome`. Non-n
|
||||
|
||||
Studio retains sanitized creation metadata.
|
||||
|
||||
HTTP and run metadata creation share `goal_thread_lock`. Lock in detached
|
||||
`_ensure_thread_metadata`; preserve await-free attachment and cancel/drain
|
||||
ownership. Its 5s budget covers admission and I/O: timeout fails strict starts;
|
||||
others omit incarnation (MCP fails closed). Tests: `test_thread_create_run_metadata_race.py`.
|
||||
|
||||
Capability Center's `business` adapter validates bundled credentials and creates
|
||||
an MCP connection. The MCP API permits only its isolated launcher and credential
|
||||
environment keys; reject arbitrary commands and manifest-policy bypasses. Keep
|
||||
|
||||
@ -1,5 +1,11 @@
|
||||
# Router contracts
|
||||
|
||||
HTTP creation holds `goal_thread_lock` from the idempotency lookup through the
|
||||
initial checkpoint write. Same-id HTTP creates must return the existing row,
|
||||
including in memory mode; retain the native memory store's overwrite contract
|
||||
and SQL duplicate-insert recovery across workers. Release on failure/cancellation
|
||||
without blocking different thread IDs. Tests: `test_thread_create_idempotency.py`.
|
||||
|
||||
`GET /api/plugins/diagnostics` uses `require_admin_user` before reading the
|
||||
canonical runtime diagnostic snapshot. Return bounded entry-point identifiers,
|
||||
known levels and fixed summaries; raw messages stay in logs. Never clear the
|
||||
|
||||
@ -965,11 +965,17 @@ async def create_thread(body: ThreadCreateRequest, request: Request) -> ThreadRe
|
||||
and an empty checkpoint (so state endpoints work immediately).
|
||||
Idempotent: returns the existing record when ``thread_id`` already exists.
|
||||
"""
|
||||
thread_id = resolve_thread_id(body.thread_id)
|
||||
async with goal_thread_lock(thread_id):
|
||||
return await _create_thread_locked(body, request, thread_id)
|
||||
|
||||
|
||||
async def _create_thread_locked(body: ThreadCreateRequest, request: Request, thread_id: str) -> ThreadResponse:
|
||||
"""Keep the idempotency read and initial writes in one lifecycle lock."""
|
||||
from app.gateway.deps import get_thread_store
|
||||
|
||||
checkpointer = get_checkpointer(request)
|
||||
thread_store = get_thread_store(request)
|
||||
thread_id = resolve_thread_id(body.thread_id)
|
||||
now = now_iso()
|
||||
thread_owner_user_id = get_trusted_internal_owner_user_id(request)
|
||||
thread_owner_kwargs = {"user_id": thread_owner_user_id} if thread_owner_user_id else {}
|
||||
@ -1006,13 +1012,13 @@ async def create_thread(body: ThreadCreateRequest, request: Request) -> ThreadRe
|
||||
raise HTTPException(status_code=404, detail="Thread not found") from None
|
||||
except IntegrityError:
|
||||
# The idempotency read above and this insert are not atomic: a
|
||||
# concurrent request for the same thread_id can commit in between, so
|
||||
# another worker's request for the same thread_id can commit in between, so
|
||||
# the SQL-backed store rejects ours on the duplicate primary key.
|
||||
# Honour the documented idempotency contract by resolving the
|
||||
# now-existing record — running the same owner reconciliation the fast
|
||||
# path does — instead of surfacing the conflict as a 500. The memory
|
||||
# store serializes same-id creates under its per-thread lock and keeps
|
||||
# its historical overwrite behavior, so it does not reach this branch.
|
||||
# path does — instead of surfacing the conflict as a 500. Same-process
|
||||
# HTTP creates share the lifecycle lock; the memory store retains its
|
||||
# historical overwrite behavior for direct callers.
|
||||
existing_record = await _resolve_existing_thread(thread_store, thread_id, thread_owner_user_id, thread_owner_kwargs)
|
||||
if existing_record is not None:
|
||||
return _existing_thread_response(thread_id, existing_record)
|
||||
|
||||
@ -164,6 +164,8 @@ _TERMINAL_RUN_STATUSES = {
|
||||
_TERMINAL_STREAM_RECOVERED = "recovered"
|
||||
_TERMINAL_STREAM_PRODUCER_EXPIRED = "producer_expired"
|
||||
|
||||
# Includes same-thread create admission and metadata I/O; timeout rejects strict
|
||||
# starts and leaves non-strict starts without verified metadata (MCP fails closed).
|
||||
_THREAD_METADATA_SETUP_TIMEOUT_SECONDS = 5.0
|
||||
|
||||
# Message metadata a scheduled launch puts on its prompt: the task, run and the
|
||||
@ -261,6 +263,18 @@ async def _ensure_thread_metadata(
|
||||
require_existing_thread: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""Ensure an admitted run's thread exists without delaying task attachment."""
|
||||
async with goal_thread_lock(record.thread_id):
|
||||
return await _ensure_thread_metadata_locked(run_ctx, record, owner_user_id=owner_user_id, require_existing_thread=require_existing_thread)
|
||||
|
||||
|
||||
async def _ensure_thread_metadata_locked(
|
||||
run_ctx: RunContext,
|
||||
record: RunRecord,
|
||||
*,
|
||||
owner_user_id: str | None,
|
||||
require_existing_thread: bool,
|
||||
) -> dict[str, Any]:
|
||||
"""Share create admission with HTTP while retaining detached task ownership."""
|
||||
thread_store = run_ctx.thread_store
|
||||
existing = await thread_store.get(record.thread_id)
|
||||
if existing is None and owner_user_id:
|
||||
|
||||
241
backend/tests/test_thread_create_idempotency.py
Normal file
241
backend/tests/test_thread_create_idempotency.py
Normal file
@ -0,0 +1,241 @@
|
||||
"""Concurrent HTTP creates preserve the first admitted thread and checkpoint."""
|
||||
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
from uuid import UUID
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from _router_auth_helpers import make_authed_test_app
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.checkpoint.base import empty_checkpoint
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.store.memory import InMemoryStore
|
||||
|
||||
from app.gateway.auth.models import User
|
||||
from app.gateway.routers import threads
|
||||
from deerflow.persistence.thread_meta import ThreadMetaRepository
|
||||
from deerflow.persistence.thread_meta.memory import MemoryThreadMetaStore
|
||||
from deerflow.runtime import goal
|
||||
|
||||
USER = User(id=UUID("00000000-0000-0000-0000-000000000001"), email="create@example.com", password_hash="x", system_role="admin")
|
||||
|
||||
|
||||
@pytest.fixture(params=["memory", "sqlite"])
|
||||
async def thread_store(request, tmp_path):
|
||||
if request.param == "memory":
|
||||
yield MemoryThreadMetaStore(InMemoryStore())
|
||||
return
|
||||
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
|
||||
|
||||
try:
|
||||
await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'threads.db'}", sqlite_dir=str(tmp_path))
|
||||
yield ThreadMetaRepository(get_session_factory())
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
def app_for(store, checkpointer):
|
||||
# Identity is a test double; authorization, owner filtering, the router,
|
||||
# metadata store and checkpoint storage retain their production paths.
|
||||
app = make_authed_test_app(user_factory=lambda: USER, bind_current_user=True)
|
||||
app.state.thread_store = store
|
||||
app.state.checkpointer = checkpointer
|
||||
app.include_router(threads.router)
|
||||
return app
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("assertion", ["metadata", "status", "checkpoint"])
|
||||
async def test_racing_creates_preserve_first_thread(thread_store, monkeypatch, assertion):
|
||||
first_read = asyncio.Event()
|
||||
second_attempt = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
original_get = thread_store.get
|
||||
original_create = thread_store.create
|
||||
real_lock = threads.goal_thread_lock
|
||||
missing_reads = 0
|
||||
lock_attempts = 0
|
||||
creates = 0
|
||||
|
||||
async def observed_get(thread_id, **kwargs):
|
||||
nonlocal missing_reads
|
||||
row = await original_get(thread_id, **kwargs)
|
||||
if thread_id == "same-thread" and row is None:
|
||||
missing_reads += 1
|
||||
if missing_reads == 1:
|
||||
first_read.set()
|
||||
await release.wait()
|
||||
else:
|
||||
# Baseline: both requests have observed no existing row.
|
||||
second_attempt.set()
|
||||
return row
|
||||
|
||||
@asynccontextmanager
|
||||
async def observed_lock(thread_id):
|
||||
nonlocal lock_attempts
|
||||
if thread_id == "same-thread":
|
||||
lock_attempts += 1
|
||||
if lock_attempts == 2:
|
||||
# Fixed path: observe the second lock attempt before it waits.
|
||||
second_attempt.set()
|
||||
async with real_lock(thread_id):
|
||||
yield
|
||||
|
||||
async def observed_create(thread_id, **kwargs):
|
||||
nonlocal creates
|
||||
row = await original_create(thread_id, **kwargs)
|
||||
creates += 1
|
||||
if creates == 1:
|
||||
await thread_store.update_status(thread_id, "busy")
|
||||
await thread_store.update_display_name(thread_id, "Existing conversation")
|
||||
return row
|
||||
|
||||
class ProgressSaver(InMemorySaver):
|
||||
seed_writes = 0
|
||||
|
||||
async def aput(self, config, checkpoint, metadata, versions):
|
||||
self.seed_writes += 1
|
||||
result = await super().aput(config, checkpoint, metadata, versions)
|
||||
if self.seed_writes == 1:
|
||||
# Interleave a real persisted checkpoint between the initial
|
||||
# seed and the late create. Model/worker execution is not claimed.
|
||||
progress = empty_checkpoint()
|
||||
progress["channel_values"] = {"messages": [HumanMessage(content="Keep this history")]}
|
||||
progress["channel_versions"] = {"messages": 1}
|
||||
await super().aput(result, progress, {"source": "update", "step": 0, "writes": {}}, {"messages": 1})
|
||||
return result
|
||||
|
||||
saver = ProgressSaver()
|
||||
monkeypatch.setattr(thread_store, "get", observed_get)
|
||||
monkeypatch.setattr(thread_store, "create", observed_create)
|
||||
monkeypatch.setattr(threads, "goal_thread_lock", observed_lock)
|
||||
app = app_for(thread_store, saver)
|
||||
tasks = []
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
try:
|
||||
tasks.append(asyncio.create_task(client.post("/api/threads", json={"thread_id": "same-thread", "metadata": {"request": "first"}})))
|
||||
await asyncio.wait_for(first_read.wait(), 5)
|
||||
tasks.append(asyncio.create_task(client.post("/api/threads", json={"thread_id": "same-thread", "metadata": {"request": "second"}})))
|
||||
await asyncio.wait_for(second_attempt.wait(), 5)
|
||||
release.set()
|
||||
responses = await asyncio.wait_for(asyncio.gather(*tasks), 5)
|
||||
assert [response.status_code for response in responses] == [200, 200]
|
||||
row = await original_get("same-thread", user_id=str(USER.id))
|
||||
checkpoint = await saver.aget_tuple({"configurable": {"thread_id": "same-thread", "checkpoint_ns": ""}})
|
||||
if assertion == "metadata":
|
||||
assert responses[0].json()["metadata"] == responses[1].json()["metadata"] == row["metadata"]
|
||||
elif assertion == "status":
|
||||
assert row["status"] == "busy"
|
||||
assert row["display_name"] == "Existing conversation"
|
||||
else:
|
||||
assert checkpoint.checkpoint["channel_values"]["messages"][0].content == "Keep this history"
|
||||
assert saver.seed_writes == 1
|
||||
finally:
|
||||
release.set()
|
||||
for task in tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
assert not goal._goal_locks._entries_by_loop
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancelled_waiter_keeps_other_threads_available(thread_store, monkeypatch):
|
||||
attempted = asyncio.Event()
|
||||
real_lock = threads.goal_thread_lock
|
||||
|
||||
@asynccontextmanager
|
||||
async def observed_lock(thread_id):
|
||||
if thread_id == "held-thread":
|
||||
attempted.set()
|
||||
async with real_lock(thread_id):
|
||||
yield
|
||||
|
||||
monkeypatch.setattr(threads, "goal_thread_lock", observed_lock)
|
||||
saver = InMemorySaver()
|
||||
app = app_for(thread_store, saver)
|
||||
task = None
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
try:
|
||||
async with real_lock("held-thread"):
|
||||
task = asyncio.create_task(client.post("/api/threads", json={"thread_id": "held-thread"}))
|
||||
await asyncio.wait_for(attempted.wait(), 5)
|
||||
# A wait on this id must not stop another conversation.
|
||||
other = await asyncio.wait_for(client.post("/api/threads", json={"thread_id": "other-thread"}), 5)
|
||||
assert other.status_code == 200
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert await thread_store.get("held-thread", user_id=str(USER.id)) is None
|
||||
retry = await asyncio.wait_for(client.post("/api/threads", json={"thread_id": "held-thread"}), 5)
|
||||
assert retry.status_code == 200
|
||||
finally:
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
assert not goal._goal_locks._entries_by_loop
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancelled_holder_allows_create_retry(thread_store, monkeypatch):
|
||||
entered = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
original_get = thread_store.get
|
||||
first = True
|
||||
|
||||
async def held_get(thread_id, **kwargs):
|
||||
nonlocal first
|
||||
row = await original_get(thread_id, **kwargs)
|
||||
if thread_id == "cancelled-create" and row is None and first:
|
||||
first = False
|
||||
entered.set()
|
||||
await release.wait()
|
||||
return row
|
||||
|
||||
monkeypatch.setattr(thread_store, "get", held_get)
|
||||
saver = InMemorySaver()
|
||||
app = app_for(thread_store, saver)
|
||||
task = None
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
try:
|
||||
task = asyncio.create_task(client.post("/api/threads", json={"thread_id": "cancelled-create"}))
|
||||
await asyncio.wait_for(entered.wait(), 5)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert await original_get("cancelled-create", user_id=str(USER.id)) is None
|
||||
release.set()
|
||||
retry = await asyncio.wait_for(client.post("/api/threads", json={"thread_id": "cancelled-create", "metadata": {"retry": True}}), 5)
|
||||
assert retry.status_code == 200
|
||||
assert retry.json()["metadata"] == {"retry": True}
|
||||
assert await saver.aget_tuple({"configurable": {"thread_id": "cancelled-create", "checkpoint_ns": ""}}) is not None
|
||||
finally:
|
||||
release.set()
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
assert not goal._goal_locks._entries_by_loop
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_checkpoint_failure_releases_creation_lock(thread_store, monkeypatch):
|
||||
saver = InMemorySaver()
|
||||
original_put = saver.aput
|
||||
|
||||
async def failed_put(*args, **kwargs):
|
||||
raise RuntimeError("injected checkpoint failure")
|
||||
|
||||
monkeypatch.setattr(saver, "aput", failed_put)
|
||||
app = app_for(thread_store, saver)
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
failed = await client.post("/api/threads", json={"thread_id": "failed-seed", "metadata": {"keep": True}})
|
||||
assert failed.status_code == 500
|
||||
assert not goal._goal_locks._entries_by_loop
|
||||
monkeypatch.setattr(saver, "aput", original_put)
|
||||
# Preserve existing retry semantics for a metadata row whose initial
|
||||
# seed failed; this fix does not add transactional seed recovery.
|
||||
retry = await asyncio.wait_for(client.post("/api/threads", json={"thread_id": "failed-seed"}), 5)
|
||||
assert retry.status_code == 200
|
||||
assert retry.json()["metadata"] == {"keep": True}
|
||||
assert not goal._goal_locks._entries_by_loop
|
||||
352
backend/tests/test_thread_create_run_metadata_race.py
Normal file
352
backend/tests/test_thread_create_run_metadata_race.py
Normal file
@ -0,0 +1,352 @@
|
||||
"""The detached run metadata writer participates in thread create admission."""
|
||||
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.store.memory import InMemoryStore
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from test_gateway_services import _stub_app_config as _base_config
|
||||
from test_thread_create_idempotency import USER, app_for
|
||||
from test_thread_create_idempotency import thread_store as _thread_store
|
||||
|
||||
from app.gateway import services
|
||||
from app.gateway.routers import threads
|
||||
from deerflow.runtime import RunManager, goal
|
||||
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
from deerflow.runtime.user_context import reset_current_user, set_current_user
|
||||
|
||||
thread_store = _thread_store
|
||||
_stub_app_config = _base_config
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("first_reader", ["run", "http"])
|
||||
async def test_detached_run_metadata_and_http_create_share_admission(thread_store, monkeypatch, first_reader):
|
||||
thread_id = "cross-entry-create"
|
||||
first_missing = asyncio.Event()
|
||||
competing_attempt = asyncio.Event()
|
||||
release_first = asyncio.Event()
|
||||
original_get = thread_store.get
|
||||
original_create = thread_store.create
|
||||
real_lock = goal.goal_thread_lock
|
||||
creates = []
|
||||
first_task = None
|
||||
|
||||
async def observed_get(identifier, **kwargs):
|
||||
row = await original_get(identifier, **kwargs)
|
||||
if identifier == thread_id and row is None:
|
||||
if asyncio.current_task() is first_task or (first_reader == "http" and first_task is None):
|
||||
first_missing.set()
|
||||
await release_first.wait()
|
||||
else:
|
||||
# Baseline reaches the second missing read; the fixed writer
|
||||
# instead reaches the lock observer before waiting for it.
|
||||
competing_attempt.set()
|
||||
return row
|
||||
|
||||
@asynccontextmanager
|
||||
async def observed_lock(identifier):
|
||||
if identifier == thread_id and first_missing.is_set():
|
||||
competing_attempt.set()
|
||||
async with real_lock(identifier):
|
||||
yield
|
||||
|
||||
saver = InMemorySaver()
|
||||
|
||||
async def observed_create(identifier, **kwargs):
|
||||
row = await original_create(identifier, **kwargs)
|
||||
creates.append(row)
|
||||
if len(creates) == 1:
|
||||
await thread_store.update_status(identifier, "busy")
|
||||
await thread_store.update_display_name(identifier, "Existing title")
|
||||
return row
|
||||
|
||||
monkeypatch.setattr(thread_store, "get", observed_get)
|
||||
monkeypatch.setattr(thread_store, "create", observed_create)
|
||||
monkeypatch.setattr(threads, "goal_thread_lock", observed_lock)
|
||||
monkeypatch.setattr(services, "goal_thread_lock", observed_lock)
|
||||
app = app_for(thread_store, saver)
|
||||
record = await RunManager().create(thread_id, metadata={"source": "run"}, kwargs={"input": {"title": "Run title"}})
|
||||
tasks = []
|
||||
|
||||
async def ensure_from_run():
|
||||
token = set_current_user(USER)
|
||||
try:
|
||||
# Production run_after_metadata owns this writer as a detached
|
||||
# task too. No graph/model execution is simulated by this test.
|
||||
return await services._ensure_thread_metadata(SimpleNamespace(thread_store=thread_store), record, owner_user_id=None)
|
||||
finally:
|
||||
reset_current_user(token)
|
||||
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
|
||||
async def post():
|
||||
return await client.post("/api/threads", json={"thread_id": thread_id, "metadata": {"source": "http"}})
|
||||
|
||||
try:
|
||||
if first_reader == "run":
|
||||
first_task = asyncio.create_task(ensure_from_run())
|
||||
tasks.append(first_task)
|
||||
await asyncio.wait_for(first_missing.wait(), 5)
|
||||
tasks.append(asyncio.create_task(post()))
|
||||
else:
|
||||
# HTTP executes in the middleware's child task, so leave the
|
||||
# first-task marker absent until its first missing read parks.
|
||||
tasks.append(asyncio.create_task(post()))
|
||||
await asyncio.wait_for(first_missing.wait(), 5)
|
||||
first_task = tasks[0]
|
||||
tasks.append(asyncio.create_task(ensure_from_run()))
|
||||
await asyncio.wait_for(competing_attempt.wait(), 5)
|
||||
release_first.set()
|
||||
results = await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), 5)
|
||||
# The SQL loser may raise IntegrityError; run_after_metadata logs
|
||||
# that existing non-fatal failure. The winning row stays intact.
|
||||
assert all(isinstance(value, (dict, httpx.Response, IntegrityError)) for value in results)
|
||||
response = next(value for value in results if isinstance(value, httpx.Response))
|
||||
assert response.status_code == 200
|
||||
row = await original_get(thread_id, user_id=str(USER.id))
|
||||
assert len(creates) == 1
|
||||
assert response.json()["metadata"] == row["metadata"] == creates[0]["metadata"]
|
||||
assert row["status"] == "busy"
|
||||
assert row["display_name"] == "Existing title"
|
||||
finally:
|
||||
release_first.set()
|
||||
for task in tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
assert not goal._goal_locks._entries_by_loop
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_start_run_detached_writer_cannot_replace_http_thread(thread_store, monkeypatch, _stub_app_config):
|
||||
"""Use native admission/task attachment; only agent execution is stubbed."""
|
||||
from app.gateway.routers.thread_runs import RunCreateRequest
|
||||
|
||||
missing = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
http_attempted = asyncio.Event()
|
||||
writer_locked = False
|
||||
writer_task = None
|
||||
original_get = thread_store.get
|
||||
original_create = thread_store.create
|
||||
original_ensure = services._ensure_thread_metadata
|
||||
real_lock = goal.goal_thread_lock
|
||||
creations = []
|
||||
|
||||
async def tracked_ensure(*args, **kwargs):
|
||||
nonlocal writer_task
|
||||
writer_task = asyncio.current_task()
|
||||
return await original_ensure(*args, **kwargs)
|
||||
|
||||
async def held_get(identifier, **kwargs):
|
||||
row = await original_get(identifier, **kwargs)
|
||||
if identifier == "native-admission" and row is None and asyncio.current_task() is writer_task:
|
||||
missing.set()
|
||||
await release.wait()
|
||||
return row
|
||||
|
||||
async def counted_create(*args, **kwargs):
|
||||
row = await original_create(*args, **kwargs)
|
||||
creations.append(row)
|
||||
return row
|
||||
|
||||
@asynccontextmanager
|
||||
async def service_lock(identifier):
|
||||
nonlocal writer_locked
|
||||
if asyncio.current_task() is writer_task:
|
||||
writer_locked = True
|
||||
async with real_lock(identifier):
|
||||
yield
|
||||
|
||||
@asynccontextmanager
|
||||
async def http_lock(identifier):
|
||||
http_attempted.set()
|
||||
async with real_lock(identifier):
|
||||
yield
|
||||
|
||||
async def no_agent(*args, **kwargs):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(services, "_ensure_thread_metadata", tracked_ensure)
|
||||
monkeypatch.setattr(services, "goal_thread_lock", service_lock)
|
||||
monkeypatch.setattr(threads, "goal_thread_lock", http_lock)
|
||||
monkeypatch.setattr(thread_store, "get", held_get)
|
||||
monkeypatch.setattr(thread_store, "create", counted_create)
|
||||
monkeypatch.setattr(services, "resolve_agent_factory", lambda _name: object())
|
||||
monkeypatch.setattr(services, "run_agent", no_agent)
|
||||
app = app_for(thread_store, InMemorySaver())
|
||||
app.state.run_manager = RunManager(store=MemoryRunStore())
|
||||
app.state.stream_bridge = SimpleNamespace()
|
||||
app.state.store = InMemoryStore()
|
||||
app.state.run_event_store = MemoryRunEventStore()
|
||||
request = SimpleNamespace(headers={}, state=SimpleNamespace(user=USER), app=app)
|
||||
record = None
|
||||
http_task = None
|
||||
token = set_current_user(USER)
|
||||
try:
|
||||
record = await services.start_run(RunCreateRequest(input={"messages": [{"role": "user", "content": "hello"}]}, metadata={"source": "run"}), "native-admission", request)
|
||||
await asyncio.wait_for(missing.wait(), 5)
|
||||
assert writer_task is not record.task
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
http_task = asyncio.create_task(client.post("/api/threads", json={"thread_id": "native-admission", "metadata": {"source": "http"}}))
|
||||
await asyncio.wait_for(http_attempted.wait(), 5)
|
||||
if not writer_locked:
|
||||
# Baseline permits this HTTP create to finish while the
|
||||
# detached writer still retains its earlier missing read.
|
||||
await asyncio.wait_for(asyncio.shield(http_task), 5)
|
||||
release.set()
|
||||
await asyncio.wait_for(asyncio.shield(record.task), 5)
|
||||
response = await asyncio.wait_for(http_task, 5)
|
||||
assert response.status_code == 200
|
||||
row = await original_get("native-admission", user_id=str(USER.id))
|
||||
assert len(creations) == 1
|
||||
assert response.json()["metadata"] == row["metadata"]
|
||||
finally:
|
||||
release.set()
|
||||
tasks = [task for task in [http_task, writer_task, record.task if record is not None else None] if task is not None]
|
||||
for task in tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
reset_current_user(token)
|
||||
assert not goal._goal_locks._entries_by_loop
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cancelled_metadata_waiter_releases_its_lock_entry(thread_store, monkeypatch):
|
||||
attempted = asyncio.Event()
|
||||
real_lock = goal.goal_thread_lock
|
||||
|
||||
@asynccontextmanager
|
||||
async def observed_lock(identifier):
|
||||
attempted.set()
|
||||
async with real_lock(identifier):
|
||||
yield
|
||||
|
||||
monkeypatch.setattr(services, "goal_thread_lock", observed_lock)
|
||||
record = await RunManager().create("cancelled-metadata")
|
||||
ctx = SimpleNamespace(thread_store=thread_store)
|
||||
task = None
|
||||
token = set_current_user(USER)
|
||||
try:
|
||||
async with real_lock(record.thread_id):
|
||||
task = asyncio.create_task(services._ensure_thread_metadata(ctx, record, owner_user_id=None))
|
||||
await asyncio.wait_for(attempted.wait(), 5)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert await thread_store.get(record.thread_id) is None
|
||||
row = await asyncio.wait_for(services._ensure_thread_metadata(ctx, record, owner_user_id=None), 5)
|
||||
assert row["thread_id"] == record.thread_id
|
||||
finally:
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
reset_current_user(token)
|
||||
assert not goal._goal_locks._entries_by_loop
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("strict", [False, True], ids=["fallback", "strict"])
|
||||
async def test_metadata_setup_budget_includes_create_admission(thread_store, monkeypatch, _stub_app_config, caplog, strict):
|
||||
from app.gateway.routers.thread_runs import RunCreateRequest
|
||||
from deerflow.runtime.runs.manager import RunStartOutcome
|
||||
from deerflow.runtime.runs.schemas import RunStatus
|
||||
|
||||
thread_id = "metadata-admission-budget"
|
||||
waiting = asyncio.Event()
|
||||
metadata_created = asyncio.Event()
|
||||
admit_metadata = asyncio.Event()
|
||||
waiter_finished = asyncio.Event()
|
||||
original_ensure = services._ensure_thread_metadata
|
||||
real_lock = goal.goal_thread_lock
|
||||
metadata_task = None
|
||||
observed_budgets = []
|
||||
agent_calls = []
|
||||
|
||||
async def tracked_ensure(*args, **kwargs):
|
||||
nonlocal metadata_task
|
||||
metadata_task = asyncio.current_task()
|
||||
metadata_created.set()
|
||||
await admit_metadata.wait()
|
||||
return await original_ensure(*args, **kwargs)
|
||||
|
||||
@asynccontextmanager
|
||||
async def observed_lock(identifier):
|
||||
is_metadata_task = asyncio.current_task() is metadata_task
|
||||
if is_metadata_task:
|
||||
waiting.set()
|
||||
try:
|
||||
async with real_lock(identifier):
|
||||
yield
|
||||
finally:
|
||||
if is_metadata_task:
|
||||
waiter_finished.set()
|
||||
|
||||
async def expire_metadata_budget(tasks, *, timeout, return_when):
|
||||
observed_budgets.append(timeout)
|
||||
# Expire only after the detached writer is waiting for admission.
|
||||
# No wall-clock sleep or five-second test delay is needed.
|
||||
await asyncio.wait_for(waiting.wait(), 5)
|
||||
done, pending = await asyncio.wait(tasks, timeout=0, return_when=return_when)
|
||||
assert not done
|
||||
return done, pending
|
||||
|
||||
async def no_agent(_bridge, manager, record, **kwargs):
|
||||
agent_calls.append((await manager.try_start(record.run_id), kwargs))
|
||||
|
||||
monkeypatch.setattr(services, "_ensure_thread_metadata", tracked_ensure)
|
||||
monkeypatch.setattr(services, "goal_thread_lock", observed_lock)
|
||||
monkeypatch.setattr(services, "asyncio", SimpleNamespace(**{**vars(asyncio), "wait": expire_metadata_budget}))
|
||||
monkeypatch.setattr(services, "resolve_agent_factory", lambda _name: object())
|
||||
monkeypatch.setattr(services, "run_agent", no_agent)
|
||||
app = app_for(thread_store, InMemorySaver())
|
||||
manager = RunManager(store=MemoryRunStore())
|
||||
app.state.run_manager = manager
|
||||
app.state.stream_bridge = SimpleNamespace()
|
||||
app.state.store = InMemoryStore()
|
||||
app.state.run_event_store = MemoryRunEventStore()
|
||||
request = SimpleNamespace(headers={}, state=SimpleNamespace(user=USER), app=app)
|
||||
record = None
|
||||
token = set_current_user(USER)
|
||||
try:
|
||||
row = await thread_store.create(thread_id, user_id=str(USER.id), metadata={"source": "http"})
|
||||
record = await services.start_run(RunCreateRequest(input={"messages": [{"role": "user", "content": "hello"}]}), thread_id, request, require_existing_thread=strict)
|
||||
await asyncio.wait_for(metadata_created.wait(), 5)
|
||||
# A same-thread operation wins admission after start_run releases its
|
||||
# caller-side lock and before the detached metadata writer enters it.
|
||||
async with real_lock(thread_id):
|
||||
admit_metadata.set()
|
||||
await asyncio.wait_for(asyncio.shield(record.task), 5)
|
||||
await asyncio.wait_for(waiter_finished.wait(), 5)
|
||||
assert metadata_task.cancelled()
|
||||
assert observed_budgets == [services._THREAD_METADATA_SETUP_TIMEOUT_SECONDS]
|
||||
assert len(agent_calls) == 1
|
||||
outcome, kwargs = agent_calls[0]
|
||||
assert "thread_incarnation" not in kwargs
|
||||
assert outcome is (RunStartOutcome.cancelled if strict else RunStartOutcome.started)
|
||||
assert record.status is (RunStatus.error if strict else RunStatus.running)
|
||||
assert (await manager.get(record.run_id)).status is record.status
|
||||
if strict:
|
||||
assert record.error == "Timed out verifying existing thread metadata"
|
||||
else:
|
||||
assert "MCP access will fail closed" in caplog.text
|
||||
assert "Timed out ensuring thread_meta" in caplog.text
|
||||
assert await thread_store.get(thread_id) == row
|
||||
assert not goal._goal_locks._entries_by_loop
|
||||
# The cancelled waiter neither writes late nor prevents a later setup.
|
||||
assert await services._ensure_thread_metadata(SimpleNamespace(thread_store=thread_store), record, owner_user_id=None, require_existing_thread=strict) == row
|
||||
finally:
|
||||
tasks = [task for task in [metadata_task, record.task if record is not None else None] if task is not None and task is not asyncio.current_task()]
|
||||
for task in tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
reset_current_user(token)
|
||||
assert not goal._goal_locks._entries_by_loop
|
||||
@ -1126,11 +1126,26 @@ def test_fast_path_trusted_claim_does_not_take_over_owned_row() -> None:
|
||||
assert final_record["user_id"] == "owner-a"
|
||||
|
||||
|
||||
def test_insert_race_concurrent_trusted_claims_have_one_winner() -> None:
|
||||
def test_insert_race_concurrent_trusted_claims_have_one_winner(monkeypatch) -> None:
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from app.gateway.auth_disabled import AUTH_SOURCE_INTERNAL
|
||||
from app.gateway.internal_auth import INTERNAL_OWNER_USER_ID_HEADER_NAME, INTERNAL_SYSTEM_ROLE
|
||||
from deerflow.runtime.keyed_lock import AsyncKeyedLockTable
|
||||
|
||||
# These INSERTs represent separate Gateways: they share SQL rows but have
|
||||
# distinct process-local lifecycle locks. Same-Gateway serialization is
|
||||
# covered with the real stores in test_thread_create_idempotency.py.
|
||||
worker_locks = [AsyncKeyedLockTable[str](), AsyncKeyedLockTable[str]()]
|
||||
lock_calls = 0
|
||||
|
||||
def worker_thread_lock(thread_id):
|
||||
nonlocal lock_calls
|
||||
lock_calls += 1
|
||||
assert lock_calls <= len(worker_locks), f"expected two simulated Gateway lock acquisitions, got {lock_calls}"
|
||||
return worker_locks[lock_calls - 1].hold(thread_id)
|
||||
|
||||
monkeypatch.setattr(threads, "goal_thread_lock", worker_thread_lock)
|
||||
|
||||
class _ConcurrentInsertRaceStore(MemoryThreadMetaStore):
|
||||
def __init__(self):
|
||||
@ -1184,6 +1199,7 @@ def test_insert_race_concurrent_trusted_claims_have_one_winner() -> None:
|
||||
assert final_record["user_id"] in owners
|
||||
assert len(failures) == 1
|
||||
assert failures[0].status_code == 500
|
||||
assert all(not table._entries_by_loop for table in worker_locks)
|
||||
|
||||
|
||||
def test_create_thread_maps_memory_owner_conflict_to_404() -> None:
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user