mirror of
https://github.com/HKUDS/OpenSpace.git
synced 2026-08-28 05:15:00 +00:00
970 lines
31 KiB
Python
970 lines
31 KiB
Python
"""Query-level background housekeeping.
|
|
|
|
Background tasks are routed through TaskManager and the TUI
|
|
``agent_task_update`` surface so startup, drain, and cancellation behavior stays
|
|
observable.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import inspect
|
|
import logging
|
|
import os
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Awaitable, Callable, Mapping
|
|
|
|
from openspace.agents.task_manager import (
|
|
StopTaskError,
|
|
TaskManager,
|
|
TaskType,
|
|
generate_task_id,
|
|
)
|
|
from openspace.services.memory.background import (
|
|
MemoryDrainResult,
|
|
drain_pending_memory_tasks,
|
|
)
|
|
from openspace.services.memory.task_scope import maybe_memory_task_scope_key
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
BackgroundEventSink = Callable[[str, dict[str, Any]], Awaitable[None] | None]
|
|
|
|
DEFAULT_IDLE_GRACE_S = 60.0
|
|
DEFAULT_SLOW_OP_DELAY_S = 10 * 60.0
|
|
DEFAULT_RECURRING_CLEANUP_INTERVAL_S = 24 * 60 * 60.0
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class BackgroundTaskRecord:
|
|
task_id: str
|
|
name: str
|
|
source: str
|
|
task_type: str
|
|
description: str
|
|
task: asyncio.Task[Any]
|
|
scope_key: str | None
|
|
session_id: str | None = None
|
|
parent_task_id: str | None = None
|
|
root_task_id: str | None = None
|
|
event_sink: BackgroundEventSink | None = None
|
|
task_manager: TaskManager | None = None
|
|
started_at: float = field(default_factory=lambda: time.time() * 1000)
|
|
status: str = "running"
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class BackgroundHousekeepingState:
|
|
started: bool = False
|
|
started_at: float = 0.0
|
|
last_interaction_at: float = field(default_factory=time.time)
|
|
idle_grace_s: float = DEFAULT_IDLE_GRACE_S
|
|
slow_op_delay_s: float = DEFAULT_SLOW_OP_DELAY_S
|
|
recurring_cleanup_interval_s: float = DEFAULT_RECURRING_CLEANUP_INTERVAL_S
|
|
slow_ops_task: asyncio.Task[Any] | None = None
|
|
recurring_task: asyncio.Task[Any] | None = None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class BackgroundDrainResult:
|
|
timeout_s: float
|
|
memory: MemoryDrainResult
|
|
tracked_pending: int = 0
|
|
|
|
@property
|
|
def pending_count(self) -> int:
|
|
return self.memory.pending_count + self.tracked_pending
|
|
|
|
@property
|
|
def timed_out(self) -> bool:
|
|
return self.pending_count > 0
|
|
|
|
def as_event_payload(self) -> dict[str, Any]:
|
|
payload = self.memory.as_event_payload()
|
|
payload["tracked_pending"] = self.tracked_pending
|
|
payload["pending_count"] = self.pending_count
|
|
return payload
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class BackgroundQueueStats:
|
|
source: str
|
|
queued_count: int = 0
|
|
running_count: int = 0
|
|
completed_count: int = 0
|
|
failed_count: int = 0
|
|
cancelled_count: int = 0
|
|
backlog_warning_count: int = 0
|
|
latest_error: str | None = None
|
|
latest_queue_lag_ms: float = 0.0
|
|
latest_backlog_count: int = 0
|
|
latest_started_at_ms: float = 0.0
|
|
latest_finished_at_ms: float = 0.0
|
|
|
|
def to_dict(self, *, backlog_warning_threshold: int | None = None) -> dict[str, Any]:
|
|
backlogged = (
|
|
backlog_warning_threshold is not None
|
|
and backlog_warning_threshold > 0
|
|
and self.queued_count >= backlog_warning_threshold
|
|
)
|
|
return {
|
|
"source": self.source,
|
|
"queued_count": self.queued_count,
|
|
"running_count": self.running_count,
|
|
"completed_count": self.completed_count,
|
|
"failed_count": self.failed_count,
|
|
"cancelled_count": self.cancelled_count,
|
|
"backlog_warning_count": self.backlog_warning_count,
|
|
"backlogged": backlogged,
|
|
"latest_error": self.latest_error,
|
|
"latest_queue_lag_ms": self.latest_queue_lag_ms,
|
|
"latest_backlog_count": self.latest_backlog_count,
|
|
"latest_started_at_ms": self.latest_started_at_ms,
|
|
"latest_finished_at_ms": self.latest_finished_at_ms,
|
|
}
|
|
|
|
|
|
class BackgroundSupervisor:
|
|
"""Small explicit supervisor around existing fire-and-forget tasks."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
max_concurrency_per_source: int = 1,
|
|
backlog_warning_threshold: int | None = 50,
|
|
event_sink: BackgroundEventSink | None = None,
|
|
) -> None:
|
|
self.max_concurrency_per_source = max(1, int(max_concurrency_per_source))
|
|
self.backlog_warning_threshold = (
|
|
None
|
|
if backlog_warning_threshold is None
|
|
else max(0, int(backlog_warning_threshold))
|
|
)
|
|
self._event_sink = event_sink
|
|
self._semaphores: dict[str, asyncio.Semaphore] = {}
|
|
self._stats: dict[str, BackgroundQueueStats] = {}
|
|
self._tasks: set[asyncio.Task[Any]] = set()
|
|
|
|
def submit(
|
|
self,
|
|
*,
|
|
source: str,
|
|
name: str,
|
|
coro_factory: Callable[[], Awaitable[Any] | Any],
|
|
context: Any | None = None,
|
|
description: str = "",
|
|
task_type: TaskType | str = "background",
|
|
timeout_s: float | None = None,
|
|
task_id: str | None = None,
|
|
) -> asyncio.Task[Any]:
|
|
source_name = str(source or "background")
|
|
resolved_task_id = task_id or generate_task_id(task_type)
|
|
queued_at_ms = time.time() * 1000
|
|
stats = self._stats_for(source_name)
|
|
stats.queued_count += 1
|
|
sink = self._event_sink or _event_sink_from_context(context)
|
|
self._maybe_emit_backlog_warning(
|
|
sink=sink,
|
|
stats=stats,
|
|
source=source_name,
|
|
name=name,
|
|
task_id=resolved_task_id,
|
|
description=description,
|
|
)
|
|
|
|
async def _runner() -> Any:
|
|
started = False
|
|
try:
|
|
semaphore = self._semaphore_for(source_name)
|
|
async with semaphore:
|
|
started = True
|
|
stats.queued_count = max(0, stats.queued_count - 1)
|
|
stats.running_count += 1
|
|
stats.latest_started_at_ms = time.time() * 1000
|
|
stats.latest_queue_lag_ms = max(
|
|
0.0,
|
|
stats.latest_started_at_ms - queued_at_ms,
|
|
)
|
|
await _emit_event(
|
|
sink,
|
|
"background_lag",
|
|
{
|
|
"task_id": resolved_task_id,
|
|
"source": source_name,
|
|
"name": name,
|
|
"description": description,
|
|
"queue_lag_ms": stats.latest_queue_lag_ms,
|
|
"queued_count": stats.queued_count,
|
|
"running_count": stats.running_count,
|
|
},
|
|
)
|
|
await _emit_event(
|
|
sink,
|
|
"background.queue_lag",
|
|
{
|
|
"task_id": resolved_task_id,
|
|
"source": source_name,
|
|
"name": name,
|
|
"description": description,
|
|
"queue_lag_ms": stats.latest_queue_lag_ms,
|
|
"queued_count": stats.queued_count,
|
|
"running_count": stats.running_count,
|
|
},
|
|
)
|
|
try:
|
|
result = coro_factory()
|
|
if inspect.isawaitable(result):
|
|
awaitable = result
|
|
else:
|
|
async def _constant() -> Any:
|
|
return result
|
|
|
|
awaitable = _constant()
|
|
if timeout_s is not None:
|
|
result = await asyncio.wait_for(
|
|
awaitable,
|
|
timeout=max(0.0, float(timeout_s)),
|
|
)
|
|
else:
|
|
result = await awaitable
|
|
except asyncio.CancelledError:
|
|
stats.cancelled_count += 1
|
|
stats.latest_finished_at_ms = time.time() * 1000
|
|
await _emit_event(
|
|
sink,
|
|
"background_task_cancelled",
|
|
{
|
|
"task_id": resolved_task_id,
|
|
"source": source_name,
|
|
"name": name,
|
|
},
|
|
)
|
|
raise
|
|
except Exception as exc:
|
|
stats.failed_count += 1
|
|
stats.latest_error = str(exc)
|
|
stats.latest_finished_at_ms = time.time() * 1000
|
|
await _emit_event(
|
|
sink,
|
|
"background_task_failed",
|
|
{
|
|
"task_id": resolved_task_id,
|
|
"source": source_name,
|
|
"name": name,
|
|
"error": str(exc),
|
|
},
|
|
)
|
|
return {"error": str(exc)}
|
|
else:
|
|
stats.completed_count += 1
|
|
stats.latest_error = None
|
|
stats.latest_finished_at_ms = time.time() * 1000
|
|
await _emit_event(
|
|
sink,
|
|
"background_task_completed",
|
|
{
|
|
"task_id": resolved_task_id,
|
|
"source": source_name,
|
|
"name": name,
|
|
},
|
|
)
|
|
return result
|
|
finally:
|
|
stats.running_count = max(0, stats.running_count - 1)
|
|
except asyncio.CancelledError:
|
|
if not started:
|
|
stats.queued_count = max(0, stats.queued_count - 1)
|
|
stats.cancelled_count += 1
|
|
stats.latest_finished_at_ms = time.time() * 1000
|
|
await _emit_event(
|
|
sink,
|
|
"background_task_cancelled",
|
|
{
|
|
"task_id": resolved_task_id,
|
|
"source": source_name,
|
|
"name": name,
|
|
},
|
|
)
|
|
raise
|
|
finally:
|
|
current = asyncio.current_task()
|
|
if current is not None:
|
|
self._tasks.discard(current)
|
|
|
|
task = asyncio.create_task(
|
|
_runner(),
|
|
name=f"openspace-background-{source_name}-{resolved_task_id}",
|
|
)
|
|
self._tasks.add(task)
|
|
return task
|
|
|
|
async def drain(self, *, timeout_s: float = 3.0) -> int:
|
|
pending = [task for task in self._tasks if not task.done()]
|
|
if not pending:
|
|
return 0
|
|
done, still_pending = await asyncio.wait(
|
|
pending,
|
|
timeout=max(0.0, float(timeout_s)),
|
|
)
|
|
for task in done:
|
|
try:
|
|
task.result()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception:
|
|
logger.debug("Background supervisor task failed during drain", exc_info=True)
|
|
return len(still_pending)
|
|
|
|
async def cancel_all(self, *, reason: str = "cancelled") -> int:
|
|
count = 0
|
|
for task in list(self._tasks):
|
|
if not task.done():
|
|
task.cancel()
|
|
count += 1
|
|
if count:
|
|
await _emit_event(
|
|
self._event_sink,
|
|
"background_supervisor_cancel_all",
|
|
{"reason": reason, "count": count},
|
|
)
|
|
return count
|
|
|
|
def status(self) -> dict[str, Any]:
|
|
threshold = self.backlog_warning_threshold
|
|
return {
|
|
"tasks": sum(1 for task in self._tasks if not task.done()),
|
|
"backlog_warning_threshold": threshold,
|
|
"backlogged_sources": [
|
|
source
|
|
for source, stats in sorted(self._stats.items())
|
|
if threshold is not None
|
|
and threshold > 0
|
|
and stats.queued_count >= threshold
|
|
],
|
|
"queues": {
|
|
source: stats.to_dict(backlog_warning_threshold=threshold)
|
|
for source, stats in sorted(self._stats.items())
|
|
},
|
|
}
|
|
|
|
def source_backlogged(self, source: str) -> bool:
|
|
threshold = self.backlog_warning_threshold
|
|
if threshold is None or threshold <= 0:
|
|
return False
|
|
stats = self._stats.get(str(source or "background"))
|
|
return bool(stats is not None and stats.queued_count >= threshold)
|
|
|
|
def _maybe_emit_backlog_warning(
|
|
self,
|
|
*,
|
|
sink: BackgroundEventSink | None,
|
|
stats: BackgroundQueueStats,
|
|
source: str,
|
|
name: str,
|
|
task_id: str,
|
|
description: str,
|
|
) -> None:
|
|
threshold = self.backlog_warning_threshold
|
|
if threshold is None or threshold <= 0 or stats.queued_count < threshold:
|
|
return
|
|
stats.backlog_warning_count += 1
|
|
stats.latest_backlog_count = stats.queued_count
|
|
asyncio.create_task(
|
|
_emit_event(
|
|
sink,
|
|
"background_backlog_high",
|
|
{
|
|
"task_id": task_id,
|
|
"source": source,
|
|
"name": name,
|
|
"description": description,
|
|
"queued_count": stats.queued_count,
|
|
"running_count": stats.running_count,
|
|
"threshold": threshold,
|
|
},
|
|
),
|
|
name=f"openspace-background-backlog-{source}-{task_id}",
|
|
)
|
|
|
|
def _stats_for(self, source: str) -> BackgroundQueueStats:
|
|
stats = self._stats.get(source)
|
|
if stats is None:
|
|
stats = BackgroundQueueStats(source=source)
|
|
self._stats[source] = stats
|
|
return stats
|
|
|
|
def _semaphore_for(self, source: str) -> asyncio.Semaphore:
|
|
semaphore = self._semaphores.get(source)
|
|
if semaphore is None:
|
|
semaphore = asyncio.Semaphore(self.max_concurrency_per_source)
|
|
self._semaphores[source] = semaphore
|
|
return semaphore
|
|
|
|
|
|
_state = BackgroundHousekeepingState()
|
|
_tasks: dict[str, BackgroundTaskRecord] = {}
|
|
_supervisor = BackgroundSupervisor()
|
|
|
|
|
|
def get_background_supervisor() -> BackgroundSupervisor:
|
|
return _supervisor
|
|
|
|
|
|
def start_background_housekeeping(
|
|
context: Any | None = None,
|
|
*,
|
|
event_sink: BackgroundEventSink | None = None,
|
|
idle_grace_s: float = DEFAULT_IDLE_GRACE_S,
|
|
slow_op_delay_s: float = DEFAULT_SLOW_OP_DELAY_S,
|
|
recurring_cleanup_interval_s: float = DEFAULT_RECURRING_CLEANUP_INTERVAL_S,
|
|
) -> BackgroundHousekeepingState:
|
|
"""Initialize background runners and idle cleanup timers.
|
|
|
|
OpenSpace returns ``void`` and exposes no cleanup handle. OS returns the state for
|
|
tests while keeping the module-level singleton behavior.
|
|
"""
|
|
|
|
if _state.started:
|
|
record_interaction()
|
|
return _state
|
|
|
|
from openspace.services.memory.extract import init_extract_memories
|
|
from openspace.services.memory.dream import init_auto_dream
|
|
from openspace.services.memory.session_memory import init_session_memory
|
|
|
|
init_session_memory()
|
|
init_extract_memories()
|
|
init_auto_dream()
|
|
|
|
_state.started = True
|
|
_state.started_at = time.time()
|
|
_state.last_interaction_at = time.time()
|
|
_state.idle_grace_s = max(0.0, float(idle_grace_s))
|
|
_state.slow_op_delay_s = max(0.0, float(slow_op_delay_s))
|
|
_state.recurring_cleanup_interval_s = max(0.0, float(recurring_cleanup_interval_s))
|
|
|
|
sink = event_sink or _event_sink_from_context(context)
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
loop = None
|
|
if loop is not None:
|
|
_state.slow_ops_task = loop.create_task(_run_delayed_idle_cleanup(sink))
|
|
if os.environ.get("USER_TYPE") == "ant":
|
|
_state.recurring_task = loop.create_task(_run_recurring_cleanup(sink))
|
|
return _state
|
|
|
|
|
|
async def run_startup_evolution_recovery(
|
|
recovery: Any,
|
|
*,
|
|
event_sink: BackgroundEventSink | None = None,
|
|
) -> Any | None:
|
|
"""Run evolution startup recovery without letting it break runtime init."""
|
|
|
|
try:
|
|
run = getattr(recovery, "run_startup_recovery", None)
|
|
if not callable(run):
|
|
return None
|
|
result = run()
|
|
payload = (
|
|
result.to_dict()
|
|
if hasattr(result, "to_dict") and callable(result.to_dict)
|
|
else {"result": result}
|
|
)
|
|
await _emit_event(event_sink, "evolution_recovery_complete", payload)
|
|
return result
|
|
except Exception as exc:
|
|
logger.debug("Evolution startup recovery failed", exc_info=True)
|
|
await _emit_event(
|
|
event_sink,
|
|
"evolution_recovery_failed",
|
|
{"error": str(exc)},
|
|
)
|
|
return None
|
|
|
|
|
|
def record_interaction() -> None:
|
|
"""Mark recent user activity for OpenSpace idle cleanup deferral."""
|
|
|
|
_state.last_interaction_at = time.time()
|
|
|
|
|
|
async def stop_background_housekeeping(
|
|
context: Any | None = None,
|
|
*,
|
|
timeout_s: float = 10.0,
|
|
event_sink: BackgroundEventSink | None = None,
|
|
cancel_pending: bool = True,
|
|
) -> BackgroundDrainResult:
|
|
"""Drain registered housekeeping work and stop timers."""
|
|
|
|
_cancel_timer(_state.slow_ops_task)
|
|
_cancel_timer(_state.recurring_task)
|
|
_state.slow_ops_task = None
|
|
_state.recurring_task = None
|
|
|
|
sink = event_sink or _event_sink_from_context(context)
|
|
result = await drain_background_tasks(
|
|
context,
|
|
timeout_s=timeout_s,
|
|
event_sink=sink,
|
|
)
|
|
if cancel_pending and result.pending_count:
|
|
await cancel_all_background_tasks(context=context, reason="shutdown")
|
|
_state.started = False
|
|
return result
|
|
|
|
|
|
async def schedule_session_memory(
|
|
context: Any,
|
|
append_system_message: Any | None = None,
|
|
) -> asyncio.Task[Any] | None:
|
|
from openspace.services.memory.session_memory import (
|
|
should_schedule_session_memory,
|
|
submit_session_memory_extraction,
|
|
)
|
|
|
|
if not should_schedule_session_memory(context):
|
|
return None
|
|
task_id = _reserve_task_id(context, "session_memory", TaskType.SESSION_MEMORY)
|
|
task = submit_session_memory_extraction(context, append_system_message)
|
|
await track_background_task(
|
|
context,
|
|
task,
|
|
name="Session Memory",
|
|
source="session_memory",
|
|
task_type=TaskType.SESSION_MEMORY,
|
|
description="Background session memory extraction",
|
|
task_id=task_id,
|
|
)
|
|
return task
|
|
|
|
|
|
async def schedule_extract_memories(
|
|
context: Any,
|
|
append_system_message: Any | None = None,
|
|
) -> asyncio.Task[Any] | None:
|
|
from openspace.services.memory.extract import (
|
|
should_schedule_extract_memories,
|
|
submit_extract_memories,
|
|
)
|
|
|
|
if not should_schedule_extract_memories(context):
|
|
return None
|
|
task_id = _reserve_task_id(context, "extract_memories", TaskType.MEMORY_EXTRACT)
|
|
task = submit_extract_memories(context, append_system_message)
|
|
await track_background_task(
|
|
context,
|
|
task,
|
|
name="Memory Extract",
|
|
source="extract_memories",
|
|
task_type=TaskType.MEMORY_EXTRACT,
|
|
description="Background memory extraction",
|
|
task_id=task_id,
|
|
)
|
|
return task
|
|
|
|
|
|
async def schedule_auto_dream(
|
|
context: Any,
|
|
append_system_message: Any | None = None,
|
|
) -> asyncio.Task[Any] | None:
|
|
from openspace.services.memory.dream import (
|
|
should_schedule_auto_dream,
|
|
submit_auto_dream,
|
|
)
|
|
|
|
if not should_schedule_auto_dream(context):
|
|
return None
|
|
task_id = _reserve_task_id(context, "auto_dream", TaskType.DREAM)
|
|
task = submit_auto_dream(context, append_system_message)
|
|
await track_background_task(
|
|
context,
|
|
task,
|
|
name="Memory Dream",
|
|
source="auto_dream",
|
|
task_type=TaskType.DREAM,
|
|
description="Memory consolidation",
|
|
task_id=task_id,
|
|
)
|
|
return task
|
|
|
|
|
|
async def track_background_task(
|
|
context: Any,
|
|
task: asyncio.Task[Any] | None,
|
|
*,
|
|
name: str,
|
|
source: str,
|
|
task_type: TaskType | str,
|
|
description: str,
|
|
task_id: str | None = None,
|
|
) -> BackgroundTaskRecord | None:
|
|
"""Track a submitted task in context, global registry, and TaskManager."""
|
|
|
|
if task is None or not isinstance(task, asyncio.Task):
|
|
return None
|
|
|
|
resolved_task_id = task_id or _reserve_task_id(context, source, task_type)
|
|
scope_key = maybe_memory_task_scope_key(context)
|
|
task_manager = _task_manager_from_context(context)
|
|
event_sink = _event_sink_from_context(context)
|
|
|
|
record = BackgroundTaskRecord(
|
|
task_id=resolved_task_id,
|
|
name=name,
|
|
source=source,
|
|
task_type=str(getattr(task_type, "value", task_type)),
|
|
description=description,
|
|
task=task,
|
|
scope_key=scope_key,
|
|
session_id=_none_or_str(_context_value(context, "session_id")),
|
|
parent_task_id=(
|
|
_none_or_str(_context_value(context, "task_id"))
|
|
or _none_or_str(_context_value(context, "parent_task_id"))
|
|
),
|
|
root_task_id=_none_or_str(_context_value(context, "task_id")),
|
|
event_sink=event_sink,
|
|
task_manager=task_manager,
|
|
)
|
|
_tasks[resolved_task_id] = record
|
|
_add_context_task(context, task)
|
|
|
|
if task_manager is not None:
|
|
await task_manager.register_external_task(
|
|
runner_task=task,
|
|
description=description,
|
|
task_type=task_type,
|
|
agent_type="memory",
|
|
agent_id=resolved_task_id,
|
|
parent_task_id=(
|
|
_context_value(context, "task_id")
|
|
or _context_value(context, "parent_task_id")
|
|
),
|
|
is_backgrounded=True,
|
|
)
|
|
else:
|
|
await _emit_task_update(record, "task_started")
|
|
|
|
def _done(done: asyncio.Task[Any]) -> None:
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
return
|
|
loop.create_task(_finalize_unmanaged_record(record, done))
|
|
|
|
task.add_done_callback(_done)
|
|
|
|
def _cleanup(done: asyncio.Task[Any]) -> None:
|
|
_discard_context_task(context, done)
|
|
if done.cancelled():
|
|
record.status = "killed"
|
|
elif done.exception() is not None:
|
|
record.status = "failed"
|
|
else:
|
|
record.status = "completed"
|
|
|
|
task.add_done_callback(_cleanup)
|
|
return record
|
|
|
|
|
|
async def drain_background_tasks(
|
|
context: Any | None = None,
|
|
*,
|
|
timeout_s: float = 3.0,
|
|
event_sink: BackgroundEventSink | None = None,
|
|
) -> BackgroundDrainResult:
|
|
"""Drain low-level memory tasks plus registered housekeeping records."""
|
|
|
|
timeout_s = max(0.0, float(timeout_s))
|
|
deadline = time.monotonic() + timeout_s
|
|
sink = event_sink or _event_sink_from_context(context)
|
|
|
|
def remaining() -> float:
|
|
return max(0.0, deadline - time.monotonic())
|
|
|
|
memory_result = await drain_pending_memory_tasks(
|
|
timeout_s=remaining(),
|
|
event_sink=sink,
|
|
context=context,
|
|
)
|
|
tracked_pending = await _drain_tracked_records(
|
|
context,
|
|
timeout_s=remaining(),
|
|
)
|
|
result = BackgroundDrainResult(
|
|
timeout_s=timeout_s,
|
|
memory=memory_result,
|
|
tracked_pending=tracked_pending,
|
|
)
|
|
if result.timed_out:
|
|
await _emit_event(
|
|
sink,
|
|
"background_housekeeping_drain_timeout",
|
|
result.as_event_payload(),
|
|
)
|
|
return result
|
|
|
|
|
|
async def cancel_background_task(task_id: str, *, reason: str = "cancelled") -> bool:
|
|
record = _tasks.get(task_id)
|
|
if record is None:
|
|
return False
|
|
if record.task.done():
|
|
return False
|
|
record.status = "killed"
|
|
if record.task_manager is not None:
|
|
try:
|
|
await record.task_manager.stop_task_or_raise(task_id)
|
|
return True
|
|
except StopTaskError:
|
|
pass
|
|
record.task.cancel()
|
|
await _emit_task_update(record, "task_stopped", {"reason": reason})
|
|
return True
|
|
|
|
|
|
async def cancel_all_background_tasks(
|
|
*,
|
|
context: Any | None = None,
|
|
reason: str = "cancelled",
|
|
) -> int:
|
|
scope_key = maybe_memory_task_scope_key(context)
|
|
count = 0
|
|
for record in list(_tasks.values()):
|
|
if record.task.done():
|
|
continue
|
|
if scope_key is not None and record.scope_key != scope_key:
|
|
continue
|
|
if await cancel_background_task(record.task_id, reason=reason):
|
|
count += 1
|
|
return count
|
|
|
|
|
|
async def _run_delayed_idle_cleanup(
|
|
event_sink: BackgroundEventSink | None,
|
|
) -> None:
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(_state.slow_op_delay_s)
|
|
if time.time() - _state.last_interaction_at < _state.idle_grace_s:
|
|
continue
|
|
await _emit_event(
|
|
event_sink,
|
|
"background_housekeeping_idle",
|
|
{"idle_for_s": time.time() - _state.last_interaction_at},
|
|
)
|
|
await _emit_event(
|
|
event_sink,
|
|
"background_housekeeping_cleanup_complete",
|
|
{
|
|
"cleanup": "os_retention_policy",
|
|
"skipped": True,
|
|
"reason": "no_openspace_retention_policy_configured",
|
|
},
|
|
)
|
|
return
|
|
except asyncio.CancelledError:
|
|
return
|
|
|
|
|
|
async def _run_recurring_cleanup(event_sink: BackgroundEventSink | None) -> None:
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(_state.recurring_cleanup_interval_s)
|
|
await _emit_event(
|
|
event_sink,
|
|
"background_housekeeping_recurring_cleanup",
|
|
{"skipped": True, "reason": "anthropic_package_cleanup_not_applicable"},
|
|
)
|
|
except asyncio.CancelledError:
|
|
return
|
|
|
|
|
|
async def _drain_tracked_records(context: Any | None, *, timeout_s: float) -> int:
|
|
scope_key = maybe_memory_task_scope_key(context)
|
|
pending = [
|
|
record.task
|
|
for record in _tasks.values()
|
|
if not record.task.done()
|
|
and (scope_key is None or record.scope_key == scope_key)
|
|
]
|
|
if not pending:
|
|
return 0
|
|
done, still_pending = await asyncio.wait(pending, timeout=max(0.0, timeout_s))
|
|
for task in done:
|
|
try:
|
|
task.result()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception:
|
|
logger.debug("Background housekeeping task failed during drain", exc_info=True)
|
|
return len(still_pending)
|
|
|
|
|
|
async def _finalize_unmanaged_record(
|
|
record: BackgroundTaskRecord,
|
|
done: asyncio.Task[Any],
|
|
) -> None:
|
|
if done.cancelled():
|
|
record.status = "killed"
|
|
await _emit_task_update(record, "task_stopped")
|
|
return
|
|
try:
|
|
result = done.result()
|
|
error = _mapping_get(result, "error")
|
|
if error:
|
|
record.status = "failed"
|
|
await _emit_task_update(record, "task_failed", {"error": str(error)})
|
|
else:
|
|
record.status = "completed"
|
|
await _emit_task_update(record, "task_completed")
|
|
except Exception as exc:
|
|
record.status = "failed"
|
|
await _emit_task_update(record, "task_failed", {"error": str(exc)})
|
|
|
|
|
|
async def _emit_task_update(
|
|
record: BackgroundTaskRecord,
|
|
event_type: str,
|
|
extra: Mapping[str, Any] | None = None,
|
|
) -> None:
|
|
status = {
|
|
"task_started": "running",
|
|
"task_completed": "completed",
|
|
"task_failed": "failed",
|
|
"task_stopped": "killed",
|
|
}.get(event_type, record.status)
|
|
payload = {
|
|
"task_id": record.task_id,
|
|
"session_id": record.session_id,
|
|
"parent_task_id": record.parent_task_id,
|
|
"source_task_id": record.root_task_id,
|
|
"agent_id": "memory",
|
|
"agent_type": "memory",
|
|
"task_type": record.task_type,
|
|
"description": record.description,
|
|
"current_operation": record.description,
|
|
"status": status,
|
|
"is_backgrounded": True,
|
|
"start_time": record.started_at,
|
|
}
|
|
if status != "running":
|
|
payload["end_time"] = time.time() * 1000
|
|
if extra:
|
|
payload.update(dict(extra))
|
|
await _emit_event(record.event_sink, event_type, payload)
|
|
await _emit_event(record.event_sink, "agent_task_update", payload)
|
|
|
|
|
|
def _reserve_task_id(
|
|
context: Any,
|
|
source: str,
|
|
task_type: TaskType | str,
|
|
) -> str:
|
|
existing = _background_task_ids(context).get(source)
|
|
if existing:
|
|
return existing
|
|
task_id = generate_task_id(task_type)
|
|
_background_task_ids(context)[source] = task_id
|
|
return task_id
|
|
|
|
|
|
def _background_task_ids(context: Any) -> dict[str, str]:
|
|
if isinstance(context, Mapping):
|
|
value = context.get("background_task_ids")
|
|
if isinstance(value, dict):
|
|
return value
|
|
value = {}
|
|
try:
|
|
context["background_task_ids"] = value # type: ignore[index]
|
|
except Exception:
|
|
pass
|
|
return value
|
|
value = getattr(context, "background_task_ids", None)
|
|
if isinstance(value, dict):
|
|
return value
|
|
return {}
|
|
|
|
|
|
def _add_context_task(context: Any, task: asyncio.Task[Any]) -> None:
|
|
tasks = _context_value(context, "background_hook_tasks")
|
|
if isinstance(tasks, set):
|
|
tasks.add(task)
|
|
|
|
|
|
def _discard_context_task(context: Any, task: asyncio.Task[Any]) -> None:
|
|
tasks = _context_value(context, "background_hook_tasks")
|
|
if isinstance(tasks, set):
|
|
tasks.discard(task)
|
|
|
|
|
|
def _task_manager_from_context(context: Any) -> TaskManager | None:
|
|
manager = _context_value(context, "task_manager")
|
|
return manager if isinstance(manager, TaskManager) else None
|
|
|
|
|
|
def _event_sink_from_context(context: Any | None) -> BackgroundEventSink | None:
|
|
if context is None:
|
|
return None
|
|
sink = _context_value(context, "event_sink")
|
|
return sink if callable(sink) else None
|
|
|
|
|
|
def _context_value(context: Any, key: str) -> Any:
|
|
if isinstance(context, Mapping):
|
|
return context.get(key)
|
|
return getattr(context, key, None)
|
|
|
|
|
|
def _none_or_str(value: Any) -> str | None:
|
|
if value is None:
|
|
return None
|
|
text = str(value)
|
|
return text if text else None
|
|
|
|
|
|
def _mapping_get(value: Any, key: str, default: Any = None) -> Any:
|
|
if isinstance(value, Mapping):
|
|
return value.get(key, default)
|
|
return getattr(value, key, default)
|
|
|
|
|
|
async def _emit_event(
|
|
event_sink: BackgroundEventSink | None,
|
|
event_type: str,
|
|
data: dict[str, Any],
|
|
) -> None:
|
|
if event_sink is None:
|
|
return
|
|
try:
|
|
result = event_sink(event_type, data)
|
|
if inspect.isawaitable(result):
|
|
await result
|
|
except Exception:
|
|
logger.debug("Background housekeeping event sink failed", exc_info=True)
|
|
|
|
|
|
def _cancel_timer(task: asyncio.Task[Any] | None) -> None:
|
|
if task is not None and not task.done():
|
|
task.cancel()
|
|
|
|
|
|
__all__ = [
|
|
"BackgroundDrainResult",
|
|
"BackgroundHousekeepingState",
|
|
"BackgroundTaskRecord",
|
|
"cancel_all_background_tasks",
|
|
"cancel_background_task",
|
|
"drain_background_tasks",
|
|
"record_interaction",
|
|
"run_startup_evolution_recovery",
|
|
"schedule_auto_dream",
|
|
"schedule_extract_memories",
|
|
"schedule_session_memory",
|
|
"start_background_housekeeping",
|
|
"stop_background_housekeeping",
|
|
"track_background_task",
|
|
]
|