Add real-provider create-time contract coverage (#30)

Squash merge PR #30 after green default CI and Live Provider Streaming.

Author: webup

src/agentseek_api/api/runs.py

@@ -1,10 +1,9 @@
 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 fastapi.responses import StreamingResponse
 
 from agentseek_api.core.auth_deps import get_current_user
 from agentseek_api.core.database import db_manager
@@ -24,9 +23,8 @@
     _protocol_stream_location,
     _stream_response_headers,
     _validate_supported_run_controls,
-    _wait_response_payload,
+    _wait_json_stream_response,
     create_run,
-    wait_run,
 )
 
 router = APIRouter(prefix="/runs", tags=["Stateless Runs"])
@@ -48,7 +46,7 @@ async def create_stateless_run(payload: RunCreateStateless, user: User = Depends
 
 @router.post(
     "/wait",
-    response_class=JSONResponse,
+    response_class=StreamingResponse,
     responses={
         200: {
             "content": {"application/json": {"schema": {}}},
@@ -59,12 +57,12 @@ async def create_stateless_run(payload: RunCreateStateless, user: User = Depends
         }
     },
 )
-async def create_stateless_run_wait(payload: RunCreateStreamingStateless, user: User = Depends(get_current_user)) -> JSONResponse:
+async def create_stateless_run_wait(payload: RunCreateStreamingStateless, user: User = Depends(get_current_user)) -> StreamingResponse:
     _normalize_stream_modes(payload.stream_mode)
     created = await create_stateless_run(payload, 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),
+    return _wait_json_stream_response(
+        run=created,
+        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}",

src/agentseek_api/api/stateless_runs.py

@@ -27,6 +27,7 @@
     load_thread_stream_events,
     parse_last_event_id,
 )
+from agentseek_api.services.sse import iter_with_sse_keepalives, sse_keepalive_comment
 from agentseek_api.services.thread_checkpoint_store import (
     checkpoint_to_payload,
     copy_checkpoints,
@@ -245,6 +246,59 @@ def _empty_thread_state_payload(thread: Thread) -> dict[str, object]:
     }
 
 
+def _checkpoint_namespace(payload: dict[str, object]) -> str:
+    checkpoint = payload.get("checkpoint")
+    if not isinstance(checkpoint, dict):
+        return ""
+    checkpoint_ns = checkpoint.get("checkpoint_ns", "")
+    return str(checkpoint_ns) if checkpoint_ns is not None else ""
+
+
+async def _find_thread_state_payload(
+    *,
+    thread_id: str,
+    user: User,
+    checkpoint_ns: str | None = None,
+) -> dict[str, object] | None:
+    session_factory = db_manager.get_session_factory()
+    async with session_factory() as session:
+        thread = await session.scalar(select(Thread).where(Thread.thread_id == thread_id, Thread.user_id == user.identity))
+        if thread is None:
+            raise HTTPException(status_code=404, detail="Thread not found")
+        runs = (
+            await session.scalars(
+                select(Run).where(Run.thread_id == thread_id, Run.user_id == user.identity).order_by(Run.created_at.asc())
+            )
+        ).all()
+    visible = _visible_checkpoint_payloads(
+        thread,
+        runs,
+        [checkpoint_to_payload(item) for item in await list_checkpoints(thread_id)],
+    )
+    if checkpoint_ns is not None:
+        visible = [payload for payload in visible if _checkpoint_namespace(payload) == checkpoint_ns]
+    if not visible:
+        return None
+    return visible[0]
+
+
+async def _get_thread_state_payload(
+    *,
+    thread_id: str,
+    user: User,
+    checkpoint_ns: str | None = None,
+) -> dict[str, object]:
+    payload = await _find_thread_state_payload(thread_id=thread_id, user=user, checkpoint_ns=checkpoint_ns)
+    if payload is not None:
+        return payload
+    session_factory = db_manager.get_session_factory()
+    async with session_factory() as session:
+        thread = await session.scalar(select(Thread).where(Thread.thread_id == thread_id, Thread.user_id == user.identity))
+        if thread is None:
+            raise HTTPException(status_code=404, detail="Thread not found")
+        return _empty_thread_state_payload(thread)
+
+
 @router.post("", response_model=ThreadRead)
 async def create_thread(payload: ThreadCreate, user: User = Depends(get_current_user)) -> ThreadRead:
     session_factory = db_manager.get_session_factory()
@@ -437,25 +491,7 @@ async def copy_thread(thread_id: str, user: User = Depends(get_current_user)) ->
 
 @router.get("/{thread_id}/state")
 async def get_thread_state(thread_id: str, user: User = Depends(get_current_user)) -> dict[str, object]:
-    session_factory = db_manager.get_session_factory()
-    async with session_factory() as session:
-        thread = await session.scalar(select(Thread).where(Thread.thread_id == thread_id, Thread.user_id == user.identity))
-        if thread is None:
-            raise HTTPException(status_code=404, detail="Thread not found")
-        runs = (
-            await session.scalars(
-                select(Run).where(Run.thread_id == thread_id, Run.user_id == user.identity).order_by(Run.created_at.asc())
-            )
-        ).all()
-    checkpoints = await list_checkpoints(thread_id)
-    visible = _visible_checkpoint_payloads(
-        thread,
-        runs,
-        [checkpoint_to_payload(item) for item in checkpoints],
-    )
-    if not visible:
-        return _empty_thread_state_payload(thread)
-    return visible[0]
+    return await _get_thread_state_payload(thread_id=thread_id, user=user)
 
 
 @router.get("/{thread_id}/history")
@@ -586,27 +622,37 @@ async def _event_iter() -> AsyncIterator[str]:
             yield f"id: {seq}\nevent: {event_name}\ndata: {json.dumps(event)}\n\n"
 
         if _uses_redis_executor():
-            async for event in _iter_persisted_thread_events(
-                thread_id=thread_id,
-                payload=payload,
-                user_id=user.identity,
-                after_seq=current_seq,
-                wait_for_future_runs=True,
+            async for event in iter_with_sse_keepalives(
+                _iter_persisted_thread_events(
+                    thread_id=thread_id,
+                    payload=payload,
+                    user_id=user.identity,
+                    after_seq=current_seq,
+                    wait_for_future_runs=True,
+                )
             ):
+                if event is None:
+                    yield sse_keepalive_comment()
+                    continue
                 seq = int(event.get("seq", 0))
                 current_seq = max(current_seq, seq)
                 event_name = str(event.get("method", "event"))
                 yield f"id: {seq}\nevent: {event_name}\ndata: {json.dumps(event)}\n\n"
             return
 
-        async for event in thread_protocol_broker.stream(
-            thread_id,
-            channels=payload.channels,
-            namespaces=payload.namespaces,
-            depth=payload.depth,
-            since=current_seq,
-            wait_for_future_runs=True,
+        async for event in iter_with_sse_keepalives(
+            thread_protocol_broker.stream(
+                thread_id=thread_id,
+                channels=payload.channels,
+                namespaces=payload.namespaces,
+                depth=payload.depth,
+                since=current_seq,
+                wait_for_future_runs=True,
+            )
         ):
+            if event is None:
+                yield sse_keepalive_comment()
+                continue
             seq = int(event.get("seq", 0))
             current_seq = max(current_seq, seq)
             event_name = str(event.get("method", "event"))
@@ -768,25 +814,35 @@ async def _event_iter() -> AsyncIterator[str]:
             yield f"id: {seq}\nevent: {method}\ndata: {body}\n\n"
 
         if _uses_redis_executor():
-            async for event in _iter_persisted_thread_events(
-                thread_id=thread_id,
-                payload=payload,
-                user_id=user.identity,
-                after_seq=current_seq,
+            async for event in iter_with_sse_keepalives(
+                _iter_persisted_thread_events(
+                    thread_id=thread_id,
+                    payload=payload,
+                    user_id=user.identity,
+                    after_seq=current_seq,
+                )
             ):
+                if event is None:
+                    yield sse_keepalive_comment()
+                    continue
                 seq = int(event.get("seq", 0))
                 method = str(event.get("method", "event"))
                 body = json.dumps(event)
                 yield f"id: {seq}\nevent: {method}\ndata: {body}\n\n"
             return
 
-        async for event in thread_protocol_broker.stream(
-            thread_id,
-            channels=payload.channels,
-            namespaces=payload.namespaces,
-            depth=payload.depth,
-            since=current_seq,
+        async for event in iter_with_sse_keepalives(
+            thread_protocol_broker.stream(
+                thread_id,
+                channels=payload.channels,
+                namespaces=payload.namespaces,
+                depth=payload.depth,
+                since=current_seq,
+            )
         ):
+            if event is None:
+                yield sse_keepalive_comment()
+                continue
             seq = int(event.get("seq", 0))
             method = str(event.get("method", "event"))
             body = json.dumps(event)

src/agentseek_api/api/threads.py

@@ -119,7 +119,7 @@ class ThreadRead(BaseModel):
 
 
 class RunCreateStateful(BaseModel):
-    model_config = ConfigDict(extra="allow")
+    model_config = ConfigDict(extra="forbid")
 
     assistant_id: str
     checkpoint: dict[str, Any] | None = None
@@ -147,7 +147,7 @@ class RunCreateStreamingStateful(RunCreateStateful):
 
 
 class RunCreateStateless(BaseModel):
-    model_config = ConfigDict(extra="allow")
+    model_config = ConfigDict(extra="forbid")
 
     assistant_id: str
     input: Any = None

src/agentseek_api/models/api.py

@@ -860,6 +860,7 @@ async def execute_run(
                 interrupt_id=str(item.get("id", "")),
                 payload=item.get("value"),
                 namespace=interrupt_namespace,
+                run_id=run_id,
             )
 
     output = entry.extract_output(result, payload)

src/agentseek_api/services/run_executor.py

@@ -18,10 +18,12 @@
     persist_run_stream_event,
     persist_thread_stream_event,
 )
+from agentseek_api.services.thread_checkpoint_store import checkpoint_to_payload, get_latest_checkpoint
 from agentseek_api.services.thread_protocol import publish_lifecycle_event, thread_protocol_broker
 
 RUN_EXECUTION_JOB_KIND = "run.execute"
 TERMINAL_RUN_STATUSES = {"success", "error", "interrupted"}
+RUN_CHECKPOINT_ID_METADATA_KEY = "__agentseek_checkpoint_id"
 
 
 @dataclass(slots=True)
@@ -173,6 +175,17 @@ async def execute_run_job(job: RunExecutionJob) -> None:
                 await _persist_thread_snapshot(job.thread_id)
                 await execution_session.refresh(db_run)
                 if not _is_cancelled_run(db_run):
+                    # A missing checkpoint lookup should not turn a successful run into a failed one.
+                    try:
+                        latest_checkpoint = await get_latest_checkpoint(job.thread_id)
+                    except Exception:  # noqa: BLE001
+                        latest_checkpoint = None
+                    if latest_checkpoint is not None:
+                        checkpoint_id = checkpoint_to_payload(latest_checkpoint)["checkpoint"]["checkpoint_id"]
+                        db_run.metadata_json = {
+                            **(db_run.metadata_json or {}),
+                            RUN_CHECKPOINT_ID_METADATA_KEY: checkpoint_id,
+                        }
                     _apply_execution_result(db_run, result)
             except Exception as exc:  # noqa: BLE001
                 await execution_session.refresh(db_run)

src/agentseek_api/services/run_jobs.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from typing import TypeVar
+
+T = TypeVar("T")
+
+DEFAULT_SSE_KEEPALIVE_INTERVAL_SECONDS = 15.0
+
+
+def sse_keepalive_comment() -> str:
+    return ": keepalive\n\n"
+
+
+async def iter_with_sse_keepalives(
+    source: AsyncIterator[T],
+    *,
+    interval_seconds: float | None = None,
+) -> AsyncIterator[T | None]:
+    interval = DEFAULT_SSE_KEEPALIVE_INTERVAL_SECONDS if interval_seconds is None else interval_seconds
+    iterator = source.__aiter__()
+    pending: asyncio.Task[T] | None = None
+    try:
+        while True:
+            if pending is None:
+                pending = asyncio.create_task(anext(iterator))
+            try:
+                item = await asyncio.wait_for(asyncio.shield(pending), timeout=interval)
+            except TimeoutError:
+                yield None
+                continue
+            except StopAsyncIteration:
+                return
+            pending = None
+            yield item
+    finally:
+        if pending is not None and not pending.done():
+            pending.cancel()
+            try:
+                await pending
+            except (asyncio.CancelledError, StopAsyncIteration):
+                pass
+        aclose = getattr(iterator, "aclose", None)
+        if callable(aclose):
+            await aclose()