wutongyuonce 3c7d3303d3
feat(gateway): paginate thread run history (#5283)
* feat(gateway): paginate thread run history (#5282)

GET /api/threads/{thread_id}/runs stays a bare array of the newest 100
runs so LangGraph SDK clients keep working. Add GET /runs/page with a
(created_at, run_id) keyset cursor so callers can walk older history.

* fix(gateway): reject one-sided run history cursors

RunManager.list_by_thread now raises if only one of before_created_at
or before_run_id is set, matching the HTTP 422. Document the per-page
sort cost on the SQL keyset query, and add the missing CHANGELOG
[#5282] link definition.

* fix(gateway): round-trip run page cursors through query strings

Emit next_before_created_at with a Z suffix so '+' is not decoded as a
space. Accept that space, and Z, when parsing. Treat blank cursor
fields as absent and reject a non-ISO before_created_at in RunManager
so a harness caller cannot silently restart at the newest page.

* style(gateway): ruff-format run page cursor files

Collapse the one-sided cursor ValueError and the two before_created_at
asserts so ruff format --check passes at line-length 240.
2026-09-09 14:56:31 +08:00

481 lines
16 KiB
Python

"""Abstract interface for run metadata storage.
RunManager depends on this interface. Implementations:
- MemoryRunStore: in-memory dict (development, tests)
- Future: RunRepository backed by SQLAlchemy ORM
All methods accept an optional user_id for user isolation.
When user_id is None, no user filtering is applied (single-user mode).
"""
from __future__ import annotations
import abc
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any
from deerflow.utils.time import coerce_iso
@dataclass(frozen=True)
class EditReplayVisibility:
hidden_source_run_ids: set[str] = field(default_factory=set)
hidden_attempt_run_ids: set[str] = field(default_factory=set)
@dataclass(frozen=True)
class LeaseRenewal:
"""Result of renewing a run lease.
``cancel_action`` carries a durable cancellation request to the owning
worker without transferring lease ownership.
"""
renewed: bool
cancel_action: str | None = None
@dataclass(frozen=True)
class StatusFinalization:
"""Result of completing a run only if cancellation has not won."""
finalized: bool
cancel_action: str | None = None
class RunIdempotencyConflict(RuntimeError):
"""A run with the requested process-wide idempotency key already exists."""
def __init__(self, existing: dict[str, Any]) -> None:
super().__init__(f"Run idempotency key already belongs to {existing.get('run_id')}")
self.existing = existing
def normalize_run_created_at_iso(value: str) -> str:
"""Make a run timestamp parseable as ISO-8601.
``Z`` becomes ``+00:00``. An unencoded ``+`` in a query string arrives as a
space (``...T00:00:00 00:00``); restore the offset ``+``.
"""
value = value.strip().replace("Z", "+00:00")
if "T" in value and " " in value and "+" not in value.split("T", 1)[1]:
date, _, rest = value.partition("T")
time_part, sep, offset = rest.rpartition(" ")
if sep and offset.replace(":", "").isdigit():
value = f"{date}T{time_part}+{offset}"
return value
def format_run_cursor_created_at(value: str) -> str:
"""UTC keyset cursor using ``Z`` so ``+`` is not decoded as space in query strings."""
dt = datetime.fromisoformat(normalize_run_created_at_iso(value))
if dt.tzinfo is None:
dt = dt.replace(tzinfo=UTC)
else:
dt = dt.astimezone(UTC)
return dt.isoformat().replace("+00:00", "Z")
def parse_run_created_at(value: object) -> datetime:
"""Parse a stored run timestamp into an aware UTC datetime for keyset order."""
iso = coerce_iso(value)
if not iso:
return datetime.min.replace(tzinfo=UTC)
try:
dt = datetime.fromisoformat(normalize_run_created_at_iso(iso))
except ValueError:
return datetime.min.replace(tzinfo=UTC)
if dt.tzinfo is None:
return dt.replace(tzinfo=UTC)
return dt.astimezone(UTC)
def run_sort_key(created_at: object, run_id: str) -> tuple[datetime, str]:
"""Total order for newest-first run listings: ``created_at`` then ``run_id``."""
return (parse_run_created_at(created_at), run_id)
def run_is_before_cursor(
created_at: object,
run_id: str,
*,
before_created_at: str | None,
before_run_id: str | None,
) -> bool:
"""Return True when ``(created_at, run_id)`` is older than the keyset cursor."""
if not before_created_at or not before_run_id:
return True
return run_sort_key(created_at, run_id) < run_sort_key(before_created_at, before_run_id)
class RunStore(abc.ABC):
@abc.abstractmethod
async def put(
self,
run_id: str,
*,
thread_id: str,
assistant_id: str | None = None,
user_id: str | None = None,
model_name: str | None = None,
status: str = "pending",
operation_kind: str = "run",
multitask_strategy: str = "reject",
metadata: dict[str, Any] | None = None,
kwargs: dict[str, Any] | None = None,
error: str | None = None,
stop_reason: str | None = None,
created_at: str | None = None,
owner_worker_id: str | None = None,
lease_expires_at: str | None = None,
idempotency_key: str | None = None,
) -> None:
pass
@abc.abstractmethod
async def get(
self,
run_id: str,
*,
user_id: str | None = None,
) -> dict[str, Any] | None:
pass
@abc.abstractmethod
async def list_by_thread(
self,
thread_id: str,
*,
user_id: str | None = None,
limit: int = 100,
before_created_at: str | None = None,
before_run_id: str | None = None,
) -> list[dict[str, Any]]:
pass
async def list_successful_regenerate_sources(
self,
thread_id: str,
*,
user_id: str | None = None,
) -> set[str]:
"""Return source run IDs superseded by successful regenerations.
Implementations must inspect the complete thread and must not apply the
normal bounded run-list limit.
"""
raise NotImplementedError
async def list_edit_regenerate_runs(
self,
thread_id: str,
*,
user_id: str | None = None,
) -> list[dict[str, Any]]:
"""Return all edit-regenerate attempt runs for one thread, oldest first."""
raise NotImplementedError
async def get_many_by_thread(
self,
thread_id: str,
run_ids: set[str],
*,
user_id: str | None = None,
) -> dict[str, dict[str, Any]]:
"""Batch-load selected runs belonging to one thread."""
raise NotImplementedError
@abc.abstractmethod
async def update_status(
self,
run_id: str,
status: str,
*,
error: str | None = None,
stop_reason: str | None = None,
) -> bool | None:
"""Update a run status.
Returns ``False`` when the store can prove no row was updated. Older or
lightweight stores may return ``None`` when they cannot report rowcount.
"""
pass
@abc.abstractmethod
async def start_run(self, run_id: str) -> bool:
"""Atomically transition a pending run to running.
Returns ``False`` when the row is missing or no longer pending.
"""
pass
@abc.abstractmethod
async def delete(self, run_id: str) -> None:
pass
async def delete_thread_operation(self, run_id: str, *, user_id: str | None) -> None:
"""Release an admitted thread operation for its recorded owner.
The default keeps legacy stores compatible: older implementations only
accepted ``run_id``. User-aware stores should override this method so
cleanup never depends on ambient request context.
"""
await self.delete(run_id)
@abc.abstractmethod
async def update_model_name(
self,
run_id: str,
model_name: str | None,
) -> None:
"""Update the model_name field for an existing run."""
pass
@abc.abstractmethod
async def update_run_completion(
self,
run_id: str,
*,
status: str,
total_input_tokens: int = 0,
total_output_tokens: int = 0,
total_tokens: int = 0,
llm_call_count: int = 0,
lead_agent_tokens: int = 0,
subagent_tokens: int = 0,
middleware_tokens: int = 0,
token_usage_by_model: dict[str, dict[str, int]] | None = None,
message_count: int = 0,
last_ai_message: str | None = None,
first_human_message: str | None = None,
error: str | None = None,
) -> bool | None:
"""Persist final completion fields.
Implementations must not replace a different terminal status. Returns
``False`` when the row is missing or already has a conflicting terminal
outcome.
"""
pass
async def update_run_progress(
self,
run_id: str,
*,
total_input_tokens: int | None = None,
total_output_tokens: int | None = None,
total_tokens: int | None = None,
llm_call_count: int | None = None,
lead_agent_tokens: int | None = None,
subagent_tokens: int | None = None,
middleware_tokens: int | None = None,
token_usage_by_model: dict[str, dict[str, int]] | None = None,
message_count: int | None = None,
last_ai_message: str | None = None,
first_human_message: str | None = None,
) -> None:
"""Persist a best-effort running snapshot without changing run status."""
return None
@abc.abstractmethod
async def list_pending(self, *, before: str | None = None) -> list[dict[str, Any]]:
pass
@abc.abstractmethod
async def list_inflight(self, *, before: str | None = None) -> list[dict[str, Any]]:
"""Return persisted runs that are still ``pending`` or ``running``."""
pass
@abc.abstractmethod
async def aggregate_tokens_by_thread(self, thread_id: str, *, include_active: bool = False) -> dict[str, Any]:
"""Aggregate token usage for completed runs in a thread.
Returns a dict with keys: total_tokens, total_input_tokens,
total_output_tokens, total_runs, by_model (model_name → {tokens, runs}),
by_caller ({lead_agent, subagent, middleware}).
"""
pass
@abc.abstractmethod
async def update_lease(
self,
run_id: str,
*,
owner_worker_id: str,
lease_expires_at: str,
) -> bool:
"""Renew the lease on an active run. Returns ``False`` when no row matched."""
pass
async def renew_lease(
self,
run_id: str,
*,
owner_worker_id: str,
lease_expires_at: str,
) -> LeaseRenewal:
"""Renew ownership and return any durable cancellation request.
The default wraps the legacy ``update_lease`` method and returns no
cancellation action, so third-party stores remain source-compatible
without adding a background read. Stores that support multi-process
cancellation must override this method to renew and observe the
request atomically.
"""
renewed = await self.update_lease(
run_id,
owner_worker_id=owner_worker_id,
lease_expires_at=lease_expires_at,
)
return LeaseRenewal(renewed=renewed)
async def request_cancel(self, run_id: str, *, action: str) -> str | None:
"""Persist the first cancellation action for an active run.
Implementations must update only ``pending`` or ``running`` rows and
return the winning action, or ``None`` when no active row matched.
"""
raise NotImplementedError
async def finalize_if_not_cancelled(
self,
run_id: str,
*,
status: str,
error: str | None = None,
stop_reason: str | None = None,
) -> StatusFinalization:
"""Atomically finalize an active run unless cancellation won.
The compatibility default is safe for stores that do not implement
durable cancellation.
"""
updated = await self.update_status(
run_id,
status,
error=error,
stop_reason=stop_reason,
)
return StatusFinalization(finalized=updated is not False)
@abc.abstractmethod
async def claim_for_takeover(
self,
run_id: str,
*,
grace_seconds: int,
error: str,
stop_reason: str | None = None,
) -> bool:
"""Atomically mark an expired-lease active run as ``error``.
Only rows whose lease has expired past *grace_seconds* (or whose
lease is NULL — pre-ownership data) are updated. The conditional
WHERE closes the race between the caller's stale read of the lease
and a concurrent heartbeat renewal by the owning worker. When
provided, *stop_reason* is persisted in the same atomic update.
Returns ``False`` when:
- the run is no longer ``pending`` / ``running``,
- the lease is still valid (owner heartbeat is alive), or
- the row doesn't exist.
"""
pass
@abc.abstractmethod
async def list_inflight_with_expired_lease(
self,
*,
before: str | None = None,
grace_seconds: int = 10,
) -> list[dict[str, Any]]:
"""Return active runs whose lease has expired (or is NULL for pre-ownership rows)."""
pass
async def create_thread_operation_atomic(
self,
run_id: str,
*,
thread_id: str,
owner_worker_id: str,
lease_expires_at: str | None,
operation_kind: str = "run",
multitask_strategy: str = "reject",
assistant_id: str | None = None,
user_id: str | None = None,
model_name: str | None = None,
metadata: dict[str, Any] | None = None,
kwargs: dict[str, Any] | None = None,
created_at: str | None = None,
grace_seconds: int = 10,
idempotency_key: str | None = None,
) -> tuple[dict[str, Any], list[dict[str, Any]]]:
"""Atomically create an active thread operation with cross-process uniqueness.
The default implementation preserves compatibility with stores that
still implement the former ``create_run_atomic`` interface. Legacy
stores support only normal run rows; internal operation kinds require
an implementation of this method.
Returns ``(new_run_dict, claimed_run_dicts)``.
Raises ``IntegrityError`` on conflict for ``reject`` strategy.
"""
legacy_impl = type(self).create_run_atomic
if legacy_impl is RunStore.create_run_atomic:
raise NotImplementedError("RunStore must implement create_thread_operation_atomic() or create_run_atomic()")
if operation_kind != "run":
raise NotImplementedError("Legacy RunStore.create_run_atomic() cannot create non-run thread operations")
if idempotency_key is not None:
raise NotImplementedError("Legacy RunStore.create_run_atomic() cannot guarantee idempotent admission")
return await self.create_run_atomic(
run_id,
thread_id=thread_id,
owner_worker_id=owner_worker_id,
lease_expires_at=lease_expires_at,
multitask_strategy=multitask_strategy,
assistant_id=assistant_id,
user_id=user_id,
model_name=model_name,
metadata=metadata,
kwargs=kwargs,
created_at=created_at,
grace_seconds=grace_seconds,
)
async def create_run_atomic(
self,
run_id: str,
*,
thread_id: str,
owner_worker_id: str,
lease_expires_at: str | None,
multitask_strategy: str = "reject",
assistant_id: str | None = None,
user_id: str | None = None,
model_name: str | None = None,
metadata: dict[str, Any] | None = None,
kwargs: dict[str, Any] | None = None,
created_at: str | None = None,
grace_seconds: int = 10,
) -> tuple[dict[str, Any], list[dict[str, Any]]]:
"""Deprecated compatibility alias for normal-run admission."""
operation_impl = type(self).create_thread_operation_atomic
if operation_impl is RunStore.create_thread_operation_atomic:
raise NotImplementedError("RunStore must implement create_thread_operation_atomic() or create_run_atomic()")
return await self.create_thread_operation_atomic(
run_id,
thread_id=thread_id,
owner_worker_id=owner_worker_id,
lease_expires_at=lease_expires_at,
operation_kind="run",
multitask_strategy=multitask_strategy,
assistant_id=assistant_id,
user_id=user_id,
model_name=model_name,
metadata=metadata,
kwargs=kwargs,
created_at=created_at,
grace_seconds=grace_seconds,
)