diff --git a/openviking/server/routers/tasks.py b/openviking/server/routers/tasks.py index b40c829891..e941f0f3d2 100644 --- a/openviking/server/routers/tasks.py +++ b/openviking/server/routers/tasks.py @@ -104,10 +104,6 @@ async def list_tasks( if pagination == "cursor": from openviking.service.task_pagination import decode_cursor, encode_cursor, page_scope - # ROOT retains the legacy system + cached visibility and its own - # persisted tasks, including submissions restored after restart. - account = SYSTEM_TASK_ACCOUNT_ID if _ctx.role == Role.ROOT else _ctx.account_id - user = SYSTEM_TASK_USER_ID if _ctx.role == Role.ROOT else _ctx.user.user_id filters = { "task_type": task_type, "status": status, @@ -118,17 +114,32 @@ async def list_tasks( scope = page_scope( account=_ctx.account_id, user=_ctx.user.user_id, role=str(_ctx.role), **filters ) - tasks = await tracker.list_page( - account_id=account, - user_id=user, - limit=limit + 1, - before=decode_cursor(cursor, scope), - include_cached=_ctx.role == Role.ROOT, - additional_owner=(_ctx.account_id, _ctx.user.user_id) + list_call = ( + tracker.list_root_tasks( + account_id=_ctx.account_id, + user_id=_ctx.user.user_id, + limit=limit + 1, + before=decode_cursor(cursor, scope), + task_type=task_type, + status=status, + resource_id=resource_id, + include_internal=include_internal, + q=q, + ) if _ctx.role == Role.ROOT - else None, - **filters, + else tracker.list_tasks( + account_id=_ctx.account_id, + user_id=_ctx.user.user_id, + limit=limit + 1, + before=decode_cursor(cursor, scope), + task_type=task_type, + status=status, + resource_id=resource_id, + include_internal=include_internal, + q=q, + ) ) + tasks = await list_call more = len(tasks) > limit items = [t.to_dict() for t in tasks[:limit]] return Response( diff --git a/openviking/service/deletion.py b/openviking/service/deletion.py index 6a86b73d2f..708a8255cd 100644 --- a/openviking/service/deletion.py +++ b/openviking/service/deletion.py @@ -319,8 +319,13 @@ async def _process(self, message: dict[str, Any]) -> Optional[str]: if scheduler is not None: await run_to_completion(lambda: scheduler.delete_tasks(account_id, user_id)) await run_to_completion(lambda: self._cancel_tasks(account_id, user_id)) - if user_id is not None: - await run_to_completion(lambda: tracker.delete_user_tasks(account_id, user_id)) + await run_to_completion( + lambda: tracker.delete_tasks( + account_id, + user_id, + allow_active_types={"user_delete"} if user_id is None else (), + ) + ) vectors = self._service.viking_fs.vector_store if vectors is not None: @@ -383,7 +388,6 @@ async def _delete_account_files(self, account_id: str) -> None: await agfs.stat(path) except Exception as exc: if is_not_found_error(exc): - get_task_tracker().forget_account_tasks(account_id) return raise raise RuntimeError("Account directory still exists after deletion") diff --git a/openviking/service/external_task_service.py b/openviking/service/external_task_service.py index c913ee8616..923db2f265 100644 --- a/openviking/service/external_task_service.py +++ b/openviking/service/external_task_service.py @@ -343,7 +343,7 @@ async def execute(self, task_id: str, account_id: str, user_id: str) -> bool: await self._execute_task(task, provider, payload, account_id, user_id) return True except asyncio.CancelledError: - if not tracker.is_cancellation_requested(task_id): + if not await tracker.is_cancellation_requested(task_id): raise await run_to_completion(lambda: self.cancel_recovered(task_id, account_id, user_id)) raise diff --git a/openviking/service/fs_service.py b/openviking/service/fs_service.py index d61a7e033f..e64d6fc858 100644 --- a/openviking/service/fs_service.py +++ b/openviking/service/fs_service.py @@ -739,7 +739,7 @@ async def _wait_for_refresh(self) -> Dict[str, Any]: ) except (TimeoutError, asyncio.TimeoutError) as exc: raise DeadlineExceededError("queue processing", None) from exc - return tracker.build_queue_status( + return await tracker.build_queue_status( context.task_id, work_ids if context.work_id is not None else None ) diff --git a/openviking/service/memory_compile.py b/openviking/service/memory_compile.py index 25ae7e491b..dd53fdc8f5 100644 --- a/openviking/service/memory_compile.py +++ b/openviking/service/memory_compile.py @@ -192,7 +192,11 @@ async def _run( ctx: RequestContext, ) -> None: tracker = get_task_tracker() - tracker.register_running_task(task_id) + await tracker.register_running_task( + task_id, + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) root_token = None if bind_root_observability_context is not None: root_attrs = create_root_span_attributes( diff --git a/openviking/service/reindex_executor.py b/openviking/service/reindex_executor.py index 92043660f6..b25d71e7a8 100644 --- a/openviking/service/reindex_executor.py +++ b/openviking/service/reindex_executor.py @@ -685,7 +685,7 @@ async def _run( ) self._apply_embedding_wait_status( counters, - get_task_tracker().build_queue_status(task_context.task_id), + await get_task_tracker().build_queue_status(task_context.task_id), ) finally: active_lease = run.lease if run is not None else lease diff --git a/openviking/service/resource_service.py b/openviking/service/resource_service.py index b8f08bc715..dfa973a8ae 100644 --- a/openviking/service/resource_service.py +++ b/openviking/service/resource_service.py @@ -2463,7 +2463,11 @@ async def create_skill_task() -> None: await task_tracker.start( task.task_id, account_id=ctx.account_id, user_id=ctx.user.user_id ) - task_tracker.register_running_task(task.task_id) + await task_tracker.register_running_task( + task.task_id, + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) try: with bind_task_context(task.task_id, ctx.account_id, ctx.user.user_id): if isinstance(data, SkillProcessingPreparation): @@ -2543,7 +2547,11 @@ async def create_skill_task() -> None: "queue.wait.duration_ms", round((time.perf_counter() - wait_start) * 1000, 3), ) - status = task_tracker.build_queue_status(task.task_id) + status = await task_tracker.build_queue_status( + task.task_id, + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) result["queue_status"] = status self._raise_queue_status_errors(status) # Cancellation settles queue entries without counting errors. @@ -2582,7 +2590,11 @@ async def cancel_skill_processing(self, task_id: str, ctx: RequestContext) -> No return # Synchronous preparation failed before the task was created. # This includes queued work that must be discarded and active handlers # that must finish their cancellation/write cleanup, not just task status. - while tracker.has_work(task_id): + while await tracker.has_work( + task_id, + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ): await asyncio.sleep(0.05) async def build_index( diff --git a/openviking/service/task_domain.py b/openviking/service/task_domain.py new file mode 100644 index 0000000000..6aa5e3d0f5 --- /dev/null +++ b/openviking/service/task_domain.py @@ -0,0 +1,256 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Persistence-neutral Task and Work records.""" + +from __future__ import annotations + +import math +import re +import time +from copy import deepcopy +from dataclasses import dataclass, field +from datetime import datetime, timezone +from enum import Enum +from typing import Any, Dict, Iterable, Optional + + +class TaskStatus(str, Enum): + PENDING = "pending" + RUNNING = "running" + CANCELLING = "cancelling" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class WorkState(str, Enum): + PENDING = "pending" + RUNNING = "running" + DONE = "done" + FAILED = "failed" + + +class WorkRecoveryMode(str, Enum): + """How a Store rebuilds Work state after process startup.""" + + QUEUE_SNAPSHOT = "queue_snapshot" + DURABLE_STORE = "durable_store" + + +TERMINAL_TASK_STATUSES = ( + TaskStatus.COMPLETED, + TaskStatus.FAILED, + TaskStatus.CANCELLED, +) +ACTIVE_TASK_STATUSES = ( + TaskStatus.PENDING, + TaskStatus.RUNNING, + TaskStatus.CANCELLING, +) +TERMINAL_WORK_STATES = (WorkState.DONE, WorkState.FAILED) +MAX_QUEUE_ERROR_DETAILS = 100 + + +class TaskWorkRejected(Exception): + """Task lifecycle state prevented Work registration.""" + + +@dataclass(frozen=True) +class TaskRetentionPolicy: + completed_seconds: float + failed_seconds: float + + def is_expired(self, status: str, updated_at: float, now: float) -> bool: + if status in ("completed", "cancelled"): + return now - updated_at > self.completed_seconds + return status == "failed" and now - updated_at > self.failed_seconds + + +@dataclass(frozen=True) +class DeletedTask: + task_id: str + account_id: str + user_id: str + + +@dataclass(frozen=True) +class DeleteExpiredResult: + deleted: tuple[DeletedTask, ...] = () + scanned: int = 0 + scan_complete: bool = False + + @property + def deleted_count(self) -> int: + return len(self.deleted) + + +@dataclass +class QueueWorkSummary: + processed: int = 0 + requeue_count: int = 0 + error_count: int = 0 + errors: list[str] = field(default_factory=list) + queue_wait_ms: float = 0.0 + execute_ms: float = 0.0 + + +@dataclass +class WorkSummary: + queues: Dict[str, QueueWorkSummary] = field(default_factory=dict) + context_count: int = 0 + + +SENSITIVE_TASK_KEYS = frozenset({"api_key", "user_key"}) +_SENSITIVE_PATTERNS = re.compile( + r"(sk-|cr_|ghp_|ntn_|xox[baprs]-|Bearer\s+)[a-zA-Z0-9._-]+", + re.IGNORECASE, +) +_MAX_ERROR_LEN = 500 + + +def sanitize_error(error: str) -> str: + sanitized = _SENSITIVE_PATTERNS.sub("[REDACTED]", error) + if len(sanitized) > _MAX_ERROR_LEN: + sanitized = sanitized[:_MAX_ERROR_LEN] + "...[truncated]" + return sanitized + + +def sanitize_task_result(result: Any) -> Any: + if isinstance(result, dict): + return { + key: sanitize_task_result(value) + for key, value in result.items() + if key not in SENSITIVE_TASK_KEYS + } + if isinstance(result, list): + return [sanitize_task_result(item) for item in result] + return result + + +@dataclass +class TaskRecord: + task_id: str + task_type: str + status: TaskStatus = TaskStatus.PENDING + created_at: float = field(default_factory=time.time) + updated_at: float = field(default_factory=time.time) + resource_id: Optional[str] = None + account_id: Optional[str] = None + user_id: Optional[str] = None + meta: Dict[str, Any] = field(default_factory=dict) + stage: Optional[str] = None + result: Optional[Dict[str, Any]] = None + error: Optional[str] = None + auth: Dict[str, Any] = field(default_factory=dict, repr=False) + execution_events: Optional[Any] = None + processing_seconds: Optional[float] = None + _extra_fields: Dict[str, Any] = field( + default_factory=dict, init=False, repr=False, compare=False + ) + + def to_dict(self, *, include_events: bool = False) -> Dict[str, Any]: + public_meta = deepcopy(self.meta) + public_meta.pop("submission_hash", None) + public_meta.pop("submission_token", None) + return { + **({"execution_events": deepcopy(self.execution_events)} if include_events else {}), + "task_id": self.task_id, + "task_type": self.task_type, + "status": self.status.value, + "created_at": self.created_at, + "updated_at": self.updated_at, + "processing_seconds": self.processing_seconds, + "resource_id": self.resource_id, + "meta": sanitize_task_result(public_meta), + "stage": self.stage, + "result": sanitize_task_result(deepcopy(self.result)), + "error": self.error, + "created_at_iso": datetime.fromtimestamp(self.created_at, tz=timezone.utc).isoformat(), + "updated_at_iso": datetime.fromtimestamp(self.updated_at, tz=timezone.utc).isoformat(), + } + + def is_terminal(self) -> bool: + return self.status in TERMINAL_TASK_STATUSES + + def is_active(self) -> bool: + return self.status in ACTIVE_TASK_STATUSES + + @staticmethod + def next_updated_at(previous: float) -> float: + return max(time.time(), math.nextafter(previous, math.inf)) + + +@dataclass +class WorkRecord: + work_id: str + task_id: str + queue_name: str + state: WorkState = WorkState.PENDING + requeue_count: int = 0 + error: Optional[str] = None + errors: list[str] = field(default_factory=list) + context_count: int = 0 + enqueued_at: float = field(default_factory=time.time) + started_at: Optional[float] = None + finished_at: Optional[float] = None + updated_at: float = field(default_factory=time.time) + + def is_open(self) -> bool: + return self.state not in TERMINAL_WORK_STATES + + +def processing_seconds( + works: Iterable[WorkRecord], *, now: Optional[float] = None +) -> Optional[float]: + """Return the wall-clock union of claimed Work intervals.""" + current = time.time() if now is None else now + intervals = sorted( + ( + work.started_at, + max( + work.started_at, + work.finished_at if work.finished_at is not None else current, + ), + ) + for work in works + if work.started_at is not None + ) + if not intervals: + return None + elapsed = 0.0 + left, right = intervals[0] + for start, end in intervals[1:]: + if start <= right: + right = max(right, end) + else: + elapsed += right - left + left, right = start, end + return elapsed + right - left + + +def summarize_works(works: Iterable[WorkRecord]) -> WorkSummary: + """Aggregate persisted Work facts without imposing an API response shape.""" + summary = WorkSummary() + for work in works: + bucket = summary.queues.setdefault(work.queue_name, QueueWorkSummary()) + bucket.requeue_count += work.requeue_count + summary.context_count += work.context_count + if work.state == WorkState.DONE: + bucket.processed += 1 + elif work.state == WorkState.FAILED: + errors = work.errors or ([work.error] if work.error else []) + bucket.error_count += len(errors) or 1 + bucket.errors.extend(errors) + if work.started_at is None: + continue + bucket.queue_wait_ms += max(0.0, (work.started_at - work.enqueued_at) * 1000.0) + if work.finished_at is not None: + bucket.execute_ms += max(0.0, (work.finished_at - work.started_at) * 1000.0) + return summary + + +def first_work_error(works: Iterable[WorkRecord]) -> Optional[str]: + return next( + (work.error for work in works if work.state == WorkState.FAILED and work.error), + None, + ) diff --git a/openviking/service/task_queue_middleware.py b/openviking/service/task_queue_middleware.py index c542cae9a3..d9826f06ff 100644 --- a/openviking/service/task_queue_middleware.py +++ b/openviking/service/task_queue_middleware.py @@ -1,11 +1,13 @@ # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. # SPDX-License-Identifier: AGPL-3.0 -"""Adapt the existing in-memory TaskWorkIndex to queue lifecycle middleware.""" +"""Bridge QueueFS lifecycle events to Store-authoritative Task Work.""" import asyncio +from typing import TYPE_CHECKING +from openviking.service.task_store import WorkRegistration +from openviking.service.task_tracker_concurrency import KeyedAsyncLockPool, run_to_completion from openviking.service.task_work_index import ( - TaskWorkIndex, TaskWorkRejected, bind_task_context, extract_task_metadata, @@ -14,8 +16,6 @@ ) from openviking.storage.queuefs.process_result import ProcessOutcome, ProcessResult from openviking.storage.queuefs.queue_middleware import ( - AckContext, - AckNext, ClearContext, ClearNext, EnqueueContext, @@ -25,123 +25,260 @@ QueueMiddleware, ) +if TYPE_CHECKING: + from openviking.service.task_tracker import TaskTracker + from openviking.service.task_work_index import QueueTaskMetadata + class TaskWorkQueueMiddleware(QueueMiddleware): - """Keep QueueFS as the durable source of truth; no new task state or store.""" + """Map tracked queue deliveries to semantic TaskTracker operations.""" + + def __init__(self) -> None: + self._enqueue_locks = KeyedAsyncLockPool() - def __init__(self, index: TaskWorkIndex) -> None: - self._index = index + @property + def _tracker(self) -> "TaskTracker": + from openviking.service.task_tracker import get_task_tracker + + return get_task_tracker() async def enqueue(self, ctx: EnqueueContext, call_next: EnqueueNext) -> str: if not isinstance(ctx.payload, dict): return await call_next(ctx) ctx.payload, metadata = prepare_task_payload(ctx.payload) ctx.work_id = metadata.work_id if metadata is not None else None - if not self._index.register(ctx.queue, metadata): + if metadata is None: + return await call_next(ctx) + if not metadata.account_id or not metadata.user_id: + raise TaskWorkRejected(f"Task {metadata.task_id} work is missing owner metadata") + + key = (metadata.account_id, metadata.user_id, metadata.task_id, metadata.work_id) + async with self._enqueue_locks.acquire(key): + return await self._enqueue_registered(ctx, call_next, metadata) + + async def _enqueue_registered( + self, + ctx: EnqueueContext, + call_next: EnqueueNext, + metadata: "QueueTaskMetadata", + ) -> str: + registration = WorkRegistration.REJECTED + + async def register() -> None: + nonlocal registration + registration = await self._tracker.register_work( + metadata.task_id, + metadata.work_id, + ctx.queue, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) + + await run_to_completion(register) + if registration is WorkRegistration.REJECTED: raise TaskWorkRejected( - f"Task {metadata.task_id} is cancelling; rejected work for {ctx.queue}" + f"Task {metadata.task_id} is unavailable; rejected work for {ctx.queue}" ) try: - result = await call_next(ctx) + message_id = await call_next(ctx) except asyncio.CancelledError: - if metadata is not None: - if ctx.committed: - record_enqueued_work(ctx.queue, metadata.work_id) - else: - await self._index.discard(ctx.queue, metadata) - raise - except BaseException as exc: - if metadata is not None: - if ctx.committed: - record_enqueued_work(ctx.queue, metadata.work_id) - else: - await self._index.fail_enqueue( - ctx.queue, - metadata, - str(exc), + if ctx.committed: + record_enqueued_work(ctx.queue, metadata.work_id) + elif registration is WorkRegistration.CREATED: + await run_to_completion( + lambda: self._tracker.discard_work( + metadata.task_id, + metadata.work_id, + account_id=metadata.account_id, + user_id=metadata.user_id, ) - record_enqueued_work(ctx.queue, metadata.work_id) + ) raise - if metadata is not None: + except BaseException as exc: if ctx.committed: record_enqueued_work(ctx.queue, metadata.work_id) else: - await self._index.discard(ctx.queue, metadata) - return result + failure_message = str(exc) or type(exc).__name__ + await run_to_completion( + lambda: self._tracker.mark_work_failed( + metadata.task_id, + metadata.work_id, + failure_message, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) + ) + record_enqueued_work(ctx.queue, metadata.work_id) + raise + + if ctx.committed: + record_enqueued_work(ctx.queue, metadata.work_id) + elif registration is WorkRegistration.CREATED: + await run_to_completion( + lambda: self._tracker.discard_work( + metadata.task_id, + metadata.work_id, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) + ) + return message_id async def process(self, ctx: ProcessContext, call_next: ProcessNext) -> ProcessResult: metadata = extract_task_metadata(ctx.message) if metadata is None: return await call_next(ctx) + if not metadata.account_id or not metadata.user_id: + raise TaskWorkRejected(f"Task {metadata.task_id} work is missing owner metadata") - active_task = asyncio.current_task() - enqueued_work = [] - with bind_task_context( + duplicate = await self._reconcile_delivery(metadata, ctx.queue) + if duplicate is not None: + return duplicate + + if await self._tracker.is_cancelling( metadata.task_id, - metadata.account_id, - metadata.user_id, - enqueued_work=enqueued_work, - work_id=metadata.work_id, + account_id=metadata.account_id, + user_id=metadata.user_id, ): - if self._index.cancellation_requested(metadata.task_id): - result = await ctx.cancel() - self._record_result(metadata, result) - return result - if active_task is not None: - self._index.register_active( - metadata.task_id, - active_task, - work_id=metadata.work_id, - ) - try: - result = await call_next(ctx) - except asyncio.CancelledError: - if not self._index.cancellation_requested(metadata.task_id): - raise - # Handler unwinding finished; settle user cancellation before ACK. - result = ProcessResult.cancelled() - except Exception as exc: - self._index.record_failure(metadata.task_id, str(exc)) - raise - finally: - if active_task is not None: - self._index.unregister_active(metadata.task_id, active_task) - self._record_result(metadata, result) + result = await ctx.cancel() + await self._persist_result(metadata, result) return result - def _record_result(self, metadata, result: ProcessResult) -> None: - outcome = { - ProcessOutcome.SUCCESS: "completed", - ProcessOutcome.FAILED: "failed", - ProcessOutcome.REQUEUED: "requeued", - ProcessOutcome.CANCELLED: "cancelled", - }[result.outcome] - self._index.record_result( - metadata, - outcome, - error=result.error, - errors=result.errors, + await self._tracker.start_work( + metadata.task_id, + metadata.work_id, + account_id=metadata.account_id, + user_id=metadata.user_id, ) - - async def ack(self, ctx: AckContext, call_next: AckNext) -> None: - prepared = None + await self._tracker.register_running_task( + metadata.task_id, + work_id=metadata.work_id, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) + enqueued_work: list[tuple[str, str]] = [] try: - if ctx.message is not None: - prepared = await self._index.prepare_ack(ctx.queue, ctx.message) - await call_next(ctx) + with bind_task_context( + metadata.task_id, + metadata.account_id, + metadata.user_id, + enqueued_work=enqueued_work, + work_id=metadata.work_id, + ): + try: + result = await call_next(ctx) + except asyncio.CancelledError: + if not await self._tracker.is_cancelling( + metadata.task_id, + account_id=metadata.account_id, + user_id=metadata.user_id, + ): + raise + result = ProcessResult.cancelled() + except Exception as exc: + await self._tracker.record_work_error( + metadata.task_id, + metadata.work_id, + str(exc) or type(exc).__name__, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) + raise + await self._persist_result(metadata, result) + return result finally: - if prepared is not None and not ctx.committed: - self._index.rollback_ack(ctx.queue, prepared) + await self._tracker.unregister_running_task(metadata.task_id) + + async def _reconcile_delivery( + self, + metadata: "QueueTaskMetadata", + queue_name: str, + ) -> ProcessResult | None: + from openviking.service.task_domain import WorkState + + state = await self._tracker.get_work_state( + metadata.task_id, + metadata.work_id, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) + if state is None: + task = await self._tracker.get( + metadata.task_id, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) + if task is not None and task.is_terminal(): + return ProcessResult.success() + await self._tracker.reconcile_delivery( + metadata.task_id, + metadata.work_id, + queue_name, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) + return None + if state in (WorkState.DONE, WorkState.FAILED): + await self._tracker.settle_work( + metadata.task_id, + metadata.work_id, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) + return ProcessResult.success() + return None + + async def _persist_result( + self, + metadata: "QueueTaskMetadata", + result: ProcessResult, + ) -> None: + if result.outcome is ProcessOutcome.FAILED: + await self._tracker.mark_work_failed( + metadata.task_id, + metadata.work_id, + result.error or "Queue handler failed", + errors=result.errors, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) + elif result.outcome is ProcessOutcome.REQUEUED: + await self._tracker.mark_work_requeued( + metadata.task_id, + metadata.work_id, + delta=1, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) + else: + vector_written = bool( + isinstance(result.value, dict) and result.value.get("vector_written") + ) + await self._tracker.mark_work_done( + metadata.task_id, + metadata.work_id, + vector_written=vector_written, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) async def clear(self, ctx: ClearContext, call_next: ClearNext) -> None: - prepared = [] - try: - for message in ctx.messages: - metadata = await self._index.prepare_ack(ctx.queue, message) - if metadata is not None: - prepared.append(metadata) - await call_next(ctx) - finally: - if not ctx.committed: - for metadata in prepared: - self._index.rollback_ack(ctx.queue, metadata) + tracked = [ + metadata + for message in ctx.messages + if (metadata := extract_task_metadata(message)) is not None + and metadata.account_id + and metadata.user_id + ] + await call_next(ctx) + if not ctx.committed: + return + for metadata in tracked: + await run_to_completion( + lambda metadata=metadata: self._tracker.settle_work( + metadata.task_id, + metadata.work_id, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) + ) diff --git a/openviking/service/task_store.py b/openviking/service/task_store.py index bcb8b4e200..3a9f9751bd 100644 --- a/openviking/service/task_store.py +++ b/openviking/service/task_store.py @@ -1,249 +1,1533 @@ # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. # SPDX-License-Identifier: AGPL-3.0 -"""Internal storage backends for TaskTracker.""" +"""Task/Work storage contract and the local file-backed implementation.""" from __future__ import annotations +import asyncio import heapq +import inspect import json -from collections.abc import Awaitable, Callable +import threading +import time +from collections import defaultdict +from collections.abc import AsyncIterator, Callable, Collection from copy import deepcopy +from dataclasses import dataclass, fields +from enum import Enum from typing import Any, Dict, List, Optional, Protocol from openviking.pyagfs import AsyncAGFSClient from openviking.pyagfs.exceptions import AGFSAlreadyExistsError, AGFSNotFoundError -from openviking.service.task_tracker_concurrency import StoreIOLimiter, run_to_completion +from openviking.service.task_domain import ( + ACTIVE_TASK_STATUSES, + MAX_QUEUE_ERROR_DETAILS, + TERMINAL_TASK_STATUSES, + DeletedTask, + DeleteExpiredResult, + QueueWorkSummary, + TaskRecord, + TaskRetentionPolicy, + TaskStatus, + TaskWorkRejected, + WorkRecord, + WorkRecoveryMode, + WorkState, + WorkSummary, + first_work_error, + processing_seconds, + sanitize_error, + summarize_works, +) +from openviking.service.task_events import append_task_event +from openviking.service.task_tracker_concurrency import ( + KeyedAsyncLockPool, + StoreIOLimiter, + run_to_completion, +) SYSTEM_TASK_ACCOUNT_ID = "_system" SYSTEM_TASK_USER_ID = "root" +TaskStatsSnapshot = Dict[str, Dict[str, int]] -class TaskStore(Protocol): - async def create(self, task: Any) -> None: ... +class WorkRegistration(str, Enum): + CREATED = "created" + EXISTING = "existing" + REJECTED = "rejected" - async def update(self, task: Any) -> None: ... + def __bool__(self) -> bool: + return self is not WorkRegistration.REJECTED - async def get( + +class TaskAlreadyActive(RuntimeError): + def __init__(self, task_type: str, resource_id: Optional[str]) -> None: + self.task_type = task_type + self.resource_id = resource_id + super().__init__(f"An active {task_type} task already exists for {resource_id}") + + +class TaskWorkStore(Protocol): + """Atomic semantic operations required by TaskTracker. + + Implementations own transaction and revision details. Callers never replace + a complete Task/Work aggregate or perform compare-and-set retries. + """ + + @property + def work_recovery_mode(self) -> WorkRecoveryMode: ... + + async def create(self, task: TaskRecord, *, reject_if_active: bool = False) -> TaskRecord: ... + + async def get_task( self, task_id: str, *, account_id: Optional[str] = None, user_id: Optional[str] = None, - ) -> Optional[Dict[str, Any]]: ... + ) -> Optional[TaskRecord]: ... async def list( - self, account_id: str, *, user_id: Optional[str] = None - ) -> List[Dict[str, Any]]: ... + self, + account_id: Optional[str] = None, + *, + user_id: Optional[str] = None, + task_type: Optional[str] = None, + status: Optional[str] = None, + resource_id: Optional[str] = None, + limit: Optional[int] = None, + before: tuple[float, str] | None = None, + include_internal: bool = True, + q: Optional[str] = None, + ) -> List[TaskRecord]: ... - async def list_page( + async def delete_tasks( self, + task_id: Optional[str] = None, + *, account_id: str, + user_id: Optional[str] = None, + allow_active_types: Collection[str] = (), + ) -> List[str]: ... + + async def delete_expired( + self, policy: TaskRetentionPolicy, *, now: float, limit: int + ) -> DeleteExpiredResult: ... + + def snapshot_task_stats(self) -> TaskStatsSnapshot: ... + + async def start_task( + self, + task_id: str, + *, + stage: Optional[str], + account_id: Optional[str], + user_id: Optional[str], + ) -> Optional[TaskRecord]: ... + + async def update_stage( + self, + task_id: str, + stage: str, *, + meta: Optional[Dict[str, Any]], + account_id: Optional[str], + user_id: Optional[str], + ) -> Optional[TaskRecord]: ... + + async def update_auth( + self, + task_id: str, + values: Dict[str, Any], + *, + account_id: str, user_id: str, - limit: int, - before: tuple[float, str] | None = None, - prune_expired: Callable[[Dict[str, Any]], Awaitable[bool]] | None = None, - io_limiter: StoreIOLimiter | None = None, - **filters: Any, - ) -> List[Dict[str, Any]]: ... + ) -> Optional[TaskRecord]: ... + + async def merge_meta( + self, + task_id: str, + values: Dict[str, Any], + *, + account_id: str, + user_id: str, + ) -> Optional[TaskRecord]: ... - async def delete( - self, task_id: str, *, account_id: str, user_id: Optional[str] = None + async def record_event( + self, + task_id: str, + kind: str, + *, + operation: Optional[str], + account_id: str, + user_id: str, + ) -> Optional[TaskRecord]: ... + + async def record_outcome( + self, + task_id: str, + *, + result: Optional[Dict[str, Any]], + error: Optional[str], + resource_id: Optional[str], + cancelled: bool, + account_id: Optional[str], + user_id: Optional[str], + ) -> Optional[TaskRecord]: ... + + async def request_cancellation( + self, + task_id: str, + *, + cancellable_types: Collection[str], + allow_terminal: bool, + account_id: Optional[str], + user_id: Optional[str], + ) -> Optional[TaskRecord]: ... + + async def mark_enqueue_complete( + self, + task_id: str, + work_id: str, + *, + account_id: str, + user_id: str, + ) -> None: ... + + async def is_enqueue_complete( + self, + task_id: str, + work_id: str, + *, + account_id: str, + user_id: str, + ) -> bool: ... + + async def add_work( + self, work: WorkRecord, *, account_id: str, user_id: str + ) -> WorkRegistration: ... + + async def reconcile_delivery( + self, work: WorkRecord, *, account_id: str, user_id: str + ) -> None: ... + + async def discard_work( + self, task_id: str, work_id: str, *, account_id: str, user_id: str + ) -> Optional[TaskRecord]: ... + + async def get_work( + self, task_id: str, work_id: str, *, account_id: str, user_id: str + ) -> Optional[WorkRecord]: ... + + async def get_works( + self, + task_id: str, + work_ids: Collection[str], + *, + account_id: Optional[str], + user_id: Optional[str], + ) -> Dict[str, WorkRecord]: ... + + async def start_work( + self, task_id: str, work_id: str, *, account_id: str, user_id: str ) -> None: ... + async def finish_work( + self, + task_id: str, + work_id: str, + state: WorkState, + *, + error: Optional[str], + errors: Optional[Collection[str]], + context_count_delta: int, + account_id: str, + user_id: str, + ) -> Optional[TaskRecord]: ... + + async def requeue_work( + self, + task_id: str, + work_id: str, + *, + delta: int, + account_id: str, + user_id: str, + ) -> None: ... + + async def record_work_error( + self, + task_id: str, + work_id: str, + error: str, + *, + account_id: str, + user_id: str, + ) -> None: ... + + async def settle_work( + self, task_id: str, work_id: str, *, account_id: str, user_id: str + ) -> Optional[TaskRecord]: ... + + async def has_open_work( + self, + task_id: str, + *, + exclude_work_id: Optional[str], + account_id: Optional[str], + user_id: Optional[str], + ) -> bool: ... + + async def summarize_works( + self, + task_id: str, + *, + work_ids: Optional[Collection[str]], + account_id: Optional[str], + user_id: Optional[str], + ) -> WorkSummary: ... + + async def filter_cancellation_requested( + self, + task_ids: Collection[str], + *, + account_id: str, + user_id: str, + ) -> set[str]: ... + + async def is_cancellation_requested( + self, + task_id: str, + *, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> bool: ... + + +# Compatibility name for integrations that only annotate the constructor. +TaskStore = TaskWorkStore + + +@dataclass +class _WorkSet: + works: Dict[str, WorkRecord] + open_count: int = 0 + + +@dataclass(frozen=True) +class _WorkMutationResult: + changed: bool + work: Optional[WorkRecord] -class PersistentTaskStore: - """Persist task records into AGFS under account-scoped system task directories.""" + +class PersistentTaskStore(TaskWorkStore): + """File Task store with process-local Work state reconstructed from QueueFS. + + Only Task records are durable. QueueFS startup recovery can reconstruct + unacknowledged Work, but terminal Work history, timings, and errors may be + incomplete after restart. + """ ROOT_PREFIX = "/local" SYSTEM_DIRNAME = "_system" TASKS_DIRNAME = "tasks" - def __init__(self, agfs: Any) -> None: + def __init__(self, agfs: Any, *, max_concurrent_io: int = 8) -> None: self._agfs = agfs if isinstance(agfs, AsyncAGFSClient) else AsyncAGFSClient(agfs) self._ensured_task_dirs: set[tuple[str, str]] = set() + self._lock = threading.Lock() + self._task_cache: Dict[tuple[str, str, str], TaskRecord] = {} + self._work_sets: Dict[tuple[str, str, str], _WorkSet] = {} + self._business_locks = KeyedAsyncLockPool() + self._record_locks = KeyedAsyncLockPool() + self._cleanup_lock = KeyedAsyncLockPool() + self._cleanup_scan: Optional[AsyncIterator[Optional[Dict[str, Any]]]] = None + self._io = StoreIOLimiter(max_concurrent_io) - async def create(self, task: Any) -> None: - await self._write_task(task) + @property + def work_recovery_mode(self) -> WorkRecoveryMode: + return WorkRecoveryMode.QUEUE_SNAPSHOT - async def update(self, task: Any) -> None: - await self._write_task(task) + async def create(self, task: TaskRecord, *, reject_if_active: bool = False) -> TaskRecord: + if not task.account_id or not task.user_id: + raise ValueError("PersistentTaskStore requires account_id and user_id") + business_key = (task.account_id, task.user_id, task.task_type, task.resource_id) + async with self._business_locks.acquire(business_key): + if reject_if_active: + candidates = await self.list(task.account_id, user_id=task.user_id) + if any( + item.task_type == task.task_type + and item.resource_id == task.resource_id + and item.status in ACTIVE_TASK_STATUSES + for item in candidates + ): + raise TaskAlreadyActive(task.task_type, task.resource_id) + owner = (task.account_id, task.user_id, task.task_id) + async with self._record_locks.acquire(owner): + existing = await self._load_task(task.task_id, task.account_id, task.user_id) + if existing is not None: + return self._copy_task(existing, self._work_sets.get(owner)) + stored = deepcopy(task) + self._append_events(stored, None) + await self._write_task(stored) + with self._lock: + self._task_cache[owner] = stored + self._work_sets[owner] = _WorkSet({}) + return deepcopy(stored) - async def get( + async def get_task( self, task_id: str, *, account_id: Optional[str] = None, user_id: Optional[str] = None, - ) -> Optional[Dict[str, Any]]: - if not account_id or not user_id: + ) -> Optional[TaskRecord]: + task = await self._load_task(task_id, account_id, user_id) + if task is None: return None - path = self._task_path(account_id, user_id, task_id) - try: - raw = await self._agfs.read(path) - except (AGFSNotFoundError, FileNotFoundError): + with self._lock: + owner = self._find_cached_owner(task_id, account_id, user_id) + if owner is None: + return None + return self._copy_task(self._task_cache[owner], self._work_sets.get(owner)) + + async def start_task( + self, + task_id: str, + *, + stage: Optional[str], + account_id: Optional[str], + user_id: Optional[str], + ) -> Optional[TaskRecord]: + def mutate(task: TaskRecord, _work_set: _WorkSet) -> bool: + if task.status not in (TaskStatus.PENDING, TaskStatus.RUNNING): + return False + task.status = TaskStatus.RUNNING + if stage is not None: + task.stage = stage + return True + + return await self._mutate_task(task_id, account_id, user_id, mutate) + + async def update_stage( + self, + task_id: str, + stage: str, + *, + meta: Optional[Dict[str, Any]], + account_id: Optional[str], + user_id: Optional[str], + ) -> Optional[TaskRecord]: + def mutate(task: TaskRecord, _work_set: _WorkSet) -> bool: + if task.status not in ACTIVE_TASK_STATUSES: + return False + task.stage = stage + if meta: + _deep_merge(task.meta, deepcopy(meta)) + return True + + return await self._mutate_task(task_id, account_id, user_id, mutate) + + async def update_auth( + self, + task_id: str, + values: Dict[str, Any], + *, + account_id: str, + user_id: str, + ) -> Optional[TaskRecord]: + def mutate(task: TaskRecord, _work_set: _WorkSet) -> bool: + if task.status not in ACTIVE_TASK_STATUSES: + return False + _deep_merge(task.auth, deepcopy(values)) + return True + + return await self._mutate_task(task_id, account_id, user_id, mutate) + + async def merge_meta( + self, + task_id: str, + values: Dict[str, Any], + *, + account_id: str, + user_id: str, + ) -> Optional[TaskRecord]: + def mutate(task: TaskRecord, _work_set: _WorkSet) -> bool: + if task.status not in (TaskStatus.PENDING, TaskStatus.RUNNING): + raise ValueError("Task is no longer active") + _deep_merge(task.meta, deepcopy(values)) + return True + + return await self._mutate_task(task_id, account_id, user_id, mutate) + + async def record_event( + self, + task_id: str, + kind: str, + *, + operation: Optional[str], + account_id: str, + user_id: str, + ) -> Optional[TaskRecord]: + def mutate(task: TaskRecord, _work_set: _WorkSet) -> bool: + if task.status not in ACTIVE_TASK_STATUSES: + return False + task.execution_events = append_task_event( + task.execution_events, + kind=kind, + status=task.status.value, + stage=sanitize_error(task.stage) if task.stage else None, + operation=operation, + ) + return True + + return await self._mutate_task(task_id, account_id, user_id, mutate) + + async def record_outcome( + self, + task_id: str, + *, + result: Optional[Dict[str, Any]], + error: Optional[str], + resource_id: Optional[str], + cancelled: bool, + account_id: Optional[str], + user_id: Optional[str], + ) -> Optional[TaskRecord]: + def mutate(task: TaskRecord, work_set: _WorkSet) -> bool: + if task.status not in ACTIVE_TASK_STATUSES: + return False + if result is not None: + task.result = deepcopy(result) + if resource_id is not None: + task.resource_id = resource_id + if error and task.error is None: + task.error = sanitize_error(error) + if cancelled: + task.status = TaskStatus.CANCELLING + enqueue_markers = task.auth.get("enqueue_complete_work_ids") + task.auth = ( + {"enqueue_complete_work_ids": deepcopy(enqueue_markers)} if enqueue_markers else {} + ) + self._try_finalize(task, work_set) + return True + + return await self._mutate_task(task_id, account_id, user_id, mutate) + + async def request_cancellation( + self, + task_id: str, + *, + cancellable_types: Collection[str], + allow_terminal: bool, + account_id: Optional[str], + user_id: Optional[str], + ) -> Optional[TaskRecord]: + def mutate(task: TaskRecord, work_set: _WorkSet) -> bool: + if task.status in (TaskStatus.CANCELLING, TaskStatus.CANCELLED): + return False + if task.task_type not in cancellable_types: + raise ValueError(f"Task type '{task.task_type}' does not support cancellation") + if task.status in TERMINAL_TASK_STATUSES and not allow_terminal: + raise ValueError(f"Task is already {task.status.value}") + task.status = TaskStatus.CANCELLING + self._try_finalize(task, work_set) + return True + + return await self._mutate_task(task_id, account_id, user_id, mutate) + + async def mark_enqueue_complete( + self, + task_id: str, + work_id: str, + *, + account_id: str, + user_id: str, + ) -> None: + def mutate(task: TaskRecord, _work_set: _WorkSet) -> bool: + if task.status not in ACTIVE_TASK_STATUSES: + return False + completed = set(task.auth.get("enqueue_complete_work_ids") or ()) + if work_id in completed: + return False + completed.add(work_id) + task.auth["enqueue_complete_work_ids"] = sorted(completed) + return True + + await self._mutate_task(task_id, account_id, user_id, mutate) + + async def is_enqueue_complete( + self, + task_id: str, + work_id: str, + *, + account_id: str, + user_id: str, + ) -> bool: + task = await self.get_task(task_id, account_id=account_id, user_id=user_id) + return bool(task and work_id in set(task.auth.get("enqueue_complete_work_ids") or ())) + + async def add_work( + self, work: WorkRecord, *, account_id: str, user_id: str + ) -> WorkRegistration: + registration = WorkRegistration.CREATED + + def mutate(task: TaskRecord, existing: Optional[WorkRecord]) -> _WorkMutationResult: + nonlocal registration + if task.is_terminal() or task.status == TaskStatus.CANCELLING: + raise TaskWorkRejected(f"Task {work.task_id} is {task.status.value}; rejected work") + if existing is not None: + registration = WorkRegistration.EXISTING + return _WorkMutationResult(False, existing) + return _WorkMutationResult(True, deepcopy(work)) + + found, _ = await self._mutate_work( + work.task_id, + work.work_id, + account_id, + user_id, + mutate, + ) + return registration if found else WorkRegistration.REJECTED + + async def reconcile_delivery(self, work: WorkRecord, *, account_id: str, user_id: str) -> None: + """Ensure a delivered QueueFS message has local Work state.""" + + def mutate(task: TaskRecord, existing: Optional[WorkRecord]) -> _WorkMutationResult: + if task.is_terminal(): + return _WorkMutationResult(False, existing) + if existing is None: + return _WorkMutationResult(True, deepcopy(work)) + if existing.is_open() and existing.queue_name != work.queue_name: + existing.queue_name = work.queue_name + return _WorkMutationResult(True, existing) + return _WorkMutationResult(False, existing) + + await self._mutate_work( + work.task_id, + work.work_id, + account_id, + user_id, + mutate, + ) + + async def discard_work( + self, task_id: str, work_id: str, *, account_id: str, user_id: str + ) -> Optional[TaskRecord]: + def mutate(_task: TaskRecord, existing: Optional[WorkRecord]) -> _WorkMutationResult: + return _WorkMutationResult(existing is not None, None) + + _, task = await self._mutate_work( + task_id, + work_id, + account_id, + user_id, + mutate, + try_finalize=True, + ) + return task + + async def get_work( + self, task_id: str, work_id: str, *, account_id: str, user_id: str + ) -> Optional[WorkRecord]: + task = await self._load_task(task_id, account_id, user_id) + if task is None: return None - return json.loads(_decode_bytes(raw)) + with self._lock: + work_set = self._work_sets.get((account_id, user_id, task_id)) + work = work_set.works.get(work_id) if work_set else None + return deepcopy(work) if work is not None else None - async def list(self, account_id: str, *, user_id: Optional[str] = None) -> List[Dict[str, Any]]: - if user_id is None: - try: - owners = await self._agfs.ls(self._task_root_dir(account_id)) - except (AGFSNotFoundError, FileNotFoundError): - return [] - tasks: List[Dict[str, Any]] = [] - for owner in owners: - if owner.get("isDir") and owner.get("name") not in (".", ".."): - tasks.extend(await self.list(account_id, user_id=owner["name"])) - return tasks - directory = self._task_dir(account_id, user_id) - try: - items = await self._agfs.ls(directory) - except (AGFSNotFoundError, FileNotFoundError): - return [] - tasks: List[Dict[str, Any]] = [] - for item in items: - path = item.get("path") or f"{directory}/{item.get('name', '')}" - if not path.endswith(".json"): - continue - try: - raw = await self._agfs.read(path) - except (AGFSNotFoundError, FileNotFoundError): - continue - tasks.append(json.loads(_decode_bytes(raw))) - return tasks + async def get_works( + self, + task_id: str, + work_ids: Collection[str], + *, + account_id: Optional[str], + user_id: Optional[str], + ) -> Dict[str, WorkRecord]: + task = await self._load_task(task_id, account_id, user_id) + if task is None or not task.account_id or not task.user_id: + return {} + wanted = set(work_ids) + with self._lock: + work_set = self._work_sets.get((task.account_id, task.user_id, task_id)) + return { + work_id: deepcopy(work) + for work_id, work in (work_set.works.items() if work_set else ()) + if work_id in wanted + } + + async def start_work( + self, task_id: str, work_id: str, *, account_id: str, user_id: str + ) -> None: + def mutate(_task: TaskRecord, work: Optional[WorkRecord]) -> _WorkMutationResult: + if work is None or not work.is_open(): + return _WorkMutationResult(False, work) + now = time.time() + work.state = WorkState.RUNNING + work.started_at = work.started_at or now + work.finished_at = None + work.updated_at = now + return _WorkMutationResult(True, work) + + await self._mutate_work(task_id, work_id, account_id, user_id, mutate) - async def list_page( + async def finish_work( self, + task_id: str, + work_id: str, + state: WorkState, + *, + error: Optional[str], + errors: Optional[Collection[str]], + context_count_delta: int, account_id: str, + user_id: str, + ) -> Optional[TaskRecord]: + if state not in (WorkState.DONE, WorkState.FAILED): + raise ValueError(f"Non-terminal Work state: {state}") + + def mutate(_task: TaskRecord, work: Optional[WorkRecord]) -> _WorkMutationResult: + if work is None or not work.is_open(): + return _WorkMutationResult(False, work) + now = time.time() + work.state = state + work.finished_at = now + work.updated_at = now + if error: + work.error = sanitize_error(error) + if state == WorkState.FAILED: + details = [sanitize_error(str(value)) for value in errors or () if str(value)] + work.errors = details or ([work.error] if work.error else []) + if context_count_delta > 0: + work.context_count += context_count_delta + return _WorkMutationResult(True, work) + + _, task = await self._mutate_work( + task_id, + work_id, + account_id, + user_id, + mutate, + try_finalize=True, + ) + return task + + async def requeue_work( + self, + task_id: str, + work_id: str, *, + delta: int, + account_id: str, user_id: str, - limit: int, + ) -> None: + def mutate(_task: TaskRecord, work: Optional[WorkRecord]) -> _WorkMutationResult: + if work is None or not work.is_open(): + return _WorkMutationResult(False, work) + work.state = WorkState.PENDING + work.requeue_count += max(delta, 0) + work.started_at = None + work.finished_at = None + work.updated_at = time.time() + return _WorkMutationResult(True, work) + + await self._mutate_work(task_id, work_id, account_id, user_id, mutate) + + async def record_work_error( + self, + task_id: str, + work_id: str, + error: str, + *, + account_id: str, + user_id: str, + ) -> None: + def mutate(_task: TaskRecord, work: Optional[WorkRecord]) -> _WorkMutationResult: + if work is None or work.error: + return _WorkMutationResult(False, work) + work.error = sanitize_error(error) + work.updated_at = time.time() + return _WorkMutationResult(True, work) + + await self._mutate_work(task_id, work_id, account_id, user_id, mutate) + + async def settle_work( + self, task_id: str, work_id: str, *, account_id: str, user_id: str + ) -> Optional[TaskRecord]: + def mutate(_task: TaskRecord, work: Optional[WorkRecord]) -> _WorkMutationResult: + if work is None: + return _WorkMutationResult(False, None) + changed = False + if work.is_open(): + work.state = WorkState.FAILED if work.error else WorkState.DONE + work.finished_at = time.time() + work.updated_at = work.finished_at + changed = True + return _WorkMutationResult(changed, work) + + _, task = await self._mutate_work( + task_id, + work_id, + account_id, + user_id, + mutate, + try_finalize=True, + finalize_when_unchanged=True, + ) + return task + + async def has_open_work( + self, + task_id: str, + *, + exclude_work_id: Optional[str] = None, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> bool: + task = await self._load_task(task_id, account_id, user_id) + if task is None or not task.account_id or not task.user_id: + return False + with self._lock: + work_set = self._work_sets.get((task.account_id, task.user_id, task_id)) + if work_set is None or work_set.open_count == 0: + return False + if exclude_work_id is None: + return True + excluded = work_set.works.get(exclude_work_id) + return work_set.open_count > int(excluded is not None and excluded.is_open()) + + async def summarize_works( + self, + task_id: str, + *, + work_ids: Optional[Collection[str]] = None, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> WorkSummary: + works = await self._selected_works(task_id, work_ids, account_id, user_id) + return summarize_works(works) + + async def filter_cancellation_requested( + self, + task_ids: Collection[str], + *, + account_id: str, + user_id: str, + ) -> set[str]: + requested = set() + for task_id in set(task_ids): + task = await self.get_task(task_id, account_id=account_id, user_id=user_id) + if task is not None and task.status in ( + TaskStatus.CANCELLING, + TaskStatus.CANCELLED, + ): + requested.add(task_id) + return requested + + async def is_cancellation_requested( + self, + task_id: str, + *, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> bool: + task = await self.get_task(task_id, account_id=account_id, user_id=user_id) + return bool(task and task.status in (TaskStatus.CANCELLING, TaskStatus.CANCELLED)) + + async def list( + self, + account_id: Optional[str] = None, + *, + user_id: Optional[str] = None, + task_type: Optional[str] = None, + status: Optional[str] = None, + resource_id: Optional[str] = None, + limit: Optional[int] = None, before: tuple[float, str] | None = None, - prune_expired: Callable[[Dict[str, Any]], Awaitable[bool]] | None = None, - io_limiter: StoreIOLimiter | None = None, - **filters: Any, - ) -> List[Dict[str, Any]]: - """Scan durable records, retaining only the requested page in memory. - - AGFS has no indexed query API. This bounds record memory, but still - performs O(n) reads; it deliberately does not load all tasks into the tracker. + include_internal: bool = True, + q: Optional[str] = None, + ) -> List[TaskRecord]: + """List matching Tasks. + + Owner-scoped reads include durable records. Without an account, the + file backend can only return records already loaded into its cache. """ from openviking.service.task_pagination import matches - async def read(operation: str, factory: Callable[[], Awaitable[Any]]) -> Any: - if io_limiter is None: - return await factory() - return await io_limiter.run(operation, lambda: run_to_completion(factory)) + if limit is not None and limit <= 0: + return [] + if account_id and user_id and limit is not None: + return await self._list_owner_page( + account_id, + user_id=user_id, + limit=limit, + before=before, + task_type=task_type, + status=status, + resource_id=resource_id, + include_internal=include_internal, + q=q, + ) + if account_id: + users = [user_id] if user_id else await self._owner_users(account_id) + for owner_user in users: + for payload in await self._list_owner_payloads(account_id, owner_user): + try: + self._put_if_absent( + _task_from_payload(payload, account_id=account_id, user_id=owner_user) + ) + except (ValueError, TypeError, KeyError): + continue + with self._lock: + tasks = [ + self._copy_task(task, self._work_sets.get(owner)) + for owner, task in self._task_cache.items() + for account, user, _ in (owner,) + if (account_id is None or account == account_id) + and (user_id is None or user == user_id) + ] + tasks = [ + task + for task in tasks + if (before is None or (task.created_at, task.task_id) < before) + and matches( + task.to_dict(), + task_type=task_type, + status=status, + resource_id=resource_id, + include_internal=include_internal, + q=q, + ) + ] + tasks.sort(key=lambda task: (task.created_at, task.task_id), reverse=True) + return tasks if limit is None else tasks[:limit] + + async def _list_owner_page( + self, + account_id: str, + *, + user_id: str, + limit: int, + before: tuple[float, str] | None = None, + task_type: Optional[str], + status: Optional[str], + resource_id: Optional[str], + include_internal: bool, + q: Optional[str], + ) -> List[TaskRecord]: + from openviking.service.task_pagination import matches - directory = self._task_dir(account_id, user_id) try: - entries = await read("list_page_ls", lambda: self._agfs.ls(directory)) + entries = await self._ls(self._task_dir(account_id, user_id)) except (AGFSNotFoundError, FileNotFoundError): return [] - heap: list = [] - for entry in entries: - path = entry.get("path") or f"{directory}/{entry.get('name', '')}" + heap: list[tuple[float, str, TaskRecord]] = [] + for item in entries: + path = ( + item.get("path") or f"{self._task_dir(account_id, user_id)}/{item.get('name', '')}" + ) if not path.endswith(".json"): continue try: - task = json.loads( - _decode_bytes( - await read("list_page_read", lambda path=path: self._agfs.read(path)) - ) - ) - except (AGFSNotFoundError, FileNotFoundError): - continue - if task.get("account_id") != account_id or task.get("user_id") != user_id: + payload = json.loads(_decode_bytes(await self._read(path))) + record = _task_from_payload(payload, account_id=account_id, user_id=user_id) + except (AGFSNotFoundError, FileNotFoundError, ValueError, TypeError, KeyError): continue - if prune_expired is not None and await prune_expired(task): - continue - key = (float(task["created_at"]), str(task["task_id"])) + with self._lock: + owner = (account_id, user_id, record.task_id) + cached = self._task_cache.get(owner) + if cached is not None: + record = self._copy_task(cached, self._work_sets.get(owner)) + key = (record.created_at, record.task_id) if before is not None and key >= before: continue - if not matches(task, **filters): + if not matches( + record.to_dict(), + task_type=task_type, + status=status, + resource_id=resource_id, + include_internal=include_internal, + q=q, + ): continue - item = (*key, task) + candidate = (*key, record) if len(heap) < limit: - heapq.heappush(heap, item) + heapq.heappush(heap, candidate) elif key > heap[0][:2]: - heapq.heapreplace(heap, item) + heapq.heapreplace(heap, candidate) return [item[2] for item in sorted(heap, reverse=True)] - async def delete(self, task_id: str, *, account_id: str, user_id: Optional[str] = None) -> None: - """Delete a task record, succeeding if it has already been removed.""" - if not user_id: + async def delete_tasks( + self, + task_id: Optional[str] = None, + *, + account_id: str, + user_id: Optional[str] = None, + allow_active_types: Collection[str] = (), + ) -> List[str]: + """Delete one Task or all Tasks in an owner scope.""" + if task_id is not None: + if not user_id: + raise ValueError("Deleting one Task requires user_id") + owner = (account_id, user_id, task_id) + async with self._record_locks.acquire(owner): + task = await self._load_task(task_id, account_id, user_id) + if task is not None: + self._validate_deletable_task( + task, self._work_sets.get(owner), allow_active_types + ) + deleted = await self._delete_task_unlocked(owner, task is not None) + return [task_id] if deleted else [] + + return await self._delete_owner_tasks( + account_id, + user_id=user_id, + allow_active_types=allow_active_types, + ) + + async def _delete_owner_tasks( + self, + account_id: str, + *, + user_id: Optional[str], + allow_active_types: Collection[str], + ) -> List[str]: + tasks = await self.list(account_id, user_id=user_id, limit=None) + deleted: list[str] = [] + for task in tasks: + if task.user_id: + deleted.extend( + await self.delete_tasks( + task.task_id, + account_id=account_id, + user_id=task.user_id, + allow_active_types=allow_active_types, + ) + ) + self._ensured_task_dirs = { + owner + for owner in self._ensured_task_dirs + if owner[0] != account_id or (user_id is not None and owner[1] != user_id) + } + return sorted(deleted) + + @staticmethod + def _validate_deletable_task( + task: TaskRecord, + work_set: Optional[_WorkSet], + allow_active_types: Collection[str], + ) -> None: + allowed = set(allow_active_types) + if task.status in ACTIVE_TASK_STATUSES and task.task_type not in allowed: + raise RuntimeError( + f"Cannot delete active task record: {task.task_id}({task.task_type})" + ) + if task.task_type in allowed: return + if work_set is not None and work_set.open_count: + raise RuntimeError(f"Cannot delete active task record: {task.task_id}") + + async def _delete_task_unlocked( + self, + owner: tuple[str, str, str], + task_exists: bool, + ) -> bool: + account_id, user_id, task_id = owner try: - await self._agfs.rm( - self._task_path(account_id, user_id, task_id), - force=True, - auto_pathlock=False, + await self._rm(self._task_path(account_id, user_id, task_id)) + except (AGFSNotFoundError, FileNotFoundError): + if not task_exists: + return False + with self._lock: + self._task_cache.pop(owner, None) + self._work_sets.pop(owner, None) + return task_exists + + async def delete_expired( + self, policy: TaskRetentionPolicy, *, now: float, limit: int + ) -> DeleteExpiredResult: + if limit <= 0: + return DeleteExpiredResult() + deleted: list[DeletedTask] = [] + scanned = 0 + scan_complete = False + async with self._cleanup_lock.acquire("scan"): + if self._cleanup_scan is None: + self._cleanup_scan = self._iter_task_payloads() + for _ in range(limit): + try: + payload = await anext(self._cleanup_scan) + except StopAsyncIteration: + self._cleanup_scan = None + scan_complete = True + break + scanned += 1 + if payload is None: + continue + try: + expired = policy.is_expired( + str(payload.get("status", "")), + float(payload.get("updated_at", 0.0)), + now, + ) + except (TypeError, ValueError): + expired = False + if not expired: + continue + owner = ( + str(payload.get("account_id") or ""), + str(payload.get("user_id") or ""), + str(payload.get("task_id") or ""), + ) + if not all(owner): + continue + account_id, user_id, task_id = owner + async with self._record_locks.acquire(owner): + task = await self._load_task(task_id, account_id, user_id) + work_set = self._work_sets.get(owner) + if task is None or (work_set is not None and work_set.open_count): + continue + if not policy.is_expired(task.status.value, task.updated_at, now): + continue + try: + await self._delete_task_unlocked(owner, task_exists=True) + except BaseException: + self._cleanup_scan = None + raise + deleted.append(DeletedTask(task_id, account_id, user_id)) + return DeleteExpiredResult(tuple(deleted), scanned, scan_complete) + + def snapshot_task_stats(self) -> TaskStatsSnapshot: + grouped: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int)) + with self._lock: + for task in self._task_cache.values(): + grouped[task.task_type][task.status.value] += 1 + return {task_type: dict(statuses) for task_type, statuses in grouped.items()} + + async def _selected_works( + self, + task_id: str, + work_ids: Optional[Collection[str]], + account_id: Optional[str], + user_id: Optional[str], + ) -> List[WorkRecord]: + task = await self._load_task(task_id, account_id, user_id) + if task is None or not task.account_id or not task.user_id: + return [] + selected = set(work_ids) if work_ids is not None else None + with self._lock: + work_set = self._work_sets.get((task.account_id, task.user_id, task_id)) + return [ + deepcopy(work) + for work in (work_set.works.values() if work_set else ()) + if selected is None or work.work_id in selected + ] + + async def _mutate_task( + self, + task_id: str, + account_id: Optional[str], + user_id: Optional[str], + mutate: Callable[[TaskRecord, _WorkSet], bool], + ) -> Optional[TaskRecord]: + task = await self._load_task(task_id, account_id, user_id) + if task is None or not task.account_id or not task.user_id: + return None + owner = (task.account_id, task.user_id, task_id) + async with self._record_locks.acquire(owner): + await self._load_task(task_id, owner[0], owner[1]) + with self._lock: + live_task = self._task_cache.get(owner) + if live_task is None: + return None + work_set = self._work_sets.setdefault(owner, _WorkSet({})) + candidate = deepcopy(live_task) + previous = deepcopy(live_task) + mutation = mutate(candidate, work_set) + if inspect.isawaitable(mutation): + raise TypeError("File store mutations must be synchronous") + if not mutation: + return self._copy_task(live_task, work_set) + candidate.updated_at = TaskRecord.next_updated_at(candidate.updated_at) + self._append_events(candidate, previous) + snapshot = deepcopy(candidate) + await self._write_task_and_publish(snapshot, owner, candidate) + with self._lock: + return self._copy_task(candidate, self._work_sets.get(owner)) + + async def _mutate_work( + self, + task_id: str, + work_id: str, + account_id: Optional[str], + user_id: Optional[str], + mutate: Callable[[TaskRecord, Optional[WorkRecord]], _WorkMutationResult], + *, + try_finalize: bool = False, + finalize_when_unchanged: bool = False, + ) -> tuple[bool, Optional[TaskRecord]]: + task = await self._load_task(task_id, account_id, user_id) + if task is None or not task.account_id or not task.user_id: + return False, None + owner = (task.account_id, task.user_id, task_id) + async with self._record_locks.acquire(owner): + await self._load_task(task_id, owner[0], owner[1]) + with self._lock: + live_task = self._task_cache.get(owner) + if live_task is None: + return False, None + work_set = self._work_sets.setdefault(owner, _WorkSet({})) + existing = work_set.works.get(work_id) + mutation = mutate(live_task, deepcopy(existing) if existing is not None else None) + if inspect.isawaitable(mutation): + raise TypeError("File store mutations must be synchronous") + if not isinstance(mutation, _WorkMutationResult): + raise TypeError("Work mutations must return _WorkMutationResult") + old_open = existing is not None and existing.is_open() + new_open = mutation.work is not None and mutation.work.is_open() + open_count = work_set.open_count + int(new_open) - int(old_open) + + candidate_task: Optional[TaskRecord] = None + candidate_work_set: Optional[_WorkSet] = None + finalized = False + if ( + try_finalize + and open_count == 0 + and (mutation.changed or finalize_when_unchanged) + ): + candidate_task = deepcopy(live_task) + previous = deepcopy(live_task) + works = dict(work_set.works) + if mutation.changed: + if mutation.work is None: + works.pop(work_id, None) + else: + works[work_id] = mutation.work + candidate_work_set = _WorkSet(works, open_count) + finalized = self._try_finalize(candidate_task, candidate_work_set) + + if finalized: + assert candidate_task is not None + assert candidate_work_set is not None + candidate_task.updated_at = TaskRecord.next_updated_at( + candidate_task.updated_at + ) + self._append_events(candidate_task, previous) + snapshot = deepcopy(candidate_task) + else: + if mutation.changed: + if mutation.work is None: + work_set.works.pop(work_id, None) + else: + work_set.works[work_id] = mutation.work + work_set.open_count = open_count + return True, None + assert candidate_task is not None + assert candidate_work_set is not None + await self._write_task_and_publish(snapshot, owner, candidate_task, candidate_work_set) + with self._lock: + return True, deepcopy(candidate_task) + + async def _write_task_and_publish( + self, + snapshot: TaskRecord, + owner: tuple[str, str, str], + candidate_task: TaskRecord, + candidate_work_set: Optional[_WorkSet] = None, + ) -> None: + write_task = asyncio.create_task(self._write_task(snapshot)) + cancellation: asyncio.CancelledError | None = None + try: + while not write_task.done(): + try: + await asyncio.shield(write_task) + except asyncio.CancelledError as exc: + if write_task.cancelled(): + raise + cancellation = exc + write_task.result() + with self._lock: + self._task_cache[owner] = candidate_task + if candidate_work_set is not None: + self._work_sets[owner] = candidate_work_set + if cancellation is not None: + raise cancellation + finally: + del write_task, cancellation + + def _try_finalize(self, task: TaskRecord, work_set: _WorkSet) -> bool: + if task.status not in ACTIVE_TASK_STATUSES or work_set.open_count: + return False + if task.status == TaskStatus.CANCELLING: + task.status = TaskStatus.CANCELLED + else: + work_error = first_work_error(work_set.works.values()) + if work_error and task.error is None: + task.error = sanitize_error(work_error) + if task.error is not None: + task.status = TaskStatus.FAILED + elif task.result is not None: + task.status = TaskStatus.COMPLETED + else: + task.error = "Task operation interrupted before its result was persisted" + task.status = TaskStatus.FAILED + if isinstance(task.result, dict): + # QueueFS reconstructs only unacknowledged Work after restart. For this + # file backend, terminal Work history and its errors may therefore be + # absent from the final projection until Work itself becomes durable. + summary = summarize_works(work_set.works.values()) + projected = { + queue_name: _queue_result(summary.queues.get(queue_name)) + for queue_name in ("Semantic", "Embedding") + } + has_projected_work = any( + work.queue_name in projected for work in work_set.works.values() + ) + if has_projected_work: + task.result["queue_status"] = projected + for queue_name, field_name in ( + ("Semantic", "semantic_status"), + ("Embedding", "vector_status"), + ): + if task.result.get(field_name) == "queued": + task.result[field_name] = ( + "failed" + if projected.get(queue_name, {}).get("error_count", 0) + else "complete" + ) + task.stage = task.status.value + task.auth = {} + task.processing_seconds = processing_seconds(work_set.works.values()) + return True + + def _copy_task(self, task: TaskRecord, work_set: Optional[_WorkSet]) -> TaskRecord: + copied = deepcopy(task) + if copied.status in ACTIVE_TASK_STATUSES and work_set is not None: + copied.processing_seconds = processing_seconds(work_set.works.values()) + return copied + + async def _load_task( + self, task_id: str, account_id: Optional[str], user_id: Optional[str] + ) -> Optional[TaskRecord]: + with self._lock: + owner = self._find_cached_owner(task_id, account_id, user_id) + if owner is not None: + return self._task_cache[owner] + if not account_id or not user_id: + return None + try: + payload = json.loads( + _decode_bytes(await self._read(self._task_path(account_id, user_id, task_id))) ) except (AGFSNotFoundError, FileNotFoundError): + return None + task = _task_from_payload(payload, account_id=account_id, user_id=user_id) + self._put_if_absent(task) + with self._lock: + return self._task_cache.get((account_id, user_id, task_id)) + + def _find_cached_owner( + self, task_id: str, account_id: Optional[str], user_id: Optional[str] + ) -> Optional[tuple[str, str, str]]: + if account_id and user_id: + owner = (account_id, user_id, task_id) + return owner if owner in self._task_cache else None + return max( + ( + owner + for owner in self._task_cache + for account, user, cached_id in (owner,) + if cached_id == task_id + and (account_id is None or account == account_id) + and (user_id is None or user == user_id) + ), + key=lambda owner: self._task_cache[owner].updated_at, + default=None, + ) + + def _put_if_absent(self, task: TaskRecord) -> None: + if not task.account_id or not task.user_id: return + owner = (task.account_id, task.user_id, task.task_id) + with self._lock: + self._task_cache.setdefault(owner, task) + self._work_sets.setdefault(owner, _WorkSet({})) - async def _write_task(self, task: Any) -> None: - account_id = getattr(task, "account_id", None) - user_id = getattr(task, "user_id", None) - if not account_id or not user_id: + @staticmethod + def _append_events(task: TaskRecord, previous: Optional[TaskRecord]) -> None: + from datetime import datetime, timezone + + recorded_at = datetime.now(timezone.utc).isoformat() + if previous is None: + kinds = ["created"] + else: + kinds = [] + if task.error is not None and previous.error is None: + kinds.append("error_recorded") + if task.status == TaskStatus.CANCELLED and previous.status != TaskStatus.CANCELLING: + task.execution_events = append_task_event( + task.execution_events, + kind="status_changed", + status=TaskStatus.CANCELLING.value, + stage=sanitize_error(previous.stage) if previous.stage else None, + recorded_at=recorded_at, + ) + if task.status != previous.status: + kinds.append("status_changed") + if task.stage != previous.stage and task.status not in TERMINAL_TASK_STATUSES: + kinds.append("stage_changed") + + stage = previous.stage if previous and task.status in TERMINAL_TASK_STATUSES else task.stage + for kind in kinds: + task.execution_events = append_task_event( + task.execution_events, + kind=kind, + status=task.status.value, + stage=sanitize_error(stage) if stage else None, + error=sanitize_error(task.error) + if kind == "error_recorded" and task.error + else None, + recorded_at=recorded_at, + ) + + async def _read(self, path: str) -> Any: + return await self._io.run("read", lambda: run_to_completion(lambda: self._agfs.read(path))) + + async def _ls(self, path: str) -> Any: + return await self._io.run("ls", lambda: run_to_completion(lambda: self._agfs.ls(path))) + + async def _rm(self, path: str) -> None: + await self._io.run( + "rm", + lambda: run_to_completion(lambda: self._agfs.rm(path, force=True, auto_pathlock=False)), + ) + + async def _write_task(self, task: TaskRecord) -> None: + if not task.account_id or not task.user_id: raise ValueError("PersistentTaskStore requires account_id and user_id") - await self._ensure_task_dir(account_id, user_id) - path = self._task_path(account_id, user_id, task.task_id) - payload = json.dumps(_task_to_payload(task), ensure_ascii=False).encode("utf-8") + await self._ensure_task_dir(task.account_id, task.user_id) + payload = json.dumps(_task_to_payload(task), ensure_ascii=False).encode() + path = self._task_path(task.account_id, task.user_id, task.task_id) try: - await self._write_task_payload(path, payload) + await self._io.run( + "write", + lambda: run_to_completion( + lambda: self._agfs.write(path, payload, auto_pathlock=False) + ), + ) except (AGFSNotFoundError, FileNotFoundError): - self._ensured_task_dirs.discard((account_id, user_id)) - await self._ensure_task_dir(account_id, user_id) - await self._write_task_payload(path, payload) + self._ensured_task_dirs.discard((task.account_id, task.user_id)) + await self._ensure_task_dir(task.account_id, task.user_id) + await self._io.run( + "write", + lambda: run_to_completion( + lambda: self._agfs.write(path, payload, auto_pathlock=False) + ), + ) async def _ensure_task_dir(self, account_id: str, user_id: str) -> None: - cache_key = (account_id, user_id) - if cache_key in self._ensured_task_dirs: + key = (account_id, user_id) + if key in self._ensured_task_dirs: return - - seen_paths: set[str] = set() for path in self._task_dir_chain(account_id, user_id): - if path in seen_paths: + try: + await self._io.run( + "mkdir", + lambda path=path: run_to_completion(lambda: self._agfs.mkdir(path)), + ) + except AGFSAlreadyExistsError: + pass + except Exception as exc: + if "already exists" not in str(exc).lower(): + raise + self._ensured_task_dirs.add(key) + + async def _owner_users(self, account_id: str) -> List[str]: + try: + return [ + str(item["name"]) + for item in await self._ls(self._task_root_dir(account_id)) + if item.get("name") + ] + except (AGFSNotFoundError, FileNotFoundError): + return [] + + async def _list_owner_payloads(self, account_id: str, user_id: str) -> List[Dict[str, Any]]: + try: + entries = await self._ls(self._task_dir(account_id, user_id)) + except (AGFSNotFoundError, FileNotFoundError): + return [] + payloads = [] + for item in entries: + path = ( + item.get("path") or f"{self._task_dir(account_id, user_id)}/{item.get('name', '')}" + ) + if not path.endswith(".json"): continue - seen_paths.add(path) - await self._mkdir_if_missing(path) - self._ensured_task_dirs.add(cache_key) + try: + value = json.loads(_decode_bytes(await self._read(path))) + if isinstance(value, dict): + payloads.append(value) + except (AGFSNotFoundError, FileNotFoundError, ValueError, TypeError): + continue + return payloads + + async def _iter_task_payloads(self) -> AsyncIterator[Optional[Dict[str, Any]]]: + seen = set() + with self._lock: + cached = list(self._task_cache.items()) + for key, task in cached: + seen.add(key) + yield _task_to_payload(task) + try: + accounts = await self._ls(self.ROOT_PREFIX) + except (AGFSNotFoundError, FileNotFoundError): + return + for account in accounts: + raw_account_id = account.get("name") + if not raw_account_id: + continue + account_id = str(raw_account_id) + for user_id in await self._owner_users(account_id): + try: + entries = await self._ls(self._task_dir(account_id, user_id)) + except (AGFSNotFoundError, FileNotFoundError): + continue + for item in entries: + path = item.get("path") or ( + f"{self._task_dir(account_id, user_id)}/{item.get('name', '')}" + ) + if not path.endswith(".json"): + continue + task_id = path.rsplit("/", 1)[-1][:-5] + if (account_id, user_id, task_id) in seen: + continue + try: + payload = json.loads(_decode_bytes(await self._read(path))) + except ( + AGFSNotFoundError, + FileNotFoundError, + ValueError, + TypeError, + ): + yield None + continue + yield payload if isinstance(payload, dict) else None def _task_dir_chain(self, account_id: str, user_id: str) -> tuple[str, ...]: return ( - self._account_dir(account_id), + f"{self.ROOT_PREFIX}/{account_id}", self._system_dir(account_id), self._task_root_dir(account_id), self._task_dir(account_id, user_id), ) - async def _write_task_payload(self, path: str, payload: bytes) -> None: - # TaskTracker is the owner of task mutations and serializes updates per - # task. PersistentTaskStore does not implement store-level revision/CAS, - # so AGFS pathlock cannot make multiple independent writers correct; it - # only adds one storage lock around every task lifecycle write. Avoid - # that extra PathLock overhead on this internal task file. - await self._agfs.write(path, payload, auto_pathlock=False) - - async def _mkdir_if_missing(self, path: str) -> None: - try: - await self._agfs.mkdir(path) - except AGFSAlreadyExistsError: - return - except Exception as exc: - if "already exists" in str(exc).lower(): - return - raise - - def _account_dir(self, account_id: str) -> str: - return f"{self.ROOT_PREFIX}/{account_id}" - def _system_dir(self, account_id: str) -> str: if account_id == SYSTEM_TASK_ACCOUNT_ID: - return self._account_dir(account_id) - return f"{self._account_dir(account_id)}/{self.SYSTEM_DIRNAME}" + return f"{self.ROOT_PREFIX}/{account_id}" + return f"{self.ROOT_PREFIX}/{account_id}/{self.SYSTEM_DIRNAME}" def _task_root_dir(self, account_id: str) -> str: return f"{self._system_dir(account_id)}/{self.TASKS_DIRNAME}" @@ -255,13 +1539,30 @@ def _task_path(self, account_id: str, user_id: str, task_id: str) -> str: return f"{self._task_dir(account_id, user_id)}/{task_id}.json" -def _task_to_payload(task: Any) -> Dict[str, Any]: - status = getattr(task, "status", None) +def _deep_merge(target: Dict[str, Any], values: Dict[str, Any]) -> None: + for key, value in values.items(): + if isinstance(value, dict) and isinstance(target.get(key), dict): + _deep_merge(target[key], value) + else: + target[key] = value + + +def _queue_result(summary: Optional[QueueWorkSummary]) -> Dict[str, Any]: + summary = summary or QueueWorkSummary() return { - **deepcopy(getattr(task, "_extra_fields", {})), + "processed": summary.processed, + "requeue_count": summary.requeue_count, + "error_count": summary.error_count, + "errors": [{"message": error} for error in summary.errors[:MAX_QUEUE_ERROR_DETAILS]], + } + + +def _task_to_payload(task: TaskRecord) -> Dict[str, Any]: + return { + **deepcopy(task._extra_fields), "task_id": task.task_id, "task_type": task.task_type, - "status": status.value if hasattr(status, "value") else status, + "status": task.status.value, "created_at": task.created_at, "updated_at": task.updated_at, "processing_seconds": task.processing_seconds, @@ -277,7 +1578,27 @@ def _task_to_payload(task: Any) -> Dict[str, Any]: } +def _task_from_payload( + payload: Dict[str, Any], + *, + account_id: Optional[str] = None, + user_id: Optional[str] = None, +) -> TaskRecord: + known = { + item.name for item in fields(TaskRecord) if item.init and not item.name.startswith("_") + } + data = {key: deepcopy(value) for key, value in payload.items() if key in known} + data["status"] = TaskStatus(data["status"]) + if account_id is not None: + data["account_id"] = account_id + if user_id is not None: + data["user_id"] = user_id + task = TaskRecord(**data) + task._extra_fields = { + key: deepcopy(value) for key, value in payload.items() if key not in known + } + return task + + def _decode_bytes(raw: Any) -> str: - if isinstance(raw, bytes): - return raw.decode("utf-8") - return str(raw) + return raw.decode() if isinstance(raw, bytes) else str(raw) diff --git a/openviking/service/task_tracker.py b/openviking/service/task_tracker.py index 201e34366c..73ef78b589 100644 --- a/openviking/service/task_tracker.py +++ b/openviking/service/task_tracker.py @@ -1,50 +1,91 @@ # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. # SPDX-License-Identifier: AGPL-3.0 -""" -Async Task Tracker for OpenViking. +"""Coordinate Task lifecycle, queue Work, waiting, and cancellation.""" -Provides a lightweight registry for tracking background operations -(e.g. session commit with wait=false). Callers receive a task_id that can be -polled via the /tasks API to check completion status, results, or errors. - -Design decisions: - - Thread-safe (QueueManager workers run in separate threads). - - TTL-based cleanup applies to the process-local cache. - - Error messages are sanitized to avoid leaking sensitive data. -""" +from __future__ import annotations import asyncio -import math import re import threading import time from contextlib import contextmanager from contextvars import ContextVar from copy import deepcopy -from dataclasses import dataclass, field, fields -from datetime import datetime, timezone -from enum import Enum -from typing import Any, Awaitable, Callable, Dict, Iterator, List, Optional, TypeVar +from dataclasses import dataclass +from typing import ( + Any, + Awaitable, + Callable, + Collection, + Dict, + Generator, + List, + Optional, + TypeVar, +) from uuid import uuid4 -from openviking.service.task_events import ( - PROCESS_EVENT_KINDS, - TaskEventHistory, - append_task_event, +from openviking.service.task_domain import ( + ACTIVE_TASK_STATUSES, + MAX_QUEUE_ERROR_DETAILS, + SENSITIVE_TASK_KEYS, + TERMINAL_TASK_STATUSES, + TERMINAL_WORK_STATES, + DeleteExpiredResult, + QueueWorkSummary, + TaskRecord, + TaskRetentionPolicy, + TaskStatus, + TaskWorkRejected, + WorkRecord, + WorkRecoveryMode, + WorkState, + WorkSummary, + sanitize_error, + sanitize_task_result, ) -from openviking.service.task_store import TaskStore -from openviking.service.task_tracker_concurrency import ( - KeyedAsyncLockPool, - StoreIOLimiter, - run_to_completion, +from openviking.service.task_events import PROCESS_EVENT_KINDS +from openviking.service.task_store import ( + SYSTEM_TASK_ACCOUNT_ID, + SYSTEM_TASK_USER_ID, + TaskAlreadyActive, + TaskWorkStore, + WorkRegistration, ) -from openviking.service.task_work_index import QueueTaskMetadata, TaskWorkIndex +from openviking.service.task_tracker_concurrency import run_to_completion +from openviking.service.task_work_index import TaskWorkIndex from openviking_cli.exceptions import DeadlineExceededError, OpenVikingError from openviking_cli.utils.logger import get_logger logger = get_logger(__name__) _OperationResult = TypeVar("_OperationResult") +__all__ = [ + "SENSITIVE_TASK_KEYS", + "TaskAlreadyActive", + "TaskRecord", + "TaskStatus", + "TaskTracker", + "get_task_tracker", + "run_tracked_operation", + "set_task_tracker", +] + +_CANCELLABLE_TASK_TYPES = { + "add_resource", + "add_skill", + "compile", + "session_commit", + "admin_reindex", + "snapshot_restore_reindex", +} +_CANCEL_POLL_INTERVAL = 30 +_TASK_WAIT_POLL_INTERVAL = 0.05 +_PUBLIC_QUEUE_NAMES = ("Semantic", "Embedding") + +_sanitize_error = sanitize_error +_sanitize_task_result = sanitize_task_result + @dataclass class _WaitBudget: @@ -56,7 +97,9 @@ class _WaitBudget: @contextmanager -def _consume_wait_budget(task_id: str, timeout: Optional[float]) -> Iterator[Optional[float]]: +def _consume_wait_budget( + task_id: str, timeout: Optional[float] +) -> Generator[Optional[float], None, None]: """Charge only time spent waiting, shared across an inline operation.""" budget = _wait_budget.get() if budget is None or budget.task_id != task_id: @@ -70,97 +113,14 @@ def _consume_wait_budget(task_id: str, timeout: Optional[float]) -> Iterator[Opt budget.remaining = max(0.0, budget.remaining - (time.monotonic() - started)) -class _CommittedMutationCancelled(asyncio.CancelledError): - """Caller cancellation observed after store and cache commit completed.""" - - -class TaskStatus(str, Enum): - """Lifecycle states of an async task.""" - - PENDING = "pending" - RUNNING = "running" - CANCELLING = "cancelling" - COMPLETED = "completed" - FAILED = "failed" - CANCELLED = "cancelled" - - -_TERMINAL_STATUSES = (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED) -_ACTIVE_STATUSES = (TaskStatus.PENDING, TaskStatus.RUNNING, TaskStatus.CANCELLING) -_PUBLIC_QUEUE_NAMES = ("Semantic", "Embedding") - -_CANCELLABLE_TASK_TYPES = { - "add_resource", - "add_skill", - "compile", - "session_commit", - "admin_reindex", - "snapshot_restore_reindex", -} - - -@dataclass -class TaskRecord: - """Immutable snapshot of an async task.""" - - task_id: str - task_type: str # e.g. "session_commit" - status: TaskStatus = TaskStatus.PENDING - created_at: float = field(default_factory=time.time) - updated_at: float = field(default_factory=time.time) - resource_id: Optional[str] = None # e.g. session_id - account_id: Optional[str] = None - user_id: Optional[str] = None - meta: Dict[str, Any] = field(default_factory=dict) - stage: Optional[str] = None - result: Optional[Dict[str, Any]] = None - error: Optional[str] = None - auth: Dict[str, Any] = field(default_factory=dict, repr=False) - execution_events: Optional[TaskEventHistory] = None - processing_seconds: Optional[float] = None - _extra_fields: Dict[str, Any] = field( - default_factory=dict, init=False, repr=False, compare=False - ) - - def to_dict(self, *, include_events: bool = False) -> Dict[str, Any]: - """Serialize for JSON response.""" - public_meta = deepcopy(self.meta) - public_meta.pop("submission_hash", None) - public_meta.pop("submission_token", None) - return { - **({"execution_events": deepcopy(self.execution_events)} if include_events else {}), - "task_id": self.task_id, - "task_type": self.task_type, - "status": self.status.value, - "created_at": self.created_at, - "updated_at": self.updated_at, - "processing_seconds": self.processing_seconds, - "resource_id": self.resource_id, - "meta": _sanitize_task_result(public_meta), - "stage": self.stage, - "result": _sanitize_task_result(deepcopy(self.result)), - "error": self.error, - "created_at_iso": datetime.fromtimestamp(self.created_at, tz=timezone.utc).isoformat(), - "updated_at_iso": datetime.fromtimestamp(self.updated_at, tz=timezone.utc).isoformat(), - } - - -# ── Singleton ── - _instance: Optional["TaskTracker"] = None _init_lock = threading.Lock() def get_task_tracker() -> "TaskTracker": - """Get the global TaskTracker singleton installed by service storage initialization.""" + """Return the tracker installed during service storage initialization.""" with _init_lock: if _instance is None: - logger.error( - "TaskTracker accessed before service storage initialization; refusing to create " - "a separate AGFS client. Ensure OpenVikingService installs the shared tracker " - "with set_task_tracker() before task APIs are used.", - stack_info=True, - ) raise RuntimeError( "TaskTracker not initialized. OpenVikingService must install the shared " "tracker with set_task_tracker() during storage initialization." @@ -168,8 +128,7 @@ def get_task_tracker() -> "TaskTracker": return _instance -def set_task_tracker(tracker: "TaskTracker") -> None: - """Replace the global TaskTracker singleton.""" +def set_task_tracker(tracker: Optional["TaskTracker"]) -> None: global _instance with _init_lock: _instance = tracker @@ -199,7 +158,7 @@ async def run_tracked_operation( return await operation(), current.task_id tracker = get_task_tracker() - created_task: Optional[TaskRecord] = existing_task + created_task = existing_task async def create_task() -> None: nonlocal created_task @@ -211,8 +170,6 @@ async def create_task() -> None: ) try: - # Keep the committed record reachable even if the caller is cancelled - # while the store write is still running. if created_task is None: await run_to_completion(create_task) assert created_task is not None @@ -222,27 +179,30 @@ async def create_task() -> None: ) except BaseException as exc: if created_task is not None: - failed_task_id = created_task.task_id failure_message = str(exc) or type(exc).__name__ await run_to_completion( lambda: tracker.fail( - failed_task_id, + created_task.task_id, failure_message, account_id=account_id, user_id=user_id, ) ) raise - enqueued_work: list[tuple[str, str]] = [] async def execute_operation() -> _OperationResult: - tracker.register_running_task(task.task_id) + running_work_id = await tracker.register_running_task( + task.task_id, + account_id=account_id, + user_id=user_id, + ) try: with bind_task_context( task.task_id, account_id, user_id, - enqueued_work=enqueued_work, + enqueued_work=[], + work_id=running_work_id, ): result = await operation() task_result = deepcopy(result) if isinstance(result, dict) else {"result": result} @@ -257,8 +217,6 @@ async def execute_operation() -> _OperationResult: if isinstance(result, dict) and not wait: result.setdefault("task_id", task.task_id) return result - except _CommittedMutationCancelled: - raise except BaseException as exc: failure_message = str(exc) or type(exc).__name__ await run_to_completion( @@ -307,86 +265,38 @@ async def execute_operation() -> _OperationResult: _wait_budget.reset(token) -# ── Sanitization ── - -_SENSITIVE_PATTERNS = re.compile( - r"(sk-|cr_|ghp_|ntn_|xox[baprs]-|Bearer\s+)[a-zA-Z0-9._-]+", - re.IGNORECASE, -) - -_MAX_ERROR_LEN = 500 -SENSITIVE_TASK_KEYS = frozenset({"api_key", "user_key"}) - - -def _sanitize_error(error: str) -> str: - """Remove potential secrets from error messages.""" - sanitized = _SENSITIVE_PATTERNS.sub("[REDACTED]", error) - if len(sanitized) > _MAX_ERROR_LEN: - sanitized = sanitized[:_MAX_ERROR_LEN] + "...[truncated]" - return sanitized - - -def _sanitize_task_result(result: Any) -> Any: - """Remove sensitive fields from task results before exposing snapshots.""" - if isinstance(result, dict): - return { - key: _sanitize_task_result(value) - for key, value in result.items() - if key not in SENSITIVE_TASK_KEYS - } - if isinstance(result, list): - return [_sanitize_task_result(item) for item in result] - return result - - -# ── TaskTracker ── - - class TaskTracker: - """Async task tracker with persistent storage and a process-local cache. + """Thin coordinator over the semantic :class:`TaskWorkStore` contract.""" - Mutations are serialized per task across caller threads and event loops. - The thread lock only protects short accesses to immutable cache snapshots. - """ - - MAX_TASKS = 10_000 - TTL_COMPLETED = 86_400 # 24 hours - TTL_FAILED = 604_800 # 7 days - CLEANUP_INTERVAL = 300 # 5 minutes + TTL_COMPLETED = 86_400 + TTL_FAILED = 604_800 + CLEANUP_INTERVAL = 300 + CLEANUP_BATCH = 256 - def __init__(self, store: TaskStore, *, max_concurrent_store_io: int = 8) -> None: + def __init__(self, store: TaskWorkStore, *, max_concurrent_store_io: int = 8) -> None: + del max_concurrent_store_io self._store = store - self._tasks: Dict[str, TaskRecord] = {} - self._lock = threading.Lock() - # The keyed registries provide process-local ordering. A deployment with - # multiple TaskTracker writers must add store-level revision/CAS first. - self._task_locks = KeyedAsyncLockPool[str]() - self._business_locks = KeyedAsyncLockPool[tuple[str, str, str, str]]() - self._store_io = StoreIOLimiter(max_concurrent_store_io) - self._cleanup_task: Optional[asyncio.Task] = None + self._cleanup_task: Optional[asyncio.Task[Any]] = None + self._cancel_poll_task: Optional[asyncio.Task[Any]] = None self._work_index = TaskWorkIndex() - self._install_work_index_callbacks() - logger.info( - "[TaskTracker] Initialized (store=%s, max_tasks=%d)", - self._store.__class__.__name__, - self.MAX_TASKS, + self._running: ContextVar[tuple[str, str, str, bool] | None] = ContextVar( + "task_tracker_running_work", default=None ) - # ── Lifecycle ── - - def _install_work_index_callbacks(self) -> None: - self._work_index.set_callbacks( - finalize_before_ack=self._finalize_before_ack, - is_cancellation_requested=self.is_cancellation_requested, - ) + @staticmethod + def _validate_owner(account_id: str, user_id: str) -> None: + if not account_id or not user_id: + raise ValueError("Task ownership requires non-empty account_id and user_id") def attach_work_index(self, work_index: TaskWorkIndex) -> None: - """Use the QueueManager-owned index as the task lifecycle authority.""" + """Replace the process-local active-handle index.""" self._work_index = work_index - self._install_work_index_callbacks() + + @property + def work_recovery_mode(self) -> WorkRecoveryMode: + return self._store.work_recovery_mode async def restore_work_tasks(self, owners: Dict[str, tuple[str, str]]) -> List[TaskRecord]: - """Restore task records referenced by rebuilt QueueFS work.""" restored = [] for task_id, (account_id, user_id) in owners.items(): task = await self.get(task_id, account_id=account_id, user_id=user_id) @@ -394,143 +304,65 @@ async def restore_work_tasks(self, owners: Dict[str, tuple[str, str]]) -> List[T restored.append(task) return restored - async def _finalize_before_ack(self, metadata: QueueTaskMetadata) -> None: - """Persist the terminal state before QueueFS removes the last recovery message.""" - await self._finalize_task( - metadata.task_id, - account_id=metadata.account_id or None, - user_id=metadata.user_id or None, - ) - - def is_cancellation_requested(self, task_id: str) -> bool: - """Fast thread-safe status check used by queue workers.""" - with self._lock: - task = self._tasks.get(task_id) - return task is not None and task.status in ( - TaskStatus.CANCELLING, - TaskStatus.CANCELLED, - ) - def start_cleanup_loop(self) -> None: - """Start the background TTL cleanup coroutine. - - Safe to call multiple times; subsequent calls are no-ops. - Must be called from within a running event loop. - """ - if self._cleanup_task is not None and not self._cleanup_task.done(): - return - self._cleanup_task = asyncio.create_task(self._cleanup_loop()) - logger.debug("[TaskTracker] Cleanup loop started") + if self._cleanup_task is None or self._cleanup_task.done(): + self._cleanup_task = asyncio.create_task(self._cleanup_loop()) + if self._cancel_poll_task is None or self._cancel_poll_task.done(): + self._cancel_poll_task = asyncio.create_task(self._cancel_poll_loop()) def stop_cleanup_loop(self) -> None: - """Cancel the background cleanup task. Safe to call if not started.""" - if self._cleanup_task is not None and not self._cleanup_task.done(): - self._cleanup_task.cancel() - logger.debug("[TaskTracker] Cleanup loop stopped") + for task in (self._cleanup_task, self._cancel_poll_task): + if task is not None and not task.done(): + task.cancel() async def _cleanup_loop(self) -> None: while True: try: await asyncio.sleep(self.CLEANUP_INTERVAL) - await self._evict_expired() + await self.delete_expired() except asyncio.CancelledError: break except Exception: - logger.exception("[TaskTracker] Cleanup error") - - async def _evict_expired(self) -> None: - """Remove expired tasks and enforce MAX_TASKS.""" - now = time.time() - with self._lock: - expired_ids = [ - task_id for task_id, task in self._tasks.items() if self._is_expired(task, now) - ] - - evicted_count = 0 - for task_id in expired_ids: - evicted_count += await self._delete_expired_task(task_id, now) - - capacity_evicted: list[tuple[str, bool]] = [] - with self._lock: - if len(self._tasks) > self.MAX_TASKS: - sorted_tasks = sorted(self._tasks.items(), key=lambda x: x[1].created_at) - excess = len(self._tasks) - self.MAX_TASKS - for tid, task in sorted_tasks[:excess]: - self._tasks.pop(tid, None) - capacity_evicted.append((tid, task.status in _TERMINAL_STATUSES)) - - for tid, terminal in capacity_evicted: - if terminal: - self._work_index.forget_task(tid) - else: - self._work_index.forget_processing(tid) - - if evicted_count: - logger.debug("[TaskTracker] Evicted %d expired tasks", evicted_count) - - def _is_expired(self, task: TaskRecord, now: float) -> bool: - age = now - task.updated_at - if task.status in (TaskStatus.COMPLETED, TaskStatus.CANCELLED): - return age > self.TTL_COMPLETED - return task.status == TaskStatus.FAILED and age > self.TTL_FAILED - - async def _delete_expired_task(self, task_id: str, now: float) -> bool: - async with self._task_locks.acquire(task_id): - task = self._cached_task(task_id) - if task is None or not self._is_expired(task, now): - return False - account_id = task.account_id - user_id = task.user_id - if not account_id or not user_id: - logger.warning( - "[TaskTracker] Cannot delete expired ownerless task %s", - task_id, - ) - return False + logger.exception("[TaskTracker] cleanup error") + + async def _cancel_poll_loop(self) -> None: + while True: try: - await self._store_io.run( - "delete", - lambda: run_to_completion( - lambda: self._store.delete( + await asyncio.sleep(_CANCEL_POLL_INTERVAL) + for ( + account_id, + user_id, + ), task_ids in self._work_index.active_tasks_by_owner().items(): + cancelling = await self._store.filter_cancellation_requested( + task_ids, + account_id=account_id, + user_id=user_id, + ) + for task_id in cancelling: + self._work_index.cancel_active( task_id, account_id=account_id, user_id=user_id, ) - ), - ) + except asyncio.CancelledError: + break except Exception: - logger.warning( - "[TaskTracker] Failed to delete expired task %s", - task_id, - exc_info=True, - ) - return False - - with self._lock: - self._tasks.pop(task_id, None) - self._work_index.forget_task(task_id) - return True - - @staticmethod - def _matches_owner( - task: TaskRecord, - account_id: Optional[str] = None, - user_id: Optional[str] = None, - ) -> bool: - """Return True when a task belongs to the requested owner filter.""" - if account_id is not None and task.account_id != account_id: - return False - if user_id is not None and task.user_id != user_id: - return False - return True - - @staticmethod - def _validate_owner(account_id: str, user_id: str) -> None: - """Reject ownerless task creation for user-originated background work.""" - if not account_id or not user_id: - raise ValueError("Task ownership requires non-empty account_id and user_id") + logger.exception("[TaskTracker] cancel-poll error") - # ── CRUD ── + async def delete_expired(self) -> DeleteExpiredResult: + result = await self._store.delete_expired( + TaskRetentionPolicy(self.TTL_COMPLETED, self.TTL_FAILED), + now=time.time(), + limit=self.CLEANUP_BATCH, + ) + for deleted in result.deleted: + logger.info( + "[TaskTracker] deleted expired task task_id=%s account_id=%s user_id=%s", + deleted.task_id, + deleted.account_id, + deleted.user_id, + ) + return result async def create( self, @@ -542,9 +374,13 @@ async def create( task_id: Optional[str] = None, meta: Optional[Dict[str, Any]] = None, auth: Optional[Dict[str, Any]] = None, + reject_if_active: bool = False, ) -> TaskRecord: - """Register a new pending task. Returns a snapshot copy.""" self._validate_owner(account_id, user_id) + if task_id is not None: + existing = await self.get(task_id, account_id=account_id, user_id=user_id) + if existing is not None: + return existing task = TaskRecord( task_id=task_id or str(uuid4()), task_type=task_type, @@ -554,31 +390,10 @@ async def create( meta=dict(meta or {}), auth=dict(auth or {}), ) - async with self._task_locks.acquire(task.task_id): - if task_id is not None: - existing = self._cached_task(task.task_id) - if existing is not None: - if not self._matches_owner(existing, task.account_id, task.user_id): - raise ValueError( - f"Task ID already belongs to another owner: {task.task_id}" - ) - return self._copy(existing) - existing = await self._load_from_store( - task.task_id, - task.account_id or "", - task.user_id, - ) - if existing is not None: - self._publish_task(existing) - return self._copy(existing) - await self._persist_and_publish("create", task) - logger.debug( - "[TaskTracker] Created task %s type=%s resource=%s", - task.task_id, - task.task_type, - task.resource_id, + stored = await run_to_completion( + lambda: self._store.create(task, reject_if_active=reject_if_active) ) - return self._copy(task) + return self._public(stored or task) async def create_if_no_running( self, @@ -588,42 +403,16 @@ async def create_if_no_running( account_id: str, user_id: str, ) -> Optional[TaskRecord]: - """Atomically check for running tasks and create a new one if none exist. - - Returns TaskRecord on success, None if a running task already exists. - This eliminates the race condition between has_running() and create(). - """ - self._validate_owner(account_id, user_id) - business_key = (account_id, user_id, task_type, resource_id) - async with self._business_locks.acquire(business_key): - self._merge_loaded_tasks(await self._load_all_from_store(account_id, user_id)) - - tasks = self._cache_snapshot() - has_active = any( - t.task_type == task_type - and t.resource_id == resource_id - and self._matches_owner(t, account_id, user_id) - and t.status in _ACTIVE_STATUSES - for t in tasks - ) - if has_active: - return None - task = TaskRecord( - task_id=str(uuid4()), - task_type=task_type, - resource_id=resource_id, + try: + return await self.create( + task_type, + resource_id, account_id=account_id, user_id=user_id, + reject_if_active=True, ) - async with self._task_locks.acquire(task.task_id): - await self._persist_and_publish("create", task) - logger.debug( - "[TaskTracker] Created task %s type=%s resource=%s", - task.task_id, - task_type, - resource_id, - ) - return self._copy(task) + except TaskAlreadyActive: + return None async def start( self, @@ -632,16 +421,7 @@ async def start( user_id: Optional[str] = None, stage: Optional[str] = None, ) -> None: - """Transition task to RUNNING.""" - async with self._task_locks.acquire(task_id): - task = await self._load_for_update(task_id, account_id, user_id) - if task and task.status in (TaskStatus.PENDING, TaskStatus.RUNNING): - updated = deepcopy(task) - updated.status = TaskStatus.RUNNING - if stage is not None: - updated.stage = stage - updated.updated_at = self._next_updated_at(task) - await self._persist_and_publish("update", updated, previous=task) + await self._store.start_task(task_id, stage=stage, account_id=account_id, user_id=user_id) async def update_stage( self, @@ -652,16 +432,9 @@ async def update_stage( *, meta: Optional[Dict[str, Any]] = None, ) -> None: - """Update task progress without changing its lifecycle status.""" - async with self._task_locks.acquire(task_id): - task = await self._load_for_update(task_id, account_id, user_id) - if task and task.status in _ACTIVE_STATUSES: - updated = deepcopy(task) - updated.stage = stage - if meta: - updated.meta.update(deepcopy(meta)) - updated.updated_at = self._next_updated_at(task) - await self._persist_and_publish("update", updated, previous=task) + await self._store.update_stage( + task_id, stage, meta=meta, account_id=account_id, user_id=user_id + ) async def update_task_auth( self, @@ -671,15 +444,8 @@ async def update_task_auth( account_id: str, user_id: str, ) -> None: - """Persist private state required to resume an active task.""" self._validate_owner(account_id, user_id) - async with self._task_locks.acquire(task_id): - task = await self._load_for_update(task_id, account_id, user_id) - if task and task.status in _ACTIVE_STATUSES: - updated = deepcopy(task) - updated.auth.update(deepcopy(values)) - updated.updated_at = self._next_updated_at(task) - await self._persist_and_publish("update", updated, previous=task) + await self._store.update_auth(task_id, values, account_id=account_id, user_id=user_id) async def record_feishu_response( self, @@ -689,16 +455,13 @@ async def record_feishu_response( account_id: str, user_id: str, ) -> None: - """Persist a child submission before polling so a source retry can resume it.""" - - async with self._task_locks.acquire(task_id): - task = await self._load_for_update(task_id, account_id, user_id) - if task is None or task.status not in (TaskStatus.PENDING, TaskStatus.RUNNING): - raise ValueError("Feishu source task is no longer active") - updated = deepcopy(task) - updated.meta.setdefault("feishu_responses", {})[entry] = response_id - updated.updated_at = self._next_updated_at(task) - await self._persist_and_publish("update", updated, previous=task) + self._validate_owner(account_id, user_id) + await self._store.merge_meta( + task_id, + {"feishu_responses": {entry: response_id}}, + account_id=account_id, + user_id=user_id, + ) async def complete( self, @@ -709,13 +472,14 @@ async def complete( *, resource_id: Optional[str] = None, ) -> None: - """Record successful completion and finalize after owned work settles.""" await self._record_outcome( task_id, - account_id, - user_id, result=result, + error=None, resource_id=resource_id, + cancelled=False, + account_id=account_id, + user_id=user_id, ) async def fail( @@ -727,13 +491,14 @@ async def fail( *, result: Optional[Dict[str, Any]] = None, ) -> None: - """Record failure and optional structured metadata, then finalize.""" await self._record_outcome( task_id, - account_id, - user_id, result=result, - error=error, + error=sanitize_error(error), + resource_id=None, + cancelled=False, + account_id=account_id, + user_id=user_id, ) async def mark_cancelled( @@ -742,12 +507,14 @@ async def mark_cancelled( account_id: Optional[str] = None, user_id: Optional[str] = None, ) -> None: - """Record an externally observed cancellation without requesting local cancellation.""" await self._record_outcome( task_id, - account_id, - user_id, - terminal_status=TaskStatus.CANCELLED, + result=None, + error=None, + resource_id=None, + cancelled=True, + account_id=account_id, + user_id=user_id, ) async def record_cancelled( @@ -756,74 +523,30 @@ async def record_cancelled( account_id: Optional[str] = None, user_id: Optional[str] = None, ) -> None: - """Record cancellation reported by the task owner without cancelling its worker.""" - async with self._task_locks.acquire(task_id): - task = await self._load_for_update(task_id, account_id, user_id) - if task and task.status in _ACTIVE_STATUSES: - updated = deepcopy(task) - updated.status = TaskStatus.CANCELLING - updated.updated_at = self._next_updated_at(task) - await self._persist_and_publish("update", updated, previous=task) - await self._finalize_task(task_id, account_id, user_id) + await self.mark_cancelled(task_id, account_id, user_id) async def _record_outcome( self, task_id: str, + *, + result: Optional[Dict[str, Any]], + error: Optional[str], + resource_id: Optional[str], + cancelled: bool, account_id: Optional[str], user_id: Optional[str], - *, - result: Optional[Dict[str, Any]] = None, - error: Optional[str] = None, - resource_id: Optional[str] = None, - terminal_status: Optional[TaskStatus] = None, ) -> None: - cancellation: asyncio.CancelledError | None = None - outcome_persisted = False - async with self._task_locks.acquire(task_id): - task = await self._load_for_update(task_id, account_id, user_id) - if task and task.status in (TaskStatus.PENDING, TaskStatus.RUNNING): - updated = deepcopy(task) - if result is not None: - updated.result = deepcopy(result) - if resource_id is not None: - updated.resource_id = resource_id - if error is not None and updated.error is None: - updated.error = _sanitize_error(error) - work_error = self._work_index.failure(task_id) - if work_error and updated.error is None: - updated.error = _sanitize_error(work_error) - if terminal_status is not None: - updated.status = terminal_status - updated.stage = terminal_status.value - updated.updated_at = self._next_updated_at(task) - enqueue_markers = updated.auth.get("enqueue_complete_work_ids") - updated.auth = ( - {"enqueue_complete_work_ids": deepcopy(enqueue_markers)} - if enqueue_markers - else {} - ) - try: - await self._persist_and_publish("update", updated, previous=task) - except _CommittedMutationCancelled as exc: - cancellation = exc - outcome_persisted = True - else: - outcome_persisted = True - if outcome_persisted: - try: - await run_to_completion( - lambda: self._finalize_task( - task_id, - account_id=account_id, - user_id=user_id, - ) - ) - except asyncio.CancelledError as exc: - cancellation = cancellation or exc - else: - await self._finalize_task(task_id, account_id=account_id, user_id=user_id) - if cancellation is not None: - raise cancellation + await run_to_completion( + lambda: self._store.record_outcome( + task_id, + result=result, + error=error, + resource_id=resource_id, + cancelled=cancelled, + account_id=account_id, + user_id=user_id, + ) + ) async def cancel( self, @@ -831,114 +554,148 @@ async def cancel( account_id: Optional[str] = None, user_id: Optional[str] = None, ) -> Optional[TaskRecord]: - """Request cooperative cancellation and return the current task snapshot.""" - return await self._cancel_task(task_id, account_id, user_id) + task = await self._request_cancellation( + task_id, + account_id=account_id, + user_id=user_id, + allow_terminal=False, + ) + if task is not None: + self._work_index.cancel_active( + task_id, account_id=task.account_id, user_id=task.user_id + ) + return task + + async def cancel_skill_update_for_rollback( + self, task_id: str, *, account_id: str, user_id: str + ) -> Optional[TaskRecord]: + task = await self.get(task_id, account_id=account_id, user_id=user_id) + if task is not None and task.task_type != "add_skill": + raise ValueError("Only Skill processing can be cancelled for package rollback") + cancelled = await self._request_cancellation( + task_id, + account_id=account_id, + user_id=user_id, + allow_terminal=True, + ) + if cancelled is not None: + self._work_index.cancel_active(task_id, account_id=account_id, user_id=user_id) + return cancelled - async def _cancel_task( + async def _request_cancellation( self, task_id: str, + *, account_id: Optional[str], user_id: Optional[str], - *, - rollback_skill_update: bool = False, + allow_terminal: bool, ) -> Optional[TaskRecord]: - cancellation: asyncio.CancelledError | None = None - cancellation_persisted = False - async with self._task_locks.acquire(task_id): - task = await self._load_for_update(task_id, account_id, user_id) - if task is None: - return None - if rollback_skill_update and task.task_type != "add_skill": - raise ValueError("Only Skill processing can be cancelled for package rollback") - if task.status == TaskStatus.CANCELLED: - return self._copy(task) - if task.status != TaskStatus.CANCELLING: - if task.task_type not in _CANCELLABLE_TASK_TYPES: - raise ValueError(f"Task type '{task.task_type}' does not support cancellation") - if task.status in _TERMINAL_STATUSES and not rollback_skill_update: - raise ValueError(f"Task is already {task.status.value}") - - updated = deepcopy(task) - updated.status = TaskStatus.CANCELLING - updated.updated_at = self._next_updated_at(task) - try: - await self._persist_and_publish("update", updated, previous=task) - except _CommittedMutationCancelled as exc: - cancellation = exc - cancellation_persisted = True - else: - cancellation_persisted = True - - async def finish_cancellation() -> None: - self._work_index.cancel_active(task_id) - await self._finalize_task( + task = await run_to_completion( + lambda: self._store.request_cancellation( task_id, + cancellable_types=_CANCELLABLE_TASK_TYPES, + allow_terminal=allow_terminal, account_id=account_id, user_id=user_id, ) + ) + return self._public(task) if task is not None else None - if cancellation_persisted: + async def register_running_task( + self, + task_id: str, + *, + work_id: Optional[str] = None, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> str: + active_task = asyncio.current_task() + task = await self._store.get_task(task_id, account_id=account_id, user_id=user_id) + if active_task is None or task is None or task.is_terminal(): + raise TaskWorkRejected(f"Task {task_id} is unavailable") + if not task.account_id or not task.user_id: + raise TaskWorkRejected(f"Task {task_id} has no owner") + owner_account_id = task.account_id + owner_user_id = task.user_id + inline = work_id is None + running_work_id = work_id or f"inline:{uuid4()}" + if inline: + registration = await self.register_work( + task_id, + running_work_id, + "Inline", + account_id=owner_account_id, + user_id=owner_user_id, + ) + if registration is WorkRegistration.REJECTED: + raise TaskWorkRejected(f"Task {task_id} rejects inline work") try: - await run_to_completion(finish_cancellation) - except asyncio.CancelledError as exc: - cancellation = cancellation or exc - else: - await finish_cancellation() - task = self._cached_task(task_id) - if task is None or not self._matches_owner(task, account_id, user_id): - return None - if cancellation is not None: - raise cancellation - return self._copy(task) - - async def cancel_skill_update_for_rollback( - self, task_id: str, *, account_id: str, user_id: str - ) -> Optional[TaskRecord]: - """Prevent even a late queue replay from writing into a restored Skill. + await self.start_work( + task_id, + running_work_id, + account_id=owner_account_id, + user_id=owner_user_id, + ) + except BaseException: + if registration is WorkRegistration.CREATED: + await run_to_completion( + lambda: self.discard_work( + task_id, + running_work_id, + account_id=owner_account_id, + user_id=owner_user_id, + ) + ) + raise + self._work_index.register_active( + task_id, + active_task, + account_id=owner_account_id, + user_id=owner_user_id, + ) + self._running.set((running_work_id, owner_account_id, owner_user_id, inline)) + if await self.is_cancellation_requested( + task_id, account_id=owner_account_id, user_id=owner_user_id + ): + active_task.cancel() + return running_work_id - A last ACK can fail after the task appears finished. Persist cancellation - for this rolled-back update even if it has reached a terminal state. - Normal public cancellation keeps its existing terminal-state rules. - """ - return await self._cancel_task(task_id, account_id, user_id, rollback_skill_update=True) + async def unregister_running_task(self, task_id: str) -> None: + active_task = asyncio.current_task() + state = self._running.get() + self._running.set(None) + if active_task is not None and state is not None: + _, account_id, user_id, _ = state + self._work_index.unregister_active( + task_id, + active_task, + account_id=account_id, + user_id=user_id, + ) + if state is not None and state[3]: + work_id, account_id, user_id, _inline = state + await run_to_completion( + lambda: self.mark_work_done( + task_id, + work_id, + account_id=account_id, + user_id=user_id, + ) + ) - async def _finalize_task( + async def get( self, task_id: str, account_id: Optional[str] = None, user_id: Optional[str] = None, - ) -> None: - """Apply the task's terminal outcome after all owned work is settled.""" - if self._work_index.has_work(task_id): - return - async with self._task_locks.acquire(task_id): - task = await self._load_for_update(task_id, account_id, user_id) - if task and task.status in _ACTIVE_STATUSES: - if self._work_index.has_work(task_id): - return - updated = deepcopy(task) - if updated.status != TaskStatus.CANCELLING: - work_error = self._work_index.failure(task_id) - if work_error and updated.error is None: - updated.error = _sanitize_error(work_error) - if updated.status == TaskStatus.CANCELLING: - updated.status = TaskStatus.CANCELLED - elif updated.error is not None: - updated.status = TaskStatus.FAILED - elif updated.result is not None: - updated.status = TaskStatus.COMPLETED - else: - updated.error = "Task operation interrupted before its result was persisted" - updated.status = TaskStatus.FAILED - work_status = self._work_index.queue_status(task_id) - if updated.result is not None and work_status: - updated.result["queue_status"] = self._public_queue_status(work_status) - updated.stage = updated.status.value - updated.updated_at = self._next_updated_at(task) - updated.auth = {} - await self._persist_and_publish("update", updated, previous=task) - self._work_index.clear_failure(task_id) - logger.info("[TaskTracker] Task %s %s", task_id, updated.status.value) + ) -> Optional[TaskRecord]: + task = await self._store.get_task(task_id, account_id=account_id, user_id=user_id) + return self._public(task) if task is not None else None + + async def get_task_auth(self, task_id: str, *, account_id: str, user_id: str) -> Dict[str, Any]: + self._validate_owner(account_id, user_id) + task = await self._store.get_task(task_id, account_id=account_id, user_id=user_id) + return deepcopy(task.auth) if task is not None else {} async def wait( self, @@ -946,21 +703,15 @@ async def wait( account_id: Optional[str] = None, user_id: Optional[str] = None, timeout: Optional[float] = None, - poll_interval: float = 0.05, + poll_interval: float = _TASK_WAIT_POLL_INTERVAL, ) -> TaskRecord: - """Wait for one task's terminal state without changing its lifecycle.""" - with _consume_wait_budget(task_id, timeout) as remaining: deadline = None if remaining is None else time.monotonic() + remaining while True: - lookup = self.get(task_id, account_id=account_id, user_id=user_id) - if remaining is not None and self._cached_task(task_id) is None: - task = await asyncio.wait_for(lookup, max(0.0, remaining)) - else: - task = await lookup + task = await self.get(task_id, account_id=account_id, user_id=user_id) if task is None: raise KeyError(f"Task not found: {task_id}") - if task.status in _TERMINAL_STATUSES: + if task.status in TERMINAL_TASK_STATUSES: return task remaining = None if deadline is None else deadline - time.monotonic() if remaining is not None and remaining <= 0: @@ -969,554 +720,592 @@ async def wait( poll_interval if remaining is None else min(poll_interval, remaining) ) - async def get_task_auth( + async def list_tasks( self, - task_id: str, + task_type: Optional[str] = None, + status: Optional[str] = None, + resource_id: Optional[str] = None, + limit: Optional[int] = 50, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + include_internal: bool = True, *, - account_id: str, - user_id: str, - ) -> Dict[str, Any]: - """Load task-owned authentication excluded from public snapshots.""" - self._validate_owner(account_id, user_id) + before: tuple[float, str] | None = None, + q: Optional[str] = None, + ) -> List[TaskRecord]: + from openviking.service.task_pagination import matches - task = self._cached_task(task_id) - if task is None: - task = await self._load_from_store(task_id, account_id, user_id) - if task is not None: - self._publish_task(task) - if task is None or not self._matches_owner(task, account_id, user_id): - return {} - return deepcopy(task.auth) - - def forget_account_tasks(self, account_id: str) -> None: - """Invalidate snapshots after executions settle and account storage is removed.""" - with self._lock: - task_ids = [ - task_id for task_id, task in self._tasks.items() if task.account_id == account_id - ] - for task_id in task_ids: - del self._tasks[task_id] - for task_id in task_ids: - self._work_index.forget_task(task_id) - - async def delete_user_tasks(self, account_id: str, user_id: str) -> int: - """Delete terminal task records for one user from storage and cache.""" - self._validate_owner(account_id, user_id) - self._merge_loaded_tasks(await self._load_all_from_store(account_id, user_id)) - tasks = [ - task - for task in self._cache_snapshot() - if self._matches_owner(task, account_id, user_id) - ] - active = [task for task in tasks if task.status in _ACTIVE_STATUSES] - if active: - raise RuntimeError( - "Cannot delete active task records: " - + ", ".join(f"{task.task_id}({task.task_type})" for task in active) + records = await self._store.list( + account_id, + user_id=user_id, + task_type=task_type, + status=status, + resource_id=resource_id, + limit=limit, + before=before, + include_internal=include_internal, + q=q, + ) + visible = { + task.task_id: task + for task in records + if (before is None or (task.created_at, task.task_id) < before) + and matches( + task.to_dict(), + task_type=task_type, + status=status, + resource_id=resource_id, + include_internal=include_internal, + q=q, ) + } + ordered = sorted( + visible.values(), + key=lambda task: (task.created_at, task.task_id), + reverse=True, + ) + selected = ordered if limit is None else ordered[:limit] + return [self._public(task) for task in selected] - deleted = 0 - for task in tasks: - async with self._task_locks.acquire(task.task_id): - current = self._cached_task(task.task_id) - if current is None or not self._matches_owner(current, account_id, user_id): - continue - if current.status in _ACTIVE_STATUSES or self._work_index.has_work(task.task_id): - raise RuntimeError( - f"Cannot delete active task record: {task.task_id}({task.task_type})" - ) - await self._store_io.run( - "delete", - lambda task_id=task.task_id: run_to_completion( - lambda: self._store.delete( - task_id, - account_id=account_id, - user_id=user_id, - ) - ), - ) - with self._lock: - self._tasks.pop(task.task_id, None) - self._work_index.forget_task(task.task_id) - deleted += 1 - return deleted - - async def wait_for_work( + async def list_root_tasks( self, - task_id: str, - work_ids: List[str], *, - timeout: Optional[float] = None, - stop_waiting: Optional[Any] = None, - ) -> None: - """Wait only for the supplied logical Work IDs to become terminal.""" - with _consume_wait_budget(task_id, timeout) as remaining: - await self._work_index.wait_for_work( - task_id, - work_ids, - timeout=remaining, - stop_waiting=stop_waiting, - ) + account_id: str, + user_id: str, + task_type: Optional[str] = None, + status: Optional[str] = None, + resource_id: Optional[str] = None, + limit: int = 50, + before: tuple[float, str] | None = None, + include_internal: bool = True, + q: Optional[str] = None, + ) -> List[TaskRecord]: + """Build the ROOT compatibility view over the FileStore limitations. - async def wait_for_task_work( + The logical meaning is "all Tasks". The file-backed implementation + cannot enumerate every account without an owner, so include the + durable system/current-owner scopes and its cache-only global view. + """ + + source_limit = limit + 1 + records = await asyncio.gather( + self.list_tasks( + task_type=task_type, + status=status, + resource_id=resource_id, + limit=source_limit, + account_id=SYSTEM_TASK_ACCOUNT_ID, + user_id=SYSTEM_TASK_USER_ID, + before=before, + include_internal=include_internal, + q=q, + ), + self.list_tasks( + task_type=task_type, + status=status, + resource_id=resource_id, + limit=source_limit, + account_id=account_id, + user_id=user_id, + before=before, + include_internal=include_internal, + q=q, + ), + self.list_tasks( + task_type=task_type, + status=status, + resource_id=resource_id, + limit=None, + before=before, + include_internal=include_internal, + q=q, + ), + ) + visible = {task.task_id: task for source in records for task in source} + ordered = sorted( + visible.values(), + key=lambda task: (task.created_at, task.task_id), + reverse=True, + ) + return ordered[:limit] + + async def has_running( + self, + task_type: str, + resource_id: str, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> bool: + tasks = await self._store.list( + account_id, + user_id=user_id, + task_type=task_type, + resource_id=resource_id, + ) + return any(task.status in ACTIVE_TASK_STATUSES for task in tasks) + + async def register_work( self, task_id: str, + work_id: str, + queue_name: str, *, - exclude_work_id: Optional[str] = None, - timeout: Optional[float] = None, + account_id: str, + user_id: str, + ) -> WorkRegistration: + work = WorkRecord(work_id=work_id, task_id=task_id, queue_name=queue_name) + + async def register() -> WorkRegistration: + try: + return await self._store.add_work(work, account_id=account_id, user_id=user_id) + except TaskWorkRejected: + return WorkRegistration.REJECTED + + return await register() + + async def discard_work( + self, task_id: str, work_id: str, *, account_id: str, user_id: str ) -> None: - """Let a root coordinator drain its Task's other queue deliveries.""" - with _consume_wait_budget(task_id, timeout) as remaining: - await self._work_index.wait_for_task_work( - task_id, exclude_work_id=exclude_work_id, timeout=remaining + await run_to_completion( + lambda: self._store.discard_work( + task_id, work_id, account_id=account_id, user_id=user_id ) + ) - async def mark_enqueue_complete( + async def reconcile_delivery( self, task_id: str, work_id: str, + queue_name: str, *, account_id: str, user_id: str, ) -> None: - """Persist that a coordinator's complete child Work set is durable.""" - async with self._task_locks.acquire(task_id): - task = await self._load_for_update(task_id, account_id, user_id) - if task is None or task.status not in _ACTIVE_STATUSES: - return - updated = deepcopy(task) - completed = set(updated.auth.get("enqueue_complete_work_ids") or ()) - completed.add(work_id) - updated.auth["enqueue_complete_work_ids"] = sorted(completed) - updated.updated_at = self._next_updated_at(task) - await self._persist_and_publish("update", updated, previous=task) + work = WorkRecord(work_id=work_id, task_id=task_id, queue_name=queue_name) + await self._store.reconcile_delivery(work, account_id=account_id, user_id=user_id) - async def is_enqueue_complete( + async def get_work_state( + self, task_id: str, work_id: str, *, account_id: str, user_id: str + ) -> Optional[WorkState]: + work = await self._store.get_work(task_id, work_id, account_id=account_id, user_id=user_id) + return work.state if work is not None else None + + async def start_work( + self, task_id: str, work_id: str, *, account_id: str, user_id: str + ) -> None: + await self._store.start_work(task_id, work_id, account_id=account_id, user_id=user_id) + + async def mark_work_done( self, task_id: str, work_id: str, *, + context_count_delta: int = 0, + vector_written: bool = False, account_id: str, user_id: str, - ) -> bool: - """Check the durable producer marker used by a lock coordinator.""" - auth = await self.get_task_auth( + ) -> None: + await self._finish_work( task_id, + work_id, + WorkState.DONE, + error=None, + errors=None, + context_count_delta=max(context_count_delta, int(vector_written)), account_id=account_id, user_id=user_id, ) - return work_id in set(auth.get("enqueue_complete_work_ids") or ()) - def build_queue_status( + async def mark_work_failed( self, task_id: str, - work_ids: Optional[List[str]] = None, - ) -> Dict[str, Dict[str, object]]: - """Project Task-scoped queue outcomes from logical Work.""" - return self._public_queue_status(self._work_index.queue_status(task_id, work_ids)) - - @staticmethod - def _public_queue_status( - status: Dict[str, Dict[str, object]], - ) -> Dict[str, Dict[str, object]]: - result: Dict[str, Dict[str, object]] = {} - for queue_name in _PUBLIC_QUEUE_NAMES: - queue_status = deepcopy( - status.get( - queue_name, - { - "processed": 0, - "requeue_count": 0, - "error_count": 0, - "errors": [], - }, - ) - ) - errors = queue_status.get("errors") - if isinstance(errors, list): - queue_status["errors"] = [ - { - **error, - "message": _sanitize_error(str(error.get("message") or "")), - } - if isinstance(error, dict) - else {"message": _sanitize_error(str(error))} - for error in errors[:100] - ] - result[queue_name] = queue_status - return result - - def get_queue_timing(self, task_id: str) -> Dict[str, Dict[str, float]]: - """Return Task-scoped queue wait and execution durations.""" - return self._work_index.queue_timing(task_id) - - def get_context_count(self, task_id: str) -> int: - """Return successful context writes recorded by Task Work.""" - return self._work_index.context_count(task_id) + work_id: str, + error: Optional[str] = None, + *, + errors: Optional[Collection[str]] = None, + account_id: str, + user_id: str, + ) -> None: + await self._finish_work( + task_id, + work_id, + WorkState.FAILED, + error=sanitize_error(error) if error else None, + errors=errors, + context_count_delta=0, + account_id=account_id, + user_id=user_id, + ) - def record_context_count(self, task_id: str, work_id: str, delta: int = 1) -> None: - """Record successful context writes against one logical Work.""" - self._work_index.record_context_count(task_id, work_id, delta) + async def mark_work_requeued( + self, + task_id: str, + work_id: str, + *, + delta: int = 1, + account_id: str, + user_id: str, + ) -> None: + await self._store.requeue_work( + task_id, + work_id, + delta=delta, + account_id=account_id, + user_id=user_id, + ) - async def record_event( + async def _finish_work( self, task_id: str, - kind: str, + work_id: str, + state: WorkState, *, + error: Optional[str], + errors: Optional[Collection[str]], + context_count_delta: int, account_id: str, user_id: str, - operation: Optional[str] = None, ) -> None: - """Record a registered process event without changing task status or stage. + await run_to_completion( + lambda: self._store.finish_work( + task_id, + work_id, + state, + error=error, + errors=errors, + context_count_delta=context_count_delta, + account_id=account_id, + user_id=user_id, + ) + ) - Add process event kinds here and their Studio translations when instrumenting - new execution points. Callers report facts, not inferred lifecycle transitions. - """ - self._validate_owner(account_id, user_id) - if kind not in PROCESS_EVENT_KINDS: - raise ValueError(f"Unknown task process event: {kind}") - if operation is not None and not re.fullmatch(r"[\w.:-]{1,128}", operation): - raise ValueError("operation must be a bounded operation identifier") + async def record_work_error( + self, + task_id: str, + work_id: str, + error: str, + *, + account_id: str, + user_id: str, + ) -> None: + await self._store.record_work_error( + task_id, + work_id, + sanitize_error(error), + account_id=account_id, + user_id=user_id, + ) - async with self._task_locks.acquire(task_id): - task = await self._load_for_update(task_id, account_id, user_id) - if task is None or task.status not in _ACTIVE_STATUSES: - return - updated = deepcopy(task) - updated.execution_events = append_task_event( - updated.execution_events, - kind=kind, - status=task.status.value, - stage=_sanitize_error(task.stage) if task.stage else None, - operation=operation, + async def settle_work( + self, task_id: str, work_id: str, *, account_id: str, user_id: str + ) -> None: + await run_to_completion( + lambda: self._store.settle_work( + task_id, work_id, account_id=account_id, user_id=user_id ) - updated.updated_at = self._next_updated_at(task) - await self._persist_and_publish("update", updated, previous=task) - - def has_work(self, task_id: str, *, exclude_work_id: Optional[str] = None) -> bool: - """Return whether a task still owns durable or active queue work.""" - return self._work_index.has_work(task_id, exclude_work_id=exclude_work_id) + ) - def has_active_producer(self, task_id: str) -> bool: - """Return whether another active handler can persist an enqueue marker.""" - return self._work_index.has_other_active(task_id, asyncio.current_task()) + async def is_cancelling( + self, + task_id: str, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> bool: + return await self.is_cancellation_requested(task_id, account_id=account_id, user_id=user_id) - def register_running_task(self, task_id: str) -> None: - """Register the current asyncio task so cancellation can interrupt it.""" - active_task = asyncio.current_task() - if active_task is not None: - self._work_index.register_active(task_id, active_task) + async def is_cancellation_requested( + self, + task_id: str, + *, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> bool: + return await self._store.is_cancellation_requested( + task_id, account_id=account_id, user_id=user_id + ) - async def unregister_running_task(self, task_id: str) -> None: - """Release direct background work and persist its terminal outcome.""" - active_task = asyncio.current_task() - if active_task is not None: - self._work_index.unregister_active(task_id, active_task) - await self._finalize_task(task_id) + async def wait_for_work( + self, + task_id: str, + work_ids: Collection[str], + *, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + timeout: Optional[float] = None, + stop_waiting: Optional[Callable[[], bool]] = None, + poll_interval: float = _TASK_WAIT_POLL_INTERVAL, + ) -> None: + wanted = set(work_ids) + if not wanted: + return + with _consume_wait_budget(task_id, timeout) as remaining: + deadline = None if remaining is None else time.monotonic() + remaining + while True: + works = await self._store.get_works( + task_id, + wanted, + account_id=account_id, + user_id=user_id, + ) + missing = wanted.difference(works) + if missing: + raise KeyError(f"Work not found for Task {task_id}: {sorted(missing)}") + if all(work.state in TERMINAL_WORK_STATES for work in works.values()): + return + if stop_waiting is not None and stop_waiting(): + raise asyncio.CancelledError("Worker stopped with queued Work") + remaining = None if deadline is None else deadline - time.monotonic() + if remaining is not None and remaining <= 0: + raise TimeoutError(f"Timed out waiting for Work in Task {task_id}") + await asyncio.sleep( + poll_interval if remaining is None else min(poll_interval, remaining) + ) - async def get( + async def wait_for_task_work( self, task_id: str, + *, + exclude_work_id: Optional[str] = None, account_id: Optional[str] = None, user_id: Optional[str] = None, - ) -> Optional[TaskRecord]: - """Look up a single task. Returns a snapshot copy (None if not found).""" - task = self._cached_task(task_id) - if task is not None: - if not self._matches_owner(task, account_id, user_id): - return None - return self._copy(task) - if account_id is None: - return None - task = await self._load_from_store(task_id, account_id, user_id) - if task is not None: - self._merge_loaded_tasks([task]) - task = self._cached_task(task_id) - if task is None or not self._matches_owner(task, account_id, user_id): - return None - return self._copy(task) - - async def _prune_persisted_expired(self, payload: Dict[str, Any]) -> bool: - """Prune cold records under the same lock and I/O budget as mutations.""" - candidate = self._record_from_payload(payload) - now = time.time() - if not self._is_expired(candidate, now): - return False - async with self._task_locks.acquire(candidate.task_id): - # The scan may race with a lifecycle write. Re-read under the same - # lock used by mutations before removing any durable state. - current = await self._store_io.run( - "get", - lambda: self._store.get( - candidate.task_id, account_id=candidate.account_id, user_id=candidate.user_id - ), - ) - if current is None: - return True - task = self._record_from_payload(current) - if not self._is_expired(task, now) or self._work_index.has_work(task.task_id): - payload.clear() - payload.update(current) - return False - await self._store_io.run( - "delete", - lambda: run_to_completion( - lambda: self._store.delete( - task.task_id, account_id=task.account_id, user_id=task.user_id - ) - ), - ) - with self._lock: - self._tasks.pop(task.task_id, None) - self._work_index.forget_task(task.task_id) - return True + timeout: Optional[float] = None, + poll_interval: float = _TASK_WAIT_POLL_INTERVAL, + ) -> None: + with _consume_wait_budget(task_id, timeout) as remaining: + deadline = None if remaining is None else time.monotonic() + remaining + while await self._store.has_open_work( + task_id, + exclude_work_id=exclude_work_id, + account_id=account_id, + user_id=user_id, + ): + remaining = None if deadline is None else deadline - time.monotonic() + if remaining is not None and remaining <= 0: + raise TimeoutError(f"Timed out waiting for Work in Task {task_id}") + await asyncio.sleep( + poll_interval if remaining is None else min(poll_interval, remaining) + ) - async def list_page( + async def mark_enqueue_complete( self, + task_id: str, + work_id: str, *, account_id: str, user_id: str, - limit: int, - before: tuple[float, str] | None = None, - include_cached: bool = False, - additional_owner: tuple[str, str] | None = None, - **filters: Any, - ) -> list[TaskRecord]: - from openviking.service.task_pagination import matches + ) -> None: + await self._store.mark_enqueue_complete( + task_id, + work_id, + account_id=account_id, + user_id=user_id, + ) - owners = {(account_id, user_id)} - if additional_owner: - owners.add(additional_owner) - records = [] - for account, user in owners: - page = await self._store.list_page( - account, - user_id=user, - limit=limit, - before=before, - prune_expired=self._prune_persisted_expired, - io_limiter=self._store_io, - **filters, - ) - records.extend(self._copy(self._record_from_payload(record)) for record in page) - if include_cached: - records.extend(self._copy(t) for t in self._cache_snapshot()) - visible = { - t.task_id: t - for t in records - if (before is None or (t.created_at, t.task_id) < before) - and matches(t.to_dict(), **filters) - } - return sorted(visible.values(), key=lambda t: (t.created_at, t.task_id), reverse=True)[ - :limit - ] + async def is_enqueue_complete( + self, + task_id: str, + work_id: str, + *, + account_id: str, + user_id: str, + ) -> bool: + return await self._store.is_enqueue_complete( + task_id, + work_id, + account_id=account_id, + user_id=user_id, + ) - async def list_tasks( + async def build_queue_status( self, - task_type: Optional[str] = None, - status: Optional[str] = None, - resource_id: Optional[str] = None, - limit: Optional[int] = 50, + task_id: str, + work_ids: Optional[Collection[str]] = None, + *, account_id: Optional[str] = None, user_id: Optional[str] = None, - include_internal: bool = True, - ) -> List[TaskRecord]: - """List tasks with optional filters. Most-recent first. Returns snapshot copies.""" - if account_id is not None: - self._merge_loaded_tasks(await self._load_all_from_store(account_id, user_id)) - source = self._cache_snapshot() - tasks = [t for t in source if self._matches_owner(t, account_id, user_id)] - if not include_internal: - tasks = [t for t in tasks if t.meta.get("internal") is not True] - if task_type: - tasks = [t for t in tasks if t.task_type == task_type] - if status: - tasks = [t for t in tasks if t.status.value == status] - if resource_id: - tasks = [t for t in tasks if t.resource_id == resource_id] - tasks.sort(key=lambda t: t.created_at, reverse=True) - return [self._copy(t) for t in tasks[:limit]] + ) -> Dict[str, Dict[str, object]]: + summary = await self.get_work_summary( + task_id, + work_ids=work_ids, + account_id=account_id, + user_id=user_id, + ) + return self.format_queue_status(summary) - async def has_running( + @staticmethod + def format_queue_status( + summary: WorkSummary, + ) -> Dict[str, Dict[str, object]]: + result: Dict[str, Dict[str, object]] = {} + for queue_name in _PUBLIC_QUEUE_NAMES: + queue = summary.queues.get(queue_name, QueueWorkSummary()) + result[queue_name] = { + "processed": queue.processed, + "requeue_count": queue.requeue_count, + "error_count": queue.error_count, + "errors": [ + {"message": sanitize_error(str(error))} + for error in queue.errors[:MAX_QUEUE_ERROR_DETAILS] + ], + } + return result + + async def get_work_summary( self, - task_type: str, - resource_id: str, + task_id: str, + work_ids: Optional[Collection[str]] = None, + *, account_id: Optional[str] = None, user_id: Optional[str] = None, - ) -> bool: - """Check if there is already a running task for the given type+resource.""" - if account_id is not None: - self._merge_loaded_tasks(await self._load_all_from_store(account_id, user_id)) - tasks = self._cache_snapshot() - return any( - t.task_type == task_type - and t.resource_id == resource_id - and self._matches_owner(t, account_id, user_id) - and t.status in _ACTIVE_STATUSES - for t in tasks + ) -> WorkSummary: + return await self._store.summarize_works( + task_id, + work_ids=work_ids, + account_id=account_id, + user_id=user_id, ) - async def _load_for_update( + @staticmethod + def format_queue_timing( + summary: WorkSummary, + ) -> Dict[str, Dict[str, float]]: + return { + queue_name.lower(): { + "queue_wait_ms": queue.queue_wait_ms, + "execute_ms": queue.execute_ms, + } + for queue_name, queue in summary.queues.items() + } + + async def get_queue_timing( self, task_id: str, - account_id: Optional[str], - user_id: Optional[str], - ) -> Optional[TaskRecord]: - task = self._cached_task(task_id) - if task is not None: - return task if self._matches_owner(task, account_id, user_id) else None - if account_id is None or user_id is None: - return None - return await self._load_from_store(task_id, account_id, user_id) + *, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> Dict[str, Dict[str, float]]: + summary = await self.get_work_summary( + task_id, + work_ids=None, + account_id=account_id, + user_id=user_id, + ) + return self.format_queue_timing(summary) - @staticmethod - def _record_from_payload(payload: Dict[str, Any]) -> TaskRecord: - known_fields = {item.name for item in fields(TaskRecord) if item.init} - data = {key: deepcopy(value) for key, value in payload.items() if key in known_fields} - data["status"] = TaskStatus(data["status"]) - record = TaskRecord(**data) - # Keep fields from other writers for the next persistence update. - record._extra_fields = { - key: deepcopy(value) for key, value in payload.items() if key not in known_fields - } - return record + async def get_context_count( + self, + task_id: str, + *, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> int: + summary = await self.get_work_summary( + task_id, + work_ids=None, + account_id=account_id, + user_id=user_id, + ) + return summary.context_count - async def _load_from_store( + async def record_event( self, task_id: str, + kind: str, + *, account_id: str, - user_id: Optional[str], - ) -> Optional[TaskRecord]: - payload = await self._store_io.run( - "get", - lambda: self._store.get(task_id, account_id=account_id, user_id=user_id), + user_id: str, + operation: Optional[str] = None, + ) -> None: + self._validate_owner(account_id, user_id) + if kind not in PROCESS_EVENT_KINDS: + raise ValueError(f"Unknown task process event: {kind}") + if operation is not None and not re.fullmatch(r"[\w.:-]{1,128}", operation): + raise ValueError("operation must be a bounded operation identifier") + await self._store.record_event( + task_id, + kind, + operation=operation, + account_id=account_id, + user_id=user_id, ) - if payload is None: - return None - return self._record_from_payload(payload) - async def _load_all_from_store( - self, account_id: str, user_id: Optional[str] - ) -> List[TaskRecord]: - return [ - self._record_from_payload(payload) - for payload in await self._store_io.run( - "list", - lambda: self._store.list(account_id, user_id=user_id), - ) - ] + async def has_work( + self, + task_id: str, + *, + exclude_work_id: Optional[str] = None, + account_id: Optional[str] = None, + user_id: Optional[str] = None, + ) -> bool: + return await self._store.has_open_work( + task_id, + exclude_work_id=exclude_work_id, + account_id=account_id, + user_id=user_id, + ) - async def _persist_and_publish( - self, operation: str, task: TaskRecord, *, previous: Optional[TaskRecord] = None - ) -> None: - # All mutation callers supply the snapshot read under this task's lock. - # Build history before the same physical write as the state transition. - if operation == "update" and previous is None: - raise ValueError("Task updates require the previous snapshot") - # Persist the final Work interval union; restored active tasks lack full timing coverage. - if previous is None or previous.status in _ACTIVE_STATUSES: - task.processing_seconds = self._work_index.processing_seconds(task.task_id) - kinds = [] - if previous is None: - kinds.append("created") - else: - if task.error is not None and previous.error is None: - kinds.append("error_recorded") - if task.status != previous.status: - kinds.append("status_changed") - if task.stage != previous.stage and task.status not in _TERMINAL_STATUSES: - kinds.append("stage_changed") - stage = previous.stage if previous and task.status in _TERMINAL_STATUSES else task.stage - recorded_at = datetime.now(timezone.utc).isoformat() - for kind in kinds: - task.execution_events = append_task_event( - task.execution_events, - kind=kind, - status=task.status.value, - stage=_sanitize_error(stage) if stage else None, - error=_sanitize_error(task.error) - if kind == "error_recorded" and task.error is not None - else None, - recorded_at=recorded_at, + def has_active_producer(self, task_id: str) -> bool: + state = self._running.get() + return self._work_index.has_other_active( + task_id, + asyncio.current_task(), + account_id=state[1] if state is not None else None, + user_id=state[2] if state is not None else None, + ) + + async def delete_tasks( + self, + account_id: str, + user_id: Optional[str] = None, + *, + allow_active_types: Collection[str] = (), + ) -> int: + if not account_id: + raise ValueError("Task deletion requires a non-empty account_id") + tasks = await self._store.list(account_id, user_id=user_id) + allowed = set(allow_active_types) + active = [ + task + for task in tasks + if task.status in ACTIVE_TASK_STATUSES and task.task_type not in allowed + ] + if active: + raise RuntimeError( + "Cannot delete active task records: " + + ", ".join(f"{task.task_id}({task.task_type})" for task in active) ) - committed = False - - async def write_and_publish() -> None: - nonlocal committed - if operation == "create": - await self._store.create(task) - else: - await self._store.update(task) - self._publish_task(task) - if operation == "create": - self._work_index.init_processing(task.task_id) - elif task.status in _TERMINAL_STATUSES: - self._work_index.forget_processing(task.task_id) - committed = True - - # Semaphore/lock waits remain cancellable. Once the store operation is - # admitted, settle the physical write and publish its cache snapshot - # before propagating caller cancellation. - try: - await self._store_io.run( - operation, - lambda: run_to_completion(write_and_publish), + for task in tasks: + if task.task_type in allowed: + continue + if await self.has_work( + task.task_id, + account_id=account_id, + user_id=task.user_id, + ): + raise RuntimeError(f"Cannot delete active task record: {task.task_id}") + deleted_task_ids = await self._store.delete_tasks( + account_id=account_id, + user_id=user_id, + allow_active_types=allowed, + ) + for task in tasks: + self._work_index.forget_task( + task.task_id, + account_id=account_id, + user_id=task.user_id, ) - except asyncio.CancelledError as exc: - if committed: - raise _CommittedMutationCancelled() from exc - raise + return len(deleted_task_ids) - def _cached_task(self, task_id: str) -> Optional[TaskRecord]: - with self._lock: - task = self._tasks.get(task_id) - return deepcopy(task) if task is not None else None - - def _cache_snapshot(self) -> List[TaskRecord]: - # Published records are replaced, never mutated. Internal readers may - # share them; public callers receive defensive copies via _copy(). - with self._lock: - return list(self._tasks.values()) - - def _publish_task(self, task: TaskRecord) -> None: - published = deepcopy(task) - with self._lock: - self._tasks[task.task_id] = published - - def _merge_loaded_tasks(self, loaded_tasks: List[TaskRecord]) -> None: - candidates = [deepcopy(task) for task in loaded_tasks] - with self._lock: - for loaded in candidates: - cached = self._tasks.get(loaded.task_id) - if cached is None or loaded.updated_at > cached.updated_at: - self._tasks[loaded.task_id] = loaded + def count(self) -> int: + return sum(sum(statuses.values()) for statuses in self.snapshot_counts_by_type().values()) - @staticmethod - def _next_updated_at(task: TaskRecord) -> float: - return max(time.time(), math.nextafter(task.updated_at, math.inf)) + def snapshot_counts_by_type(self) -> Dict[str, Dict[str, int]]: + return self._store.snapshot_task_stats() - def _copy(self, task: TaskRecord) -> TaskRecord: - """Return a defensive copy of a TaskRecord.""" + @staticmethod + def _public(task: TaskRecord) -> TaskRecord: copied = deepcopy(task) - if task.status in _ACTIVE_STATUSES: - copied.processing_seconds = self._work_index.processing_seconds(task.task_id) - copied.meta = _sanitize_task_result(copied.meta) - copied.result = _sanitize_task_result(copied.result) + copied.meta = sanitize_task_result(copied.meta) + copied.result = sanitize_task_result(copied.result) copied.auth = {} return copied - - def count(self) -> int: - """Return total task count.""" - with self._lock: - return len(self._tasks) - - def snapshot_counts_by_type(self) -> Dict[str, Dict[str, int]]: - """Return a snapshot of task counts grouped by task_type and status.""" - from collections import defaultdict - - grouped: Dict[str, Dict[str, int]] = defaultdict(lambda: defaultdict(int)) - with self._lock: - tasks = list(self._tasks.values()) - for t in tasks: - grouped[t.task_type][t.status.value] += 1 - return {k: dict(v) for k, v in grouped.items()} diff --git a/openviking/service/task_work_index.py b/openviking/service/task_work_index.py index da65f42fbd..d51242bb2e 100644 --- a/openviking/service/task_work_index.py +++ b/openviking/service/task_work_index.py @@ -1,32 +1,23 @@ # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. # SPDX-License-Identifier: AGPL-3.0 -"""Runtime index of logical queue Work owned by tracked Tasks. - -QueueFS remains the durable delivery source of truth. This index separates -logical Work state from delivery counts so a replacement delivery can retain -the same Work ID without allowing the old delivery's ACK to finish the Task. -""" +"""Process-local execution context and cancellable handles for Task Work.""" from __future__ import annotations import asyncio import json import threading -import time from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass -from typing import Any, Awaitable, Callable, Dict, Generator, Iterable, Mapping, Optional +from typing import Any, Dict, Generator, Mapping, Optional from uuid import uuid4 +from openviking.service.task_domain import TaskWorkRejected as TaskWorkRejected + TASK_WORK_ID_FIELD = "_task_work_id" TASK_ACCOUNT_ID_FIELD = "_task_account_id" TASK_USER_ID_FIELD = "_task_user_id" -_TERMINAL_WORK_STATUSES = frozenset({"completed", "failed", "cancelled"}) - - -class TaskWorkRejected(Exception): - """Task cancellation prevented descendant queue work from being created.""" @dataclass(frozen=True) @@ -46,35 +37,6 @@ class QueueTaskMetadata: user_id: str = "" -@dataclass -class WorkTiming: - """Timestamps for one logical Work across all of its deliveries.""" - - enqueued_at: float - started_at: Optional[float] = None - finished_at: Optional[float] = None - - -@dataclass -class WorkRecord: - """Mutable process-local state for one logical Work.""" - - task_id: str - work_id: str - queue_name: str - status: str = "pending" - delivery_count: int = 0 - requeue_count: int = 0 - error: Optional[str] = None - errors: Optional[list[str]] = None - context_count: int = 0 - timing: Optional[WorkTiming] = None - - @property - def terminal(self) -> bool: - return self.status in _TERMINAL_WORK_STATUSES - - _current_task_context: ContextVar[Optional[TaskExecutionContext]] = ContextVar( "openviking_task_execution_context", default=None, @@ -266,494 +228,97 @@ def bind_task_context_from_message(message: Any) -> Generator[None, None, None]: class TaskWorkIndex: - """Thread-safe, rebuildable index of logical Work and active handlers.""" + """Process-local asyncio handles used to accelerate cancellation.""" def __init__(self) -> None: self._lock = threading.Lock() - self._work: Dict[str, Dict[str, WorkRecord]] = {} - self._delivery_counts: Dict[str, int] = {} - self._active: Dict[str, set[asyncio.Task[Any]]] = {} - self._failures: Dict[str, str] = {} - self._work_timings: Dict[str, Dict[str, WorkTiming]] = {} - self._active_work_ids: Dict[tuple[str, asyncio.Task[Any]], str] = {} - self._finalize_before_ack: Optional[Callable[[QueueTaskMetadata], Awaitable[None]]] = None - self._is_cancellation_requested: Optional[Callable[[str], bool]] = None - - def set_callbacks( - self, - *, - finalize_before_ack: Callable[[QueueTaskMetadata], Awaitable[None]], - is_cancellation_requested: Callable[[str], bool], - ) -> None: - self._finalize_before_ack = finalize_before_ack - self._is_cancellation_requested = is_cancellation_requested - - def rebuild( - self, - snapshots: Mapping[str, Iterable[Any]], - ) -> Dict[str, tuple[str, str]]: - """Rebuild unacknowledged deliveries and return Task owners to restore.""" - work: Dict[str, Dict[str, WorkRecord]] = {} - delivery_counts: Dict[str, int] = {} - owners: Dict[str, tuple[str, str]] = {} - now = time.time() - for queue_name, messages in snapshots.items(): - for message in messages: - metadata = extract_task_metadata(message) - if metadata is None: - continue - records = work.setdefault(metadata.task_id, {}) - record = records.setdefault( - metadata.work_id, - WorkRecord( - task_id=metadata.task_id, - work_id=metadata.work_id, - queue_name=queue_name, - timing=WorkTiming(enqueued_at=now), - ), - ) - record.delivery_count += 1 - delivery_counts[metadata.task_id] = delivery_counts.get(metadata.task_id, 0) + 1 - if metadata.account_id and metadata.user_id: - owners[metadata.task_id] = (metadata.account_id, metadata.user_id) - with self._lock: - self._work = work - self._delivery_counts = delivery_counts - self._failures = {} - return owners - - def _add_delivery_locked(self, task_id: str) -> None: - self._delivery_counts[task_id] = self._delivery_counts.get(task_id, 0) + 1 - - def _remove_delivery_locked(self, task_id: str) -> None: - remaining = self._delivery_counts.get(task_id, 0) - 1 - if remaining > 0: - self._delivery_counts[task_id] = remaining - else: - self._delivery_counts.pop(task_id, None) - - def register(self, queue_name: str, metadata: Optional[QueueTaskMetadata]) -> bool: - """Reject cancelled Work or provisionally register one delivery.""" - if metadata is None: - return True - with self._lock: - if self.cancellation_requested(metadata.task_id): - return False - records = self._work.setdefault(metadata.task_id, {}) - record = records.get(metadata.work_id) - now = time.time() - if record is None: - record = WorkRecord( - task_id=metadata.task_id, - work_id=metadata.work_id, - queue_name=queue_name, - timing=WorkTiming(enqueued_at=now), - ) - records[metadata.work_id] = record - elif record.queue_name != queue_name: - raise ValueError( - f"Work {metadata.work_id} already belongs to queue {record.queue_name}" - ) - record.delivery_count += 1 - self._add_delivery_locked(metadata.task_id) - timings = self._work_timings.get(metadata.task_id) - if timings is not None: - record.timing = timings.setdefault(metadata.work_id, record.timing) - return True - - def _remove_delivery( - self, - queue_name: str, - metadata: QueueTaskMetadata, - ) -> tuple[bool, bool]: - """Remove one delivery and report whether its Task became idle.""" - with self._lock: - record = self._work.get(metadata.task_id, {}).get(metadata.work_id) - removed = bool( - record is not None and record.queue_name == queue_name and record.delivery_count > 0 - ) - if removed: - record.delivery_count -= 1 - self._remove_delivery_locked(metadata.task_id) - became_idle = removed and not self._has_work_locked(metadata.task_id) - return removed, became_idle - - def _has_delivery_locked(self, task_id: str, exclude_work_id: Optional[str] = None) -> bool: - if exclude_work_id is not None: - record = self._work.get(task_id, {}).get(exclude_work_id) - excluded = record.delivery_count if record is not None else 0 - return self._delivery_counts.get(task_id, 0) > excluded - return self._delivery_counts.get(task_id, 0) > 0 - - def _has_work_locked(self, task_id: str, exclude_work_id: Optional[str] = None) -> bool: - if self._has_delivery_locked(task_id, exclude_work_id): - return True - if exclude_work_id is not None: - return False - return bool(self._active.get(task_id)) - - async def discard(self, queue_name: str, message: Any) -> None: - """Discard a provisionally registered delivery that did not commit.""" - metadata = ( - message if isinstance(message, QueueTaskMetadata) else extract_task_metadata(message) - ) - if metadata is None: - return - removed, became_idle = self._remove_delivery(queue_name, metadata) - if removed: - with self._lock: - record = self._work.get(metadata.task_id, {}).get(metadata.work_id) - if record is not None and record.delivery_count == 0 and not record.terminal: - self._work[metadata.task_id].pop(metadata.work_id, None) - timings = self._work_timings.get(metadata.task_id) - if timings is not None: - timings.pop(metadata.work_id, None) - callback = self._finalize_before_ack - if became_idle and callback is not None: - await callback(metadata) - - async def fail_enqueue( - self, - queue_name: str, - metadata: QueueTaskMetadata, - error: str, - ) -> None: - """Keep a failed logical Work visible after its delivery failed to persist.""" - self.record_result(metadata, "failed", error=error) - removed, became_idle = self._remove_delivery(queue_name, metadata) - if not removed: - return - callback = self._finalize_before_ack - if became_idle and callback is not None: - await callback(metadata) - - async def prepare_ack( - self, - queue_name: str, - message: Any, - ) -> Optional[QueueTaskMetadata]: - """Provisionally remove one delivery and finalize before the last ACK.""" - metadata = ( - message if isinstance(message, QueueTaskMetadata) else extract_task_metadata(message) - ) - if metadata is None: - return None - - removed, became_idle = self._remove_delivery(queue_name, metadata) - if not removed: - return None - if not became_idle: - return metadata - - callback = self._finalize_before_ack - if callback is None: - return metadata - try: - await callback(metadata) - except BaseException: - self.rollback_ack(queue_name, metadata) - raise - return metadata - - def rollback_ack(self, queue_name: str, metadata: QueueTaskMetadata) -> None: - """Restore a provisionally removed delivery when finalization or ACK fails.""" - with self._lock: - record = self._work.get(metadata.task_id, {}).get(metadata.work_id) - if record is None: - record = WorkRecord( - task_id=metadata.task_id, - work_id=metadata.work_id, - queue_name=queue_name, - timing=WorkTiming(enqueued_at=time.time()), - ) - self._work.setdefault(metadata.task_id, {})[metadata.work_id] = record - record.delivery_count += 1 - self._add_delivery_locked(metadata.task_id) - - def has_work(self, task_id: str, exclude_work_id: Optional[str] = None) -> bool: - with self._lock: - return self._has_work_locked(task_id, exclude_work_id) + self._active: Dict[tuple[str, str, str], set[asyncio.Task[Any]]] = {} def has_other_active( self, task_id: str, current_task: Optional[asyncio.Task[Any]], + *, + account_id: Optional[str] = None, + user_id: Optional[str] = None, ) -> bool: """Return whether another in-process handler can finish producing Work.""" with self._lock: - return any(task is not current_task for task in self._active.get(task_id, ())) - - def cancellation_requested(self, task_id: str) -> bool: - callback = self._is_cancellation_requested - return bool(callback is not None and callback(task_id)) + return any( + task is not current_task + for (account, user, active_task_id), tasks in self._active.items() + if active_task_id == task_id + and (account_id is None or account == account_id) + and (user_id is None or user == user_id) + for task in tasks + ) - def record_result( + def forget_task( self, - metadata: QueueTaskMetadata, - outcome: str, + task_id: str, *, - error: Optional[str] = None, - errors: Optional[Iterable[str]] = None, + account_id: Optional[str] = None, + user_id: Optional[str] = None, ) -> None: - """Apply one delivery result to its logical Work monotonically.""" with self._lock: - record = self._work.get(metadata.task_id, {}).get(metadata.work_id) - if record is None or record.terminal: - return - if outcome == "requeued": - record.requeue_count += 1 - record.status = "pending" - return - if outcome not in _TERMINAL_WORK_STATUSES: - raise ValueError(f"Unsupported Work outcome: {outcome}") - record.status = outcome - record.error = error if outcome == "failed" and error else None - if outcome == "failed": - details = [str(value) for value in errors or () if str(value)] - record.errors = details or ([error] if error else []) - if record.timing is not None: - record.timing.finished_at = max( - record.timing.started_at or record.timing.enqueued_at, - time.time(), - ) - if outcome == "failed" and error: - self._failures.setdefault(metadata.task_id, error) - - def record_context_count(self, task_id: str, work_id: str, delta: int = 1) -> None: - """Add a result count to one Work without changing its lifecycle.""" - if delta <= 0: - return - with self._lock: - record = self._work.get(task_id, {}).get(work_id) - if record is not None: - record.context_count += delta - - def record_failure(self, task_id: str, error: str) -> None: - """Record the first failed Work owned by a Task.""" - with self._lock: - self._failures.setdefault(task_id, error) - - def failure(self, task_id: str) -> Optional[str]: - with self._lock: - return self._failures.get(task_id) + for key in [ + key + for key in self._active + if key[2] == task_id + and (account_id is None or key[0] == account_id) + and (user_id is None or key[1] == user_id) + ]: + self._active.pop(key, None) - def clear_failure(self, task_id: str) -> None: - with self._lock: - self._failures.pop(task_id, None) - - async def wait_for_work( + def register_active( self, task_id: str, - work_ids: Iterable[str], + active_task: asyncio.Task[Any], *, - timeout: Optional[float] = None, - poll_interval: float = 0.05, - stop_waiting: Optional[Callable[[], bool]] = None, + account_id: str, + user_id: str, ) -> None: - """Wait until exactly the supplied logical Work IDs are terminal.""" - requested = frozenset(str(work_id) for work_id in work_ids) - if not requested: - return - - deadline = None if timeout is None else time.monotonic() + timeout - while True: - with self._lock: - records = self._work.get(task_id, {}) - missing = requested.difference(records) - done = not missing and all(records[work_id].terminal for work_id in requested) - if done: - return - if missing: - raise KeyError(f"Unknown Work for Task {task_id}: {sorted(missing)}") - if stop_waiting is not None and stop_waiting(): - raise asyncio.CancelledError("Worker stopped with queued Work") - remaining = None if deadline is None else deadline - time.monotonic() - if remaining is not None and remaining <= 0: - raise TimeoutError(f"Timed out waiting for Work in Task {task_id}") - await asyncio.sleep( - poll_interval if remaining is None else min(poll_interval, remaining) - ) + with self._lock: + self._active.setdefault((account_id, user_id, task_id), set()).add(active_task) - async def wait_for_task_work( + def unregister_active( self, task_id: str, + active_task: asyncio.Task[Any], *, - exclude_work_id: Optional[str] = None, - timeout: Optional[float] = None, - poll_interval: float = 0.05, + account_id: str, + user_id: str, ) -> None: - """Drain Task deliveries, including newly enqueued and recovered Work. - - Only the Task coordinator may use this barrier, excluding its own Work. - Inline coordinators have no delivery to exclude. Other workers must wait - for explicit Work IDs to avoid waiting on each other. - """ - deadline = None if timeout is None else time.monotonic() + timeout - while True: - with self._lock: - pending = self._has_delivery_locked(task_id, exclude_work_id) - if not pending: - return - remaining = None if deadline is None else deadline - time.monotonic() - if remaining is not None and remaining <= 0: - raise TimeoutError(f"Timed out waiting for Work in Task {task_id}") - await asyncio.sleep( - poll_interval if remaining is None else min(poll_interval, remaining) - ) - - def queue_status( - self, - task_id: str, - work_ids: Optional[Iterable[str]] = None, - ) -> Dict[str, Dict[str, object]]: - """Project queue counters from logical Work records.""" - selected = None if work_ids is None else frozenset(str(value) for value in work_ids) - with self._lock: - records = [ - record - for work_id, record in self._work.get(task_id, {}).items() - if selected is None or work_id in selected - ] - result: Dict[str, Dict[str, object]] = {} - for record in records: - status = result.setdefault( - record.queue_name, - { - "processed": 0, - "requeue_count": 0, - "error_count": 0, - "errors": [], - }, - ) - if record.status == "completed": - status["processed"] += 1 - status["requeue_count"] += record.requeue_count - if record.status == "failed": - errors = record.errors or ([record.error] if record.error else []) - status["error_count"] += len(errors) or 1 - status["errors"].extend({"message": message} for message in errors) - return result - - def queue_timing(self, task_id: str) -> Dict[str, Dict[str, float]]: - """Project queue wait and execution timing in milliseconds.""" with self._lock: - result: Dict[str, Dict[str, float]] = {} - for record in self._work.get(task_id, {}).values(): - timing = record.timing - metrics = result.setdefault( - record.queue_name.lower(), - {"queue_wait_ms": 0.0, "execute_ms": 0.0}, - ) - if timing is None or timing.started_at is None: - continue - metrics["queue_wait_ms"] += max( - 0.0, (timing.started_at - timing.enqueued_at) * 1000.0 - ) - if timing.finished_at is not None: - metrics["execute_ms"] += max( - 0.0, (timing.finished_at - timing.started_at) * 1000.0 - ) - return result - - def context_count(self, task_id: str) -> int: - with self._lock: - return sum(record.context_count for record in self._work.get(task_id, {}).values()) - - def init_processing(self, task_id: str) -> None: - with self._lock: - timings = self._work_timings.setdefault(task_id, {}) - for work_id, record in self._work.get(task_id, {}).items(): - if record.timing is not None: - timings.setdefault(work_id, record.timing) - - def processing_seconds(self, task_id: str) -> Optional[float]: - with self._lock: - timings = self._work_timings.get(task_id) - if timings is None: - return None - now = time.time() - intervals: list[tuple[float, float]] = [] - for timing in timings.values(): - if timing.started_at is None: - continue - end = timing.finished_at if timing.finished_at is not None else now - intervals.append((timing.started_at, max(timing.started_at, end))) - intervals.sort() - if not intervals: - return None - - elapsed = 0.0 - current_start, current_end = intervals[0] - for start, end in intervals[1:]: - if start <= current_end: - current_end = max(current_end, end) - else: - elapsed += current_end - current_start - current_start, current_end = start, end - return elapsed + current_end - current_start - - def forget_processing(self, task_id: str) -> None: - with self._lock: - self._work_timings.pop(task_id, None) + key = (account_id, user_id, task_id) + entries = self._active.get(key) + if entries is not None: + entries.discard(active_task) + if not entries: + self._active.pop(key, None) - def forget_task(self, task_id: str) -> None: - """Release all process-local state after the owning Task is removed.""" + def active_tasks_by_owner(self) -> Dict[tuple[str, str], set[str]]: with self._lock: - self._work.pop(task_id, None) - self._delivery_counts.pop(task_id, None) - self._active.pop(task_id, None) - self._failures.pop(task_id, None) - self._work_timings.pop(task_id, None) - active_keys = [key for key in self._active_work_ids if key[0] == task_id] - for key in active_keys: - self._active_work_ids.pop(key, None) + grouped: Dict[tuple[str, str], set[str]] = {} + for account_id, user_id, task_id in self._active: + grouped.setdefault((account_id, user_id), set()).add(task_id) + return grouped - def register_active( + def cancel_active( self, task_id: str, - active_task: asyncio.Task[Any], *, - work_id: Optional[str] = None, + account_id: Optional[str] = None, + user_id: Optional[str] = None, ) -> None: with self._lock: - cancelled = self.cancellation_requested(task_id) - if not cancelled: - self._active.setdefault(task_id, set()).add(active_task) - now = time.time() - timing_id = work_id or f"inline:{id(active_task)}" - record = self._work.get(task_id, {}).get(timing_id) - if record is not None: - timing = record.timing - if timing is not None and timing.started_at is None: - timing.started_at = now - if not record.terminal: - record.status = "running" - timings = self._work_timings.get(task_id) - if timings is not None: - timing = timings.setdefault(timing_id, WorkTiming(enqueued_at=now)) - if timing.started_at is None: - timing.started_at = now - timing.finished_at = None - self._active_work_ids[(task_id, active_task)] = timing_id - if cancelled: - active_task.get_loop().call_soon(active_task.cancel) - - def unregister_active(self, task_id: str, active_task: asyncio.Task[Any]) -> None: - with self._lock: - timing_id = self._active_work_ids.pop((task_id, active_task), None) - timings = self._work_timings.get(task_id) - if timing_id is not None and timings is not None: - timing = timings.get(timing_id) - if timing is not None and timing.started_at is not None: - timing.finished_at = max(timing.started_at, time.time()) - entries = self._active.get(task_id) - if entries is not None: - entries.discard(active_task) - if not entries: - self._active.pop(task_id, None) - - def cancel_active(self, task_id: str) -> None: - with self._lock: - active = list(self._active.get(task_id, ())) + active = [ + task + for (account, user, active_task_id), tasks in self._active.items() + if active_task_id == task_id + and (account_id is None or account == account_id) + and (user_id is None or user == user_id) + for task in tasks + ] for task in active: task.get_loop().call_soon_threadsafe(task.cancel) diff --git a/openviking/session/session.py b/openviking/session/session.py index 847f6999e8..db8ed545d1 100644 --- a/openviking/session/session.py +++ b/openviking/session/session.py @@ -1197,7 +1197,7 @@ async def _ensure_phase1_ready(self, archive_uri: str) -> bool: from openviking.service.task_tracker import get_task_tracker tracker = get_task_tracker() - if not task_id or not tracker.has_work(str(task_id)): + if not task_id or not await tracker.has_work(str(task_id)): error = "Phase 1 has no QueueFS work to resume" await self._write_failed_marker( archive_uri, @@ -2516,7 +2516,7 @@ async def _run_long_term_memory_extraction( telemetry.set_error("session.commit.phase2", "CANCELLED", "session commit cancelled") snapshot = telemetry.finish("cancelled") _publish_telemetry_summary_best_effort(snapshot) - if tracker.is_cancellation_requested(task_id): + if await tracker.is_cancellation_requested(task_id): await self._write_failed_marker( archive_uri, stage="cancelled", @@ -2836,7 +2836,7 @@ async def _can_run_archive(self, archive_index: int) -> bool: from openviking.service.task_tracker import get_task_tracker tracker = get_task_tracker() - if tracker.has_work(str(task_id)): + if await tracker.has_work(str(task_id)): return False error = "Session commit queue work is missing" diff --git a/openviking/storage/collection_schemas.py b/openviking/storage/collection_schemas.py index d10fea436e..dfffc8e81f 100644 --- a/openviking/storage/collection_schemas.py +++ b/openviking/storage/collection_schemas.py @@ -21,7 +21,6 @@ from openviking.models.embedder.base import embed_compat from openviking.models.model_error import ModelErrorKind, model_error_kind from openviking.server.identity import RequestContext, Role -from openviking.service.task_tracker import get_task_tracker from openviking.service.task_tracker_concurrency import run_to_completion from openviking.storage.acl import ACL_GRANT_FIELDS, ACL_MODE_FIELD, AclMode from openviking.storage.errors import ( @@ -701,19 +700,16 @@ async def on_dequeue(self, data: Optional[Dict[str, Any]]) -> ProcessResult: if self._vikingdb.is_closing: logger.debug("Skip embedding dequeue during shutdown") self._merge_request_stats(embedding_msg.telemetry_id, processed=1) - self._record_request_success(embedding_msg) return ProcessResult.success() if embedding_msg.action is IndexAction.NONE: self._merge_request_stats(embedding_msg.telemetry_id, processed=1) - self._record_request_success(embedding_msg) logger.debug("Skipping no-op embedding message: %s", embedding_msg.id) return ProcessResult.success() if embedding_msg.action not in {IndexAction.UPSERT, IndexAction.MERGE}: result = await self._apply_non_embedding_operation(embedding_msg, ctx) self._merge_request_stats(embedding_msg.telemetry_id, processed=1) - self._record_request_success(embedding_msg) return ProcessResult.success(result) if not isinstance(embedding_msg.message, (str, list)): @@ -721,7 +717,6 @@ async def on_dequeue(self, data: Optional[Dict[str, Any]]) -> ProcessResult: f"Skipping unsupported message type: {type(embedding_msg.message)}" ) self._merge_request_stats(embedding_msg.telemetry_id, processed=1) - self._record_request_success(embedding_msg) return ProcessResult.success(data) # Circuit breaker: if API is known-broken, re-enqueue and wait @@ -942,7 +937,6 @@ async def on_dequeue(self, data: Optional[Dict[str, Any]]) -> ProcessResult: if self._vikingdb.is_closing: logger.debug(f"Skip embedding write during shutdown: {db_err}") self._merge_request_stats(embedding_msg.telemetry_id, processed=1) - self._record_request_success(embedding_msg) return ProcessResult.success() error_msg = self._embedding_error_msg( embedding_msg, @@ -960,7 +954,6 @@ async def on_dequeue(self, data: Optional[Dict[str, Any]]) -> ProcessResult: if self._vikingdb.is_closing: logger.debug(f"Skip embedding write during shutdown: {db_err}") self._merge_request_stats(embedding_msg.telemetry_id, processed=1) - self._record_request_success(embedding_msg) return ProcessResult.success() error_msg = self._embedding_error_msg( embedding_msg, @@ -973,20 +966,9 @@ async def on_dequeue(self, data: Optional[Dict[str, Any]]) -> ProcessResult: return ProcessResult.failed(error_msg) self._merge_request_stats(embedding_msg.telemetry_id, processed=1) - self._record_request_success( - embedding_msg, - vector_written=bool(record_id), - ) - return ProcessResult.success(inserted_data) + return ProcessResult.success({"vector_written": bool(record_id)}) except asyncio.CancelledError: - if ( - embedding_msg is not None - and embedding_msg.context_data.get("context_type") == ContextType.SKILL.value - ): - # Active cancellation does not call on_cancelled in NamedQueue. - # Settle only after any already-started vector write has exited. - self._record_request_success(embedding_msg) raise except Exception as e: error_msg = self._embedding_error_msg( @@ -1021,18 +1003,15 @@ async def on_dequeue(self, data: Optional[Dict[str, Any]]) -> ProcessResult: async def on_cancelled(self, data: Optional[Dict[str, Any]]) -> ProcessResult: """Settle request-scoped waiting when a queued embedding is cancelled.""" - embedding_msg: Optional[EmbeddingMsg] = None try: if data: payload = data.get("data", data) if isinstance(payload, str): payload = json.loads(payload) - embedding_msg = EmbeddingMsg.from_dict(payload) + EmbeddingMsg.from_dict(payload) except (KeyError, TypeError, ValueError) as exc: return ProcessResult.failed(str(exc)) - if embedding_msg is not None: - self._record_request_success(embedding_msg) return ProcessResult.cancelled() async def _reenqueue_embedding_msg(self, msg: EmbeddingMsg) -> None: @@ -1040,15 +1019,3 @@ async def _reenqueue_embedding_msg(self, msg: EmbeddingMsg) -> None: await run_to_completion(lambda: self._vikingdb.enqueue_embedding_msg(msg)) else: await self._vikingdb.enqueue_embedding_msg(msg) - - @staticmethod - def _record_request_success( - embedding_msg: EmbeddingMsg, - *, - vector_written: bool = False, - ) -> None: - if vector_written and embedding_msg.task_id and embedding_msg._task_work_id: - get_task_tracker().record_context_count( - embedding_msg.task_id, - embedding_msg._task_work_id, - ) diff --git a/openviking/storage/content_write.py b/openviking/storage/content_write.py index 9f3455c8b0..850cde9741 100644 --- a/openviking/storage/content_write.py +++ b/openviking/storage/content_write.py @@ -2075,7 +2075,7 @@ async def _wait_for_request( await tracker.wait_for_work(context.task_id, work_ids, timeout=None) except (TimeoutError, asyncio.TimeoutError) as exc: raise DeadlineExceededError("queue processing", None) from exc - return tracker.build_queue_status( + return await tracker.build_queue_status( context.task_id, work_ids if context.work_id is not None else None ) diff --git a/openviking/storage/queuefs/add_resource_processor.py b/openviking/storage/queuefs/add_resource_processor.py index e0dd262662..6085ad1b94 100644 --- a/openviking/storage/queuefs/add_resource_processor.py +++ b/openviking/storage/queuefs/add_resource_processor.py @@ -302,14 +302,19 @@ async def _set_stage(stage: str) -> None: msg.task_id, exclude_work_id=metadata.work_id if metadata is not None else None, ) - result["queue_status"] = tracker.build_queue_status(msg.task_id) + work_summary = await tracker.get_work_summary( + msg.task_id, + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) + result["queue_status"] = tracker.format_queue_status(work_summary) if replay_result is None: - result["context_count"] = tracker.get_context_count(msg.task_id) + result["context_count"] = work_summary.context_count record_resource_queue_metrics( telemetry=telemetry, telemetry_id=telemetry_id, root_uri=result.get("root_uri"), - queue_timing=tracker.get_queue_timing(msg.task_id), + queue_timing=tracker.format_queue_timing(work_summary), ) telemetry.set("resource.total.duration_ms", telemetry.elapsed_ms()) diff --git a/openviking/storage/queuefs/external_task_processor.py b/openviking/storage/queuefs/external_task_processor.py index 47a8ea14b1..6f296094ae 100644 --- a/openviking/storage/queuefs/external_task_processor.py +++ b/openviking/storage/queuefs/external_task_processor.py @@ -8,7 +8,11 @@ from typing import Any, Dict, Optional from openviking.service.external_task_service import ExternalTaskService -from openviking.service.task_work_index import TaskWorkRejected +from openviking.service.task_work_index import ( + TASK_WORK_ID_FIELD, + TaskWorkRejected, + extract_task_metadata, +) from openviking.storage.queuefs import QueueManager, get_queue_manager from openviking.storage.queuefs.named_queue import DequeueHandlerBase from openviking.storage.queuefs.process_result import ProcessResult @@ -45,11 +49,15 @@ async def on_dequeue(self, data: Optional[Dict[str, Any]]) -> ProcessResult: processed = await self._service.execute(task_id, account_id, user_id) if not processed: # Register a new delivery before the old one is ACKed. Keeping the - # task ID preserves cancellation and task-owned work across rotation. + # Work ID preserves one logical Work across delivery rotation. + replacement = {"task_id": task_id, "account_id": account_id, "user_id": user_id} + metadata = extract_task_metadata(data) + if metadata is not None: + replacement[TASK_WORK_ID_FIELD] = metadata.work_id try: await get_queue_manager().enqueue( QueueManager.EXTERNAL_TASK, - {"task_id": task_id, "account_id": account_id, "user_id": user_id}, + replacement, ) except TaskWorkRejected: # Cancellation only needs the current delivery to be ACKed. diff --git a/openviking/storage/queuefs/queue_manager.py b/openviking/storage/queuefs/queue_manager.py index 1907850d1b..6a29f9b6a4 100644 --- a/openviking/storage/queuefs/queue_manager.py +++ b/openviking/storage/queuefs/queue_manager.py @@ -11,7 +11,6 @@ import traceback from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Set, Union -from openviking.service.task_work_index import TaskWorkIndex from openviking.utils.async_client_cache import LoopScopedAsyncClientCache from openviking_cli.utils.logger import get_logger @@ -136,13 +135,12 @@ def __init__( self._queue_stop_events: Dict[str, threading.Event] = {} self._embedding_worker_stopped = threading.Event() self._poll_interval = 0.2 - self._task_work_index = TaskWorkIndex() self._vlm_resolver: Optional["VLMResolver"] = None # Import at composition time to avoid a service <-> queue package cycle. from openviking.service.task_queue_middleware import TaskWorkQueueMiddleware self._middlewares: tuple[QueueMiddleware, ...] = ( - TaskWorkQueueMiddleware(self._task_work_index), + TaskWorkQueueMiddleware(), *middlewares, ) @@ -167,10 +165,28 @@ def start(self) -> None: logger.info(f"[QueueManager] mount_point={self.mount_point} Started") async def prepare_task_tracking(self, tracker: Any) -> list[Any]: - """Rebuild task work from QueueFS before any consumer starts.""" + """Bootstrap Work only for stores that reconstruct it from QueueFS.""" + from openviking.service.task_domain import WorkRecoveryMode + from openviking.service.task_work_index import extract_task_metadata + + if tracker.work_recovery_mode is WorkRecoveryMode.DURABLE_STORE: + return [] + snapshots = {name: await queue.snapshot() for name, queue in self._queues.items()} - owners = self._task_work_index.rebuild(snapshots) - tracker.attach_work_index(self._task_work_index) + owners: Dict[str, tuple[str, str]] = {} + for queue_name, messages in snapshots.items(): + for message in messages: + metadata = extract_task_metadata(message) + if metadata is None or not metadata.account_id or not metadata.user_id: + continue + owners[metadata.task_id] = (metadata.account_id, metadata.user_id) + await tracker.reconcile_delivery( + metadata.task_id, + metadata.work_id, + queue_name, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) return await tracker.restore_work_tasks(owners) def setup_standard_queues( diff --git a/openviking/storage/queuefs/semantic_processor.py b/openviking/storage/queuefs/semantic_processor.py index 222494fdb8..3ed8656528 100644 --- a/openviking/storage/queuefs/semantic_processor.py +++ b/openviking/storage/queuefs/semantic_processor.py @@ -598,9 +598,11 @@ async def _process_message( account_id=task_context.account_id, user_id=task_context.user_id, ) - siblings_pending = tracker.has_work( + siblings_pending = await tracker.has_work( task_context.task_id, exclude_work_id=task_context.work_id, + account_id=task_context.account_id, + user_id=task_context.user_id, ) else: enqueue_complete = True diff --git a/openviking/storage/viking_fs/_snapshot.py b/openviking/storage/viking_fs/_snapshot.py index b88c352cef..1f0e169956 100644 --- a/openviking/storage/viking_fs/_snapshot.py +++ b/openviking/storage/viking_fs/_snapshot.py @@ -858,7 +858,11 @@ async def _run_restore_rebuild_tracked( from openviking.service.task_work_index import bind_task_context tracker = get_task_tracker() - tracker.register_running_task(task_id) + await tracker.register_running_task( + task_id, + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) try: await tracker.start( task_id, diff --git a/tests/agfs/test_viking_fs_git.py b/tests/agfs/test_viking_fs_git.py index 94ed4253b1..72d071dafa 100644 --- a/tests/agfs/test_viking_fs_git.py +++ b/tests/agfs/test_viking_fs_git.py @@ -21,9 +21,11 @@ AGFSNotSupportedError, ) from openviking.server.identity import RequestContext, Role +from openviking.service.task_store import PersistentTaskStore from openviking.storage.viking_fs import VikingFS from openviking_cli.exceptions import NotFoundError from openviking_cli.session.user_id import UserIdentifier +from tests.test_task_tracker import _FakeAgfs ragfs_python = pytest.importorskip("ragfs_python") @@ -635,26 +637,16 @@ def factory(*, vlm_resolver, vector_config_resolver): monkeypatch.setattr(reindex_mod, "ReindexExecutor", factory) -class _MemTaskStore: +class _MemTaskStore(PersistentTaskStore): def __init__(self, *, fail_creates: int = 0): + super().__init__(_FakeAgfs()) self.fail_creates = fail_creates - async def create(self, task): + async def create(self, task, *, reject_if_active=False): if self.fail_creates: self.fail_creates -= 1 raise OSError("task store temporarily unavailable") - - async def update(self, task): - return None - - async def get(self, task_id, *, account_id=None, user_id=None): - return None - - async def list(self, account_id, *, user_id=None): - return [] - - async def delete(self, task_id, *, account_id, user_id=None): - return None + return await super().create(task, reject_if_active=reject_if_active) @pytest.mark.asyncio diff --git a/tests/server/test_admin_api.py b/tests/server/test_admin_api.py index ac43dd73cc..f8c8d1cafe 100644 --- a/tests/server/test_admin_api.py +++ b/tests/server/test_admin_api.py @@ -40,6 +40,10 @@ SYSTEM_TASK_USER_ID, ) from openviking.service.task_tracker import get_task_tracker +from openviking.service.task_work_index import ( + TaskWorkRejected, + bind_task_context, +) from openviking.session.memory.account_templates import ( EDITABLE_MEMORY_TEMPLATE_FIELDS, account_memory_template_path, @@ -319,8 +323,6 @@ async def test_delete_account_files_disables_automatic_pathlock(monkeypatch): agfs.stat = AsyncMock(side_effect=AGFSNotFoundError(path)) service = Mock() service.viking_fs._async_agfs = agfs - tracker = Mock() - monkeypatch.setattr("openviking.service.deletion.get_task_tracker", lambda: tracker) deletion = DeletionService( service=service, manager=Mock(), @@ -330,7 +332,6 @@ async def test_delete_account_files_disables_automatic_pathlock(monkeypatch): await deletion._delete_account_files("account-to-delete") agfs.rm.assert_awaited_once_with(path, recursive=True, auto_pathlock=False) - tracker.forget_account_tasks.assert_called_once_with("account-to-delete") @pytest_asyncio.fixture @@ -1036,10 +1037,12 @@ async def test_account_memory_templates_range_expressions_reach_file_body( fs = lightweight_admin_app.state.fake_service.viking_fs snapshot = await resolve_account_memory_registry(fs, account_id, MemoryTypeRegistry()) schema = snapshot.get("events") - context = ExtractContext([ - Message(id="m1", role="user", parts=[TextPart("Selected source message")]), - Message(id="m2", role="user", parts=[TextPart("Unrelated source message")]), - ]) + context = ExtractContext( + [ + Message(id="m1", role="user", parts=[TextPart("Selected source message")]), + Message(id="m2", role="user", parts=[TextPart("Unrelated source message")]), + ] + ) rendered = MemoryFileUtils.write( MemoryFile(memory_type="events", extra_fields={"ranges": "0", "summary": "summary"}), content_template=schema.content_template, @@ -1362,29 +1365,33 @@ async def test_account_memory_templates_locked_fields_cannot_be_modified( for key, value in defaults.items() if key not in editable ] - invalid.extend([ - {"unknown": "ignored?"}, - {"Description": "wrong key"}, - {"_updated_at": "forged"}, - {"_account_content_template": False}, - {"deployment_content_template": "forged"}, - {"enabled": 1}, - {"fields": [{"name": "renamed_field", "description": "changed"}]}, - {"description": None}, - {"description": " \n"}, - {"fields": [None]}, - ]) + invalid.extend( + [ + {"unknown": "ignored?"}, + {"Description": "wrong key"}, + {"_updated_at": "forged"}, + {"_account_content_template": False}, + {"deployment_content_template": "forged"}, + {"enabled": 1}, + {"fields": [{"name": "renamed_field", "description": "changed"}]}, + {"description": None}, + {"description": " \n"}, + {"fields": [None]}, + ] + ) for field in defaults["fields"]: invalid.extend( {"fields": [{"name": field["name"], key: None if value is not None else "changed"}]} for key, value in field.items() if key not in {"name", "description"} ) - invalid.extend([ - {"fields": [{"name": field["name"], "unknown": "ignored?"}]}, - {"fields": [{"name": field["name"], "description": ""}]}, - {"fields": [{"name": field["name"]}, {"name": field["name"]}]}, - ]) + invalid.extend( + [ + {"fields": [{"name": field["name"], "unknown": "ignored?"}]}, + {"fields": [{"name": field["name"], "description": ""}]}, + {"fields": [{"name": field["name"]}, {"name": field["name"]}]}, + ] + ) if memory_type == "events": invalid.extend( {"fields": [{"name": name, "description": "changed"}]} for name in ("goal", "ranges") @@ -2378,18 +2385,29 @@ async def fail_vectors(account_id, *, ctx): account_id=SYSTEM_TASK_ACCOUNT_ID, user_id=SYSTEM_TASK_USER_ID, ) - await queue.enqueue({ - "task_id": user_cleanup.task_id, - "account_id": acct, - "user_id": "alice", - "target": {"account_id": acct, "user_id": "bob"}, - }) - await queue.enqueue({ - "task_id": late_task_id, - "account_id": SYSTEM_TASK_ACCOUNT_ID, - "user_id": SYSTEM_TASK_USER_ID, - "target": {"account_id": acct, "user_id": "bob"}, - }) + with bind_task_context(user_cleanup.task_id, acct, "alice"): + with pytest.raises(TaskWorkRejected, match="unavailable"): + await queue.enqueue( + { + "task_id": user_cleanup.task_id, + "account_id": acct, + "user_id": "alice", + "target": {"account_id": acct, "user_id": "bob"}, + } + ) + with bind_task_context( + late_task_id, + SYSTEM_TASK_ACCOUNT_ID, + SYSTEM_TASK_USER_ID, + ): + await queue.enqueue( + { + "task_id": late_task_id, + "account_id": SYSTEM_TASK_ACCOUNT_ID, + "user_id": SYSTEM_TASK_USER_ID, + "target": {"account_id": acct, "user_id": "bob"}, + } + ) skipped = await _wait_for_task(admin_client, late_task_id) assert skipped["status"] == "completed" assert skipped["result"] == {"deleted": not recreated, "stale": True} @@ -3521,10 +3539,12 @@ async def test_server_restart_requires_root_and_valid_revision( server_overrides = {"host": cli_host} lightweight_admin_app.state.server_config_overrides = server_overrides path.write_text( - json.dumps({ - "server": server, - "vlm": {"provider": "openai", "model": "gpt-4o", "api_key": "test-key"}, - }) + json.dumps( + { + "server": server, + "vlm": {"provider": "openai", "model": "gpt-4o", "api_key": "test-key"}, + } + ) ) monkeypatch.setattr(OpenVikingConfigSingleton, "_config_file", path) monkeypatch.setattr(OpenVikingConfigSingleton, "_config_file_content", path.read_bytes()) diff --git a/tests/server/test_bot_proxy_auth.py b/tests/server/test_bot_proxy_auth.py index 7e909871cc..ad1e113d6b 100644 --- a/tests/server/test_bot_proxy_auth.py +++ b/tests/server/test_bot_proxy_auth.py @@ -371,9 +371,9 @@ async def request(self, method, url, headers, json): } private_payload = {"args": private_args} if private_args else {} assert task.to_dict()["meta"]["request"].get("args", {}) == public_args - stored = await store.get(task.task_id, **owner) + stored = await store.get_task(task.task_id, **owner) assert stored is not None - assert stored["meta"]["request"].get("args", {}) == public_args + assert stored.meta["request"].get("args", {}) == public_args tracker = TaskTracker(store) restored = await tracker.get(task.task_id, **owner) @@ -483,6 +483,7 @@ async def read(self, path): queue._async_agfs = backend store = PersistentTaskStore(MockLocalAGFS(root_path=tmp_path)) tracker = TaskTracker(store) + monkeypatch.setattr("openviking.service.task_tracker._instance", tracker) monkeypatch.setattr(external_task_service_module, "get_task_tracker", lambda: tracker) monkeypatch.setattr(external_task_service_module, "get_queue_manager", lambda: manager) monkeypatch.setattr( @@ -559,13 +560,21 @@ async def state(task): include_events=False, _ctx=ctx, ) - assert queried.result == (await state(first)).to_dict() + queried_task = dict(queried.result) + current_task = (await state(first)).to_dict() + assert queried_task.pop("processing_seconds") >= 0 + assert current_task.pop("processing_seconds") >= 0 + assert queried_task == current_task assert queried.result["meta"]["request"].get("args", {}) == public_args await queue.dequeue() # Same target rotates, including across users. assert (await state(waiting)).status == TaskStatus.PENDING assert waiting.task_id not in submitted - assert tracker.has_work(waiting.task_id) # Requeue survives ACK of the old delivery. + assert await tracker.has_work( + waiting.task_id, + account_id=waiting.account_id, + user_id=waiting.user_id, + ) # Requeue survives ACK of the old delivery. await queue.dequeue() assert (await state(independent)).status == TaskStatus.COMPLETED await tracker.cancel(cancelled.task_id, account_id="acct", user_id="alice") @@ -584,6 +593,7 @@ async def state(task): with pytest.raises(asyncio.CancelledError): await worker tracker = TaskTracker(store) + monkeypatch.setattr("openviking.service.task_tracker._instance", tracker) tasks = ExternalTaskService() tasks.register(CompileService(service._config, tasks, SimpleNamespace())) queue.set_dequeue_handler(ExternalTaskProcessor(tasks)) diff --git a/tests/server/test_compile_studio.py b/tests/server/test_compile_studio.py index 3ef5f044c9..c7423c0ec7 100644 --- a/tests/server/test_compile_studio.py +++ b/tests/server/test_compile_studio.py @@ -1,6 +1,7 @@ """Studio task pagination and submission retry contracts.""" import asyncio +import json import time from types import SimpleNamespace from unittest.mock import AsyncMock @@ -19,20 +20,20 @@ @pytest.mark.asyncio async def test_page_reads_past_200_and_isolates_owner(): - store = PersistentTaskStore(_FakeAgfs()) + agfs = _FakeAgfs() + store = PersistentTaskStore(agfs) tracker = TaskTracker(store) owner = {"account_id": "acme", "user_id": "alice"} for index in range(205): - record = await tracker.create( + await tracker.create( "compile", task_id=f"cmp_{index:03}", meta={"request": {"to": "viking://resources/wiki"}}, **owner, ) - record.created_at = 100.0 - await store.update(record) await tracker.create("compile", task_id="foreign", account_id="acme", user_id="bob") # Empty cache: pagination must read durable records without repopulating all tasks. + store = PersistentTaskStore(agfs) tracker = TaskTracker(store) filters = { "task_type": "compile", @@ -44,7 +45,7 @@ async def test_page_reads_past_200_and_isolates_owner(): seen = [] before = None while True: - page = await tracker.list_page(**owner, limit=30, before=before, **filters) + page = await tracker.list_tasks(**owner, limit=30, before=before, **filters) if not page: break seen.extend(task.task_id for task in page) @@ -52,7 +53,6 @@ async def test_page_reads_past_200_and_isolates_owner(): assert len(seen) == len(set(seen)) == 205 assert seen[0] == "cmp_204" assert "foreign" not in seen - assert not tracker._tasks def test_cursor_is_bound_to_identity_and_filters(): @@ -69,7 +69,8 @@ def test_cursor_is_bound_to_identity_and_filters(): async def test_concurrent_submission_retry_and_restart(monkeypatch): import openviking.service.external_task_service as module - store = PersistentTaskStore(_FakeAgfs()) + agfs = _FakeAgfs() + store = PersistentTaskStore(agfs) tracker = TaskTracker(store) monkeypatch.setattr(module, "get_task_tracker", lambda: tracker) monkeypatch.setattr(module, "get_queue_manager", lambda: SimpleNamespace(enqueue=AsyncMock())) @@ -109,7 +110,8 @@ async def test_http_contract_preserves_legacy_list_and_retries(monkeypatch): from openviking_cli.exceptions import OpenVikingError from openviking_cli.utils.config.open_viking_config import CompileApiConfig - store = PersistentTaskStore(_FakeAgfs()) + agfs = _FakeAgfs() + store = PersistentTaskStore(agfs) tracker = TaskTracker(store) monkeypatch.setattr(external, "get_task_tracker", lambda: tracker) monkeypatch.setattr(task_router, "get_task_tracker", lambda: tracker) @@ -189,6 +191,7 @@ async def on_error(request, error): assert retry_after_cancel.json()["result"]["task_id"] == task_id assert "submission_hash" not in retry_after_cancel.json()["result"]["meta"] ctx = RequestContext(user=UserIdentifier("acme", "alice"), role=Role.ROOT) + store = PersistentTaskStore(agfs) tracker = TaskTracker(store) root_page = await client.get( "/api/v1/tasks", params={"pagination": "cursor", "task_type": "compile"} @@ -224,7 +227,8 @@ async def on_error(request, error): @pytest.mark.asyncio async def test_cursor_prunes_cold_expired_records_without_loading_cache(): - store = PersistentTaskStore(_FakeAgfs()) + agfs = _FakeAgfs() + store = PersistentTaskStore(agfs) tracker = TaskTracker(store) owner = {"account_id": "acme", "user_id": "alice"} for task_id, status, days in [ @@ -235,12 +239,16 @@ async def test_cursor_prunes_cold_expired_records_without_loading_cache(): ("active", TaskStatus.RUNNING, 8), ]: task = await tracker.create("compile", task_id=task_id, **owner) - task.status = status - task.updated_at = time.time() - days * 86400 - await store.update(task) + path = f"/local/acme/_system/tasks/alice/{task.task_id}.json" + payload = json.loads(agfs.files[path]) + payload["status"] = status.value + payload["updated_at"] = time.time() - days * 86400 + agfs.files[path] = json.dumps(payload).encode() + store = PersistentTaskStore(agfs) tracker = TaskTracker(store, max_concurrent_store_io=1) + await tracker.delete_expired() page = await asyncio.wait_for( - tracker.list_page( + tracker.list_tasks( **owner, limit=30, task_type="compile", @@ -252,49 +260,41 @@ async def test_cursor_prunes_cold_expired_records_without_loading_cache(): timeout=3, ) assert {task.task_id for task in page} == {"recent-failed", "active"} - assert tracker.count() == 0 - assert {task["task_id"] for task in await store.list("acme", user_id="alice")} == { + assert {task.task_id for task in await store.list("acme", user_id="alice")} == { "recent-failed", "active", } @pytest.mark.asyncio -async def test_cursor_cleanup_rechecks_task_updated_after_scan(monkeypatch): - store = PersistentTaskStore(_FakeAgfs()) +async def test_cursor_cleanup_rechecks_task_updated_after_scan(): + agfs = _FakeAgfs() + store = PersistentTaskStore(agfs) tracker = TaskTracker(store, max_concurrent_store_io=1) owner = {"account_id": "acme", "user_id": "alice"} task = await tracker.create("compile", task_id="updated-after-scan", **owner) - task.status = TaskStatus.COMPLETED - task.updated_at = time.time() - 2 * 86400 - await store.update(task) + path = f"/local/acme/_system/tasks/alice/{task.task_id}.json" + payload = json.loads(agfs.files[path]) + payload["status"] = TaskStatus.COMPLETED.value + payload["updated_at"] = time.time() - 2 * 86400 + agfs.files[path] = json.dumps(payload).encode() + store = PersistentTaskStore(agfs) tracker = TaskTracker(store, max_concurrent_store_io=1) - prune = tracker._prune_persisted_expired - scanned = asyncio.Event() - - async def signal_before_prune(payload): - scanned.set() - return await prune(payload) - - monkeypatch.setattr(tracker, "_prune_persisted_expired", signal_before_prune) - async with tracker._task_locks.acquire(task.task_id): - reading = asyncio.create_task( - tracker.list_page( - **owner, - limit=30, - task_type="compile", - status=None, - resource_id=None, - include_internal=False, - q=None, - ) - ) - await asyncio.wait_for(scanned.wait(), timeout=3) - task.updated_at = time.time() - # A scan waiting for this task lock must not hold the sole I/O slot. - await asyncio.wait_for( - tracker._store_io.run("update", lambda: store.update(task)), timeout=3 - ) - page = await asyncio.wait_for(reading, timeout=3) + record_lock = ("acme", "alice", task.task_id) + async with store._record_locks.acquire(record_lock): + cleanup = asyncio.create_task(tracker.delete_expired()) + await asyncio.sleep(0) + payload["updated_at"] = time.time() + agfs.files[path] = json.dumps(payload).encode() + await asyncio.wait_for(cleanup, timeout=3) + page = await tracker.list_tasks( + **owner, + limit=30, + task_type="compile", + status=None, + resource_id=None, + include_internal=False, + q=None, + ) assert [record.task_id for record in page] == [task.task_id] - assert await store.get(task.task_id, **owner) is not None + assert await store.get_task(task.task_id, **owner) is not None diff --git a/tests/server/test_content_batch_write.py b/tests/server/test_content_batch_write.py index 8f71a95f01..97dfc6f932 100644 --- a/tests/server/test_content_batch_write.py +++ b/tests/server/test_content_batch_write.py @@ -9,7 +9,7 @@ from openviking.server.identity import RequestContext, Role from openviking.service.task_store import PersistentTaskStore from openviking.service.task_tracker import TaskTracker -from openviking.service.task_work_index import TaskWorkIndex, bind_task_context +from openviking.service.task_work_index import bind_task_context from openviking.session.memory.dataclass import MemoryFile from openviking.session.memory.utils import MemoryFileUtils from openviking.storage.content_write import ContentWriteCoordinator @@ -131,7 +131,6 @@ def _coordinator_queue_for_unit_tests(monkeypatch, request): ) monkeypatch.setattr(content_write_module, "get_queue_manager", lambda: queue_manager) tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) - tracker.attach_work_index(TaskWorkIndex()) monkeypatch.setattr("openviking.service.task_tracker._instance", tracker) async def run_tracked_operation(_task_type, *, operation, **_kwargs): diff --git a/tests/service/test_resource_service_watch.py b/tests/service/test_resource_service_watch.py index f35a9813e5..0d28137829 100644 --- a/tests/service/test_resource_service_watch.py +++ b/tests/service/test_resource_service_watch.py @@ -16,6 +16,7 @@ from openviking.server.identity import RequestContext, Role from openviking.service import resource_service as resource_service_module from openviking.service.resource_service import ResourceService +from openviking.service.task_domain import WorkSummary from openviking.service.task_tracker import TaskStatus from openviking.service.task_work_index import TASK_WORK_ID_FIELD from openviking.storage.content_write import ContentWriteCoordinator @@ -1026,9 +1027,9 @@ async def test_add_resource_processor_records_paused_watch_result( fail=AsyncMock(), get_task_auth=AsyncMock(return_value={}), wait_for_task_work=AsyncMock(), - build_queue_status=Mock(return_value={}), - get_context_count=Mock(return_value=0), - get_queue_timing=Mock(return_value={}), + get_work_summary=AsyncMock(return_value=WorkSummary()), + format_queue_status=Mock(return_value={}), + format_queue_timing=Mock(return_value={}), get=AsyncMock(return_value=SimpleNamespace(status=TaskStatus.CANCELLING)), ) monkeypatch.setattr( diff --git a/tests/service/test_skill_processing_cancellation.py b/tests/service/test_skill_processing_cancellation.py index eba7e505ff..0abca48f11 100644 --- a/tests/service/test_skill_processing_cancellation.py +++ b/tests/service/test_skill_processing_cancellation.py @@ -9,9 +9,9 @@ from openviking.server.identity import RequestContext, Role from openviking.service.resource_service import ResourceService -from openviking.service.task_store import PersistentTaskStore +from openviking.service.task_store import PersistentTaskStore, WorkRegistration from openviking.service.task_tracker import TaskStatus, TaskTracker -from openviking.service.task_work_index import QueueTaskMetadata, TaskWorkIndex, get_task_context +from openviking.service.task_work_index import get_task_context from openviking.telemetry import ( OperationTelemetry, bind_telemetry, @@ -38,7 +38,7 @@ async def test_skill_rollback_persists_cancellation_for_finished_work(monkeypatc restarted = TaskTracker(store=store) restored = await restarted.get(task.task_id, **owner) assert restored.status == TaskStatus.CANCELLED - assert restarted.is_cancellation_requested(task.task_id) + assert await restarted.is_cancellation_requested(task.task_id, **owner) if outcome == "fail": assert restored.error == "failed" @@ -46,21 +46,23 @@ async def test_skill_rollback_persists_cancellation_for_finished_work(monkeypatc @pytest.mark.asyncio async def test_skill_rollback_prevents_late_ack_failure_from_replaying_new_index(monkeypatch): tracker = TaskTracker(store=PersistentTaskStore(_FakeAgfs())) - work = TaskWorkIndex() - tracker.attach_work_index(work) monkeypatch.setattr("openviking.service.task_tracker.get_task_tracker", lambda: tracker) ctx = RequestContext(user=UserIdentifier.the_default_user(), role=Role.ROOT) owner = {"account_id": ctx.account_id, "user_id": ctx.user.user_id} task = await tracker.create("add_skill", **owner) - child = QueueTaskMetadata(task.task_id, "last-vector", ctx.account_id, ctx.user.user_id) - work.register("Embedding", child) + assert ( + await tracker.register_work(task.task_id, "last-vector", "Embedding", **owner) + is WorkRegistration.CREATED + ) await tracker.complete(task.task_id, {}, **owner) - await work.prepare_ack("Embedding", child) + await tracker.mark_work_done(task.task_id, "last-vector", **owner) assert (await tracker.get(task.task_id, **owner)).status == TaskStatus.COMPLETED await ResourceService.__new__(ResourceService).cancel_skill_processing(task.task_id, ctx) - work.rollback_ack("Embedding", child) - assert work.cancellation_requested(task.task_id) - assert not work.register("Embedding", child) + assert await tracker.is_cancellation_requested(task.task_id, **owner) + assert ( + await tracker.register_work(task.task_id, "last-vector", "Embedding", **owner) + is WorkRegistration.REJECTED + ) @pytest.mark.asyncio @@ -69,10 +71,10 @@ async def test_cancel_during_skill_task_creation_settles_persisted_record(monkey started, resume = asyncio.Event(), asyncio.Event() original_create = store.create - async def slow_create(task): + async def slow_create(task, **kwargs): started.set() await resume.wait() - await original_create(task) + await original_create(task, **kwargs) monkeypatch.setattr(store, "create", slow_create) tracker = TaskTracker(store=store) @@ -97,8 +99,8 @@ async def slow_create(task): await asyncio.wait_for(adding, 1) record = await tracker.get("create-test", **owner) assert record.status == TaskStatus.FAILED - assert (await store.get("create-test", **owner))["status"] == "failed" - assert not tracker.has_work(record.task_id) + assert (await store.get_task("create-test", **owner)).status is TaskStatus.FAILED + assert not await tracker.has_work(record.task_id, **owner) service._skill_processor.process_skill.assert_not_awaited() finally: resume.set() @@ -108,17 +110,17 @@ async def slow_create(task): @pytest.mark.asyncio async def test_skill_success_remains_completed_when_queue_work_finishes_first(monkeypatch): tracker = TaskTracker(store=PersistentTaskStore(_FakeAgfs())) - index = TaskWorkIndex() - tracker.attach_work_index(index) monkeypatch.setattr("openviking.service.task_tracker.get_task_tracker", lambda: tracker) ctx = RequestContext(user=UserIdentifier.the_default_user(), role=Role.ROOT) owner = {"account_id": ctx.account_id, "user_id": ctx.user.user_id} async def process_skill(**kwargs): - work = QueueTaskMetadata(get_task_context().task_id, "vector", **owner) - assert index.register("Embedding", work) - index.record_result(work, "completed") - await index.prepare_ack("Embedding", work) + task_id = get_task_context().task_id + assert ( + await tracker.register_work(task_id, "vector", "Embedding", **owner) + is WorkRegistration.CREATED + ) + await tracker.mark_work_done(task_id, "vector", **owner) return {"uri": "viking://agent/skills/demo"} service = ResourceService.__new__(ResourceService) diff --git a/tests/service/test_task_processing_time.py b/tests/service/test_task_processing_time.py index b7d138884a..db78581cee 100644 --- a/tests/service/test_task_processing_time.py +++ b/tests/service/test_task_processing_time.py @@ -7,70 +7,99 @@ import pytest -from openviking.service.task_work_index import ( - QueueTaskMetadata, - TaskWorkIndex, - bind_task_context, -) +from openviking.service.task_domain import WorkRecord, processing_seconds +from openviking.service.task_store import PersistentTaskStore +from openviking.service.task_tracker import TaskTracker +from openviking.service.task_work_index import bind_task_context from openviking.storage.queuefs.semantic_executor import ( SemanticTreeExecutor, SemanticTreeScheduler, ) +from tests.test_task_tracker import _FakeAgfs + +OWNER = {"account_id": "account", "user_id": "user"} + + +def test_work_timestamps_union_worker_occupancy_and_exclude_queue_time(): + works = [ + WorkRecord( + work_id="first", + task_id="task", + queue_name="Semantic", + started_at=10, + finished_at=15, + ), + WorkRecord( + work_id="second", + task_id="task", + queue_name="Embedding", + started_at=12, + finished_at=18, + ), + WorkRecord( + work_id="third", + task_id="task", + queue_name="Inline", + started_at=40, + finished_at=43, + ), + ] + + assert processing_seconds(works, now=100) == 11 + + +async def _create_work( + tracker: TaskTracker, + task_id: str, + work_id: str = "worker", +) -> None: + task = await tracker.create("add_resource", task_id=task_id, **OWNER) + await tracker.start(task.task_id, **OWNER) + await tracker.register_work(task_id, work_id, "Semantic", **OWNER) + await tracker.start_work(task_id, work_id, **OWNER) -def test_work_timestamps_union_worker_occupancy_and_exclude_queue_time(monkeypatch): +@pytest.mark.asyncio +async def test_requeued_work_restarts_execution_time(monkeypatch): now = [0.0] - monkeypatch.setattr("openviking.service.task_work_index.time.time", lambda: now[0]) - index = TaskWorkIndex() - index.init_processing("task") - first = QueueTaskMetadata("task", "first") - second = QueueTaskMetadata("task", "second") - assert index.register("Semantic", first) - assert index.processing_seconds("task") is None - - now[0] = 10 # Initial queue wait. - first_worker = object() - index.register_active("task", first_worker, work_id=first.work_id) - now[0] = 12 - assert index.register("Embedding", second) - second_worker = object() - index.register_active("task", second_worker, work_id=second.work_id) - now[0] = 15 - index.unregister_active("task", first_worker) - now[0] = 18 - index.unregister_active("task", second_worker) - assert index.processing_seconds("task") == 8 - now[0] = 40 # No worker: waiting for another queue. - assert index.processing_seconds("task") == 8 - third_worker = object() - index.register_active("task", third_worker) - now[0] = 43 - assert index.processing_seconds("task") == 11 - index.unregister_active("task", third_worker) - now[0] = 100 - assert index.processing_seconds("task") == 11 + monkeypatch.setattr("openviking.service.task_store.time.time", lambda: now[0]) + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) + + await _create_work(tracker, "task") + now[0] = 30 + assert (await tracker.get("task", **OWNER)).processing_seconds == 30 + + await tracker.mark_work_requeued("task", "worker", **OWNER) + assert (await tracker.get("task", **OWNER)).processing_seconds is None + + now[0] = 50 + await tracker.start_work("task", "worker", **OWNER) + now[0] = 55 + assert (await tracker.get("task", **OWNER)).processing_seconds == 5 + + await tracker.mark_work_done("task", "worker", **OWNER) + assert (await tracker.get("task", **OWNER)).processing_seconds == 5 @pytest.mark.asyncio async def test_saturated_dag_scheduler_counts_claimed_handler_time(monkeypatch): now = [0.0] - monkeypatch.setattr( - "openviking.service.task_work_index.time", - SimpleNamespace(time=lambda: now[0]), - ) - index = TaskWorkIndex() + monkeypatch.setattr("openviking.service.task_store.time.time", lambda: now[0]) + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) scheduler = SemanticTreeScheduler(max_workers=1) monkeypatch.setattr( "openviking.storage.queuefs.semantic_executor.get_semantic_tree_scheduler", lambda _: scheduler, ) - monkeypatch.setattr("openviking.storage.queuefs.semantic_executor.get_viking_fs", lambda: None) + monkeypatch.setattr( + "openviking.storage.queuefs.semantic_executor.get_viking_fs", lambda: None + ) entered = {name: asyncio.Event() for name in ["a", "b"]} release = {name: asyncio.Event() for name in ["a", "b"]} queued = asyncio.Event() executors = {} for name in ["a", "b"]: - index.init_processing(name) + await _create_work(tracker, name) with bind_task_context(name, "account", "user"): executor = SemanticTreeExecutor(None, "resource", 1, None) executors[name] = executor @@ -90,30 +119,25 @@ def queue_b(*args, **kwargs): monkeypatch.setattr(executors["b"], "_schedule_dir", queue_b) async def coordinate(name): - worker = asyncio.current_task() - assert worker is not None - index.register_active(name, worker) try: await executors[name].run("viking://resources/" + name) finally: - index.unregister_active(name, worker) + await tracker.mark_work_done(name, "worker", **OWNER) first = asyncio.create_task(coordinate("a")) await entered["a"].wait() second = asyncio.create_task(coordinate("b")) await queued.wait() now[0] = 20 - assert index.processing_seconds("a") == 20 - # B's QueueFS handler already owns a worker while its internal DAG node - # waits for the shared scheduler. - assert index.processing_seconds("b") == 20 + assert (await tracker.get("a", **OWNER)).processing_seconds == 20 + assert (await tracker.get("b", **OWNER)).processing_seconds == 20 release["a"].set() await entered["b"].wait() now[0] = 25 - assert index.processing_seconds("b") == 25 + assert (await tracker.get("b", **OWNER)).processing_seconds == 25 release["b"].set() await asyncio.gather(first, second) - assert index.processing_seconds("b") == 25 + assert (await tracker.get("b", **OWNER)).processing_seconds == 25 for worker in list(scheduler._workers): worker.cancel() await asyncio.gather(*scheduler._workers, return_exceptions=True) @@ -122,31 +146,25 @@ async def coordinate(name): @pytest.mark.asyncio async def test_downstream_work_wait_remains_worker_occupied_time(monkeypatch): now = [0.0] - monkeypatch.setattr( - "openviking.service.task_work_index.time", - SimpleNamespace(time=lambda: now[0]), - ) - index = TaskWorkIndex() - index.init_processing("task") - metadata = QueueTaskMetadata("task", "embedding") - assert index.register("Embedding", metadata) - worker = asyncio.current_task() - assert worker is not None - index.register_active("task", worker) - try: - now[0] = 3 - waiter = asyncio.create_task( - index.wait_for_work("task", ["embedding"], poll_interval=0.001) + monkeypatch.setattr("openviking.service.task_store.time.time", lambda: now[0]) + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) + await _create_work(tracker, "task", "coordinator") + await tracker.register_work("task", "embedding", "Embedding", **OWNER) + + now[0] = 3 + waiter = asyncio.create_task( + tracker.wait_for_work( + "task", ["embedding"], poll_interval=0.001, **OWNER ) - await asyncio.sleep(0.01) - now[0] = 20 - assert index.processing_seconds("task") == 20 - index.record_result(metadata, "completed") - await waiter - now[0] = 22 - assert index.processing_seconds("task") == 22 - finally: - index.unregister_active("task", worker) + ) + await asyncio.sleep(0.01) + now[0] = 20 + assert (await tracker.get("task", **OWNER)).processing_seconds == 20 + await tracker.mark_work_done("task", "embedding", **OWNER) + await waiter + now[0] = 22 + assert (await tracker.get("task", **OWNER)).processing_seconds == 22 + await tracker.mark_work_done("task", "coordinator", **OWNER) @pytest.mark.asyncio @@ -157,19 +175,18 @@ async def test_semantic_retry_cooldown_counts_full_worker_occupied_time( from openviking.storage.queuefs.semantic_processor import SemanticProcessor now = [0.0] - monkeypatch.setattr( - "openviking.service.task_work_index.time", - SimpleNamespace(time=lambda: now[0]), + monkeypatch.setattr("openviking.service.task_store.time.time", lambda: now[0]) + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) + await _create_work(tracker, "task") + message = SimpleNamespace( + uri="viking://resources/retry-test", account_id="account" ) - index = TaskWorkIndex() - index.init_processing("task") - message = SimpleNamespace(uri="viking://resources/retry-test", account_id="account") enqueued = [] async def sleep(delay): assert delay == 30 now[0] += delay - assert index.processing_seconds("task") == 33 + assert (await tracker.get("task", **OWNER)).processing_seconds == 33 async def enqueue(msg): assert msg is message @@ -190,15 +207,11 @@ async def custom(queue_arg, msg): processor = SimpleNamespace( _account_breaker=lambda _account_id: SimpleNamespace(retry_after=30) ) - worker = asyncio.current_task() - index.register_active("task", worker) - try: - now[0] = 3 - await SemanticProcessor._reenqueue_semantic_msg( - processor, message, enqueue=custom if custom_enqueue else None - ) - finally: - index.unregister_active("task", worker) + now[0] = 3 + await SemanticProcessor._reenqueue_semantic_msg( + processor, message, enqueue=custom if custom_enqueue else None + ) + await tracker.mark_work_done("task", "worker", **OWNER) assert enqueued == [message] - assert index.processing_seconds("task") == 35 + assert (await tracker.get("task", **OWNER)).processing_seconds == 35 diff --git a/tests/service/test_task_queue_middleware.py b/tests/service/test_task_queue_middleware.py index aadc545837..ba447addd6 100644 --- a/tests/service/test_task_queue_middleware.py +++ b/tests/service/test_task_queue_middleware.py @@ -1,18 +1,19 @@ # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. # SPDX-License-Identifier: AGPL-3.0 -"""Legacy task-index behavior through the middleware boundary.""" +"""Task Store behavior through the QueueFS middleware boundary.""" import asyncio import json -import threading from unittest.mock import AsyncMock import pytest from openviking.pyagfs import AsyncAGFSClient +from openviking.service.task_domain import TaskStatus, WorkState from openviking.service.task_queue_middleware import TaskWorkQueueMiddleware +from openviking.service.task_store import PersistentTaskStore +from openviking.service.task_tracker import TaskTracker, set_task_tracker from openviking.service.task_work_index import ( - TaskWorkIndex, TaskWorkRejected, bind_task_context, extract_task_metadata, @@ -20,42 +21,51 @@ ) from openviking.storage.queuefs.named_queue import DequeueHandlerBase, NamedQueue from openviking.storage.queuefs.process_result import ProcessOutcome, ProcessResult -from openviking.storage.queuefs.queue_middleware import QueueMiddleware -from openviking.storage.queuefs.session_commit_processor import SessionCommitProcessor from openviking.telemetry import OperationTelemetry, bind_telemetry from openviking.utils.log_correlation import log_correlation +from tests.test_task_tracker import _FakeAgfs + +OWNER = {"account_id": "account", "user_id": "user"} @pytest.fixture -def tracked_queue(monkeypatch): +async def tracked_queue(monkeypatch): transport = AsyncMock(spec=AsyncAGFSClient) transport.write.return_value = "message-1" - # Tests override transport.read.return_value to feed /messages snapshots; - # status reads must still return the backend JSON shape. transport.read.return_value = b"0" async def read(path): if path.endswith("/status"): return b'{"pending":0,"processing":0}' - value = transport.read.return_value - return value + return transport.read.return_value transport.read.side_effect = read monkeypatch.setattr( "openviking.storage.queuefs.named_queue.AsyncAGFSClient", lambda _client: transport ) - index = TaskWorkIndex() - finalize = AsyncMock() - cancelled = set() - index.set_callbacks( - finalize_before_ack=finalize, - is_cancellation_requested=lambda task_id: task_id in cancelled, + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) + set_task_tracker(tracker) + queue = NamedQueue( + object(), + "/queue", + "Embedding", + middlewares=[TaskWorkQueueMiddleware()], ) - queue = NamedQueue(object(), "/queue", "Test", middlewares=[TaskWorkQueueMiddleware(index)]) - return queue, index, transport, finalize, cancelled + try: + yield queue, tracker, transport + finally: + set_task_tracker(None) -async def enqueue_task(queue, transport, task_id="task-1"): +async def create_task(tracker: TaskTracker, task_id: str = "task-1") -> None: + task = await tracker.get(task_id, **OWNER) + if task is None: + task = await tracker.create("add_resource", task_id=task_id, **OWNER) + await tracker.start(task.task_id, **OWNER) + + +async def enqueue_task(queue, tracker, transport, task_id="task-1"): + await create_task(tracker, task_id) with bind_task_context(task_id, "account", "user"): await queue.enqueue({"value": 1}) return {"id": "message-1", "data": transport.write.await_args.args[1].decode()} @@ -64,10 +74,7 @@ async def enqueue_task(queue, transport, task_id="task-1"): def test_log_correlation_uses_bound_task_and_telemetry() -> None: telemetry = OperationTelemetry(operation="add_resource_job", enabled=True) - with ( - bind_task_context("task-123", "account", "user"), - bind_telemetry(telemetry), - ): + with bind_task_context("task-123", "account", "user"), bind_telemetry(telemetry): assert log_correlation(message_id="semantic-456") == ( f"task_id=task-123 telemetry_id={telemetry.telemetry_id} message_id=semantic-456" ) @@ -80,35 +87,34 @@ def test_log_correlation_accepts_explicit_queue_ids_without_context() -> None: ) -async def test_enqueue_registers_before_write_and_ack_finalizes_before_delete(tracked_queue): - queue, index, transport, finalize, _ = tracked_queue +async def test_enqueue_registers_work_before_transport_write(tracked_queue): + queue, tracker, transport = tracked_queue + await create_task(tracker) async def write(path, data): assert path.endswith("/enqueue") - assert index.has_work("task-1") + assert await tracker.has_work("task-1", **OWNER) return "message-1" transport.write.side_effect = write - message = await enqueue_task(queue, transport) + message = await enqueue_task(queue, tracker, transport) metadata = extract_task_metadata(message) + + assert metadata is not None assert metadata.task_id == "task-1" assert metadata.account_id == "account" assert metadata.user_id == "user" + assert await tracker.get_work_state( + metadata.task_id, metadata.work_id, **OWNER + ) is WorkState.PENDING assert get_task_context() is None - async def ack_write(path, data): - assert path.endswith("/ack") - assert not index.has_work("task-1") - finalize.assert_awaited_once_with(metadata) - return "" - - transport.write.side_effect = ack_write - await queue.ack("message-1", message) - assert not index.has_work("task-1") - async def test_enqueue_keeps_content_owner_separate_from_task_owner(tracked_queue): - queue, _, transport, _, _ = tracked_queue + queue, tracker, transport = tracked_queue + task_owner = {"account_id": "account", "user_id": "task-owner"} + task = await tracker.create("add_resource", task_id="task-1", **task_owner) + await tracker.start(task.task_id, **task_owner) payload = { "task_id": "task-1", "account_id": "account", @@ -134,458 +140,196 @@ async def test_enqueue_keeps_content_owner_separate_from_task_owner(tracked_queu assert persisted["_task_account_id"] == "account" assert persisted["_task_user_id"] == "task-owner" assert extract_task_metadata(persisted).user_id == "task-owner" + metadata = extract_task_metadata(persisted) + assert await tracker.get_work_state( + metadata.task_id, + metadata.work_id, + **task_owner, + ) is WorkState.PENDING async def test_enqueue_does_not_stamp_task_owner_without_authoritative_context(tracked_queue): - queue, _, transport, _, _ = tracked_queue + queue, _, transport = tracked_queue + transport.write.reset_mock() with bind_task_context("task-1", "", ""): - await queue.enqueue( - { - "task_id": "task-1", - "account_id": "account", - "user_id": "content-owner", - } - ) + with pytest.raises(TaskWorkRejected, match="missing owner metadata"): + await queue.enqueue( + { + "task_id": "task-1", + "account_id": "account", + "user_id": "content-owner", + } + ) - persisted = json.loads(transport.write.await_args.args[1]) - assert persisted["account_id"] == "account" - assert persisted["user_id"] == "content-owner" - assert "_task_account_id" not in persisted - assert "_task_user_id" not in persisted + transport.write.assert_not_awaited() async def test_enqueue_failure_records_failed_work(tracked_queue): - queue, index, transport, finalize, _ = tracked_queue + queue, tracker, transport = tracked_queue + await create_task(tracker) transport.write.side_effect = OSError("write failed") enqueued_work = [] + with bind_task_context( - "task-1", - "account", - "user", - enqueued_work=enqueued_work, - ): - with pytest.raises(OSError, match="write failed"): - await queue.enqueue({}) - assert not index.has_work("task-1") + "task-1", "account", "user", enqueued_work=enqueued_work + ), pytest.raises(OSError, match="write failed"): + await queue.enqueue({}) + payload = json.loads(transport.write.await_args.args[1]) - assert enqueued_work == [("Test", payload["_task_work_id"])] - assert index.queue_status("task-1")["Test"] == { - "processed": 0, - "requeue_count": 0, - "error_count": 1, - "errors": [{"message": "write failed"}], - } - finalize.assert_awaited_once() + assert enqueued_work == [("Embedding", payload["_task_work_id"])] + status = await tracker.build_queue_status("task-1", **OWNER) + assert status["Embedding"]["error_count"] == 1 + task = await tracker.get("task-1", **OWNER) + assert task is not None and task.status is TaskStatus.FAILED + assert task.error == "write failed" -async def test_cancelled_enqueue_discards_registered_work(tracked_queue): - queue, index, transport, finalize, _ = tracked_queue +async def test_cancelled_enqueue_discards_uncommitted_work(tracked_queue): + queue, tracker, transport = tracked_queue + await create_task(tracker) transport.write.side_effect = asyncio.CancelledError() - with bind_task_context("task-1", "account", "user"): - with pytest.raises(asyncio.CancelledError): - await queue.enqueue({}) - assert not index.has_work("task-1") - assert index.queue_status("task-1") == {} - finalize.assert_awaited_once() - -async def test_cancelled_task_rejects_enqueue_without_write(tracked_queue): - queue, index, transport, _, cancelled = tracked_queue - cancelled.add("task-1") with bind_task_context("task-1", "account", "user"): - with pytest.raises(TaskWorkRejected, match="cancelling"): - await queue.enqueue({}) - transport.write.assert_not_awaited() - assert not index.has_work("task-1") - - -@pytest.mark.parametrize("operation", ["ack", "clear"]) -@pytest.mark.parametrize("failure", [OSError("write failed"), asyncio.CancelledError()]) -async def test_failed_ack_or_clear_restores_work(tracked_queue, operation, failure): - queue, index, transport, finalize, _ = tracked_queue - message = await enqueue_task(queue, transport) - transport.read.return_value = json.dumps([message]).encode() - transport.write.reset_mock() - transport.write.side_effect = failure - - async def invoke(): - if operation == "ack": - return await queue.ack("message-1", message) - return await queue.clear() - - if isinstance(failure, asyncio.CancelledError): with pytest.raises(asyncio.CancelledError): - await invoke() - elif operation == "clear": - assert await invoke() is False - else: - assert await invoke() is None - finalize.assert_awaited_once() - assert index.has_work("task-1") - - -@pytest.mark.parametrize("operation", ["ack", "clear"]) -async def test_finalization_failure_prevents_transport_delete(tracked_queue, operation): - queue, index, transport, finalize, _ = tracked_queue - message = await enqueue_task(queue, transport) - transport.read.return_value = json.dumps([message]).encode() - transport.write.reset_mock() - finalize.side_effect = OSError("task store unavailable") - if operation == "ack": - await queue.ack("message-1", message) - else: - assert await queue.clear() is False - transport.write.assert_not_awaited() - assert index.has_work("task-1") - + await queue.enqueue({}) -@pytest.mark.parametrize("operation", ["enqueue", "ack", "clear"]) -async def test_inner_short_circuit_compensates_uncommitted_index_change(tracked_queue, operation): - queue, index, transport, _, _ = tracked_queue + assert not await tracker.has_work("task-1", **OWNER) + task = await tracker.get("task-1", **OWNER) + assert task is not None and task.status is TaskStatus.FAILED + assert task.error == "Task operation interrupted before its result was persisted" - class ShortCircuit(QueueMiddleware): - async def invoke(self, ctx, call_next): - return "deduplicated" - setattr(ShortCircuit, operation, ShortCircuit.invoke) - queue = NamedQueue( - object(), "/queue", "Test", middlewares=[TaskWorkQueueMiddleware(index), ShortCircuit()] - ) - message = await enqueue_task(queue, transport) if operation != "enqueue" else None +async def test_terminal_task_rejects_enqueue_without_write(tracked_queue): + queue, tracker, transport = tracked_queue + await create_task(tracker) + await tracker.complete("task-1", {"done": True}, **OWNER) transport.write.reset_mock() - if operation == "enqueue": - with bind_task_context("task-1", "account", "user"): - assert await queue.enqueue({}) == "deduplicated" - assert not index.has_work("task-1") - elif operation == "ack": - await queue.ack("message-1", message) - assert index.has_work("task-1") - else: - transport.read.return_value = json.dumps([message]).encode() - await queue.clear() - assert index.has_work("task-1") - transport.write.assert_not_awaited() - -@pytest.mark.parametrize("operation", ["enqueue", "ack", "clear"]) -async def test_post_commit_error_does_not_undo_index_change(tracked_queue, operation): - queue, index, transport, _, _ = tracked_queue - - class FailAfter(QueueMiddleware): - async def invoke(self, ctx, call_next): - await call_next(ctx) - raise RuntimeError("after commit") + with bind_task_context("task-1", "account", "user"): + with pytest.raises(TaskWorkRejected, match="unavailable"): + await queue.enqueue({}) - setattr(FailAfter, operation, FailAfter.invoke) - queue = NamedQueue( - object(), "/queue", "Test", middlewares=[TaskWorkQueueMiddleware(index), FailAfter()] - ) - if operation == "enqueue": - with bind_task_context("task-1", "account", "user"): - with pytest.raises(RuntimeError, match="after commit"): - await queue.enqueue({}) - assert index.has_work("task-1") - else: - message = await enqueue_task(queue, transport) - transport.read.return_value = json.dumps([message]).encode() - if operation == "ack": - await queue.ack("message-1", message) - else: - assert await queue.clear() is False - assert not index.has_work("task-1") - - -async def test_clear_rolls_back_all_prepared_messages(tracked_queue): - queue, index, transport, finalize, _ = tracked_queue - first = await enqueue_task(queue, transport, "first") - second = await enqueue_task(queue, transport, "second") - transport.read.return_value = json.dumps([first, second]).encode() - - async def finish(metadata): - if metadata.task_id == "second": - raise OSError("second task failed") - - finalize.side_effect = finish - transport.write.reset_mock() - assert await queue.clear() is False transport.write.assert_not_awaited() - assert index.has_work("first") and index.has_work("second") - finalize.side_effect = None - assert await queue.clear() - assert not index.has_work("first") and not index.has_work("second") -async def test_process_result_tracks_children_and_errors(tracked_queue): - queue, index, transport, _, _ = tracked_queue - message = await enqueue_task(queue, transport) - contexts = [] - correlations = [] +async def test_successful_process_records_context_and_vector_count(tracked_queue): + queue, tracker, transport = tracked_queue + message = await enqueue_task(queue, tracker, transport) + metadata = extract_task_metadata(message) + seen = [] class Handler(DequeueHandlerBase): async def on_dequeue(self, data): - with bind_task_context("task-1", "account", "user"): - contexts.append(get_task_context()) - correlations.append(log_correlation(telemetry_id="tm-1", message_id="semantic-1")) - await queue.enqueue({"child": True}) - return ProcessResult.failed("failed work") + seen.append(get_task_context()) + return ProcessResult.success({"vector_written": True}) + await tracker.complete("task-1", {"done": True}, **OWNER) queue.set_dequeue_handler(Handler()) - assert (await queue.process_dequeued(message)).outcome is ProcessOutcome.FAILED - assert contexts[0].task_id == "task-1" - assert contexts[0].work_id == extract_task_metadata(message).work_id - assert correlations == ["task_id=task-1 telemetry_id=tm-1 message_id=semantic-1"] + result = await queue.process_dequeued(message) + + assert result.outcome is ProcessOutcome.SUCCESS + assert seen[0].task_id == "task-1" + assert seen[0].work_id == metadata.work_id assert get_task_context() is None - assert index.failure("task-1") == "failed work" - child_payload = next( - call.args[1] - for call in reversed(transport.write.await_args_list) - if call.args[0].endswith("/enqueue") - ) - child = {"id": "child", "data": child_payload.decode()} - assert index.has_work("task-1") - await queue.ack("message-1", message) - assert index.has_work("task-1") - await queue.ack("child", child) - assert not index.has_work("task-1") - status = await queue.get_status() - assert status.error_count == 1 - assert status.in_progress == 0 + assert await tracker.get_context_count("task-1", **OWNER) == 1 + task = await tracker.get("task-1", **OWNER) + assert task is not None and task.status is TaskStatus.COMPLETED -async def test_handler_exception_records_task_failure_and_propagates(tracked_queue): - queue, index, transport, _, _ = tracked_queue - message = await enqueue_task(queue, transport) +async def test_failed_process_preserves_multiple_errors(tracked_queue): + queue, tracker, transport = tracked_queue + message = await enqueue_task(queue, tracker, transport) class Handler(DequeueHandlerBase): async def on_dequeue(self, data): - raise RuntimeError("handler failed") + return ProcessResult.failed("combined", errors=["first", "second"]) + await tracker.complete("task-1", {"done": True}, **OWNER) queue.set_dequeue_handler(Handler()) - with pytest.raises(RuntimeError, match="handler failed"): - await queue.process_dequeued(message) - assert index.failure("task-1") == "handler failed" - assert index.has_work("task-1") - + result = await queue.process_dequeued(message) -async def test_concurrent_results_update_each_task(tracked_queue): - queue, index, transport, _, _ = tracked_queue - first = await enqueue_task(queue, transport, "first") - second = await enqueue_task(queue, transport, "second") - both_started = asyncio.Event() - started = 0 + assert result.outcome is ProcessOutcome.FAILED + status = await tracker.build_queue_status("task-1", **OWNER) + assert status["Embedding"]["error_count"] == 2 + task = await tracker.get("task-1", **OWNER) + assert task is not None and task.status is TaskStatus.FAILED + assert task.error == "combined" - class Handler(DequeueHandlerBase): - async def on_dequeue(self, data): - nonlocal started - started += 1 - if started == 2: - both_started.set() - await both_started.wait() - task_id = get_task_context().task_id - return ProcessResult.failed(f"{task_id} failed") - queue.set_dequeue_handler(Handler()) - await asyncio.wait_for( - asyncio.gather(queue.process_dequeued(first), queue.process_dequeued(second)), - timeout=1, - ) - assert index.failure("first") == "first failed" - assert index.failure("second") == "second failed" - - -async def test_cancel_before_handler_uses_legacy_discard_callback(tracked_queue): - queue, index, transport, _, cancelled = tracked_queue - message = await enqueue_task(queue, transport) - cancelled.add("task-1") - calls = [] +async def test_requeue_keeps_work_open_for_redelivery(tracked_queue): + queue, tracker, transport = tracked_queue + message = await enqueue_task(queue, tracker, transport) + outcomes = [ProcessResult.requeued(), ProcessResult.success()] class Handler(DequeueHandlerBase): async def on_dequeue(self, data): - pytest.fail("cancelled task reached handler") - - async def on_cancelled(self, data): - calls.append(data) - return await super().on_cancelled(data) + return outcomes.pop(0) queue.set_dequeue_handler(Handler()) - assert (await queue.process_dequeued(message)).outcome is ProcessOutcome.CANCELLED - assert calls == [message] - await queue.ack("message-1", message) - assert not index.has_work("task-1") - status = await queue.get_status() - assert status.processed == 1 and status.in_progress == 0 - - -@pytest.mark.parametrize("user_cancel", [True, False]) -async def test_active_cancellation_distinguishes_user_from_shutdown(tracked_queue, user_cancel): - queue, index, transport, _, cancelled = tracked_queue - message = await enqueue_task(queue, transport) + first = await queue.process_dequeued(message) metadata = extract_task_metadata(message) - started = asyncio.Event() - - class Handler(DequeueHandlerBase): - async def on_dequeue(self, data): - started.set() - await asyncio.Event().wait() + assert first.outcome is ProcessOutcome.REQUEUED + assert await tracker.get_work_state( + "task-1", metadata.work_id, **OWNER + ) is WorkState.PENDING + assert await tracker.has_work("task-1", **OWNER) - queue.set_dequeue_handler(Handler()) - running = asyncio.create_task(queue.process_dequeued(message)) - await asyncio.wait_for(started.wait(), timeout=1) - try: - if user_cancel: - cancelled.add("task-1") - index.cancel_active("task-1") - result = await asyncio.wait_for(running, timeout=1) - assert result.outcome is ProcessOutcome.CANCELLED - await index.wait_for_work("task-1", [metadata.work_id], timeout=0.1) - await queue.ack("message-1", message) - assert not index.has_work("task-1") - assert (await queue.get_status()).processed == 1 - else: - running.cancel() - with pytest.raises(asyncio.CancelledError): - await running - assert index.has_work("task-1") - finally: - if not running.done(): - running.cancel() - await asyncio.gather(running, return_exceptions=True) + await tracker.complete("task-1", {"done": True}, **OWNER) + second = await queue.process_dequeued(message) + assert second.outcome is ProcessOutcome.SUCCESS + assert not await tracker.has_work("task-1", **OWNER) -async def test_shared_middleware_keeps_concurrent_contexts_isolated(tracked_queue): - queue, index, transport, _, _ = tracked_queue - first = await enqueue_task(queue, transport, "first") - second = await enqueue_task(queue, transport, "second") - both_started = asyncio.Event() - seen = [] +async def test_clear_settles_committed_work(tracked_queue): + queue, tracker, transport = tracked_queue + message = await enqueue_task(queue, tracker, transport) + transport.read.return_value = json.dumps([message]).encode() + transport.write.reset_mock() + transport.write.return_value = "" - class Handler(DequeueHandlerBase): - async def on_dequeue(self, data): - before = get_task_context().task_id - seen.append(before) - if len(seen) == 2: - both_started.set() - await both_started.wait() - assert get_task_context().task_id == before - return ProcessResult.success(data) + assert await queue.clear() is True - queue.set_dequeue_handler(Handler()) - await asyncio.wait_for( - asyncio.gather(queue.process_dequeued(first), queue.process_dequeued(second)), timeout=1 - ) - assert set(seen) == {"first", "second"} - assert get_task_context() is None - await queue.ack("first", first) - await queue.ack("second", second) - assert not index.has_work("first") and not index.has_work("second") + transport.write.assert_awaited_once_with("/queue/Embedding/clear", b"") + assert not await tracker.has_work("task-1", **OWNER) -async def test_awaited_child_result_is_recorded_before_ack(tracked_queue): - queue, index, transport, finalize, _ = tracked_queue - message = await enqueue_task(queue, transport) +async def test_cancelled_task_uses_handler_cancellation_path(tracked_queue): + queue, tracker, transport = tracked_queue + message = await enqueue_task(queue, tracker, transport) + calls = [] class Handler(DequeueHandlerBase): async def on_dequeue(self, data): - async def child(): - await asyncio.sleep(0) - return ProcessResult.failed("child failure") - - return await asyncio.create_task(child()) + pytest.fail("cancelled task reached handler") - async def finish(metadata): - assert index.failure(metadata.task_id) == "child failure" + async def on_cancelled(self, data): + calls.append(data) + return ProcessResult.cancelled() - finalize.side_effect = finish queue.set_dequeue_handler(Handler()) - await queue.process_dequeued(message) - finalize.assert_not_awaited() - await queue.ack("message-1", message) - finalize.assert_awaited_once() + await tracker.cancel("task-1", **OWNER) + result = await queue.process_dequeued(message) + + assert result.outcome is ProcessOutcome.CANCELLED + assert calls == [message] + task = await tracker.get("task-1", **OWNER) + assert task is not None and task.status is TaskStatus.CANCELLED -async def test_cancel_cleanup_interruption_does_not_ack(tracked_queue): - queue, index, transport, _, cancelled = tracked_queue - message = await enqueue_task(queue, transport) - cancelled.add("task-1") - transport.write.reset_mock() +async def test_terminal_redelivery_is_idempotent(tracked_queue): + queue, tracker, transport = tracked_queue + message = await enqueue_task(queue, tracker, transport) class Handler(DequeueHandlerBase): async def on_dequeue(self, data): - pytest.fail("cancelled work must not execute") - - async def on_cancelled(self, data): - raise asyncio.CancelledError + return ProcessResult.success() + await tracker.complete("task-1", {"done": True}, **OWNER) queue.set_dequeue_handler(Handler()) - with pytest.raises(asyncio.CancelledError): - await queue.process_dequeued(message) - assert index.has_work("task-1") - transport.write.assert_not_awaited() - - -async def test_cross_loop_cancellation_waits_for_handler_cleanup(tracked_queue): - queue, index, transport, _, cancelled = tracked_queue - started = threading.Event() - cleanup_started = threading.Event() - cleanup_release = threading.Event() - request_loop = asyncio.get_running_loop() - - class Session: - async def exists(self): - return True - - async def load(self): - pass - - async def resume_queued_commit(self, msg): - assert asyncio.get_running_loop() is not request_loop - started.set() - try: - await asyncio.Event().wait() - finally: - cleanup_started.set() - await asyncio.to_thread(cleanup_release.wait) - - class Service: - def session(self, ctx, session_id, session_uri=None): - return Session() - - from openviking.storage.queuefs.session_commit_msg import SessionCommitMsg - from openviking_cli.session.user_id import UserIdentifier - - msg = SessionCommitMsg( - task_id="task-1", - session_id="session-1", - session_uri="viking://session/session-1", - archive_uri="viking://session/session-1/history/archive-1", - user=UserIdentifier("account", "user").to_dict(), - ) - with bind_task_context("task-1", "account", "user"): - await queue.enqueue(msg.to_dict()) - message = {"id": "m", "data": transport.write.await_args.args[1].decode()} - transport.write.reset_mock() - queue.set_dequeue_handler(SessionCommitProcessor(Service())) - - def consume(): - return asyncio.run(queue.process_dequeued(message)) + await queue.process_dequeued(message) + duplicate = await queue.process_dequeued(message) - worker = asyncio.create_task(asyncio.to_thread(consume)) - try: - assert await asyncio.to_thread(started.wait, 3) - cancelled.add("task-1") - index.cancel_active("task-1") - assert await asyncio.to_thread(cleanup_started.wait, 3) - assert not worker.done() - transport.write.assert_not_awaited() - finally: - cleanup_release.set() - if not cleanup_started.is_set(): - cancelled.add("task-1") - index.cancel_active("task-1") - result = await asyncio.wait_for(worker, timeout=3) - assert result.outcome is ProcessOutcome.CANCELLED - transport.write.assert_not_awaited() - await queue.ack("m", message) - assert not index.has_work("task-1") - transport.write.assert_awaited_once_with("/queue/Test/ack", b"m") + assert duplicate.outcome is ProcessOutcome.SUCCESS + status = await tracker.build_queue_status("task-1", **OWNER) + assert status["Embedding"]["processed"] == 1 diff --git a/tests/service/test_task_work_contract.py b/tests/service/test_task_work_contract.py index 5bfc8dd9f6..712a759829 100644 --- a/tests/service/test_task_work_contract.py +++ b/tests/service/test_task_work_contract.py @@ -1,153 +1,301 @@ # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. # SPDX-License-Identifier: AGPL-3.0 +"""Observable contract tests for Task/Work stores.""" import asyncio +import time +from copy import deepcopy import pytest -from openviking.service.task_store import PersistentTaskStore -from openviking.service.task_tracker import TaskStatus, TaskTracker -from openviking.service.task_work_index import QueueTaskMetadata, TaskWorkIndex +from openviking.service.task_domain import ( + TaskRecord, + TaskRetentionPolicy, + TaskStatus, + WorkRecord, + WorkRecoveryMode, + WorkState, +) +from openviking.service.task_store import ( + PersistentTaskStore, + TaskWorkRejected, + WorkRegistration, +) +from openviking.service.task_tracker import TaskTracker from tests.test_task_tracker import _FakeAgfs +OWNER = {"account_id": "account", "user_id": "user"} + + +@pytest.fixture +def store() -> PersistentTaskStore: + return PersistentTaskStore(_FakeAgfs()) + + +async def create_running_task( + store: PersistentTaskStore, + task_id: str = "task", + task_type: str = "content_write", +) -> TaskRecord: + task = TaskRecord(task_id=task_id, task_type=task_type, **OWNER) + await store.create(task) + started = await store.start_task(task_id, stage=None, **OWNER) + assert started is not None + return started + + +def test_file_store_declares_queue_snapshot_recovery(store): + assert store.work_recovery_mode is WorkRecoveryMode.QUEUE_SNAPSHOT + + +@pytest.mark.asyncio +async def test_failed_task_write_does_not_mutate_live_cache(): + class FailingAgfs(_FakeAgfs): + def write(self, path, data, *args, **kwargs): + if b'"status": "running"' in data: + raise OSError("write failed") + return super().write(path, data, *args, **kwargs) + + store = PersistentTaskStore(FailingAgfs()) + task = TaskRecord(task_id="task", task_type="content_write", **OWNER) + await store.create(task) + + with pytest.raises(OSError, match="write failed"): + await store.start_task("task", stage="started", **OWNER) + + unchanged = await store.get_task("task", **OWNER) + assert unchanged is not None + assert unchanged.status is TaskStatus.PENDING + assert unchanged.stage is None + + +@pytest.mark.asyncio +async def test_failed_finalization_write_does_not_publish_terminal_work(): + class FailingAgfs(_FakeAgfs): + def write(self, path, data, *args, **kwargs): + if b'"status": "completed"' in data: + raise OSError("write failed") + return super().write(path, data, *args, **kwargs) + + store = PersistentTaskStore(FailingAgfs()) + tracker = TaskTracker(store) + await create_running_task(store) + await tracker.register_work("task", "work", "Semantic", **OWNER) + await tracker.complete("task", {"done": True}, **OWNER) + + with pytest.raises(OSError, match="write failed"): + await tracker.mark_work_done("task", "work", **OWNER) + + task = await tracker.get("task", **OWNER) + assert task is not None and task.status is TaskStatus.RUNNING + assert await tracker.get_work_state("task", "work", **OWNER) is WorkState.PENDING + assert await tracker.has_work("task", **OWNER) + @pytest.mark.asyncio -async def test_wait_for_work_observes_only_explicit_ids(): - index = TaskWorkIndex() - first = QueueTaskMetadata("task", "first") - unrelated = QueueTaskMetadata("task", "unrelated") - assert index.register("Semantic", first) - assert index.register("Embedding", unrelated) +async def test_work_registration_is_idempotent_and_rejects_terminal_task(store): + await create_running_task(store) + work = WorkRecord(work_id="work", task_id="task", queue_name="Semantic") - index.record_result(first, "completed") - await index.wait_for_work("task", ["first"], timeout=0) + assert await store.add_work(work, **OWNER) is WorkRegistration.CREATED + assert await store.add_work(work, **OWNER) is WorkRegistration.EXISTING - assert index.has_work("task") + await store.finish_work( + "task", + "work", + WorkState.DONE, + error=None, + errors=None, + context_count_delta=0, + **OWNER, + ) + task = await store.get_task("task", **OWNER) + assert task is not None and task.status is TaskStatus.FAILED + with pytest.raises(TaskWorkRejected): + await store.add_work( + WorkRecord(work_id="later", task_id="task", queue_name="Semantic"), + **OWNER, + ) @pytest.mark.asyncio -@pytest.mark.parametrize("recovered", [False, True]) -async def test_task_work_wait_includes_recovered_and_later_enqueued_work(recovered): - index = TaskWorkIndex() - root = QueueTaskMetadata("task", "root") - semantic = QueueTaskMetadata("task", "semantic") - embedding = QueueTaskMetadata("task", "embedding") - if recovered: - index.rebuild( - { - "Root": [{"task_id": "task", "_task_work_id": "root"}], - "Semantic": [{"task_id": "task", "_task_work_id": "semantic"}], - } +async def test_wait_for_work_observes_only_explicit_ids(store): + tracker = TaskTracker(store) + await create_running_task(store) + for work_id in ("first", "unrelated"): + assert ( + await tracker.register_work("task", work_id, "Semantic", **OWNER) + is WorkRegistration.CREATED ) - else: - assert index.register("Root", root) - assert index.register("Semantic", semantic) + + await tracker.mark_work_done("task", "first", **OWNER) + await tracker.wait_for_work("task", ["first"], timeout=0, **OWNER) + + assert await tracker.has_work("task", **OWNER) + + +@pytest.mark.asyncio +async def test_wait_for_task_work_includes_restored_and_later_work(store): + tracker = TaskTracker(store) + await create_running_task(store) + await tracker.reconcile_delivery("task", "semantic", "Semantic", **OWNER) + waiting = asyncio.create_task( - index.wait_for_task_work("task", exclude_work_id="root", poll_interval=0.001) + tracker.wait_for_task_work( + "task", + exclude_work_id="root", + poll_interval=0.001, + **OWNER, + ) ) try: await asyncio.sleep(0) assert not waiting.done() - assert index.register("Embedding", embedding) - index.record_result(semantic, "completed") - await index.prepare_ack("Semantic", semantic) + await tracker.register_work("task", "embedding", "Embedding", **OWNER) + await tracker.mark_work_done("task", "semantic", **OWNER) with pytest.raises(TimeoutError): - await index.wait_for_task_work("task", exclude_work_id="root", timeout=0) - index.record_result(embedding, "completed") - await index.prepare_ack("Embedding", embedding) + await tracker.wait_for_task_work("task", exclude_work_id="root", timeout=0, **OWNER) + await tracker.mark_work_failed("task", "embedding", errors=["embedding failed"], **OWNER) await asyncio.wait_for(waiting, timeout=1) - assert index.has_work("task") finally: waiting.cancel() await asyncio.gather(waiting, return_exceptions=True) + status = await tracker.build_queue_status("task", **OWNER) + assert status["Embedding"]["error_count"] == 1 + @pytest.mark.asyncio -async def test_requeue_keeps_logical_work_until_replacement_ack(): - index = TaskWorkIndex() - work = QueueTaskMetadata("task", "work") - assert index.register("Semantic", work) - assert index.register("Semantic", work) - - index.record_result(work, "requeued") - await index.prepare_ack("Semantic", work) - assert index.has_work("task") - - index.record_result(work, "completed") - await index.wait_for_work("task", ["work"], timeout=0.1) - assert index.queue_status("task")["Semantic"] == { +async def test_requeue_is_non_terminal_and_preserves_logical_work(store): + tracker = TaskTracker(store) + await create_running_task(store) + await tracker.register_work("task", "work", "Semantic", **OWNER) + await tracker.start_work("task", "work", **OWNER) + + await tracker.mark_work_requeued("task", "work", **OWNER) + assert await tracker.has_work("task", **OWNER) + assert await tracker.get_work_state("task", "work", **OWNER) is WorkState.PENDING + + await tracker.start_work("task", "work", **OWNER) + await tracker.mark_work_done("task", "work", **OWNER) + status = await tracker.build_queue_status("task", **OWNER) + assert status["Semantic"] == { "processed": 1, "requeue_count": 1, "error_count": 0, "errors": [], } - await index.prepare_ack("Semantic", work) - assert not index.has_work("task") - -@pytest.mark.parametrize("errors", [None, ["first", "second"]]) -def test_terminal_work_result_is_monotonic_and_idempotent(errors): - index = TaskWorkIndex() - work = QueueTaskMetadata("task", "work") - assert index.register("Embedding", work) - messages = errors or ["first"] - error = "\n".join(messages) +@pytest.mark.asyncio +async def test_terminal_work_result_is_monotonic_and_counts_context(store): + tracker = TaskTracker(store) + await create_running_task(store) + await tracker.register_work("task", "work", "Embedding", **OWNER) - index.record_result(work, "failed", error=error, errors=errors) - index.record_result(work, "completed") - index.record_result(work, "failed", error="later failure") + await tracker.mark_work_failed("task", "work", errors=["first", "second"], **OWNER) + await tracker.mark_work_done("task", "work", context_count_delta=3, **OWNER) + await tracker.mark_work_failed("task", "work", "later failure", **OWNER) - assert index.failure("task") == error - assert index.queue_status("task")["Embedding"] == { + status = await tracker.build_queue_status("task", **OWNER) + assert status["Embedding"] == { "processed": 0, "requeue_count": 0, - "error_count": len(messages), - "errors": [{"message": message} for message in messages], + "error_count": 2, + "errors": [{"message": "first"}, {"message": "second"}], } + assert await tracker.get_context_count("task", **OWNER) == 0 @pytest.mark.asyncio -async def test_enqueue_failure_finalizes_task_with_same_error(): - tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) - index = TaskWorkIndex() - tracker.attach_work_index(index) - task = await tracker.create("content_write", account_id="account", user_id="user") - await tracker.start(task.task_id, account_id="account", user_id="user") - work = QueueTaskMetadata(task.task_id, "embedding", "account", "user") - assert index.register("Embedding", work) +async def test_task_outcome_finalizes_after_all_work_settles(store): + tracker = TaskTracker(store) + await create_running_task(store) + await tracker.register_work("task", "semantic", "Semantic", **OWNER) + await tracker.register_work("task", "embedding", "Embedding", **OWNER) - await index.fail_enqueue("Embedding", work, "embedding enqueue returned false") + await tracker.complete("task", {"done": True}, **OWNER) + active = await tracker.get("task", **OWNER) + assert active is not None and active.status is TaskStatus.RUNNING - failed = await tracker.get(task.task_id, account_id="account", user_id="user") - assert failed.status is TaskStatus.FAILED - assert failed.error == "embedding enqueue returned false" + await tracker.mark_work_done("task", "semantic", **OWNER) + await tracker.mark_work_done("task", "embedding", vector_written=True, **OWNER) + completed = await tracker.get("task", **OWNER) + assert completed is not None and completed.status is TaskStatus.COMPLETED + assert completed.result is not None + assert completed.result["queue_status"]["Semantic"]["processed"] == 1 + assert await tracker.get_context_count("task", **OWNER) == 1 @pytest.mark.asyncio -async def test_deleting_task_releases_retained_work_status(): - tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) - index = TaskWorkIndex() - tracker.attach_work_index(index) - owner = {"account_id": "account", "user_id": "user"} - task = await tracker.create("content_write", **owner) - await tracker.start(task.task_id, **owner) - work = QueueTaskMetadata(task.task_id, "semantic", "account", "user") - assert index.register("Semantic", work) - - await tracker.complete( - task.task_id, - {"done": True, "queue_status": tracker.build_queue_status(task.task_id)}, - **owner, +async def test_persisted_queue_status_caps_error_details_without_losing_count(store): + tracker = TaskTracker(store) + await create_running_task(store) + await tracker.register_work("task", "embedding", "Embedding", **OWNER) + await tracker.complete("task", {"done": True}, **OWNER) + + await tracker.mark_work_failed( + "task", + "embedding", + errors=[f"failure {index}" for index in range(150)], + **OWNER, ) - index.record_result(work, "completed") - await index.prepare_ack("Semantic", work) - assert tracker.build_queue_status(task.task_id)["Semantic"]["processed"] == 1 - completed = await tracker.get(task.task_id, **owner) - assert completed.result["queue_status"]["Semantic"]["processed"] == 1 - assert await tracker.delete_user_tasks(**owner) == 1 - assert tracker.build_queue_status(task.task_id) == { + completed = await tracker.get("task", **OWNER) + assert completed is not None and completed.result is not None + embedding = completed.result["queue_status"]["Embedding"] + assert embedding["error_count"] == 150 + assert len(embedding["errors"]) == 100 + + +@pytest.mark.asyncio +async def test_delete_expired_limits_reads_and_skips_cached_records(): + class CountingAgfs(_FakeAgfs): + def __init__(self): + super().__init__() + self.read_count = 0 + + def read(self, *args, **kwargs): + value = super().read(*args, **kwargs) + self.read_count += 1 + return value + + agfs = CountingAgfs() + writer = PersistentTaskStore(agfs) + for index in range(4): + await writer.create( + TaskRecord( + task_id=f"task-{index}", + task_type="content_write", + **OWNER, + ) + ) + + store = PersistentTaskStore(agfs) + for task_id in ("task-0", "task-1"): + assert await store.get_task(task_id, **OWNER) is not None + agfs.read_count = 0 + policy = TaskRetentionPolicy(completed_seconds=3600, failed_seconds=3600) + + first = await store.delete_expired(policy, now=time.time(), limit=2) + assert first.scanned == 2 + assert agfs.read_count == 0 + + second = await store.delete_expired(policy, now=time.time(), limit=2) + assert second.scanned == 2 + assert agfs.read_count == 2 + + +@pytest.mark.asyncio +async def test_deleting_task_releases_work_projection(store): + tracker = TaskTracker(store) + await create_running_task(store) + await tracker.register_work("task", "work", "Semantic", **OWNER) + await tracker.mark_work_done("task", "work", **OWNER) + + assert await tracker.delete_tasks(**OWNER) == 1 + assert not await tracker.has_work("task", **OWNER) + assert await tracker.build_queue_status("task", **OWNER) == { "Semantic": { "processed": 0, "requeue_count": 0, @@ -161,3 +309,77 @@ async def test_deleting_task_releases_retained_work_status(): "errors": [], }, } + + +@pytest.mark.asyncio +async def test_cancelled_persisted_mutation_updates_read_view_after_write_finishes(): + class DelayedWriteStore(PersistentTaskStore): + def __init__(self) -> None: + super().__init__(_FakeAgfs()) + self.delay_writes = False + self.entered_write = asyncio.Event() + self.release_write = asyncio.Event() + + async def _write_task(self, task: TaskRecord) -> None: + if self.delay_writes: + self.entered_write.set() + await self.release_write.wait() + await super()._write_task(task) + + store = DelayedWriteStore() + tracker = TaskTracker(store) + await create_running_task(store) + store.delay_writes = True + + completing = asyncio.create_task(tracker.complete("task", {"done": True}, **OWNER)) + await asyncio.wait_for(store.entered_write.wait(), timeout=1) + completing.cancel() + store.release_write.set() + + with pytest.raises(asyncio.CancelledError): + await completing + + task = await tracker.get("task", **OWNER) + assert task is not None + assert task.status is TaskStatus.COMPLETED + assert task.result == {"done": True} + + +@pytest.mark.asyncio +async def test_delete_tasks_rechecks_liveness_under_record_lock(monkeypatch, store): + tracker = TaskTracker(store) + running = await create_running_task(store) + stale_terminal = deepcopy(running) + stale_terminal.status = TaskStatus.COMPLETED + stale_terminal.stage = TaskStatus.COMPLETED.value + stale_terminal.result = {"done": True} + + async def stale_list(*args, **kwargs): + return [deepcopy(stale_terminal)] + + monkeypatch.setattr(store, "list", stale_list) + + with pytest.raises(RuntimeError, match="Cannot delete active task record"): + await tracker.delete_tasks(**OWNER) + + task = await tracker.get("task", **OWNER) + assert task is not None + assert task.status is TaskStatus.RUNNING + + +@pytest.mark.asyncio +async def test_delete_tasks_removes_all_users_in_account_only(store): + tracker = TaskTracker(store) + owners = [ + {"account_id": "account", "user_id": "alice"}, + {"account_id": "account", "user_id": "bob"}, + {"account_id": "other", "user_id": "alice"}, + ] + for index, owner in enumerate(owners): + task = await tracker.create("content_write", task_id=f"task-{index}", **owner) + await tracker.complete(task.task_id, {"done": True}, **owner) + + assert await tracker.delete_tasks("account") == 2 + assert await tracker.get("task-0", **owners[0]) is None + assert await tracker.get("task-1", **owners[1]) is None + assert await tracker.get("task-2", **owners[2]) is not None diff --git a/tests/session/conftest.py b/tests/session/conftest.py index 7e8eef5c73..88ae682a58 100644 --- a/tests/session/conftest.py +++ b/tests/session/conftest.py @@ -32,10 +32,13 @@ async def client( tracker = get_task_tracker() original_has_work = tracker.has_work + async def has_work(task_id, **kwargs): + return task_id in queued_commit_tasks or await original_has_work(task_id, **kwargs) + monkeypatch.setattr( tracker, "has_work", - lambda task_id: task_id in queued_commit_tasks or original_has_work(task_id), + has_work, ) async def enqueue_with_session_commit_fallback(queue_name, data): diff --git a/tests/session/memory/test_consolidation_compile.py b/tests/session/memory/test_consolidation_compile.py index 085733b1cb..13365cf1ad 100644 --- a/tests/session/memory/test_consolidation_compile.py +++ b/tests/session/memory/test_consolidation_compile.py @@ -16,6 +16,7 @@ _memory_type_from_target, _peer_id_from_memory_uri, ) +from openviking.service.task_store import PersistentTaskStore from openviking.service.task_tracker import TaskStatus, TaskTracker from openviking.service.task_work_index import get_task_context from openviking.session.memory.consolidation_context_provider import ( @@ -39,6 +40,7 @@ from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils from openviking_cli.exceptions import InvalidArgumentError, NotFoundError from openviking_cli.session.user_id import UserIdentifier +from tests.test_task_tracker import _FakeAgfs def _ctx(user_id: str = "u1") -> RequestContext: @@ -168,14 +170,7 @@ async def test_root_compile_creates_one_task_and_one_loop(monkeypatch, language_ language_config.output_language_override = "en" vlm = SimpleNamespace(model="test", get_completion_async=AsyncMock(return_value="sdk.commit()")) language_config.vlm = SimpleNamespace(get_vlm_instance=lambda: vlm) - store = SimpleNamespace( - create=AsyncMock(), - update=AsyncMock(), - get=AsyncMock(return_value=None), - list=AsyncMock(return_value=[]), - delete=AsyncMock(), - ) - tracker = TaskTracker(store) + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) monkeypatch.setattr("openviking.service.memory_compile.get_task_tracker", lambda: tracker) uris = [ f"{target}/entities/person/alice.md", @@ -216,7 +211,6 @@ async def stat(uri, ctx, skip_count=False): await asyncio.gather(*list(runner._running)) final = await tracker.get(accepted.task_id, account_id="acc", user_id="u1") - store.create.assert_awaited_once() vlm.get_completion_async.assert_awaited_once() fs.glob.assert_awaited_once() assert fs.glob.await_args.kwargs["uri"] == target @@ -235,7 +229,11 @@ async def stat(uri, ctx, skip_count=False): for name in ("entities", "preferences", "profile"): assert name in prompt assert "create_events" not in prompt - assert not tracker.has_work(accepted.task_id) + assert not await tracker.has_work( + accepted.task_id, + account_id="acc", + user_id="u1", + ) @pytest.mark.asyncio @@ -736,14 +734,7 @@ async def test_compile_task_apply_outcomes( language_config.output_language_override = "en" language_config.vlm = SimpleNamespace(get_vlm_instance=lambda: SimpleNamespace(model="test")) ctx = _ctx() - store = SimpleNamespace( - create=AsyncMock(), - update=AsyncMock(), - get=AsyncMock(return_value=None), - list=AsyncMock(return_value=[]), - delete=AsyncMock(), - ) - tracker = TaskTracker(store) + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) task = await tracker.create("compile", account_id=ctx.account_id, user_id=ctx.user.user_id) monkeypatch.setattr("openviking.service.memory_compile.get_task_tracker", lambda: tracker) monkeypatch.setattr( @@ -803,7 +794,11 @@ async def test_compile_task_apply_outcomes( assert final.error.startswith("MEMORY_CONSOLIDATION_FAILED:") else: assert final.error is None - assert not tracker.has_work(task.task_id) + assert not await tracker.has_work( + task.task_id, + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) if has_operations: apply.assert_awaited_once() viking_fs._async_agfs.pathlock_release.assert_awaited_once_with(lease) @@ -821,14 +816,7 @@ async def test_compile_task_apply_outcomes( @pytest.mark.asyncio @pytest.mark.parametrize("memory_type", [None, "entities"]) async def test_cancel_running_compile_interrupts_before_memory_write(monkeypatch, memory_type): - store = SimpleNamespace( - create=AsyncMock(), - update=AsyncMock(), - get=AsyncMock(return_value=None), - list=AsyncMock(return_value=[]), - delete=AsyncMock(), - ) - tracker = TaskTracker(store) + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) ctx = _ctx() task = await tracker.create( "compile", @@ -869,7 +857,11 @@ async def consolidate(**kwargs): final = await tracker.get(task.task_id, account_id=ctx.account_id, user_id=ctx.user.user_id) assert final.status == TaskStatus.CANCELLED - assert not tracker.has_work(task.task_id) + assert not await tracker.has_work( + task.task_id, + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) memory_write.assert_not_awaited() diff --git a/tests/session/test_session_context.py b/tests/session/test_session_context.py index 866c57bc73..efc4725434 100644 --- a/tests/session/test_session_context.py +++ b/tests/session/test_session_context.py @@ -51,10 +51,14 @@ async def client( queued_commit_tasks: set[str] = set() tracker = get_task_tracker() original_has_work = tracker.has_work + + async def has_work(task_id, **kwargs): + return task_id in queued_commit_tasks or await original_has_work(task_id, **kwargs) + monkeypatch.setattr( tracker, "has_work", - lambda task_id: task_id in queued_commit_tasks or original_has_work(task_id), + has_work, ) async def enqueue_with_session_commit_fallback(queue_name, data): diff --git a/tests/storage/test_collection_schemas.py b/tests/storage/test_collection_schemas.py index da75841ddc..5cc66f6536 100644 --- a/tests/storage/test_collection_schemas.py +++ b/tests/storage/test_collection_schemas.py @@ -1259,11 +1259,7 @@ async def upsert(self, data, *, ctx, options=UpsertOptions()): assert result.outcome is ProcessOutcome.SUCCESS assert result.error is None - assert result.value == { - **queue_data["payload"]["context_data"], - "id": vector_record_id("default", "viking://resources/sample", 2), - "vector": [0.1, 0.2], - } + assert result.value == {"vector_written": True} assert "data" in captured assert captured["data"]["parent_uri"] == "viking://resources" @@ -1324,17 +1320,8 @@ async def upsert(self, data, *, ctx, options=UpsertOptions()): @pytest.mark.asyncio -async def test_embedding_handler_generates_id_and_counts_written_context(monkeypatch): - from openviking.service.task_tracker import TaskTracker - from openviking.service.task_work_index import QueueTaskMetadata, TaskWorkIndex - from openviking.storage.vector_ids import vector_record_id - +async def test_embedding_handler_generates_id_and_reports_written_context(monkeypatch): captured = {} - tracker = TaskTracker(AsyncMock()) - index = TaskWorkIndex() - tracker.attach_work_index(index) - index.register("Embedding", QueueTaskMetadata("task-1", "work-1")) - monkeypatch.setattr("openviking.storage.collection_schemas.get_task_tracker", lambda: tracker) class _CapturingVikingDB: is_closing = False @@ -1355,11 +1342,10 @@ async def upsert(self, data, *, ctx, options=UpsertOptions()): queue_data["_task_work_id"] = "work-1" payload["data"] = json.dumps(queue_data) - await TextEmbeddingHandler(_CapturingVikingDB()).on_dequeue(payload) + result = await TextEmbeddingHandler(_CapturingVikingDB()).on_dequeue(payload) assert captured["data"]["id"] == vector_record_id("acct-new-id", "viking://resources/sample", 2) - assert tracker.get_context_count("task-1") == 1 - assert tracker.get_context_count("unrelated-task") == 0 + assert result.value == {"vector_written": True} def test_context_collection_uses_acl_mode_and_excludes_parent_uri(): diff --git a/tests/storage/test_content_write_processing_mode.py b/tests/storage/test_content_write_processing_mode.py index ef8e9a5892..fd21409688 100644 --- a/tests/storage/test_content_write_processing_mode.py +++ b/tests/storage/test_content_write_processing_mode.py @@ -10,7 +10,7 @@ from openviking.server.identity import RequestContext, Role from openviking.service.task_store import PersistentTaskStore from openviking.service.task_tracker import TaskTracker -from openviking.service.task_work_index import TaskWorkIndex, bind_task_context +from openviking.service.task_work_index import bind_task_context from openviking.session.memory.utils.memory_file_utils import MemoryFileUtils from openviking.storage import content_write as content_write_module from openviking.storage.abstract_overview import ( @@ -31,7 +31,6 @@ @pytest.fixture(autouse=True) def run_content_write_operations_inline(monkeypatch): tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) - tracker.attach_work_index(TaskWorkIndex()) monkeypatch.setattr("openviking.service.task_tracker._instance", tracker) async def run_tracked_operation(_task_type, *, operation, **_kwargs): diff --git a/tests/storage/test_embedding_msg.py b/tests/storage/test_embedding_msg.py index 0880409aac..b3eaf7d64e 100644 --- a/tests/storage/test_embedding_msg.py +++ b/tests/storage/test_embedding_msg.py @@ -303,6 +303,16 @@ def test_legacy_upsert_discards_historically_ignored_fields(): assert "record_ids" not in serialized assert "update_fields" not in serialized assert "initial_fields" not in serialized + assert set(restored.to_dict()) == { + "id", + "telemetry_id", + "task_id", + "_task_work_id", + "queue_enqueued_at", + "failure_requeue_count", + "action", + "payload", + } def test_legacy_update_fields_message_normalizes_to_field_patch_payload(): diff --git a/tests/storage/test_queue_manager.py b/tests/storage/test_queue_manager.py index 10f70eed98..48547f3f96 100644 --- a/tests/storage/test_queue_manager.py +++ b/tests/storage/test_queue_manager.py @@ -14,11 +14,8 @@ import pytest from openviking.pyagfs import AsyncAGFSClient -from openviking.service.task_work_index import ( - QueueTaskMetadata, - bind_task_context, - record_enqueued_work, -) +from openviking.service.task_domain import WorkRecoveryMode +from openviking.service.task_work_index import bind_task_context, record_enqueued_work from openviking.storage.queuefs.named_queue import DequeueHandlerBase, NamedQueue from openviking.storage.queuefs.process_result import ProcessResult from openviking.storage.queuefs.queue_manager import QueueManager @@ -86,10 +83,13 @@ def has_dequeue_handler(self): async def test_skill_shutdown_releases_lock_after_embedding_worker_exits( transport, monkeypatch, concurrency ): + from openviking.service.task_store import PersistentTaskStore + from openviking.service.task_tracker import TaskTracker from openviking.service.task_tracker_concurrency import run_to_completion from openviking.storage.queuefs.semantic_executor import SemanticTreeStats from openviking.storage.queuefs.semantic_msg import SemanticMsg from openviking.storage.queuefs.semantic_processor import SemanticProcessor + from tests.test_task_tracker import _FakeAgfs msg = SemanticMsg( uri="viking://agent/skills/demo", @@ -99,9 +99,14 @@ async def test_skill_shutdown_releases_lock_after_embedding_worker_exits( ) started, queued, release = threading.Event(), threading.Event(), threading.Event() events = [] - work_index = None - active_work = QueueTaskMetadata("skill-task", "active") - queued_work = QueueTaskMetadata("skill-task", "queued") + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) + monkeypatch.setattr("openviking.service.task_tracker._instance", tracker) + task = await tracker.create( + "add_skill", task_id="skill-task", account_id="acc", user_id="alice" + ) + await tracker.start(task.task_id, account_id="acc", user_id="alice") + active_work = "active" + queued_work = "queued" class Lease: lock = {"lease_ref": "demo"} @@ -116,11 +121,22 @@ def __init__(self, **kwargs): pass async def run(self, uri): - assert work_index is not None - assert work_index.register("Embedding", active_work) - assert work_index.register("Embedding", queued_work) - record_enqueued_work("Embedding", active_work.work_id) - record_enqueued_work("Embedding", queued_work.work_id) + await tracker.register_work( + task.task_id, + active_work, + "Embedding", + account_id="acc", + user_id="alice", + ) + await tracker.register_work( + task.task_id, + queued_work, + "Embedding", + account_id="acc", + user_id="alice", + ) + record_enqueued_work("Embedding", active_work) + record_enqueued_work("Embedding", queued_work) queued.set() def get_stats(self): @@ -132,8 +148,12 @@ async def finish(): while not release.is_set(): await asyncio.sleep(0.01) events.append("write-finished") - assert work_index is not None - work_index.record_result(active_work, "completed") + await tracker.mark_work_done( + task.task_id, + active_work, + account_id="acc", + user_id="alice", + ) await run_to_completion(finish) return ProcessResult.success() @@ -156,7 +176,9 @@ async def dequeue_raw(self): async def process_dequeued(self, data): if self.name == QueueManager.SEMANTIC: - with bind_task_context("skill-task", "", "", enqueued_work=[]): + with bind_task_context( + "skill-task", "acc", "alice", enqueued_work=[] + ): return await self.handler(data) return await self.handler(data) @@ -183,11 +205,6 @@ async def dequeue(self): manager = QueueManager( object(), max_concurrent_semantic=concurrency, max_concurrent_embedding=concurrency ) - work_index = manager._task_work_index - monkeypatch.setattr( - "openviking.storage.queuefs.semantic_work.get_task_tracker", - lambda: SimpleNamespace(wait_for_work=work_index.wait_for_work), - ) manager._poll_interval = 0.001 manager.set_vlm_resolver(SimpleNamespace(get_vlm=AsyncMock(return_value=SimpleNamespace()))) manager.setup_standard_queues(object(), start=False) @@ -214,7 +231,9 @@ async def dequeue(self): assert all(not thread.is_alive() for thread in threads) assert events.index("write-finished") < events.index("lock-released") assert f"ack-{manager.SEMANTIC}" not in events - assert work_index.has_work("skill-task") # Queued vectors await restart. + assert await tracker.has_work( + task.task_id, account_id="acc", user_id="alice" + ) finally: release.set() if stopping is not None: @@ -267,13 +286,11 @@ async def enqueue(self, ctx, call_next): manager = QueueManager(object(), middlewares=middlewares) middlewares.clear() existing = manager.get_queue("existing", allow_create=True) - with bind_task_context("task", "account", "user"): - await existing.enqueue({}) - middlewares.append(Recorder()) - future = manager.get_queue("future", allow_create=True) - await future.enqueue({}) + await existing.enqueue({}) + middlewares.append(Recorder()) + future = manager.get_queue("future", allow_create=True) + await future.enqueue({}) assert calls == ["existing", "future"] - assert manager._task_work_index.has_work("task") assert transport.write.await_count == 2 @@ -334,30 +351,34 @@ async def read(path): transport.write.assert_awaited_once_with("/queue/Test/ack", b"message-1") -async def test_bootstrap_restores_legacy_messages_into_middleware_index(transport): +async def test_bootstrap_restores_legacy_messages_into_store(transport): manager = QueueManager(object()) - queue = manager.get_queue("Test", allow_create=True) + manager.get_queue("Test", allow_create=True) message = { "id": "legacy-message", "data": json.dumps({"task_id": "task", "account_id": "account", "user_id": "user"}), } transport.read.return_value = json.dumps([message]).encode() - attached = [] + restored_work = [] class Tracker: - def attach_work_index(self, index): - attached.append(index) - assert index.has_work("task") + work_recovery_mode = WorkRecoveryMode.QUEUE_SNAPSHOT + + async def reconcile_delivery( + self, task_id, work_id, queue_name, *, account_id, user_id + ): + restored_work.append( + (task_id, work_id, queue_name, account_id, user_id) + ) async def restore_work_tasks(self, owners): assert owners == {"task": ("account", "user")} return ["restored"] assert await manager.prepare_task_tracking(Tracker()) == ["restored"] - assert attached == [manager._task_work_index] - await queue.ack("legacy-message", message) - assert not attached[0].has_work("task") - transport.write.assert_awaited_once_with("/queue/Test/ack", b"legacy-message") + assert restored_work == [ + ("task", "queuefs:legacy-message", "Test", "account", "user") + ] async def test_bootstrap_restores_task_owner_from_cross_user_child_work(transport): @@ -377,13 +398,34 @@ async def test_bootstrap_restores_task_owner_from_cross_user_child_work(transpor ), } transport.read.return_value = json.dumps([message]).encode() + restored_work = [] class Tracker: - def attach_work_index(self, index): - assert index.has_work("task") + work_recovery_mode = WorkRecoveryMode.QUEUE_SNAPSHOT + + async def reconcile_delivery( + self, task_id, work_id, queue_name, *, account_id, user_id + ): + restored_work.append( + (task_id, work_id, queue_name, account_id, user_id) + ) async def restore_work_tasks(self, owners): assert owners == {"task": ("account", "task-owner")} return ["restored"] assert await manager.prepare_task_tracking(Tracker()) == ["restored"] + assert restored_work == [ + ("task", "work", "Test", "account", "task-owner") + ] + + +async def test_bootstrap_skips_queue_snapshot_for_durable_work_store(transport): + manager = QueueManager(object()) + manager.get_queue("Test", allow_create=True) + + class Tracker: + work_recovery_mode = WorkRecoveryMode.DURABLE_STORE + + assert await manager.prepare_task_tracking(Tracker()) == [] + transport.read.assert_not_awaited() diff --git a/tests/storage/test_reindex_msg.py b/tests/storage/test_reindex_msg.py index 5bb6883d39..b529efae29 100644 --- a/tests/storage/test_reindex_msg.py +++ b/tests/storage/test_reindex_msg.py @@ -63,16 +63,13 @@ def test_reindex_msg_rejects_missing_descriptor_fields(): @pytest.mark.asyncio async def test_reindex_processor_terminalizes_root_after_descendants(monkeypatch): - from openviking.service.task_store import PersistentTaskStore + from openviking.service.task_store import PersistentTaskStore, WorkRegistration from openviking.service.task_tracker import TaskTracker - from openviking.service.task_work_index import QueueTaskMetadata, TaskWorkIndex from openviking.storage.queuefs.reindex_processor import ReindexProcessor from tests.test_task_tracker import _FakeAgfs owner = {"account_id": "account-1", "user_id": "user-1"} tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) - index = TaskWorkIndex() - tracker.attach_work_index(index) monkeypatch.setattr( "openviking.storage.queuefs.reindex_processor.get_task_tracker", Mock(return_value=tracker) ) @@ -89,10 +86,14 @@ async def test_reindex_processor_terminalizes_root_after_descendants(monkeypatch processor = ReindexProcessor(viking_fs) processor._adopt_or_reacquire = AsyncMock(return_value={"lease": "root"}) task = await tracker.create("admin_reindex", resource_id="viking://resources/demo", **owner) - root_work = QueueTaskMetadata(task.task_id, "root-work", **owner) - child_work = QueueTaskMetadata(task.task_id, "child-work", **owner) - assert index.register("Reindex", root_work) - assert index.register("Embedding", child_work) + assert ( + await tracker.register_work(task.task_id, "root-work", "Reindex", **owner) + is WorkRegistration.CREATED + ) + assert ( + await tracker.register_work(task.task_id, "child-work", "Embedding", **owner) + is WorkRegistration.CREATED + ) entered_descendant_wait = asyncio.Event() original_wait_for_task_work = tracker.wait_for_task_work @@ -119,14 +120,12 @@ async def wait_for_task_work(*args, **kwargs): assert not processing.done() assert (await tracker.get(task.task_id, **owner)).status is TaskStatus.RUNNING - index.record_result(child_work, "completed") - await index.prepare_ack("Embedding", child_work) + await tracker.mark_work_done(task.task_id, "child-work", **owner) outcome = await asyncio.wait_for(processing, 1) assert outcome.outcome.value == "success" assert (await tracker.get(task.task_id, **owner)).status is TaskStatus.RUNNING - index.record_result(root_work, "completed") - await index.prepare_ack("Reindex", root_work) + await tracker.mark_work_done(task.task_id, "root-work", **owner) completed = await tracker.wait(task.task_id, **owner, timeout=0) assert completed.status is TaskStatus.COMPLETED assert completed.result["rebuilt_records"] == 3 diff --git a/tests/storage/test_semantic_executor_stats.py b/tests/storage/test_semantic_executor_stats.py index ee48eb8d1a..f64edbe4b9 100644 --- a/tests/storage/test_semantic_executor_stats.py +++ b/tests/storage/test_semantic_executor_stats.py @@ -11,7 +11,6 @@ from openviking.server.identity import RequestContext, Role from openviking.service.task_queue_middleware import TaskWorkQueueMiddleware from openviking.service.task_work_index import ( - TaskWorkIndex, TaskWorkRejected, bind_task_context, get_task_context, @@ -468,20 +467,11 @@ async def test_semantic_executor_shares_node_scheduler_across_roots(monkeypatch) @pytest.mark.asyncio async def test_task_work_rejection_does_not_stop_shared_semantic_worker(): - work_index = TaskWorkIndex() - - async def finalize_before_ack(_metadata): - return None - - work_index.set_callbacks( - finalize_before_ack=finalize_before_ack, - is_cancellation_requested=lambda _task_id: True, - ) embedding_queue = NamedQueue( None, "/queue", "Embedding", - middlewares=[TaskWorkQueueMiddleware(work_index)], + middlewares=[TaskWorkQueueMiddleware()], ) embedding_queue._initialized = True unrelated_ran = asyncio.Event() diff --git a/tests/storage/test_semantic_processor_lock_ownership.py b/tests/storage/test_semantic_processor_lock_ownership.py index 534042adf6..a3f5368e9a 100644 --- a/tests/storage/test_semantic_processor_lock_ownership.py +++ b/tests/storage/test_semantic_processor_lock_ownership.py @@ -167,7 +167,6 @@ async def test_shared_scheduler_drains_each_messages_embeddings_before_unlock( from openviking.service.task_store import PersistentTaskStore from openviking.service.task_tracker import TaskStatus, TaskTracker from openviking.service.task_work_index import ( - TaskWorkIndex, bind_task_context, extract_task_metadata, ) @@ -210,8 +209,6 @@ async def pathlock_release(self, lease): monkeypatch.setattr(f"openviking.storage.queuefs.{module}.get_viking_fs", lambda: fs) agfs = _FakeAgfs() tracker = TaskTracker(PersistentTaskStore(agfs)) - index = TaskWorkIndex() - tracker.attach_work_index(index) monkeypatch.setattr("openviking.service.task_tracker._instance", tracker) entered_embedding_drain = asyncio.Event() original_wait_for_work = tracker.wait_for_work @@ -227,7 +224,7 @@ async def wait_for_work(*args, **kwargs): cast(AGFSSyncClientProtocol, agfs), "/queue", "Embedding", - middlewares=[TaskWorkQueueMiddleware(index)], + middlewares=[TaskWorkQueueMiddleware()], ) enqueued = asyncio.Event() @@ -296,8 +293,12 @@ async def vectorize(uri, **kwargs): assert metadata.task_id != source.task_id if task: assert metadata.task_id == task.task_id - index.record_result(metadata, "completed") - await index.prepare_ack("Embedding", metadata) + await tracker.mark_work_done( + metadata.task_id, + metadata.work_id, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) result = await asyncio.wait_for(worker, 1) assert result.outcome is ( ProcessOutcome.FAILED if fail_generation else ProcessOutcome.SUCCESS @@ -319,6 +320,11 @@ async def vectorize(uri, **kwargs): if payload: metadata = extract_task_metadata(json.loads(payload)) if metadata: - index.record_result(metadata, "completed") + await tracker.mark_work_done( + metadata.task_id, + metadata.work_id, + account_id=metadata.account_id, + user_id=metadata.user_id, + ) worker.cancel() await asyncio.gather(worker, return_exceptions=True) diff --git a/tests/storage/test_semantic_processor_permanent_storage_error.py b/tests/storage/test_semantic_processor_permanent_storage_error.py index 859af3855a..eaefba68f2 100644 --- a/tests/storage/test_semantic_processor_permanent_storage_error.py +++ b/tests/storage/test_semantic_processor_permanent_storage_error.py @@ -13,7 +13,9 @@ from openviking.models.model_error import ModelError, ModelErrorKind from openviking.pyagfs.exceptions import AGFSNotADirectoryError from openviking.service.task_queue_middleware import TaskWorkQueueMiddleware -from openviking.service.task_work_index import TaskWorkIndex +from openviking.service.task_store import PersistentTaskStore +from openviking.service.task_tracker import TaskStatus, TaskTracker +from openviking.service.task_work_index import bind_task_context from openviking.storage.errors import LockAcquisitionError from openviking.storage.queuefs.named_queue import NamedQueue from openviking.storage.queuefs.process_result import ProcessOutcome @@ -22,6 +24,7 @@ from openviking.storage.queuefs.semantic_processor import SemanticProcessor from openviking.storage.viking_fs import SyncDiff from openviking.utils.circuit_breaker import CircuitBreaker, CircuitBreakerOpen +from tests.test_task_tracker import _FakeAgfs @pytest.fixture @@ -123,22 +126,33 @@ async def test_queue_acks_failed_semantic_work(semantic_env, monkeypatch, wrappe monkeypatch.setattr( "openviking.storage.queuefs.named_queue.AsyncAGFSClient", lambda agfs: backend ) - index = TaskWorkIndex() - finalize = AsyncMock() - index.set_callbacks( - finalize_before_ack=finalize, is_cancellation_requested=lambda task_id: False - ) + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) + monkeypatch.setattr("openviking.service.task_tracker._instance", tracker) queue = NamedQueue( object(), "/queue", "Semantic", dequeue_handler=env.processor, - middlewares=[TaskWorkQueueMiddleware(index)], + middlewares=[TaskWorkQueueMiddleware()], ) payload = msg.to_dict() if task_owned: payload["task_id"] = "ingestion-task" - await queue.enqueue(payload) + task = await tracker.create( + "add_resource", + task_id="ingestion-task", + account_id=msg.account_id, + user_id=msg.user_id, + ) + await tracker.start( + task.task_id, + account_id=msg.account_id, + user_id=msg.user_id, + ) + with bind_task_context(task.task_id, msg.account_id, msg.user_id): + await queue.enqueue(payload) + else: + await queue.enqueue(payload) queued_payload = backend.write.call_args.args[1].decode() backend.read = AsyncMock( side_effect=[ @@ -146,7 +160,14 @@ async def test_queue_acks_failed_semantic_work(semantic_env, monkeypatch, wrappe json.dumps({"pending": 0, "processing": 0}).encode(), ] ) - assert index.has_work("ingestion-task") is task_owned + assert ( + await tracker.has_work( + "ingestion-task", + account_id=msg.account_id, + user_id=msg.user_id, + ) + is task_owned + ) await asyncio.wait_for(queue.dequeue(), timeout=1) @@ -155,13 +176,20 @@ async def test_queue_acks_failed_semantic_work(semantic_env, monkeypatch, wrappe assert status.in_progress == status.processed == status.requeue_count == 0 assert status.error_count == 1 assert [entry.message for entry in status.errors] == [str(error)] - assert not index.has_work("ingestion-task") # Includes active consumer registration. - assert index.failure("ingestion-task") == (str(error) if task_owned else None) + assert not await tracker.has_work( + "ingestion-task", + account_id=msg.account_id, + user_id=msg.user_id, + ) if task_owned: - finalize.assert_awaited_once() - assert finalize.call_args.args[0].task_id == "ingestion-task" - else: - finalize.assert_not_awaited() + failed = await tracker.get( + "ingestion-task", + account_id=msg.account_id, + user_id=msg.user_id, + ) + assert failed is not None + assert failed.status is TaskStatus.FAILED + assert failed.error == str(error) env.retry_queue.enqueue.assert_not_awaited() env.pathlock.pathlock_release.assert_awaited_once_with(env.lease) diff --git a/tests/storage/test_skill_package_cancellation.py b/tests/storage/test_skill_package_cancellation.py index a23ddcdaf7..a79e63ed33 100644 --- a/tests/storage/test_skill_package_cancellation.py +++ b/tests/storage/test_skill_package_cancellation.py @@ -12,8 +12,9 @@ from openviking.models.embedder.base import EmbedResult from openviking.server.identity import RequestContext, Role from openviking.service.task_queue_middleware import TaskWorkQueueMiddleware +from openviking.service.task_store import PersistentTaskStore +from openviking.service.task_tracker import TaskTracker from openviking.service.task_work_index import ( - TaskWorkIndex, bind_task_context, get_task_context, prepare_task_payload, @@ -29,6 +30,7 @@ from openviking_cli.session.user_id import UserIdentifier from tests.storage.test_collection_schemas import _DummyConfig, _DummyEmbedder from tests.storage.test_semantic_executor_skip_files import _FakeVikingFS +from tests.test_task_tracker import _FakeAgfs def _ctx(): @@ -189,16 +191,12 @@ async def upsert(data, **kwargs): "account_id": "acc", }, ) - index = TaskWorkIndex() - cancelling = set() - index.set_callbacks( - finalize_before_ack=AsyncMock(), - is_cancellation_requested=lambda task_id: task_id in cancelling, - ) - monkeypatch.setattr( - "openviking.storage.collection_schemas.get_task_tracker", - lambda: SimpleNamespace(record_context_count=index.record_context_count), + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) + monkeypatch.setattr("openviking.service.task_tracker._instance", tracker) + task = await tracker.create( + "add_skill", task_id="skill-task", account_id="acc", user_id="alice" ) + await tracker.start(task.task_id, account_id="acc", user_id="alice") with bind_task_context("skill-task", "acc", "alice"): payload, metadata = prepare_task_payload(msg.to_dict()) queue = NamedQueue( @@ -206,19 +204,27 @@ async def upsert(data, **kwargs): "/queue", "Embedding", dequeue_handler=handler, - middlewares=[TaskWorkQueueMiddleware(index)], + middlewares=[TaskWorkQueueMiddleware()], ) envelope = {"data": json.dumps(payload)} - assert index.register(queue.name, metadata) + assert metadata is not None + await tracker.register_work( + task.task_id, + metadata.work_id, + queue.name, + account_id="acc", + user_id="alice", + ) worker = asyncio.create_task(queue.process_dequeued(envelope)) try: await asyncio.wait_for(started.wait(), 1) - cancelling.add("skill-task") - index.cancel_active("skill-task") + await tracker.cancel(task.task_id, account_id="acc", user_id="alice") await asyncio.sleep(0.03) if during_write: assert not worker.done() - assert index.has_work("skill-task") + assert await tracker.has_work( + task.task_id, account_id="acc", user_id="alice" + ) assert not writes else: await asyncio.wait_for(worker, 1) @@ -227,23 +233,21 @@ async def upsert(data, **kwargs): release.set() result = await asyncio.wait_for(worker, 1) assert result.outcome is ProcessOutcome.CANCELLED - await index.prepare_ack(queue.name, envelope) - assert not index.has_work("skill-task") + assert not await tracker.has_work( + task.task_id, account_id="acc", user_id="alice" + ) assert len(writes) == int(during_write) @pytest.mark.asyncio async def test_semantic_cancel_drains_embeddings_before_releasing_package(monkeypatch): - index = TaskWorkIndex() - child = None - index.set_callbacks( - finalize_before_ack=AsyncMock(), - is_cancellation_requested=lambda _task_id: False, - ) - monkeypatch.setattr( - "openviking.storage.queuefs.semantic_work.get_task_tracker", - lambda: SimpleNamespace(wait_for_work=index.wait_for_work), + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) + monkeypatch.setattr("openviking.service.task_tracker._instance", tracker) + task = await tracker.create( + "add_skill", task_id="skill-task", account_id="acc", user_id="alice" ) + await tracker.start(task.task_id, account_id="acc", user_id="alice") + child = None msg = SemanticMsg( uri="viking://agent/skills/demo", context_type="skill", @@ -258,11 +262,15 @@ def __init__(self, **kwargs): async def run(self, uri): nonlocal child - from openviking.service.task_work_index import QueueTaskMetadata - - child = QueueTaskMetadata("skill-task", "embedding") - assert index.register("Embedding", child) - record_enqueued_work("Embedding", child.work_id) + child = "embedding" + await tracker.register_work( + task.task_id, + child, + "Embedding", + account_id="acc", + user_id="alice", + ) + record_enqueued_work("Embedding", child) started.set() def get_stats(self): @@ -298,7 +306,12 @@ async def close(self): assert events == [] events.append("embedding-settled") assert child is not None - index.record_result(child, "completed") + await tracker.mark_work_done( + task.task_id, + child, + account_id="acc", + user_id="alice", + ) with pytest.raises(asyncio.CancelledError): await asyncio.wait_for(worker, 1) assert events == ["embedding-settled", "released"] diff --git a/tests/test_task_tracker.py b/tests/test_task_tracker.py index b322bfad39..ca5d223b66 100644 --- a/tests/test_task_tracker.py +++ b/tests/test_task_tracker.py @@ -14,7 +14,11 @@ from openviking.pyagfs.exceptions import AGFSAlreadyExistsError, AGFSNotFoundError from openviking.server.identity import RequestContext, Role from openviking.service.session_service import SessionService -from openviking.service.task_store import PersistentTaskStore, _task_to_payload +from openviking.service.task_store import ( + PersistentTaskStore, + _task_from_payload, + _task_to_payload, +) from openviking.service.task_tracker import ( TaskStatus, TaskTracker, @@ -23,11 +27,7 @@ run_tracked_operation, set_task_tracker, ) -from openviking.service.task_work_index import ( - QueueTaskMetadata, - TaskWorkIndex, - get_task_context, -) +from openviking.service.task_work_index import get_task_context from openviking_cli.exceptions import DeadlineExceededError from openviking_cli.session.user_id import UserIdentifier @@ -549,62 +549,53 @@ async def test_sanitize_preserves_safe_error(): @pytest.mark.parametrize("already_missing", [False, True]) async def test_evict_expired_completed(already_missing): agfs = _FakeAgfs() - tracker = TaskTracker(store=PersistentTaskStore(agfs)) - t = await tracker.create("session_commit", **_owner_kwargs()) - await tracker.start(t.task_id) - await tracker.complete(t.task_id, {}) - assert await tracker._store.get(t.task_id, **_owner_kwargs()) is not None + writer = TaskTracker(store=PersistentTaskStore(agfs)) + task = await writer.create("session_commit", **_owner_kwargs()) + await writer.start(task.task_id) + await writer.complete(task.task_id, {}) + path = f"/local/acme/_system/tasks/alice/{task.task_id}.json" if already_missing: agfs.files.clear() - # Simulate old timestamp (access internal state; get() returns defensive copies) - tracker._tasks[t.task_id].updated_at = time.time() - tracker.TTL_COMPLETED - 1 - await tracker._evict_expired() - assert await tracker.get(t.task_id) is None - assert await tracker._store.get(t.task_id, **_owner_kwargs()) is None + else: + payload = json.loads(agfs.files[path]) + payload["updated_at"] = time.time() - writer.TTL_COMPLETED - 1 + agfs.files[path] = json.dumps(payload).encode() + + tracker = TaskTracker(store=PersistentTaskStore(agfs)) + await tracker.delete_expired() + assert await tracker.get(task.task_id, **_owner_kwargs()) is None -async def test_evict_keeps_cached_task_when_persistent_delete_fails(): +async def test_delete_expired_retries_after_persistent_delete_fails(): agfs = _FakeAgfs() + writer = TaskTracker(store=PersistentTaskStore(agfs)) + task = await writer.create("session_commit", **_owner_kwargs()) + await writer.start(task.task_id) + await writer.complete(task.task_id, {}) + path = f"/local/acme/_system/tasks/alice/{task.task_id}.json" + payload = json.loads(agfs.files[path]) + payload["updated_at"] = time.time() - writer.TTL_COMPLETED - 1 + agfs.files[path] = json.dumps(payload).encode() tracker = TaskTracker(store=PersistentTaskStore(agfs)) - t = await tracker.create("session_commit", **_owner_kwargs()) - await tracker.start(t.task_id) - await tracker.complete(t.task_id, {}) - tracker._tasks[t.task_id].updated_at = time.time() - tracker.TTL_COMPLETED - 1 agfs.fail_rm = True - await tracker._evict_expired() - - assert await tracker.get(t.task_id) is not None - assert await tracker._store.get(t.task_id, **_owner_kwargs()) is not None + with pytest.raises(OSError, match="simulated delete failure"): + await tracker.delete_expired() + assert await tracker.get(task.task_id, **_owner_kwargs()) is not None agfs.fail_rm = False - await tracker._evict_expired() + assert (await tracker.delete_expired()).deleted_count == 1 - assert await tracker.get(t.task_id) is None - assert await tracker._store.get(t.task_id, **_owner_kwargs()) is None + assert await tracker.get(task.task_id, **_owner_kwargs()) is None assert all(call["ctx"]["disable_auto_pathlock"] == "true" for call in agfs.rm_calls) async def test_evict_keeps_recent_completed(tracker: TaskTracker): - t = await tracker.create("session_commit", **_owner_kwargs()) - await tracker.start(t.task_id) - await tracker.complete(t.task_id, {}) - await tracker._evict_expired() - assert await tracker.get(t.task_id) is not None - - -async def test_evict_fifo_when_over_limit(tracker: TaskTracker): - tracker.MAX_TASKS = 5 - tasks = [] - for i in range(7): - tasks.append(await tracker.create("session_commit", resource_id=f"s{i}", **_owner_kwargs())) - await tracker._evict_expired() - assert tracker.count() == 5 - # Oldest should be gone - assert await tracker.get(tasks[0].task_id) is None - assert await tracker.get(tasks[1].task_id) is None - # Newest should remain - assert await tracker.get(tasks[6].task_id) is not None + task = await tracker.create("session_commit", **_owner_kwargs()) + await tracker.start(task.task_id) + await tracker.complete(task.task_id, {}) + assert (await tracker.delete_expired()).deleted_count == 0 + assert await tracker.get(task.task_id) is not None # ── Singleton ── @@ -629,21 +620,19 @@ async def test_get_task_tracker_requires_service_initialization(): async def test_wait_timeout_preserves_recorded_operation_result(tracker): - index = TaskWorkIndex() - tracker.attach_work_index(index) set_task_tracker(tracker) state = {} async def operation(): context = get_task_context() state["task_id"] = context.task_id - work = QueueTaskMetadata( + await tracker.register_work( context.task_id, "child-work", - context.account_id, - context.user_id, + "Semantic", + account_id=context.account_id, + user_id=context.user_id, ) - assert index.register("Semantic", work) return {"status": "success", "root_uri": "resource-result"} with pytest.raises(DeadlineExceededError) as caught: @@ -666,26 +655,58 @@ async def operation(): } +async def test_inline_operation_can_wait_for_task_work_without_waiting_for_itself( + tracker, +): + set_task_tracker(tracker) + + async def operation(): + context = get_task_context() + assert context is not None + assert context.work_id is not None + await tracker.wait_for_task_work( + context.task_id, + exclude_work_id=context.work_id, + timeout=0, + **_owner_kwargs(), + ) + return {"status": "success"} + + result, task_id = await run_tracked_operation( + "content_write", + **_owner_kwargs(), + operation=operation, + wait=True, + timeout=0, + ) + + assert result == {"status": "success"} + assert (await tracker.get(task_id, **_owner_kwargs())).status is TaskStatus.COMPLETED + + async def test_accept_failed_result_returns_structured_partial_result(tracker): - index = TaskWorkIndex() - tracker.attach_work_index(index) set_task_tracker(tracker) workers = [] async def operation(): context = get_task_context() - work = QueueTaskMetadata( + await tracker.register_work( context.task_id, "embedding-work", - context.account_id, - context.user_id, + "Embedding", + account_id=context.account_id, + user_id=context.user_id, ) - assert index.register("Embedding", work) async def fail_work(): await asyncio.sleep(0) - index.record_result(work, "failed", error="embedding failed") - await index.prepare_ack("Embedding", work) + await tracker.mark_work_failed( + context.task_id, + "embedding-work", + "embedding failed", + account_id=context.account_id, + user_id=context.user_id, + ) workers.append(asyncio.create_task(fail_work())) return { @@ -850,9 +871,30 @@ async def operation(): @pytest.mark.parametrize("timeout", [0, 0.01]) async def test_wait_timeout_excludes_operation_time_for_completed_task(tracker, timeout): set_task_tracker(tracker) + workers = [] async def operation(): await asyncio.sleep(0.02) + if timeout: + context = get_task_context() + await tracker.register_work( + context.task_id, + "semantic", + "Semantic", + account_id="acme", + user_id="alice", + ) + + async def complete_work(): + await asyncio.sleep(0.001) + await tracker.mark_work_done( + context.task_id, + "semantic", + account_id="acme", + user_id="alice", + ) + + workers.append(asyncio.create_task(complete_work())) return {"status": "success"} result, task_id = await run_tracked_operation( @@ -869,27 +911,22 @@ async def operation(): assert completed.status is TaskStatus.COMPLETED assert completed.result["status"] == "success" assert completed.result["task_id"] == task_id + await asyncio.gather(*workers) async def test_multiple_queue_waits_share_budget_and_leave_task_pollable(tracker): - index = TaskWorkIndex() - tracker.attach_work_index(index) set_task_tracker(tracker) state = {} async def operation(): task_id = get_task_context().task_id state["task_id"] = task_id - first = QueueTaskMetadata(task_id, "first", "acme", "alice") - assert index.register("Semantic", first) + await tracker.register_work(task_id, "first", "Semantic", **_owner_kwargs()) # Exhaust the budget on a caught wait, then make the first Work complete. with pytest.raises(TimeoutError): - await tracker.wait_for_work(task_id, [first.work_id], timeout=0.01) - index.record_result(first, "completed") - await index.prepare_ack("Semantic", first) - second = QueueTaskMetadata(task_id, "second", "acme", "alice") - assert index.register("Embedding", second) - state["second"] = second + await tracker.wait_for_work(task_id, ["first"], timeout=0.01) + await tracker.mark_work_done(task_id, "first", **_owner_kwargs()) + await tracker.register_work(task_id, "second", "Embedding", **_owner_kwargs()) with pytest.raises(TimeoutError): await tracker.wait_for_task_work(task_id, timeout=60) return {"status": "success"} @@ -909,8 +946,7 @@ async def operation(): task = await tracker.get(state["task_id"], **_owner_kwargs()) assert task.status is TaskStatus.RUNNING assert task.error is None - index.record_result(state["second"], "completed") - await index.prepare_ack("Embedding", state["second"]) + await tracker.mark_work_done(task.task_id, "second", **_owner_kwargs()) assert ( await tracker.wait(task.task_id, **_owner_kwargs(), timeout=0) ).status is TaskStatus.COMPLETED @@ -956,7 +992,7 @@ async def test_persistent_task_extensions_survive_updates(extra): agfs.files[path] = json.dumps(payload).encode() original = agfs.files[path] - reader = TaskTracker(store=store) + reader = TaskTracker(store=PersistentTaskStore(agfs)) loaded = await reader.get(task.task_id, **_owner_kwargs()) assert loaded is not None assert loaded.status == TaskStatus.PENDING @@ -971,7 +1007,9 @@ async def test_persistent_task_extensions_survive_updates(extra): assert saved["status"] == "completed" assert saved["result"] == {"ok": True} assert {key: saved[key] for key in extra} == extra - reloaded = await TaskTracker(store=store).get(task.task_id, **_owner_kwargs()) + reloaded = await TaskTracker(store=PersistentTaskStore(agfs)).get( + task.task_id, **_owner_kwargs() + ) assert reloaded is not None assert reloaded.status == TaskStatus.COMPLETED assert {key: _task_to_payload(reloaded)[key] for key in extra} == extra @@ -986,7 +1024,7 @@ async def test_task_extension_payload_is_defensively_copied(): "meta": {"values": [2]}, } original = deepcopy(payload) - record = TaskTracker._record_from_payload(payload) + record = _task_from_payload(payload) record._extra_fields["future_field"]["values"].append(3) record.meta["values"].append(4) assert payload == original @@ -1000,7 +1038,7 @@ async def test_task_extension_payload_is_defensively_copied(): @pytest.mark.parametrize("status", ["unknown", None]) async def test_task_extensions_do_not_hide_invalid_status(status): with pytest.raises(ValueError): - TaskTracker._record_from_payload( + _task_from_payload( {"task_id": "task-1", "task_type": "session_commit", "status": status, "extra": 1} ) @@ -1040,10 +1078,10 @@ async def test_persistent_store_writes_task_record_json(): assert all(call["ctx"]["disable_auto_pathlock"] == "true" for call in agfs.write_calls) -async def test_persistent_store_keeps_tasktracker_tasks_dict(): +async def test_tracker_reads_tasks_from_persistent_store(): tracker = TaskTracker(store=PersistentTaskStore(_FakeAgfs())) task = await tracker.create("session_commit", **_owner_kwargs()) - assert task.task_id in tracker._tasks + assert await tracker.get(task.task_id, **_owner_kwargs()) == task async def test_persistent_store_survives_tracker_reset(): @@ -1178,21 +1216,16 @@ async def test_execution_events_survive_restart_and_preserve_reported_facts(): @pytest.mark.parametrize("outcome", ["complete", "fail"]) async def test_terminal_event_waits_for_owned_work(tracker, outcome): - from openviking.service.task_work_index import QueueTaskMetadata, TaskWorkIndex - owner = _owner_kwargs() task = await tracker.create("add_resource", **owner) await tracker.start(task.task_id, **owner) - work_index = TaskWorkIndex() - tracker.attach_work_index(work_index) - child = QueueTaskMetadata(task.task_id, "child-work", "acme", "alice") - work_index.register("Semantic", child) + await tracker.register_work(task.task_id, "child-work", "Semantic", **owner) await getattr(tracker, outcome)( task.task_id, {"done": True} if outcome == "complete" else "error", **owner ) snapshot = await tracker.get(task.task_id, **owner) assert snapshot.status == TaskStatus.RUNNING - await work_index.prepare_ack("Semantic", child) + await tracker.mark_work_done(task.task_id, "child-work", **owner) snapshot = await tracker.get(task.task_id, **owner) last = snapshot.execution_events["items"][-1] assert last["kind"] == "status_changed" @@ -1207,7 +1240,7 @@ async def test_legacy_tasks_do_not_get_invented_history(): payload = json.loads(agfs.files[path]) payload.pop("execution_events") agfs.files[path] = json.dumps(payload).encode() - reader = TaskTracker(store=store) + reader = TaskTracker(store=PersistentTaskStore(agfs)) assert (await reader.get(task.task_id, **_owner_kwargs())).execution_events is None await reader.start(task.task_id, **_owner_kwargs()) history = (await reader.get(task.task_id, **_owner_kwargs())).execution_events @@ -1258,18 +1291,20 @@ async def test_process_events_respect_owner_and_terminal_boundaries(tracker): ).execution_events == cancelled.execution_events -async def test_processing_time_persists_and_restored_active_tasks_are_unknown(tracker, monkeypatch): +async def test_processing_time_persists_and_restored_active_tasks_are_unknown(monkeypatch): now = [0.0] - monkeypatch.setattr("openviking.service.task_work_index.time.time", lambda: now[0]) + monkeypatch.setattr("openviking.service.task_store.time.time", lambda: now[0]) + agfs = _FakeAgfs() + tracker = TaskTracker(PersistentTaskStore(agfs)) task = await tracker.create("add_resource", **_owner_kwargs()) await tracker.start(task.task_id, **_owner_kwargs()) now[0] = 10 - tracker.register_running_task(task.task_id) + await tracker.register_running_task(task.task_id, **_owner_kwargs()) now[0] = 13 assert (await tracker.get(task.task_id)).processing_seconds == 3 # A stage update neither resets nor stops the processing clock. await tracker.update_stage(task.task_id, "semantic", **_owner_kwargs()) - restarted = TaskTracker(store=tracker._store) + restarted = TaskTracker(store=PersistentTaskStore(agfs)) assert (await restarted.get(task.task_id, **_owner_kwargs())).processing_seconds is None # Waiting, I/O, and backoff after a worker claims the task remain part of # worker-occupied wall time. @@ -1281,5 +1316,7 @@ async def test_processing_time_persists_and_restored_active_tasks_are_unknown(tr finished = await tracker.get(task.task_id) assert finished.status == TaskStatus.COMPLETED assert finished.to_dict()["processing_seconds"] == 22 - restored = await TaskTracker(store=tracker._store).get(task.task_id, **_owner_kwargs()) + restored = await TaskTracker(store=PersistentTaskStore(agfs)).get( + task.task_id, **_owner_kwargs() + ) assert restored.processing_seconds == 22 diff --git a/tests/test_task_tracker_concurrency.py b/tests/test_task_tracker_concurrency.py index 342c4566d5..7292d3efd4 100644 --- a/tests/test_task_tracker_concurrency.py +++ b/tests/test_task_tracker_concurrency.py @@ -5,149 +5,38 @@ import gc import threading import weakref -from copy import deepcopy -from typing import Any import pytest -class _ControllableTaskStore: - def __init__(self) -> None: - self.payloads: dict[str, dict[str, Any]] = {} - self.create_calls = 0 - self.update_started: dict[str, asyncio.Event] = {} - self.update_release: dict[str, asyncio.Event] = {} - self.update_errors: dict[str, Exception] = {} - self.list_started: asyncio.Event | None = None - self.list_release: asyncio.Event | None = None - self.get_calls = 0 - self.operation_loops: list[asyncio.AbstractEventLoop] = [] - - @staticmethod - def _payload(task: Any) -> dict[str, Any]: - return { - "task_id": task.task_id, - "task_type": task.task_type, - "status": task.status.value, - "created_at": task.created_at, - "updated_at": task.updated_at, - "resource_id": task.resource_id, - "account_id": task.account_id, - "user_id": task.user_id, - "meta": deepcopy(task.meta), - "stage": task.stage, - "result": deepcopy(task.result), - "error": task.error, - "execution_events": deepcopy(task.execution_events), - } - - async def create(self, task: Any) -> None: - self.operation_loops.append(asyncio.get_running_loop()) - self.create_calls += 1 - self.payloads[task.task_id] = self._payload(task) - - async def update(self, task: Any) -> None: - self.operation_loops.append(asyncio.get_running_loop()) - started = self.update_started.get(task.task_id) - if started is not None: - started.set() - release = self.update_release.get(task.task_id) - if release is not None: - await release.wait() - error = self.update_errors.get(task.task_id) - if error is not None: - raise error - self.payloads[task.task_id] = self._payload(task) - - async def get( - self, - task_id: str, - *, - account_id: str | None = None, - user_id: str | None = None, - ) -> dict[str, Any] | None: - self.operation_loops.append(asyncio.get_running_loop()) - self.get_calls += 1 - payload = self.payloads.get(task_id) - if payload is None: - return None - if account_id is not None and payload["account_id"] != account_id: - return None - if user_id is not None and payload["user_id"] != user_id: - return None - return deepcopy(payload) - - async def list( - self, - account_id: str, - *, - user_id: str | None = None, - ) -> list[dict[str, Any]]: - self.operation_loops.append(asyncio.get_running_loop()) - snapshot = [ - deepcopy(payload) - for payload in self.payloads.values() - if payload["account_id"] == account_id - and (user_id is None or payload["user_id"] == user_id) - ] - if self.list_started is not None: - self.list_started.set() - if self.list_release is not None: - await self.list_release.wait() - return snapshot - - async def delete( - self, - task_id: str, - *, - account_id: str, - user_id: str | None = None, - ) -> None: - self.payloads.pop(task_id, None) - - -class _BlockingCreateTaskStore(_ControllableTaskStore): - def __init__(self, expected_concurrency: int) -> None: - super().__init__() - self.expected_concurrency = expected_concurrency - self.create_release = asyncio.Event() - self.expected_entered = asyncio.Event() - self.create_inflight = 0 - self.max_create_inflight = 0 - - async def create(self, task: Any) -> None: - self.create_inflight += 1 - self.max_create_inflight = max(self.max_create_inflight, self.create_inflight) - if self.create_inflight == self.expected_concurrency: - self.expected_entered.set() - try: - await self.create_release.wait() - await super().create(task) - finally: - self.create_inflight -= 1 - - -class _ThreadBlockingUpdateTaskStore(_ControllableTaskStore): - def __init__(self) -> None: - super().__init__() - self.thread_update_started = threading.Event() - self.thread_update_release = threading.Event() - self._block_next_update = True +@pytest.mark.asyncio +async def test_work_index_groups_and_cancels_active_tasks_by_owner(): + from openviking.service.task_work_index import TaskWorkIndex + + index = TaskWorkIndex() + first = asyncio.create_task(asyncio.Event().wait()) + second = asyncio.create_task(asyncio.Event().wait()) + index.register_active( + "shared-id", first, account_id="account-a", user_id="user" + ) + index.register_active( + "shared-id", second, account_id="account-b", user_id="user" + ) - async def update(self, task: Any) -> None: - if not self._block_next_update: - await super().update(task) - return - self._block_next_update = False - payload = self._payload(task) + assert index.active_tasks_by_owner() == { + ("account-a", "user"): {"shared-id"}, + ("account-b", "user"): {"shared-id"}, + } - def blocking_write() -> None: - self.thread_update_started.set() - self.thread_update_release.wait() - self.payloads[task.task_id] = payload + index.cancel_active( + "shared-id", account_id="account-a", user_id="user" + ) + await asyncio.gather(first, return_exceptions=True) + assert first.cancelled() + assert not second.done() - self.operation_loops.append(asyncio.get_running_loop()) - await asyncio.to_thread(blocking_write) + second.cancel() + await asyncio.gather(second, return_exceptions=True) @pytest.mark.asyncio @@ -164,32 +53,12 @@ async def work() -> str: work_loop = asyncio.get_running_loop() return "done" - def run_from_foreign_loop() -> str: - return asyncio.run(dispatcher.run(work)) - - result = await asyncio.to_thread(run_from_foreign_loop) + result = await asyncio.to_thread(lambda: asyncio.run(dispatcher.run(work))) assert result == "done" assert work_loop is owner_loop -@pytest.mark.asyncio -async def test_owner_loop_dispatcher_uses_explicit_owner_loop(): - from openviking.service.task_tracker_concurrency import OwnerLoopDispatcher - - owner_loop = asyncio.get_running_loop() - dispatcher = await asyncio.to_thread(OwnerLoopDispatcher, owner_loop) - work_loop: asyncio.AbstractEventLoop | None = None - - async def work() -> None: - nonlocal work_loop - work_loop = asyncio.get_running_loop() - - await asyncio.to_thread(lambda: asyncio.run(dispatcher.run(work))) - - assert work_loop is owner_loop - - @pytest.mark.asyncio async def test_owner_loop_dispatcher_propagates_foreign_cancellation(): from openviking.service.task_tracker_concurrency import OwnerLoopDispatcher @@ -228,45 +97,36 @@ async def invoke() -> None: await asyncio.wait_for(work_cancelled.wait(), timeout=1) -def test_owner_loop_dispatcher_rejects_calls_after_owner_loop_closes(): +def test_owner_loop_dispatcher_rejects_closed_or_stopped_owner_loop(): from openviking.service.task_tracker_concurrency import OwnerLoopDispatcher - dispatcher = OwnerLoopDispatcher() + closed_dispatcher = OwnerLoopDispatcher() - async def bind() -> None: - dispatcher.bind_current_loop() + async def bind_closed() -> None: + closed_dispatcher.bind_current_loop() - asyncio.run(bind()) + asyncio.run(bind_closed()) - async def invoke() -> None: + async def invoke_closed() -> None: with pytest.raises(RuntimeError, match="owner event loop is closed"): - await dispatcher.run(lambda: asyncio.sleep(0)) - - asyncio.run(invoke()) - + await closed_dispatcher.run(lambda: asyncio.sleep(0)) -def test_owner_loop_dispatcher_rejects_stopped_owner_loop(): - from openviking.service.task_tracker_concurrency import OwnerLoopDispatcher + asyncio.run(invoke_closed()) - dispatcher = OwnerLoopDispatcher() owner_loop = asyncio.new_event_loop() + stopped_dispatcher = OwnerLoopDispatcher(owner_loop) owner_loop.run_until_complete(asyncio.sleep(0)) - owner_loop.run_until_complete(_bind_dispatcher(dispatcher)) - async def invoke() -> None: + async def invoke_stopped() -> None: with pytest.raises(RuntimeError, match="owner event loop is not running"): - await dispatcher.run(lambda: asyncio.sleep(0)) + await stopped_dispatcher.run(lambda: asyncio.sleep(0)) try: - asyncio.run(invoke()) + asyncio.run(invoke_stopped()) finally: owner_loop.close() -async def _bind_dispatcher(dispatcher: Any) -> None: - dispatcher.bind_current_loop() - - def test_owner_loop_dispatcher_concurrent_first_calls_share_winning_loop(): from openviking.service.task_tracker_concurrency import OwnerLoopDispatcher @@ -298,97 +158,6 @@ def invoke() -> None: assert results[0] == results[1] -def test_owner_loop_dispatcher_fails_if_loop_stops_during_dispatch(): - from openviking.service.task_tracker_concurrency import OwnerLoopDispatcher - - dispatcher = OwnerLoopDispatcher() - owner_loop = asyncio.new_event_loop() - owner_ready = threading.Event() - stop_requested = threading.Event() - let_stop_callback_return = threading.Event() - first_run_stopped = threading.Event() - cleanup_queued_callbacks = threading.Event() - - def owner_thread_main() -> None: - asyncio.set_event_loop(owner_loop) - - def bind() -> None: - dispatcher.bind_current_loop() - owner_ready.set() - - owner_loop.call_soon(bind) - owner_loop.run_forever() - first_run_stopped.set() - cleanup_queued_callbacks.wait() - owner_loop.call_soon(owner_loop.stop) - owner_loop.run_forever() - owner_loop.close() - - owner_thread = threading.Thread(target=owner_thread_main) - owner_thread.start() - assert owner_ready.wait(timeout=1) - - def stop_inside_owner_callback() -> None: - owner_loop.stop() - stop_requested.set() - let_stop_callback_return.wait() - - owner_loop.call_soon_threadsafe(stop_inside_owner_callback) - assert stop_requested.wait(timeout=1) - - async def invoke() -> None: - async def release_callback() -> None: - await asyncio.sleep(0.05) - let_stop_callback_return.set() - - release_task = asyncio.create_task(release_callback()) - with pytest.raises(RuntimeError, match="owner event loop stopped"): - await asyncio.wait_for( - dispatcher.run(lambda: asyncio.sleep(0)), - timeout=1, - ) - await release_task - - asyncio.run(invoke()) - assert first_run_stopped.wait(timeout=1) - cleanup_queued_callbacks.set() - owner_thread.join(timeout=2) - assert not owner_thread.is_alive() - - -def test_owner_loop_dispatcher_prefers_completed_result_over_simultaneous_stop(): - from openviking.service.task_tracker_concurrency import OwnerLoopDispatcher - - dispatcher = OwnerLoopDispatcher() - owner_loop = asyncio.new_event_loop() - owner_ready = threading.Event() - - def owner_thread_main() -> None: - asyncio.set_event_loop(owner_loop) - - def bind() -> None: - dispatcher.bind_current_loop() - owner_ready.set() - - owner_loop.call_soon(bind) - owner_loop.run_forever() - owner_loop.close() - - owner_thread = threading.Thread(target=owner_thread_main) - owner_thread.start() - assert owner_ready.wait(timeout=1) - - async def finish_and_stop() -> str: - owner_loop.stop() - return "committed" - - result = asyncio.run(dispatcher.run(finish_and_stop)) - - owner_thread.join(timeout=2) - assert result == "committed" - assert not owner_thread.is_alive() - - @pytest.mark.asyncio async def test_store_io_limiter_bounds_concurrency(): from openviking.service.task_tracker_concurrency import StoreIOLimiter @@ -428,7 +197,7 @@ def test_store_io_limiter_rejects_non_positive_limit(): @pytest.mark.asyncio @pytest.mark.parametrize("cancel", [False, True]) -async def test_run_to_completion_releases_failed_request_and_preserves_cancellation(cancel): +async def test_run_to_completion_releases_request_and_preserves_cancellation(cancel): from openviking.service.task_tracker_concurrency import run_to_completion started = asyncio.Event() @@ -442,7 +211,7 @@ async def fail_later() -> None: await release.wait() raise RuntimeError("late store failure") - async def request(): + async def request() -> weakref.ReferenceType[RequestPayload]: payload = RequestPayload() try: await run_to_completion(fail_later) @@ -451,10 +220,9 @@ async def request(): if not cancel: assert str(exc) == "late store failure" else: - pytest.fail("The operation must propagate its failure or caller cancellation") + pytest.fail("The operation must propagate failure or caller cancellation") return weakref.ref(payload) - # Finished requests must be released without waiting for cyclic GC. gc_was_enabled = gc.isenabled() gc.disable() try: @@ -503,15 +271,10 @@ async def work() -> None: @pytest.mark.asyncio -async def test_keyed_lock_registry_is_cleaned_after_use_and_cancellation(): +async def test_keyed_lock_registry_is_cleaned_after_waiter_cancellation(): from openviking.service.task_tracker_concurrency import KeyedAsyncLockPool pool = KeyedAsyncLockPool[str]() - - async with pool.acquire("completed"): - assert pool.entry_count == 1 - assert pool.entry_count == 0 - blocker_entered = asyncio.Event() release_blocker = asyncio.Event() @@ -521,13 +284,12 @@ async def blocker() -> None: await release_blocker.wait() async def waiter() -> None: - await blocker_entered.wait() async with pool.acquire("cancelled"): pass blocker_task = asyncio.create_task(blocker()) - waiter_task = asyncio.create_task(waiter()) await blocker_entered.wait() + waiter_task = asyncio.create_task(waiter()) await asyncio.sleep(0) waiter_task.cancel() @@ -538,272 +300,14 @@ async def waiter() -> None: await blocker_task assert pool.entry_count == 0 - # Cancellation can arrive after release hands over the slot, but before - # the waiting coroutine resumes. The slot must still become available. - async with pool.acquire("cancelled"): - granted_waiter = asyncio.create_task(waiter()) - await asyncio.sleep(0) - granted_waiter.cancel() - with pytest.raises(asyncio.CancelledError): - await granted_waiter - await asyncio.wait_for(waiter(), timeout=1) - assert pool.entry_count == 0 - @pytest.mark.asyncio -async def test_task_tracker_different_task_updates_do_not_block_each_other(): - from openviking.service.task_tracker import TaskStatus, TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - first = await tracker.create("add_resource", account_id="acme", user_id="alice") - second = await tracker.create("add_resource", account_id="acme", user_id="alice") - store.update_started[first.task_id] = asyncio.Event() - store.update_release[first.task_id] = asyncio.Event() - - first_update = asyncio.create_task(tracker.start(first.task_id)) - await store.update_started[first.task_id].wait() - - await asyncio.wait_for(tracker.start(second.task_id), timeout=1) - second_snapshot = await tracker.get(second.task_id) - assert second_snapshot is not None - assert second_snapshot.status == TaskStatus.RUNNING - - store.update_release[first.task_id].set() - await first_update - - -@pytest.mark.asyncio -async def test_task_tracker_cache_hit_get_is_not_blocked_by_in_flight_update(): - from openviking.service.task_tracker import TaskStatus, TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - store.update_started[task.task_id] = asyncio.Event() - store.update_release[task.task_id] = asyncio.Event() - - update = asyncio.create_task(tracker.start(task.task_id)) - await store.update_started[task.task_id].wait() - - snapshot = await asyncio.wait_for(tracker.get(task.task_id), timeout=0.1) - assert snapshot is not None - assert snapshot.status == TaskStatus.PENDING - - store.update_release[task.task_id].set() - await update - - -@pytest.mark.asyncio -async def test_task_tracker_failed_update_does_not_contaminate_cached_snapshot(): - from openviking.service.task_tracker import TaskStatus, TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - store.update_errors[task.task_id] = RuntimeError("store unavailable") - - with pytest.raises(RuntimeError, match="store unavailable"): - await tracker.start(task.task_id) - - snapshot = await tracker.get(task.task_id) - assert snapshot is not None - assert snapshot.status == TaskStatus.PENDING - assert snapshot.execution_events == task.execution_events - assert store.payloads[task.task_id]["execution_events"] == task.execution_events - - -@pytest.mark.asyncio -async def test_cancelled_thread_write_settles_before_later_same_task_mutation(): - from openviking.service.task_tracker import TaskStatus, TaskTracker - - store = _ThreadBlockingUpdateTaskStore() - tracker = TaskTracker(store=store) - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - - starting = asyncio.create_task(tracker.start(task.task_id)) - await asyncio.to_thread(store.thread_update_started.wait) - starting.cancel() - completing = asyncio.create_task(tracker.complete(task.task_id, {"done": True})) - await asyncio.sleep(0) - assert not completing.done() - - store.thread_update_release.set() - with pytest.raises(asyncio.CancelledError): - await starting - await completing - - snapshot = await tracker.get(task.task_id) - assert snapshot is not None - assert snapshot.status == TaskStatus.COMPLETED - assert store.payloads[task.task_id]["status"] == TaskStatus.COMPLETED.value - assert store.payloads[task.task_id]["execution_events"] == snapshot.execution_events - assert [event["status"] for event in snapshot.execution_events["items"]] == [ - "pending", - "running", - "completed", - ] - assert tracker._task_locks.entry_count == 0 - assert tracker._store_io.inflight == 0 - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("outcome", "expected_status"), - [("complete", "completed"), ("fail", "failed")], -) -async def test_cancelled_outcome_finishes_terminal_transition(outcome, expected_status): +async def test_file_store_concurrent_fixed_id_create_is_idempotent(): + from openviking.service.task_store import PersistentTaskStore from openviking.service.task_tracker import TaskTracker + from tests.test_task_tracker import _FakeAgfs - store = _ThreadBlockingUpdateTaskStore() - tracker = TaskTracker(store=store) - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - - if outcome == "complete": - operation = asyncio.create_task(tracker.complete(task.task_id, {"done": True})) - else: - operation = asyncio.create_task(tracker.fail(task.task_id, "failed on purpose")) - await asyncio.to_thread(store.thread_update_started.wait) - operation.cancel() - store.thread_update_release.set() - - with pytest.raises(asyncio.CancelledError): - await operation - - snapshot = await tracker.get(task.task_id) - assert snapshot is not None - assert snapshot.status.value == expected_status - assert store.payloads[task.task_id]["status"] == expected_status - assert tracker._task_locks.entry_count == 0 - assert tracker._store_io.inflight == 0 - - -@pytest.mark.asyncio -async def test_cancelled_cancel_request_still_cancels_active_work_and_finalizes(): - from openviking.service.task_tracker import TaskStatus, TaskTracker - - store = _ThreadBlockingUpdateTaskStore() - tracker = TaskTracker(store=store) - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - work_started = asyncio.Event() - work_cancelled = asyncio.Event() - - async def active_work() -> None: - tracker.register_running_task(task.task_id) - work_started.set() - try: - await asyncio.Future() - except asyncio.CancelledError: - work_cancelled.set() - raise - finally: - await tracker.unregister_running_task(task.task_id) - - worker = asyncio.create_task(active_work()) - await work_started.wait() - cancelling = asyncio.create_task(tracker.cancel(task.task_id)) - await asyncio.to_thread(store.thread_update_started.wait) - cancelling.cancel() - store.thread_update_release.set() - - with pytest.raises(asyncio.CancelledError): - await cancelling - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(worker, timeout=1) - await asyncio.wait_for(work_cancelled.wait(), timeout=1) - - snapshot = await tracker.get(task.task_id) - assert snapshot is not None - assert snapshot.status == TaskStatus.CANCELLED - assert store.payloads[task.task_id]["status"] == TaskStatus.CANCELLED.value - assert tracker._task_locks.entry_count == 0 - assert tracker._store_io.inflight == 0 - - -@pytest.mark.asyncio -async def test_cancel_request_cancelled_before_store_admission_does_not_cancel_work(): - from openviking.service.task_tracker import TaskStatus, TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store, max_concurrent_store_io=1) - blocker = await tracker.create("add_resource", account_id="acme", user_id="alice") - target = await tracker.create("add_resource", account_id="acme", user_id="alice") - await tracker.start(target.task_id) - store.update_started[blocker.task_id] = asyncio.Event() - store.update_release[blocker.task_id] = asyncio.Event() - - blocking_update = asyncio.create_task(tracker.start(blocker.task_id)) - await store.update_started[blocker.task_id].wait() - - work_started = asyncio.Event() - work_cancelled = asyncio.Event() - - async def active_work() -> None: - tracker.register_running_task(target.task_id) - work_started.set() - try: - await asyncio.Future() - except asyncio.CancelledError: - work_cancelled.set() - raise - - worker = asyncio.create_task(active_work()) - await work_started.wait() - cancelling = asyncio.create_task(tracker.cancel(target.task_id)) - for _ in range(100): - if tracker._task_locks.entry_count == 2: - break - await asyncio.sleep(0.01) - - cancelling.cancel() - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(cancelling, timeout=1) - - snapshot = await tracker.get(target.task_id) - assert snapshot is not None - assert snapshot.status == TaskStatus.RUNNING - assert not work_cancelled.is_set() - assert not worker.done() - - worker.cancel() - with pytest.raises(asyncio.CancelledError): - await worker - store.update_release[blocker.task_id].set() - await blocking_update - - -@pytest.mark.asyncio -async def test_same_task_late_mutations_cannot_revert_terminal_status(): - from openviking.service.task_tracker import TaskStatus, TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - await tracker.start(task.task_id) - store.update_started[task.task_id] = asyncio.Event() - store.update_release[task.task_id] = asyncio.Event() - - completing = asyncio.create_task(tracker.complete(task.task_id, {"done": True})) - await store.update_started[task.task_id].wait() - late_start = asyncio.create_task(tracker.start(task.task_id, stage="late-start")) - late_stage = asyncio.create_task(tracker.update_stage(task.task_id, "late-stage")) - - store.update_release[task.task_id].set() - await asyncio.gather(completing, late_start, late_stage) - - snapshot = await tracker.get(task.task_id) - assert snapshot is not None - assert snapshot.status == TaskStatus.COMPLETED - assert snapshot.stage == "completed" - - -@pytest.mark.asyncio -async def test_concurrent_create_with_fixed_task_id_is_idempotent(): - from openviking.service.task_tracker import TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) first, second = await asyncio.gather( tracker.create( "add_resource", @@ -820,268 +324,16 @@ async def test_concurrent_create_with_fixed_task_id_is_idempotent(): ) assert first.task_id == second.task_id == "fixed" - assert store.create_calls == 1 - with pytest.raises(ValueError, match="already belongs to another owner"): - await tracker.create( - "add_resource", - task_id="fixed", - account_id="acme", - user_id="bob", - ) - - -@pytest.mark.asyncio -async def test_task_tracker_cache_hit_does_not_read_store(): - from openviking.service.task_tracker import TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - - snapshot = await tracker.get(task.task_id, account_id="acme", user_id="alice") - - assert snapshot is not None - assert snapshot.task_id == task.task_id - assert store.get_calls == 0 - - -@pytest.mark.asyncio -async def test_task_tracker_list_does_not_replace_newer_cache_with_stale_store_data(): - from openviking.service.task_tracker import TaskStatus, TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - stale_payload = deepcopy(store.payloads[task.task_id]) - - await tracker.start(task.task_id) - store.payloads[task.task_id] = stale_payload - - snapshots = await tracker.list_tasks(account_id="acme", user_id="alice") - - assert len(snapshots) == 1 - assert snapshots[0].status == TaskStatus.RUNNING - - -@pytest.mark.asyncio -async def test_blocked_store_list_does_not_block_cached_get(): - from openviking.service.task_tracker import TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - store.list_started = asyncio.Event() - store.list_release = asyncio.Event() - - listing = asyncio.create_task(tracker.list_tasks(account_id="acme", user_id="alice")) - await store.list_started.wait() - - snapshot = await asyncio.wait_for(tracker.get(task.task_id), timeout=0.1) - assert snapshot is not None - assert snapshot.task_id == task.task_id - - store.list_release.set() - await listing - - -@pytest.mark.asyncio -async def test_stale_in_flight_list_does_not_overwrite_completed_cache(): - from openviking.service.task_tracker import TaskStatus, TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - await tracker.start(task.task_id) - store.list_started = asyncio.Event() - store.list_release = asyncio.Event() - - listing = asyncio.create_task(tracker.list_tasks(account_id="acme", user_id="alice")) - await store.list_started.wait() - await tracker.complete(task.task_id, {"root_uri": "viking://resources/done"}) - - store.list_release.set() - await listing - - snapshot = await tracker.get(task.task_id) - assert snapshot is not None - assert snapshot.status == TaskStatus.COMPLETED - - -@pytest.mark.asyncio -async def test_clock_rollback_cannot_make_stale_store_snapshot_look_newer(monkeypatch): - from openviking.service.task_tracker import TaskStatus, TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - await tracker.start(task.task_id) - stale_payload = deepcopy(store.payloads[task.task_id]) - monkeypatch.setattr("openviking.service.task_tracker.time.time", lambda: 0.0) - - await tracker.complete(task.task_id, {"done": True}) - store.payloads[task.task_id] = stale_payload - await tracker.list_tasks(account_id="acme", user_id="alice") - - snapshot = await tracker.get(task.task_id) - assert snapshot is not None - assert snapshot.status == TaskStatus.COMPLETED - assert snapshot.updated_at > stale_payload["updated_at"] - - -@pytest.mark.asyncio -async def test_task_tracker_store_io_limit_allows_bounded_parallelism(): - from openviking.service.task_tracker import TaskTracker - - store = _BlockingCreateTaskStore(expected_concurrency=3) - tracker = TaskTracker(store=store, max_concurrent_store_io=3) - - creates = [ - asyncio.create_task(tracker.create("add_resource", account_id="acme", user_id="alice")) - for _ in range(10) - ] - await asyncio.wait_for(store.expected_entered.wait(), timeout=1) - await asyncio.sleep(0) - - assert store.max_create_inflight == 3 - store.create_release.set() - await asyncio.gather(*creates) - assert tracker._store_io.max_observed_inflight == 3 - assert tracker._store_io.inflight == 0 - - -@pytest.mark.asyncio -async def test_cancelled_create_waiting_for_store_slot_has_no_side_effect(): - from openviking.service.task_tracker import TaskTracker - - store = _BlockingCreateTaskStore(expected_concurrency=1) - tracker = TaskTracker(store=store, max_concurrent_store_io=1) - - first = asyncio.create_task(tracker.create("add_resource", account_id="acme", user_id="alice")) - await asyncio.wait_for(store.expected_entered.wait(), timeout=1) - - waiting = asyncio.create_task( - tracker.create("add_resource", account_id="acme", user_id="alice") - ) - for _ in range(100): - if tracker._task_locks.entry_count == 2: - break - await asyncio.sleep(0.01) - assert tracker._task_locks.entry_count == 2 - - waiting.cancel() - await asyncio.sleep(0.05) - cancelled_before_release = waiting.done() - store.create_release.set() - with pytest.raises(asyncio.CancelledError): - await waiting - - created = await first - - assert cancelled_before_release - assert store.create_calls == 1 - assert set(store.payloads) == {created.task_id} - assert tracker._task_locks.entry_count == 0 - assert tracker._store_io.inflight == 0 - - -@pytest.mark.asyncio -async def test_cancelled_create_waiting_for_task_lock_finishes_before_lock_release(): - from openviking.service.task_tracker import TaskTracker - - store = _BlockingCreateTaskStore(expected_concurrency=1) - tracker = TaskTracker(store=store) - create_args = { - "task_type": "add_resource", - "task_id": "fixed", - "account_id": "acme", - "user_id": "alice", - } - - first = asyncio.create_task(tracker.create(**create_args)) - await asyncio.wait_for(store.expected_entered.wait(), timeout=1) - waiting = asyncio.create_task(tracker.create(**create_args)) - await asyncio.sleep(0) - - waiting.cancel() - await asyncio.sleep(0.05) - cancelled_before_release = waiting.done() - store.create_release.set() - with pytest.raises(asyncio.CancelledError): - await waiting - - await first - assert cancelled_before_release - assert store.create_calls == 1 - - -@pytest.mark.asyncio -async def test_task_tracker_public_mutation_stays_on_calling_loop(): - from openviking.service.task_tracker import TaskStatus, TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - request_loop = asyncio.get_running_loop() - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - - await asyncio.to_thread(lambda: asyncio.run(tracker.start(task.task_id))) - - with tracker._lock: - tracker._tasks.pop(task.task_id) - loaded = await asyncio.to_thread( - lambda: asyncio.run(tracker.get(task.task_id, account_id="acme", user_id="alice")) - ) - - snapshot = await tracker.get(task.task_id) - assert loaded is not None - assert snapshot is not None - assert snapshot.status == TaskStatus.RUNNING - assert store.operation_loops[0] is request_loop - assert all(loop is not request_loop for loop in store.operation_loops[1:]) - - -@pytest.mark.asyncio -async def test_cancelling_foreign_mutation_cleans_task_lock_and_store_slot(): - from openviking.service.task_tracker import TaskStatus, TaskTracker - - store = _ThreadBlockingUpdateTaskStore() - tracker = TaskTracker(store=store) - task = await tracker.create("add_resource", account_id="acme", user_id="alice") - cancel_foreign = threading.Event() - - def run_from_foreign_loop() -> None: - async def invoke() -> None: - mutation = asyncio.create_task(tracker.start(task.task_id)) - await asyncio.to_thread(cancel_foreign.wait) - mutation.cancel() - with pytest.raises(asyncio.CancelledError): - await mutation - - asyncio.run(invoke()) - - foreign_call = asyncio.create_task(asyncio.to_thread(run_from_foreign_loop)) - assert await asyncio.to_thread(store.thread_update_started.wait, 3) - cancel_foreign.set() - await asyncio.sleep(0.05) - returned_before_store_settled = foreign_call.done() - store.thread_update_release.set() - await foreign_call - - assert not returned_before_store_settled - snapshot = await tracker.get(task.task_id) - assert snapshot is not None - assert snapshot.status == TaskStatus.RUNNING - assert tracker._task_locks.entry_count == 0 - assert tracker._store_io.inflight == 0 - assert store.payloads[task.task_id]["status"] == TaskStatus.RUNNING.value + assert len(await tracker.list_tasks(account_id="acme", user_id="alice")) == 1 @pytest.mark.asyncio -async def test_create_if_no_running_is_unique_per_business_key(): +async def test_file_store_create_if_no_running_is_unique_per_business_key(): + from openviking.service.task_store import PersistentTaskStore from openviking.service.task_tracker import TaskTracker + from tests.test_task_tracker import _FakeAgfs - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - + tracker = TaskTracker(PersistentTaskStore(_FakeAgfs())) results = await asyncio.gather( *[ tracker.create_if_no_running( @@ -1094,67 +346,4 @@ async def test_create_if_no_running_is_unique_per_business_key(): ] ) - created = [task for task in results if task is not None] - assert len(created) == 1 - assert tracker._business_locks.entry_count == 0 - assert tracker._task_locks.entry_count == 0 - - -@pytest.mark.asyncio -async def test_create_if_no_running_uses_independent_business_keys(): - from openviking.service.task_tracker import TaskTracker - - store = _BlockingCreateTaskStore(expected_concurrency=2) - tracker = TaskTracker(store=store) - - first_call = asyncio.create_task( - tracker.create_if_no_running( - "add_resource", - "viking://resources/one", - account_id="acme", - user_id="alice", - ) - ) - second_call = asyncio.create_task( - tracker.create_if_no_running( - "add_resource", - "viking://resources/two", - account_id="acme", - user_id="alice", - ) - ) - await asyncio.wait_for(store.expected_entered.wait(), timeout=1) - assert store.max_create_inflight == 2 - store.create_release.set() - first, second = await asyncio.gather(first_call, second_call) - - assert first is not None - assert second is not None - assert first.task_id != second.task_id - - -@pytest.mark.asyncio -async def test_task_and_business_lock_registries_do_not_grow_with_completed_calls(): - from openviking.service.task_tracker import TaskTracker - - store = _ControllableTaskStore() - tracker = TaskTracker(store=store) - - tasks = await asyncio.gather( - *[tracker.create("add_resource", account_id="acme", user_id="alice") for _ in range(100)] - ) - await asyncio.gather(*(tracker.start(task.task_id) for task in tasks)) - await asyncio.gather( - *[ - tracker.create_if_no_running( - "reindex", - f"viking://resources/{index}", - account_id="acme", - user_id="alice", - ) - for index in range(20) - ] - ) - - assert tracker._task_locks.entry_count == 0 - assert tracker._business_locks.entry_count == 0 + assert sum(task is not None for task in results) == 1 diff --git a/tests/unit/session/test_session_commit_resume.py b/tests/unit/session/test_session_commit_resume.py index 90adac7c7a..500ed3cd05 100644 --- a/tests/unit/session/test_session_commit_resume.py +++ b/tests/unit/session/test_session_commit_resume.py @@ -11,10 +11,12 @@ import pytest from openviking.message import Message, TextPart, ToolPart +from openviking.service.task_store import PersistentTaskStore from openviking.service.task_tracker import TaskStatus, TaskTracker, set_task_tracker from openviking.session.memory_policy import MemoryPolicy from openviking.session.session import Session from openviking.storage.queuefs.session_commit_msg import SessionCommitMsg +from tests.test_task_tracker import _FakeAgfs @pytest.mark.parametrize("enabled", [False, True]) @@ -49,24 +51,9 @@ def test_legacy_queued_work_keeps_old_default_without_migrating_user_policy(): assert message.memory_policy["memory_types"] == policy["memory_types"] -class _TaskStore: +class _TaskStore(PersistentTaskStore): def __init__(self): - self.tasks = {} - - async def create(self, task): - self.tasks[task.task_id] = task - - async def update(self, task): - self.tasks[task.task_id] = task - - async def get(self, task_id, *, account_id=None, user_id=None): - return None - - async def list(self, account_id, *, user_id=None): - return [] - - async def delete(self, task_id, *, account_id, user_id=None): - self.tasks.pop(task_id, None) + super().__init__(_FakeAgfs()) class _MemoryVikingFS: