diff --git a/backend/app/gateway/langgraph_auth.py b/backend/app/gateway/langgraph_auth.py index d93acd13df..1064487079 100644 --- a/backend/app/gateway/langgraph_auth.py +++ b/backend/app/gateway/langgraph_auth.py @@ -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, diff --git a/backend/tests/test_langgraph_auth.py b/backend/tests/test_langgraph_auth.py index 804eebc7d7..ae4e49ae0e 100644 --- a/backend/tests/test_langgraph_auth.py +++ b/backend/tests/test_langgraph_auth.py @@ -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": { diff --git a/backend/tests/test_langgraph_studio_routes.py b/backend/tests/test_langgraph_studio_routes.py index a3c4aba467..11634149d0 100644 --- a/backend/tests/test_langgraph_studio_routes.py +++ b/backend/tests/test_langgraph_studio_routes.py @@ -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, ):