fix: align create-time run contract with official LangGraph API (#29)

* fix: align create-time run contract with langgraph api

* fix: repair create-time stream continuation contract

* fix: align join and stream continuation contract

* test: fix join expectation in live compat suite

* fix: reject unsupported run control fields

Author: webup

src/agentseek_api/api/runs.py

@@ -1,16 +1,33 @@
 from datetime import UTC, datetime
+from typing import Any
 
 from sqlalchemy import select
 
 from fastapi import APIRouter, Depends, Response
+from fastapi.responses import JSONResponse, StreamingResponse
 
 from agentseek_api.core.auth_deps import get_current_user
 from agentseek_api.core.database import db_manager
 from agentseek_api.core.orm import Run, Thread
-from agentseek_api.models.api import RunCreate, RunRead, RunsCancelRequest, ThreadCreate
+from agentseek_api.models.api import (
+    RunCreateStateless,
+    RunCreateStreamingStateless,
+    RunRead,
+    RunsCancelRequest,
+    ThreadCreate,
+)
 from agentseek_api.models.auth import User
 from agentseek_api.services.thread_service import create_thread_for_user
-from agentseek_api.api.runs import create_run, create_run_stream, wait_run
+from agentseek_api.api.runs import (
+    _build_create_run_stream_response,
+    _normalize_stream_modes,
+    _protocol_stream_location,
+    _stream_response_headers,
+    _validate_supported_run_controls,
+    _wait_response_payload,
+    create_run,
+    wait_run,
+)
 
 router = APIRouter(prefix="/runs", tags=["Stateless Runs"])
 
@@ -23,27 +40,69 @@ async def _best_effort_delete_for_runs(run_ids: list[str]) -> None:
 
 
 @router.post("", response_model=RunRead)
-async def create_stateless_run(payload: RunCreate, user: User = Depends(get_current_user)) -> RunRead:
+async def create_stateless_run(payload: RunCreateStateless, user: User = Depends(get_current_user)) -> RunRead:
+    _validate_supported_run_controls(payload, stateless=True)
     thread = await create_thread_for_user(payload=ThreadCreate(metadata={"stateless": True}), user=user)
     return await create_run(thread.thread_id, payload, user)
 
 
-@router.post("/wait", response_model=RunRead)
-async def create_stateless_run_wait(payload: RunCreate, user: User = Depends(get_current_user)) -> RunRead:
+@router.post(
+    "/wait",
+    response_class=JSONResponse,
+    responses={
+        200: {
+            "content": {"application/json": {"schema": {}}},
+            "headers": {
+                "Location": {"schema": {"type": "string"}},
+                "Content-Location": {"schema": {"type": "string"}},
+            },
+        }
+    },
+)
+async def create_stateless_run_wait(payload: RunCreateStreamingStateless, user: User = Depends(get_current_user)) -> JSONResponse:
+    _normalize_stream_modes(payload.stream_mode)
     created = await create_stateless_run(payload, user)
-    if created.status in {"success", "error", "interrupted"}:
-        return created
-    return await wait_run(created.thread_id, created.run_id, user)
-
-
-@router.post("/stream")
-async def create_stateless_run_stream(payload: RunCreate, user: User = Depends(get_current_user)):
+    final_run = created if created.status in {"success", "error", "interrupted"} else await wait_run(created.thread_id, created.run_id, user)
+    return JSONResponse(
+        await _wait_response_payload(final_run, user=user),
+        headers=_stream_response_headers(
+            location=f"/threads/{created.thread_id}/runs/{created.run_id}/join",
+            content_location=f"/threads/{created.thread_id}/runs/{created.run_id}",
+        ),
+    )
+
+
+@router.post(
+    "/stream",
+    response_class=StreamingResponse,
+    responses={
+        200: {
+            "content": {"text/event-stream": {"schema": {"type": "string"}}},
+            "headers": {
+                "Location": {"schema": {"type": "string"}},
+                "Content-Location": {"schema": {"type": "string"}},
+            },
+        }
+    },
+)
+async def create_stateless_run_stream(payload: RunCreateStreamingStateless, user: User = Depends(get_current_user)):
+    stream_modes = _normalize_stream_modes(payload.stream_mode)
+    _validate_supported_run_controls(payload, stateless=True)
     thread = await create_thread_for_user(payload=ThreadCreate(metadata={"stateless": True}), user=user)
-    return await create_run_stream(thread.thread_id, payload, user)
+    created = await create_run(thread.thread_id, payload, user)
+    return _build_create_run_stream_response(
+        thread_id=thread.thread_id,
+        created=created,
+        user=user,
+        stream_modes=stream_modes,
+        after_seq=0,
+        location=_protocol_stream_location(thread_id=thread.thread_id, run_id=created.run_id, stream_modes=stream_modes),
+        content_location=f"/threads/{thread.thread_id}/runs/{created.run_id}",
+    )
 
 
 @router.post("/batch", response_model=list[RunRead])
-async def create_run_batch(payload: list[RunCreate], user: User = Depends(get_current_user)) -> list[RunRead]:
+async def create_run_batch(payload: list[RunCreateStateless], user: User = Depends(get_current_user)) -> list[RunRead]:
     return [await create_stateless_run(item, user) for item in payload]
 
 

src/agentseek_api/api/stateless_runs.py

@@ -1,5 +1,6 @@
 from datetime import datetime
 from typing import Any
+from typing import Literal
 
 from pydantic import BaseModel, ConfigDict, Field
 
@@ -98,13 +99,75 @@ class ThreadRead(BaseModel):
     status: str = "idle"
 
 
-class RunCreate(BaseModel):
+RunStreamMode = Literal[
+    "values",
+    "messages",
+    "messages-tuple",
+    "tasks",
+    "checkpoints",
+    "updates",
+    "events",
+    "debug",
+    "custom",
+]
+RunInterrupt = Literal["*"] | list[str]
+RunDurability = Literal["sync", "async", "exit"]
+RunOnDisconnect = Literal["cancel", "continue"]
+RunOnCompletion = Literal["delete", "keep"]
+RunMultitaskStrategy = Literal["reject", "rollback", "interrupt", "enqueue"]
+RunIfNotExists = Literal["create", "reject"]
+
+
+class RunCreateStateful(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
     assistant_id: str
-    input: Any
+    checkpoint: dict[str, Any] | None = None
+    input: Any = None
+    command: dict[str, Any] | None = None
     metadata: dict[str, Any] = Field(default_factory=dict)
     config: dict[str, Any] = Field(default_factory=dict)
     context: dict[str, Any] = Field(default_factory=dict)
-    multitask_strategy: str = "enqueue"
+    webhook: str | None = None
+    interrupt_before: RunInterrupt | None = None
+    interrupt_after: RunInterrupt | None = None
+    stream_mode: RunStreamMode | list[RunStreamMode] | None = Field(default_factory=lambda: ["values"])
+    stream_subgraphs: bool = False
+    stream_resumable: bool = False
+    feedback_keys: list[str] | None = None
+    multitask_strategy: RunMultitaskStrategy = "enqueue"
+    if_not_exists: RunIfNotExists = "reject"
+    after_seconds: float | None = None
+    checkpoint_during: bool = False
+    durability: RunDurability = "async"
+
+
+class RunCreateStreamingStateful(RunCreateStateful):
+    on_disconnect: RunOnDisconnect = "continue"
+
+
+class RunCreateStateless(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    assistant_id: str
+    input: Any = None
+    command: dict[str, Any] | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    config: dict[str, Any] = Field(default_factory=dict)
+    context: dict[str, Any] = Field(default_factory=dict)
+    webhook: str | None = None
+    stream_mode: RunStreamMode | list[RunStreamMode] | None = Field(default_factory=lambda: ["values"])
+    feedback_keys: list[str] | None = None
+    stream_subgraphs: bool = False
+    stream_resumable: bool = False
+    on_completion: RunOnCompletion = "keep"
+    after_seconds: float | None = None
+    checkpoint_during: bool = False
+    durability: RunDurability = "async"
+
+
+class RunCreateStreamingStateless(RunCreateStateless):
+    on_disconnect: RunOnDisconnect = "continue"
 
 
 class RunsCancelRequest(BaseModel):

src/agentseek_api/models/api.py

@@ -128,8 +128,9 @@ class _OpenMessage:
         open_blocks: dict[int, str] = field(default_factory=dict)
         text_contents: dict[int, str] = field(default_factory=dict)
 
-    def __init__(self, *, thread_id: str) -> None:
+    def __init__(self, *, thread_id: str, run_id: str) -> None:
         self.thread_id = thread_id
+        self.run_id = run_id
         self._open_message_ids: dict[str, _ProtocolMessageStreamState._OpenMessage] = {}
         self.saw_live_messages = False
 
@@ -148,6 +149,7 @@ def _finish_blocks(
                 self.thread_id,
                 index=index,
                 namespace=effective_namespace,
+                run_id=self.run_id,
             )
             del state.open_blocks[index]
 
@@ -166,6 +168,7 @@ def _publish_text_block(
                 index=index,
                 content={"type": "text", "text": ""},
                 namespace=effective_namespace,
+                run_id=self.run_id,
             )
             state.open_blocks[index] = "text"
         previous_text = state.text_contents.get(index, "")
@@ -178,6 +181,7 @@ def _publish_text_block(
                 index=index,
                 delta={"type": "text-delta", "text": delta_text},
                 namespace=effective_namespace,
+                run_id=self.run_id,
             )
         state.text_contents[index] = text
 
@@ -197,13 +201,15 @@ def _publish_nontext_block(
                 index=index,
                 content=block,
                 namespace=effective_namespace,
+                run_id=self.run_id,
             )
             if final:
                 publish_content_block_finish(
                     self.thread_id,
                     index=index,
                     content=block,
                     namespace=effective_namespace,
+                    run_id=self.run_id,
                 )
                 return
             state.open_blocks[index] = str(block.get("type", "block"))
@@ -214,13 +220,15 @@ def _publish_nontext_block(
             index=index,
             delta=block,
             namespace=effective_namespace,
+            run_id=self.run_id,
         )
         if final:
             publish_content_block_finish(
                 self.thread_id,
                 index=index,
                 content=block,
                 namespace=effective_namespace,
+                run_id=self.run_id,
             )
             del state.open_blocks[index]
 
@@ -239,6 +247,7 @@ def publish_blocks(
                 message_id=message_id,
                 role=role,
                 namespace=namespace,
+                run_id=self.run_id,
             )
             state = self._OpenMessage(
                 role=role,
@@ -317,7 +326,7 @@ def finish_all(self, *, namespace: list[str] | None = None) -> None:
             state = self._open_message_ids.pop(message_id)
             message_namespace = state.namespace or namespace
             self._finish_blocks(state, namespace=message_namespace)
-            publish_message_complete(self.thread_id, namespace=message_namespace)
+            publish_message_complete(self.thread_id, namespace=message_namespace, run_id=self.run_id)
 
     async def afinish_blocks(
         self,
@@ -334,6 +343,7 @@ async def afinish_blocks(
                 self.thread_id,
                 index=index,
                 namespace=effective_namespace,
+                run_id=self.run_id,
             )
             del state.open_blocks[index]
 
@@ -352,6 +362,7 @@ async def apublish_text_block(
                 index=index,
                 content={"type": "text", "text": ""},
                 namespace=effective_namespace,
+                run_id=self.run_id,
             )
             state.open_blocks[index] = "text"
         previous_text = state.text_contents.get(index, "")
@@ -364,6 +375,7 @@ async def apublish_text_block(
                 index=index,
                 delta={"type": "text-delta", "text": delta_text},
                 namespace=effective_namespace,
+                run_id=self.run_id,
             )
         state.text_contents[index] = text
 
@@ -383,13 +395,15 @@ async def apublish_nontext_block(
                 index=index,
                 content=block,
                 namespace=effective_namespace,
+                run_id=self.run_id,
             )
             if final:
                 await apublish_content_block_finish(
                     self.thread_id,
                     index=index,
                     content=block,
                     namespace=effective_namespace,
+                    run_id=self.run_id,
                 )
                 return
             state.open_blocks[index] = str(block.get("type", "block"))
@@ -400,13 +414,15 @@ async def apublish_nontext_block(
             index=index,
             delta=block,
             namespace=effective_namespace,
+            run_id=self.run_id,
         )
         if final:
             await apublish_content_block_finish(
                 self.thread_id,
                 index=index,
                 content=block,
                 namespace=effective_namespace,
+                run_id=self.run_id,
             )
             del state.open_blocks[index]
 
@@ -425,6 +441,7 @@ async def apublish_blocks(
                 message_id=message_id,
                 role=role,
                 namespace=namespace,
+                run_id=self.run_id,
             )
             state = self._OpenMessage(
                 role=role,
@@ -503,7 +520,7 @@ async def afinish_all(self, *, namespace: list[str] | None = None) -> None:
             state = self._open_message_ids.pop(message_id)
             message_namespace = state.namespace or namespace
             await self.afinish_blocks(state, namespace=message_namespace)
-            await apublish_message_complete(self.thread_id, namespace=message_namespace)
+            await apublish_message_complete(self.thread_id, namespace=message_namespace, run_id=self.run_id)
 
 
 def _protocol_blocks_for_message(message: BaseMessage) -> list[dict[str, Any]]:
@@ -734,7 +751,7 @@ async def execute_run(
     result: Any = None
     interrupt_chunk: Any = None
     interrupt_namespace: list[str] | None = None
-    protocol_messages = _ProtocolMessageStreamState(thread_id=thread_id)
+    protocol_messages = _ProtocolMessageStreamState(thread_id=thread_id, run_id=run_id)
     async for stream_event in graph.astream_events(invocation, config, version="v2"):
         protocol_namespace = _protocol_namespace_for_event(stream_event)
         for event_name, event_payload in _translate_stream_events(stream_event):
@@ -804,7 +821,12 @@ async def execute_run(
             if isinstance(normalized_chunk, dict):
                 normalized_chunk.pop("__interrupt__", None)
                 if normalized_chunk:
-                    await apublish_updates_event(thread_id, values=normalized_chunk, namespace=protocol_namespace)
+                    await apublish_updates_event(
+                        thread_id,
+                        values=normalized_chunk,
+                        namespace=protocol_namespace,
+                        run_id=run_id,
+                    )
         if stream_event.get("event") == "on_chain_end" and _is_root_stream_event(stream_event):
             data = stream_event.get("data", {})
             if isinstance(data, dict) and "output" in data:
@@ -818,7 +840,12 @@ async def execute_run(
                         else:
                             await apublish_message_transcript(thread_id, run_id=run_id, messages=messages)
                     await protocol_messages.afinish_all()
-                    await apublish_values_event(thread_id, values=normalized_result, namespace=protocol_namespace)
+                    await apublish_values_event(
+                        thread_id,
+                        values=normalized_result,
+                        namespace=protocol_namespace,
+                        run_id=run_id,
+                    )
 
     if interrupt_chunk is not None:
         if isinstance(result, dict):