mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-19 19:16:17 +00:00
16f154f32b
* fix(gateway): preserve owner isolation when thread metadata is missing Follow-up to the #5448 review P1 (post-merge finding): owner_check=True also authorizes threads whose meta row is missing (legacy compatibility) or NULL-owner (shared/pre-auth data). _run_scope_user_id returned None for every trusted internal caller, which dropped the only remaining per-user filter on those threads and let an internal caller acting for owner A read owner B's persisted runs. _run_scope_user_id now takes the thread_id and consults the thread meta store: when an existing meta row establishes ownership, the authorized thread's runs are still read unfiltered (merged #5448 semantics, including owner-header-less internal callers); when the meta row is missing or NULL-owner, the filter falls back to the acting owner's raw stamp (the exact value start_run writes) — or the synthetic "default" identity without an owner header — so cross-user runs stay hidden. Isolation coverage uses the real MemoryThreadMetaStore with no metadata row (and a NULL-owner row) plus another user's persisted run: /runs and /runs/page must be empty and /runs/{run_id} must 404 for internal callers, while an established-ownership thread keeps the unfiltered read. * fix(gateway): gate run-scoped sub-resource reads for internal callers Review follow-up on #5484: the P1 owner-isolation class remained reachable through run-scoped sibling reads that apply no per-user filter at all — /runs/{run_id}/messages, /events, /join, /stream and /workspace-changes query by (thread_id, run_id) directly, so on missing/NULL-owner threads an internal caller acting for owner A could still read owner B's run content by id (verified 200 at the previous head). - Extract _thread_ownership_established (shared meta-row check) and add _require_run_visible_to_scope: for internal callers on threads without established ownership, the run's own user_id stamp must match the acting owner's raw value (or the legacy "default" stamp) or the read 404s. Established-ownership threads and every non-internal caller keep their existing thread-scoped semantics. - Wire the gate into join, stream, messages, events and workspace-changes; reword the now-stale messages comment to track the new scoping semantics. Regression tests: sub-resource reads 404 for a mismatched internal owner while the matching owner reads them normally, and the owner-less fallback branch (synthetic "default" filter on missing-meta threads) is pinned. Red confirmed against the pre-gate head. * fix(gateway): gate cancel and artifact archive for internal callers Review follow-up on #5484 round 2: POST /cancel resolved runs unscoped (require_existing=True only closes the missing-meta case — NULL-owner meta rows still pass), so an internal caller acting for a different owner could interrupt another owner's active run on a shared thread while /join and /stream were already gated. The archive manifest and download pair likewise leaked the other owner's delivered-file count and a 200-vs-409 delivery oracle on NULL-owner threads (missing-meta threads were already denied by require_existing=True). All three routes now call _require_run_visible_to_scope; its docstring records the extended coverage. NULL-owner-thread regression tests pin: a mismatched internal owner gets 404 from cancel, manifest and archive download, while the acting owner reaches the real conflict path (409 on a terminal run) and reads the manifest (file_count 2). * fix(gateway): tolerate state-less request stand-ins in the scope helpers The new owner-isolation gate and _run_scope_user_id read request.state directly, which crashed the FakeRequest-based unit suites for the run events, workspace-changes and scope endpoints (backend-unit-tests shards 1/2/4 on #5484). Read the state object defensively first: a request without state is simply not an internal caller, so those paths keep their pre-gate semantics. * fix(gateway): scope the thread token-usage aggregate by owner Review follow-up on #5484 round 4: GET /{thread_id}/token-usage called aggregate_tokens_by_thread(thread_id) with no user filter at all, so on missing/NULL-owner threads an internal caller acting for owner A read owner B's spend, model names, run count and (with include_active=true) live activity; the NULL-owner variant reached browser sessions too. build_context_usage's latest-model lookup was unfiltered as well. aggregate_tokens_by_thread gains an optional user_id (mirroring list_by_thread: explicit None = unfiltered, AUTO resolves the contextvar) in the memory store, the SQL repository and the store base; build_context_usage/_resolve_thread_model_name thread the scope through the latest-run lookup; the token-usage endpoint passes _run_scope_user_id's value. Established-ownership threads aggregate unfiltered as before; shared/missing-meta threads narrow to the acting identity. Stale helper-test comment reworded after the #5482 merge adaptation. * test(gateway): pin the unfiltered aggregate on established-ownership threads Review follow-up on #5484 round 5: the established-ownership branch of the token-usage scoping (store receives user_id=None) was the only unpinned half of the contract — the round-4 call-assertions never set app.state.thread_store, so their None came from the user-less stand-in path. test_token_usage_unfiltered_on_established_ownership_for_ internal_callers seeds an established meta row plus runs stamped by two different identities and asserts the totals fold (166 = 111 + 55); together with the isolation tests it now catches both failure modes (always-stamp narrowing and always-None leak).
878 lines
36 KiB
Python
878 lines
36 KiB
Python
"""SQLAlchemy-backed RunStore implementation.
|
|
|
|
Each method acquires and releases its own short-lived session.
|
|
Run status updates happen from background workers that may live
|
|
minutes -- we don't hold connections across long execution.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import Any
|
|
|
|
from sqlalchemy import and_, case, or_, select, update
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from deerflow.persistence.run.model import RunChangeClockRow, RunRow
|
|
from deerflow.runtime.runs.store.base import (
|
|
LeaseRenewal,
|
|
RunIdempotencyConflict,
|
|
RunStore,
|
|
StatusFinalization,
|
|
normalize_run_created_at_iso,
|
|
)
|
|
from deerflow.runtime.user_context import AUTO, _AutoSentinel, resolve_user_id
|
|
from deerflow.utils.time import coerce_iso
|
|
|
|
|
|
def _lease_expired_or_null(lease_col, cutoff: datetime):
|
|
"""SQLAlchemy filter: True when the lease is NULL or has expired past *cutoff*."""
|
|
return or_(lease_col.is_(None), lease_col < cutoff)
|
|
|
|
|
|
class RunRepository(RunStore):
|
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
|
self._sf = session_factory
|
|
|
|
@staticmethod
|
|
async def _next_change_seq(session: AsyncSession) -> int:
|
|
dialect = session.bind.dialect.name
|
|
if dialect == "postgresql":
|
|
from sqlalchemy.dialects.postgresql import insert
|
|
elif dialect == "sqlite":
|
|
from sqlalchemy.dialects.sqlite import insert
|
|
else: # pragma: no cover - configured databases are SQLite/Postgres
|
|
raise RuntimeError(f"unsupported run-change clock dialect: {dialect}")
|
|
await session.execute(insert(RunChangeClockRow).values(id=1, value=0).on_conflict_do_nothing(index_elements=[RunChangeClockRow.id]))
|
|
value = await session.scalar(update(RunChangeClockRow).where(RunChangeClockRow.id == 1).values(value=RunChangeClockRow.value + 1).returning(RunChangeClockRow.value))
|
|
if value is None:
|
|
raise RuntimeError("run-change clock did not return a position")
|
|
return int(value)
|
|
|
|
@staticmethod
|
|
def _normalize_model_name(model_name: str | None) -> str | None:
|
|
"""Normalize model_name for storage: strip whitespace, truncate to 128 chars."""
|
|
if model_name is None:
|
|
return None
|
|
if not isinstance(model_name, str):
|
|
model_name = str(model_name)
|
|
normalized = model_name.strip()
|
|
if len(normalized) > 128:
|
|
normalized = normalized[:128]
|
|
return normalized
|
|
|
|
@staticmethod
|
|
def _safe_json(obj: Any) -> Any:
|
|
"""Ensure obj is JSON-serializable. Falls back to model_dump() or str()."""
|
|
if obj is None:
|
|
return None
|
|
if isinstance(obj, (str, int, float, bool)):
|
|
return obj
|
|
if isinstance(obj, dict):
|
|
return {k: RunRepository._safe_json(v) for k, v in obj.items()}
|
|
if isinstance(obj, (list, tuple)):
|
|
return [RunRepository._safe_json(v) for v in obj]
|
|
if hasattr(obj, "model_dump"):
|
|
try:
|
|
return obj.model_dump()
|
|
except Exception:
|
|
pass
|
|
if hasattr(obj, "dict"):
|
|
try:
|
|
return obj.dict()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
json.dumps(obj)
|
|
return obj
|
|
except (TypeError, ValueError):
|
|
return str(obj)
|
|
|
|
@staticmethod
|
|
def _row_to_dict(row: RunRow) -> dict[str, Any]:
|
|
d = row.to_dict()
|
|
# Remap JSON columns to match RunStore interface
|
|
d["metadata"] = d.pop("metadata_json", {})
|
|
d["kwargs"] = d.pop("kwargs_json", {})
|
|
# Convert datetime to ISO string for consistency with MemoryRunStore.
|
|
# SQLite drops tzinfo on read despite ``DateTime(timezone=True)`` —
|
|
# ``coerce_iso`` normalizes naive datetimes as UTC.
|
|
for key in ("created_at", "updated_at", "lease_expires_at", "cancel_requested_at"):
|
|
val = d.get(key)
|
|
if isinstance(val, datetime):
|
|
d[key] = coerce_iso(val)
|
|
return d
|
|
|
|
async def put(
|
|
self,
|
|
run_id,
|
|
*,
|
|
thread_id,
|
|
assistant_id=None,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
model_name: str | None = None,
|
|
status="pending",
|
|
operation_kind: str = "run",
|
|
multitask_strategy="reject",
|
|
metadata=None,
|
|
kwargs=None,
|
|
error=None,
|
|
stop_reason: str | None = None,
|
|
created_at=None,
|
|
follow_up_to_run_id=None,
|
|
owner_worker_id: str | None = None,
|
|
lease_expires_at: str | None = None,
|
|
idempotency_key: str | None = None,
|
|
):
|
|
"""Insert or update a run row.
|
|
|
|
``RunManager`` retries ``put`` after transient SQLite failures. Making
|
|
this operation idempotent prevents a successful-but-unacknowledged first
|
|
commit from turning the retry into a primary-key failure.
|
|
"""
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.put")
|
|
now = datetime.now(UTC)
|
|
created = datetime.fromisoformat(created_at) if created_at else now
|
|
lease_dt = datetime.fromisoformat(lease_expires_at) if lease_expires_at else None
|
|
values = {
|
|
"thread_id": thread_id,
|
|
"assistant_id": assistant_id,
|
|
"user_id": resolved_user_id,
|
|
"model_name": self._normalize_model_name(model_name),
|
|
"status": status,
|
|
"operation_kind": operation_kind,
|
|
"multitask_strategy": multitask_strategy,
|
|
"metadata_json": self._safe_json(metadata) or {},
|
|
"kwargs_json": self._safe_json(kwargs) or {},
|
|
"error": error,
|
|
"stop_reason": stop_reason,
|
|
"follow_up_to_run_id": follow_up_to_run_id,
|
|
"owner_worker_id": owner_worker_id,
|
|
"lease_expires_at": lease_dt,
|
|
"idempotency_key": idempotency_key,
|
|
"updated_at": now,
|
|
}
|
|
async with self._sf() as session:
|
|
values["change_seq"] = await self._next_change_seq(session)
|
|
row = await session.get(RunRow, run_id)
|
|
if row is None:
|
|
session.add(RunRow(run_id=run_id, created_at=created, **values))
|
|
else:
|
|
for key, value in values.items():
|
|
setattr(row, key, value)
|
|
await session.commit()
|
|
|
|
async def get(
|
|
self,
|
|
run_id,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
):
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.get")
|
|
async with self._sf() as session:
|
|
row = await session.get(RunRow, run_id)
|
|
if row is None:
|
|
return None
|
|
if resolved_user_id is not None and row.user_id != resolved_user_id:
|
|
return None
|
|
return self._row_to_dict(row)
|
|
|
|
async def list_changed(
|
|
self,
|
|
*,
|
|
after_change_seq: int,
|
|
after_run_id: str,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
limit: int = 100,
|
|
) -> list[dict[str, Any]]:
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.list_changed")
|
|
stmt = select(RunRow).where(
|
|
RunRow.operation_kind == "run",
|
|
or_(
|
|
RunRow.change_seq > after_change_seq,
|
|
and_(RunRow.change_seq == after_change_seq, RunRow.run_id > after_run_id),
|
|
),
|
|
)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RunRow.user_id == resolved_user_id)
|
|
stmt = stmt.order_by(RunRow.change_seq.asc(), RunRow.run_id.asc()).limit(limit)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(row) for row in result.scalars()]
|
|
|
|
async def list_by_thread(
|
|
self,
|
|
thread_id,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
limit=100,
|
|
before_created_at: str | None = None,
|
|
before_run_id: str | None = None,
|
|
):
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.list_by_thread")
|
|
stmt = select(RunRow).where(RunRow.thread_id == thread_id, RunRow.operation_kind == "run")
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RunRow.user_id == resolved_user_id)
|
|
if before_created_at and before_run_id:
|
|
cursor_dt = datetime.fromisoformat(normalize_run_created_at_iso(before_created_at))
|
|
if cursor_dt.tzinfo is None:
|
|
cursor_dt = cursor_dt.replace(tzinfo=UTC)
|
|
else:
|
|
cursor_dt = cursor_dt.astimezone(UTC)
|
|
stmt = stmt.where(
|
|
or_(
|
|
RunRow.created_at < cursor_dt,
|
|
and_(RunRow.created_at == cursor_dt, RunRow.run_id < before_run_id),
|
|
)
|
|
)
|
|
# Keyset pages filter on (created_at, run_id) after thread_id. Existing
|
|
# indexes are (thread_id) and (thread_id, status), so each page still
|
|
# sorts matching rows. A covering (thread_id, created_at, run_id) index
|
|
# is a follow-up if deep paging shows up in profiles.
|
|
stmt = stmt.order_by(RunRow.created_at.desc(), RunRow.run_id.desc()).limit(limit)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(r) for r in result.scalars()]
|
|
|
|
async def list_successful_regenerate_sources(
|
|
self,
|
|
thread_id,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
):
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.list_successful_regenerate_sources")
|
|
source = RunRow.metadata_json["regenerate_from_run_id"].as_string()
|
|
stmt = select(source).where(
|
|
RunRow.thread_id == thread_id,
|
|
RunRow.operation_kind == "run",
|
|
RunRow.status == "success",
|
|
source.is_not(None),
|
|
source != "",
|
|
)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RunRow.user_id == resolved_user_id)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return {value for value in result.scalars() if isinstance(value, str) and value}
|
|
|
|
async def list_edit_regenerate_runs(
|
|
self,
|
|
thread_id,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
):
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.list_edit_regenerate_runs")
|
|
replay_kind = RunRow.metadata_json["replay_kind"].as_string()
|
|
source = RunRow.metadata_json["regenerate_from_run_id"].as_string()
|
|
stmt = select(RunRow).where(
|
|
RunRow.thread_id == thread_id,
|
|
replay_kind == "edit",
|
|
source.is_not(None),
|
|
source != "",
|
|
)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RunRow.user_id == resolved_user_id)
|
|
stmt = stmt.order_by(RunRow.created_at.asc())
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(row) for row in result.scalars()]
|
|
|
|
async def get_many_by_thread(
|
|
self,
|
|
thread_id,
|
|
run_ids,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
):
|
|
if not run_ids:
|
|
return {}
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.get_many_by_thread")
|
|
stmt = select(RunRow).where(RunRow.thread_id == thread_id, RunRow.operation_kind == "run", RunRow.run_id.in_(run_ids))
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RunRow.user_id == resolved_user_id)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return {row.run_id: self._row_to_dict(row) for row in result.scalars()}
|
|
|
|
async def update_status(self, run_id, status, *, error=None, stop_reason=None) -> bool:
|
|
values: dict[str, Any] = {"status": status, "updated_at": datetime.now(UTC)}
|
|
if error is not None:
|
|
values["error"] = error
|
|
if stop_reason is not None:
|
|
values["stop_reason"] = stop_reason
|
|
# Guard: only transition rows that are still active. ``interrupted`` is
|
|
# included because the rollback path goes ``running → interrupted``
|
|
# (cancel acknowledged) then ``interrupted → error`` (task finalize).
|
|
# ``error`` and ``success`` remain locked so a peer's takeover (or a
|
|
# completed run) cannot be overwritten by a late writer.
|
|
async with self._sf() as session:
|
|
values["change_seq"] = await self._next_change_seq(session)
|
|
result = await session.execute(update(RunRow).where(RunRow.run_id == run_id, RunRow.status.in_(("pending", "running", "interrupted"))).values(**values))
|
|
await session.commit()
|
|
return result.rowcount != 0
|
|
|
|
async def start_run(self, run_id: str) -> bool:
|
|
"""Start only a still-pending run; cancelled rows must not be resurrected."""
|
|
async with self._sf() as session:
|
|
change_seq = await self._next_change_seq(session)
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
RunRow.status == "pending",
|
|
)
|
|
.values(status="running", updated_at=datetime.now(UTC), change_seq=change_seq)
|
|
)
|
|
await session.commit()
|
|
return result.rowcount != 0
|
|
|
|
async def update_model_name(self, run_id, model_name):
|
|
async with self._sf() as session:
|
|
change_seq = await self._next_change_seq(session)
|
|
await session.execute(
|
|
update(RunRow)
|
|
.where(RunRow.run_id == run_id)
|
|
.values(
|
|
model_name=self._normalize_model_name(model_name),
|
|
updated_at=datetime.now(UTC),
|
|
change_seq=change_seq,
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
async def delete(
|
|
self,
|
|
run_id,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
):
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.delete")
|
|
async with self._sf() as session:
|
|
row = await session.get(RunRow, run_id)
|
|
if row is None:
|
|
return
|
|
if resolved_user_id is not None and row.user_id != resolved_user_id:
|
|
return
|
|
await session.delete(row)
|
|
await session.commit()
|
|
|
|
async def delete_thread_operation(self, run_id: str, *, user_id: str | None) -> None:
|
|
"""Release a reservation using its captured owner, not request context."""
|
|
await self.delete(run_id, user_id=user_id)
|
|
|
|
async def list_pending(self, *, before=None):
|
|
if before is None:
|
|
before_dt = datetime.now(UTC)
|
|
elif isinstance(before, datetime):
|
|
before_dt = before
|
|
else:
|
|
before_dt = datetime.fromisoformat(before)
|
|
stmt = select(RunRow).where(RunRow.operation_kind == "run", RunRow.status == "pending", RunRow.created_at <= before_dt).order_by(RunRow.created_at.asc())
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(r) for r in result.scalars()]
|
|
|
|
async def list_inflight(self, *, before=None):
|
|
"""Return persisted active runs for startup recovery."""
|
|
if before is None:
|
|
before_dt = datetime.now(UTC)
|
|
elif isinstance(before, datetime):
|
|
before_dt = before
|
|
else:
|
|
before_dt = datetime.fromisoformat(before)
|
|
stmt = (
|
|
select(RunRow)
|
|
.where(
|
|
RunRow.status.in_(("pending", "running")),
|
|
RunRow.created_at <= before_dt,
|
|
)
|
|
.order_by(RunRow.created_at.asc())
|
|
)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(r) for r in result.scalars()]
|
|
|
|
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:
|
|
"""Update status + token usage + convenience fields on run completion.
|
|
|
|
Returns ``False`` when the row is missing or already has a conflicting
|
|
terminal outcome.
|
|
"""
|
|
values: dict[str, Any] = {
|
|
"status": status,
|
|
"total_input_tokens": total_input_tokens,
|
|
"total_output_tokens": total_output_tokens,
|
|
"total_tokens": total_tokens,
|
|
"llm_call_count": llm_call_count,
|
|
"lead_agent_tokens": lead_agent_tokens,
|
|
"subagent_tokens": subagent_tokens,
|
|
"middleware_tokens": middleware_tokens,
|
|
"token_usage_by_model": self._safe_json(token_usage_by_model) or {},
|
|
"message_count": message_count,
|
|
"updated_at": datetime.now(UTC),
|
|
}
|
|
if last_ai_message is not None:
|
|
values["last_ai_message"] = last_ai_message[:2000]
|
|
if first_human_message is not None:
|
|
values["first_human_message"] = first_human_message[:2000]
|
|
if error is not None:
|
|
values["error"] = error
|
|
allowed_sources = ["pending", "running"]
|
|
if status not in allowed_sources:
|
|
allowed_sources.append(status)
|
|
if status == "error" and "interrupted" not in allowed_sources:
|
|
allowed_sources.append("interrupted")
|
|
async with self._sf() as session:
|
|
values["change_seq"] = await self._next_change_seq(session)
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
RunRow.status.in_(tuple(allowed_sources)),
|
|
)
|
|
.values(**values)
|
|
)
|
|
await session.commit()
|
|
return result.rowcount != 0
|
|
|
|
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:
|
|
"""Update token usage + convenience fields while a run is still active."""
|
|
values: dict[str, Any] = {"updated_at": datetime.now(UTC)}
|
|
optional_counters = {
|
|
"total_input_tokens": total_input_tokens,
|
|
"total_output_tokens": total_output_tokens,
|
|
"total_tokens": total_tokens,
|
|
"llm_call_count": llm_call_count,
|
|
"lead_agent_tokens": lead_agent_tokens,
|
|
"subagent_tokens": subagent_tokens,
|
|
"middleware_tokens": middleware_tokens,
|
|
"message_count": message_count,
|
|
}
|
|
for key, value in optional_counters.items():
|
|
if value is not None:
|
|
values[key] = value
|
|
if token_usage_by_model is not None:
|
|
values["token_usage_by_model"] = self._safe_json(token_usage_by_model) or {}
|
|
if last_ai_message is not None:
|
|
values["last_ai_message"] = last_ai_message[:2000]
|
|
if first_human_message is not None:
|
|
values["first_human_message"] = first_human_message[:2000]
|
|
async with self._sf() as session:
|
|
await session.execute(update(RunRow).where(RunRow.run_id == run_id, RunRow.status == "running").values(**values))
|
|
await session.commit()
|
|
|
|
async def aggregate_tokens_by_thread(
|
|
self,
|
|
thread_id: str,
|
|
*,
|
|
include_active: bool = False,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> dict[str, Any]:
|
|
"""Aggregate token usage for a thread.
|
|
|
|
``by_model`` is reduced in Python from each row's ``token_usage_by_model``
|
|
JSON column so subagent / middleware tokens land on the model that
|
|
actually produced them (issue #3645). Rows written before that column
|
|
existed fall back to ``RunRow.model_name`` + ``RunRow.total_tokens``,
|
|
preserving the legacy lead-only behavior instead of dropping the data.
|
|
|
|
Headline totals (``total_tokens``, ``total_input_tokens``,
|
|
``total_output_tokens``) and the ``by_caller`` bucket are summed from
|
|
their own columns and are therefore unaffected by the JSON column being
|
|
empty.
|
|
"""
|
|
statuses = ("success", "error", "running") if include_active else ("success", "error")
|
|
_completed = RunRow.status.in_(statuses)
|
|
_thread = RunRow.thread_id == thread_id
|
|
_run_operation = RunRow.operation_kind == "run"
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.aggregate_tokens_by_thread")
|
|
|
|
stmt = select(
|
|
RunRow.model_name,
|
|
RunRow.total_tokens,
|
|
RunRow.total_input_tokens,
|
|
RunRow.total_output_tokens,
|
|
RunRow.lead_agent_tokens,
|
|
RunRow.subagent_tokens,
|
|
RunRow.middleware_tokens,
|
|
RunRow.token_usage_by_model,
|
|
).where(_thread, _run_operation, _completed)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RunRow.user_id == resolved_user_id)
|
|
|
|
async with self._sf() as session:
|
|
rows = (await session.execute(stmt)).all()
|
|
|
|
total_tokens = total_input = total_output = total_runs = 0
|
|
lead_agent = subagent = middleware = 0
|
|
by_model: dict[str, dict] = {}
|
|
for r in rows:
|
|
total_runs += 1
|
|
total_tokens += r.total_tokens
|
|
total_input += r.total_input_tokens
|
|
total_output += r.total_output_tokens
|
|
lead_agent += r.lead_agent_tokens
|
|
subagent += r.subagent_tokens
|
|
middleware += r.middleware_tokens
|
|
|
|
# ``or {}`` covers rows written before ``token_usage_by_model``
|
|
# existed (the column is NULL on a manual ALTER ADD COLUMN without
|
|
# backfill); fresh rows always carry the journal-produced dict.
|
|
usage_by_model = r.token_usage_by_model or {}
|
|
if usage_by_model:
|
|
for model, usage in usage_by_model.items():
|
|
entry = by_model.setdefault(model, {"tokens": 0, "runs": 0})
|
|
entry["tokens"] += usage.get("total_tokens", 0)
|
|
entry["runs"] += 1
|
|
else:
|
|
model = r.model_name or "unknown"
|
|
entry = by_model.setdefault(model, {"tokens": 0, "runs": 0})
|
|
entry["tokens"] += r.total_tokens
|
|
entry["runs"] += 1
|
|
|
|
return {
|
|
"total_tokens": total_tokens,
|
|
"total_input_tokens": total_input,
|
|
"total_output_tokens": total_output,
|
|
"total_runs": total_runs,
|
|
"by_model": by_model,
|
|
"by_caller": {
|
|
"lead_agent": lead_agent,
|
|
"subagent": subagent,
|
|
"middleware": middleware,
|
|
},
|
|
}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Multi-worker run ownership methods
|
|
# ------------------------------------------------------------------
|
|
|
|
async def update_lease(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
owner_worker_id: str,
|
|
lease_expires_at: str,
|
|
) -> bool:
|
|
lease_dt = datetime.fromisoformat(lease_expires_at)
|
|
values: dict[str, Any] = {
|
|
"owner_worker_id": owner_worker_id,
|
|
"lease_expires_at": lease_dt,
|
|
"updated_at": datetime.now(UTC),
|
|
}
|
|
async with self._sf() as session:
|
|
result = await session.execute(update(RunRow).where(RunRow.run_id == run_id, RunRow.owner_worker_id == owner_worker_id, RunRow.status.in_(("pending", "running"))).values(**values))
|
|
await session.commit()
|
|
return result.rowcount != 0
|
|
|
|
async def renew_lease(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
owner_worker_id: str,
|
|
lease_expires_at: str,
|
|
) -> LeaseRenewal:
|
|
"""Renew the owner lease and read cancellation intent atomically."""
|
|
lease_dt = datetime.fromisoformat(lease_expires_at)
|
|
async with self._sf() as session:
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
RunRow.owner_worker_id == owner_worker_id,
|
|
RunRow.status.in_(("pending", "running")),
|
|
)
|
|
.values(
|
|
lease_expires_at=lease_dt,
|
|
updated_at=datetime.now(UTC),
|
|
)
|
|
.returning(RunRow.run_id, RunRow.cancel_action)
|
|
)
|
|
row = result.first()
|
|
await session.commit()
|
|
if row is None:
|
|
return LeaseRenewal(renewed=False)
|
|
return LeaseRenewal(renewed=True, cancel_action=row.cancel_action)
|
|
|
|
async def request_cancel(self, run_id: str, *, action: str) -> str | None:
|
|
"""Atomically persist the first cancellation action on an active run."""
|
|
if action not in ("interrupt", "rollback"):
|
|
raise ValueError(f"Unsupported cancellation action: {action}")
|
|
now = datetime.now(UTC)
|
|
async with self._sf() as session:
|
|
change_seq = await self._next_change_seq(session)
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
RunRow.status.in_(("pending", "running")),
|
|
)
|
|
.values(
|
|
cancel_action=case(
|
|
(RunRow.cancel_action.is_(None), action),
|
|
else_=RunRow.cancel_action,
|
|
),
|
|
cancel_requested_at=case(
|
|
(RunRow.cancel_requested_at.is_(None), now),
|
|
else_=RunRow.cancel_requested_at,
|
|
),
|
|
updated_at=now,
|
|
change_seq=change_seq,
|
|
)
|
|
.returning(RunRow.cancel_action)
|
|
)
|
|
row = result.first()
|
|
await session.commit()
|
|
return row.cancel_action if row is not None else None
|
|
|
|
async def finalize_if_not_cancelled(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
status: str,
|
|
error: str | None = None,
|
|
stop_reason: str | None = None,
|
|
) -> StatusFinalization:
|
|
"""Atomically let completion win only before cancellation."""
|
|
values: dict[str, Any] = {
|
|
"status": status,
|
|
"updated_at": datetime.now(UTC),
|
|
}
|
|
if error is not None:
|
|
values["error"] = error
|
|
if stop_reason is not None:
|
|
values["stop_reason"] = stop_reason
|
|
|
|
async with self._sf() as session:
|
|
values["change_seq"] = await self._next_change_seq(session)
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
RunRow.status.in_(("pending", "running")),
|
|
RunRow.cancel_action.is_(None),
|
|
)
|
|
.values(**values)
|
|
.returning(RunRow.run_id)
|
|
)
|
|
if result.first() is not None:
|
|
await session.commit()
|
|
return StatusFinalization(finalized=True)
|
|
|
|
current = await session.execute(select(RunRow.cancel_action).where(RunRow.run_id == run_id))
|
|
cancel_action = current.scalar_one_or_none()
|
|
await session.commit()
|
|
return StatusFinalization(
|
|
finalized=False,
|
|
cancel_action=cancel_action,
|
|
)
|
|
|
|
async def claim_for_takeover(
|
|
self,
|
|
run_id: str,
|
|
*,
|
|
grace_seconds: int,
|
|
error: str,
|
|
stop_reason: str | None = None,
|
|
) -> bool:
|
|
cutoff = datetime.now(UTC) - timedelta(seconds=grace_seconds)
|
|
values: dict[str, Any] = {
|
|
"status": "error",
|
|
"error": error,
|
|
"updated_at": datetime.now(UTC),
|
|
}
|
|
if stop_reason is not None:
|
|
values["stop_reason"] = stop_reason
|
|
async with self._sf() as session:
|
|
values["change_seq"] = await self._next_change_seq(session)
|
|
result = await session.execute(
|
|
update(RunRow)
|
|
.where(
|
|
RunRow.run_id == run_id,
|
|
RunRow.status.in_(("pending", "running")),
|
|
_lease_expired_or_null(RunRow.lease_expires_at, cutoff),
|
|
)
|
|
.values(**values)
|
|
)
|
|
await session.commit()
|
|
return result.rowcount != 0
|
|
|
|
async def list_inflight_with_expired_lease(
|
|
self,
|
|
*,
|
|
before: str | None = None,
|
|
grace_seconds: int = 10,
|
|
) -> list[dict[str, Any]]:
|
|
if before is None:
|
|
before_dt = datetime.now(UTC)
|
|
elif isinstance(before, datetime):
|
|
before_dt = before
|
|
else:
|
|
before_dt = datetime.fromisoformat(before)
|
|
cutoff = datetime.now(UTC) - timedelta(seconds=grace_seconds)
|
|
stmt = (
|
|
select(RunRow)
|
|
.where(
|
|
RunRow.status.in_(("pending", "running")),
|
|
RunRow.created_at <= before_dt,
|
|
_lease_expired_or_null(RunRow.lease_expires_at, cutoff),
|
|
)
|
|
.order_by(RunRow.created_at.asc())
|
|
)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(r) for r in result.scalars()]
|
|
|
|
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 a run with cross-process thread-uniqueness.
|
|
|
|
- For ``reject``: INSERT, let the partial unique index enforce
|
|
single-active-run. Returns ``(row_dict, [])`` on success, raises
|
|
``IntegrityError`` on conflict.
|
|
- For ``interrupt`` / ``rollback``: SELECT FOR UPDATE inflight
|
|
rows for the thread, cancel them (unless their lease is still valid),
|
|
then INSERT the new row — all in one transaction. Returns
|
|
``(row_dict, claimed_row_dicts)``.
|
|
|
|
Returns:
|
|
Tuple of ``(new_run_dict, claimed_run_dicts)``.
|
|
"""
|
|
from deerflow.runtime.runs.manager import ConflictError
|
|
|
|
resolved_user_id = resolve_user_id(user_id or AUTO, method_name="RunRepository.create_thread_operation_atomic")
|
|
now = datetime.now(UTC)
|
|
created = datetime.fromisoformat(created_at) if created_at else now
|
|
lease_dt = datetime.fromisoformat(lease_expires_at) if lease_expires_at else None
|
|
cutoff = now - timedelta(seconds=grace_seconds)
|
|
|
|
values = {
|
|
"thread_id": thread_id,
|
|
"assistant_id": assistant_id,
|
|
"user_id": resolved_user_id,
|
|
"model_name": self._normalize_model_name(model_name),
|
|
"status": "pending",
|
|
"operation_kind": operation_kind,
|
|
"multitask_strategy": multitask_strategy,
|
|
"metadata_json": self._safe_json(metadata) or {},
|
|
"kwargs_json": self._safe_json(kwargs) or {},
|
|
"owner_worker_id": owner_worker_id,
|
|
"lease_expires_at": lease_dt,
|
|
"idempotency_key": idempotency_key,
|
|
"created_at": created,
|
|
"updated_at": now,
|
|
}
|
|
|
|
async with self._sf() as session:
|
|
# Keep the global clock -> run-row lock order used by every other
|
|
# mutator. One position covers this atomic set of changes; run_id
|
|
# provides deterministic ordering within the position.
|
|
change_seq = await self._next_change_seq(session)
|
|
claimed: list[dict[str, Any]] = []
|
|
|
|
if multitask_strategy in ("interrupt", "rollback"):
|
|
stmt = (
|
|
select(RunRow)
|
|
.where(
|
|
RunRow.thread_id == thread_id,
|
|
RunRow.status.in_(("pending", "running")),
|
|
)
|
|
.with_for_update()
|
|
)
|
|
result = await session.execute(stmt)
|
|
for row in result.scalars():
|
|
lease_expired = False
|
|
if row.lease_expires_at is not None:
|
|
# SQLite drops tzinfo on read despite
|
|
# ``DateTime(timezone=True)`` (see ``_row_to_dict``).
|
|
# Treat naive values as UTC — same convention as
|
|
# ``coerce_iso`` — so the Python-side comparison
|
|
# against the aware ``cutoff`` does not raise
|
|
# ``TypeError: can't compare offset-naive and
|
|
# offset-aware datetimes`` when heartbeat is enabled
|
|
# on SQLite.
|
|
row_lease = row.lease_expires_at
|
|
if row_lease.tzinfo is None:
|
|
row_lease = row_lease.replace(tzinfo=UTC)
|
|
lease_expired = row_lease < cutoff
|
|
if row_lease >= cutoff and row.owner_worker_id != owner_worker_id:
|
|
# Live run owned by another worker — we cannot
|
|
# interrupt it and the partial unique index would
|
|
# reject our INSERT anyway. Surface as
|
|
# ConflictError so the caller gets a clean signal
|
|
# instead of a retry loop on IntegrityError.
|
|
raise ConflictError(f"Thread {thread_id} already has an active run owned by another worker")
|
|
if row.operation_kind != "run" and not lease_expired:
|
|
raise ConflictError(f"Thread {thread_id} has an active checkpoint write")
|
|
row.status = "interrupted"
|
|
row.error = "Cancelled by newer run"
|
|
row.owner_worker_id = owner_worker_id
|
|
row.updated_at = now
|
|
row.change_seq = change_seq
|
|
claimed.append(self._row_to_dict(row))
|
|
|
|
values["change_seq"] = change_seq
|
|
session.add(RunRow(run_id=run_id, **values))
|
|
try:
|
|
await session.commit()
|
|
except IntegrityError as exc:
|
|
await session.rollback()
|
|
if idempotency_key is not None:
|
|
existing = (await session.execute(select(RunRow).where(RunRow.idempotency_key == idempotency_key))).scalar_one_or_none()
|
|
if existing is not None:
|
|
raise RunIdempotencyConflict(self._row_to_dict(existing)) from exc
|
|
raise
|
|
|
|
new_row = await session.get(RunRow, run_id)
|
|
return self._row_to_dict(new_row), claimed
|