github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-10-27 21:50:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(langgraph): avoid orphan threads on rejected runs (#5762)

Browse Source 
* fix(langgraph): avoid orphan threads on rejected runs

* test(langgraph): cover assistant admission branches

* test(langgraph): drop ineffective context assertions

---------

Co-authored-by: CorgiBoyG <CorgiBoyG@users.noreply.github.com>
This commit is contained in:
RongJie G 2026-09-24 08:42:10 +08:00 committed by GitHub
parent f6269fb6b1
commit a723ee6d3e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 202 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
backend/app/gateway/langgraph_auth.py
View File

@ -101,6 +101,7 @@ async def _read_standalone_thread(thread_id, ctx) -> dict | None:
async def _ensure_standalone_thread_incarnation(
thread_id,
assistant_id,
ctx,
*,
create_if_missing: bool,
@ -124,12 +125,26 @@ async def _ensure_standalone_thread_incarnation(
raise RuntimeError("Standalone LangGraph thread has an invalid incarnation")
return stored
from langgraph_api.utils import AuthContext as RuntimeAuthContext
from langgraph_runtime.database import connect
from langgraph_runtime.ops import Threads
from langgraph_runtime.ops import Assistants, Threads
token = _allow_thread_incarnation_write.set(True)
try:
async with connect() as conn:
auth_token = RuntimeAuthContext.set(None)
try:
assistant_rows = await Assistants.get(conn, assistant_id, ctx=None)
assistant = await anext(assistant_rows, None)
finally:
RuntimeAuthContext.reset(auth_token)
if assistant is None:
raise HTTPException(status_code=404, detail="Assistant not found")
if assistant.get("metadata", {}).get("created_by") != "system":
authorized_rows = await Assistants.get(conn, assistant_id, ctx=ctx)
if await anext(authorized_rows, None) is None:
raise HTTPException(status_code=404, detail="Assistant not found")
rows = await Threads.put(
conn,
thread_id,
@ -173,6 +188,7 @@ async def _bind_standalone_run_incarnation(ctx, value: dict) -> None:
}
incarnation = await _ensure_standalone_thread_incarnation(
thread_id,
value["assistant_id"],
ctx,
create_if_missing=value.get("if_not_exists") == "create",
creation_metadata=creation_metadata,

103
backend/tests/test_langgraph_auth.py
View File

@ -7,6 +7,7 @@ Validates that the LangGraph auth layer enforces the same rules as Gateway:
import asyncio
import os
import sys
from contextvars import ContextVar
from datetime import timedelta
from pathlib import Path
from types import ModuleType, SimpleNamespace
@ -14,6 +15,7 @@ from unittest.mock import AsyncMock, patch
from uuid import uuid4
import pytest
from starlette.exceptions import HTTPException
os.environ.setdefault("AUTH_JWT_SECRET", "test-secret-key-for-langgraph-auth-testing-min-32")
@ -269,6 +271,7 @@ def test_run_admission_uses_persisted_incarnation_and_scrubs_client_values():
thread_id = uuid4()
run_id = uuid4()
value = {
"assistant_id": uuid4(),
"thread_id": thread_id,
"run_id": run_id,
"metadata": {
@ -343,6 +346,7 @@ def test_existing_legacy_thread_uses_explicit_none_without_backfill():
):
incarnation = asyncio.run(
auth_module._ensure_standalone_thread_incarnation(
uuid4(),
uuid4(),
_make_ctx("user-a", action="create_run"),
create_if_missing=False,
@ -352,7 +356,7 @@ def test_existing_legacy_thread_uses_explicit_none_without_backfill():
assert incarnation is None
def test_implicit_create_accepts_legacy_thread_created_by_mixed_version_peer():
def _fake_standalone_runtime(assistant_results, *, put_metadata=None):
class _Connection:
async def __aenter__(self):
return object()
@ -360,16 +364,52 @@ def test_implicit_create_accepts_legacy_thread_created_by_mixed_version_peer():
async def __aexit__(self, *_args):
return None
async def _put_rows():
yield {"metadata": {"legacy": True}}
async def _rows(value):
if value is not None:
yield value
put = AsyncMock(return_value=_put_rows())
remaining_assistants = iter(assistant_results)
auth_context_observations = []
async def _get_assistant(*_args, **_kwargs):
auth_context_observations.append(runtime_auth_context.get())
return _rows(next(remaining_assistants))
async def _put_thread(*_args, **_kwargs):
return _rows({"metadata": put_metadata})
get_assistant = AsyncMock(side_effect=_get_assistant)
put = AsyncMock(side_effect=_put_thread)
runtime_package = ModuleType("langgraph_runtime")
runtime_package.__path__ = []
database_module = ModuleType("langgraph_runtime.database")
database_module.connect = _Connection
ops_module = ModuleType("langgraph_runtime.ops")
ops_module.Assistants = SimpleNamespace(get=get_assistant)
ops_module.Threads = SimpleNamespace(put=put)
utils_module = ModuleType("langgraph_api.utils")
runtime_auth_context = ContextVar("test_runtime_auth_context", default="ambient")
utils_module.AuthContext = runtime_auth_context
return SimpleNamespace(
modules={
"langgraph_runtime": runtime_package,
"langgraph_runtime.database": database_module,
"langgraph_runtime.ops": ops_module,
"langgraph_api.utils": utils_module,
},
get_assistant=get_assistant,
put=put,
runtime_auth_context=runtime_auth_context,
auth_context_observations=auth_context_observations,
)
def test_implicit_create_accepts_legacy_thread_created_by_mixed_version_peer():
runtime = _fake_standalone_runtime(
[{"metadata": {"created_by": "system"}}],
put_metadata={"legacy": True},
)
with (
patch.object(
@ -379,15 +419,12 @@ def test_implicit_create_accepts_legacy_thread_created_by_mixed_version_peer():
),
patch.dict(
sys.modules,
{
"langgraph_runtime": runtime_package,
"langgraph_runtime.database": database_module,
"langgraph_runtime.ops": ops_module,
},
runtime.modules,
),
):
incarnation = asyncio.run(
auth_module._ensure_standalone_thread_incarnation(
uuid4(),
uuid4(),
_make_ctx("user-a", action="create_run"),
create_if_missing=True,
@ -395,12 +432,55 @@ def test_implicit_create_accepts_legacy_thread_created_by_mixed_version_peer():
)
assert incarnation is None
put.assert_awaited_once()
runtime.put.assert_awaited_once()
assert runtime.get_assistant.await_args.kwargs["ctx"] is None
assert runtime.auth_context_observations == [None]
@pytest.mark.parametrize(
"assistant_results",
[
[None],
[{"metadata": {"created_by": "user"}}, None],
],
ids=["missing", "unauthorized"],
)
def test_implicit_create_rejects_unavailable_assistant_before_thread_put(assistant_results):
runtime = _fake_standalone_runtime(assistant_results)
ctx = _make_ctx("user-a", action="create_run")
with (
patch.object(
auth_module,
"_read_standalone_thread",
AsyncMock(return_value=None),
),
patch.dict(sys.modules, runtime.modules),
pytest.raises(HTTPException) as exc_info,
):
asyncio.run(
auth_module._ensure_standalone_thread_incarnation(
uuid4(),
uuid4(),
ctx,
create_if_missing=True,
)
)
assert exc_info.value.status_code == 404
runtime.put.assert_not_awaited()
assert runtime.get_assistant.await_args_list[0].kwargs["ctx"] is None
if len(assistant_results) == 2:
assert runtime.get_assistant.await_args_list[1].kwargs["ctx"] is ctx
assert runtime.auth_context_observations == [None, "ambient"]
else:
assert runtime.auth_context_observations == [None]
@pytest.mark.parametrize("persisted_incarnation", [None, "versioned-incarnation"])
def test_run_admission_preserves_persisted_incarnation_value(persisted_incarnation):
value = {
"assistant_id": uuid4(),
"thread_id": uuid4(),
"metadata": {},
"kwargs": {},
@ -426,6 +506,7 @@ def test_run_admission_preserves_persisted_incarnation_value(persisted_incarnati
def test_run_admission_does_not_invent_incarnation_for_missing_rejected_thread():
value = {
"assistant_id": uuid4(),
"thread_id": uuid4(),
"metadata": {THREAD_INCARNATION_CONTEXT_KEY: "attacker"},
"kwargs": {
@ -453,6 +534,7 @@ def test_run_admission_does_not_invent_incarnation_for_missing_rejected_thread()
def test_run_admission_rejects_invalid_persisted_incarnation():
value = {
"assistant_id": uuid4(),
"thread_id": uuid4(),
"metadata": {},
"kwargs": {},
@ -484,6 +566,7 @@ def test_run_admission_rejects_invalid_persisted_incarnation():
def test_temporary_run_gets_explicit_legacy_incarnation():
value = {
"assistant_id": uuid4(),
"thread_id": None,
"metadata": {},
"kwargs": {

92
backend/tests/test_langgraph_studio_routes.py
View File

@ -496,6 +496,98 @@ def test_studio_implicit_thread_creation_preserves_searchable_metadata(
assert response.json()["metadata"] == metadata
def test_rejected_implicit_run_does_not_leave_an_orphan_thread(
studio_client: httpx.Client,
):
assistant_id = str(uuid4())
with httpx.Client(
base_url=studio_client.base_url,
timeout=10,
trust_env=False,
) as owner_client:
response = owner_client.post(
"/assistants",
json={"assistant_id": assistant_id, "graph_id": "test_graph"},
)
assert response.status_code == 200, response.text
thread_id = str(uuid4())
response = studio_client.post(
f"/threads/{thread_id}/runs/wait",
json={
"assistant_id": assistant_id,
"if_not_exists": "create",
"input": {"messages": []},
"context": {"preserved_context_probe": "kept"},
},
)
assert response.status_code == 404, response.text
response = studio_client.get(f"/threads/{thread_id}")
assert response.status_code == 404, response.text
def test_assistant_owner_can_implicitly_create_thread(
studio_client: httpx.Client,
):
assistant_id = str(uuid4())
thread_id = str(uuid4())
with httpx.Client(
base_url=studio_client.base_url,
timeout=10,
trust_env=False,
) as owner_client:
response = owner_client.post(
"/assistants",
json={"assistant_id": assistant_id, "graph_id": "test_graph"},
)
assert response.status_code == 200, response.text
response = owner_client.post(
f"/threads/{thread_id}/runs/wait",
json={
"assistant_id": assistant_id,
"if_not_exists": "create",
"input": {"messages": []},
"context": {"preserved_context_probe": "kept"},
},
)
assert response.status_code == 200, response.text
response = owner_client.get(f"/threads/{thread_id}")
assert response.status_code == 200, response.text
metadata = response.json()["metadata"]
assert metadata["user_id"] != "langgraph-studio-user"
assert metadata["thread_incarnation"]
def test_non_studio_user_can_implicitly_create_for_system_assistant(
studio_client: httpx.Client,
):
thread_id = str(uuid4())
with httpx.Client(
base_url=studio_client.base_url,
timeout=10,
trust_env=False,
) as client:
scope = _run_scope(
client,
thread_id,
if_not_exists="create",
)
response = client.get(f"/threads/{thread_id}")
assert response.status_code == 200, response.text
metadata = response.json()["metadata"]
assert metadata["user_id"] != "langgraph-studio-user"
assert scope == mcp_session_scope_key(
user_id=metadata["user_id"],
thread_id=thread_id,
thread_incarnation=metadata["thread_incarnation"],
)
def test_studio_stateless_runs_get_distinct_mcp_scopes(
studio_client: httpx.Client,
):