From 08c8725ef9cf15fd06a0b4741a08831f60f72744 Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 12:17:15 +0800 Subject: [PATCH 1/7] feat: extract agent bus scheduling kernel --- README.md | 28 +- agent_core/components/agent_bus/__init__.py | 89 +- agent_core/components/agent_bus/agent_comm.py | 228 ++ agent_core/components/agent_bus/bus.py | 2176 +++++++++++++++++ agent_core/components/agent_bus/fan_in.py | 431 ++++ agent_core/components/agent_bus/models.py | 246 ++ agent_core/components/agent_bus/runtime.py | 501 ++++ .../components/agent_bus/spawn_guard.py | 341 +++ .../components/finalization/__init__.py | 54 +- agent_core/components/finalization/budget.py | 263 ++ agent_core/models/agent_message.py | 3 +- agent_core/models/event.py | 230 ++ agent_core/models/task_budget.py | 133 + agent_core/protocols.py | 48 +- agent_core/runtime/resources/__init__.py | 22 + agent_core/runtime/resources/llm.py | 54 + agent_core/runtime/resources/manager.py | 188 ++ agent_core/runtime/resources/permissions.py | 45 + agent_core/scheduling/__init__.py | 26 +- agent_core/scheduling/scheduler.py | 503 ++++ agent_core/scheduling/workflow_defaults.py | 38 + agent_core/tool.py | 282 +++ docs/agent-bus-boundary.md | 26 + tests/test_agent_bus_event_sink_injection.py | 122 + tests/test_agent_bus_session_invariant.py | 418 ++++ tests/test_agent_bus_session_queue.py | 414 ++++ tests/test_agent_bus_waiting.py | 265 ++ tests/test_finalization_budget.py | 246 ++ tests/test_resource_manager.py | 240 ++ tests/test_scheduler.py | 140 ++ tests/test_spawn_guard.py | 437 ++++ 31 files changed, 8223 insertions(+), 14 deletions(-) create mode 100644 agent_core/components/agent_bus/agent_comm.py create mode 100644 agent_core/components/agent_bus/bus.py create mode 100644 agent_core/components/agent_bus/fan_in.py create mode 100644 agent_core/components/agent_bus/models.py create mode 100644 agent_core/components/agent_bus/runtime.py create mode 100644 agent_core/components/agent_bus/spawn_guard.py create mode 100644 agent_core/components/finalization/budget.py create mode 100644 agent_core/models/event.py create mode 100644 agent_core/models/task_budget.py create mode 100644 agent_core/runtime/resources/__init__.py create mode 100644 agent_core/runtime/resources/llm.py create mode 100644 agent_core/runtime/resources/manager.py create mode 100644 agent_core/runtime/resources/permissions.py create mode 100644 agent_core/scheduling/scheduler.py create mode 100644 agent_core/scheduling/workflow_defaults.py create mode 100644 agent_core/tool.py create mode 100644 docs/agent-bus-boundary.md create mode 100644 tests/test_agent_bus_event_sink_injection.py create mode 100644 tests/test_agent_bus_session_invariant.py create mode 100644 tests/test_agent_bus_session_queue.py create mode 100644 tests/test_agent_bus_waiting.py create mode 100644 tests/test_finalization_budget.py create mode 100644 tests/test_resource_manager.py create mode 100644 tests/test_scheduler.py create mode 100644 tests/test_spawn_guard.py diff --git a/README.md b/README.md index 7f561df..de11eda 100644 --- a/README.md +++ b/README.md @@ -41,8 +41,12 @@ Version `0.1.x` contains the converged foundation layer: stream-repetition detection; - portable loop observers for budget/context/wall-clock guards, finalization, repetition detection, SSE events, task boards, and bounded trajectories; -- shared agent-bus stop signals/pools, filesystem skill loaders, session - history, renewable wall-time leases, and language helpers. +- complete agent-bus communication, spawn control, reusable sessions, + result recovery, fan-in policy, shared pools and stop signals; +- portable tools, task/finalization budgets, role-scoped resource management, + and backend-injected DAG scheduling; +- filesystem skill loaders, session history, renewable wall-time leases, and + language helpers. The initial extraction is based on the already-merged integration branches: @@ -68,9 +72,10 @@ Code belongs in AgentCore when it: - accepts product behavior through typed inputs or explicit hooks; - has tests that run without either product repository installed. -Provider clients, session affinity, user/session retention policy, checkpoint -association, workflow node implementations, sandbox mounting/authorization, -UI history, and product-specific observers remain in their product. +Provider clients, session affinity, durable process/event-store implementations, +user/session retention policy, checkpoint association, workflow node +implementations, sandbox mounting/authorization, UI history, and +product-specific observers remain in their product. ## Development @@ -147,10 +152,15 @@ edit both products' core copies, that is evidence it belongs here. MiniDAG/graph construction, middleware/observers, agent-bus coordination primitives (shared pools and stop signals), skills, session history, models, and scoped registries. -8. **Agent-bus scheduling** (next): extract the shared communication models, - runtime, spawn guard, bus, and fan-in orchestration behind explicit product - hooks. Product tool assembly and workflow-specific policy remain downstream. -9. Remove product compatibility facades after downstream imports have moved to +8. **Agent-bus scheduling** (complete in AgentCore): communication models, + durable result-message recovery, runtime, spawn guard, reusable sessions, + bus/fan-in orchestration, resource management, task budgets, and the + backend-injected DAG scheduler. Product tool assembly, persistence backends, + and workflow policy remain downstream. +9. **Provider substrate** (next): converge the shared provider transports, + fallback engine, prompt cache, usage metering, and non-blocking streaming + behind product configuration and session-affinity adapters. +10. Remove product compatibility facades after downstream imports have moved to `agent_core`. Each slice must leave product CI green and must not depend on a floating branch. diff --git a/agent_core/components/agent_bus/__init__.py b/agent_core/components/agent_bus/__init__.py index 889cbef..5e17779 100644 --- a/agent_core/components/agent_bus/__init__.py +++ b/agent_core/components/agent_bus/__init__.py @@ -1 +1,88 @@ -"""Portable agent coordination primitives.""" +"""AgentBus — multi-agent dispatch, communication, and spawn control. + +Also serves as the back-compat entry point for ``agent_core.components.agent_bus`` +(the former flat-file module). Re-exports the public API from submodules. +""" + +from agent_core.components.agent_bus.agent_comm import AgentComm, DeliveryMode +from agent_core.components.agent_bus.bus import AgentBus +from agent_core.components.agent_bus.fan_in import ( + COMPLETE_STOP_REASONS, + INCOMPLETE_STOP_REASONS, + ORCHESTRATOR_AGENT_NAME, + CompletionInfo, + CompletionPolicy, + FanInBatch, + classify_completion, + format_report_block, + format_status_line, + format_status_report_block, + process_collected, + session_state, +) +from agent_core.components.agent_bus.models import ( + CollectResult, + DepthLimitExceeded, + JobEntry, + PendingSessionTask, + SessionWaitOutcome, + SubAgentResult, + SubAgentRuntimeSpec, + SubAgentSession, + SubTask, +) +from agent_core.components.agent_bus.runtime import ( + adapt_default_session_result, + adapt_default_subagent_result, + build_default_subagent_loop_config, + build_default_subagent_observers, + build_session_loop_config, + close_session_boundary_aborted, + resolve_session_observers, +) +from agent_core.components.agent_bus.shared_pool import SharedArtifactPool +from agent_core.components.agent_bus.spawn_guard import ( + BudgetExhausted, + SpawnDepthExceeded, + SpawnGuard, + SpawnReservation, +) + +__all__ = [ + "COMPLETE_STOP_REASONS", + "INCOMPLETE_STOP_REASONS", + "ORCHESTRATOR_AGENT_NAME", + "AgentBus", + "AgentComm", + "BudgetExhausted", + "CollectResult", + "CompletionInfo", + "CompletionPolicy", + "DeliveryMode", + "DepthLimitExceeded", + "FanInBatch", + "JobEntry", + "PendingSessionTask", + "SessionWaitOutcome", + "SharedArtifactPool", + "SpawnDepthExceeded", + "SpawnGuard", + "SpawnReservation", + "SubAgentResult", + "SubAgentRuntimeSpec", + "SubAgentSession", + "SubTask", + "adapt_default_session_result", + "adapt_default_subagent_result", + "build_default_subagent_loop_config", + "build_default_subagent_observers", + "build_session_loop_config", + "classify_completion", + "close_session_boundary_aborted", + "format_report_block", + "format_status_line", + "format_status_report_block", + "process_collected", + "resolve_session_observers", + "session_state", +] diff --git a/agent_core/components/agent_bus/agent_comm.py b/agent_core/components/agent_bus/agent_comm.py new file mode 100644 index 0000000..44401bc --- /dev/null +++ b/agent_core/components/agent_bus/agent_comm.py @@ -0,0 +1,228 @@ +"""AgentComm — inter-agent communication service. + +Sends typed AgentMessages, persists them as KernelEvents, +and publishes to EventBus for real-time streaming. + +Phase D (Issue #24): DeliveryMode (TRIGGER/QUEUE) for message routing. +""" + +from __future__ import annotations + +import asyncio +import itertools +import logging +from enum import StrEnum +from typing import Protocol, runtime_checkable + +from agent_core.events import EventType +from agent_core.models.agent_message import AgentMessage +from agent_core.models.event import KernelEvent +from agent_core.protocols import EventReader, EventSink +from agent_core.runtime.events import EventBus +from agent_core.types import TaskId + +logger = logging.getLogger(__name__) + + +@runtime_checkable +class _AgentCommEventStore(EventSink, EventReader, Protocol): + """AgentComm needs both append (EventSink) and cursored reads + (EventReader) — combined here because Python lacks an intersection + type. Public Protocols stay minimal in ``core.protocols``. + + The combination has a stronger semantic contract than ``EventSink`` + alone: ``append`` retains the event, assigns and returns a unique + integer id, and sequential awaited appends receive increasing ids. + The id — not ``timestamp`` — is the total order. Concurrent appends + have no defined relative order, only that every reader observes the + same resulting order. + + Ordering is not guaranteed across live writer processes sharing one + JSONL path. Nor is losslessness guaranteed for a capacity-bounded + in-memory store, which evicts its oldest retained events. + """ + + +class DeliveryMode(StrEnum): + """Message delivery mode. + + TRIGGER: persist + hot queue + EventBus broadcast. + Use when receiver must act immediately (e.g., assertion → critic). + QUEUE: persist + hot queue only (no broadcast). + Use for status updates the receiver pulls when ready. + """ + + TRIGGER = "trigger" + QUEUE = "queue" + + +class AgentComm: + """Sends inter-agent messages and logs them as events. + + All messages are persisted to EventStore (truth source). + Hot queues are optional in-memory acceleration. + DeliveryMode controls whether EventBus broadcast fires. + """ + + def __init__( + self, + event_store: _AgentCommEventStore, + event_bus: EventBus, + ) -> None: + self._event_store = event_store + self._event_bus = event_bus + # Each message type gets an independent cursor. A specialized + # consumer (for example AgentBus result recovery) must not advance + # the generic inbox cursor or consume status/control messages that + # belong to another component. + self._cursors: dict[tuple[str, str | None, str | None], int] = {} + self._hot_queues: dict[str, asyncio.Queue[KernelEvent]] = {} + + async def send( + self, + msg: AgentMessage, + mode: DeliveryMode = DeliveryMode.QUEUE, + ) -> KernelEvent: + """Send an agent message. + + Always persists to EventStore (truth source). + TRIGGER mode additionally broadcasts via EventBus for immediate wakeup. + """ + event = KernelEvent( + task_id=TaskId(msg.task_id), + event_type=EventType.AGENT_MESSAGE, + from_agent=msg.from_agent, + to_agent=msg.to_agent, + message_type=msg.message_type, + correlation_id=msg.content.get("correlation_id"), + payload={ + "message_id": msg.id, + "from_agent": msg.from_agent, + "to_agent": msg.to_agent, + "message_type": msg.message_type, + "content": msg.content, + "parent_id": msg.parent_id, + "correlation_id": msg.content.get("correlation_id"), + "delivery_mode": mode.value, + }, + ) + + # 1. Persist (always — EventStore is truth source) + persisted = await self._event_store.append(event) + + # 2. Hot queue (always — non-durable acceleration) + queue = self._hot_queues.get(msg.to_agent) + if queue is not None: + try: + queue.put_nowait(persisted) + except asyncio.QueueFull: + logger.debug( + "Hot queue for %s is full; consumer uses EventStore", + msg.to_agent, + ) + + # 3. Broadcast (TRIGGER only — immediate wakeup signal) + if mode == DeliveryMode.TRIGGER: + await self._event_bus.publish( + event.event_type, event.payload, + ) + + logger.debug( + "Agent message %s: %s → %s [%s] mode=%s", + msg.id, msg.from_agent, msg.to_agent, + msg.message_type, mode.value, + ) + return persisted + + async def consume( + self, + agent_id: str, + *, + task_id: str | None = None, + limit: int = 50, + message_type: str | None = None, + ) -> list[KernelEvent]: + """Consume undelivered messages for an agent using an idempotent cursor. + + Cursor is per ``(agent_id, task_id, message_type)``. Typed consumers + therefore cannot skip or discard unrelated inbox traffic. Restarting + from 0 replays all messages in the selected stream. + + Ordering (L2-2): messages come back in send order within a task. + That rests entirely on the stronger ``_AgentCommEventStore`` + contract above and the "Ordering" section of + ``core.protocols.EventReader``. Three consequences callers rely on: + + * a sender that awaits each ``send`` sees its own messages here + in exactly that order — never reordered, never with a later + message ahead of an earlier one; + * the cursor only moves forward, so each message reaches this + ``(agent, task, message_type)`` reader exactly once even when + ``limit`` truncates the batch — the remainder arrives on the + next call, still in order; + * ``message_type`` is part of the cursor key on purpose: a + shared cursor would let one consumer's read carry another + consumer's mail past unread. + + ``tests/state/test_event_ordering_contract.py`` pins all three + against every retained, readable store implementation. + """ + cursor_key = (agent_id, task_id, message_type) + after_id = self._cursors.get(cursor_key, 0) + # Preserve compatibility with third-party EventReader implementations + # written before typed inbox queries existed. + if message_type is None: + events = await self._event_store.get_events_for_agent( + agent_id, + after_id=after_id, + limit=limit, + task_id=task_id, + ) + else: + events = await self._event_store.get_events_for_agent( + agent_id, + after_id=after_id, + limit=limit, + task_id=task_id, + message_type=message_type, + ) + if events: + ids = [int(event.id) for event in events] + if any(current <= previous for previous, current in itertools.pairwise(ids)): + # Advancing to max(ids) can permanently skip lower ids that + # a misordered, limited query left outside this batch. Keep + # the cursor untouched and fail loudly instead of returning + # a partial stream with silent data loss. + logger.error( + "Event store returned %d events for %s outside strict " + "ascending id order (%s); cursor remains at %d", + len(ids), agent_id, ids, after_id, + ) + raise RuntimeError( + "EventReader must return events in strictly ascending id order" + ) + self._cursors[cursor_key] = ids[-1] + return events + + def reset_cursor( + self, + agent_id: str, + task_id: str | None = None, + *, + message_type: str | None = None, + ) -> None: + """Reset one typed cursor, or every cursor for an agent/task pair.""" + if message_type is not None: + self._cursors.pop((agent_id, task_id, message_type), None) + return + for key in list(self._cursors): + if key[:2] == (agent_id, task_id): + self._cursors.pop(key, None) + + def hot_queue_for( + self, agent_id: str, + ) -> asyncio.Queue[KernelEvent]: + """Return the non-durable hot queue for an agent.""" + return self._hot_queues.setdefault( + agent_id, asyncio.Queue(maxsize=256), + ) diff --git a/agent_core/components/agent_bus/bus.py b/agent_core/components/agent_bus/bus.py new file mode 100644 index 0000000..8d3529e --- /dev/null +++ b/agent_core/components/agent_bus/bus.py @@ -0,0 +1,2176 @@ +# pyright: reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownParameterType=false, reportUnknownVariableType=false, reportMissingParameterType=false, reportMissingTypeArgument=false, reportUnnecessaryContains=false, reportUnnecessaryIsInstance=false +"""AgentBus — async job model for parallel sub-agent dispatch. + +Phase A (Issue #24): submit/collect/abort async job model. +Phase B (Issue #24): SpawnGuard budget-aware concurrency/depth/token control. +dispatch_parallel is preserved as backward-compatible convenience wrapper. + +Phase 0 thin-kernel guardrail: +- Do not add new research-specific observer defaults here. +- Do not add new domain-shaped result fields here. +- Do not add new workflow-specific prompt/result assembly here. + +AgentBus is a migration target toward a generic subtask runtime. Research +observer stacks and result adapters should move to workflow-owned assembly. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import inspect +import json +import logging +import time +from collections import deque +from collections.abc import Awaitable, Callable +from pathlib import Path +from typing import Any + +from agent_core.components.agent_bus.models import ( + CollectResult, + DepthLimitExceeded, + JobEntry, + PendingSessionTask, + SessionWaitOutcome, + SubAgentResult, + SubAgentRuntimeSpec, + SubAgentSession, + SubTask, +) +from agent_core.components.agent_bus.runtime import SUBTASK_RESULT_MESSAGE_TYPE +from agent_core.components.agent_bus.runtime import ( + adapt_default_session_result as _adapt_default_session_result, +) +from agent_core.components.agent_bus.runtime import ( + adapt_default_subagent_result as _adapt_default_subagent_result, +) +from agent_core.components.agent_bus.runtime import ( + build_default_subagent_loop_config as _build_default_subagent_loop_config, +) +from agent_core.components.agent_bus.runtime import ( + build_default_subagent_observers as _build_default_subagent_observers, +) +from agent_core.components.agent_bus.runtime import ( + build_session_loop_config as _build_session_loop_config, +) +from agent_core.components.agent_bus.runtime import ( + close_session_boundary_aborted as _close_session_boundary_aborted, +) +from agent_core.components.agent_bus.runtime import ( + emit_session_task_completed as _emit_session_task_completed, +) +from agent_core.components.agent_bus.runtime import ( + emit_session_task_submitted as _emit_session_task_submitted, +) +from agent_core.components.agent_bus.runtime import ( + message_content_to_result as _message_content_to_result, +) +from agent_core.components.agent_bus.runtime import ( + rehydrate_orphaned_jobs as _rehydrate_orphaned_jobs, +) +from agent_core.components.agent_bus.runtime import ( + resolve_session_observers as _resolve_session_observers, +) +from agent_core.components.agent_bus.runtime import ( + safe_final_content as _safe_final_content, +) +from agent_core.components.agent_bus.runtime import ( + safe_metadata as _safe_metadata, +) +from agent_core.components.agent_bus.shared_pool import SharedArtifactPool +from agent_core.components.agent_bus.spawn_guard import SpawnGuard +from agent_core.components.observers.wall_clock_guard import WallClockGuard +from agent_core.events import EventType +from agent_core.loop_types import BaseObserver +from agent_core.messages import ( + Message, + assistant_msg, + system_msg, + user_msg, +) +from agent_core.models.pipeline_spec import SubAgentProfile +from agent_core.protocols import EventSink +from agent_core.runtime.loop.agent_loop import run_agent_loop +from agent_core.runtime.loop.message_trimmer import ( + MessageTrimmer, + NullTrimmer, + find_final_assistant, + trim_and_remap_boundaries, +) +from agent_core.runtime.registries import services as registry +from agent_core.runtime.registries.agents import AgentRegistry +from agent_core.tool import Tool + +logger = logging.getLogger(__name__) + + +PauseCheckFn = Callable[[], Awaitable[bool]] +PauseCheckFactory = Callable[[str], PauseCheckFn | None] + + +def _with_wall_clock_guard( + observers: list[Any] | None, + guard: SpawnGuard | None, +) -> list[Any]: + """Append a soft wall-clock deadline derived from ``guard.timeout_s``. + + Both spawn paths cap a sub-agent with a hard ``asyncio.wait_for(coro, + guard.timeout_s)``. That is a cancel from OUTSIDE the loop, so the + coroutine dies at an arbitrary ``await``: the loop never reaches an exit + branch (``stopped_by`` stays ``""``), the trajectory ``end`` event is + never written, and the runtime's pre-absorb ``force_final_answer`` hook — + which sits *after* the ``wait_for`` — is skipped. The sub-agent's whole + run is then discarded as ``(empty report)``. Measured at **52.8% of + sub-agents** on bc200 × agent-team; see :class:`WallClockGuard`. + + Both the soft deadline here and the hard one below read the same + ``guard.timeout_s``, which is the point of routing them through one + helper — they cannot drift apart, and neither can the two call sites. + Attached at this layer rather than inside each workflow's + ``observers_builder`` so every fan-out workflow inherits it. + """ + if guard is None or guard.timeout_s <= 0: + return list(observers or []) + return [ + *(observers or []), + WallClockGuard(budget_s=float(guard.timeout_s)), + ] + + +def _strip_job_suffix(task_id: str) -> str: + """Strip ``.job.N`` suffixes iteratively to recover the root task id. + + Sub-agent spawn sites pass ``job_id`` (e.g. ``"root-1.job.2"``) + rather than the root task id — that's the identifier threaded + through LoopConfig. But pause_check must bind to the REAL task + managed by ProcessManager, which exists only for the root. At + depth >= 2 this is what makes grandsub-agents observe root pause + instead of silently running their full max_turns. + """ + root = task_id + while ".job." in root: + root = root.rsplit(".job.", 1)[0] + return root + + +def _legacy_pause_check_factory(task_id: str) -> PauseCheckFn | None: + """Product-neutral default: no process-manager pause integration.""" + del task_id + return None + + +def _safe_session_filename(session_id: str) -> str: + """Return a filesystem-safe representation of ``session_id``.""" + return session_id.replace("/", "_").replace(":", "_").replace(" ", "_") + + +class _SessionActivityObserver(BaseObserver): + """Capture a bounded worker event trail for interactive team UIs.""" + + critical = True + _DETAIL_LIMIT = 8_000 + + def __init__(self, session: SubAgentSession) -> None: + self.session = session + self._sequence = 0 + self._thinking_by_turn: dict[int, dict[str, Any]] = {} + + def _append( + self, kind: str, title: str, detail: str, *, turn: int = 0, + is_error: bool = False, + ) -> dict[str, Any]: + detail = str(detail or "").strip() + if len(detail) > self._DETAIL_LIMIT: + detail = detail[:self._DETAIL_LIMIT] + "\n… truncated in live view" + self._sequence += 1 + event = { + "id": f"{self.session.session_id}:{self._sequence}", + "kind": kind, + "title": title, + "detail": detail, + "turn": turn, + "is_error": is_error, + "at": time.monotonic(), + } + self.session.activity_events.append(event) + return event + + async def on_llm_delta(self, ctx: Any) -> None: + delta = str(getattr(ctx, "thinking_delta", "") or "") + if not delta: + return + turn = int(getattr(ctx, "turn", 0) or 0) + event = self._thinking_by_turn.get(turn) + if event is None: + event = self._append("thinking", "thinking", "", turn=turn) + self._thinking_by_turn[turn] = event + detail = str(event.get("detail") or "") + delta + if len(detail) > self._DETAIL_LIMIT: + detail = detail[:self._DETAIL_LIMIT] + "\n… truncated in live view" + event["detail"] = detail + + async def on_llm_response(self, ctx: Any) -> None: + thinking = str(getattr(ctx, "thinking", "") or "").strip() + if thinking: + turn = int(getattr(ctx, "turn", 0) or 0) + event = self._thinking_by_turn.get(turn) + if event is None: + self._thinking_by_turn[turn] = self._append( + "thinking", "thinking", thinking, turn=turn, + ) + else: + event["detail"] = thinking[:self._DETAIL_LIMIT] + text = str(getattr(ctx, "ai_text", "") or "").strip() + if text: + self._append( + "message", "assistant", text, + turn=int(getattr(ctx, "turn", 0) or 0), + ) + + async def on_tool_call(self, ctx: Any, tool_call: dict[str, Any]) -> None: + name = str(tool_call.get("name") or "tool") + args = tool_call.get("args") or {} + try: + detail = json.dumps(args, ensure_ascii=False, indent=2, default=str) + except Exception: + detail = str(args) + self._append( + "tool_call", name, detail, + turn=int(getattr(ctx, "turn", 0) or 0), + ) + + async def on_tool_result(self, ctx: Any, result: Any) -> None: + self._append( + "tool_error" if bool(getattr(result, "is_error", False)) else "tool_result", + str(getattr(result, "name", "tool") or "tool"), + str(getattr(result, "result", "") or ""), + turn=int(getattr(ctx, "turn", 0) or 0), + is_error=bool(getattr(result, "is_error", False)), + ) + + + + +def _offload_dropped_messages( + session: SubAgentSession, + dropped: list[Message], + out_dir: Path, +) -> None: + """Write dropped messages to a per-session JSONL file, best-effort. + + Each message is serialised individually to avoid materialising the full + history as a single dict — this keeps peak memory proportional to + the largest single message rather than the entire dropped batch. + Native ``Message`` is a plain ``TypedDict``, so it serialises directly + with ``json.dumps`` (no langchain ``message_to_dict`` conversion needed). + Failures are logged but never re-raised so the trim always proceeds. + """ + import json + + try: + session_dir = out_dir / _safe_session_filename(session.session_id) + session_dir.mkdir(parents=True, exist_ok=True) + out_path = session_dir / f"task_{session.total_task_count:03d}.jsonl" + with out_path.open("w", encoding="utf-8") as f: + for msg in dropped: + f.write(json.dumps(msg, ensure_ascii=False) + "\n") + session.offloaded_history_path = out_path + except Exception: + logger.exception( + "eager_trim: failed to offload dropped messages for %s; trim still proceeds", + session.session_id, + ) + + +# ── AgentBus ──────────────────────────────────────────────────────────── + + +async def _run_within_context(setup, job_id, item, inner): + """Await *inner* (the sub-agent ``run_agent_loop`` coroutine) inside the + context manager returned by ``setup(job_id, item)``. + + Runs in the sub-agent's own asyncio task, so any contextvar the CM sets + (e.g. a per-agent sandbox) is scoped to that sub-agent and reset on exit. + """ + with setup(job_id, item): + return await inner + + +class AgentBus: + """Async job model for sub-agent dispatch. + + Core API: submit() / collect() / abort(). + dispatch_parallel() is a backward-compatible convenience wrapper. + + If a SpawnGuard is attached, submit() enforces concurrency/depth/token + limits. Without a guard, only basic depth checking applies. + """ + + def __init__( + self, + *, + event_sink: EventSink | None = None, + pause_check_factory: PauseCheckFactory | None = None, + agent_registry: AgentRegistry | None = None, + resource_manager: Any = None, + session_history_dir: Path | None = None, + ) -> None: + """Construct an AgentBus. + + Args: + event_sink: Optional ``core.protocols.EventSink``. When set, + the bus emits all telemetry through this sink directly + instead of doing a global ``registry.get(EventSink)`` + lookup on every call. + pause_check_factory: Optional ``(root_task_id) -> PauseCheckFn`` + builder. + agent_registry: Optional ``AgentRegistry``. When set, the bus + resolves agent prompts from this instance directly + instead of doing ``registry.get(AgentRegistry)``. Phase 4 + D13 surgery 3 — lets the SDK assembly inject a + workflow-scoped registry without polluting the global + service container. + resource_manager: Optional ``ResourceManager``. When set, the + bus resolves per-role LLMs / tools through this instance + instead of via ``registry.get(ResourceManager)``. Phase 4 + D13 surgery 3. + session_history_dir: Optional base directory for offloading + dropped messages to disk after eager trim. When ``None``, + only the in-memory trim is performed (memory still reclaimed). + + ``None`` for any of the kwargs keeps the legacy registry fallback so + existing zero-arg ``AgentBus()`` callers continue to work. + """ + self._jobs: dict[str, JobEntry] = {} + self._counter: int = 0 + self._spawn_guard: SpawnGuard | None = None + self._sub_agent_profiles: dict[str, dict[str, SubAgentProfile]] = {} + self._sessions: dict[str, SubAgentSession] = {} + # Per-task evidence / assertion pool — populated when main agents + # (or collect_reports) harvest completed SubAgentResults, drained + # by the pipeline node at phase end for the report node to consume. + self._task_aggregates: dict[str, dict[str, list[dict[str, Any]]]] = {} + # L1-0 idempotency: job ids already handed to a fan-in caller, per + # task. The in-memory queue pop is destructive, so without this the + # message-based reconciliation below would re-deliver a result that + # had already been drained and the main agent would see the same + # ```` twice. This is also the dedupe table L2-4 needs. + self._delivered_jobs: dict[str, set[str]] = {} + # Agents that are addressed by ``subtask_result`` messages for a + # task, learned from ``spawn_context`` at dispatch. Keeps the + # reconciliation read cheap in-process; after a restart it is empty + # and recipients are recovered from the log instead. + self._result_recipients: dict[str, set[str]] = {} + # Results pulled out of the message channel but not yet handed to a + # caller. ``AgentComm.consume`` advances its cursor past the whole + # batch it read, so returning only the first message would silently + # drop the rest; they wait here instead — the message-channel + # counterpart of ``session.pending_results``. + self._recovered_results: dict[ + str, deque[tuple[str, SubAgentResult]] + ] = {} + self._event_sink_injected: EventSink | None = event_sink + self._pause_check_factory: PauseCheckFactory | None = pause_check_factory + self._agent_registry_injected: AgentRegistry | None = agent_registry + self._resource_manager_injected: Any = resource_manager + self._session_history_dir: Path | None = session_history_dir + + def _event_sink(self) -> Any: + """Resolve the active event sink. + + Returns the constructor-injected sink when available; otherwise + falls back to the global ``registry.get_optional(EventSink)`` + lookup, then the legacy concrete ``EventStore`` registration, + so legacy callers keep working. + """ + if self._event_sink_injected is not None: + return self._event_sink_injected + event_sink = registry.get_optional(EventSink) + if event_sink is not None: + return event_sink + return registry.get_optional_by_type_name("EventStore") + + def _agent_registry(self, *, required: bool) -> Any: + """Resolve the active ``AgentRegistry``. + + Returns the constructor-injected registry when set; otherwise + falls back to the global service registry. ``required=True`` + raises (via ``registry.get``) when no registry is available; + ``required=False`` returns ``None`` instead. + """ + if self._agent_registry_injected is not None: + return self._agent_registry_injected + if required: + return registry.get(AgentRegistry) + return registry.get_optional(AgentRegistry) + + def _resource_manager(self, *, required: bool) -> Any: + """Resolve the active ``ResourceManager`` (lazy import to avoid + a kernel→runtime.resources eager dependency at module load). + + + """ + if self._resource_manager_injected is not None: + return self._resource_manager_injected + from agent_core.runtime.resources.manager import ResourceManager + if required: + return registry.get(ResourceManager) + return registry.get_optional(ResourceManager) + + def _task_pause_check_or_none(self, task_id: str | None) -> PauseCheckFn | None: + """Resolve a per-task pause check closure. + + Strips ``.job.N`` suffixes to recover the root task id, then: + - if a ``pause_check_factory`` was injected, calls it directly; + - otherwise falls back to the legacy lazy import of + ``agent_core.runtime.pause_check.make_task_pause_check``. + + Returns ``None`` when ``task_id`` is empty (orphan/shadow runs) + or when both paths come up empty. + """ + if not task_id: + return None + root = _strip_job_suffix(task_id) + if self._pause_check_factory is not None: + try: + return self._pause_check_factory(root) + except Exception: + logger.warning( + "pause_check_factory raised for task=%s; running without pause probe", + root, exc_info=True, + ) + return None + return _legacy_pause_check_factory(root) + + def register_sub_agent_profiles( + self, node_id: str, profiles: dict[str, SubAgentProfile] + ) -> None: + """Register sub-agent profiles for a node (called at compile time).""" + self._sub_agent_profiles[node_id] = profiles + + def get_sub_agent_profile( + self, node_id: str, profile_name: str + ) -> SubAgentProfile | None: + """Look up a named sub-agent profile for a node.""" + node_profiles = self._sub_agent_profiles.get(node_id, {}) + return node_profiles.get(profile_name) + + def set_spawn_guard(self, guard: SpawnGuard) -> None: + """Attach a SpawnGuard for budget-aware spawn control.""" + self._spawn_guard = guard + + @property + def spawn_guard(self) -> SpawnGuard | None: + return self._spawn_guard + + def _next_job_id(self, parent_task_id: str) -> str: + self._counter += 1 + return f"{parent_task_id}.job.{self._counter}" + + # ── submit: non-blocking ──────────────────────────────────────────── + + async def submit( + self, + parent_task_id: str, + item: SubTask, + *, + shared_evidence: SharedArtifactPool | None = None, + max_turns: int = 8, + current_depth: int = 0, + max_depth: int = 2, + estimated_tokens: int = 0, + runtime_spec: SubAgentRuntimeSpec | None = None, + spawn_context: dict[str, Any] | None = None, + ) -> str: + """Submit a sub-agent job. Returns job_id immediately. + + If a SpawnGuard is attached, enforces concurrency/depth/token + limits. Without a guard, only basic depth checking applies. + + Raises: + DepthLimitExceeded: depth >= max_depth (or SpawnGuard depth) + SpawnDepthExceeded: SpawnGuard depth limit + BudgetExhausted: SpawnGuard token budget exceeded + """ + guard = self._spawn_guard + + # Use SpawnGuard depth limit if stricter + effective_max_depth = max_depth + if guard and guard.max_depth < effective_max_depth: + effective_max_depth = guard.max_depth + + if current_depth >= effective_max_depth: + raise DepthLimitExceeded( + f"Sub-agent depth limit exceeded: " + f"current_depth={current_depth}, " + f"max_depth={effective_max_depth}" + ) + + job_id = self._next_job_id(parent_task_id) + agent_registry = self._agent_registry(required=True) + dispatch_depth = current_depth + 1 + + # Layer 5+3: SpawnGuard pre-check (depth + budget — non-blocking) + # Layer 4 (concurrency) is deferred to _run_and_finalize + if guard: + await guard.pre_check( + job_id, dispatch_depth, estimated_tokens, + ) + + system_prompt = ( + item.system_prompt + or agent_registry.get_prompt_for(item.role_id) + ) + runtime = runtime_spec or SubAgentRuntimeSpec() + + async def _run() -> SubAgentResult: + # Declared outside try/except so exception handlers can + # salvage partial metadata on failure (see the ``agent_result`` + # assignment inside the try block). + agent_result: Any = None + try: + event_store = self._event_sink() + resource_mgr = self._resource_manager(required=True) + sub_llm = resource_mgr.get_llm(item.role_id) + sub_tools = resource_mgr.get_tools_for_role( + item.role_id, + ) + + if runtime.config_builder is not None: + sub_config = runtime.config_builder( + job_id, item, max_turns, + ) + else: + sub_config = _build_default_subagent_loop_config( + job_id, item, max_turns, + ) + + if runtime.observers_builder is not None: + # Non-session dispatch — no task_index ordinal; pass 0. + sub_observers = runtime.observers_builder( + job_id, item, 0, + ) + else: + sub_observers = _build_default_subagent_observers( + event_store=event_store, + job_id=job_id, + ) + + # codex #6 / Log Schema v1 §2.4: stamp spawn_context onto + # observers that surface it (WorkerTraceFileObserver + # writes it as a top-level field on the sub trace doc). + if spawn_context: + for _obs in sub_observers: + _setter = getattr(_obs, "set_extension_data", None) + if _setter is None: + continue + try: + _setter(spawn_context=dict(spawn_context)) + except Exception as _exc: + logger.debug( + "spawn_context stamp failed on %s: %s", + type(_obs).__name__, _exc, + ) + + sub_observers = _with_wall_clock_guard(sub_observers, guard) + + coro = run_agent_loop( + system_prompt=system_prompt, + user_message=item.question, + llm=sub_llm, + tools=sub_tools, + config=sub_config, + observers=sub_observers, + model_profile=runtime.model_profile, + history_policy=runtime.history_policy, + pause_check=self._task_pause_check_or_none(parent_task_id), + ) + if runtime.context_setup is not None: + coro = _run_within_context( + runtime.context_setup, job_id, item, coro, + ) + # Layer 2: Timeout from SpawnGuard + timeout = float(guard.timeout_s) if guard else None + if timeout: + agent_result = await asyncio.wait_for( + coro, timeout=timeout, + ) + else: + agent_result = await coro + if runtime.result_adapter is not None: + adapted = runtime.result_adapter(agent_result, job_id, item) + if inspect.isawaitable(adapted): + adapted = await adapted + return adapted + return _adapt_default_subagent_result( + agent_result, job_id, item, + ) + except asyncio.CancelledError: + return SubAgentResult( + question=item.question, + role_id=item.role_id, + final_content=_safe_final_content(agent_result), + success=False, + error="aborted", + error_class="CancelledError", + job_id=job_id, + metadata=_safe_metadata(agent_result), + ) + except TimeoutError: + logger.warning( + "Sub-agent job %s timed out", job_id, + ) + return SubAgentResult( + question=item.question, + role_id=item.role_id, + final_content=_safe_final_content(agent_result), + success=False, + error="timeout", + error_class="TimeoutError", + job_id=job_id, + metadata=_safe_metadata(agent_result), + ) + except Exception as exc: + logger.warning( + "Sub-agent job %s failed: %s", job_id, exc + ) + return SubAgentResult( + question=item.question, + role_id=item.role_id, + final_content=_safe_final_content(agent_result), + success=False, + error=str(exc), + error_class=type(exc).__name__, + job_id=job_id, + metadata=_safe_metadata(agent_result), + ) + + async def _run_and_finalize() -> SubAgentResult: + entry = self._jobs[job_id] + # Layer 4: Concurrency gate (may queue here) + if guard: + await guard.acquire_slot(job_id) + entry.status = "running" + try: + result = await _run() + entry.result = result + entry.completed_at = time.monotonic() + if result.error == "aborted": + entry.status = "aborted" + elif result.success: + entry.status = "completed" + else: + entry.status = "failed" + return result + finally: + # Layer 1: RAII release — always release guard slot + if guard: + guard.release(job_id) + + task = asyncio.create_task(_run_and_finalize(), name=job_id) + entry = JobEntry( + job_id=job_id, + parent_task_id=parent_task_id, + item=item, + task=task, + status="submitted", + submitted_at=time.monotonic(), + ) + self._jobs[job_id] = entry + + event_store = self._event_sink() + if event_store: + await event_store.append( + task_id=parent_task_id, + event_type=EventType.AGENT_ACTION, + payload={ + "trace_type": "agent_submitted", + "parent_task_id": parent_task_id, + "job_id": job_id, + "question": item.question, + "role_id": item.role_id, + "depth": dispatch_depth, + }, + agent_role="system", + ) + + logger.info( + "Submitted job %s (question=%s, role=%s, depth=%d)", + job_id, item.question[:60], item.role_id, dispatch_depth, + ) + return job_id + + # ── collect: wait with timeout ────────────────────────────────────── + + async def collect( + self, + job_ids: list[str], + *, + timeout: float = 1800, + ) -> CollectResult: + """Wait for jobs to complete. Returns partial results on timeout. + + Default 1800s matches the ``collect_results`` tool's user-facing + default — sub-agents on hard BrowseComp / FrontierScience runs + regularly need 10–20 min, and a shorter primitive default would + silently strand work for any caller that forgets to override. + """ + tasks_to_wait: dict[asyncio.Task, str] = {} + result = CollectResult() + + for jid in job_ids: + entry = self._jobs.get(jid) + if entry is None: + continue + if entry.result is not None: + # Already finished + if entry.status == "failed" or not entry.result.success: + result.failed.append(entry.result) + else: + result.completed.append(entry.result) + else: + if entry.status == "queued" or entry.task is None: + result.pending.append(jid) + else: + tasks_to_wait[entry.task] = jid + + if tasks_to_wait: + done, pending = await asyncio.wait( + tasks_to_wait.keys(), + timeout=timeout, + return_when=asyncio.ALL_COMPLETED, + ) + + for task in done: + jid = tasks_to_wait[task] + entry = self._jobs.get(jid) + if entry and entry.result: + if entry.result.success: + result.completed.append(entry.result) + else: + result.failed.append(entry.result) + + for task in pending: + jid = tasks_to_wait[task] + result.pending.append(jid) + + event_store = self._event_sink() + if event_store and job_ids: + parent = self._jobs.get(job_ids[0]) + parent_task_id = parent.parent_task_id if parent else "" + await event_store.append( + task_id=parent_task_id, + event_type=EventType.AGENT_ACTION, + payload={ + "trace_type": "agent_collected", + "job_ids": job_ids, + "completed": len(result.completed), + "failed": len(result.failed), + "pending": len(result.pending), + }, + agent_role="system", + ) + + return result + + # ── abort: graceful cancel ────────────────────────────────────────── + + async def abort(self, job_id: str) -> str: + """Abort a job. Returns final status string. + + Idempotent: aborting a completed job returns its actual status. + Aborting a queued (not-yet-dispatched) job removes it from its + session's pending queue, releases its guard reservation, and + marks it aborted without touching the running task. + """ + entry = self._jobs.get(job_id) + if entry is None: + return "not_found" + + # Already finished — don't override + if entry.status in ("completed", "failed", "aborted"): + return entry.status + + if entry.status == "queued" or entry.task is None: + # Queued task: scrub from the session's pending deque, + # release any guard reservation, and finalise as aborted. + self._purge_queued_job(job_id) + self._mark_job_aborted(job_id, entry) + else: + # Running task: cancel the asyncio task and let + # _run_and_finalize observe the CancelledError. + entry.task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await entry.task + + # If the task finished between cancel() and await (race), + # respect the actual result + if entry.result and entry.result.success: + entry.status = "completed" + elif entry.status not in ("completed", "failed"): + self._mark_job_aborted(job_id, entry) + + event_store = self._event_sink() + if event_store: + await event_store.append( + task_id=entry.parent_task_id, + event_type=EventType.AGENT_ACTION, + payload={ + "trace_type": "agent_aborted", + "job_id": job_id, + "final_status": entry.status, + }, + agent_role="system", + ) + + logger.info("Aborted job %s → %s", job_id, entry.status) + return entry.status + + async def cancel_for_parent(self, parent_task_id: str) -> int: + """Cancel all pending/running jobs for a parent task. + + Called between agent mode rounds to prevent orphaned sub-agents + from holding semaphore slots and blocking the next round. Both + running tasks (asyncio cancel) and queued tasks (purge + status + flip) are handled — neither survives the round boundary. + + Returns the number of jobs cancelled. + """ + cancelled = 0 + for jid, entry in list(self._jobs.items()): + if entry.parent_task_id != parent_task_id: + continue + if entry.status in ("completed", "failed", "aborted"): + continue + if entry.status == "queued" or entry.task is None: + self._purge_queued_job(jid) + self._mark_job_aborted(jid, entry) + cancelled += 1 + continue + entry.task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await entry.task + if entry.status not in ("completed", "failed", "aborted"): + self._mark_job_aborted(jid, entry) + cancelled += 1 + if cancelled: + logger.info( + "Cancelled %d orphaned sub-agent jobs for %s", + cancelled, parent_task_id, + ) + return cancelled + + def _purge_queued_job(self, job_id: str) -> None: + """Remove ``job_id`` from whichever session has it queued. + + Releases the guard's pre-check reservation if a guard is + attached — queued tasks reserved budget at submit time and + won't reach the ``_run_and_finalize`` finally that normally + releases. + """ + for session in self._sessions.values(): + queue = session.pending_tasks + if not queue: + continue + kept: deque[PendingSessionTask] = deque() + removed = False + for pending in queue: + if pending.job_id == job_id: + removed = True + continue + kept.append(pending) + if removed: + session.pending_tasks = kept + if self._spawn_guard is not None: + self._spawn_guard.release(job_id) + return + + def _mark_job_aborted( + self, job_id: str, entry: JobEntry | None = None, + ) -> None: + """Mark a job aborted and release any SpawnGuard reservation. + + This is intentionally idempotent. A task can be cancelled before + its coroutine reaches the ``finally`` block that normally releases + the guard reservation; queued tasks never reach that block at all. + Calling ``release`` here covers both without double-freeing slots. + """ + if self._spawn_guard is not None: + self._spawn_guard.release(job_id) + entry = entry or self._jobs.get(job_id) + if entry is None or entry.status in ("completed", "failed", "aborted"): + return + entry.status = "aborted" + entry.completed_at = time.monotonic() + if entry.result is not None: + return + item = entry.item + entry.result = SubAgentResult( + question=str(getattr(item, "question", "") or ""), + role_id=str(getattr(item, "role_id", "") or ""), + final_content="", + success=False, + error="aborted", + error_class="CancelledError", + job_id=job_id, + ) + + # ── get_job_status: query job state ───────────────────────────────── + + def get_job_status(self, job_id: str) -> str | None: + """Return current status of a job, or None if not found.""" + entry = self._jobs.get(job_id) + return entry.status if entry else None + + # ── Persistent sessions (AgentTeam-style multi-task sub-agents) ───── + + def _make_session_id(self, task_id: str, name: str) -> str: + return f"{task_id}::{name}" + + async def create_session( + self, + *, + task_id: str, + name: str, + role_id: str, + system_prompt: str | None = None, + tools_override: list[Tool] | None = None, + llm_override: Any = None, + trimmer: MessageTrimmer | None = None, + max_turns: int = 100, + tool_result_max_chars: int | None = None, + runtime_spec: SubAgentRuntimeSpec | None = None, + llm_timeout: int | None = None, + ) -> str: + """Create (or return existing) persistent sub-agent session. + + Idempotent ONLY when the caller asks for exactly the same session: + same ``(task_id, name, role_id, system_prompt)``. If a session with + the same ``(task_id, name)`` exists but with a different role or + system prompt, raise ``ValueError`` rather than silently reusing + the old one — agent names are often LLM-generated and a clashing + specialization is almost always a bug (e.g. a fresh verifier + getting turned into a reused researcher). + + ``system_prompt`` falls back to ``AgentRegistry.get_prompt_for(role_id)``. + ``tools_override`` / ``llm_override`` fall back to ResourceManager. + ``trimmer`` defaults to ``NullTrimmer``. + + ``llm_timeout`` (seconds) sets the per-LLM-call timeout for every + task submitted to this session. ``None`` (default) defers to + ``LoopConfig.llm_timeout`` (180s). Set this when the session's + model has high response latency variance (e.g. a strong reasoning + model used for review). Identity check (idempotent reuse) ignores + this field — it is treated as a runtime knob, not part of the + session's identity. + """ + prompt = system_prompt + if prompt is None: + agent_registry = self._agent_registry(required=False) + if agent_registry is not None: + prompt = agent_registry.get_prompt_for(role_id) + if not prompt: + prompt = "" + + session_id = self._make_session_id(task_id, name) + existing = self._sessions.get(session_id) + if existing is not None: + if existing.role_id != role_id or existing.system_prompt != prompt: + raise ValueError( + f"Session name {name!r} already exists for task " + f"{task_id!r} with a different role or system_prompt " + f"(existing role={existing.role_id!r}, requested " + f"role={role_id!r}). Pick a unique name." + ) + logger.debug( + "create_session: returning existing session %s", session_id, + ) + return session_id + + llm = llm_override + tools = tools_override + if llm is None or tools is None: + resource_mgr = self._resource_manager(required=False) + if llm is None and resource_mgr is not None: + llm = resource_mgr.get_llm(role_id) + if tools is None and resource_mgr is not None: + tools = resource_mgr.get_tools_for_role(role_id) + if tools is None: + tools = [] + + session = SubAgentSession( + session_id=session_id, + task_id=task_id, + name=name, + role_id=role_id, + system_prompt=prompt, + tools=tools, + llm=llm, + trimmer=trimmer or NullTrimmer(), + max_turns=max_turns, + tool_result_max_chars=tool_result_max_chars, + llm_timeout=llm_timeout, + runtime_spec=runtime_spec, + ) + self._sessions[session_id] = session + + event_store = self._event_sink() + if event_store is not None: + await event_store.append( + task_id=task_id, + event_type=EventType.AGENT_ACTION, + payload={ + "trace_type": "session_created", + "session_id": session_id, + "name": name, + "role_id": role_id, + "agent": name, + "action": "session_created", + "detail": ( + f"Created sub-agent '{name}' (role={role_id})" + ), + }, + agent_role="system", + ) + + logger.info( + "Created session %s (role=%s, max_turns=%d)", + session_id, role_id, max_turns, + ) + return session_id + + async def submit_task_to_session( + self, + session_id: str, + task_prompt: str, + *, + max_turns: int | None = None, + observers: list[Any] | None = None, + estimated_tokens: int = 0, + runtime_spec: SubAgentRuntimeSpec | None = None, + spawn_context: dict[str, Any] | None = None, + task_metadata: dict[str, Any] | None = None, + ) -> str: + """Queue a task on an existing session. Non-blocking: returns job_id. + + Sessions execute tasks **strictly serially** — the session's + message list / boundary list cannot be safely mutated by two + concurrent runs. When the session is already busy, the new task + joins ``session.pending_tasks`` (FIFO) and runs as soon as the + current task finalises. This is transparent to the caller: the + returned ``job_id`` is valid immediately and ``collect_reports`` + sees queued tasks once they start running. + + When a ``SpawnGuard`` is attached to the bus, session tasks flow + through the same five-layer protection as ``submit()``: depth + + token-budget pre-check (immediate), concurrency semaphore (per- + run, gated inside the task), per-job wall-time timeout, and RAII + slot release. This prevents a bad decomposition from fanning out + unbounded background work on the session path. + + Raises ``KeyError`` if the session doesn't exist. + Raises ``SpawnDepthExceeded`` / ``BudgetExhausted`` via the + guard's pre-check (queued tasks reserve budget at submit time + too — aborting a queued task releases the reservation). + """ + session = self._sessions.get(session_id) + if session is None: + raise KeyError(f"Session {session_id!r} not found") + + guard = self._spawn_guard + # Sessions are spawned by the main agent (depth 0) → session tasks + # are depth 1. The swarm_sub role doesn't itself carry + # delegate_subtask / create_subagent so there's no further + # nesting in the current workflow, but we still respect the guard + # so future recursive workflows don't silently bypass it. + dispatch_depth = 1 + + job_id = self._next_job_id(session.task_id) + + # Layer 5+3: SpawnGuard pre-check at submit time (not at dispatch) + # so queued tasks still raise BudgetExhausted / SpawnDepthExceeded + # synchronously at the API boundary — preserves the legacy + # behaviour that callers see budget errors immediately rather + # than only via ``collect_reports`` after the queue drains. + if guard is not None: + await guard.pre_check(job_id, dispatch_depth, estimated_tokens) + + pending = PendingSessionTask( + job_id=job_id, + task_prompt=task_prompt, + max_turns=max_turns, + observers=observers, + estimated_tokens=estimated_tokens, + runtime_spec=runtime_spec, + spawn_context=dict(spawn_context) if spawn_context else None, + task_metadata=dict(task_metadata or {}), + ) + + # Busy = a current task is in submitted/running state. ``queued`` + # current_job_id can't happen (we only set current_job_id at + # dispatch time below), but check defensively in case future code + # paths set the field early. + busy = False + if session.current_job_id is not None: + current_entry = self._jobs.get(session.current_job_id) + if current_entry is not None and current_entry.status in ( + "submitted", "running", + ): + busy = True + + if busy: + # Park in the queue with a placeholder JobEntry so consumers + # that look up by ``job_id`` (status queries, abort paths) + # see the task before it actually starts running. + entry = JobEntry( + job_id=job_id, + parent_task_id=session.task_id, + item=SubTask( + question=task_prompt, role_id=session.role_id, + system_prompt=session.system_prompt, + metadata=dict(task_metadata or {}), + ), + task=None, + status="queued", + submitted_at=time.monotonic(), + ) + self._jobs[job_id] = entry + session.pending_tasks.append(pending) + logger.debug( + "Session %s busy → queued job %s (%d ahead)", + session_id, job_id, len(session.pending_tasks) - 1, + ) + return job_id + + # Free path: dispatch immediately. + await self._dispatch_session_task(session, pending) + return job_id + + async def _dispatch_session_task( + self, + session: SubAgentSession, + pending: PendingSessionTask, + ) -> None: + """Spin up the asyncio task for ``pending`` on this session. + + Opens the new task boundary (so the trimmer sees only closed + boundaries when seeding ``initial_messages``), registers / + upgrades the ``JobEntry``, and emits the + ``session_task_submitted`` SSE event. Called from the free + branch of ``submit_task_to_session`` for fresh tasks and from + ``_drain_session_queue`` when the previous task finalises. + """ + guard = self._spawn_guard + job_id = pending.job_id + task_prompt = pending.task_prompt + observers = pending.observers + max_turns = pending.max_turns + runtime_spec = pending.runtime_spec + session_id = session.session_id + + # Build initial_messages from trimmed prior history, before we + # mutate session.messages so the trimmer sees the clean state. + trimmed = session.trimmer.trim( + session.messages, session.task_boundaries, + ) + initial_messages: list[Message] = [ + system_msg(session.system_prompt), *trimmed, + ] + + # Open a new boundary — the task prompt index is the current tail. + session.total_task_count += 1 + boundary_start = len(session.messages) + session.messages.append(user_msg(task_prompt)) + session.task_boundaries.append((boundary_start, None)) + + effective_max_turns = ( + max_turns if max_turns is not None else session.max_turns + ) + active_runtime = ( + runtime_spec or session.runtime_spec or SubAgentRuntimeSpec() + ) + session_item = SubTask( + question=task_prompt, + role_id=session.role_id, + system_prompt=session.system_prompt, + metadata=dict(pending.task_metadata), + ) + + # Snapshot for _run closure so later mutations don't leak. + # ``run_agent_loop`` only appends a HumanMessage when ``user_message`` + # is non-empty (resume path re-enters with ``user_message=""``). We + # always pass ``task_prompt`` here which is non-empty by contract, + # but the explicit branch keeps the accounting correct if a future + # caller ever invokes this path with an empty prompt. + prefix_count = len(initial_messages) + (1 if task_prompt else 0) + + async def _run() -> SubAgentResult: + # ``raw_result`` lets the exception handlers salvage any + # partial metadata (evidence, assertions, react_steps) when + # the failure came AFTER the agent loop completed — e.g. + # the result_adapter raised, or the boundary bookkeeping + # hit an assertion. Without this, sub-agents that did real + # web_search work but tripped on post-loop bookkeeping + # would leak their evidence silently. + raw_result: Any = None + try: + if active_runtime.config_builder is not None: + loop_config = active_runtime.config_builder( + job_id, session_item, effective_max_turns, + ) + else: + loop_config = _build_session_loop_config( + session, job_id, effective_max_turns, + ) + + if active_runtime.observers_builder is not None: + # Pass task_index as the 1-based ordinal within the session + # (session.total_task_count was already incremented before + # this _run closure was constructed). + loop_observers = active_runtime.observers_builder( + job_id, session_item, session.total_task_count, + ) + else: + loop_observers = _resolve_session_observers( + observers, session, + ) + + # codex #6 / Log Schema v1 §2.4: stamp the spawn_context + # onto any observer with ``set_extension_data`` (today: + # WorkerTraceFileObserver) so the sub-agent's trace file + # carries the delegation lineage at top-level. SFT reflow + # reads ``doc["spawn_context"]`` when merging per-sub + # docs into the canonical ``runs[i].sub_agents[]`` shape. + if pending.spawn_context: + for _obs in loop_observers: + _setter = getattr(_obs, "set_extension_data", None) + if _setter is None: + continue + try: + _setter(spawn_context=dict(pending.spawn_context)) + except Exception as _exc: + logger.debug( + "spawn_context stamp failed on %s: %s", + type(_obs).__name__, _exc, + ) + + loop_observers.append(_SessionActivityObserver(session)) + loop_observers = _with_wall_clock_guard(loop_observers, guard) + + coro = run_agent_loop( + system_prompt=session.system_prompt, + user_message=task_prompt, + initial_messages=initial_messages, + llm=session.llm, + tools=session.tools, + config=loop_config, + observers=loop_observers, + model_profile=active_runtime.model_profile, + history_policy=active_runtime.history_policy, + pause_check=self._task_pause_check_or_none(session.task_id), + ) + if active_runtime.context_setup is not None: + coro = _run_within_context( + active_runtime.context_setup, job_id, session_item, coro, + ) + # Layer 2: per-job wall-time timeout (from guard). + timeout = float(guard.timeout_s) if guard else None + if timeout and timeout > 0: + raw_result = await asyncio.wait_for(coro, timeout=timeout) + else: + raw_result = await coro + + # Pre-absorb hook: let the runtime spec inject a forced + # final turn (e.g. swarm's ``force_final_answer``) + # BEFORE we copy raw_result.messages into the session. + # Anything appended here lands inside the closing + # boundary and feeds session.last_report below. + if active_runtime.force_finalizer is not None: + try: + forced = active_runtime.force_finalizer( + raw_result, session_item, + ) + if asyncio.iscoroutine(forced): + forced = await forced + if forced is not None: + raw_result = forced + except Exception as exc: + logger.warning( + "force_finalizer failed for session %s " + "task #%d: %s — continuing with un-forced " + "result", + session.session_id, + session.total_task_count, + exc, + ) + + # Absorb new turns back into session.messages. + new_turns = list(raw_result.messages[prefix_count:]) + session.messages.extend(new_turns) + + # Boundary closure invariant: the slice (start..end) must + # contain at least one clean AIMessage so the trimmer can + # surface a "final report" on reuse. When the loop exits + # without one (max_turns / budget / llm_error and the + # spec didn't supply / didn't successfully run a + # force_finalizer), synthesise a stub from + # raw_result.final_content or a static reason marker. + final_text = (raw_result.final_content or "").strip() + start_idx, _ = session.task_boundaries[-1] + provisional_end = len(session.messages) - 1 + if find_final_assistant( + session.messages, start_idx + 1, provisional_end, + ) is None: + stopped = ( + getattr(raw_result, "stopped_by", "") + or "unknown" + ) + stub_text = ( + final_text + or f"[task ended without a clean final answer " + f"(stopped_by={stopped})]" + ) + session.messages.append(assistant_msg(stub_text)) + + # Close the boundary for this task. + end_idx = len(session.messages) - 1 + session.task_boundaries[-1] = (start_idx, end_idx) + # Fall back to the previous last_report when the new + # run produced nothing — keeps cross-agent + # references stable mid-session. + session.last_report = ( + final_text + or session.last_report + ) + + # Bubble evidence + assertions harvested by observers up + # to the caller so collect_reports / main_agent_node + # can surface them to the report node + frontend. + if active_runtime.result_adapter is not None: + adapted = active_runtime.result_adapter( + raw_result, job_id, session_item, + ) + if inspect.isawaitable(adapted): + adapted = await adapted + return adapted + return _adapt_default_session_result( + raw_result, job_id, session, task_prompt, + ) + except asyncio.CancelledError: + _close_session_boundary_aborted(session) + return SubAgentResult( + question=task_prompt, + role_id=session.role_id, + final_content=_safe_final_content(raw_result), + success=False, + error="aborted", + error_class="CancelledError", + job_id=job_id, + metadata=_safe_metadata(raw_result), + ) + except TimeoutError: + logger.warning("Session task %s timed out", job_id) + _close_session_boundary_aborted(session) + return SubAgentResult( + question=task_prompt, + role_id=session.role_id, + final_content=_safe_final_content(raw_result), + success=False, + error="timeout", + error_class="TimeoutError", + job_id=job_id, + metadata=_safe_metadata(raw_result), + ) + except Exception as exc: + logger.warning( + "Session task %s failed: %s", job_id, exc, + ) + _close_session_boundary_aborted(session) + return SubAgentResult( + question=task_prompt, + role_id=session.role_id, + final_content=_safe_final_content(raw_result), + success=False, + error=str(exc), + error_class=type(exc).__name__, + job_id=job_id, + metadata=_safe_metadata(raw_result), + ) + + async def _run_and_finalize() -> SubAgentResult: + entry = self._jobs[job_id] + # Layer 4: Concurrency gate (may queue when max_parallel reached). + if guard is not None: + await guard.acquire_slot(job_id) + entry.status = "running" + try: + result = await _run() + entry.result = result + entry.completed_at = time.monotonic() + if result.error == "aborted": + entry.status = "aborted" + elif result.success: + entry.status = "completed" + else: + entry.status = "failed" + if pending.task_metadata.get("can_publish") is True: + stopped_by = str(result.metadata.get("stopped_by") or "") + tool_calls = int(result.metadata.get("tool_calls_count", 0) or 0) + # Any ``no_tool`` ending earns a credit, not just a session + # with zero tool calls. The common degenerate shape is a + # publisher that runs a few commands and THEN emits its + # write as a fenced block — tool_calls_count > 0, yet the + # deliverable was never written, which is exactly the + # sticky-cap deadlock the credit exists to break. + if stopped_by == "no_tool": + session.publisher_no_tool_failures += 1 + elif tool_calls > 0: + # Session demonstrated working tool use: drop the + # accumulated credit so it does not permanently inflate + # this session's task cap for the rest of the run. + session.publisher_no_tool_failures = 0 + # Enqueue for wait_any_session regardless of success/failure + # so callers can observe errors (matching AgentTeam behavior + # where failed agents also emit a ). + session.pending_results.append(result) + await _emit_session_task_completed( + session, job_id, result, + event_sink=self._event_sink_injected, + spawn_context=pending.spawn_context, + ) + return result + finally: + # Layer 1: RAII release — always free the guard slot, even + # if the inner loop crashed before entry.status was set. + if guard is not None: + guard.release(job_id) + if entry.submitted_at: + completed_at = entry.completed_at or time.monotonic() + session.last_task_elapsed_s = max( + 0.0, completed_at - entry.submitted_at, + ) + if session.current_job_id == job_id: + session.current_job_id = None + # Eager trim: compress completed boundaries before the next + # task starts reading session.messages. Must be awaited here + # (not fire-and-forget) to avoid a race with _drain_session_queue + # reading stale messages for initial_messages construction. + try: + await self._eager_trim_and_offload(session) + except Exception: + logger.exception( + "eager_trim failed for session %s; continuing", + session.session_id, + ) + # Hand off to the next queued task (if any) so the + # session continues draining without the main agent + # having to re-call ``assign_task``. Failures are logged + # but never re-raised — a broken successor must not + # corrupt the just-finalised result. + try: + await self._drain_session_queue(session) + except Exception as exc: + logger.warning( + "Session %s: queue drain after %s failed: %s", + session.session_id, job_id, exc, + ) + + # Either upgrade an already-registered queued JobEntry or create + # a fresh one (free path). + existing = self._jobs.get(job_id) + if existing is not None: + existing.status = "submitted" + existing.submitted_at = existing.submitted_at or time.monotonic() + entry = existing + else: + entry = JobEntry( + job_id=job_id, + parent_task_id=session.task_id, + item=SubTask( + question=task_prompt, role_id=session.role_id, + system_prompt=session.system_prompt, + metadata=dict(pending.task_metadata), + ), + task=None, + status="submitted", + submitted_at=time.monotonic(), + ) + self._jobs[job_id] = entry + session.current_job_id = job_id + try: + task = asyncio.create_task( + _run_and_finalize(), name=f"session:{session_id}:{job_id}", + ) + except BaseException: + if session.current_job_id == job_id: + session.current_job_id = None + if existing is None: + self._jobs.pop(job_id, None) + raise + entry.task = task + # L1-0: remember who this job's result message will address, so the + # reconciliation read stays a cursor consume rather than a log scan. + self._note_result_recipient(session.task_id, pending.spawn_context) + + await _emit_session_task_submitted( + session, job_id, task_prompt, + event_sink=self._event_sink_injected, + ) + + async def _eager_trim_and_offload(self, session: SubAgentSession) -> None: + """Compress completed boundaries immediately after a task finishes. + + Must be awaited before ``_drain_session_queue`` to prevent a race + where the next queued task reads stale (untrimmed) messages as its + initial_messages. For single-task specialist sessions, this reclaims + the full intermediate history (tool calls + results) right after the + task completes, leaving only [task_prompt, final_ai]. + """ + new_messages, new_boundaries = trim_and_remap_boundaries( + session.messages, session.task_boundaries, + ) + if new_messages is session.messages: + # trim_and_remap_boundaries returns the same object when there are + # no completed boundaries — nothing to do. + return + + old_len = len(session.messages) + if self._session_history_dir is not None: + kept_ids = {id(m) for m in new_messages} + dropped = [m for m in session.messages if id(m) not in kept_ids] + await asyncio.to_thread( + _offload_dropped_messages, + session, + dropped, + self._session_history_dir, + ) + + session.messages = new_messages + session.task_boundaries = new_boundaries + logger.debug( + "eager_trim: session %s messages %d → %d (dropped %d)", + session.session_id, + old_len, + len(new_messages), + old_len - len(new_messages), + ) + + async def _drain_session_queue( + self, + session: SubAgentSession, + ) -> None: + """Pop and dispatch the next queued task on this session, if any. + + Called from ``_run_and_finalize`` after the current task clears + ``current_job_id``. Only dispatches one task at a time — the + next queued task will trigger the same drain when *it* finishes, + keeping execution strictly serial. + """ + if not session.pending_tasks: + return + next_pending = session.pending_tasks.popleft() + # If the queued JobEntry was aborted while parked, skip it and + # try the next one. ``abort_job`` removes the queued PendingTask + # before flipping status, but cancel_for_parent doesn't touch + # the queue first; defensive check covers both. + entry = self._jobs.get(next_pending.job_id) + if entry is not None and entry.status == "aborted": + await self._drain_session_queue(session) + return + await self._dispatch_session_task(session, next_pending) + + async def _reconcile_terminal_session_jobs(self, task_id: str) -> None: + """Turn terminal tasks that never published a result into failures. + + A task can be cancelled before ``_run_and_finalize`` enters its + ``try/finally`` block (or fail in wrapper setup). In that case the + asyncio task is done, but ``current_job_id`` and the job status remain + ``submitted``/``running`` forever. The old wait path treated that + stale state as a timeout and returned immediately, making callers say + they had waited 30 minutes when only milliseconds elapsed. + """ + sessions_to_drain: list[SubAgentSession] = [] + for session in self._sessions.values(): + if session.task_id != task_id or session.current_job_id is None: + continue + job_id = session.current_job_id + entry = self._jobs.get(job_id) + task = entry.task if entry is not None else None + if entry is None or task is None or not task.done(): + continue + + result = entry.result + if result is None and not task.cancelled(): + try: + task_result = task.result() + except (asyncio.CancelledError, Exception) as exc: + error = str(exc) or type(exc).__name__ + error_class = type(exc).__name__ + else: + if isinstance(task_result, SubAgentResult): + result = task_result + error = "" + error_class = "" + else: + error = "sub-agent task ended without a report" + error_class = "MissingSubAgentResult" + elif result is None: + error = "sub-agent task was cancelled before publishing a report" + error_class = "CancelledError" + else: + error = "" + error_class = "" + + if result is None: + result = SubAgentResult( + question=str(getattr(entry.item, "question", "") or ""), + role_id=str(getattr(entry.item, "role_id", "") or ""), + final_content="", + success=False, + error=error, + error_class=error_class, + job_id=job_id, + ) + _close_session_boundary_aborted(session) + + entry.result = result + entry.completed_at = entry.completed_at or time.monotonic() + entry.status = ( + "completed" if result.success + else "aborted" if result.error_class == "CancelledError" + else "failed" + ) + if not any(item.job_id == job_id for item in session.pending_results): + session.pending_results.append(result) + if entry.submitted_at: + session.last_task_elapsed_s = max( + 0.0, entry.completed_at - entry.submitted_at, + ) + if session.current_job_id == job_id: + session.current_job_id = None + if self._spawn_guard is not None: + self._spawn_guard.release(job_id) + sessions_to_drain.append(session) + logger.warning( + "Reconciled terminal sub-agent job without a published report: " + "%s (%s)", job_id, result.error_class or entry.status, + ) + + for session in sessions_to_drain: + await self._drain_session_queue(session) + + + async def wait_any_session_detailed( + self, task_id: str, *, timeout: float = 1800.0, + ) -> SessionWaitOutcome: + """Wait for one result and preserve why the wait ended.""" + started = time.monotonic() + ready = self._pop_ready_result(task_id) + if ready is not None: + return SessionWaitOutcome(ready, "ready", time.monotonic() - started) + + await self._reconcile_terminal_session_jobs(task_id) + ready = self._pop_ready_result(task_id) + if ready is not None: + return SessionWaitOutcome(ready, "ready", time.monotonic() - started) + + pending: set[asyncio.Task[SubAgentResult]] = set() + for session in self._sessions.values(): + if session.task_id != task_id or session.current_job_id is None: + continue + entry = self._jobs.get(session.current_job_id) + if ( + entry is not None + and entry.status in ("submitted", "running") + and entry.task is not None + and not entry.task.done() + ): + pending.add(entry.task) + + if not pending: + recovered = await self._reconcile_from_messages(task_id) + reason = "ready" if recovered is not None else "no_pending" + return SessionWaitOutcome( + recovered, reason, time.monotonic() - started, + ) + + done, _ = await asyncio.wait( + pending, + timeout=max(0.0, float(timeout)), + return_when=asyncio.FIRST_COMPLETED, + ) + if done: + await self._reconcile_terminal_session_jobs(task_id) + ready = self._pop_ready_result(task_id) + if ready is None: + ready = await self._reconcile_from_messages(task_id) + if ready is not None: + return SessionWaitOutcome( + ready, "ready", time.monotonic() - started, + ) + return SessionWaitOutcome( + None, + "unpublished" if done else "timeout", + time.monotonic() - started, + ) + + async def wait_any_session( + self, task_id: str, *, timeout: float = 1800.0, + ) -> tuple[str, SubAgentResult] | None: + """Return the next completed task, or ``None`` when no result arrives.""" + outcome = await self.wait_any_session_detailed(task_id, timeout=timeout) + return outcome.result + + def _pop_ready_result( + self, task_id: str, + ) -> tuple[str, SubAgentResult] | None: + for session in self._sessions.values(): + if session.task_id != task_id: + continue + if session.pending_results: + result = session.pending_results.pop(0) + self._mark_delivered(task_id, result.job_id) + return session.session_id, result + return None + + # ── L1-0: message-channel reconciliation ──────────────────────────── + + def _mark_delivered(self, task_id: str, job_id: str) -> None: + """Record that ``job_id``'s result has been handed to a caller.""" + if job_id: + self._delivered_jobs.setdefault(task_id, set()).add(job_id) + + def _note_result_recipient( + self, task_id: str, spawn_context: dict[str, Any] | None, + ) -> None: + """Remember who ``subtask_result`` messages for this task address.""" + if not spawn_context: + return + parent = str(spawn_context.get("parent_agent_id") or "").strip() + if parent: + self._result_recipients.setdefault(task_id, set()).add(parent) + + async def _discover_result_recipients(self, task_id: str) -> set[str]: + """Recipients for ``task_id``, from memory or else from the log. + + In-process the dispatch path has already recorded them. After a + restart it has not, and the only surviving record is the messages + themselves — their ``to_agent`` is the recipient, so one scan of + the task's message events recovers the set. + """ + known = self._result_recipients.get(task_id) + if known: + return set(known) + sink = self._event_sink() + if sink is None or not hasattr(sink, "get_events"): + return set() + try: + events = await sink.get_events( + task_id, event_type=EventType.AGENT_MESSAGE, + ) + except Exception: + logger.debug("recipient discovery failed for %s", task_id, exc_info=True) + return set() + found = { + str(event.to_agent) + for event in events + if getattr(event, "message_type", None) == SUBTASK_RESULT_MESSAGE_TYPE + and getattr(event, "to_agent", None) + } + if found: + self._result_recipients.setdefault(task_id, set()).update(found) + return found + + async def _reconcile_from_messages( + self, task_id: str, + ) -> tuple[str, SubAgentResult] | None: + """Recover one undelivered result from the message channel. + + Reads through ``AgentComm``'s type-isolated + ``(agent, task, subtask_result)`` cursor, so repeated calls do not + re-read the same results and unrelated status/control messages remain + available to their own consumers. Skips job ids the in-memory path + already delivered. Returns the same + ``(session_id, result)`` shape as the in-memory path — callers and + the ```` rendering downstream cannot tell the difference. + + One consume call drains a whole batch and moves the cursor past all + of it, so every usable message is buffered and one is returned per + call; otherwise a batch of N would deliver 1 and lose N-1. + + Entirely soft: no ``AgentComm``, no messages, or any read failure + all fall back to today's ``None``. + """ + buffered = self._recovered_results.get(task_id) + if buffered: + return buffered.popleft() + + # Lazy import: ``agent_comm`` is a sibling module and this keeps the + # bus importable without it, matching how ``runtime.py`` reaches for + # the same class. + from agent_core.components.agent_bus.agent_comm import AgentComm + + comm = registry.get_optional(AgentComm) + if comm is None: + return None + recipients = await self._discover_result_recipients(task_id) + if not recipients: + return None + delivered = self._delivered_jobs.get(task_id, set()) + pending_recovered = self._recovered_results.setdefault(task_id, deque()) + for agent_id in sorted(recipients): + try: + events = await comm.consume( + agent_id, + task_id=task_id, + message_type=SUBTASK_RESULT_MESSAGE_TYPE, + ) + except Exception: + logger.debug( + "message reconciliation failed for %s/%s", + task_id, agent_id, exc_info=True, + ) + continue + for event in events: + payload = event.payload or {} + if payload.get("message_type") != SUBTASK_RESULT_MESSAGE_TYPE: + continue + content = payload.get("content") or {} + job_id = str(content.get("correlation_id") or "") + if not job_id or job_id in delivered: + continue + raw = content.get("result") + if not isinstance(raw, dict): + continue + self._mark_delivered(task_id, job_id) + delivered = self._delivered_jobs.get(task_id, set()) + pending_recovered.append(( + str(content.get("session_id") or ""), + _message_content_to_result(raw), + )) + logger.info( + "Recovered sub-agent result %s from the message channel", + job_id, + ) + return pending_recovered.popleft() if pending_recovered else None + + def get_session(self, session_id: str) -> SubAgentSession | None: + return self._sessions.get(session_id) + + def list_sessions_for_task(self, task_id: str) -> list[SubAgentSession]: + return [s for s in self._sessions.values() if s.task_id == task_id] + + def current_job_metadata(self, session_id: str) -> dict[str, Any]: + """Task metadata of the job this session is running now. + + Empty when the session is unknown, idle, or its job entry has already + been reaped. Callers use it to tell *what kind* of work a stop or a + cancellation would interrupt — ``can_publish``, for one, marks the job + that produces the run's deliverable. A copy, so a caller inspecting a + live job cannot mutate the dispatched task's own metadata. + """ + session = self._sessions.get(session_id) + if session is None or session.current_job_id is None: + return {} + entry = self._jobs.get(session.current_job_id) + if entry is None: + return {} + return dict(getattr(entry.item, "metadata", None) or {}) + + def describe_sessions_for_task(self, task_id: str) -> list[dict[str, Any]]: + """Return a small, UI-safe snapshot of every sub-agent session. + + The snapshot deliberately excludes prompts, reports, and model data; + it is suitable for frequent progress rendering while collect_reports + blocks. Durations use the job's monotonic submission clock. + """ + now = time.monotonic() + snapshots: list[dict[str, Any]] = [] + for session in self.list_sessions_for_task(task_id): + entry = ( + self._jobs.get(session.current_job_id) + if session.current_job_id is not None else None + ) + if entry is not None: + status = entry.status + active = entry.status in ("queued", "submitted", "running") + # A terminal entry can still be the session's ``current_job_id`` + # for the moment ``_run_and_finalize`` spends in its finally + # block; measure it to ``completed_at`` so that window does not + # publish a duration that keeps growing. + until = now if active else (entry.completed_at or now) + elapsed = ( + max(0.0, until - entry.submitted_at) + if entry.submitted_at else 0.0 + ) + else: + active = False + # ``last_task_elapsed_s`` describes work that finished, so it + # only belongs to the states that mean "finished". A session + # merely holding a queue has not started that task yet. + elapsed = 0.0 + if session.pending_results: + status = "ready" + elapsed = session.last_task_elapsed_s + elif session.pending_tasks: + status = "queued" + elif session.total_task_count: + status = "idle" + elapsed = session.last_task_elapsed_s + else: + status = "unassigned" + snapshots.append({ + "session_id": session.session_id, + "name": session.name, + "role_id": session.role_id, + "status": status, + # ``active`` says whether ``elapsed_s`` is still counting. + # Without it a consumer cannot tell "ran for 12s and stopped" + # from "has been running 12s", and re-deriving the duration + # from its own clock makes finished workers tick forever. + "active": active, + "elapsed_s": elapsed, + "queued": len(session.pending_tasks), + "completed": len(session.pending_results), + "events": [dict(event) for event in session.activity_events], + }) + return snapshots + + + def accumulate_task_metadata( + self, + task_id: str, + **payload: list[dict[str, Any]] | None, + ) -> None: + """Merge named list-payloads into the task's metadata pool. + + Each keyword argument names a bucket (e.g. ``evidence_cards=[...], + assertions=[...]``). Empty or ``None`` payloads are skipped so + callers can pass optional harvests without guarding each one. + + The pool is drained via ``drain_task_metadata`` when the caller + (typically a pipeline node) wants to hand the corpus downstream. + """ + merged = {k: v for k, v in payload.items() if v} + if not merged: + return + bucket = self._task_aggregates.setdefault(task_id, {}) + for key, items in merged.items(): + bucket.setdefault(key, []).extend(items) + + def drain_task_metadata( + self, task_id: str, + ) -> dict[str, list[dict[str, Any]]]: + """Return + clear accumulated metadata lists for a task. + + Returns an empty dict if nothing was accumulated; callers that + expect specific keys should use ``.get(key, [])``. + """ + return self._task_aggregates.pop(task_id, {}) + + async def cleanup_task( + self, task_id: str, *, cancel_timeout_s: float = 10.0, + ) -> int: + """Cancel all running session tasks under ``task_id`` and drop sessions. + + If a sub-agent's asyncio.Task cannot be cancelled within + ``cancel_timeout_s`` (e.g. blocked in a non-cancellable subprocess + or a long socket read), we detach it, log a warning, and move on. + That prevents the pipeline's post-loop cleanup from hanging + indefinitely on stuck sub-agents, which previously caused the + ``main_agent → report`` transition to stall until the eval + driver killed the task. + + Returns the number of sessions cleaned up. Also clears any residual + aggregate pool so re-runs don't carry over state. + + Rescue: ``session.pending_results`` are merged into + ``_task_aggregates`` via ``accumulate_task_metadata`` before + sessions are dropped, so a post-cleanup ``drain_task_metadata`` + still surfaces sub-agent reports the main agent never claimed + (e.g. ``force_final_answer`` short-circuited a slow sub-agent). + """ + self._task_aggregates.pop(task_id, None) + + rescued_evidence: list[dict[str, Any]] = [] + rescued_assertions: list[dict[str, Any]] = [] + + def rescue_pending_results(session: SubAgentSession) -> None: + while session.pending_results: + result = session.pending_results.pop(0) + md = result.metadata or {} + rescued_evidence.extend(md.get("evidence_cards") or []) + rescued_assertions.extend(md.get("assertions") or []) + + def abort_pending_tasks(session: SubAgentSession) -> None: + while session.pending_tasks: + pending = session.pending_tasks.popleft() + self._mark_job_aborted( + pending.job_id, + self._jobs.get(pending.job_id), + ) + + cleaned = 0 + stuck: list[str] = [] + for sid in list(self._sessions): + session = self._sessions.get(sid) + if session is None or session.task_id != task_id: + continue + rescue_pending_results(session) + # Prevent the running task's ``finally`` from draining queued + # successors while cleanup is tearing this session down. + abort_pending_tasks(session) + if session.current_job_id is not None: + current_job_id = session.current_job_id + entry = self._jobs.get(session.current_job_id) + if entry is not None and entry.status in ( + "submitted", "running", + ) and entry.task is not None: + entry.task.cancel() + try: + await asyncio.wait_for( + entry.task, timeout=cancel_timeout_s, + ) + except (TimeoutError, asyncio.CancelledError): + if not entry.task.done(): + stuck.append(current_job_id) + except Exception: + pass + if entry.status not in ("completed", "failed", "aborted"): + self._mark_job_aborted(current_job_id, entry) + elif entry is not None and entry.status not in ( + "completed", "failed", "aborted", + ): + self._mark_job_aborted(current_job_id, entry) + # A cancelled job can catch CancelledError, build a + # SubAgentResult, and append it to pending_results while + # cleanup is awaiting entry.task. Rescue again immediately + # before dropping the session so that metadata is not lost. + rescue_pending_results(session) + del self._sessions[sid] + cleaned += 1 + if rescued_evidence or rescued_assertions: + self.accumulate_task_metadata( + task_id, + evidence_cards=rescued_evidence, + assertions=rescued_assertions, + ) + logger.info( + "cleanup_task: rescued %d evidence + %d assertions from " + "unconsumed pending_results (task=%s)", + len(rescued_evidence), len(rescued_assertions), task_id, + ) + if cleaned: + logger.info( + "cleanup_task: dropped %d sessions for %s", cleaned, task_id, + ) + if stuck: + logger.warning( + "cleanup_task: %d session job(s) did not cancel within %ss — " + "detaching so pipeline can progress: %s", + len(stuck), cancel_timeout_s, stuck, + ) + # These maps are task-scoped runtime acceleration only. cleanup_task + # is the terminal lifecycle boundary used by the SDK/API workflows; + # retaining their keys would make a long-lived AgentBus grow once per + # completed task. Durable messages remain in EventStore and can still + # be replayed by a fresh runtime after a process restart. + self._delivered_jobs.pop(task_id, None) + self._result_recipients.pop(task_id, None) + self._recovered_results.pop(task_id, None) + return cleaned + + # ── dispatch_parallel: backward-compatible convenience ────────────── + + async def dispatch_parallel( + self, + parent_task_id: str, + items: list[SubTask], + shared_evidence: SharedArtifactPool | None = None, + max_turns: int = 8, + current_depth: int = 0, + max_depth: int = 2, + runtime_spec: SubAgentRuntimeSpec | None = None, + ) -> list[SubAgentResult]: + """Submit all items and collect all results. + + Backward-compatible wrapper around submit + collect. + Existing callers (fan_out_node, tests) work unchanged. + """ + if current_depth >= max_depth: + raise DepthLimitExceeded( + f"Sub-agent depth limit exceeded: " + f"current_depth={current_depth}, " + f"max_depth={max_depth}" + ) + + if not items: + return [] + + evidence_pool = shared_evidence or SharedArtifactPool() + + event_store = self._event_sink() + t0 = time.monotonic() + + if event_store: + await event_store.append( + task_id=parent_task_id, + event_type=EventType.AGENT_ACTION, + payload={ + "trace_type": "agent_delegated", + "parent_task_id": parent_task_id, + "subtasks": [ + {"question": item.question, "role": item.role_id} + for item in items + ], + "depth": current_depth + 1, + }, + agent_role="system", + ) + + # Submit all + job_ids: list[str] = [] + for item in items: + jid = await self.submit( + parent_task_id, + item, + shared_evidence=evidence_pool, + max_turns=max_turns, + current_depth=current_depth, + max_depth=max_depth, + runtime_spec=runtime_spec, + ) + job_ids.append(jid) + + collect_result = await self.collect(job_ids, timeout=1800) + + if event_store: + all_results = collect_result.completed + collect_result.failed + await event_store.append( + task_id=parent_task_id, + event_type=EventType.AGENT_ACTION, + payload={ + "trace_type": "agent_completed", + "parent_task_id": parent_task_id, + "results_count": len(all_results), + "successful_results": sum( + 1 for r in all_results if r.success + ), + "total_evidence": evidence_pool.count(), + "elapsed_ms": int( + (time.monotonic() - t0) * 1000 + ), + }, + agent_role="system", + ) + + return collect_result.completed + collect_result.failed + + # ── rehydration: recover from restart ─────────────────────────────── + + async def rehydrate_active_jobs(self) -> int: + """Mark jobs whose asyncio.Task was lost across a restart as orphaned.""" + return await _rehydrate_orphaned_jobs( + self._jobs, event_sink=self._event_sink_injected, + ) diff --git a/agent_core/components/agent_bus/fan_in.py b/agent_core/components/agent_bus/fan_in.py new file mode 100644 index 0000000..d68a0c8 --- /dev/null +++ b/agent_core/components/agent_bus/fan_in.py @@ -0,0 +1,431 @@ +# pyright: reportUnknownVariableType=false +"""Shared report-formatting helpers for sub-agent fan-in paths.""" + +from __future__ import annotations + +import re +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal + +from agent_core.components.agent_bus.models import SubAgentResult + +if TYPE_CHECKING: + from agent_core.components.agent_bus.bus import AgentBus + + +# Stop reasons that mean "agent finished cleanly via a terminal tool". +# Anything outside this set + INCOMPLETE_STOP_REASONS + ``"paused"`` +# falls through to a generic ``incomplete`` label with the raw +# stopped_by value (sanitized) as the ``reason`` attribute, which keeps +# unknown observer-driven stops from being silently mislabelled as +# complete. +COMPLETE_STOP_REASONS: frozenset[str] = frozenset({ + "", + "completed", + "final_answer", + "submit_report", +}) + +# Stop reasons that mean "agent was truncated mid-task" — the report +# may be partial. ``force_final_answer()`` may rescue a best-effort +# plain-text conclusion, but the structural status is still +# incomplete and the main agent should treat it as such. +# +# This set is the curated label table consumed by +# :func:`classify_completion` (status="incomplete" + a specific NOTE), AND +# the trigger list for ``force_final_answer()``: both rescue gates in +# ``subagent_runtime`` read ``stopped_by not in INCOMPLETE_STOP_REASONS`` +# and return early, so it is an allowlist. A new observer-driven stop +# reason that is not added here still classifies as ``incomplete`` (via the +# catch-all at the bottom of :func:`classify_completion`) but silently +# loses its forced-final answer, which is the whole point of stopping a +# looping agent early. Add the reason here when you add the stop. +INCOMPLETE_STOP_REASONS: frozenset[str] = frozenset({ + "max_turns", + "max_attempts", + "llm_error", + "no_tool", + "budget_exhausted", + # Ran out of wall clock rather than tokens — emitted by WallClockGuard + # (a sub-agent's SpawnGuard slot), WallClockDeadlineObserver (the whole + # run's deadline), and agent_loop's mid-turn wall refusal. + "wall_deadline", + "context_limit_reached", + "cross_turn_repetition", + "repeated_tool_calls", + # The output cap cut every continuation off mid-sentence. Distinct from + # ``no_tool``: the agent never chose to stop, so its report is unfinished + # rather than merely answer-less. + "response_truncated", + "exception", +}) + +_INCOMPLETE_NOTES: dict[str, str] = { + "max_turns": "agent reached the max_turns budget; report is partial", + "max_attempts": ( + "agent exceeded the per-turn LLM retry budget; report is partial" + ), + "llm_error": "sub-agent's LLM call failed; report is partial", + "no_tool": "agent stopped without producing a final answer", + "response_truncated": ( + "agent's replies kept hitting the output token limit; report is partial" + ), + "budget_exhausted": "agent exhausted its token budget; report is partial", + "wall_deadline": ( + "agent ran out of wall-clock time; report is partial" + ), + "context_limit_reached": ( + "agent ran out of model context window mid-task; report is partial" + ), + "cross_turn_repetition": ( + "agent stopped after repeating itself across turns; " + "report is a best-effort partial" + ), + "repeated_tool_calls": ( + "agent stopped after re-issuing the same tool call with identical " + "arguments; report is a best-effort partial" + ), + "exception": ( + "agent terminated with an unhandled exception; report is partial" + ), +} + + +CompletionStatus = Literal["complete", "incomplete", "failed", "paused"] +SessionState = Literal[ + "no_subagents", + "ready_to_collect", + "running", + "no_work_queued", + "all_collected", +] + + +# Attribute values must be stable short tokens: XML-attribute-safe, +# log-grep-safe, prompt-template-safe. Raw error strings (with quotes, +# newlines, angle brackets, …) belong in the NOTE / body, never in the +# ``reason`` attribute. +_ATTR_TOKEN_RE = re.compile(r"[^a-z0-9_.-]+") + + +@dataclass(frozen=True) +class CompletionPolicy: + """Workflow-specific stop-reason labels for the shared fan-in mechanics.""" + + complete_reasons: frozenset[str] + incomplete_notes: Mapping[str, str] + include_error_class: bool = True + + +DEFAULT_COMPLETION_POLICY = CompletionPolicy( + complete_reasons=COMPLETE_STOP_REASONS, + incomplete_notes=_INCOMPLETE_NOTES, +) + + +def _safe_reason(value: str) -> str: + """Sanitize an arbitrary string into an XML-attribute-safe token.""" + token = (value or "unknown").strip().lower() + token = _ATTR_TOKEN_RE.sub("_", token).strip("_") + return token[:80] or "unknown" + + +@dataclass(frozen=True) +class CompletionInfo: + """Translation of sub-agent stop reason into a UI-facing label. + + ``reason`` is ``""`` for clean completion; otherwise a sanitized + short tag (``"max_turns"``, ``"timeout"``, …). ``note`` is rendered + as ``[NOTE: .]`` at the top of the report body when non-empty. + """ + + status: CompletionStatus + reason: str + note: str + + +def classify_completion( + result: SubAgentResult, + *, + policy: CompletionPolicy = DEFAULT_COMPLETION_POLICY, +) -> CompletionInfo: + """Inspect ``result`` and return how the report should be labelled. + + Decision tree: + + 1. ``success=False`` → ``failed`` + 2. ``stopped_by="paused"`` → ``paused`` + 3. ``stopped_by`` in :data:`INCOMPLETE_STOP_REASONS` → ``incomplete`` + 4. ``stopped_by`` in :data:`COMPLETE_STOP_REASONS` → ``complete`` + 5. **Anything else** → ``incomplete`` with the sanitized stop + reason. Catch-all for observer-driven stops we don't have a + curated note for; silently labelling them ``complete`` would be + exactly the F2 failure mode this helper exists to prevent. + """ + md = result.metadata or {} + stopped_by = (md.get("stopped_by") or "").strip() + + if not result.success: + err = (result.error or "unknown error").strip() + cls = (result.error_class or "").strip() if policy.include_error_class else "" + display = f"{cls}: {err}" if cls else err + return CompletionInfo( + status="failed", + reason=_safe_reason(err), + note=f"agent failed mid-task ({display}); output is partial", + ) + + if stopped_by == "paused": + return CompletionInfo( + status="paused", + reason="paused", + note="agent paused (resumable from checkpoint)", + ) + + if stopped_by in policy.incomplete_notes: + return CompletionInfo( + status="incomplete", + reason=stopped_by, + note=policy.incomplete_notes[stopped_by], + ) + + if stopped_by in policy.complete_reasons: + return CompletionInfo(status="complete", reason="", note="") + + return CompletionInfo( + status="incomplete", + reason=_safe_reason(stopped_by), + note=( + f"agent stopped by observer reason `{stopped_by}`; " + "report may be partial" + ), + ) + + +def format_report_block( + name: str, + result: SubAgentResult, + info: CompletionInfo | None = None, + *, + policy: CompletionPolicy = DEFAULT_COMPLETION_POLICY, +) -> str: + """Render one ```` block. + + Pass a pre-computed ``info`` to skip a redundant + :func:`classify_completion` call when the caller already inspected + the status (the fan-in path needs both the rendered block and the + status bucket). + """ + if info is None: + info = classify_completion(result, policy=policy) + body = (result.final_content or "(empty report)").strip() + if info.note: + body = f"[NOTE: {info.note}.]\n{body}" + if info.reason: + return ( + f'\n{body}\n' + ) + return f'\n{body}\n' + + +# Synthetic agent name for orchestrator-level status notices. The +# status short-circuits in ``collect_reports`` (all_collected / +# no_work_queued / …) used to return bare ``[status] …`` text; SDK +# consumers that parse ``tool_finished.result_preview`` for +# ```` blocks choked on those payloads and pasted them back to +# end users verbatim with a parse-failure marker. Wrapping the notice +# in the same envelope keeps every collect_reports return parseable by +# a single code path. +ORCHESTRATOR_AGENT_NAME = "orchestrator" + + +def format_status_report_block(reason: str, body: str) -> str: + """Render an orchestrator status notice as a ```` block. + + Same envelope as :func:`format_report_block` so both consumers — + the main agent reading the tool result and SDK callers parsing + ``tool_finished.result_preview`` — handle status-only returns with + the report-block code path instead of free-form text. + ``status="complete"`` keeps the attribute within the documented + vocabulary (``complete|incomplete|failed|paused``); the actual + state token travels in ``reason``. + """ + return ( + f'\n{body.strip()}\n' + ) + + +@dataclass +class FanInBatch: + """Aggregated outcome of draining a batch of sub-agent results. + + ``blocks`` are rendered ```` strings ready to join with + ``"\\n\\n"``. The other fields feed + :func:`format_status_line` (paused bucket, partial-output count) and + callers that want the cumulative evidence/assertion totals harvested + from this batch. + """ + + blocks: list[str] = field(default_factory=list) + paused_names: set[str] = field(default_factory=set) + incomplete_count: int = 0 + evidence_count: int = 0 + assertion_count: int = 0 + + +def process_collected( + bus: AgentBus, + task_id: str, + collected: list[tuple[str, SubAgentResult]], + *, + policy: CompletionPolicy = DEFAULT_COMPLETION_POLICY, +) -> FanInBatch: + """Single-pass fan-in: classify, render, harvest evidence/assertions. + + ``classify_completion`` runs exactly once per result; the same + :class:`CompletionInfo` is then threaded into + :func:`format_report_block` to avoid a redundant second pass. + + Evidence / assertions are harvested EVEN on failure — a sub-agent + that ran 20 web_searches then crashed in the result adapter still + contributed real evidence we don't want to silently drop. + """ + batch = FanInBatch() + for session_id, result in collected: + name = session_id.split("::", 1)[-1] + info = classify_completion(result, policy=policy) + if info.status == "paused": + batch.paused_names.add(name) + if info.status in ("incomplete", "failed"): + batch.incomplete_count += 1 + batch.blocks.append(format_report_block(name, result, info)) + + ev = list(result.metadata.get("evidence_cards", [])) + asserts = list(result.metadata.get("assertions", [])) + if ev or asserts: + bus.accumulate_task_metadata( + task_id, + evidence_cards=ev, + assertions=asserts, + ) + batch.evidence_count += len(ev) + batch.assertion_count += len(asserts) + + return batch + + +def session_state(sessions: list[Any]) -> SessionState: + """Classify the swarm's overall session state for ``collect_reports``. + + Used to give the main agent an actionable next-step hint instead of + a generic "wait more" line. + """ + if not sessions: + return "no_subagents" + if any(getattr(s, "pending_results", None) for s in sessions): + return "ready_to_collect" + # A queued task on any session counts as "running": the next + # ``wait_any_session`` call will eventually surface its report, + # so the main agent should keep waiting rather than concluding + # all_collected. + if any( + getattr(s, "current_job_id", None) is not None + or getattr(s, "pending_tasks", None) + for s in sessions + ): + return "running" + all_unassigned = all( + getattr(s, "total_task_count", 0) == 0 + and not getattr(s, "last_report", "") + for s in sessions + ) + if all_unassigned: + return "no_work_queued" + return "all_collected" + + +def format_status_line( + bus: AgentBus, + task_id: str, + *, + paused_names: set[str] | None = None, + incomplete_count: int = 0, +) -> str: + """One-line summary of sub-agent state for the main agent. + + Deliberately does NOT name idle sessions — idle here means + "already fanned in" or "freshly created, awaiting assignment", + neither of which is actionable. Showing them tends to mislead the + main agent into re-issuing tasks or polling unnecessarily. + + When no session is running, ready_to_collect, or paused, the line + collapses to ``no_work_queued`` (next step: ``assign_task``) or + ``all_collected`` (next step: synthesize / follow-up). + """ + paused_names = paused_names or set() + sessions = bus.list_sessions_for_task(task_id) + if not sessions: + return "[status] no sub-agents" + + running: list[str] = [] + ready: list[str] = [] + paused: list[str] = [] + total_task_count = 0 + has_last_report = False + + for s in sessions: + total_task_count += getattr(s, "total_task_count", 0) + if getattr(s, "last_report", ""): + has_last_report = True + if s.name in paused_names: + paused.append(s.name) + elif ( + getattr(s, "current_job_id", None) is not None + or getattr(s, "pending_tasks", None) + ): + queued = len(getattr(s, "pending_tasks", []) or []) + label = f"{s.name}+{queued}q" if queued else s.name + running.append(label) + elif getattr(s, "pending_results", None): + ready.append(s.name) + + bits: list[str] = [] + if running: + bits.append(f"running={running}") + if ready: + bits.append(f"ready_to_collect={ready}") + if paused: + bits.append(f"paused={paused}") + + if not bits: + if total_task_count == 0 and not has_last_report: + bits.append("no_work_queued") + else: + bits.append("all_collected") + + if incomplete_count: + bits.append(f"incomplete_this_batch={incomplete_count}") + + return "[status] " + " ".join(bits) + + +__all__ = [ + "COMPLETE_STOP_REASONS", + "DEFAULT_COMPLETION_POLICY", + "INCOMPLETE_STOP_REASONS", + "ORCHESTRATOR_AGENT_NAME", + "CompletionInfo", + "CompletionPolicy", + "CompletionStatus", + "FanInBatch", + "SessionState", + "classify_completion", + "format_report_block", + "format_status_line", + "format_status_report_block", + "process_collected", + "session_state", +] diff --git a/agent_core/components/agent_bus/models.py b/agent_core/components/agent_bus/models.py new file mode 100644 index 0000000..1124c07 --- /dev/null +++ b/agent_core/components/agent_bus/models.py @@ -0,0 +1,246 @@ +# pyright: reportUnknownVariableType=false +"""Data models for AgentBus. + +Keeping these models separate lets ``agent_bus.py`` focus on lifecycle logic +while preserving the same import surface for callers via re-export. +""" + +from __future__ import annotations + +import asyncio +from collections import deque +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal + +from agent_core.messages import Message +from agent_core.runtime.loop.message_trimmer import MessageTrimmer, TaskBoundary +from agent_core.tool import Tool + + +class DepthLimitExceeded(RuntimeError): + """Raised when a sub-agent dispatch would exceed the configured depth.""" + + +@dataclass(slots=True) +class SubTask: + """Single sub-agent task.""" + + question: str + role_id: str = "researcher" + system_prompt: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class SubAgentResult: + """Normalized result returned from a sub-agent execution. + + ``metadata`` is the workflow-agnostic artifact bag — runtime specs + and result adapters stash domain-specific payloads (research + evidence cards, assertions, tool trails, …) there. The kernel data + model itself carries no domain vocabulary; the legacy dedicated + ``evidence_cards`` / ``assertions`` fields were removed as part of + the thin-kernel migration. + """ + + question: str + role_id: str + final_content: str + success: bool + error: str | None = None + error_class: str | None = None + job_id: str = "" + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class JobEntry: + """Tracks a single async sub-agent job. + + ``task`` is ``None`` while the job is still queued — sessions execute + tasks strictly serially, so submitted tasks can sit behind a running + one without an asyncio.Task spawned yet. The ``"queued"`` status + marks that state; transitions to ``"submitted"`` (then ``"running"``) + when the bus dequeues and dispatches it. + """ + + job_id: str + parent_task_id: str + item: SubTask + task: asyncio.Task[SubAgentResult] | None = None + status: Literal[ + "queued", "submitted", "running", "completed", "failed", "aborted" + ] = "submitted" + result: SubAgentResult | None = None + submitted_at: float = 0.0 + completed_at: float | None = None + + +@dataclass(slots=True) +class PendingSessionTask: + """A session task waiting in the queue behind the running one. + + Holds every input ``submit_task_to_session`` was called with so the + bus can dispatch the queued task identically once the previous one + finishes — same boundary semantics, same runtime spec, same guard + bookkeeping. Created at submit time, drained FIFO when + ``session.current_job_id`` clears. + """ + + job_id: str + task_prompt: str + max_turns: int | None = None + observers: list[Any] | None = None + estimated_tokens: int = 0 + runtime_spec: SubAgentRuntimeSpec | None = None + # ``spawn_context`` carries the parent's run_id, the verbatim delegation + # prompt, allowed tools snapshot, depth, and budget. Bus stamps it onto + # the sub-agent's trace observer at dispatch time so a downstream + # training-data pipeline can reconstruct the delegation graph + # (parent → child + the exact instruction). + spawn_context: dict[str, Any] | None = None + task_metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class CollectResult: + """Aggregated result from collect().""" + + completed: list[SubAgentResult] = field(default_factory=list) + pending: list[str] = field(default_factory=list) + failed: list[SubAgentResult] = field(default_factory=list) + + +@dataclass(slots=True) +class SessionWaitOutcome: + """Detailed outcome of waiting for one reusable sub-agent session. + + ``wait_any_session`` keeps its historical ``result | None`` API, while + orchestration tools use this richer form so ``None`` is not automatically + (and incorrectly) described as a full timeout. + + ``reason`` distinguishes the three empty-handed outcomes, which need + different advice for the coordinator: + + ``no_pending`` + Nothing was waitable when the wait began — no time passed, so the + requested timeout must not be reported as elapsed. + ``timeout`` + A live task was waited on and did not finish within the budget. + ``unpublished`` + A task *did* finish during the wait but produced no collectable + result. Real time elapsed here, so this is not a ``no_pending``. + """ + + result: tuple[str, SubAgentResult] | None + reason: Literal["ready", "timeout", "no_pending", "unpublished"] + elapsed_s: float + + +@dataclass(slots=True) +class SubAgentRuntimeSpec: + """Optional runtime injection for generic sub-agent execution.""" + + config_builder: Callable[[str, SubTask, int], Any] | None = None + # ``task_index`` is the 1-based ordinal of this task within the + # session (== ``session.total_task_count`` at submit time). Lets + # observers tell first-run from reuse without consulting the bus. + observers_builder: Callable[ + [str, SubTask, int], list[Any] + ] | None = None + # Sync or async — bus awaits when the return is a coroutine. + result_adapter: Callable[ + [Any, str, SubTask], + SubAgentResult | Awaitable[SubAgentResult], + ] | None = None + model_profile: Any = None + history_policy: Any = None + # Invoked AFTER ``run_agent_loop`` returns and BEFORE the bus + # absorbs ``raw_result.messages`` back into ``session.messages``. + # Sync or async. Mutates ``raw_result`` in place (or returns a + # replacement). Anything appended to ``result.messages`` flows + # naturally into ``session.messages`` and the closing boundary; + # rewriting ``result.final_content`` propagates to + # ``session.last_report``. Use this for "no-tool recovery" + # logic (e.g. swarm's ``force_final_answer``) — putting it here + # rather than in ``result_adapter`` ensures session bookkeeping + # sees the rescue. + force_finalizer: Callable[ + [Any, SubTask], + Any | Awaitable[Any], + ] | None = None + # Sync context manager ``setup(job_id, item)`` wrapping the whole + # sub-agent ``run_agent_loop`` coroutine. Runs in the sub-agent's own + # asyncio task, so any contextvar it sets (e.g. agent_team's per-sub + # bwrap sandbox binding ``/workspace`` + ``/inputs`` (ro) + ``/outputs`` + # (rw)) is scoped to that sub-agent and reset on exit. ``None`` → no + # wrapping (agent_team leases its sandbox through its own mechanism). + context_setup: Callable[[str, SubTask], Any] | None = None + + +@dataclass +class SubAgentSession: + """A durable sub-agent session that accumulates history across tasks. + + Boundary invariant + ------------------ + Each *closed* boundary ``(start, end)`` in ``task_boundaries`` is + guaranteed to contain at least one ``AIMessage`` with no + ``tool_calls`` between ``start + 1`` and ``end`` inclusive. The + bus enforces this on boundary closure: when the loop exits without + one, the bus appends a synthetic ``AIMessage`` populated from + ``raw_result.final_content`` (or a deterministic stub naming + ``stopped_by``) and points ``end`` at that synthetic message. + + This lets ``TaskBoundaryTrimmer`` always surface a "final report" + for completed tasks; without it, max_turns / llm_error / aborted + sub-agents would be silently dropped from the trimmed history on + reuse. + """ + + session_id: str + task_id: str + name: str + role_id: str + system_prompt: str + tools: list[Tool] + llm: Any + trimmer: MessageTrimmer + max_turns: int = 100 + tool_result_max_chars: int | None = None + # Per-LLM-call timeout (seconds) used by the agent loop's LLM client. + # ``None`` means defer to ``LoopConfig.llm_timeout`` (default 180s). + # Set this when the session's model has high response latency variance + # (e.g. slow auditor / strong reasoning model used for review). + llm_timeout: int | None = None + messages: list[Message] = field(default_factory=list) + task_boundaries: list[TaskBoundary] = field(default_factory=list) + total_task_count: int = 0 + current_job_id: str | None = None + # Wall-clock seconds the most recently finished task took. Kept so a + # progress snapshot can report a *stable* duration for an idle session: + # once ``current_job_id`` clears there is no job entry left to measure, + # and recomputing from "now" makes finished workers appear to keep + # running. + last_task_elapsed_s: float = 0.0 + last_report: str = "" + runtime_spec: SubAgentRuntimeSpec | None = None + pending_results: list[SubAgentResult] = field(default_factory=list) + # Tasks queued behind the currently-running one. Sessions execute + # serially (the messages list / boundary list cannot be safely + # mutated by two concurrent runs), so a second ``submit_task_to_session`` + # call enqueues here instead of running in parallel. The queue + # drains FIFO whenever ``current_job_id`` clears. + pending_tasks: deque[PendingSessionTask] = field(default_factory=deque) + # Bounded live event trail consumed by local UIs while a task runs. + # Full trajectories remain owned by workflow observers; this is only the + # recent thinking/tool activity needed for an expandable team overview. + activity_events: deque[dict[str, Any]] = field( + default_factory=lambda: deque(maxlen=40), + ) + publisher_no_tool_failures: int = 0 + # Path written by _eager_trim_and_offload for the most-recently dropped + # message batch. Debug only — None when offload is disabled or skipped. + offloaded_history_path: Path | None = None diff --git a/agent_core/components/agent_bus/runtime.py b/agent_core/components/agent_bus/runtime.py new file mode 100644 index 0000000..88f8642 --- /dev/null +++ b/agent_core/components/agent_bus/runtime.py @@ -0,0 +1,501 @@ +# pyright: reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportAttributeAccessIssue=false +"""Generic runtime helpers for AgentBus. + +This module keeps the generic default config / observer / result-adaptation +helpers out of ``agent_bus.py`` so the bus focuses on job and session +lifecycle management. Telemetry emission and post-restart job rehydration +also live here — they are mostly glue to ``EventStore`` and share no +state with the ``AgentBus`` class itself. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from agent_core.components.agent_bus.models import JobEntry, SubAgentResult, SubAgentSession +from agent_core.events import EventType +from agent_core.messages import assistant_msg +from agent_core.protocols import EventSink +from agent_core.runtime.loop.message_trimmer import find_final_assistant +from agent_core.runtime.registries import services as registry + +logger = logging.getLogger(__name__) + + +def build_default_subagent_loop_config( + job_id: str, + item: Any, + max_turns: int, +) -> Any: + """Build the generic default LoopConfig for async sub-agent jobs.""" + from agent_core.loop_types import LoopConfig + + return LoopConfig( + max_turns=max_turns, + task_id=job_id, + role_id=item.role_id, + ) + + +def build_default_subagent_observers( + *, + event_store: Any | None, + job_id: str, +) -> list[Any]: + """Build the generic default observer stack for async sub-agent jobs.""" + del event_store, job_id + return [] + + +def adapt_default_subagent_result( + agent_result: Any, + job_id: str, + item: Any, +) -> Any: + """Adapt an AgentLoopResult into a generic SubAgentResult. + + The kernel default forwards the full loop metadata bag untouched — + any domain-specific fields (evidence_cards, assertions, …) live + inside ``metadata`` and only appear when workflow observers + populated them. + """ + metadata = dict(getattr(agent_result, "metadata", {}) or {}) + return SubAgentResult( + question=item.question, + role_id=item.role_id, + final_content=getattr(agent_result, "final_content", "") or "", + success=True, + job_id=job_id, + metadata=metadata, + ) + + +def adapt_default_session_result( + agent_result: Any, + job_id: str, + session: Any, + task_prompt: str, +) -> Any: + """Adapt a session loop result into a generic SubAgentResult.""" + metadata = dict(getattr(agent_result, "metadata", {}) or {}) + return SubAgentResult( + question=task_prompt, + role_id=session.role_id, + final_content=getattr(agent_result, "final_content", "") or "", + success=True, + job_id=job_id, + metadata=metadata, + ) + + +def build_session_loop_config( + session: Any, + job_id: str, + max_turns: int, +) -> Any: + """Build LoopConfig for a session task.""" + from agent_core.loop_types import LoopConfig + + kwargs: dict[str, Any] = { + "max_turns": max_turns, + "task_id": session.task_id, + "role_id": session.role_id, + "tool_result_max_chars": session.tool_result_max_chars, + } + if getattr(session, "llm_timeout", None) is not None: + kwargs["llm_timeout"] = session.llm_timeout + return LoopConfig(**kwargs) + + +def resolve_session_observers( + observers: list[Any] | None, + session: Any | None = None, +) -> list[Any]: + """Resolve observers for a session task. + + The kernel default is intentionally empty. Workflow-specific observer + stacks should be injected through ``runtime_spec`` or explicit callers. + """ + if observers is not None: + return observers + del session + return [] + + +def close_session_boundary_aborted(session: Any) -> None: + """Close the current session task boundary when cancelled/errored. + + Maintains the SubAgentSession boundary invariant: when no clean + AIMessage exists in the in-flight slice, append an abort stub so + the trimmer can still surface "task #N happened, here's what we + have" when the agent is reused. + """ + if not session.task_boundaries: + return + start, end = session.task_boundaries[-1] + if end is not None: + return + + if find_final_assistant( + session.messages, start + 1, len(session.messages) - 1, + ) is None: + session.messages.append( + assistant_msg("[task aborted before producing a final answer]"), + ) + + tail = max(start, len(session.messages) - 1) + session.task_boundaries[-1] = (start, tail) + + +async def emit_session_task_submitted( + session: SubAgentSession, + job_id: str, + task_prompt: str, + *, + event_sink: Any = None, +) -> None: + """Record the submit side of a session-task lifecycle to EventStore. + + ``event_sink`` is an optional ``core.protocols.EventSink`` injected + by ``AgentBus``. Empty falls back to the + global registry lookup so existing direct callers keep working. + """ + event_store = event_sink if event_sink is not None else registry.get_optional(EventSink) + if event_store is None: + return + await event_store.append( + task_id=session.task_id, + event_type=EventType.AGENT_ACTION, + payload={ + "trace_type": "session_task_submitted", + "session_id": session.session_id, + "job_id": job_id, + "task_count": session.total_task_count, + "is_reuse": session.total_task_count > 1, + "role_id": session.role_id, + "agent": session.name, + "action": "assign_task", + "detail": ( + f"{'Reusing' if session.total_task_count > 1 else 'Starting'} " + f"sub-agent '{session.name}' " + f"(task #{session.total_task_count}): " + f"{task_prompt[:140]}" + ), + }, + agent_role="system", + ) + + +def safe_metadata(raw_result: Any) -> dict[str, Any]: + """Best-effort extraction of ``metadata`` from a partial AgentLoopResult. + + Called from exception handlers where ``raw_result`` may be ``None`` + (agent loop never completed) or a completed result whose downstream + adapter raised. Never raises — returns ``{}`` on any oddity. + + The evidence harvested inside an agent loop is the most expensive + side-effect of a research run (it cost real web_search/web_fetch + API calls). Losing it because a post-loop bookkeeping step + crashed is the worst possible UX: the user still paid for the + calls but sees nothing in the evidence graph. + """ + if raw_result is None: + return {} + try: + meta = getattr(raw_result, "metadata", None) + if isinstance(meta, dict): + return dict(meta) + except Exception: + pass + return {} + + +def safe_final_content(raw_result: Any) -> str: + """Safely pull ``final_content`` off a possibly-None result.""" + if raw_result is None: + return "" + try: + value = getattr(raw_result, "final_content", "") or "" + return str(value) + except Exception: + return "" + + +SUBTASK_RESULT_MESSAGE_TYPE = "subtask_result" +"""``AgentMessage.message_type`` carrying a sub-agent's result home. + +Open string by design — ``models/agent_message.py`` documents that any +agent may define its own types, and L1-3 will add a registry that +validates them rather than a hardcoded enum. +""" + + +def result_to_message_content(result: SubAgentResult) -> dict[str, Any]: + """Serialise a ``SubAgentResult`` for transport, losslessly. + + All eight fields (``models.py``'s dataclass) travel, because the + receiving side rebuilds the object and hands it to the *same* + ``fan_in.process_collected``. Dropping any of them would change the + rendered ````: ``final_content`` is the body, ``success`` and + ``error_class`` drive ``classify_completion``, and ``metadata`` holds + both ``stopped_by`` (the status label) and the harvested evidence. + """ + return { + "question": result.question, + "role_id": result.role_id, + "final_content": result.final_content, + "success": result.success, + "error": result.error, + "error_class": result.error_class, + "job_id": result.job_id, + "metadata": dict(result.metadata or {}), + } + + +def message_content_to_result(content: dict[str, Any]) -> SubAgentResult: + """Rebuild a ``SubAgentResult`` from :func:`result_to_message_content`. + + Defensive about every field: a message may have been written by an + older build, and a partially-shaped payload should still produce a + renderable report rather than raising inside the collect path. + """ + return SubAgentResult( + question=str(content.get("question") or ""), + role_id=str(content.get("role_id") or ""), + final_content=str(content.get("final_content") or ""), + success=bool(content.get("success")), + error=content.get("error") or None, + error_class=content.get("error_class") or None, + job_id=str(content.get("job_id") or ""), + metadata=dict(content.get("metadata") or {}), + ) + + +async def send_subtask_result_message( + session: SubAgentSession, + job_id: str, + result: SubAgentResult, + spawn_context: dict[str, Any] | None, +) -> None: + """Hand a sub-agent's result back to its parent as a message. + + Why this exists (L1-0): the result used to travel only through + ``AgentBus`` in-process state — ``entry.result`` plus a destructive + ``session.pending_results`` pop. The event stream recorded the + *fact* of completion (``session_task_completed``) but never the + content, so a hand-back could not be replayed, and a double drain + had no idempotency key. Sending it as a real message makes the + hand-back addressable (``to_agent``), correlated + (``correlation_id == job_id``), causally linked (``parent_id`` is the + ``delegate_subtask`` tool call that spawned it), and durable. + + The join keys are **reused, not invented**: ``spawn_context`` already + carries the parent→sub lineage that Log Schema v1 §2.4 requires, so + trace-side and message-side reconstruction stay on one set of keys. + + ``spawn_context`` values are passed through **verbatim**, including an + empty ``parent_run_id`` — agent_team's coordinator carries no + ``run_id`` in scope metadata, so ``assign_task`` stamps ``""`` (v1-doc + P7). The trace layer normalises that to ``run_1`` at its v1-translation + boundary (``worker_trace._sub_run_to_v1_sub_agent``); doing it here too + would put the same rule in two places, and inventing a value in the + kernel would leak a trace-layer convention into the bus. Join on + ``to_agent`` (the parent agent id), which is always populated — + ``parent_run_id`` is only reliable after that normalisation. + + Soft-dependency and fully swallowed: with no ``AgentComm`` + registered, or no ``spawn_context`` to address, behaviour is exactly + what it was before this function existed. + """ + if not spawn_context: + return + parent_agent_id = str(spawn_context.get("parent_agent_id") or "").strip() + if not parent_agent_id: + return + try: + from agent_core.components.agent_bus.agent_comm import AgentComm + from agent_core.models.agent_message import AgentMessage + + comm = registry.get_optional(AgentComm) + if comm is None: + return + await comm.send( + AgentMessage( + task_id=session.task_id, + from_agent=session.name, + to_agent=parent_agent_id, + message_type=SUBTASK_RESULT_MESSAGE_TYPE, + # The spawning tool call — the causal parent of this reply. + parent_id=( + str(spawn_context.get("spawned_by_tool_call_id") or "") or None + ), + content={ + # AgentComm lifts this onto the KernelEvent column, which + # is what makes the message queryable by job. + "correlation_id": job_id, + "session_id": session.session_id, + "result": result_to_message_content(result), + "spawn_context": { + "parent_run_id": spawn_context.get("parent_run_id") or "", + "parent_turn": spawn_context.get("parent_turn") or 0, + "spawned_by_llm_call_id": ( + spawn_context.get("spawned_by_llm_call_id") or "" + ), + }, + }, + ), + ) + except Exception: + # Same discipline as the telemetry below: the delivery channel is + # not allowed to break the session lifecycle. + logger.debug("subtask_result message not sent for %s", job_id, exc_info=True) + + +async def emit_session_task_completed( + session: SubAgentSession, + job_id: str, + result: SubAgentResult, + *, + event_sink: Any = None, + spawn_context: dict[str, Any] | None = None, +) -> None: + """Record the completion side of a session-task lifecycle. + + Swallows failures — telemetry must not break the session lifecycle. + List-valued metadata entries are summarised as ``{key}_count`` to + keep the event payload compact regardless of workflow-specific + metadata shape. + + ``event_sink`` accepts a ``core.protocols.EventSink`` injected by + ``AgentBus``; falls back to the global + registry when omitted. + + ``spawn_context`` is the delegation lineage recorded at dispatch + (Log Schema v1 §2.4). When present, and when an ``AgentComm`` is + registered, the result is **also** sent as a ``subtask_result`` + message so the hand-back is addressable, correlated and replayable + — see :func:`send_subtask_result_message`. + """ + # Message first: it has its own store via AgentComm, so an absent + # telemetry sink must not suppress it. + await send_subtask_result_message(session, job_id, result, spawn_context) + ev_store = event_sink if event_sink is not None else registry.get_optional(EventSink) + if ev_store is None: + return + try: + metadata_counts = { + f"{k}_count": len(v) + for k, v in result.metadata.items() + if isinstance(v, list) + } + detail_parts = [ + f"{len(v)} {k}" + for k, v in result.metadata.items() + if isinstance(v, list) and v + ] + detail = f"'{session.name}' returned" + if detail_parts: + detail += " " + ", ".join(detail_parts) + if not result.success and result.error: + # Make the failure legible in both the event stream and + # any downstream log summarizer — the previous shape left + # "success: false" with no reason and every sub-agent + # failure looked identical. + cls = (result.error_class or "").strip() + err_short = str(result.error)[:200] + detail += f" (failed: {cls}: {err_short})" if cls else f" (failed: {err_short})" + payload: dict[str, Any] = { + "trace_type": "session_task_completed", + "session_id": session.session_id, + "job_id": job_id, + "success": result.success, + "agent": session.name, + "action": "report_returned", + "detail": detail, + **metadata_counts, + } + if result.error: + payload["error"] = str(result.error)[:500] + if result.error_class: + payload["error_class"] = result.error_class + await ev_store.append( + task_id=session.task_id, + event_type=EventType.AGENT_ACTION, + payload=payload, + agent_role="system", + ) + except Exception: + pass + + +async def rehydrate_orphaned_jobs( + jobs: dict[str, JobEntry], *, event_sink: Any = None, +) -> int: + """Recover job registry after process restart. + + Reads ``EventStore`` for ``agent_submitted`` events that have no + matching ``agent_completed`` / ``agent_aborted`` / ``agent_collected`` + event. Marks those orphaned jobs as failed (their asyncio.Tasks are + gone) by appending an ``agent_orphaned`` event for each. + + Returns the number of orphaned jobs recorded. No-op if no EventStore + is registered. ``event_sink`` lets ``AgentBus`` thread its + constructor-injected sink through; empty + falls back to the global registry. + """ + event_store = event_sink if event_sink is not None else registry.get_optional(EventSink) + if event_store is None: + return 0 + + all_events = await event_store.get_all_events( + event_type=EventType.AGENT_ACTION, + ) + + submitted: dict[str, dict[str, Any]] = {} + completed_ids: set[str] = set() + + for evt in all_events: + payload = evt.payload or {} + trace_type = payload.get("trace_type", "") + job_id = payload.get("job_id", "") + + if trace_type == "agent_submitted" and job_id: + submitted[job_id] = payload + elif trace_type in ( + "agent_completed", "agent_aborted", "agent_collected", + ): + if job_id: + completed_ids.add(job_id) + for jid in payload.get("job_ids", []): + completed_ids.add(jid) + + orphaned = 0 + for job_id, payload in submitted.items(): + if job_id in completed_ids: + continue + if job_id in jobs: + continue # already tracked (shouldn't happen post-restart) + + orphaned += 1 + parent_task_id = payload.get("parent_task_id", "") + logger.warning( + "Rehydration: marking orphaned job %s as failed (parent=%s)", + job_id, parent_task_id, + ) + await event_store.append( + task_id=parent_task_id, + event_type=EventType.AGENT_ACTION, + payload={ + "trace_type": "agent_orphaned", + "job_id": job_id, + "reason": "process_restart", + }, + agent_role="system", + ) + + if orphaned: + logger.info( + "Rehydration: %d orphaned jobs marked as failed", orphaned, + ) + return orphaned diff --git a/agent_core/components/agent_bus/spawn_guard.py b/agent_core/components/agent_bus/spawn_guard.py new file mode 100644 index 0000000..256c5fa --- /dev/null +++ b/agent_core/components/agent_bus/spawn_guard.py @@ -0,0 +1,341 @@ +# pyright: reportDeprecated=false, reportMissingTypeArgument=false, reportUnknownParameterType=false, reportUnknownVariableType=false +"""Guard sub-agent spawning by depth, concurrency, budget, and wall time. + +Reservations are released on both completion and failure. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import time +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from dataclasses import dataclass, field + +from agent_core.models.task_budget import TaskBudget + +logger = logging.getLogger(__name__) + +# Ceiling the long-running callers (a worker bootstrap, the benchmark +# ``kernel_adapter``) pin explicitly — 5 hours, deliberately far above +# ``SpawnGuard.DEFAULT_SUB_AGENT_TIMEOUT_S``. See +# :func:`resolve_sub_agent_timeout_s` for why it is safe to be this loose and +# ``benchmarks/kernel_adapter.py`` for the measurements. +PINNED_SUB_AGENT_TIMEOUT_S = 18_000 + +_SUB_AGENT_TIMEOUT_ENV = "SUB_AGENT_TIMEOUT_S" + + +def resolve_sub_agent_timeout_s( + default: int = PINNED_SUB_AGENT_TIMEOUT_S, +) -> int: + """Resolve the per-sub-agent wall-clock ceiling from the environment. + + Single source of truth for the ``SUB_AGENT_TIMEOUT_S`` override so the + literal is not duplicated across call sites, and so a malformed value + cannot take down the process it is read in. A worker bootstrap reads + this at start-up, where a bare ``int(os.environ[...])`` would turn + a templated-but-empty ``SUB_AGENT_TIMEOUT_S=`` into a boot failure. + + Non-positive values are rejected rather than honoured: ``0`` would disable + the hard ``asyncio.wait_for`` while also skipping ``WallClockGuard``, and a + negative would make ``wait_for`` fail every sub-agent instantly. + """ + raw = os.environ.get(_SUB_AGENT_TIMEOUT_ENV, "").strip() + if not raw: + return default + try: + value = int(raw) + except ValueError: + logger.warning( + "Invalid %s=%r (not an integer); using %d", + _SUB_AGENT_TIMEOUT_ENV, raw, default, + ) + return default + if value <= 0: + logger.warning( + "Ignoring %s=%d (must be > 0); using %d", + _SUB_AGENT_TIMEOUT_ENV, value, default, + ) + return default + return value + +# ── Exceptions ────────────────────────────────────────────────────────── + + +class SpawnDepthExceeded(RuntimeError): + """Depth limit reached — cannot spawn deeper.""" + + +class BudgetExhausted(RuntimeError): + """Token/cost budget insufficient for this spawn.""" + + +# ── SpawnReservation ──────────────────────────────────────────────────── + + +@dataclass +class SpawnReservation: + """RAII handle for a single spawn slot. + + Holds a semaphore slot + pre-charged tokens. + Must be released via SpawnGuard.release() or the async context manager. + + ``acquired`` flips to ``True`` only after :meth:`SpawnGuard.acquire_slot` + returns with the semaphore in hand. The bus.py ``finally`` path may + call ``release`` after ``pre_check`` succeeded but before + ``acquire_slot`` finished (cancellation / timeout); without this + flag, ``release`` would hand back a slot it never owned and inflate + ``max_parallel`` over time. + """ + + job_id: str + depth: int + estimated_tokens: int + acquired_at: float = field(default_factory=time.monotonic) + acquired: bool = False + + +# ── SpawnGuard ────────────────────────────────────────────────────────── + + +class SpawnGuard: + """Budget-aware spawn controller for AgentBus. + + Enforces five layers of protection: + 1. Depth limit + 2. Concurrency limit (semaphore — queues, never rejects) + 3. Token budget (rejects if insufficient) + 4. Timeout (enforced by caller via asyncio.wait_for) + 5. RAII reservation (auto-release on exception) + + Thread-safe within a single asyncio event loop. + """ + + # Default sub-agent timeout: 90 minutes. Deliberately generous — slow + # models (e.g. Qwen 397B, R1, Sonnet thinking) routinely spend 400-800s + # per turn on real research, so a tight default silently masks legitimate + # work as ``(agent failed: timeout)`` and triggers cascading respawns. + # Callers that want a tighter ceiling should pass ``sub_agent_timeout_s`` + # explicitly. + # + # Set above the usual one-hour ceiling because long research calls may + # opus5-mix3 (4417 sub-agent trajectories). **52.8% of sub-agents were + # killed by this ceiling** and returned ``(empty report)``, discarding + # ~30 turns of work each; the main agent then respawned from scratch + # (the "cascading respawns" this comment already warned about). + # + # The wall clock was NOT spent in tools — median ``bash`` 0.8s, + # ``web_search`` 1.6s, ``web_fetch`` 6.3s; tools were only 12-15% of + # the 58 min. **85-88% was waiting on the LLM**: ~100s per turn × 30 + # turns. And that latency was queueing, not compute — single-stream + # output throughput measured 12.9 tok/s under CONCURRENCY=64 vs + # 26.1 tok/s on a lighter run, and latency *fell* as prompts grew + # (turn 1-5: 14k prompt → 112s; turn 41+: 71k prompt → 60s). + # + # So the real fix for that run was cutting eval concurrency; this raise + # is headroom for genuinely long sub-tasks, deliberately kept at +50% + # rather than 2× so a wedged sub-agent still cannot squat a slot for + # hours. The clean-exit half of the fix is + # :class:`~agent_core.components.observers.wall_clock_guard.\ + # WallClockGuard`, which ``AgentBus`` attaches from this same + # ``timeout_s`` so the loop stops itself before the hard + # ``asyncio.wait_for`` can cancel it. + DEFAULT_SUB_AGENT_TIMEOUT_S = 5400 + + def __init__( + self, + budget: TaskBudget | None = None, + sub_agent_timeout_s: int | None = None, + ) -> None: + b = budget or TaskBudget() + self._max_depth = b.max_depth + self._max_tokens = b.max_tokens + self._max_parallel = b.max_parallel + # Sub-agent timeout is fully decoupled from ``TaskBudget.max_wall_time_s`` + # (the root task's wall budget): they answer different questions + # and the previous ``b.max_wall_time_s or DEFAULT`` fallback turned + # the root default (300s) into a silent sub-agent kill switch. + # Use the caller-supplied value or the class-wide default; never + # inherit from the root budget. + self._timeout_s = ( + sub_agent_timeout_s + if sub_agent_timeout_s is not None + else self.DEFAULT_SUB_AGENT_TIMEOUT_S + ) + self._semaphore = asyncio.Semaphore(max(b.max_parallel, 1)) + self._tokens_reserved: int = 0 + self._tokens_actual: int = 0 + self._active: dict[str, SpawnReservation] = {} + self._total_spawns: int = 0 + self._lock = asyncio.Lock() + + # ── Properties ────────────────────────────────────────────────────── + + @property + def max_depth(self) -> int: + return self._max_depth + + @property + def max_parallel(self) -> int: + return self._max_parallel + + @property + def timeout_s(self) -> int: + return self._timeout_s + + @property + def tokens_reserved(self) -> int: + return self._tokens_reserved + + @property + def tokens_actual(self) -> int: + return self._tokens_actual + + @property + def active_count(self) -> int: + return len(self._active) + + @property + def total_spawns(self) -> int: + return self._total_spawns + + @property + def remaining_tokens(self) -> int: + return max(0, self._max_tokens - self._tokens_reserved) + + # ── pre_check + acquire_slot + acquire (combined) ───────────────── + + async def pre_check( + self, + job_id: str, + depth: int, + estimated_tokens: int = 0, + ) -> None: + """Non-blocking pre-check: depth + budget. Called at submit time. + + Raises: + SpawnDepthExceeded: if depth >= max_depth + BudgetExhausted: if estimated_tokens > remaining budget + """ + # Layer 5: Depth check + if depth >= self._max_depth: + raise SpawnDepthExceeded( + f"Spawn depth {depth} >= max {self._max_depth}" + ) + + # Layer 3: Token budget check (under lock) + async with self._lock: + remaining = self._max_tokens - self._tokens_reserved + if estimated_tokens > 0 and estimated_tokens > remaining: + raise BudgetExhausted( + f"Need {estimated_tokens} tokens, " + f"only {remaining} remaining " + f"(reserved {self._tokens_reserved} " + f"of {self._max_tokens})" + ) + self._tokens_reserved += estimated_tokens + + # Register reservation (without semaphore yet) + reservation = SpawnReservation( + job_id=job_id, + depth=depth, + estimated_tokens=estimated_tokens, + ) + self._active[job_id] = reservation + self._total_spawns += 1 + + logger.debug( + "SpawnGuard.pre_check(%s): depth=%d, est_tokens=%d, " + "reserved=%d/%d", + job_id, depth, estimated_tokens, + self._tokens_reserved, self._max_tokens, + ) + + async def acquire_slot(self, job_id: str) -> None: + """Acquire concurrency slot. May block (queue). Called at run time. + + Layer 4: Concurrency limit (semaphore — queues, never rejects). + """ + await self._semaphore.acquire() + reservation = self._active.get(job_id) + if reservation is not None: + reservation.acquired = True + logger.debug( + "SpawnGuard.acquire_slot(%s): active=%d/%d", + job_id, len(self._active), self._max_parallel, + ) + + async def acquire( + self, + job_id: str, + depth: int, + estimated_tokens: int = 0, + ) -> SpawnReservation: + """Combined pre_check + acquire_slot. Blocks on concurrency. + + Use pre_check + acquire_slot separately for non-blocking submit. + This combined method is for the RAII context manager and tests. + """ + await self.pre_check(job_id, depth, estimated_tokens) + await self.acquire_slot(job_id) + return self._active[job_id] + + def release( + self, + job_id: str, + actual_tokens: int = 0, + ) -> None: + """Release a spawn slot. Corrects token estimate with actual usage. + + Safe to call multiple times (idempotent). + """ + reservation = self._active.pop(job_id, None) + if reservation is None: + return + + self._tokens_reserved -= reservation.estimated_tokens + self._tokens_actual += actual_tokens + + if reservation.acquired: + self._semaphore.release() + + logger.debug( + "SpawnGuard.release(%s): est=%d, actual=%d, " + "active=%d/%d, slot_returned=%s", + job_id, reservation.estimated_tokens, actual_tokens, + len(self._active), self._max_parallel, reservation.acquired, + ) + + @asynccontextmanager + async def reservation( + self, + job_id: str, + depth: int, + estimated_tokens: int = 0, + ) -> AsyncIterator[SpawnReservation]: + """RAII context manager for acquire/release. + + Guarantees release even if the spawned task raises. + """ + res = await self.acquire(job_id, depth, estimated_tokens) + try: + yield res + finally: + self.release(job_id) + + # ── Stats ─────────────────────────────────────────────────────────── + + def stats(self) -> dict: + """Return current guard state for debugging/SSE.""" + return { + "active": len(self._active), + "max_parallel": self._max_parallel, + "max_depth": self._max_depth, + "tokens_reserved": self._tokens_reserved, + "tokens_actual": self._tokens_actual, + "max_tokens": self._max_tokens, + "total_spawns": self._total_spawns, + } diff --git a/agent_core/components/finalization/__init__.py b/agent_core/components/finalization/__init__.py index a8cb9a6..c0c1638 100644 --- a/agent_core/components/finalization/__init__.py +++ b/agent_core/components/finalization/__init__.py @@ -1 +1,53 @@ -"""Portable final-answer recovery primitives.""" +"""Portable finalization budget and recovery primitives.""" + +from agent_core.components.finalization.budget import ( + LEGACY_TASK_WALL_TIME_ENVS, + REPORT_WALL_TIME_ENV, + RESEARCH_WALL_TIME_ENV, + TASK_WALL_TIME_ENV, + ResearchWall, + check_wall_feasibility, + nonnegative_seconds, + positive_seconds, + remaining_phase_budget_s, + resolve_report_wall_time, + resolve_research_wall, + soft_wall_deadline_s, +) +from agent_core.components.finalization.recovery import ( + COMMON_RECOVERY_NUDGE_PREFIXES, + RECOVERY_CONTEXT_MAX_CHARS, + RECOVERY_ITEM_MAX_CHARS, + build_recovery_context, + chat_with_fallback_budget, + fallback_leg_count, + has_malformed_tool_protocol, + is_recovery_nudge, + minimal_best_effort_answer, + truncate_text_to_tokens, +) + +__all__ = [ + "COMMON_RECOVERY_NUDGE_PREFIXES", + "LEGACY_TASK_WALL_TIME_ENVS", + "RECOVERY_CONTEXT_MAX_CHARS", + "RECOVERY_ITEM_MAX_CHARS", + "REPORT_WALL_TIME_ENV", + "RESEARCH_WALL_TIME_ENV", + "TASK_WALL_TIME_ENV", + "ResearchWall", + "build_recovery_context", + "chat_with_fallback_budget", + "check_wall_feasibility", + "fallback_leg_count", + "has_malformed_tool_protocol", + "is_recovery_nudge", + "minimal_best_effort_answer", + "nonnegative_seconds", + "positive_seconds", + "remaining_phase_budget_s", + "resolve_report_wall_time", + "resolve_research_wall", + "soft_wall_deadline_s", + "truncate_text_to_tokens", +] diff --git a/agent_core/components/finalization/budget.py b/agent_core/components/finalization/budget.py new file mode 100644 index 0000000..555cb11 --- /dev/null +++ b/agent_core/components/finalization/budget.py @@ -0,0 +1,263 @@ +"""Wall-clock arithmetic for workflows with a research-only deadline. + +A ``TASK_WALL_TIME_MODE = "soft_research"`` workflow stops its agent loop +*early* so the downstream finalization stage (reporter or ``force_final_answer``) +can still complete inside whatever hard ceiling the platform imposes. Getting +that split right needs two numbers, not one: + +``research_deadline_s`` + when the in-loop :class:`WallClockDeadlineObserver` stops research. + +``hard_total_s`` + the externally enforced ceiling the whole task must fit inside — the + ``MIROHARNESS_TASK_WALL_TIME_S`` env value (enforced in live ``serve`` by + ``RenewableWallTimeLease``) or a legacy total-task ``wall_deadline_s``. + ``0`` means "no known external ceiling". + +The reserve subtracted from ``hard_total_s`` is only ever a *plan*: +``soft_wall_deadline_s`` floors research at half the wall so a generous reserve +can never starve it, which means on a short wall the reserve is smaller than +requested. :func:`remaining_phase_budget_s` is the matching safety net — the +finalization stage clamps its own ceiling to the time that is actually left +when it starts, so it fails open with a real answer instead of being killed +mid-call by the external ceiling. +""" + +from __future__ import annotations + +import logging +import math +import os +import time +from dataclasses import dataclass +from typing import Any + +logger = logging.getLogger(__name__) + +TASK_WALL_TIME_ENV = "AGENT_CORE_TASK_WALL_TIME_S" +LEGACY_TASK_WALL_TIME_ENVS = ( + "MIROHARNESS_TASK_WALL_TIME_S", + "FRONTIER_AGENT_TASK_WALL_TIME_S", +) +RESEARCH_WALL_TIME_ENV = "RESEARCH_WALL_TIME" +REPORT_WALL_TIME_ENV = "REPORT_WALL_TIME" + + +def positive_seconds(raw: object, *, label: str) -> float | None: + """Parse a positive duration; invalid/disabled values contribute no cap.""" + if raw is None or raw == "": + return None + try: + value = float(raw) # type: ignore[arg-type] + except (TypeError, ValueError): + logger.warning("Invalid %s=%r; ignoring it", label, raw) + return None + if not math.isfinite(value): + logger.warning("Non-finite %s=%r; ignoring it", label, raw) + return None + return value if value > 0 else None + + +def nonnegative_seconds(raw: object, *, default: float, label: str) -> float: + """Parse a non-negative duration with a tolerant profile fallback.""" + if raw is None or raw == "": + return float(default) + try: + value = float(raw) # type: ignore[arg-type] + except (TypeError, ValueError): + logger.warning("Invalid %s=%r; using %.0fs", label, raw, default) + return float(default) + if not math.isfinite(value) or value < 0: + logger.warning("Invalid non-negative %s=%r; using %.0fs", label, raw, default) + return float(default) + return value + + +def resolve_report_wall_time( + agent_cfg: dict[str, Any], + *, + default: float, + label_prefix: str, + env_var: str = REPORT_WALL_TIME_ENV, +) -> float: + """Resolve finalization reserve, with an optional deployment override. + + ``REPORT_WALL_TIME`` maps to the profile's internal + ``wall_deadline_reserve_s`` field. The environment wins when it contains a + valid non-negative duration; unset/invalid values preserve the profile. + """ + profile_s = nonnegative_seconds( + agent_cfg.get("wall_deadline_reserve_s"), + default=default, + label=f"{label_prefix} wall_deadline_reserve_s", + ) + raw_env = os.environ.get(env_var) + if raw_env is None or raw_env == "": + return profile_s + return nonnegative_seconds(raw_env, default=profile_s, label=env_var) + + +def soft_wall_deadline_s(total_s: float, reserve_s: float) -> float: + """Convert a hard task wall into a research deadline with finalize grace. + + The ``total_s * 0.5`` floor keeps a generous reserve from starving research + on a short wall. When the floor binds, the reserve that actually survives is + smaller than ``reserve_s`` — callers MUST pair this with + :func:`remaining_phase_budget_s` so the finalization stage shrinks to the + time that is really left instead of assuming it got its full reserve. + """ + return max(total_s - reserve_s, total_s * 0.5) + + +@dataclass(frozen=True) +class ResearchWall: + """Resolved research deadline plus the external ceiling it came from.""" + + research_deadline_s: float + """When the in-loop observer stops research. ``0`` = no in-loop wall.""" + + hard_total_s: float + """Externally enforced whole-task ceiling. ``0`` = none known.""" + + +def resolve_research_wall( + agent_cfg: dict[str, Any], + *, + reserve_s: float, + label_prefix: str, + env_var: str = TASK_WALL_TIME_ENV, +) -> ResearchWall: + """Resolve the research deadline from profile budget and platform wall. + + ``research_wall_time_s`` is already a research-only budget and is not + shortened. A valid ``RESEARCH_WALL_TIME`` replaces that profile budget. + The legacy ``wall_deadline_s`` and the operational task-wall env value are + total-task ceilings, so finalize grace is subtracted before comparing them — + and they are also what :attr:`ResearchWall.hard_total_s` reports. + """ + research_budget = "research_wall_time_s" in agent_cfg + if research_budget: + profile_raw = agent_cfg.get("research_wall_time_s") + elif "wall_deadline_s" in agent_cfg: + profile_raw = agent_cfg.get("wall_deadline_s") + else: + profile_raw = None + profile_s = positive_seconds( + profile_raw, + label=f"{label_prefix} research wall deadline", + ) + research_env_raw = os.environ.get(RESEARCH_WALL_TIME_ENV) + research_env_s = positive_seconds( + research_env_raw, + label=RESEARCH_WALL_TIME_ENV, + ) + if research_env_s is not None: + # This env is an explicit phase-budget override, not another ceiling: + # it may intentionally lengthen as well as shorten the YAML default. + profile_s = research_env_s + research_budget = True + resolved_env = env_var + if env_var == TASK_WALL_TIME_ENV and not os.environ.get(env_var): + resolved_env = next( + (name for name in LEGACY_TASK_WALL_TIME_ENVS if os.environ.get(name)), + env_var, + ) + env_s = positive_seconds(os.environ.get(resolved_env), label=resolved_env) + + profile_deadline_s = profile_s + if profile_deadline_s is not None and not research_budget: + profile_deadline_s = soft_wall_deadline_s(profile_deadline_s, reserve_s) + env_deadline_s = ( + soft_wall_deadline_s(env_s, reserve_s) if env_s is not None else None + ) + candidates = [ + value + for value in (profile_deadline_s, env_deadline_s) + if value is not None + ] + + # Only total-task values are hard ceilings; a research-only profile budget + # is not one, because the reporter is deliberately outside it. + hard_candidates = [ + value + for value in (None if research_budget else profile_s, env_s) + if value is not None + ] + return ResearchWall( + research_deadline_s=min(candidates) if candidates else 0.0, + hard_total_s=min(hard_candidates) if hard_candidates else 0.0, + ) + + +def remaining_phase_budget_s( + requested_s: float, + deadline_monotonic_s: float | None, + *, + minimum_s: float = 1.0, +) -> float: + """Clamp a finalization phase ceiling to the time that is actually left. + + ``deadline_monotonic_s`` is a :func:`time.monotonic` instant, normally + ``node_start + hard_total_s``. Returns ``requested_s`` unchanged when no + external ceiling is known. ``minimum_s`` keeps the phase from being handed a + zero/negative timeout — it still gets one short attempt and then fails open + to its baseline answer, which is strictly better than being cancelled by the + external ceiling with nothing to show. + """ + if deadline_monotonic_s is None: + return requested_s + remaining = deadline_monotonic_s - time.monotonic() + return max(min(float(requested_s), remaining), float(minimum_s)) + + +def check_wall_feasibility( + *, + hard_total_s: float, + research_deadline_s: float, + tool_timeout_s: float, + landing_budget_s: float, + label_prefix: str, +) -> bool: + """Warn when no schedule can honour the external ceiling; return feasibility. + + :func:`remaining_phase_budget_s` shrinks the finalization phase to fit, but + it cannot shrink a tool call that is already running. When + ``research_deadline_s + tool_timeout_s`` alone exceeds ``hard_total_s``, a + single tool started just before the research deadline blows the wall no + matter what the finalization stage does — the config itself is the problem + (usually ``tool_timeout_s`` larger than half the wall). Say so loudly rather + than letting the run get killed with no answer and no explanation. + """ + if hard_total_s <= 0 or research_deadline_s <= 0: + return True + worst_case_s = research_deadline_s + max(tool_timeout_s, 0.0) + if worst_case_s <= hard_total_s: + return True + logger.warning( + "%s: wall-time config cannot guarantee a final answer — research stops " + "at %.0fs and one late tool call may run to %.0fs, past the %.0fs hard " + "ceiling, leaving no room for the %.0fs finalization phase. Lower " + "tool_timeout_s (to at most half the wall) or raise the wall.", + label_prefix, + research_deadline_s, + worst_case_s, + hard_total_s, + landing_budget_s, + ) + return False + + +__all__ = [ + "LEGACY_TASK_WALL_TIME_ENVS", + "REPORT_WALL_TIME_ENV", + "RESEARCH_WALL_TIME_ENV", + "TASK_WALL_TIME_ENV", + "ResearchWall", + "check_wall_feasibility", + "nonnegative_seconds", + "positive_seconds", + "remaining_phase_budget_s", + "resolve_report_wall_time", + "resolve_research_wall", + "soft_wall_deadline_s", +] diff --git a/agent_core/models/agent_message.py b/agent_core/models/agent_message.py index a0345dd..e0d6cdd 100644 --- a/agent_core/models/agent_message.py +++ b/agent_core/models/agent_message.py @@ -8,6 +8,7 @@ from __future__ import annotations from datetime import UTC, datetime +from typing import Any from uuid import uuid4 from pydantic import BaseModel, Field @@ -21,6 +22,6 @@ class AgentMessage(BaseModel): from_agent: str # role_id of sender to_agent: str # role_id of recipient message_type: str # free-form: "assertion", "dispute", "delegation", etc. - content: dict = Field(default_factory=dict) # payload varies by message_type + content: dict[str, Any] = Field(default_factory=dict) # payload varies by message_type parent_id: str | None = None # links to triggering message for chain tracing timestamp: datetime = Field(default_factory=lambda: datetime.now(UTC)) diff --git a/agent_core/models/event.py b/agent_core/models/event.py new file mode 100644 index 0000000..5cc9b20 --- /dev/null +++ b/agent_core/models/event.py @@ -0,0 +1,230 @@ +"""Immutable kernel event model.""" + +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + +from agent_core.types import EventId, PromptId, SessionId, StepId, TaskId + +# SDK Agent Protocol namespace prefix — must stay in sync with +# the protocol layer's own event namespace (not imported from there: that +# layer sits above agent_core in the stack). +_SDK_EVENT_NS: str = "response.swarm" + + +class KernelEvent(BaseModel): + """An immutable event in the OS event log. + + Every action in FrontierAgent produces an event. Events are append-only + and support replay for debugging and audit. + + ``event_type`` is stored as a plain string. Callers can pass any + ``str``-Enum member (``EventType``, or any domain-specific event enum a + caller defines) and pydantic will coerce it to its underlying + string value, so the kernel does not need to know about every + domain enum that writes to the event log. + """ + + model_config = ConfigDict(frozen=True) + + id: EventId = Field(default=EventId("")) + task_id: TaskId + session_id: SessionId | None = None + prompt_id: PromptId | None = None + step_id: StepId | None = None + event_type: str + timestamp: datetime = Field(default_factory=lambda: datetime.now(UTC)) + payload: dict[str, Any] = Field(default_factory=dict) + from_agent: str | None = None + to_agent: str | None = None + message_type: str | None = None + correlation_id: str | None = None + retry_count: int = 0 + degrade_from: str | None = None + degrade_to: str | None = None + + # Internal events not shown in the frontend activity feed + _SKIP_EVENTS = {"task_created", "task_status_changed", "tool_called", "tool_result"} + + # Mapping from kernel/domain event strings to frontend SSEEventType + _TYPE_MAP: dict[str, str] = { + "phase_transition": "phase_change", + "agent_action": "agent_action", + "agent_tool_call": "agent_action", + "agent_message": "agent_action", + "error": "error", + "report_generated": "completed", + } + + # Protocol envelope event types — emitted by the optional external + # protocol observer (see ``workflows/_shared/sdk_shim``) and bridged into + # the API by whatever emitter that deployment wires up. They carry the + # full envelope in ``payload``; SSE surfaces them verbatim so the + # frontend can drive the DAG / log view directly off the protocol + # contract without an intermediate translation table. + # + # All events now share the ``response.swarm.*`` namespace, so the + # prefix check below is sufficient — the legacy bare-type set is + # empty (kept as a frozenset for forward compat if a non-namespaced + # type ever shows up again). + _SDK_PROTOCOL_RESERVED: frozenset[str] = frozenset() + + def _is_sdk_protocol_event(self, evt_value: str) -> bool: + if evt_value.startswith(_SDK_EVENT_NS + "."): + return True + return evt_value in self._SDK_PROTOCOL_RESERVED + + def to_sse_event(self) -> dict[str, Any] | None: + """Convert to SSE-friendly dict matching frontend event interface. + + Frontend expects: { type: string, timestamp: string, data: {...} } + Returns None for internal events that should not be streamed. + """ + evt_value = str(self.event_type) + + # ── SDK Agent Protocol v1 envelope — pass through verbatim ───── + if self._is_sdk_protocol_event(evt_value): + data = dict(self.payload) + data.setdefault("source_task_id", str(self.task_id)) + return { + "type": evt_value, + "timestamp": self.timestamp.isoformat(), + "data": data, + } + + if evt_value == "task_status_changed" and self.payload.get("new_status") == "suspended": + return { + "type": "suspended", + "timestamp": self.timestamp.isoformat(), + "data": { + "task_id": str(self.task_id), + "summary": self.payload.get("message", "Task paused"), + "status": "suspended", + "source_task_id": self.payload.get("source_task_id", str(self.task_id)), + }, + } + if evt_value == "task_status_changed" and self.payload.get("new_status") == "aborted": + return { + "type": "aborted", + "timestamp": self.timestamp.isoformat(), + "data": { + "task_id": str(self.task_id), + "summary": self.payload.get("message", "Task aborted"), + "status": "aborted", + "source_task_id": self.payload.get("source_task_id", str(self.task_id)), + }, + } + + # Skip internal lifecycle events + if evt_value in self._SKIP_EVENTS: + return None + + sse_type = self._TYPE_MAP.get(evt_value, "agent_action") + + # Build data according to the expected frontend interface + source_task_id = self.payload.get("source_task_id", str(self.task_id)) + + if sse_type == "phase_change": + data = { + "phase": self.payload.get("phase", ""), + "message": self.payload.get("message", f"Entering {self.payload.get('phase', 'unknown')} phase"), + "source_task_id": source_task_id, + } + elif sse_type == "error": + data = { + "message": self.payload.get("error", str(self.payload)), + "source_task_id": source_task_id, + } + elif sse_type == "completed": + data = { + "task_id": str(self.task_id), + "summary": self.payload.get("summary", "Research completed"), + "status": self.payload.get("status", "completed"), + "source_task_id": source_task_id, + } + elif evt_value == "routing_decision": + data = { + "agent": self.payload.get("agent", "system"), + "action": "routing_decision", + "detail": self.payload.get("hints", {}).get("reason", "routing_decision"), + "initial_macro": self.payload.get("initial_macro"), + "budget": self.payload.get("budget", {}), + "escalation_policy": self.payload.get("escalation_policy", {}), + "hints": self.payload.get("hints", {}), + "features": self.payload.get("features", {}), + "source_task_id": source_task_id, + } + else: + # Special frontend event types based on trace_type + trace_type = self.payload.get("trace_type", "") + # Whitelist: events whose payload should be preserved verbatim + # (top-level ``type`` mirrors ``trace_type``). Without this, the + # generic agent_action branch below collapses the payload into a + # human-readable detail and drops the structured fields. + # + # Fan-out phase entries: the run phase emits + # ``dag_finalize_start`` + ``heavy_converge_start``; the + # reporter emits user-safe progress events. + if trace_type in ( + # ReAct + verification + skill traces. + "verification_step", "verification_complete", + "react_think", "react_tool_call", "skill_loaded", + # swarm_heavy: user-safe phase boundary + reporter status events. + # ``dag.phase_started`` / ``dag.phase_cancelled`` are the + # pipeline-agnostic phase brackets; + # ``dag_finalize_*`` brackets the DAG + # finalize sub-phase. + "dag.phase_started", "dag.phase_cancelled", + "dag.phase_timeout", + "dag_finalize_start", + "dag_finalize.progress", "heavy_converge_start", + "verify.started", "verify.progress", "verify.done", + "verify.degraded", + "outline.started", "outline.progress", "outline.submitted", + "report.started", "report.progress", "report.submitted", + "report.degraded", + "report.citations.started", "report.citations.refining", + "report.citations.ready", "stt.progress", + ): + trace_data = dict(self.payload) + trace_data.setdefault("source_task_id", source_task_id) + return { + "type": trace_type, + "timestamp": self.timestamp.isoformat(), + "data": trace_data, + } + + # agent_action — build a human-readable detail + agent = self.payload.get("agent", self.from_agent or self.payload.get("from_agent", "system")) + action = self.payload.get("action", evt_value) + # Prefer specific detail fields over raw payload dump + if "detail" in self.payload: + detail = self.payload["detail"] + elif self.payload.get("tool_name"): + detail = f"Called {self.payload['tool_name']}" + if self.payload.get("input_data"): + detail += f": {self.payload['input_data'][:100]}" + elif self.message_type or self.payload.get("message_type"): + from_a = self.from_agent or self.payload.get("from_agent", "?") + to_a = self.to_agent or self.payload.get("to_agent", "?") + msg_type = self.message_type or self.payload.get("message_type", evt_value) + detail = f"{from_a} → {to_a}: {msg_type}" + elif "output_preview" in self.payload: + detail = self.payload["output_preview"][:2000] + else: + detail = action + data = { + "agent": agent, + "action": action, + "detail": detail, + "source_task_id": source_task_id, + } + + return { + "type": sse_type, + "timestamp": self.timestamp.isoformat(), + "data": data, + } diff --git a/agent_core/models/task_budget.py b/agent_core/models/task_budget.py new file mode 100644 index 0000000..809c8bc --- /dev/null +++ b/agent_core/models/task_budget.py @@ -0,0 +1,133 @@ +"""Budget models for the orchestrator runtime. + +TaskBudget constrains cost, depth, and expensive primitives. +BudgetState tracks consumption at runtime. +BudgetCharge is the unit of budget consumption per primitive call. +""" + +from __future__ import annotations + +from pydantic import BaseModel, Field + + +class TaskBudget(BaseModel): + """Allocated budget for a single research task. + + Uses soft limits (max_debate_rounds=0 means "prefer not to", + not "absolutely forbidden") to avoid recreating the old + topology-commit problem. + """ + + max_tokens: int = 500_000 + max_cost_usd: float | None = None + max_wall_time_s: int = 300 + max_parallel: int = 3 + max_depth: int = 1 # v1: always 1 (no true recursion) + max_search_calls: int = 20 + max_verify_passes: int = 5 + max_debate_rounds: int = 0 + default_model_tier: str = "medium" # "light" | "medium" | "strong" + role_tiers: dict[str, str] = Field(default_factory=dict) + + +class BudgetCharge(BaseModel): + """Unit of budget consumption from a single primitive call.""" + + primitive: str + llm_calls: int = 0 + search_calls: int = 0 + verify_passes: int = 0 + debate_rounds: int = 0 + tokens: int = 0 + cost_usd: float = 0.0 + wall_time_s: int = 0 + + +class BudgetState(BaseModel): + """Runtime budget tracking — allocated vs consumed.""" + + allocated: TaskBudget = Field(default_factory=TaskBudget) + llm_calls_used: int = 0 + search_calls_used: int = 0 + verify_passes_used: int = 0 + debate_rounds_used: int = 0 + tokens_used: int = 0 + cost_usd_used: float = 0.0 + wall_time_s_used: int = 0 + exhausted: bool = False + warnings: list[str] = Field(default_factory=list) + + def charge(self, c: BudgetCharge) -> None: + """Apply a charge and check exhaustion.""" + self.llm_calls_used += c.llm_calls + self.search_calls_used += c.search_calls + self.verify_passes_used += c.verify_passes + self.debate_rounds_used += c.debate_rounds + self.tokens_used += c.tokens + self.cost_usd_used += c.cost_usd + self.wall_time_s_used += c.wall_time_s + self._check_exhaustion() + + def _check_exhaustion(self) -> None: + a = self.allocated + if self.search_calls_used >= a.max_search_calls: + self.exhausted = True + self.warnings.append("search_calls exhausted") + if self.verify_passes_used >= a.max_verify_passes: + self.warnings.append("verify_passes exhausted") + if self.tokens_used >= a.max_tokens: + self.exhausted = True + self.warnings.append("token budget exhausted") + if a.max_cost_usd and self.cost_usd_used >= a.max_cost_usd: + self.exhausted = True + self.warnings.append("cost budget exhausted") + if a.max_wall_time_s > 0 and self.wall_time_s_used >= a.max_wall_time_s: + self.exhausted = True + self.warnings.append("wall time exhausted") + + def escalate(self, new_budget: TaskBudget) -> list[str]: + """Upgrade budget allocation in-place. Returns list of what changed. + + Only upgrades fields that are higher in new_budget than current. + Preserves consumption counters — just raises the ceilings. + """ + changes: list[str] = [] + a = self.allocated + for field in ( + "max_tokens", "max_wall_time_s", "max_parallel", + "max_search_calls", "max_verify_passes", "max_debate_rounds", + ): + old_val = getattr(a, field) + new_val = getattr(new_budget, field) + if new_val > old_val: + setattr(a, field, new_val) + changes.append(f"{field}: {old_val} → {new_val}") + # Model tier upgrade + tier_order = {"light": 0, "medium": 1, "strong": 2} + if tier_order.get(new_budget.default_model_tier, 0) > tier_order.get(a.default_model_tier, 0): + old_tier = a.default_model_tier + a.default_model_tier = new_budget.default_model_tier + changes.append(f"model_tier: {old_tier} → {new_budget.default_model_tier}") + # Role tiers: merge (only upgrade, never downgrade) + for role, tier in new_budget.role_tiers.items(): + current = a.role_tiers.get(role, "light") + if tier_order.get(tier, 0) > tier_order.get(current, 0): + a.role_tiers[role] = tier + changes.append(f"role_tier[{role}]: {current} → {tier}") + # Reset exhaustion if we have headroom now + if changes: + self.exhausted = False + self.warnings = [w for w in self.warnings if "exhausted" not in w] + return changes + + @property + def remaining_search_calls(self) -> int: + return max(0, self.allocated.max_search_calls - self.search_calls_used) + + @property + def remaining_verify_passes(self) -> int: + return max(0, self.allocated.max_verify_passes - self.verify_passes_used) + + @property + def remaining_debate_rounds(self) -> int: + return max(0, self.allocated.max_debate_rounds - self.debate_rounds_used) diff --git a/agent_core/protocols.py b/agent_core/protocols.py index a8584f6..9be32d8 100644 --- a/agent_core/protocols.py +++ b/agent_core/protocols.py @@ -3,13 +3,51 @@ from __future__ import annotations import time -from collections.abc import Mapping, Sequence +from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import dataclass, field from typing import Any, Protocol, runtime_checkable from agent_core.llm import LLMClient +@runtime_checkable +class EventSink(Protocol): + """Append-only event destination used by portable runtime components.""" + + async def append( + self, + task_id: Any = "", + event_type: Any = None, + payload: dict[str, Any] | None = None, + agent_role: str = "system", + ) -> Any: ... + + def replay(self, task_id: str) -> AsyncIterator[Any]: ... + + +@runtime_checkable +class EventReader(Protocol): + """Cursored event reads required by durable inter-agent messaging.""" + + async def get_events( + self, + task_id: Any, + event_type: Any = None, + after_id: int = 0, + limit: int | None = None, + ) -> list[Any]: ... + + async def get_events_for_agent( + self, + to_agent: str, + after_id: int = 0, + limit: int = 50, + *, + task_id: str | Any | None = None, + message_type: str | None = None, + ) -> list[Any]: ... + + @runtime_checkable class TraceSink(Protocol): async def log_llm_call( @@ -107,6 +145,11 @@ async def run_on_error( ) -> Exception | None: ... +@runtime_checkable +class LLMWrapper(Protocol): + def wrap_llm(self, llm: LLMClient, *, role_id: str) -> LLMClient: ... + + @runtime_checkable class SubAgentProfileRegistry(Protocol): def register_sub_agent_profiles( @@ -148,8 +191,11 @@ def reload(self) -> None: ... __all__ = [ + "EventReader", + "EventSink", "ExecutionMiddleware", "LLMResourceProvider", + "LLMWrapper", "PhaseContext", "PhaseMiddlewareChain", "Skill", diff --git a/agent_core/runtime/resources/__init__.py b/agent_core/runtime/resources/__init__.py new file mode 100644 index 0000000..614e4b4 --- /dev/null +++ b/agent_core/runtime/resources/__init__.py @@ -0,0 +1,22 @@ +"""Resource management — tool permissions and LLM routing per agent role.""" + +from agent_core.runtime.resources.llm import resolve_base_llm_for_role, wrap_llm_with_middleware +from agent_core.runtime.resources.manager import ResourceManager +from agent_core.runtime.resources.permissions import ( + get_allowed_tools_for_role, + get_effective_tools_for_role, +) +from agent_core.runtime.tool_permission import ( + ToolPermissionContext, + from_execution_policy, +) + +__all__ = [ + "ResourceManager", + "ToolPermissionContext", + "from_execution_policy", + "get_allowed_tools_for_role", + "get_effective_tools_for_role", + "resolve_base_llm_for_role", + "wrap_llm_with_middleware", +] diff --git a/agent_core/runtime/resources/llm.py b/agent_core/runtime/resources/llm.py new file mode 100644 index 0000000..dbb2792 --- /dev/null +++ b/agent_core/runtime/resources/llm.py @@ -0,0 +1,54 @@ +"""LLM-resolution helpers for ResourceManager.""" + +from __future__ import annotations + +from collections.abc import Callable + +from agent_core.llm import LLMClient +from agent_core.models.agent_definition import AgentDefinition +from agent_core.protocols import LLMWrapper +from agent_core.runtime.registries import services as service_registry + + +def resolve_base_llm_for_role( + *, + default_llm: LLMClient, + role_id: str | None, + cache: dict[str, LLMClient], + role_llm_factory: Callable[[AgentDefinition], LLMClient] | None = None, +) -> LLMClient: + """Resolve the base LLM for a role without middleware wrapping.""" + if role_id is None: + return default_llm + + if role_id in cache: + return cache[role_id] + + from agent_core.runtime.registries.agents import AgentRegistry + + try: + agent_reg = service_registry.get(AgentRegistry) + defn = agent_reg.get(role_id) + if defn.model and role_llm_factory is not None: + role_llm = role_llm_factory(defn) + cache[role_id] = role_llm + return role_llm + except (KeyError, RuntimeError, ImportError): + pass + + return default_llm + + +def wrap_llm_with_middleware( + llm: LLMClient, + *, + role_id: str, +) -> LLMClient: + """Wrap the LLM with an optional LLM wrapper service.""" + try: + wrapper = service_registry.get_optional(LLMWrapper) + if wrapper is not None: + return wrapper.wrap_llm(llm, role_id=role_id) + except (ImportError, RuntimeError): + pass + return llm diff --git a/agent_core/runtime/resources/manager.py b/agent_core/runtime/resources/manager.py new file mode 100644 index 0000000..f31ec27 --- /dev/null +++ b/agent_core/runtime/resources/manager.py @@ -0,0 +1,188 @@ +"""Resource manager — tool permissions and LLM routing. + +Tool permissions are driven by AgentRegistry (dynamic roles) and per-call +``ToolPermissionContext`` (built from ``NodeExecutionPolicy`` by pipeline +nodes). The manager itself knows nothing about workflow phases. +""" + +from __future__ import annotations + +from collections.abc import Callable + +from agent_core.errors import PermissionDenied +from agent_core.llm import LLMClient +from agent_core.models.agent_definition import AgentDefinition +from agent_core.runtime.resources.llm import ( + resolve_base_llm_for_role, + wrap_llm_with_middleware, +) +from agent_core.runtime.resources.permissions import get_effective_tools_for_role +from agent_core.runtime.tool_permission import ToolPermissionContext +from agent_core.tool import Tool + + +class ResourceManager: + """Manages tool permissions and LLM routing per agent role.""" + + def __init__( + self, + llm: LLMClient, + tools: dict[str, Tool], + *, + role_llm_factory: Callable[[AgentDefinition], LLMClient] | None = None, + ) -> None: + self._llm = llm + self._tools = tools # name → tool instance + self._role_llms: dict[str, LLMClient] = {} # cached per-role LLMs + self._role_llm_factory = role_llm_factory + # Process-wide allow/deny policy layered under every per-call + # ``permission_context``. Driven by config injection (Agent Protocol + # ``Request.config['tools']`` / agent_team profile ``tools:``) so an + # operator can switch a tool (e.g. ``web_search``) off for a run + # without editing role definitions. ``None`` = no global restriction. + self._global_policy: ToolPermissionContext | None = None + + @property + def llm(self) -> LLMClient: + return self._llm + + @property + def global_tool_policy(self) -> ToolPermissionContext | None: + """The active process-wide tool policy, or ``None`` if unset.""" + return self._global_policy + + def set_global_tool_policy( + self, policy: ToolPermissionContext | None, + ) -> None: + """Install (or clear) the process-wide tool allow/deny policy. + + Callers set this unconditionally per run — passing an empty policy + or ``None`` clears it — so a prior run's policy never leaks into the + next on a reused manager. + """ + if policy is not None and policy.is_empty(): + policy = None + self._global_policy = policy + + def _effective_context( + self, permission_context: ToolPermissionContext | None, + ) -> ToolPermissionContext | None: + """Layer the global policy under a per-call permission context.""" + if self._global_policy is None: + return permission_context + if permission_context is None: + return self._global_policy + return self._global_policy.merge(permission_context) + + @property + def all_tools(self) -> dict[str, Tool]: + return dict(self._tools) + + def get_tools_for_role( + self, + role_id: str, + *, + permission_context: ToolPermissionContext | None = None, + ) -> list[Tool]: + """Return the tools permitted for a role under the given permission context.""" + allowed = get_effective_tools_for_role( + role_id=role_id, + permission_context=self._effective_context(permission_context), + ) + return [t for name, t in self._tools.items() if name in allowed] + + def get_tool_for_role( + self, + role_id: str, + tool_name: str, + *, + permission_context: ToolPermissionContext | None = None, + ) -> Tool | None: + """Return a single permitted tool, or None if unavailable for the role.""" + if not self.check_permission( + role_id, + tool_name, + permission_context=permission_context, + ): + return None + return self._tools.get(tool_name) + + def get_tool_names_for_role( + self, + role_id: str, + *, + permission_context: ToolPermissionContext | None = None, + ) -> list[str]: + allowed = get_effective_tools_for_role( + role_id=role_id, + permission_context=self._effective_context(permission_context), + ) + return [name for name in self._tools if name in allowed] + + def check_permission( + self, + role_id: str, + tool_name: str, + *, + permission_context: ToolPermissionContext | None = None, + ) -> bool: + """Check if a tool is permitted for a role under the given permission context.""" + allowed = get_effective_tools_for_role( + role_id=role_id, + permission_context=self._effective_context(permission_context), + ) + return tool_name in allowed + + def require_permission( + self, + role_id: str, + tool_name: str, + *, + permission_context: ToolPermissionContext | None = None, + ) -> None: + if not self.check_permission( + role_id, + tool_name, + permission_context=permission_context, + ): + raise PermissionDenied(role_id, tool_name) + + def get_llm(self, role_id: str | None = None) -> LLMClient: + """Get LLM for a role. Uses per-role model if configured, else default. + + If an LLMMiddlewareChain is registered, wraps the returned LLM + in an LLMProxy so all chat/stream calls pass through the chain. + This is transparent to callers — pipeline nodes need zero changes. + """ + base_llm = self._resolve_base_llm(role_id) + return wrap_llm_with_middleware( + base_llm, + role_id=role_id or "default", + ) + + def get_raw_llm(self, role_id: str | None = None) -> LLMClient: + """Get LLM WITHOUT middleware wrapping. Use for internal operations + (e.g., summarization) to avoid infinite recursion.""" + return self._resolve_base_llm(role_id) + + def set_role_llm(self, role_id: str, llm: LLMClient) -> None: + """Register a base LLM (no middleware) under ``role_id``. + + Downstream ``get_llm(role_id)`` calls return this LLM wrapped with + middleware. Used by workflow nodes that build per-task LLMs from a + profile (the ``agent_team`` main agent registers ``swarm_main``) so + peer nodes in the same DAG can pick up the same LLM via + ``ctx.stream_llm(role="swarm_main")`` instead of falling + back to the default LLM, which may not be reachable on the + profile's gateway / routing group. + """ + self._role_llms[role_id] = llm + + def _resolve_base_llm(self, role_id: str | None) -> LLMClient: + """Resolve the base LLM for a role (no middleware).""" + return resolve_base_llm_for_role( + default_llm=self._llm, + role_id=role_id, + cache=self._role_llms, + role_llm_factory=self._role_llm_factory, + ) diff --git a/agent_core/runtime/resources/permissions.py b/agent_core/runtime/resources/permissions.py new file mode 100644 index 0000000..4b5b1cf --- /dev/null +++ b/agent_core/runtime/resources/permissions.py @@ -0,0 +1,45 @@ +"""Tool-permission helpers for ResourceManager.""" + +from __future__ import annotations + +import logging + +from agent_core.errors import ServiceNotRegistered +from agent_core.runtime.registries import services as service_registry +from agent_core.runtime.tool_permission import ToolPermissionContext + +logger = logging.getLogger(__name__) + + +def get_allowed_tools_for_role(role_id: str) -> set[str]: + """Return the role's allowed tool names from AgentRegistry. + + Fail-closed: unknown role or unavailable registry returns an empty set. + """ + from agent_core.runtime.registries.agents import AgentRegistry + + try: + agent_reg = service_registry.get(AgentRegistry) + return set(agent_reg.get_tools_for(role_id)) + except (KeyError, RuntimeError, ServiceNotRegistered): + return set() + + +def get_effective_tools_for_role( + *, + role_id: str, + permission_context: ToolPermissionContext | None = None, +) -> set[str]: + """Return the effective tool set for a role under the active policy.""" + role_tools = get_allowed_tools_for_role(role_id) + if permission_context is None: + return role_tools + effective = permission_context.filter(role_tools) + if effective != role_tools: + blocked = role_tools - effective + if blocked: + logger.debug( + "Tool policy blocked tools for role '%s': %s", + role_id, blocked, + ) + return effective diff --git a/agent_core/scheduling/__init__.py b/agent_core/scheduling/__init__.py index 2646a0e..c9ed5c6 100644 --- a/agent_core/scheduling/__init__.py +++ b/agent_core/scheduling/__init__.py @@ -1 +1,25 @@ -"""Portable pipeline and topology registries.""" +"""Portable scheduling engine and pipeline/topology registries.""" + +from agent_core.scheduling.scheduler import ( + ProcessManager, + Scheduler, + SchedulerTask, + TaskWallTimeExceeded, + resolve_wall_time_s, +) +from agent_core.scheduling.workflow_defaults import ( + clear_workflow_defaults, + get_workflow_default, + register_workflow_defaults, +) + +__all__ = [ + "ProcessManager", + "Scheduler", + "SchedulerTask", + "TaskWallTimeExceeded", + "clear_workflow_defaults", + "get_workflow_default", + "register_workflow_defaults", + "resolve_wall_time_s", +] diff --git a/agent_core/scheduling/scheduler.py b/agent_core/scheduling/scheduler.py new file mode 100644 index 0000000..09f5673 --- /dev/null +++ b/agent_core/scheduling/scheduler.py @@ -0,0 +1,503 @@ +# pyright: reportUnknownArgumentType=false, reportUnknownVariableType=false +"""DAG Scheduler — wraps MiniDAG execution with OS semantics.""" + +from __future__ import annotations + +import asyncio +import logging +import os +from collections.abc import AsyncIterator +from typing import Any, Protocol +from typing import Any as _GraphType # MiniDAGRunner via duck typing + +from agent_core.events import EventType +from agent_core.protocols import EventSink +from agent_core.runtime.dag.graph_builder import resolve_state_type +from agent_core.types import TaskId, TaskStatus + +logger = logging.getLogger(__name__) + + +_WALL_TIME_ENVS = ( + "AGENT_CORE_TASK_WALL_TIME_S", + "MIROHARNESS_TASK_WALL_TIME_S", + "FRONTIER_AGENT_TASK_WALL_TIME_S", +) + + +class SchedulerTask(Protocol): + status: TaskStatus + thread_id: str + pipeline_id: str | None + context: dict[str, Any] + + +class ProcessManager(Protocol): + async def get_task(self, task_id: TaskId) -> SchedulerTask: ... + + async def update_status(self, task_id: TaskId, status: TaskStatus) -> Any: ... + + async def abort_task(self, task_id: TaskId, reason: str) -> Any: ... + + async def set_error(self, task_id: TaskId, error: str) -> Any: ... + + +class TaskWallTimeExceeded(RuntimeError): + """Task exceeded its configured wall-time budget. + + Raised by :meth:`Scheduler.execute` when ``astream`` does not finish + within ``wall_time_s`` (caller-supplied or read from + ``FRONTIER_AGENT_TASK_WALL_TIME_S``). The scheduler treats this as a + budget abort rather than an implementation failure. + """ + + +def resolve_wall_time_s( + explicit: int | None, + pipeline_id: str | None = None, +) -> int | None: + """Resolve the effective wall-time budget. + + Priority: caller-supplied ``explicit`` > workflow wall-time mode > + ``FRONTIER_AGENT_TASK_WALL_TIME_S`` env var > workflow default (looked up by + ``pipeline_id``) > ``None`` (no timeout). + + Convention ``0 = unlimited`` is preserved at every layer so a + workflow can declare "this benchmark needs no wall-time" (e.g. SWE) + in its ``defaults.py`` and any caller / env value still overrides + when supplied. + + Returns ``None`` (no deadline) or a positive int. + """ + if explicit is not None: + return explicit if explicit > 0 else None + + # Some workflows own a phase-aware *research* deadline internally so they + # can stop research and let a reporter finish. For those workflows the env + # value is consumed by the in-loop observer and must not also become the + # graph deadline. The workflow default remains an absolute backstop, + # however, so a missing profile/env deadline can never make the graph + # unbounded. An explicit caller value remains the highest-priority ceiling. + from agent_core.scheduling.workflow_defaults import get_workflow_default + + wall_mode = get_workflow_default(pipeline_id, "TASK_WALL_TIME_MODE") + if wall_mode == "soft_research": + workflow_default = get_workflow_default(pipeline_id, "TASK_WALL_TIME_S") + if isinstance(workflow_default, int): + return workflow_default if workflow_default > 0 else None + return None + + env_name = next( + (name for name in _WALL_TIME_ENVS if os.environ.get(name, "").strip()), + None, + ) + if env_name is not None: + raw = os.environ[env_name].strip() + try: + parsed = int(raw) + except ValueError: + logger.warning("Invalid %s=%r; ignoring", env_name, raw) + else: + return parsed if parsed > 0 else None + + # No caller value, no env override — fall back to the workflow's + # declared default. Lazy-imported so the scheduler stays workflow- + # agnostic at import time. + workflow_default = get_workflow_default(pipeline_id, "TASK_WALL_TIME_S") + if isinstance(workflow_default, int): + return workflow_default if workflow_default > 0 else None + return None + + +class Scheduler: + """Executes research task DAGs via the configured graph runner.""" + + def __init__( + self, + graph: _GraphType, + process_manager: ProcessManager, + event_store: EventSink, + *, + graph_builder: Any | None = None, + pipeline_registry: Any | None = None, + checkpointer: Any | None = None, + ) -> None: + self._default_graph = graph + self._pm = process_manager + self._event_store = event_store + self._graph_builder = graph_builder + self._pipeline_registry = pipeline_registry + self._checkpointer = checkpointer + self._compiled_cache: dict[str, _GraphType] = {} + + def _get_graph(self, pipeline_id: str | None = None) -> _GraphType: + """Get compiled graph for a pipeline, or the default. + + All pipelines (including "deep_research") are now built via + DynamicGraphBuilder for uniform middleware wrapping. + The _default_graph is the middleware-wrapped deep_research graph + built at app startup. + """ + if not pipeline_id: + return self._default_graph + + # Check if we have a dynamically compiled version + if pipeline_id in self._compiled_cache: + return self._compiled_cache[pipeline_id] + + # Pre-built defaults (already middleware-wrapped at startup). Benchmark + # sessions may construct Scheduler with no default graph and rely on + # dynamic PipelineSpec compilation instead. + if pipeline_id in ("deep_research", "react_research") and self._default_graph is not None: + return self._default_graph + + if not self._pipeline_registry or not self._graph_builder: + return self._default_graph + + spec = self._pipeline_registry.get(pipeline_id) + self._register_pipeline_agents(spec) + # The spec declares its own state type (a TypedDict with Annotated + # fan-in reducers). A spec that declares nothing now falls back to + # AgentCore's ``BaseTaskState`` — the universal fields plus their + # fan-in reducers — rather than the bare ``dict`` this used to get. + state_type = resolve_state_type(spec) + compiled = self._graph_builder.compile( + spec, state_type=state_type, checkpointer=self._checkpointer, + ) + self._compiled_cache[pipeline_id] = compiled + + return self._compiled_cache[pipeline_id] + + def _register_pipeline_agents(self, spec: Any) -> None: + """Auto-register agent definitions bundled with a PipelineSpec.""" + from agent_core.runtime.registries import services as kernel_registry + from agent_core.runtime.registries.agents import AgentRegistry + + if not hasattr(spec, "agent_definitions") or not spec.agent_definitions: + return + try: + agent_reg = kernel_registry.get(AgentRegistry) + for defn in spec.agent_definitions: + if not agent_reg.has(defn.role_id): + agent_reg.register(defn) + logger.info("Auto-registered agent '%s' from pipeline '%s'", defn.role_id, spec.pipeline_id) + except Exception as e: + logger.warning("Failed to auto-register agents: %s", e) + + async def execute( + self, + task_id: TaskId, + input_data: dict[str, Any] | None, + *, + pipeline_id: str | None = None, + config_override: dict[str, Any] | None = None, + execution_token: str | None = None, + wall_time_s: int | None = None, + ) -> AsyncIterator[tuple[str, Any]]: + """Execute a research task, yielding stream events. + + Wraps the graph runner's `astream()` with OS-level lifecycle management. + + ``wall_time_s`` is a hard wall-time ceiling for graph execution. + Resolution priority: caller-supplied value > workflow wall-time mode > + ``FRONTIER_AGENT_TASK_WALL_TIME_S`` env var > workflow default + (``workflows//defaults.py:TASK_WALL_TIME_S``, looked up + by ``pipeline_id``) > no deadline. ``0`` means unlimited at + every layer (matches the existing HLE budget convention). Workflows + declaring ``TASK_WALL_TIME_MODE="soft_research"`` consume the env/profile + budget inside their research loop and exclude reporter finalization from + this scheduler-wide hard ceiling unless the caller explicitly supplies + ``wall_time_s``. + On expiry the scheduler marks the task ABORTED and raises + :class:`TaskWallTimeExceeded`. + """ + task = await self._pm.get_task(task_id) + if task.status != TaskStatus.RUNNING: + await self._pm.update_status(task_id, TaskStatus.RUNNING) + + graph = self._get_graph(pipeline_id) + + config = config_override or { + "configurable": { + "thread_id": task.thread_id, + }, + } + + wall_time = resolve_wall_time_s(wall_time_s, pipeline_id=pipeline_id) + + try: + async for chunk in self._astream_with_wall_time( + graph, input_data, config, + task_id=task_id, wall_time=wall_time, + ): + latest_task = await self._pm.get_task(task_id) + if self._is_stale_execution(latest_task, execution_token): + logger.info( + "Task %s execution token %s is stale before delivering chunk; stopping runner", + task_id, + execution_token, + ) + return + + # chunk is {"node_name": {state_updates}} + # Note: middleware hooks are applied inside node wrappers + # (see DynamicGraphBuilder.wrap_with_middleware), NOT here. + # By the time we receive a chunk, the node has already executed + # with full middleware coverage (before/after/error). + for node_name in chunk: + phase = node_name + logger.info("Task %s: phase → %s", task_id, phase) + + yield "updates", chunk + + latest_task = await self._pm.get_task(task_id) + if self._is_stale_execution(latest_task, execution_token): + logger.info( + "Task %s execution token %s was superseded after chunk; stopping runner", + task_id, + execution_token, + ) + return + if latest_task.status == TaskStatus.SUSPENDED: + logger.info("Task %s suspended after phase boundary; stopping execution", task_id) + return + if latest_task.status == TaskStatus.ABORTED: + logger.info("Task %s aborted after phase boundary; stopping execution", task_id) + return + + latest_task = await self._pm.get_task(task_id) + if self._is_stale_execution(latest_task, execution_token): + logger.info("Task %s execution token %s is stale at finalization; stopping runner", task_id, execution_token) + return + if latest_task.status == TaskStatus.SUSPENDED: + logger.info("Task %s suspended; skipping completion transition", task_id) + return + if latest_task.status == TaskStatus.ABORTED: + logger.info("Task %s finished after abort; skipping completion transition", task_id) + return + if latest_task.status == TaskStatus.COMPLETED: + logger.info("Task %s already completed by another runner; skipping duplicate completion", task_id) + return + + final_state = await self._get_state_for_config(graph, config) + if not self._has_terminal_output(final_state): + raise RuntimeError( + self._format_no_terminal_error(task_id, final_state) + ) + + latest_task = await self._pm.get_task(task_id) + if self._is_stale_execution(latest_task, execution_token): + logger.info( + "Task %s execution token %s was superseded after final state read; skipping completion", + task_id, + execution_token, + ) + return + if latest_task.status == TaskStatus.COMPLETED: + logger.info("Task %s already completed after final state read; skipping duplicate completion", task_id) + return + + await self._pm.update_status(task_id, TaskStatus.COMPLETED) + await self._event_store.append( + task_id=task_id, + event_type=EventType.REPORT_GENERATED, + payload={"summary": "Research completed successfully"}, + ) + + except Exception as e: + try: + latest_task = await self._pm.get_task(task_id) + except Exception: + latest_task = None + + if latest_task and self._is_stale_execution(latest_task, execution_token): + logger.info("Task %s stale execution token %s raised after being superseded; suppressing failure", task_id, execution_token) + return + if latest_task and latest_task.status == TaskStatus.COMPLETED: + logger.info("Task %s raised after another runner already completed it; suppressing failure", task_id) + return + if latest_task and latest_task.status == TaskStatus.SUSPENDED: + logger.info("Task %s suspended during execution; suppressing failure transition", task_id) + return + if latest_task and latest_task.status == TaskStatus.ABORTED: + logger.info("Task %s aborted during execution; suppressing failure transition", task_id) + return + if isinstance(e, TaskWallTimeExceeded): + logger.warning("Task %s wall-time exceeded; aborting", task_id) + await self._pm.abort_task(task_id, str(e)) + raise + + logger.exception("Task %s failed: %s", task_id, e) + await self._pm.set_error(task_id, str(e)) + raise + + async def _astream_with_wall_time( + self, + graph: _GraphType, + input_data: dict[str, Any] | None, + config: dict[str, Any], + *, + task_id: TaskId, + wall_time: int | None, + ) -> AsyncIterator[dict[str, Any]]: + """Yield ``graph.astream`` chunks under a wall-time deadline. + + ``wall_time=None`` is the no-deadline fast path. On expiry we + raise :class:`TaskWallTimeExceeded`; the outer ``execute`` arm + does the abort + status transition. + """ + loop = asyncio.get_running_loop() + deadline = loop.time() + wall_time if wall_time is not None else None + stream = graph.astream( + input_data, config=config, stream_mode="updates", + ) + try: + while True: + try: + if deadline is None: + chunk = await stream.__anext__() + else: + remaining = deadline - loop.time() + if remaining <= 0: + raise TaskWallTimeExceeded( + f"task {task_id} wall-time exceeded " + f"after {wall_time}s" + ) + chunk = await asyncio.wait_for( + stream.__anext__(), timeout=remaining, + ) + except StopAsyncIteration: + return + except TimeoutError as e: + raise TaskWallTimeExceeded( + f"task {task_id} wall-time exceeded " + f"after {wall_time}s" + ) from e + yield chunk + finally: + aclose = getattr(stream, "aclose", None) + if aclose is not None: + await aclose() + + async def get_state(self, task_id: TaskId) -> dict[str, Any]: + """Get the current state snapshot for a task. + + Tries the task's actual pipeline graph first, falls back to default, + then falls back to raw checkpoint read. + """ + task = await self._pm.get_task(task_id) + pid = getattr(task, "pipeline_id", None) + if pid == "auto" or pid is None: + pid = "deep_research" + config = {"configurable": {"thread_id": task.thread_id}} + + # Try with the specific pipeline's graph + for try_pid in [pid, "deep_research"] if pid != "deep_research" else ["deep_research"]: + try: + graph = self._get_graph(try_pid) + snapshot = await graph.aget_state(config) + if snapshot and snapshot.values: + return snapshot.values + except Exception: + continue + + # Try all compiled pipeline graphs (covers any pipeline not captured above, + # e.g. swe_solver compiled on demand; MiniDAGRunner.aget_state is a pure + # in-memory dict lookup so this is safe and cheap) + for cached_graph in self._compiled_cache.values(): + try: + snapshot = await cached_graph.aget_state(config) + if snapshot and snapshot.values: + return snapshot.values + except Exception: + continue + + # Last resort: read checkpoint directly from checkpointer + if self._checkpointer: + try: + checkpoint_tuple = await self._checkpointer.aget_tuple(config) + if checkpoint_tuple and checkpoint_tuple.checkpoint: + channel_values = checkpoint_tuple.checkpoint.get("channel_values", {}) + if channel_values: + return channel_values + except Exception as e: + logger.warning("Direct checkpoint read failed for %s: %s", task_id, e) + + return {} + + async def _get_state_for_config( + self, + graph: _GraphType | None, + config: dict[str, Any], + ) -> dict[str, Any]: + """Best-effort state read for a specific checkpoint config.""" + if graph is not None: + try: + snapshot = await graph.aget_state(config) + if snapshot and snapshot.values: + return snapshot.values + except Exception: + pass + + if self._checkpointer: + try: + checkpoint_tuple = await self._checkpointer.aget_tuple(config) + if checkpoint_tuple and checkpoint_tuple.checkpoint: + channel_values = checkpoint_tuple.checkpoint.get("channel_values", {}) + if "__root__" in channel_values and isinstance(channel_values["__root__"], dict): + return channel_values["__root__"] + if channel_values: + return channel_values + except Exception: + pass + return {} + + @staticmethod + def _has_terminal_output(state: dict[str, Any]) -> bool: + """A task is terminal only when it has a reader-facing output. + + ``report`` is the canonical field for multi-phase research + pipelines. Single-solver topologies (solo, debate) sometimes + populate ``final_content`` instead — both are valid terminal + outputs for the user-facing purpose. + """ + for key in ("report", "final_content"): + value = state.get(key) + if isinstance(value, str) and value.strip(): + return True + if isinstance(value, dict) and value: + return True + return False + + @staticmethod + def _format_no_terminal_error( + task_id: Any, state: dict[str, Any], + ) -> str: + """Build a human-actionable error for empty pipeline output. + + Surfaces the phase the task was in, any accumulated errors, + and truncates long values so the message stays readable when + it lands in a DB column and ultimately a UI banner. + """ + phase = state.get("current_phase") or "?" + bits = [ + f"Task {task_id} produced no final answer at phase '{phase}'." + ] + errors = state.get("errors") + if isinstance(errors, list) and errors: + # Cap at 3 error strings, each ≤200 chars. Anything longer + # is probably a stack trace already logged by the node. + shortened = [str(e)[:200] for e in errors[:3]] + bits.append("Errors: " + " | ".join(shortened)) + iteration = state.get("iteration_count") + if isinstance(iteration, int) and iteration > 0: + bits.append(f"iterations={iteration}") + return " ".join(bits) + + @staticmethod + def _is_stale_execution(task: Any, execution_token: str | None) -> bool: + """Detect whether a newer runner has taken over this task.""" + if not execution_token: + return False + context = getattr(task, "context", {}) or {} + return context.get("active_execution_token") != execution_token diff --git a/agent_core/scheduling/workflow_defaults.py b/agent_core/scheduling/workflow_defaults.py new file mode 100644 index 0000000..dc944f3 --- /dev/null +++ b/agent_core/scheduling/workflow_defaults.py @@ -0,0 +1,38 @@ +"""Product-neutral workflow default registry.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +_WORKFLOW_DEFAULTS: dict[str, dict[str, Any]] = {} + + +def register_workflow_defaults( + pipeline_ids: str | tuple[str, ...], + defaults: Mapping[str, Any], +) -> None: + """Register an immutable snapshot of product-owned workflow defaults.""" + ids = (pipeline_ids,) if isinstance(pipeline_ids, str) else pipeline_ids + snapshot = dict(defaults) + for pipeline_id in ids: + _WORKFLOW_DEFAULTS[pipeline_id] = snapshot + + +def clear_workflow_defaults() -> None: + """Clear registrations, primarily for isolated tests and app teardown.""" + _WORKFLOW_DEFAULTS.clear() + + +def get_workflow_default(pipeline_id: str | None, attr: str) -> Any: + """Return one registered product default, or ``None`` when absent.""" + if not pipeline_id: + return None + return _WORKFLOW_DEFAULTS.get(pipeline_id, {}).get(attr) + + +__all__ = [ + "clear_workflow_defaults", + "get_workflow_default", + "register_workflow_defaults", +] diff --git a/agent_core/tool.py b/agent_core/tool.py new file mode 100644 index 0000000..7d09266 --- /dev/null +++ b/agent_core/tool.py @@ -0,0 +1,282 @@ +# pyright: reportMissingTypeArgument=false, reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownParameterType=false, reportUnknownVariableType=false +"""Tool definitions — function + OpenAI function-schema.""" + +from __future__ import annotations + +import inspect +import re +import typing +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any, Union, get_args, get_origin, overload + +ToolFn = Callable[..., Awaitable[Any]] + + +@dataclass +class Tool: + """A callable tool with an OpenAI-compatible JSON-schema signature.""" + + name: str + description: str + parameters: dict[str, Any] # JSON schema (OpenAI ``function.parameters``) + func: ToolFn + metadata: dict[str, Any] = field(default_factory=dict) + + async def ainvoke(self, args: dict[str, Any]) -> Any: + """Run the tool with keyword args. ``func`` must be async.""" + return await self.func(**args) + + def to_openai_schema(self) -> dict[str, Any]: + """One entry of OpenAI's ``tools=`` array. + + A workflow that must reproduce an exact served byte-shape (e.g. the + FastMCP-aligned profiles) pins the final dict in + ``metadata["openai_schema"]``; this method returns it verbatim then. + """ + if (spec := self.metadata.get("openai_schema")) is not None: + return spec + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": self.parameters, + }, + } + + +_PRIMITIVE_TYPES: dict[Any, str] = { + str: "string", int: "integer", float: "number", bool: "boolean", + type(None): "null", +} + + +def _schema_for_type(t: Any) -> dict[str, Any]: + """Build a JSON-schema fragment for a Python type hint.""" + if t is Any: + return {} + if t in _PRIMITIVE_TYPES: + return {"type": _PRIMITIVE_TYPES[t]} + + # Pydantic models are the preferred way to describe structured tool + # arguments. Keep the dependency duck-typed here: ``core.tool`` does not + # need to import Pydantic, while any BaseModel subclass still contributes + # its full object schema (properties, required fields, descriptions, and + # ``additionalProperties`` policy) to the surrounding tool schema. + model_json_schema = getattr(t, "model_json_schema", None) + if callable(model_json_schema): + model_schema = model_json_schema() + if isinstance(model_schema, dict): + return model_schema + + origin = get_origin(t) + if origin is None: + # Unannotated / unknown — accept anything. + return {"type": "string"} + + if origin in (list, set, tuple): + args = get_args(t) + if args: + return {"type": "array", "items": _schema_for_type(args[0])} + return {"type": "array"} + + if origin is dict: + return {"type": "object"} + + if origin in (Union, typing.Union): # type: ignore[attr-defined] + non_none = [a for a in get_args(t) if a is not type(None)] + nullable = len(non_none) != len(get_args(t)) + schema: dict[str, Any] + if len(non_none) == 1: + schema = _schema_for_type(non_none[0]) + else: + schema = {"anyOf": [_schema_for_type(a) for a in non_none]} + if nullable: + # ``default: null`` is FastMCP's idiom; many proxies accept it. + schema.setdefault("default", None) + return schema + + # Python 3.10+ ``X | Y`` resolves to ``types.UnionType``; treated the + # same as ``typing.Union`` above when get_origin returns it. + if origin is type(int | str): # types.UnionType + # A runtime subscript of typing.Union (not an annotation), so PEP 604 + # `X | Y` syntax does not apply here. + return _schema_for_type(Union[get_args(t)]) # noqa: UP007 + + return {"type": "string"} + + +def _hoist_schema_defs( + schema: Any, + definitions: dict[str, Any], +) -> None: + """Move nested JSON Schema definitions to the parameters document root. + + Pydantic emits references such as ``#/$defs/Inner``. Those references are + document-root-relative, so leaving ``$defs`` embedded in an argument or + array-item schema produces a dangling reference. Collect every definition + while the inferred parameters document is assembled. + """ + if isinstance(schema, list): + for item in schema: + _hoist_schema_defs(item, definitions) + return + if not isinstance(schema, dict): + return + + local_defs = schema.pop("$defs", {}) + if isinstance(local_defs, dict): + for name, definition in local_defs.items(): + existing = definitions.get(name) + if existing is not None and existing != definition: + raise ValueError( + "Conflicting inferred JSON Schema definition " + f"{name!r}; provide an explicit tool parameters schema." + ) + definitions[name] = definition + for definition in local_defs.values(): + _hoist_schema_defs(definition, definitions) + + for value in schema.values(): + _hoist_schema_defs(value, definitions) + + +_ARG_HEADER_RE = re.compile(r"^\s*(Args|Arguments|Parameters):\s*$", re.MULTILINE) +_RETURNS_HEADER_RE = re.compile( + r"^\s*(Returns|Yields|Raises|Examples):\s*$", re.MULTILINE, +) + + +def _parse_docstring(doc: str) -> tuple[str, dict[str, str]]: + """Return ``(summary, {arg_name: description})`` from a Google-style docstring.""" + if not doc: + return "", {} + # ``inspect.cleandoc`` normalises indentation (strips the common leading + # whitespace shared by all lines after the first), so parsing is + # independent of how deeply the function — and therefore its docstring — + # is indented in source: a nested / method-level ``Args:`` block now + # parses the same as a module-level one. The previous fixed-width + # ``line.startswith(" " * 12)`` heuristic silently dropped EVERY arg + # description for tools whose docstring sat at >=12-space indent. + text = inspect.cleandoc(doc) + arg_match = _ARG_HEADER_RE.search(text) + if not arg_match: + return text, {} + summary = text[:arg_match.start()].rstrip() + body = text[arg_match.end():] + end = _RETURNS_HEADER_RE.search(body) + if end: + body = body[:end.start()] + + # An arg line is one at the Args-block base indent; more-indented lines + # fold in as continuations of the current arg. + body_lines = [ln for ln in body.splitlines() if ln.strip()] + base = min((len(ln) - len(ln.lstrip()) for ln in body_lines), default=0) + descriptions: dict[str, str] = {} + current = "" # name of the arg whose description is being accumulated + for line in body.splitlines(): + stripped = line.strip() + if not stripped: + continue + indent = len(line) - len(line.lstrip()) + m = re.match(r"^([A-Za-z_][\w]*)(?:\s*\([^)]+\))?\s*:\s*(.*)", stripped) + if m and indent <= base: + # Both groups are mandatory subpatterns, so they always + # participate; the ``or ""`` fallbacks are unreachable at runtime. + current = m.group(1) or "" + descriptions[current] = m.group(2) or "" + elif current: + descriptions[current] = f"{descriptions[current]} {stripped}".strip() + return summary, descriptions + + +def _infer_parameters( + fn: Callable[..., Any], + arg_docs: dict[str, str], +) -> dict[str, Any]: + """Build a JSON schema from a function's annotations + docstring.""" + hints = typing.get_type_hints(fn) + hints.pop("return", None) + sig = inspect.signature(fn) + properties: dict[str, Any] = {} + required: list[str] = [] + definitions: dict[str, Any] = {} + for name, param in sig.parameters.items(): + if name in {"self", "cls"} or param.kind in ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ): + continue + schema = _schema_for_type(hints.get(name, str)) + _hoist_schema_defs(schema, definitions) + if desc := arg_docs.get(name): + schema["description"] = desc + if param.default is inspect.Parameter.empty: + required.append(name) + elif "default" not in schema: + schema["default"] = param.default + properties[name] = schema + out: dict[str, Any] = {"type": "object", "properties": properties} + if required: + out["required"] = required + if definitions: + out["$defs"] = definitions + return out + + +@overload +def tool(fn: ToolFn, /) -> Tool: ... + + +@overload +def tool( + *, + name: str | None = ..., + description: str | None = ..., + parameters: dict[str, Any] | None = ..., +) -> Callable[[ToolFn], Tool]: ... + + +def tool( + fn: ToolFn | None = None, + *, + name: str | None = None, + description: str | None = None, + parameters: dict[str, Any] | None = None, +) -> Tool | Callable[[ToolFn], Tool]: + """Decorator that wraps an async function as a :class:`Tool`. + + The two overloads above distinguish the bare ``@tool`` form (which produces + a :class:`Tool`) from the ``@tool(...)`` factory form (which produces the + decorator). Without them every decorated symbol in the codebase infers as + the union of both, so any attribute access on one — ``.func``, + ``.description``, ``.name`` — is an error at each of the ~30 use sites. + + Usage:: + + @tool + async def web_search(query: str) -> str: + ... + + @tool(name="run_python", description="Run code in a sandbox.") + async def _runner(code: str) -> str: + ... + + When ``parameters`` is omitted, a JSON schema is inferred from the + function's type hints + Google-style docstring. + """ + + def _wrap(f: ToolFn) -> Tool: + summary, arg_docs = _parse_docstring(f.__doc__ or "") + return Tool( + name=name or f.__name__, + description=description or summary or "", + parameters=parameters or _infer_parameters(f, arg_docs), + func=f, + ) + + return _wrap(fn) if fn is not None else _wrap + + +__all__ = ["Tool", "ToolFn", "tool"] diff --git a/docs/agent-bus-boundary.md b/docs/agent-bus-boundary.md new file mode 100644 index 0000000..5339cb7 --- /dev/null +++ b/docs/agent-bus-boundary.md @@ -0,0 +1,26 @@ +# Agent-bus and scheduling boundary + +AgentCore owns the product-neutral multi-agent scheduling state machine: + +- typed inter-agent messages and cursored durable delivery; +- spawn depth, concurrency, token, and wall-time guards; +- parallel jobs and reusable serial sessions; +- cancellation, queue draining, result recovery, and idempotent fan-in; +- completion classification and report formatting; +- role-scoped tools/LLMs and fail-closed tool permissions; +- DAG lifecycle scheduling through structural backend protocols. + +The implementation accepts product decisions through constructor inputs, +runtime specs, callbacks, and protocols. In particular, products inject the +event store, process manager, pause-check factory, per-role LLM factory, +workflow defaults, observers, result adapters, and execution context setup. + +Products continue to own tool assembly, database schemas and persistence, +workflow-specific evidence harvesting/completion policy, UI rendering, +sandbox allocation, endpoint configuration, and session-retention policy. +Neither product package may be imported by AgentCore. + +The two legacy wall-time environment prefixes remain accepted after the +portable `AGENT_CORE_TASK_WALL_TIME_S` name so deployments can migrate without +a flag day. Product workflow defaults are registered explicitly rather than +discovered by importing a `workflows` package from the shared kernel. diff --git a/tests/test_agent_bus_event_sink_injection.py b/tests/test_agent_bus_event_sink_injection.py new file mode 100644 index 0000000..220edf1 --- /dev/null +++ b/tests/test_agent_bus_event_sink_injection.py @@ -0,0 +1,122 @@ +"""Phase 4 D13 surgery 1 — ``AgentBus(event_sink=...)`` injection. + +Lets the SDK assembly hand a custom EventSink straight to the bus +without going through the global service registry. The legacy +zero-arg constructor still works (falls back to +``registry.get_optional(EventStore)``). +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import patch + +import pytest + +from agent_core.components.agent_bus.bus import AgentBus + + +class _FakeSink: + """Minimal EventSink Protocol stub recording every append.""" + + def __init__(self) -> None: + self.appended: list[dict[str, Any]] = [] + + async def append( + self, + task_id: Any = "", + event_type: Any = None, + payload: dict[str, Any] | None = None, + agent_role: str = "system", + ) -> None: + self.appended.append({ + "task_id": task_id, + "event_type": event_type, + "payload": payload or {}, + "agent_role": agent_role, + }) + + def replay(self, task_id: str): # pragma: no cover — unused in this test + async def _empty(): + if False: + yield None + return _empty() + + +def test_default_constructor_falls_back_to_registry() -> None: + """Zero-arg construction still works via registry lookup.""" + bus = AgentBus() + # ``_event_sink()`` returns whatever ``registry.get_optional`` returns, + # which is None in a clean test environment. + assert bus._event_sink() is None + + +def test_injected_sink_is_returned_directly() -> None: + """Injected sink bypasses the registry entirely.""" + sink = _FakeSink() + bus = AgentBus(event_sink=sink) + assert bus._event_sink() is sink + + +def test_injected_sink_used_when_registry_also_has_one() -> None: + """Injected sink wins over registry when both are present.""" + sink = _FakeSink() + other = _FakeSink() + bus = AgentBus(event_sink=sink) + + with patch( + "agent_core.components.agent_bus.bus.registry.get_optional", + return_value=other, + ): + assert bus._event_sink() is sink + + +def test_injected_agent_registry_returns_directly() -> None: + """D13.3 — injected AgentRegistry bypasses the registry.""" + from agent_core.runtime.registries.agents import AgentRegistry + reg = AgentRegistry() + bus = AgentBus(agent_registry=reg) + assert bus._agent_registry(required=False) is reg + assert bus._agent_registry(required=True) is reg + + +def test_injected_resource_manager_returns_directly() -> None: + """D13.3 — injected ResourceManager bypasses the registry.""" + sentinel = object() + bus = AgentBus(resource_manager=sentinel) + assert bus._resource_manager(required=False) is sentinel + assert bus._resource_manager(required=True) is sentinel + + +@pytest.mark.asyncio +async def test_session_task_lifecycle_uses_injected_sink() -> None: + """End-to-end: lifecycle helpers route through the injected sink.""" + from agent_core.components.agent_bus.models import SubAgentResult, SubAgentSession + from agent_core.components.agent_bus.runtime import ( + emit_session_task_completed, + emit_session_task_submitted, + ) + from agent_core.runtime.loop.message_trimmer import NullTrimmer + + sink = _FakeSink() + session = SubAgentSession( + session_id="s1", task_id="root", name="alice", + role_id="researcher", system_prompt="", tools=[], llm=None, + trimmer=NullTrimmer(), + ) + + await emit_session_task_submitted( + session, "j1", "do the thing", event_sink=sink, + ) + await emit_session_task_completed( + session, "j1", + SubAgentResult( + question="q", role_id="researcher", + success=True, final_content="done", + ), + event_sink=sink, + ) + + assert len(sink.appended) == 2 + assert sink.appended[0]["payload"]["trace_type"] == "session_task_submitted" + assert sink.appended[1]["payload"]["trace_type"] == "session_task_completed" diff --git a/tests/test_agent_bus_session_invariant.py b/tests/test_agent_bus_session_invariant.py new file mode 100644 index 0000000..1783a23 --- /dev/null +++ b/tests/test_agent_bus_session_invariant.py @@ -0,0 +1,418 @@ +"""Tests for AgentBus session invariants: + +* ``force_finalizer`` mutates raw_result before bus absorbs into session. +* Closed boundary always has a clean final assistant message. +* Aborted boundary also has a clean final assistant message. +* ``observers_builder`` receives ``task_index``. +* ``session_task_submitted`` SSE payload includes ``is_reuse`` + ``role_id``. +""" + +from __future__ import annotations + +import contextlib +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from agent_core.components.agent_bus import ( + AgentBus, + SubAgentRuntimeSpec, +) +from agent_core.loop_types import AgentLoopResult +from agent_core.messages import ( + assistant_msg, + is_assistant_msg, + text_of, + user_msg, +) +from agent_core.runtime.loop.message_trimmer import ( + NullTrimmer, + find_final_assistant, +) +from agent_core.runtime.registries import services as registry + + +def _tool_call(call_id: str, name: str, arguments: str = "{}") -> dict: + """Build an OpenAI-wire tool_call dict.""" + return { + "id": call_id, + "type": "function", + "function": {"name": name, "arguments": arguments}, + } + + +class _FakeEventStore: + def __init__(self) -> None: + self.append = AsyncMock() + self.get_events = AsyncMock(return_value=[]) + + +@pytest.fixture(autouse=True) +def clear_registry(): + registry.clear() + yield + registry.clear() + + +async def _make_bus_with_session( + *, + name: str = "agent_a", + task_id: str = "task-1", + role_id: str = "researcher", + system_prompt: str = "You are a researcher.", + trimmer=None, + runtime_spec: SubAgentRuntimeSpec | None = None, +): + bus = AgentBus() + sid = await bus.create_session( + task_id=task_id, + name=name, + role_id=role_id, + system_prompt=system_prompt, + llm_override=MagicMock(), + tools_override=[], + trimmer=trimmer, + runtime_spec=runtime_spec, + ) + return bus, sid + + +# ── force_finalizer ─────────────────────────────────────────────────────── + + +class TestForceFinalizer: + @pytest.mark.asyncio + async def test_appends_to_session_messages_and_last_report(self, monkeypatch): + """force_finalizer mutating raw_result.messages + final_content + must land in session.messages + session.last_report.""" + + async def fake_loop(**kwargs): + im = list(kwargs.get("initial_messages") or []) + um = kwargs.get("user_message", "") + return AgentLoopResult( + messages=[*im, user_msg(um), assistant_msg("thinking", tool_calls=[_tool_call("c1", "search")])], + final_content="", # no clean final from the loop + stopped_by="max_turns", + metadata={}, + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", fake_loop, + ) + + async def force_finalizer(raw_result, _item): + raw_result.messages.append(assistant_msg("FORCED")) + raw_result.final_content = "FORCED" + return raw_result + + spec = SubAgentRuntimeSpec(force_finalizer=force_finalizer) + bus, sid = await _make_bus_with_session(runtime_spec=spec) + + await bus.submit_task_to_session(sid, "task 1") + await bus.wait_any_session("task-1", timeout=2) + + session = bus.get_session(sid) + assert session.last_report == "FORCED" + # The forced assistant message must be the boundary's final message. + start, end = session.task_boundaries[0] + assert end is not None + assert is_assistant_msg(session.messages[end]) + assert text_of(session.messages[end].get("content")) == "FORCED" + # The trimmer's invariant helper finds it. + ai = find_final_assistant(session.messages, start + 1, end) + assert ai is not None and text_of(ai.get("content")) == "FORCED" + + +# ── Closed-boundary invariant ───────────────────────────────────────────── + + +class TestClosedBoundaryInvariant: + @pytest.mark.asyncio + async def test_synth_stub_when_loop_only_has_tool_calls(self, monkeypatch): + """No clean final assistant message from the loop AND no + force_finalizer → bus appends a stub assistant message so the + boundary still has one.""" + + async def fake_loop(**kwargs): + im = list(kwargs.get("initial_messages") or []) + um = kwargs.get("user_message", "") + return AgentLoopResult( + messages=[*im, user_msg(um), assistant_msg("", tool_calls=[_tool_call("c1", "search")])], + final_content="", + stopped_by="max_turns", + metadata={}, + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", fake_loop, + ) + + bus, sid = await _make_bus_with_session() + await bus.submit_task_to_session(sid, "task 1") + await bus.wait_any_session("task-1", timeout=2) + + session = bus.get_session(sid) + start, end = session.task_boundaries[0] + assert end is not None + # Stub must exist at end, with the stopped_by reason mentioned. + assert is_assistant_msg(session.messages[end]) + assert "max_turns" in text_of(session.messages[end].get("content")) + # find_final_assistant locates it. + assert find_final_assistant(session.messages, start + 1, end) is not None + + @pytest.mark.asyncio + async def test_no_synth_stub_when_loop_has_clean_final(self, monkeypatch): + """If the loop already produced a clean final assistant message, bus + does NOT append a synthetic one.""" + + async def fake_loop(**kwargs): + im = list(kwargs.get("initial_messages") or []) + um = kwargs.get("user_message", "") + return AgentLoopResult( + messages=[*im, user_msg(um), assistant_msg("answer here")], + final_content="answer here", + stopped_by="terminal_tool", + metadata={}, + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", fake_loop, + ) + + bus, sid = await _make_bus_with_session() + await bus.submit_task_to_session(sid, "task 1") + await bus.wait_any_session("task-1", timeout=2) + + session = bus.get_session(sid) + # Boundary length is exactly 2 (user_msg + assistant_msg), no stub. + start, end = session.task_boundaries[0] + assert end - start == 1 + assert text_of(session.messages[end].get("content")) == "answer here" + + @pytest.mark.asyncio + async def test_last_report_falls_back_to_previous_when_empty(self, monkeypatch): + """When raw_result.final_content is empty and force_finalizer + doesn't fix it, last_report keeps its previous value (so + cross-agent never goes blank mid-session).""" + + responses = iter([ + AgentLoopResult( + messages=[user_msg("t1"), assistant_msg("A1")], + final_content="A1", + metadata={}, + ), + AgentLoopResult( + messages=[ + user_msg("t2"), + assistant_msg("", tool_calls=[_tool_call("c1", "x")]), + ], + final_content="", + stopped_by="max_turns", + metadata={}, + ), + ]) + + async def fake_loop(**kwargs): + base = next(responses) + im = list(kwargs.get("initial_messages") or []) + return AgentLoopResult( + messages=im + base.messages, + final_content=base.final_content, + stopped_by=base.stopped_by, + metadata=base.metadata, + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", fake_loop, + ) + + bus, sid = await _make_bus_with_session(trimmer=NullTrimmer()) + await bus.submit_task_to_session(sid, "t1") + await bus.wait_any_session("task-1", timeout=2) + assert bus.get_session(sid).last_report == "A1" + + await bus.submit_task_to_session(sid, "t2") + await bus.wait_any_session("task-1", timeout=2) + # Empty new final → keep "A1". + assert bus.get_session(sid).last_report == "A1" + + @pytest.mark.asyncio + async def test_last_report_falls_back_when_final_is_whitespace_only( + self, monkeypatch, + ): + """Whitespace-only final_content must NOT clobber last_report.""" + + responses = iter([ + AgentLoopResult( + messages=[user_msg("t1"), assistant_msg("A1")], + final_content="A1", + metadata={}, + ), + AgentLoopResult( + messages=[ + user_msg("t2"), + assistant_msg(" \n "), + ], + final_content=" \n ", # whitespace only + stopped_by="terminal_tool", + metadata={}, + ), + ]) + + async def fake_loop(**kwargs): + base = next(responses) + im = list(kwargs.get("initial_messages") or []) + return AgentLoopResult( + messages=im + base.messages, + final_content=base.final_content, + stopped_by=base.stopped_by, + metadata=base.metadata, + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", fake_loop, + ) + + bus, sid = await _make_bus_with_session(trimmer=NullTrimmer()) + await bus.submit_task_to_session(sid, "t1") + await bus.wait_any_session("task-1", timeout=2) + assert bus.get_session(sid).last_report == "A1" + + await bus.submit_task_to_session(sid, "t2") + await bus.wait_any_session("task-1", timeout=2) + # Whitespace-only new final → keep "A1", don't overwrite. + assert bus.get_session(sid).last_report == "A1" + + +# ── Aborted boundary invariant ──────────────────────────────────────────── + + +class TestAbortedBoundaryInvariant: + @pytest.mark.asyncio + async def test_cancelled_task_still_has_clean_final(self, monkeypatch): + """When the loop coroutine is cancelled, the bus closes the + boundary via close_session_boundary_aborted — that helper must + also append an abort stub so trimmer doesn't lose the task.""" + + import asyncio + + gate = asyncio.Event() + cancel_token = asyncio.Event() + + async def gated_loop(**kwargs): + cancel_token.set() + await gate.wait() # block forever; we'll cancel + return AgentLoopResult( + messages=[], final_content="", metadata={}, + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", gated_loop, + ) + + bus, sid = await _make_bus_with_session() + await bus.submit_task_to_session(sid, "task 1") + await cancel_token.wait() + + # Cancel the in-flight task via the bus's tracked task object. + session = bus.get_session(sid) + job = bus._jobs[session.current_job_id].task + job.cancel() + # Drain the cancellation + with contextlib.suppress(asyncio.CancelledError, Exception): + await asyncio.wait_for(job, timeout=2) + + # Boundary must be closed with at least one clean assistant message. + start, end = session.task_boundaries[0] + assert end is not None + assert find_final_assistant(session.messages, start + 1, end) is not None + # Stub mentions abort. + assert any( + is_assistant_msg(m) + and "aborted" in text_of(m.get("content")).lower() + for m in session.messages[start + 1: end + 1] + ) + + # Cleanup so pytest doesn't warn. + gate.set() + + +# ── observers_builder receives task_index ───────────────────────────────── + + +class TestObserversBuilderTaskIndex: + @pytest.mark.asyncio + async def test_builder_called_with_incrementing_task_index(self, monkeypatch): + async def fake_loop(**kwargs): + im = list(kwargs.get("initial_messages") or []) + um = kwargs.get("user_message", "") + return AgentLoopResult( + messages=[*im, user_msg(um), assistant_msg("ok")], + final_content="ok", + metadata={}, + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", fake_loop, + ) + + seen: list[int] = [] + + def observers_builder(_job_id, _item, task_index): + seen.append(task_index) + return [] + + spec = SubAgentRuntimeSpec(observers_builder=observers_builder) + bus, sid = await _make_bus_with_session(runtime_spec=spec) + + for prompt in ("t1", "t2", "t3"): + await bus.submit_task_to_session(sid, prompt) + await bus.wait_any_session("task-1", timeout=2) + + assert seen == [1, 2, 3] + + +# ── SSE payload extension ───────────────────────────────────────────────── + + +class TestSessionTaskSubmittedPayload: + @pytest.mark.asyncio + async def test_payload_has_is_reuse_and_role_id(self, monkeypatch): + async def fake_loop(**kwargs): + im = list(kwargs.get("initial_messages") or []) + um = kwargs.get("user_message", "") + return AgentLoopResult( + messages=[*im, user_msg(um), assistant_msg("ok")], + final_content="ok", + metadata={}, + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", fake_loop, + ) + + event_store = _FakeEventStore() + from agent_core.protocols import EventSink + + registry.register(EventSink, event_store) + + bus, sid = await _make_bus_with_session() + + await bus.submit_task_to_session(sid, "t1") + await bus.wait_any_session("task-1", timeout=2) + await bus.submit_task_to_session(sid, "t2") + await bus.wait_any_session("task-1", timeout=2) + + submitted = [ + call.kwargs["payload"] + for call in event_store.append.await_args_list + if call.kwargs.get("payload", {}).get("trace_type") + == "session_task_submitted" + ] + assert len(submitted) == 2 + assert submitted[0]["is_reuse"] is False + assert submitted[0]["role_id"] == "researcher" + assert submitted[0]["task_count"] == 1 + assert submitted[1]["is_reuse"] is True + assert submitted[1]["role_id"] == "researcher" + assert submitted[1]["task_count"] == 2 diff --git a/tests/test_agent_bus_session_queue.py b/tests/test_agent_bus_session_queue.py new file mode 100644 index 0000000..c41e1c7 --- /dev/null +++ b/tests/test_agent_bus_session_queue.py @@ -0,0 +1,414 @@ +"""Session-task queue: behaviour around assign-while-busy. + +Sessions execute serially (the message list / boundary list cannot be +mutated by two concurrent runs), but a second submit must NOT raise — +it queues FIFO and runs as soon as the predecessor finalises. Without +this, BrowseComp main agents see "already has a running task" errors +and fall back to spawning near-duplicate sessions, killing the reuse +benefit. Trial trace evidence: 38/112 trials hit the old error and ~50% +of those rejections led to a fresh-session spawn instead of reuse. +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import MagicMock + +import pytest + +from agent_core.components.agent_bus import AgentBus +from agent_core.components.agent_bus.spawn_guard import SpawnGuard +from agent_core.loop_types import AgentLoopResult +from agent_core.messages import assistant_msg, user_msg +from agent_core.models.task_budget import TaskBudget +from agent_core.runtime.registries import services as registry + + +@pytest.fixture(autouse=True) +def _clear_registry(): + registry.clear() + yield + registry.clear() + + +async def _bus_session() -> tuple[AgentBus, str]: + bus = AgentBus() + sid = await bus.create_session( + task_id="t-q", + name="lit", + role_id="researcher", + system_prompt="You are a researcher.", + llm_override=MagicMock(), + tools_override=[], + ) + return bus, sid + + +def _gated_loop_factory(gate: asyncio.Event): + async def loop(**kwargs): + await gate.wait() + im = list(kwargs.get("initial_messages") or []) + um = kwargs.get("user_message", "") + return AgentLoopResult( + messages=[*im, user_msg(um), assistant_msg(f"r:{um}")], + final_content=f"r:{um}", + metadata={}, + ) + return loop + + +# ────────────────────────────────────────────────────────────────────────── +# FIFO + multi-queue +# ────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_three_queued_drain_in_fifo_order(monkeypatch): + gate = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + _gated_loop_factory(gate), + ) + bus, sid = await _bus_session() + + j1 = await bus.submit_task_to_session(sid, "T1") + j2 = await bus.submit_task_to_session(sid, "T2") + j3 = await bus.submit_task_to_session(sid, "T3") + + session = bus.get_session(sid) + assert session.current_job_id == j1 + assert [p.job_id for p in session.pending_tasks] == [j2, j3] + assert bus.get_job_status(j2) == "queued" + assert bus.get_job_status(j3) == "queued" + + gate.set() + results = [] + for _ in range(3): + r = await bus.wait_any_session("t-q", timeout=2) + assert r is not None + results.append(r[1].final_content) + + # FIFO order preserved. + assert results == ["r:T1", "r:T2", "r:T3"] + assert not session.pending_tasks + assert session.current_job_id is None + + +# ────────────────────────────────────────────────────────────────────────── +# Drain still happens after current task FAILS +# ────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_queue_drains_when_current_fails(monkeypatch): + """A failed run still triggers queue drain — otherwise a single + flaky task would strand all queued follow-ups.""" + fail_gate = asyncio.Event() + success_gate = asyncio.Event() + seen: list[str] = [] + + async def loop(**kwargs): + um = kwargs.get("user_message", "") + seen.append(um) + if um == "T1": + await fail_gate.wait() + raise RuntimeError("inner loop boom") + await success_gate.wait() + im = list(kwargs.get("initial_messages") or []) + return AgentLoopResult( + messages=[*im, user_msg(um), assistant_msg(f"r:{um}")], + final_content=f"r:{um}", + metadata={}, + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", loop, + ) + bus, sid = await _bus_session() + + j1 = await bus.submit_task_to_session(sid, "T1") + j2 = await bus.submit_task_to_session(sid, "T2") + + fail_gate.set() + success_gate.set() + + # Drain via wait_any_session twice. + r1 = await bus.wait_any_session("t-q", timeout=2) + r2 = await bus.wait_any_session("t-q", timeout=2) + by_job = {r[1].job_id: r[1] for r in (r1, r2)} + + assert by_job[j1].success is False + assert "inner loop boom" in (by_job[j1].error or "") + assert by_job[j2].success is True + assert by_job[j2].final_content == "r:T2" + assert seen == ["T1", "T2"] + + +# ────────────────────────────────────────────────────────────────────────── +# Aborting a queued (not-yet-running) job +# ────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_abort_queued_job_purges_from_queue(monkeypatch): + gate = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + _gated_loop_factory(gate), + ) + bus, sid = await _bus_session() + + await bus.submit_task_to_session(sid, "T1") + j2 = await bus.submit_task_to_session(sid, "T2") + j3 = await bus.submit_task_to_session(sid, "T3") + + final = await bus.abort(j2) + assert final == "aborted" + session = bus.get_session(sid) + assert [p.job_id for p in session.pending_tasks] == [j3] + assert bus.get_job_status(j2) == "aborted" + # Aborted entry should produce a synthetic "aborted" SubAgentResult. + aborted_entry = bus._jobs[j2] + assert aborted_entry.result is not None + assert aborted_entry.result.success is False + assert aborted_entry.result.error == "aborted" + + gate.set() + # Only j1 + j3 should surface; j2 must not. + r1 = await bus.wait_any_session("t-q", timeout=2) + r2 = await bus.wait_any_session("t-q", timeout=2) + surfaced = {r[1].final_content for r in (r1, r2)} + assert surfaced == {"r:T1", "r:T3"} + # Third call returns None (everything drained). + assert await bus.wait_any_session("t-q", timeout=0.1) is None + + +@pytest.mark.asyncio +async def test_cancel_for_parent_clears_running_and_queued(monkeypatch): + gate = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + _gated_loop_factory(gate), + ) + bus, sid = await _bus_session() + + j1 = await bus.submit_task_to_session(sid, "T1") + j2 = await bus.submit_task_to_session(sid, "T2") + j3 = await bus.submit_task_to_session(sid, "T3") + + cancelled = await bus.cancel_for_parent("t-q") + assert cancelled == 3 + assert bus.get_job_status(j1) == "aborted" + assert bus.get_job_status(j2) == "aborted" + assert bus.get_job_status(j3) == "aborted" + assert not bus.get_session(sid).pending_tasks + # No surviving asyncio tasks expected — gate-set should be a no-op. + gate.set() + await asyncio.sleep(0) + + +# ────────────────────────────────────────────────────────────────────────── +# Status reporting + spawn guard +# ────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_collect_treats_queued_session_job_as_pending(monkeypatch): + """Queued session tasks have no asyncio.Task yet, so collect() must + report them as pending rather than passing None into asyncio.wait.""" + gate = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + _gated_loop_factory(gate), + ) + bus, sid = await _bus_session() + + await bus.submit_task_to_session(sid, "T1") + j2 = await bus.submit_task_to_session(sid, "T2") + + result = await bus.collect([j2], timeout=0.01) + + assert result.pending == [j2] + assert result.completed == [] + assert result.failed == [] + + gate.set() + for _ in range(2): + await bus.wait_any_session("t-q", timeout=2) + + +@pytest.mark.asyncio +async def test_session_state_running_when_only_queued(monkeypatch): + """If a task is queued (no current_job_id yet), the swarm is not + 'all_collected' — collect_reports must keep waiting.""" + from agent_core.components.agent_bus.fan_in import session_state + + gate = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + _gated_loop_factory(gate), + ) + bus, sid = await _bus_session() + await bus.submit_task_to_session(sid, "T1") + await bus.submit_task_to_session(sid, "T2") + + sessions = bus.list_sessions_for_task("t-q") + # T1 running + T2 queued → state should be "running". + assert session_state(sessions) == "running" + + gate.set() + # Drain so pytest doesn't warn about pending tasks. + for _ in range(2): + await bus.wait_any_session("t-q", timeout=2) + + +@pytest.mark.asyncio +async def test_queued_jobs_release_guard_reservation_on_abort(monkeypatch): + """Queued tasks reserve budget at submit time. Aborting must release + it — otherwise repeated submit/abort cycles would silently exhaust + the task budget.""" + gate = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + _gated_loop_factory(gate), + ) + + bus = AgentBus() + budget = TaskBudget(max_tokens=10_000, max_depth=3, max_parallel=4) + guard = SpawnGuard(budget=budget, sub_agent_timeout_s=5) + bus.set_spawn_guard(guard) + + sid = await bus.create_session( + task_id="t-budget", + name="lit", + role_id="researcher", + system_prompt="You are a researcher.", + llm_override=MagicMock(), + tools_override=[], + ) + + j1 = await bus.submit_task_to_session(sid, "T1", estimated_tokens=3_000) + j2 = await bus.submit_task_to_session(sid, "T2", estimated_tokens=3_000) + reserved_before = guard.tokens_reserved + assert reserved_before >= 6_000 + + await bus.abort(j2) + # Aborting the queued j2 must drop its reservation. + assert guard.tokens_reserved < reserved_before + + gate.set() + await bus.wait_any_session("t-budget", timeout=2) + _ = j1 # silence unused + + +@pytest.mark.asyncio +async def test_cleanup_task_aborts_queued_jobs_and_releases_guard(monkeypatch): + """cleanup_task drops sessions at task end. Queued jobs must be + explicitly aborted first; otherwise their SpawnGuard reservations leak + after the session object is deleted.""" + gate = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + _gated_loop_factory(gate), + ) + + bus = AgentBus() + guard = SpawnGuard( + budget=TaskBudget(max_tokens=10_000, max_depth=3, max_parallel=2), + sub_agent_timeout_s=5, + ) + bus.set_spawn_guard(guard) + + sid = await bus.create_session( + task_id="t-clean", + name="lit", + role_id="researcher", + system_prompt="You are a researcher.", + llm_override=MagicMock(), + tools_override=[], + ) + + j1 = await bus.submit_task_to_session(sid, "T1", estimated_tokens=3_000) + j2 = await bus.submit_task_to_session(sid, "T2", estimated_tokens=3_000) + assert bus.get_job_status(j2) == "queued" + assert guard.tokens_reserved >= 6_000 + + cleaned = await bus.cleanup_task("t-clean", cancel_timeout_s=0.1) + + assert cleaned == 1 + assert bus.get_session(sid) is None + assert bus.get_job_status(j1) == "aborted" + assert bus.get_job_status(j2) == "aborted" + assert guard.tokens_reserved == 0 + assert guard.active_count == 0 + gate.set() + + +# ────────────────────────────────────────────────────────────────────────── +# codex #6: spawn_context plumbing +# ────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_submit_task_threads_spawn_context_into_pending(monkeypatch): + """``submit_task_to_session(spawn_context=...)`` must land the dict + on ``PendingSessionTask.spawn_context`` so ``_dispatch_session_task`` + can stamp it onto the sub trace observer at dispatch time. This is + the Log Schema v1 §2.4 wiring for SFT data attribution. + """ + gate = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + _gated_loop_factory(gate), + ) + bus, sid = await _bus_session() + + spawn_ctx = { + "parent_run_id": "heavy_run_1", + "parent_agent_id": "swarm_main", + "delegation_prompt": "find NVIDIA Q3 datacenter revenue", + "allowed_tools": ["web_search"], + "depth": 1, + "budget": {"max_turns": 8}, + } + + # First task runs immediately (free path); second queues. The + # queued task is the easy assertion target since its ``pending`` + # field is still readable on the session. + await bus.submit_task_to_session(sid, "T1") + await bus.submit_task_to_session(sid, "T2", spawn_context=spawn_ctx) + + session = bus.get_session(sid) + # The queued task carries the spawn_context verbatim. + assert len(session.pending_tasks) == 1 + queued = session.pending_tasks[0] + assert queued.spawn_context == spawn_ctx + # Defensive copy: caller mutation must not leak into the bus. + spawn_ctx["depth"] = 999 + assert queued.spawn_context["depth"] == 1 + + gate.set() + # Drain so the test doesn't leave dangling tasks. + for _ in range(2): + await bus.wait_any_session("t-q", timeout=2) + + +@pytest.mark.asyncio +async def test_submit_task_without_spawn_context_stays_none(monkeypatch): + """No spawn_context arg → ``pending.spawn_context`` stays ``None`` + so callers that don't wire the field keep the legacy behaviour.""" + gate = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + _gated_loop_factory(gate), + ) + bus, sid = await _bus_session() + + await bus.submit_task_to_session(sid, "T1") + await bus.submit_task_to_session(sid, "T2") + + session = bus.get_session(sid) + assert session.pending_tasks[0].spawn_context is None + gate.set() + for _ in range(2): + await bus.wait_any_session("t-q", timeout=2) diff --git a/tests/test_agent_bus_waiting.py b/tests/test_agent_bus_waiting.py new file mode 100644 index 0000000..8517e20 --- /dev/null +++ b/tests/test_agent_bus_waiting.py @@ -0,0 +1,265 @@ +from __future__ import annotations + +import asyncio +import contextlib +from types import SimpleNamespace + +import pytest + +from agent_core.components.agent_bus import ( + AgentBus, + JobEntry, + SubAgentSession, + SubTask, +) +from agent_core.loop_types import AgentLoopResult +from agent_core.messages import assistant_msg, system_msg, user_msg +from agent_core.runtime.loop.message_trimmer import NullTrimmer + + +def _session(task_id: str = "root") -> SubAgentSession: + return SubAgentSession( + session_id=f"{task_id}::worker", + task_id=task_id, + name="worker", + role_id="researcher", + system_prompt="test", + tools=[], + llm=None, + trimmer=NullTrimmer(), + ) +async def test_session_spawn_registers_job_before_eager_task_starts( + monkeypatch, +) -> None: + """The TUI's eager task factory must not outrun AgentBus bookkeeping.""" + async def fake_run_agent_loop(**kwargs): + ctx = SimpleNamespace( + turn=1, + thinking="I should inspect the ClinVar source before generating data.", + ai_text="I will query the source and validate the response.", + ) + for observer in kwargs["observers"]: + await observer.on_llm_delta(SimpleNamespace( + turn=1, thinking_delta="I should inspect ", delta="", + )) + await observer.on_llm_delta(SimpleNamespace( + turn=1, thinking_delta="the ClinVar source first.", delta="", + )) + await observer.on_llm_response(ctx) + await observer.on_tool_call(ctx, { + "name": "web_search", "args": {"query": "ClinVar API"}, + }) + await observer.on_tool_result(ctx, SimpleNamespace( + name="web_search", result="found source", is_error=False, + )) + return AgentLoopResult( + messages=[ + system_msg(kwargs["system_prompt"]), + user_msg(kwargs["user_message"]), + assistant_msg("finished"), + ], + final_content="finished", + stopped_by="final_answer", + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + fake_run_agent_loop, + ) + loop = asyncio.get_running_loop() + previous_factory = loop.get_task_factory() + loop.set_task_factory(asyncio.eager_task_factory) + try: + bus = AgentBus() + session_id = await bus.create_session( + task_id="root", + name="worker", + role_id="researcher", + system_prompt="test", + tools_override=[], + llm_override=object(), + ) + job_id = await bus.submit_task_to_session(session_id, "do work") + outcome = await bus.wait_any_session_detailed("root", timeout=0) + finally: + loop.set_task_factory(previous_factory) + + assert outcome.reason == "ready" + assert outcome.result is not None + _, result = outcome.result + assert result.success + assert result.final_content == "finished" + assert bus._jobs[job_id].status == "completed" + assert bus._sessions[session_id].current_job_id is None + snapshot = bus.describe_sessions_for_task("root")[0] + assert [event["kind"] for event in snapshot["events"]] == [ + "thinking", "message", "tool_call", "tool_result", + ] + assert snapshot["events"][0]["detail"].startswith("I should inspect") + assert len([e for e in snapshot["events"] if e["kind"] == "thinking"]) == 1 + + +@pytest.mark.asyncio +async def test_wait_reconciles_cancelled_task_instead_of_fake_timeout() -> None: + bus = AgentBus() + session = _session() + job_id = "root::worker::1" + task = asyncio.create_task(asyncio.sleep(60)) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + session.current_job_id = job_id + session.total_task_count = 1 + bus._sessions[session.session_id] = session + bus._jobs[job_id] = JobEntry( + job_id=job_id, + parent_task_id="root", + item=SubTask(question="work", role_id="researcher"), + task=task, + status="running", + ) + + outcome = await bus.wait_any_session_detailed("root", timeout=1800) + + assert outcome.reason == "ready" + assert outcome.elapsed_s < 1 + assert outcome.result is not None + _, result = outcome.result + assert not result.success + assert result.error_class == "CancelledError" + assert session.current_job_id is None + assert bus._jobs[job_id].status == "aborted" + + +@pytest.mark.asyncio +async def test_wait_reports_actual_timeout_only_for_live_task() -> None: + bus = AgentBus() + session = _session() + job_id = "root::worker::1" + task = asyncio.create_task(asyncio.sleep(60)) + session.current_job_id = job_id + session.total_task_count = 1 + bus._sessions[session.session_id] = session + bus._jobs[job_id] = JobEntry( + job_id=job_id, + parent_task_id="root", + item=SubTask(question="work", role_id="researcher"), + task=task, + status="running", + ) + + try: + outcome = await bus.wait_any_session_detailed("root", timeout=0.02) + finally: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + assert outcome.reason == "timeout" + assert outcome.result is None + assert 0.01 <= outcome.elapsed_s < 1 + assert session.current_job_id == job_id + + +@pytest.mark.asyncio +async def test_wait_reports_unpublished_when_a_task_ends_empty_handed() -> None: + """A task that finished during the wait is not "nothing to wait for". + + Real time elapsed, so telling the coordinator to discount the wait (what + ``no_pending`` means) would be the same wrong-elapsed-time steering in the + opposite direction. + """ + bus = AgentBus() + session = _session() + job_id = "root::worker::1" + + async def _finish_without_publishing() -> None: + await asyncio.sleep(0.02) + # Mimic a wrapper that completes but leaves the session's bookkeeping + # untouched, so reconciliation has nothing to hand back either. + bus._jobs[job_id].status = "completed" + session.current_job_id = None + + task = asyncio.create_task(_finish_without_publishing()) + session.current_job_id = job_id + session.total_task_count = 1 + bus._sessions[session.session_id] = session + bus._jobs[job_id] = JobEntry( + job_id=job_id, + parent_task_id="root", + item=SubTask(question="work", role_id="researcher"), + task=task, # type: ignore[arg-type] + status="running", + ) + + outcome = await bus.wait_any_session_detailed("root", timeout=5) + + assert outcome.reason == "unpublished" + assert outcome.result is None + assert outcome.elapsed_s >= 0.02 + + +@pytest.mark.asyncio +async def test_session_snapshot_freezes_a_finished_worker_duration() -> None: + """The UI reads ``elapsed_s`` verbatim, so it must stop growing. + + ``active`` is what tells a renderer whether the number is still counting; + without it a finished worker's row ticks up for as long as the fan-in + blocks. + """ + async def fake_run_agent_loop(**kwargs): + return AgentLoopResult( + messages=[ + system_msg(kwargs["system_prompt"]), + user_msg(kwargs["user_message"]), + assistant_msg("done"), + ], + final_content="done", + stopped_by="final_answer", + ) + + bus = AgentBus() + session_id = await bus.create_session( + task_id="root", + name="market_research", + role_id="agent_team_sub", + system_prompt="test", + tools_override=[], + llm_override=object(), + ) + + unassigned = bus.describe_sessions_for_task("root")[0] + assert unassigned["status"] == "unassigned" + assert unassigned["active"] is False + assert unassigned["elapsed_s"] == 0.0 + assert unassigned["role_id"] == "agent_team_sub" + + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + fake_run_agent_loop, + ) + await bus.submit_task_to_session(session_id, "do work") + assert await bus.wait_any_session("root", timeout=5) is not None + + finished = bus.describe_sessions_for_task("root")[0] + assert finished["status"] == "idle" + assert finished["active"] is False + frozen = finished["elapsed_s"] + + await asyncio.sleep(0.05) + assert bus.describe_sessions_for_task("root")[0]["elapsed_s"] == frozen + + +@pytest.mark.asyncio +async def test_wait_distinguishes_no_pending_work_from_timeout() -> None: + bus = AgentBus() + session = _session() + bus._sessions[session.session_id] = session + + outcome = await bus.wait_any_session_detailed("root", timeout=1800) + + assert outcome.reason == "no_pending" + assert outcome.result is None + assert outcome.elapsed_s < 1 diff --git a/tests/test_finalization_budget.py b/tests/test_finalization_budget.py new file mode 100644 index 0000000..9ffc036 --- /dev/null +++ b/tests/test_finalization_budget.py @@ -0,0 +1,246 @@ +"""Wall-clock split between research and the finalization phase. + +The invariant this file protects: research stop + a full late-tool overrun + the +finalization phase must fit inside the externally enforced ceiling. It does NOT +hold on the research deadline alone, because ``soft_wall_deadline_s`` floors +research at half the wall — on a short wall the planned reserve is silently +larger than what survives. ``remaining_phase_budget_s`` is what closes the gap, +so these tests pin both halves together. +""" + +from __future__ import annotations + +import time + +import pytest + +from agent_core.components.finalization import ( + check_wall_feasibility, + remaining_phase_budget_s, + resolve_report_wall_time, + resolve_research_wall, + soft_wall_deadline_s, +) + +# ── hard_total_s: which config keys are real external ceilings ──────────── + +def test_env_wall_is_a_hard_ceiling(monkeypatch): + monkeypatch.setenv("MIROHARNESS_TASK_WALL_TIME_S", "7200") + wall = resolve_research_wall({}, reserve_s=1800, label_prefix="t") + assert wall.hard_total_s == 7200 + assert wall.research_deadline_s == 5400 + + +def test_research_only_budget_is_not_a_hard_ceiling(monkeypatch): + """The reporter deliberately runs *outside* ``research_wall_time_s``.""" + monkeypatch.delenv("MIROHARNESS_TASK_WALL_TIME_S", raising=False) + wall = resolve_research_wall( + {"research_wall_time_s": 9000}, reserve_s=1800, label_prefix="t", + ) + assert wall.hard_total_s == 0 + assert wall.research_deadline_s == 9000 + + +def test_legacy_wall_deadline_s_is_a_hard_ceiling(monkeypatch): + monkeypatch.delenv("MIROHARNESS_TASK_WALL_TIME_S", raising=False) + wall = resolve_research_wall( + {"wall_deadline_s": 3600}, reserve_s=600, label_prefix="t", + ) + assert wall.hard_total_s == 3600 + assert wall.research_deadline_s == 3000 + + +def test_shortest_positive_ceiling_wins(monkeypatch): + monkeypatch.setenv("MIROHARNESS_TASK_WALL_TIME_S", "3600") + wall = resolve_research_wall( + {"wall_deadline_s": 9000}, reserve_s=600, label_prefix="t", + ) + assert wall.hard_total_s == 3600 + + +def test_no_ceiling_anywhere(monkeypatch): + monkeypatch.delenv("MIROHARNESS_TASK_WALL_TIME_S", raising=False) + wall = resolve_research_wall({}, reserve_s=600, label_prefix="t") + assert wall.hard_total_s == 0 + assert wall.research_deadline_s == 0 + + +# ── operator-friendly phase overrides ──────────────────────────────────── + +def test_research_wall_env_replaces_profile_budget(monkeypatch): + monkeypatch.delenv("MIROHARNESS_TASK_WALL_TIME_S", raising=False) + monkeypatch.setenv("RESEARCH_WALL_TIME", "300") + + wall = resolve_research_wall( + {"research_wall_time_s": 9000}, reserve_s=1800, label_prefix="t", + ) + + assert wall.research_deadline_s == 300 + assert wall.hard_total_s == 0 + + +def test_research_wall_env_can_intentionally_extend_profile(monkeypatch): + monkeypatch.delenv("MIROHARNESS_TASK_WALL_TIME_S", raising=False) + monkeypatch.setenv("RESEARCH_WALL_TIME", "12000") + + wall = resolve_research_wall( + {"research_wall_time_s": 9000}, reserve_s=1800, label_prefix="t", + ) + + assert wall.research_deadline_s == 12000 + + +@pytest.mark.parametrize("raw", ["", "0", "-1", "not-a-number"]) +def test_invalid_research_wall_env_preserves_profile(monkeypatch, raw): + monkeypatch.delenv("MIROHARNESS_TASK_WALL_TIME_S", raising=False) + monkeypatch.setenv("RESEARCH_WALL_TIME", raw) + + wall = resolve_research_wall( + {"research_wall_time_s": 9000}, reserve_s=1800, label_prefix="t", + ) + + assert wall.research_deadline_s == 9000 + + +def test_task_wall_still_caps_research_env(monkeypatch): + monkeypatch.setenv("RESEARCH_WALL_TIME", "9000") + monkeypatch.setenv("MIROHARNESS_TASK_WALL_TIME_S", "7200") + + wall = resolve_research_wall({}, reserve_s=1800, label_prefix="t") + + assert wall.research_deadline_s == 5400 + assert wall.hard_total_s == 7200 + + +def test_report_wall_env_replaces_profile_reserve(monkeypatch): + monkeypatch.setenv("REPORT_WALL_TIME", "600") + + reserve = resolve_report_wall_time( + {"wall_deadline_reserve_s": 1800}, default=180, label_prefix="t", + ) + + assert reserve == 600 + + +@pytest.mark.parametrize("raw", ["", "-1", "not-a-number"]) +def test_invalid_report_wall_env_preserves_profile(monkeypatch, raw): + monkeypatch.setenv("REPORT_WALL_TIME", raw) + + reserve = resolve_report_wall_time( + {"wall_deadline_reserve_s": 1800}, default=180, label_prefix="t", + ) + + assert reserve == 1800 + + +# ── remaining_phase_budget_s ────────────────────────────────────────────── + +def test_no_deadline_leaves_the_requested_budget_alone(): + assert remaining_phase_budget_s(1800, None) == 1800 + + +def test_budget_is_clamped_to_the_time_left(): + deadline = time.monotonic() + 100 + assert remaining_phase_budget_s(1800, deadline) == pytest.approx(100, abs=2) + + +def test_budget_is_not_extended_past_the_request(): + deadline = time.monotonic() + 9000 + assert remaining_phase_budget_s(1800, deadline) == 1800 + + +def test_expired_deadline_still_gets_one_short_attempt(): + """Zero would abort before the fail-open baseline is ever emitted.""" + assert remaining_phase_budget_s(1800, time.monotonic() - 500) == 1.0 + assert remaining_phase_budget_s(1800, time.monotonic() - 500, minimum_s=5) == 5 + + +# ── the invariant the 0.5 floor used to break ───────────────────────────── + +@pytest.mark.parametrize("hard_wall_s", [3600, 5400, 7200, 10800]) +def test_research_plus_tool_overrun_plus_finalization_fits_the_wall( + monkeypatch, hard_wall_s, +): + """Worst case must not exceed the ceiling. + + Reproduces the online shape: ``tool_timeout_s: 1800`` and a 1800s reporter + phase. Below 7200s the ``total * 0.5`` floor keeps research alive at the + reserve's expense, so the finalization phase — not research — is what has + to give. Holds while ``tool_timeout_s <= hard_wall / 2``; the case below + that bound is covered by + :func:`test_tool_timeout_over_half_the_wall_is_reported_infeasible`. + """ + tool_timeout_s = 1800.0 + requested_phase_s = 1800.0 + reserve_s = max(3600.0, tool_timeout_s + requested_phase_s) + + monkeypatch.setenv("MIROHARNESS_TASK_WALL_TIME_S", str(hard_wall_s)) + wall = resolve_research_wall( + {}, reserve_s=reserve_s, label_prefix="t", + ) + + # Research always keeps at least half the wall — a generous reserve must + # never starve it. + assert wall.research_deadline_s >= hard_wall_s * 0.5 + + # Worst case: the deadline observer only fires at turn end, so a tool + # started just before it can burn its whole timeout. + started = time.monotonic() + phase_start = wall.research_deadline_s + tool_timeout_s + effective_phase_s = remaining_phase_budget_s( + requested_phase_s, + started + wall.hard_total_s - phase_start, + ) + worst_case_finish = phase_start + effective_phase_s + assert worst_case_finish <= hard_wall_s + 1.0 + + +def test_soft_deadline_floor_is_why_the_clamp_is_required(): + """Pin the floor's behaviour so the pairing above cannot be dropped.""" + # Reserve fits: research gets the full remainder. + assert soft_wall_deadline_s(10800, 3600) == 7200 + # Reserve does not fit: research is floored and the reserve is truncated + # from 3600s to 1800s without any error. + assert soft_wall_deadline_s(3600, 3600) == 1800 + + +# ── the residual case no arithmetic can fix ─────────────────────────────── + +def test_tool_timeout_over_half_the_wall_is_reported_infeasible(caplog): + """A tool already running cannot be shrunk — the config must change.""" + with caplog.at_level("WARNING"): + feasible = check_wall_feasibility( + hard_total_s=1800, + research_deadline_s=900, + tool_timeout_s=1800, + landing_budget_s=1800, + label_prefix="stateful", + ) + assert feasible is False + assert "cannot guarantee a final answer" in caplog.text + assert "tool_timeout_s" in caplog.text + + +def test_feasible_config_is_silent(caplog): + with caplog.at_level("WARNING"): + feasible = check_wall_feasibility( + hard_total_s=10800, + research_deadline_s=7200, + tool_timeout_s=1800, + landing_budget_s=1800, + label_prefix="stateful", + ) + assert feasible is True + assert caplog.text == "" + + +def test_no_external_ceiling_is_always_feasible(caplog): + with caplog.at_level("WARNING"): + assert check_wall_feasibility( + hard_total_s=0, + research_deadline_s=9000, + tool_timeout_s=1800, + landing_budget_s=1800, + label_prefix="stateful", + ) is True + assert caplog.text == "" diff --git a/tests/test_resource_manager.py b/tests/test_resource_manager.py new file mode 100644 index 0000000..a527940 --- /dev/null +++ b/tests/test_resource_manager.py @@ -0,0 +1,240 @@ +"""Tests for the ResourceManager kernel module.""" +from unittest.mock import MagicMock + +import pytest + +from agent_core.errors import PermissionDenied +from agent_core.runtime.registries import services as service_registry +from agent_core.runtime.registries.agents import AgentDefinition, AgentRegistry +from agent_core.runtime.resources.manager import ResourceManager +from agent_core.runtime.tool_permission import ToolPermissionContext + + +@pytest.fixture +def resource_mgr_with_registry(): + """ResourceManager backed by a real AgentRegistry with known roles.""" + service_registry.clear() + + # Build a minimal AgentRegistry with the 3 canonical roles + agent_reg = AgentRegistry() + agent_reg.register(AgentDefinition( + role_id="researcher", display_name="Researcher", + system_prompt="test", allowed_tools=["web_search", "web_fetch", "bash"], + color="#3b82f6", icon="search", + )) + agent_reg.register(AgentDefinition( + role_id="summarizer", display_name="Summarizer", + system_prompt="test", allowed_tools=["read_file"], + color="#a855f7", icon="file", + )) + + # Register in service_registry so ResourceManager._get_allowed_tools can find it + service_registry.register(AgentRegistry, agent_reg) + + llm = MagicMock() + tools = { + "web_search": MagicMock(name="web_search"), + "web_fetch": MagicMock(name="web_fetch"), + "bash": MagicMock(name="bash"), + "read_file": MagicMock(name="read_file"), + } + for name, tool in tools.items(): + tool.name = name + mgr = ResourceManager(llm=llm, tools=tools) + yield mgr + service_registry.clear() + + +def test_researcher_has_all_tools(resource_mgr_with_registry): + tools = resource_mgr_with_registry.get_tools_for_role("researcher") + assert len(tools) == 3 + + +def test_summarizer_only_has_read_file(resource_mgr_with_registry): + tools = resource_mgr_with_registry.get_tools_for_role("summarizer") + assert len(tools) == 1 + + +def test_permission_check(resource_mgr_with_registry): + assert resource_mgr_with_registry.check_permission("researcher", "web_search") + assert not resource_mgr_with_registry.check_permission("summarizer", "web_search") + assert not resource_mgr_with_registry.check_permission("unknown-role", "web_search") + + +def test_require_permission_raises(resource_mgr_with_registry): + with pytest.raises(PermissionDenied): + resource_mgr_with_registry.require_permission("summarizer", "web_search") + + +def test_unknown_role_has_no_tools(resource_mgr_with_registry): + assert resource_mgr_with_registry.get_tools_for_role("unknown-role") == [] + assert resource_mgr_with_registry.get_tool_names_for_role("unknown-role") == [] + assert resource_mgr_with_registry.get_tool_for_role("unknown-role", "web_search") is None + + +def test_permission_context_filters_tools(resource_mgr_with_registry): + tools = resource_mgr_with_registry.get_tools_for_role( + "researcher", + permission_context=ToolPermissionContext.from_iterables( + allow_names={"bash"}, + ), + ) + assert [t.name for t in tools] == ["bash"] + + +def test_get_tool_for_role_honors_permission_context(resource_mgr_with_registry): + tool = resource_mgr_with_registry.get_tool_for_role( + "researcher", + "bash", + permission_context=ToolPermissionContext.from_iterables( + allow_names={"bash"}, + ), + ) + assert tool is not None + denied = resource_mgr_with_registry.get_tool_for_role( + "researcher", + "web_search", + permission_context=ToolPermissionContext.from_iterables( + allow_names={"bash"}, + ), + ) + assert denied is None + + +def test_registry_unavailable_is_fail_closed(): + """Fail-closed: no AgentRegistry → no tools allowed for any role.""" + service_registry.clear() + llm = MagicMock() + tools = { + "web_search": MagicMock(name="web_search"), + "web_fetch": MagicMock(name="web_fetch"), + } + mgr = ResourceManager(llm=llm, tools=tools) + + assert len(mgr.get_tools_for_role("any-role")) == 0 + assert not mgr.check_permission("any-role", "web_search") + with pytest.raises(PermissionDenied): + mgr.require_permission("any-role", "web_search") + + +def test_no_permission_context_returns_all_role_tools(resource_mgr_with_registry): + tools = resource_mgr_with_registry.get_tools_for_role("researcher") + assert {t.name for t in tools} == {"web_search", "web_fetch", "bash"} + + +def test_set_role_llm_overrides_default_for_subsequent_get(resource_mgr_with_registry): + """``set_role_llm`` lets a workflow node register a per-task LLM + (e.g. swarm main_agent loading qwen35_397b) so peer nodes pick it + up via ``get_llm(role_id)`` instead of falling back to the + ResourceManager's default LLM.""" + profile_llm = MagicMock(name="qwen-397b") + resource_mgr_with_registry.set_role_llm("swarm_main", profile_llm) + assert resource_mgr_with_registry.get_llm("swarm_main") is profile_llm + + +def test_set_role_llm_does_not_leak_into_other_roles(resource_mgr_with_registry): + """Registering one role's LLM doesn't change resolution for other + roles — they still hit the default LLM.""" + profile_llm = MagicMock(name="qwen-397b") + resource_mgr_with_registry.set_role_llm("swarm_main", profile_llm) + other = resource_mgr_with_registry.get_llm("researcher") + assert other is not profile_llm + assert other is resource_mgr_with_registry.llm + + +def test_set_role_llm_idempotent_last_writer_wins(resource_mgr_with_registry): + """Repeated registration (e.g. heavy_mode K parallel main_agents + all loading the same profile) is harmless: the latest registration + wins, no exceptions raised.""" + first = MagicMock(name="qwen-397b-v1") + second = MagicMock(name="qwen-397b-v2") + resource_mgr_with_registry.set_role_llm("swarm_main", first) + resource_mgr_with_registry.set_role_llm("swarm_main", second) + assert resource_mgr_with_registry.get_llm("swarm_main") is second + + +# ── Global tool policy (config injection) ───────────────────────────────── + +def test_global_deny_removes_tool_for_all_roles(resource_mgr_with_registry): + """``{web_search: false}`` config → web_search disappears from every + role that would otherwise have it.""" + from agent_core.runtime.tool_permission import ( + from_config_map, + ) + + resource_mgr_with_registry.set_global_tool_policy( + from_config_map({"web_search": False}), + ) + names = {t.name for t in resource_mgr_with_registry.get_tools_for_role("researcher")} + assert names == {"web_fetch", "bash"} + assert not resource_mgr_with_registry.check_permission("researcher", "web_search") + + +def test_global_allowlist_intersects_role_tools(resource_mgr_with_registry): + """``{bash: true}`` config → only bash survives for any role.""" + from agent_core.runtime.tool_permission import ( + from_config_map, + ) + + resource_mgr_with_registry.set_global_tool_policy( + from_config_map({"bash": True}), + ) + assert [t.name for t in resource_mgr_with_registry.get_tools_for_role("researcher")] == ["bash"] + # A role without bash ends up with nothing. + assert resource_mgr_with_registry.get_tools_for_role("summarizer") == [] + + +def test_global_policy_layers_under_per_call_context(resource_mgr_with_registry): + """A global deny and a per-call allowlist both apply — deny wins.""" + from agent_core.runtime.tool_permission import ( + from_config_map, + ) + + resource_mgr_with_registry.set_global_tool_policy( + from_config_map({"web_search": False}), + ) + # Per-call asks for web_search + bash; global deny strips web_search. + tools = resource_mgr_with_registry.get_tools_for_role( + "researcher", + permission_context=ToolPermissionContext.from_iterables( + allow_names={"web_search", "bash"}, + ), + ) + assert [t.name for t in tools] == ["bash"] + + +def test_clearing_global_policy_restores_tools(resource_mgr_with_registry): + """Setting an empty/None policy clears it so the next run is unaffected.""" + from agent_core.runtime.tool_permission import ( + from_config_map, + ) + + resource_mgr_with_registry.set_global_tool_policy( + from_config_map({"web_search": False}), + ) + # Empty config map → no-op context → policy cleared. + resource_mgr_with_registry.set_global_tool_policy(from_config_map({})) + assert resource_mgr_with_registry.global_tool_policy is None + names = {t.name for t in resource_mgr_with_registry.get_tools_for_role("researcher")} + assert names == {"web_search", "web_fetch", "bash"} + + +def test_set_global_policy_replaces_not_accumulates(resource_mgr_with_registry): + """``set_global_tool_policy`` is last-writer-wins, not additive — this is + the cross-run safety contract the SDK driver relies on (it sets the policy + unconditionally at the start of every run, so run N-1's deny never leaks + into run N).""" + from agent_core.runtime.tool_permission import ( + from_config_map, + ) + + # Run 1 denies web_search. + resource_mgr_with_registry.set_global_tool_policy( + from_config_map({"web_search": False}), + ) + # Run 2 denies web_fetch instead — must NOT also keep run 1's web_search deny. + resource_mgr_with_registry.set_global_tool_policy( + from_config_map({"web_fetch": False}), + ) + names = {t.name for t in resource_mgr_with_registry.get_tools_for_role("researcher")} + assert names == {"web_search", "bash"} # web_search back; only web_fetch denied diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py new file mode 100644 index 0000000..4adbf92 --- /dev/null +++ b/tests/test_scheduler.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any + +import pytest + +from agent_core.scheduling.scheduler import ( + Scheduler, + TaskWallTimeExceeded, + resolve_wall_time_s, +) +from agent_core.scheduling.workflow_defaults import ( + clear_workflow_defaults, + register_workflow_defaults, +) +from agent_core.types import TaskId, TaskStatus + + +@dataclass +class _Task: + status: TaskStatus = TaskStatus.CREATED + thread_id: str = "thread" + pipeline_id: str | None = None + context: dict[str, Any] = field(default_factory=dict) + + +class _ProcessManager: + def __init__(self) -> None: + self.task = _Task() + self.error: str | None = None + self.abort_reason: str | None = None + + async def get_task(self, task_id: TaskId) -> _Task: + del task_id + return self.task + + async def update_status(self, task_id: TaskId, status: TaskStatus) -> None: + del task_id + self.task.status = status + + async def abort_task(self, task_id: TaskId, reason: str) -> None: + del task_id + self.abort_reason = reason + self.task.status = TaskStatus.ABORTED + + async def set_error(self, task_id: TaskId, error: str) -> None: + del task_id + self.error = error + self.task.status = TaskStatus.FAILED + + +class _EventSink: + def __init__(self) -> None: + self.events: list[dict[str, Any]] = [] + + async def append(self, **event: Any) -> None: + self.events.append(event) + + +class _ReportGraph: + async def astream(self, *_args: Any, **_kwargs: Any): + yield {"report": {"report": "done"}} + + async def aget_state(self, _config: dict[str, Any]): + return SimpleNamespace(values={"report": "done"}) + + +class _EmptyGraph: + async def astream(self, *_args: Any, **_kwargs: Any): + yield {"research": {"current_phase": "research"}} + + async def aget_state(self, _config: dict[str, Any]): + return SimpleNamespace(values={"current_phase": "research"}) + + +class _HangingGraph: + async def astream(self, *_args: Any, **_kwargs: Any): + yield {"research": {"current_phase": "research"}} + await asyncio.sleep(60) + + +def test_wall_time_resolution_supports_portable_and_legacy_env(monkeypatch): + monkeypatch.setenv("MIROHARNESS_TASK_WALL_TIME_S", "120") + assert resolve_wall_time_s(None) == 120 + monkeypatch.setenv("AGENT_CORE_TASK_WALL_TIME_S", "30") + assert resolve_wall_time_s(None) == 30 + assert resolve_wall_time_s(10) == 10 + assert resolve_wall_time_s(0) is None + + +def test_workflow_defaults_are_product_injected(monkeypatch): + for name in ( + "AGENT_CORE_TASK_WALL_TIME_S", + "MIROHARNESS_TASK_WALL_TIME_S", + "FRONTIER_AGENT_TASK_WALL_TIME_S", + ): + monkeypatch.delenv(name, raising=False) + clear_workflow_defaults() + register_workflow_defaults( + ("research", "research-report"), + {"TASK_WALL_TIME_S": 900, "TASK_WALL_TIME_MODE": "soft_research"}, + ) + assert resolve_wall_time_s(None, "research") == 900 + clear_workflow_defaults() + + +@pytest.mark.asyncio +async def test_scheduler_completes_only_with_terminal_output(): + pm = _ProcessManager() + sink = _EventSink() + scheduler = Scheduler(_ReportGraph(), pm, sink) + chunks = [chunk async for chunk in scheduler.execute(TaskId("t"), {})] + assert chunks == [("updates", {"report": {"report": "done"}})] + assert pm.task.status == TaskStatus.COMPLETED + assert len(sink.events) == 1 + + +@pytest.mark.asyncio +async def test_scheduler_records_missing_terminal_output(): + pm = _ProcessManager() + scheduler = Scheduler(_EmptyGraph(), pm, _EventSink()) + with pytest.raises(RuntimeError, match="produced no final answer"): + async for _ in scheduler.execute(TaskId("t"), {}): + pass + assert pm.task.status == TaskStatus.FAILED + assert pm.error is not None + + +@pytest.mark.asyncio +async def test_scheduler_aborts_a_hung_graph(): + pm = _ProcessManager() + scheduler = Scheduler(_HangingGraph(), pm, _EventSink()) + with pytest.raises(TaskWallTimeExceeded): + async for _ in scheduler.execute(TaskId("t"), {}, wall_time_s=1): + pass + assert pm.task.status == TaskStatus.ABORTED + assert pm.abort_reason is not None diff --git a/tests/test_spawn_guard.py b/tests/test_spawn_guard.py new file mode 100644 index 0000000..38c27db --- /dev/null +++ b/tests/test_spawn_guard.py @@ -0,0 +1,437 @@ +"""Tests for SpawnGuard (Issue #24 Phase B). + +Covers: acquire/release, concurrency, depth, budget, timeout, +RAII context manager, and AgentBus integration. +""" + +from __future__ import annotations + +import asyncio +import contextlib +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from agent_core.components.agent_bus import ( + AgentBus, + DepthLimitExceeded, + SubTask, +) +from agent_core.components.agent_bus.spawn_guard import ( + BudgetExhausted, + SpawnDepthExceeded, + SpawnGuard, +) +from agent_core.loop_types import AgentLoopResult +from agent_core.models.agent_definition import AgentDefinition +from agent_core.models.task_budget import TaskBudget +from agent_core.protocols import EventSink +from agent_core.runtime.registries import services as registry +from agent_core.runtime.registries.agents import AgentRegistry +from agent_core.runtime.resources.manager import ResourceManager + + +class _FakeEventStore: + def __init__(self) -> None: + self.append = AsyncMock() + self.get_events = AsyncMock(return_value=[]) + + +@pytest.fixture(autouse=True) +def clear_registry(): + registry.clear() + yield + registry.clear() + + +def _setup_registry(): + agent_reg = AgentRegistry() + agent_reg.register(AgentDefinition( + role_id="researcher", + display_name="Researcher", + system_prompt="test", + allowed_tools=[], + color="#000", icon="agent", + )) + registry.register(AgentRegistry, agent_reg) + _es = _FakeEventStore() + registry.register(EventSink, _es) + resource_mgr = MagicMock(spec=ResourceManager) + resource_mgr.get_llm.return_value = MagicMock() + resource_mgr.get_tools_for_role.return_value = [] + registry.register(ResourceManager, resource_mgr) + + +# ── SpawnGuard standalone tests ───────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_acquire_release_basic(): + """Basic acquire + release cycle.""" + guard = SpawnGuard(TaskBudget(max_parallel=2, max_depth=3)) + + res = await guard.acquire("job-1", depth=1, estimated_tokens=1000) + assert res.job_id == "job-1" + assert res.depth == 1 + assert guard.active_count == 1 + assert guard.tokens_reserved == 1000 + + guard.release("job-1") + assert guard.active_count == 0 + assert guard.tokens_reserved == 0 + + +@pytest.mark.asyncio +async def test_acquire_depth_exceeded(): + """Should raise SpawnDepthExceeded at max depth.""" + guard = SpawnGuard(TaskBudget(max_depth=2)) + + with pytest.raises(SpawnDepthExceeded): + await guard.acquire("job-1", depth=2) + + +@pytest.mark.asyncio +async def test_acquire_budget_exhausted(): + """Should raise BudgetExhausted when tokens exceed remaining.""" + guard = SpawnGuard(TaskBudget(max_tokens=5000)) + + # First acquire takes most of the budget + await guard.acquire("job-1", depth=0, estimated_tokens=4000) + + # Second acquire exceeds remaining (1000) + with pytest.raises(BudgetExhausted): + await guard.acquire("job-2", depth=0, estimated_tokens=2000) + + guard.release("job-1") + + +@pytest.mark.asyncio +async def test_concurrency_limit_queues(): + """Should queue (not reject) when concurrency limit reached.""" + guard = SpawnGuard(TaskBudget(max_parallel=1, max_depth=5)) + + await guard.acquire("job-1", depth=0) + assert guard.active_count == 1 + + # Second acquire should block (we test by using a timeout) + acquired = asyncio.Event() + + async def try_acquire(): + await guard.acquire("job-2", depth=0) + acquired.set() + + task = asyncio.create_task(try_acquire()) + await asyncio.sleep(0.05) + assert not acquired.is_set() # Still blocked + + # Release first slot — second should proceed + guard.release("job-1") + await asyncio.sleep(0.05) + assert acquired.is_set() + assert guard.active_count == 1 + + guard.release("job-2") + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +async def test_release_idempotent(): + """release() called twice should not error.""" + guard = SpawnGuard(TaskBudget(max_parallel=2)) + + await guard.acquire("job-1", depth=0, estimated_tokens=500) + guard.release("job-1") + guard.release("job-1") # Second call — should be noop + assert guard.active_count == 0 + + +@pytest.mark.asyncio +async def test_release_corrects_token_estimate(): + """release() with actual_tokens should adjust accounting.""" + guard = SpawnGuard(TaskBudget(max_tokens=10000)) + + await guard.acquire("job-1", depth=0, estimated_tokens=3000) + assert guard.tokens_reserved == 3000 + + guard.release("job-1", actual_tokens=2000) + assert guard.tokens_reserved == 0 # Estimate removed + assert guard.tokens_actual == 2000 # Actual recorded + + +@pytest.mark.asyncio +async def test_raii_context_manager(): + """reservation() should auto-release on exit.""" + guard = SpawnGuard(TaskBudget(max_parallel=2, max_depth=3)) + + async with guard.reservation("job-1", depth=1, estimated_tokens=500) as res: + assert res.job_id == "job-1" + assert guard.active_count == 1 + + assert guard.active_count == 0 + + +@pytest.mark.asyncio +async def test_raii_releases_on_exception(): + """reservation() should release even if body raises.""" + guard = SpawnGuard(TaskBudget(max_parallel=2, max_depth=3)) + + with pytest.raises(ValueError): + async with guard.reservation("job-1", depth=1): + raise ValueError("boom") + + assert guard.active_count == 0 + + +@pytest.mark.asyncio +async def test_release_without_acquire_slot_does_not_leak_semaphore(): + """Regression: ``release`` after ``pre_check`` only must not return + a slot we never owned. + + Reproduces the scenario where ``bus.py:_run_and_finalize`` ran + ``pre_check`` (registering the reservation), then was cancelled + before reaching ``acquire_slot`` so the semaphore was never taken, + and finally hit its ``finally`` block which always calls ``release``. + Pre-fix the unconditional ``self._semaphore.release()`` inflated the + counter past ``max_parallel``, breaking the concurrency invariant. + Post-fix the reservation's ``acquired`` flag is ``False`` so the + semaphore is left intact. + """ + guard = SpawnGuard(TaskBudget(max_parallel=1, max_depth=5)) + + # Consume the only slot. + await guard.acquire("job-1", depth=0) + + # Simulate "pre_check ran, acquire_slot was cancelled before being + # awaited, finally block now calls release" — the exact path + # bus.py walks when its outer task is cancelled mid-await. + await guard.pre_check("job-2", depth=0, estimated_tokens=500) + guard.release("job-2") # finally-block fallback + + # job-1 still owns the only slot. A bystander acquire must still + # block — pre-fix this resolved instantly because the bogus + # ``self._semaphore.release()`` bumped the counter to 1 free. + bystander = asyncio.create_task(guard.acquire("job-3", depth=0)) + await asyncio.sleep(0.05) + assert not bystander.done(), ( + "release() leaked a semaphore slot — concurrency limit broken" + ) + + # Real release of job-1 must let the waiter through, proving the + # semaphore counter is still consistent after the dance. + guard.release("job-1") + await asyncio.wait_for(bystander, timeout=0.5) + guard.release("job-3") + assert guard.active_count == 0 + + +@pytest.mark.asyncio +async def test_release_after_acquire_slot_returns_semaphore_normally(): + """Sanity check: when ``acquire_slot`` did succeed, ``release`` must + still hand the slot back. Guards against an over-eager ``acquired`` + check that would also mask normal releases.""" + guard = SpawnGuard(TaskBudget(max_parallel=1, max_depth=5)) + + res = await guard.acquire("job-1", depth=0) + assert res.acquired is True + + guard.release("job-1") + + # Slot is free again — bystander resolves without blocking. + await asyncio.wait_for( + guard.acquire("job-2", depth=0), timeout=0.5, + ) + guard.release("job-2") + + +@pytest.mark.asyncio +async def test_stats(): + """stats() should return current state.""" + guard = SpawnGuard(TaskBudget( + max_parallel=3, max_depth=2, max_tokens=50000, + )) + await guard.acquire("job-1", depth=0, estimated_tokens=1000) + + s = guard.stats() + assert s["active"] == 1 + assert s["max_parallel"] == 3 + assert s["max_depth"] == 2 + assert s["tokens_reserved"] == 1000 + assert s["max_tokens"] == 50000 + assert s["total_spawns"] == 1 + + guard.release("job-1") + + +# ── AgentBus + SpawnGuard integration ─────────────────────────────────── + + +@pytest.mark.asyncio +async def test_agent_bus_with_spawn_guard_enforces_concurrency(monkeypatch): + """AgentBus with SpawnGuard should respect max_parallel.""" + _setup_registry() + + started = [] + gate = asyncio.Event() + + async def gated_agent_loop(**kwargs): + task_id = "" + cfg = kwargs.get("config") + if cfg: + task_id = getattr(cfg, "task_id", "") + started.append(task_id) + await gate.wait() + return AgentLoopResult( + messages=[], final_content="ok", metadata={}, + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + gated_agent_loop, + ) + + guard = SpawnGuard(TaskBudget(max_parallel=1, max_depth=3)) + bus = AgentBus() + bus.set_spawn_guard(guard) + + # Submit 2 jobs — only 1 should start (concurrency=1) + jid1 = await bus.submit("t", SubTask("q1")) + jid2 = await bus.submit("t", SubTask("q2")) + await asyncio.sleep(0.1) + + # Only 1 should have started (other queued at semaphore) + assert len(started) == 1 + + # Unblock — both should eventually complete + gate.set() + result = await bus.collect([jid1, jid2], timeout=5) + assert len(result.completed) == 2 + + +@pytest.mark.asyncio +async def test_agent_bus_with_spawn_guard_depth_limit(monkeypatch): + """AgentBus should enforce SpawnGuard depth limit.""" + _setup_registry() + + guard = SpawnGuard(TaskBudget(max_depth=1)) + bus = AgentBus() + bus.set_spawn_guard(guard) + + # depth=1 should be rejected (max_depth=1 means depth<1 only) + with pytest.raises(DepthLimitExceeded): + await bus.submit("t", SubTask("q"), current_depth=1, max_depth=5) + + +@pytest.mark.asyncio +async def test_agent_bus_with_spawn_guard_budget_exhausted(monkeypatch): + """AgentBus should reject when SpawnGuard budget exhausted.""" + _setup_registry() + + guard = SpawnGuard(TaskBudget(max_tokens=1000, max_depth=3)) + bus = AgentBus() + bus.set_spawn_guard(guard) + + with pytest.raises(BudgetExhausted): + await bus.submit( + "t", SubTask("q"), estimated_tokens=2000, + ) + + +@pytest.mark.asyncio +async def test_agent_bus_spawn_guard_releases_on_completion(monkeypatch): + """SpawnGuard slot should be released when job finishes.""" + _setup_registry() + + async def fast_agent_loop(**kwargs): + return AgentLoopResult( + messages=[], final_content="ok", metadata={}, + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + fast_agent_loop, + ) + + guard = SpawnGuard(TaskBudget(max_parallel=2, max_depth=3)) + bus = AgentBus() + bus.set_spawn_guard(guard) + + jid = await bus.submit("t", SubTask("q")) + await bus.collect([jid], timeout=5) + + # Guard slot should be released + assert guard.active_count == 0 + assert guard.total_spawns == 1 + + +@pytest.mark.asyncio +async def test_agent_bus_spawn_guard_releases_on_failure(monkeypatch): + """SpawnGuard slot should be released when job fails.""" + _setup_registry() + + async def failing_agent_loop(**kwargs): + raise RuntimeError("fail") + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + failing_agent_loop, + ) + + guard = SpawnGuard(TaskBudget(max_parallel=2, max_depth=3)) + bus = AgentBus() + bus.set_spawn_guard(guard) + + jid = await bus.submit("t", SubTask("q")) + await bus.collect([jid], timeout=5) + + assert guard.active_count == 0 + + +@pytest.mark.asyncio +async def test_agent_bus_spawn_guard_releases_on_abort(monkeypatch): + """SpawnGuard slot should be released on abort.""" + _setup_registry() + + async def slow_agent_loop(**kwargs): + await asyncio.sleep(100) + return AgentLoopResult(messages=[], final_content="", metadata={}) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + slow_agent_loop, + ) + + guard = SpawnGuard(TaskBudget(max_parallel=2, max_depth=3)) + bus = AgentBus() + bus.set_spawn_guard(guard) + + jid = await bus.submit("t", SubTask("q")) + await asyncio.sleep(0.05) + assert guard.active_count == 1 + + await bus.abort(jid) + # Give finalize a moment to run + await asyncio.sleep(0.05) + assert guard.active_count == 0 + + +@pytest.mark.asyncio +async def test_agent_bus_without_spawn_guard_still_works(monkeypatch): + """AgentBus without SpawnGuard should work as before.""" + _setup_registry() + + async def fast_agent_loop(**kwargs): + return AgentLoopResult( + messages=[], final_content="ok", metadata={}, + ) + + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", + fast_agent_loop, + ) + + bus = AgentBus() # No spawn guard + jid = await bus.submit("t", SubTask("q")) + result = await bus.collect([jid], timeout=5) + assert len(result.completed) == 1 From 93e89cf464733c8f5210db62e6a2ce7748e4e3f3 Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 12:18:17 +0800 Subject: [PATCH 2/7] feat: inject scheduler wall-time policy --- agent_core/scheduling/__init__.py | 2 ++ agent_core/scheduling/scheduler.py | 9 +++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/agent_core/scheduling/__init__.py b/agent_core/scheduling/__init__.py index c9ed5c6..62cff1f 100644 --- a/agent_core/scheduling/__init__.py +++ b/agent_core/scheduling/__init__.py @@ -5,6 +5,7 @@ Scheduler, SchedulerTask, TaskWallTimeExceeded, + WallTimeResolver, resolve_wall_time_s, ) from agent_core.scheduling.workflow_defaults import ( @@ -18,6 +19,7 @@ "Scheduler", "SchedulerTask", "TaskWallTimeExceeded", + "WallTimeResolver", "clear_workflow_defaults", "get_workflow_default", "register_workflow_defaults", diff --git a/agent_core/scheduling/scheduler.py b/agent_core/scheduling/scheduler.py index 09f5673..99946e2 100644 --- a/agent_core/scheduling/scheduler.py +++ b/agent_core/scheduling/scheduler.py @@ -6,7 +6,7 @@ import asyncio import logging import os -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Callable from typing import Any, Protocol from typing import Any as _GraphType # MiniDAGRunner via duck typing @@ -42,6 +42,9 @@ async def abort_task(self, task_id: TaskId, reason: str) -> Any: ... async def set_error(self, task_id: TaskId, error: str) -> Any: ... +WallTimeResolver = Callable[[int | None, str | None], int | None] + + class TaskWallTimeExceeded(RuntimeError): """Task exceeded its configured wall-time budget. @@ -121,6 +124,7 @@ def __init__( graph_builder: Any | None = None, pipeline_registry: Any | None = None, checkpointer: Any | None = None, + wall_time_resolver: WallTimeResolver = resolve_wall_time_s, ) -> None: self._default_graph = graph self._pm = process_manager @@ -128,6 +132,7 @@ def __init__( self._graph_builder = graph_builder self._pipeline_registry = pipeline_registry self._checkpointer = checkpointer + self._wall_time_resolver = wall_time_resolver self._compiled_cache: dict[str, _GraphType] = {} def _get_graph(self, pipeline_id: str | None = None) -> _GraphType: @@ -223,7 +228,7 @@ async def execute( }, } - wall_time = resolve_wall_time_s(wall_time_s, pipeline_id=pipeline_id) + wall_time = self._wall_time_resolver(wall_time_s, pipeline_id) try: async for chunk in self._astream_with_wall_time( From 2672df25724011a4adebc2268a6f0c08b29cdc91 Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 12:20:50 +0800 Subject: [PATCH 3/7] feat: add product composition hooks for agent bus --- agent_core/components/agent_bus/__init__.py | 6 ++++- agent_core/components/agent_bus/bus.py | 26 +++++++++++++++++---- 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/agent_core/components/agent_bus/__init__.py b/agent_core/components/agent_bus/__init__.py index 5e17779..540574d 100644 --- a/agent_core/components/agent_bus/__init__.py +++ b/agent_core/components/agent_bus/__init__.py @@ -5,7 +5,10 @@ """ from agent_core.components.agent_bus.agent_comm import AgentComm, DeliveryMode -from agent_core.components.agent_bus.bus import AgentBus +from agent_core.components.agent_bus.bus import ( + AgentBus, + configure_default_pause_check_factory, +) from agent_core.components.agent_bus.fan_in import ( COMPLETE_STOP_REASONS, INCOMPLETE_STOP_REASONS, @@ -79,6 +82,7 @@ "build_session_loop_config", "classify_completion", "close_session_boundary_aborted", + "configure_default_pause_check_factory", "format_report_block", "format_status_line", "format_status_report_block", diff --git a/agent_core/components/agent_bus/bus.py b/agent_core/components/agent_bus/bus.py index 8d3529e..2a97128 100644 --- a/agent_core/components/agent_bus/bus.py +++ b/agent_core/components/agent_bus/bus.py @@ -107,6 +107,15 @@ PauseCheckFn = Callable[[], Awaitable[bool]] PauseCheckFactory = Callable[[str], PauseCheckFn | None] +_default_pause_check_factory: PauseCheckFactory | None = None + + +def configure_default_pause_check_factory( + factory: PauseCheckFactory | None, +) -> None: + """Configure the host pause integration used by zero-argument buses.""" + global _default_pause_check_factory + _default_pause_check_factory = factory def _with_wall_clock_guard( @@ -156,8 +165,9 @@ def _strip_job_suffix(task_id: str) -> str: def _legacy_pause_check_factory(task_id: str) -> PauseCheckFn | None: """Product-neutral default: no process-manager pause integration.""" - del task_id - return None + if _default_pause_check_factory is None: + return None + return _default_pause_check_factory(task_id) def _safe_session_filename(session_id: str) -> str: @@ -415,8 +425,16 @@ def _resource_manager(self, *, required: bool) -> Any: return self._resource_manager_injected from agent_core.runtime.resources.manager import ResourceManager if required: - return registry.get(ResourceManager) - return registry.get_optional(ResourceManager) + direct = registry.get_optional(ResourceManager) + if direct is not None: + return direct + legacy = registry.get_optional_by_type_name("ResourceManager") + if legacy is None: + return registry.get(ResourceManager) + return legacy + return registry.get_optional(ResourceManager) or registry.get_optional_by_type_name( + "ResourceManager" + ) def _task_pause_check_or_none(self, task_id: str | None) -> PauseCheckFn | None: """Resolve a per-task pause check closure. From df16017e556eb4f424d3d52beb34e9295a995778 Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 12:24:03 +0800 Subject: [PATCH 4/7] feat: configure product agent bus defaults --- agent_core/components/agent_bus/__init__.py | 4 +++ agent_core/components/agent_bus/bus.py | 38 +++++++++++++++++++-- 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/agent_core/components/agent_bus/__init__.py b/agent_core/components/agent_bus/__init__.py index 540574d..ece7575 100644 --- a/agent_core/components/agent_bus/__init__.py +++ b/agent_core/components/agent_bus/__init__.py @@ -7,7 +7,9 @@ from agent_core.components.agent_bus.agent_comm import AgentComm, DeliveryMode from agent_core.components.agent_bus.bus import ( AgentBus, + configure_default_event_sink_resolver, configure_default_pause_check_factory, + configure_default_session_activity, ) from agent_core.components.agent_bus.fan_in import ( COMPLETE_STOP_REASONS, @@ -82,7 +84,9 @@ "build_session_loop_config", "classify_completion", "close_session_boundary_aborted", + "configure_default_event_sink_resolver", "configure_default_pause_check_factory", + "configure_default_session_activity", "format_report_block", "format_status_line", "format_status_report_block", diff --git a/agent_core/components/agent_bus/bus.py b/agent_core/components/agent_bus/bus.py index 2a97128..2066749 100644 --- a/agent_core/components/agent_bus/bus.py +++ b/agent_core/components/agent_bus/bus.py @@ -108,6 +108,9 @@ PauseCheckFn = Callable[[], Awaitable[bool]] PauseCheckFactory = Callable[[str], PauseCheckFn | None] _default_pause_check_factory: PauseCheckFactory | None = None +EventSinkResolver = Callable[[], EventSink | None] +_default_event_sink_resolver: EventSinkResolver | None = None +_default_session_activity = True def configure_default_pause_check_factory( @@ -118,6 +121,20 @@ def configure_default_pause_check_factory( _default_pause_check_factory = factory +def configure_default_event_sink_resolver( + resolver: EventSinkResolver | None, +) -> None: + """Configure host lookup for a product-specific EventSink protocol key.""" + global _default_event_sink_resolver + _default_event_sink_resolver = resolver + + +def configure_default_session_activity(enabled: bool) -> None: + """Enable or disable the portable live session activity observer.""" + global _default_session_activity + _default_session_activity = bool(enabled) + + def _with_wall_clock_guard( observers: list[Any] | None, guard: SpawnGuard | None, @@ -167,7 +184,10 @@ def _legacy_pause_check_factory(task_id: str) -> PauseCheckFn | None: """Product-neutral default: no process-manager pause integration.""" if _default_pause_check_factory is None: return None - return _default_pause_check_factory(task_id) + try: + return _default_pause_check_factory(task_id) + except (ImportError, RuntimeError): + return None def _safe_session_filename(session_id: str) -> str: @@ -325,6 +345,7 @@ def __init__( agent_registry: AgentRegistry | None = None, resource_manager: Any = None, session_history_dir: Path | None = None, + session_activity: bool | None = None, ) -> None: """Construct an AgentBus. @@ -385,6 +406,11 @@ def __init__( self._agent_registry_injected: AgentRegistry | None = agent_registry self._resource_manager_injected: Any = resource_manager self._session_history_dir: Path | None = session_history_dir + self._session_activity = ( + _default_session_activity + if session_activity is None + else bool(session_activity) + ) def _event_sink(self) -> Any: """Resolve the active event sink. @@ -399,6 +425,13 @@ def _event_sink(self) -> Any: event_sink = registry.get_optional(EventSink) if event_sink is not None: return event_sink + if _default_event_sink_resolver is not None: + try: + event_sink = _default_event_sink_resolver() + except (ImportError, RuntimeError): + event_sink = None + if event_sink is not None: + return event_sink return registry.get_optional_by_type_name("EventStore") def _agent_registry(self, *, required: bool) -> Any: @@ -1273,7 +1306,8 @@ async def _run() -> SubAgentResult: type(_obs).__name__, _exc, ) - loop_observers.append(_SessionActivityObserver(session)) + if self._session_activity: + loop_observers.append(_SessionActivityObserver(session)) loop_observers = _with_wall_clock_guard(loop_observers, guard) coro = run_agent_loop( From 383d746c2e30797f046d0f2098ad8b1c67efc5ac Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 12:26:07 +0800 Subject: [PATCH 5/7] fix: resolve product event sink for session telemetry --- agent_core/components/agent_bus/bus.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/agent_core/components/agent_bus/bus.py b/agent_core/components/agent_bus/bus.py index 2066749..9b48de8 100644 --- a/agent_core/components/agent_bus/bus.py +++ b/agent_core/components/agent_bus/bus.py @@ -1488,7 +1488,7 @@ async def _run_and_finalize() -> SubAgentResult: session.pending_results.append(result) await _emit_session_task_completed( session, job_id, result, - event_sink=self._event_sink_injected, + event_sink=self._event_sink(), spawn_context=pending.spawn_context, ) return result @@ -1567,7 +1567,7 @@ async def _run_and_finalize() -> SubAgentResult: await _emit_session_task_submitted( session, job_id, task_prompt, - event_sink=self._event_sink_injected, + event_sink=self._event_sink(), ) async def _eager_trim_and_offload(self, session: SubAgentSession) -> None: @@ -2224,5 +2224,5 @@ async def dispatch_parallel( async def rehydrate_active_jobs(self) -> int: """Mark jobs whose asyncio.Task was lost across a restart as orphaned.""" return await _rehydrate_orphaned_jobs( - self._jobs, event_sink=self._event_sink_injected, + self._jobs, event_sink=self._event_sink(), ) From fcbe5cd9a8d950191abf5339b26c0a56e3f4c91e Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 12:40:23 +0800 Subject: [PATCH 6/7] fix: preserve portable scheduler and fan-in contracts --- agent_core/components/agent_bus/fan_in.py | 5 ++ agent_core/scheduling/scheduler.py | 72 +++++++++++++++++++++++ tests/test_scheduler.py | 38 ++++++++++++ 3 files changed, 115 insertions(+) diff --git a/agent_core/components/agent_bus/fan_in.py b/agent_core/components/agent_bus/fan_in.py index d68a0c8..c577024 100644 --- a/agent_core/components/agent_bus/fan_in.py +++ b/agent_core/components/agent_bus/fan_in.py @@ -59,6 +59,7 @@ # rather than merely answer-less. "response_truncated", "exception", + "thrash_no_progress", }) _INCOMPLETE_NOTES: dict[str, str] = { @@ -89,6 +90,10 @@ "exception": ( "agent terminated with an unhandled exception; report is partial" ), + "thrash_no_progress": ( + "coordinator was spinning on sub-agent create/assign without new " + "information; stopped and synthesised from collected reports" + ), } diff --git a/agent_core/scheduling/scheduler.py b/agent_core/scheduling/scheduler.py index 99946e2..6feb62f 100644 --- a/agent_core/scheduling/scheduler.py +++ b/agent_core/scheduling/scheduler.py @@ -189,6 +189,14 @@ def _register_pipeline_agents(self, spec: Any) -> None: except Exception as e: logger.warning("Failed to auto-register agents: %s", e) + def execute_dynamic_spec(self, spec: Any) -> None: + """Register a dynamically generated pipeline specification.""" + if ( + self._pipeline_registry + and not self._pipeline_registry.has(spec.pipeline_id) + ): + self._pipeline_registry.register(spec) + async def execute( self, task_id: TaskId, @@ -384,6 +392,45 @@ async def _astream_with_wall_time( if aclose is not None: await aclose() + async def get_resume_config( + self, + task_id: TaskId, + *, + pipeline_id: str | None = None, + ) -> tuple[dict[str, Any], dict[str, Any]]: + """Return the effective checkpoint configuration and state for resume. + + If the latest checkpoint contains an error write, its parent is used so + the failed node is retried instead of being treated as terminal. + """ + task = await self._pm.get_task(task_id) + resume_config = {"configurable": {"thread_id": task.thread_id}} + + if not self._checkpointer: + return resume_config, {} + + latest = await self._checkpointer.aget_tuple(resume_config) + if not latest: + return resume_config, {} + + effective_config = resume_config + pending_writes = latest.pending_writes or [] + has_error_write = any( + len(item) >= 2 and item[1] == "__error__" + for item in pending_writes + ) + if has_error_write and latest.parent_config: + effective_config = latest.parent_config + + graph = None + try: + graph = self._get_graph(pipeline_id) + except Exception: + graph = None + + state = await self._get_state_for_config(graph, effective_config) + return effective_config, state + async def get_state(self, task_id: TaskId) -> dict[str, Any]: """Get the current state snapshot for a task. @@ -506,3 +553,28 @@ def _is_stale_execution(task: Any, execution_token: str | None) -> bool: return False context = getattr(task, "context", {}) or {} return context.get("active_execution_token") != execution_token + + async def get_history(self, task_id: TaskId) -> list[dict[str, Any]]: + """Return checkpoint history in a product-neutral representation.""" + task = await self._pm.get_task(task_id) + graph = self._get_graph(getattr(task, "pipeline_id", None)) + config = {"configurable": {"thread_id": task.thread_id}} + history: list[dict[str, Any]] = [] + async for snapshot in graph.aget_state_history(config): + snapshot_values = getattr(snapshot, "values", {}) or {} + snapshot_next = getattr(snapshot, "next", ()) or () + snapshot_tasks = getattr(snapshot, "tasks", ()) or () + history.append( + { + "values": snapshot_values, + "next": list(snapshot_next), + "tasks": [ + { + "name": getattr(item, "name", ""), + "id": getattr(item, "id", ""), + } + for item in snapshot_tasks + ], + } + ) + return history diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 4adbf92..56b5c3f 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -7,6 +7,7 @@ import pytest +from agent_core.components.agent_bus.fan_in import INCOMPLETE_STOP_REASONS from agent_core.scheduling.scheduler import ( Scheduler, TaskWallTimeExceeded, @@ -82,6 +83,17 @@ async def astream(self, *_args: Any, **_kwargs: Any): await asyncio.sleep(60) +class _PipelineRegistry: + def __init__(self) -> None: + self.specs: list[Any] = [] + + def has(self, pipeline_id: str) -> bool: + return any(spec.pipeline_id == pipeline_id for spec in self.specs) + + def register(self, spec: Any) -> None: + self.specs.append(spec) + + def test_wall_time_resolution_supports_portable_and_legacy_env(monkeypatch): monkeypatch.setenv("MIROHARNESS_TASK_WALL_TIME_S", "120") assert resolve_wall_time_s(None) == 120 @@ -99,6 +111,32 @@ def test_workflow_defaults_are_product_injected(monkeypatch): ): monkeypatch.delenv(name, raising=False) clear_workflow_defaults() + + +def test_scheduler_registers_dynamic_specs_once(): + registry = _PipelineRegistry() + scheduler = Scheduler( + _ReportGraph(), + _ProcessManager(), + _EventSink(), + pipeline_registry=registry, + ) + spec = SimpleNamespace(pipeline_id="dynamic") + scheduler.execute_dynamic_spec(spec) + scheduler.execute_dynamic_spec(spec) + assert registry.specs == [spec] + + +@pytest.mark.asyncio +async def test_resume_config_without_checkpointer_uses_task_thread(): + scheduler = Scheduler(_ReportGraph(), _ProcessManager(), _EventSink()) + config, state = await scheduler.get_resume_config(TaskId("t")) + assert config == {"configurable": {"thread_id": "thread"}} + assert state == {} + + +def test_agent_team_thrash_is_an_incomplete_stop_reason(): + assert "thrash_no_progress" in INCOMPLETE_STOP_REASONS register_workflow_defaults( ("research", "research-report"), {"TASK_WALL_TIME_S": 900, "TASK_WALL_TIME_MODE": "soft_research"}, From ef143b93834574d0d8af3fedc71c073a98218cbc Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 13:31:06 +0800 Subject: [PATCH 7/7] fix: close agent-bus contract gaps found in review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four defects in the extracted scheduling kernel, all of the same shape: code carried over from the source implementation kept an assumption its original environment satisfied, while the portable abstraction changed. SpawnGuard reservations leaked on any failed submit. `pre_check` reserves tokens and registers the job, but the RAII release lives in `_run_and_finalize`, which does not exist until `create_task` has run — so an unregistered role (`get_prompt_for` raising) or a failing session dispatch leaked the reservation permanently, decaying `remaining_tokens` and the concurrency count for the life of the bus. Resolve the system prompt before reserving, register the job entry before spawning, and release explicitly on every path that can raise in between. Sessions get the same treatment on both the free and the queued dispatch — the latter matters because `_drain_session_queue` swallows dispatch failures by design, so nothing upstream is left to hand the reservation back. `AgentComm.consume` ordered its durable cursor with `int(event.id)`, but `EventId` is an opaque string and `new_event_id()` returns uuid4 hex. Add an explicit `KernelEvent.seq` ordinal and read it through `event_ordinal`, falling back to a decimal id so stores that expose their integer primary key keep working. A store that stamps neither now raises `EventStoreContractError` naming the fix instead of a bare `ValueError`, and `AgentBus` logs the lost durable-recovery capability at ERROR rather than swallowing it at debug — that blanket `except Exception` is why the whole L1-0 path could be dead with nothing in the logs. `_AgentCommEventStore` declared `EventSink.append(task_id, event_type, payload, agent_role)` but was called with a whole `KernelEvent`, so a host implementing the documented protocol bound the event to `task_id` and returned None. Messaging genuinely needs the richer append — `from_agent` / `to_agent` / `message_type` / `correlation_id` are the fields the reader queries on — so declare that signature and drop the `EventSink` inheritance. `EventSink` stays exactly as it was for telemetry; the two roles are now distinct and documented. `resolve_research_wall` selected its input by key presence, so `wall_deadline_s` was unreadable whenever `research_wall_time_s` existed: `hard_total_s` came back 0, `check_wall_feasibility` passed unconditionally, and the total-task ceiling was silently unenforced. Read both keys independently and select on value validity. Contract tests cover all four, including the store-ordering contract the `consume` docstring referenced but that did not exist in this package. Validation: ruff clean, pyright 0 errors, 849 passed (was 812), `uv build`, plus a smoke run of all four paths against the built wheel. Co-Authored-By: Claude Opus 5 (1M context) --- agent_core/components/agent_bus/__init__.py | 9 +- agent_core/components/agent_bus/agent_comm.py | 80 +++- agent_core/components/agent_bus/bus.py | 115 +++++- agent_core/components/finalization/budget.py | 76 ++-- agent_core/models/event.py | 8 + docs/agent-bus-boundary.md | 43 ++ ...test_agent_bus_result_recovery_contract.py | 96 +++++ tests/test_agent_bus_spawn_reservation.py | 236 +++++++++++ tests/test_agent_comm_store_contract.py | 371 ++++++++++++++++++ tests/test_finalization_budget.py | 62 +++ 10 files changed, 1026 insertions(+), 70 deletions(-) create mode 100644 tests/test_agent_bus_result_recovery_contract.py create mode 100644 tests/test_agent_bus_spawn_reservation.py create mode 100644 tests/test_agent_comm_store_contract.py diff --git a/agent_core/components/agent_bus/__init__.py b/agent_core/components/agent_bus/__init__.py index ece7575..afddf8e 100644 --- a/agent_core/components/agent_bus/__init__.py +++ b/agent_core/components/agent_bus/__init__.py @@ -4,7 +4,12 @@ (the former flat-file module). Re-exports the public API from submodules. """ -from agent_core.components.agent_bus.agent_comm import AgentComm, DeliveryMode +from agent_core.components.agent_bus.agent_comm import ( + AgentComm, + DeliveryMode, + EventStoreContractError, + event_ordinal, +) from agent_core.components.agent_bus.bus import ( AgentBus, configure_default_event_sink_resolver, @@ -65,6 +70,7 @@ "CompletionPolicy", "DeliveryMode", "DepthLimitExceeded", + "EventStoreContractError", "FanInBatch", "JobEntry", "PendingSessionTask", @@ -87,6 +93,7 @@ "configure_default_event_sink_resolver", "configure_default_pause_check_factory", "configure_default_session_activity", + "event_ordinal", "format_report_block", "format_status_line", "format_status_report_block", diff --git a/agent_core/components/agent_bus/agent_comm.py b/agent_core/components/agent_bus/agent_comm.py index 44401bc..40411a8 100644 --- a/agent_core/components/agent_bus/agent_comm.py +++ b/agent_core/components/agent_bus/agent_comm.py @@ -17,31 +17,81 @@ from agent_core.events import EventType from agent_core.models.agent_message import AgentMessage from agent_core.models.event import KernelEvent -from agent_core.protocols import EventReader, EventSink +from agent_core.protocols import EventReader from agent_core.runtime.events import EventBus from agent_core.types import TaskId logger = logging.getLogger(__name__) +class EventStoreContractError(RuntimeError): + """The event store violated the :class:`_AgentCommEventStore` contract. + + Distinct from a transient read failure on purpose: a caller that treats + read errors as soft (``AgentBus`` result recovery does) must still + surface a misconfigured store loudly, because every future read will + fail the same way. + """ + + +def event_ordinal(event: KernelEvent) -> int: + """Return the store-assigned total-order ordinal of ``event``. + + Prefers the explicit :attr:`KernelEvent.seq`. Falls back to a decimal + ``id`` so a store that exposes its integer primary key as the event id + keeps working unchanged. + + Anything else cannot carry a total order — an unset id, or the opaque + hex of :func:`agent_core.types.new_event_id`, whose values are neither + numeric nor monotonic. Those are rejected with an actionable + :class:`EventStoreContractError` rather than the bare ``ValueError`` + ``int()`` would raise, because the fix is a store change and the + message has to say so. + """ + if event.seq is not None: + return event.seq + raw = str(event.id) + if raw.isdecimal(): + return int(raw) + raise EventStoreContractError( + "Event store must stamp every appended event with a total-order " + "ordinal: set KernelEvent.seq to a monotonically increasing integer " + "(a decimal KernelEvent.id is also accepted). Got " + f"id={raw!r}, seq=None — an opaque id such as new_event_id() cannot " + "order a durable cursor, so AgentComm.consume cannot advance safely." + ) + + @runtime_checkable -class _AgentCommEventStore(EventSink, EventReader, Protocol): - """AgentComm needs both append (EventSink) and cursored reads - (EventReader) — combined here because Python lacks an intersection - type. Public Protocols stay minimal in ``core.protocols``. - - The combination has a stronger semantic contract than ``EventSink`` - alone: ``append`` retains the event, assigns and returns a unique - integer id, and sequential awaited appends receive increasing ids. - The id — not ``timestamp`` — is the total order. Concurrent appends - have no defined relative order, only that every reader observes the - same resulting order. +class _AgentCommEventStore(EventReader, Protocol): + """AgentComm needs a whole-event append plus cursored reads + (:class:`~agent_core.protocols.EventReader`) — combined here because + Python lacks an intersection type. Public Protocols stay minimal in + ``core.protocols``. + + ``append`` takes a :class:`KernelEvent`, *not* the flattened + ``(task_id, event_type, payload, agent_role)`` of + :class:`~agent_core.protocols.EventSink`: an inter-agent message carries + ``from_agent`` / ``to_agent`` / ``message_type`` / ``correlation_id`` as + first-class fields that the reader queries on, and ``EventSink`` cannot + express them. A host that implements only ``EventSink`` is a valid + telemetry sink for ``AgentBus`` but is *not* a valid AgentComm store. + + The combination has a stronger semantic contract than either half: + ``append`` retains the event and returns it stamped with a total-order + ordinal (see :func:`event_ordinal`), and sequential awaited appends + receive increasing ordinals. That ordinal — not ``timestamp`` — is the + total order, and ``EventReader``'s ``after_id`` cursor is expressed in + it. Concurrent appends have no defined relative order, only that every + reader observes the same resulting order. Ordering is not guaranteed across live writer processes sharing one JSONL path. Nor is losslessness guaranteed for a capacity-bounded in-memory store, which evicts its oldest retained events. """ + async def append(self, event: KernelEvent) -> KernelEvent: ... + class DeliveryMode(StrEnum): """Message delivery mode. @@ -164,8 +214,8 @@ async def consume( shared cursor would let one consumer's read carry another consumer's mail past unread. - ``tests/state/test_event_ordering_contract.py`` pins all three - against every retained, readable store implementation. + ``tests/test_agent_comm_store_contract.py`` pins all three, plus + the ordinal requirement itself, against a reference store. """ cursor_key = (agent_id, task_id, message_type) after_id = self._cursors.get(cursor_key, 0) @@ -187,7 +237,7 @@ async def consume( message_type=message_type, ) if events: - ids = [int(event.id) for event in events] + ids = [event_ordinal(event) for event in events] if any(current <= previous for previous, current in itertools.pairwise(ids)): # Advancing to max(ids) can permanently skip lower ids that # a misordered, limited query left outside this batch. Keep diff --git a/agent_core/components/agent_bus/bus.py b/agent_core/components/agent_bus/bus.py index 9b48de8..6768f9b 100644 --- a/agent_core/components/agent_bus/bus.py +++ b/agent_core/components/agent_bus/bus.py @@ -562,6 +562,18 @@ async def submit( agent_registry = self._agent_registry(required=True) dispatch_depth = current_depth + 1 + # Resolved before the guard reserves anything: an unregistered role + # makes ``get_prompt_for`` raise, and a reservation taken first would + # never come back — the RAII release lives in ``_run_and_finalize``, + # which does not exist until the asyncio task below is created. A + # leaked reservation is permanent, so it decays ``remaining_tokens`` + # and the concurrency count for the life of the bus. + system_prompt = ( + item.system_prompt + or agent_registry.get_prompt_for(item.role_id) + ) + runtime = runtime_spec or SubAgentRuntimeSpec() + # Layer 5+3: SpawnGuard pre-check (depth + budget — non-blocking) # Layer 4 (concurrency) is deferred to _run_and_finalize if guard: @@ -569,12 +581,6 @@ async def submit( job_id, dispatch_depth, estimated_tokens, ) - system_prompt = ( - item.system_prompt - or agent_registry.get_prompt_for(item.role_id) - ) - runtime = runtime_spec or SubAgentRuntimeSpec() - async def _run() -> SubAgentResult: # Declared outside try/except so exception handlers can # salvage partial metadata on failure (see the ``agent_result`` @@ -719,16 +725,27 @@ async def _run_and_finalize() -> SubAgentResult: if guard: guard.release(job_id) - task = asyncio.create_task(_run_and_finalize(), name=job_id) + # Register the entry before spawning: ``_run_and_finalize`` looks + # itself up in ``self._jobs``, and the reservation has no RAII owner + # until that task exists — so anything that fails in between has to + # hand the reservation back explicitly. entry = JobEntry( job_id=job_id, parent_task_id=parent_task_id, item=item, - task=task, + task=None, status="submitted", submitted_at=time.monotonic(), ) self._jobs[job_id] = entry + try: + task = asyncio.create_task(_run_and_finalize(), name=job_id) + except BaseException: + self._jobs.pop(job_id, None) + if guard: + guard.release(job_id) + raise + entry.task = task event_store = self._event_sink() if event_store: @@ -1138,15 +1155,6 @@ async def submit_task_to_session( dispatch_depth = 1 job_id = self._next_job_id(session.task_id) - - # Layer 5+3: SpawnGuard pre-check at submit time (not at dispatch) - # so queued tasks still raise BudgetExhausted / SpawnDepthExceeded - # synchronously at the API boundary — preserves the legacy - # behaviour that callers see budget errors immediately rather - # than only via ``collect_reports`` after the queue drains. - if guard is not None: - await guard.pre_check(job_id, dispatch_depth, estimated_tokens) - pending = PendingSessionTask( job_id=job_id, task_prompt=task_prompt, @@ -1158,6 +1166,42 @@ async def submit_task_to_session( task_metadata=dict(task_metadata or {}), ) + # Layer 5+3: SpawnGuard pre-check at submit time (not at dispatch) + # so queued tasks still raise BudgetExhausted / SpawnDepthExceeded + # synchronously at the API boundary — preserves the legacy + # behaviour that callers see budget errors immediately rather + # than only via ``collect_reports`` after the queue drains. + if guard is not None: + await guard.pre_check(job_id, dispatch_depth, estimated_tokens) + + # Past the pre-check the reservation has no RAII owner yet: it passes + # to ``session.pending_tasks`` (released when the queue drains and the + # task runs) or to ``_dispatch_session_task`` (which releases on its + # own failures). Anything that raises in between must hand it back + # here, or it leaks for the life of the bus — permanently shrinking + # ``remaining_tokens`` and the concurrency count. + try: + return await self._queue_or_dispatch_session_task(session, pending) + except BaseException: + if guard is not None: + guard.release(job_id) + raise + + async def _queue_or_dispatch_session_task( + self, + session: SubAgentSession, + pending: PendingSessionTask, + ) -> str: + """Park ``pending`` behind the running task, or dispatch it now. + + Split out of :meth:`submit_task_to_session` so one wrapper there can + release the guard reservation on any failure along this path. + """ + job_id = pending.job_id + session_id = session.session_id + task_prompt = pending.task_prompt + task_metadata = pending.task_metadata + # Busy = a current task is in submitted/running state. ``queued`` # current_job_id can't happen (we only set current_job_id at # dispatch time below), but check defensively in case future code @@ -1211,7 +1255,28 @@ async def _dispatch_session_task( ``session_task_submitted`` SSE event. Called from the free branch of ``submit_task_to_session`` for fresh tasks and from ``_drain_session_queue`` when the previous task finalises. + + Releases the guard reservation on failure. Until ``_run_and_finalize`` + exists as an asyncio task, no ``finally`` owns that reservation — and + this method can raise well before then (``trimmer.trim`` on a + pathological history, the ``SubTask`` build, a cancelled dispatch). + ``_drain_session_queue`` swallows dispatch failures by design, so + without this the queued task's reservation would leak silently. """ + try: + await self._dispatch_session_task_inner(session, pending) + except BaseException: + guard = self._spawn_guard + if guard is not None: + guard.release(pending.job_id) + raise + + async def _dispatch_session_task_inner( + self, + session: SubAgentSession, + pending: PendingSessionTask, + ) -> None: + """Body of :meth:`_dispatch_session_task`; see its docstring.""" guard = self._spawn_guard job_id = pending.job_id task_prompt = pending.task_prompt @@ -1861,7 +1926,10 @@ async def _reconcile_from_messages( # Lazy import: ``agent_comm`` is a sibling module and this keeps the # bus importable without it, matching how ``runtime.py`` reaches for # the same class. - from agent_core.components.agent_bus.agent_comm import AgentComm + from agent_core.components.agent_bus.agent_comm import ( + AgentComm, + EventStoreContractError, + ) comm = registry.get_optional(AgentComm) if comm is None: @@ -1878,6 +1946,17 @@ async def _reconcile_from_messages( task_id=task_id, message_type=SUBTASK_RESULT_MESSAGE_TYPE, ) + except EventStoreContractError: + # Not a transient read failure: every future read fails the + # same way, so durable recovery is off until the store is + # fixed. Debug level would hide that permanently. + logger.error( + "durable result recovery is unavailable for %s/%s: the " + "configured event store does not stamp a total-order " + "ordinal on appended events", + task_id, agent_id, exc_info=True, + ) + continue except Exception: logger.debug( "message reconciliation failed for %s/%s", diff --git a/agent_core/components/finalization/budget.py b/agent_core/components/finalization/budget.py index 555cb11..f6866a2 100644 --- a/agent_core/components/finalization/budget.py +++ b/agent_core/components/finalization/budget.py @@ -129,63 +129,67 @@ def resolve_research_wall( ) -> ResearchWall: """Resolve the research deadline from profile budget and platform wall. - ``research_wall_time_s`` is already a research-only budget and is not - shortened. A valid ``RESEARCH_WALL_TIME`` replaces that profile budget. - The legacy ``wall_deadline_s`` and the operational task-wall env value are - total-task ceilings, so finalize grace is subtracted before comparing them — - and they are also what :attr:`ResearchWall.hard_total_s` reports. + Three inputs, resolved *independently* — a profile may legitimately carry + more than one, and each answers a different question: + + ``research_wall_time_s`` + already a research-only budget, so it is not shortened by the + finalize reserve and is **not** a hard ceiling (the reporter runs + outside it deliberately). A valid ``RESEARCH_WALL_TIME`` replaces it. + + ``wall_deadline_s`` / the task-wall env value + total-task ceilings. Finalize grace is subtracted before they compete + for ``research_deadline_s``, and they are what + :attr:`ResearchWall.hard_total_s` reports. + + Selection is by *value validity*, never by key presence: a profile that + sets ``research_wall_time_s`` to an empty or non-numeric value still has + its ``wall_deadline_s`` ceiling honoured, and a profile that sets both + keeps the ceiling enforced instead of letting the research budget hide it. """ - research_budget = "research_wall_time_s" in agent_cfg - if research_budget: - profile_raw = agent_cfg.get("research_wall_time_s") - elif "wall_deadline_s" in agent_cfg: - profile_raw = agent_cfg.get("wall_deadline_s") - else: - profile_raw = None - profile_s = positive_seconds( - profile_raw, - label=f"{label_prefix} research wall deadline", + research_s = positive_seconds( + agent_cfg.get("research_wall_time_s"), + label=f"{label_prefix} research_wall_time_s", + ) + profile_total_s = positive_seconds( + agent_cfg.get("wall_deadline_s"), + label=f"{label_prefix} wall_deadline_s", ) - research_env_raw = os.environ.get(RESEARCH_WALL_TIME_ENV) research_env_s = positive_seconds( - research_env_raw, + os.environ.get(RESEARCH_WALL_TIME_ENV), label=RESEARCH_WALL_TIME_ENV, ) if research_env_s is not None: # This env is an explicit phase-budget override, not another ceiling: # it may intentionally lengthen as well as shorten the YAML default. - profile_s = research_env_s - research_budget = True + research_s = research_env_s resolved_env = env_var if env_var == TASK_WALL_TIME_ENV and not os.environ.get(env_var): resolved_env = next( (name for name in LEGACY_TASK_WALL_TIME_ENVS if os.environ.get(name)), env_var, ) - env_s = positive_seconds(os.environ.get(resolved_env), label=resolved_env) + env_total_s = positive_seconds(os.environ.get(resolved_env), label=resolved_env) - profile_deadline_s = profile_s - if profile_deadline_s is not None and not research_budget: - profile_deadline_s = soft_wall_deadline_s(profile_deadline_s, reserve_s) - env_deadline_s = ( - soft_wall_deadline_s(env_s, reserve_s) if env_s is not None else None - ) - candidates = [ - value - for value in (profile_deadline_s, env_deadline_s) - if value is not None + # Only total-task values are hard ceilings; a research-only budget is not + # one, because the reporter is deliberately outside it. + total_candidates = [ + value for value in (profile_total_s, env_total_s) if value is not None ] - - # Only total-task values are hard ceilings; a research-only profile budget - # is not one, because the reporter is deliberately outside it. - hard_candidates = [ + candidates = [ value - for value in (None if research_budget else profile_s, env_s) + for value in ( + research_s, + *( + soft_wall_deadline_s(total_s, reserve_s) + for total_s in total_candidates + ), + ) if value is not None ] return ResearchWall( research_deadline_s=min(candidates) if candidates else 0.0, - hard_total_s=min(hard_candidates) if hard_candidates else 0.0, + hard_total_s=min(total_candidates) if total_candidates else 0.0, ) diff --git a/agent_core/models/event.py b/agent_core/models/event.py index 5cc9b20..b531773 100644 --- a/agent_core/models/event.py +++ b/agent_core/models/event.py @@ -31,6 +31,14 @@ class KernelEvent(BaseModel): model_config = ConfigDict(frozen=True) id: EventId = Field(default=EventId("")) + # Store-assigned total-order ordinal, stamped by ``append`` when the event + # is retained. ``None`` means "not persisted, or the store does not expose + # an ordinal" — durable cursors (``AgentComm.consume``) require it, so see + # ``agent_core.components.agent_bus.agent_comm.event_ordinal`` for how it + # is read and what the accepted fallback is. Deliberately separate from + # ``id``: ``EventId`` is an opaque string, and an opaque id cannot order a + # cursor. + seq: int | None = Field(default=None) task_id: TaskId session_id: SessionId | None = None prompt_id: PromptId | None = None diff --git a/docs/agent-bus-boundary.md b/docs/agent-bus-boundary.md index 5339cb7..4b84235 100644 --- a/docs/agent-bus-boundary.md +++ b/docs/agent-bus-boundary.md @@ -24,3 +24,46 @@ The two legacy wall-time environment prefixes remain accepted after the portable `AGENT_CORE_TASK_WALL_TIME_S` name so deployments can migrate without a flag day. Product workflow defaults are registered explicitly rather than discovered by importing a `workflows` package from the shared kernel. + +## What a host owes the event store + +Two distinct obligations sit behind one word ("event store"), and conflating +them is the easiest way to wire this package up wrong: + +`protocols.EventSink` — telemetry. +: Flat `append(task_id, event_type, payload, agent_role)`. This is what + `AgentBus(event_sink=...)` takes, and what the session lifecycle events + (`session_task_submitted` / `session_task_completed`) and the submit/abort + trace events go through. A host that implements only this is a complete, + valid telemetry sink. + +`agent_comm._AgentCommEventStore` — durable messaging. +: `append(event: KernelEvent) -> KernelEvent` plus `protocols.EventReader`. + `AgentComm` needs the whole event because an inter-agent message carries + `from_agent` / `to_agent` / `message_type` / `correlation_id` as + first-class fields that the reader queries on — the flat `EventSink` + signature cannot express them. An `EventSink` is therefore **not** an + AgentComm store, and passing one produces a garbage record rather than a + message. + +The messaging store additionally owes a **total-order ordinal** on every +retained event: `KernelEvent.seq`, a monotonically increasing integer, or a +decimal `KernelEvent.id` for stores whose integer primary key already serves as +the id. `EventReader`'s `after_id` cursor is expressed in that ordinal, and it — +not `timestamp` — defines the order. `EventId` is deliberately an opaque +string, so `types.new_event_id()` (uuid4 hex: neither numeric nor monotonic) +cannot order a cursor; a store that stamps only that fails +`AgentComm.consume` with `EventStoreContractError`, and `AgentBus` logs the +lost durable-recovery capability at ERROR while still degrading to the +in-memory path. `tests/test_agent_comm_store_contract.py` is the executable +form of this section. + +## Wall-time keys are independent, not alternatives + +`resolve_research_wall` reads `research_wall_time_s` (a research-only budget, +not a ceiling — the reporter runs outside it by design) and `wall_deadline_s` / +`AGENT_CORE_TASK_WALL_TIME_S` (total-task ceilings, reported as +`ResearchWall.hard_total_s`) independently, and selects on **value validity** +rather than key presence. A profile may carry both; the ceiling stays enforced +either way, and an empty or non-numeric `research_wall_time_s` does not discard +the ceiling beside it. diff --git a/tests/test_agent_bus_result_recovery_contract.py b/tests/test_agent_bus_result_recovery_contract.py new file mode 100644 index 0000000..539f320 --- /dev/null +++ b/tests/test_agent_bus_result_recovery_contract.py @@ -0,0 +1,96 @@ +"""Durable result recovery must not fail silently on a bad event store. + +``AgentBus._reconcile_from_messages`` is deliberately soft — no AgentComm, no +messages, or a transient read failure all fall back to the in-memory path. That +softness used to swallow a *permanent* fault too: a store with no total-order +ordinal raised ``ValueError`` out of ``int(event.id)``, the blanket +``except Exception`` logged it at debug level, and the whole L1-0 recovery path +was dead with nothing in the logs at default verbosity. + +So the two behaviours pinned here are complementary: recovery still works +end-to-end against a conforming store, and a non-conforming one is reported at +ERROR while still degrading to ``None`` rather than taking the run down. +""" + +from __future__ import annotations + +import logging + +import pytest + +from agent_core.components.agent_bus.agent_comm import AgentComm +from agent_core.components.agent_bus.bus import AgentBus +from agent_core.models.agent_message import AgentMessage +from agent_core.runtime.events import EventBus +from agent_core.runtime.registries import services as registry +from tests.test_agent_comm_store_contract import OpaqueIdStore, ReferenceStore + +RESULT = "subtask_result" + + +@pytest.fixture(autouse=True) +def _clear_registry(): + registry.clear() + yield + registry.clear() + + +def _result_msg(task_id: str, job_id: str) -> AgentMessage: + return AgentMessage( + task_id=task_id, + from_agent="worker", + to_agent="lead", + message_type=RESULT, + content={ + "correlation_id": job_id, + "session_id": "s1", + "result": { + "question": "q", "role_id": "researcher", + "success": True, "final_content": "done", + }, + }, + ) + + +async def _bus_with_store(store) -> AgentBus: + comm = AgentComm(store, EventBus()) + registry.register(AgentComm, comm) + await comm.send(_result_msg("t-rec", "t-rec.job.1")) + bus = AgentBus(event_sink=store) + bus._note_result_recipient("t-rec", {"parent_agent_id": "lead"}) + return bus + + +@pytest.mark.asyncio +async def test_conforming_store_recovers_the_result(): + bus = await _bus_with_store(ReferenceStore()) + + recovered = await bus._reconcile_from_messages("t-rec") + + assert recovered is not None + session_id, result = recovered + assert session_id == "s1" + assert result.final_content == "done" + assert result.success is True + + +@pytest.mark.asyncio +async def test_recovered_result_is_not_delivered_twice(): + bus = await _bus_with_store(ReferenceStore()) + + assert await bus._reconcile_from_messages("t-rec") is not None + assert await bus._reconcile_from_messages("t-rec") is None + + +@pytest.mark.asyncio +async def test_ordinal_less_store_is_reported_at_error_not_swallowed(caplog): + bus = await _bus_with_store(OpaqueIdStore()) + + with caplog.at_level(logging.ERROR, logger="agent_core.components.agent_bus.bus"): + recovered = await bus._reconcile_from_messages("t-rec") + + # Still soft: the run continues on the in-memory path. + assert recovered is None + errors = [r for r in caplog.records if r.levelno >= logging.ERROR] + assert errors, "a permanent store fault must not be logged at debug level" + assert "ordinal" in errors[0].getMessage() diff --git a/tests/test_agent_bus_spawn_reservation.py b/tests/test_agent_bus_spawn_reservation.py new file mode 100644 index 0000000..9fc5db9 --- /dev/null +++ b/tests/test_agent_bus_spawn_reservation.py @@ -0,0 +1,236 @@ +"""SpawnGuard reservations must not outlive a failed submit. + +``pre_check`` reserves tokens *and* registers the job in ``SpawnGuard._active`` +before anything owns the job. The RAII release lives in ``_run_and_finalize``, +which does not exist until ``asyncio.create_task`` has run — so every failure +between the pre-check and that point leaks the reservation permanently. + +A leak is not a transient glitch: ``remaining_tokens`` and ``active_count`` +never recover, so a bus that has rejected N bad submits behaves as if N +sub-agents were running forever, and eventually refuses legitimate work with +``BudgetExhausted`` or blocks on a semaphore nobody holds. + +Three entry points reserve: ``submit`` (unregistered role), and both session +paths — the free dispatch and the queued dispatch that ``_drain_session_queue`` +performs later (whose failures are swallowed by design, which is exactly why +the release has to happen inside ``_dispatch_session_task``). +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import MagicMock + +import pytest + +from agent_core.components.agent_bus import AgentBus, SubTask +from agent_core.components.agent_bus.spawn_guard import SpawnGuard +from agent_core.loop_types import AgentLoopResult +from agent_core.messages import assistant_msg, user_msg +from agent_core.models.agent_definition import AgentDefinition +from agent_core.models.task_budget import TaskBudget +from agent_core.runtime.registries import services as registry +from agent_core.runtime.registries.agents import AgentRegistry +from agent_core.runtime.resources.manager import ResourceManager + +BUDGET = TaskBudget(max_parallel=2, max_depth=4, max_tokens=30_000) + + +@pytest.fixture(autouse=True) +def _clear_registry(): + registry.clear() + yield + registry.clear() + + +def _registry_with(*role_ids: str) -> AgentRegistry: + reg = AgentRegistry() + for role_id in role_ids: + reg.register(AgentDefinition( + role_id=role_id, display_name=role_id, system_prompt="p", + allowed_tools=[], color="#000", icon="agent", + )) + registry.register(AgentRegistry, reg) + resource_mgr = MagicMock(spec=ResourceManager) + resource_mgr.get_llm.return_value = MagicMock() + resource_mgr.get_tools_for_role.return_value = [] + registry.register(ResourceManager, resource_mgr) + return reg + + +def _assert_guard_is_clean(guard: SpawnGuard) -> None: + assert guard.active_count == 0 + assert guard.tokens_reserved == 0 + assert guard.remaining_tokens == BUDGET.max_tokens + + +def _loop_factory(gate: asyncio.Event): + async def loop(**kwargs): + await gate.wait() + um = kwargs.get("user_message", "") + return AgentLoopResult( + messages=[*list(kwargs.get("initial_messages") or []), + user_msg(um), assistant_msg(f"r:{um}")], + final_content=f"r:{um}", + metadata={}, + ) + return loop + + +async def _session(bus: AgentBus, task_id: str = "t-res") -> str: + return await bus.create_session( + task_id=task_id, name="lit", role_id="researcher", + system_prompt="You are a researcher.", + llm_override=MagicMock(), tools_override=[], + ) + + +# ── submit() ────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_unregistered_role_releases_the_reservation(): + _registry_with("researcher") + bus = AgentBus() + guard = SpawnGuard(BUDGET) + bus.set_spawn_guard(guard) + + with pytest.raises(KeyError): + await bus.submit( + "t1", SubTask(question="q", role_id="nope"), estimated_tokens=10_000, + ) + + _assert_guard_is_clean(guard) + + +@pytest.mark.asyncio +async def test_repeated_failed_submits_do_not_decay_the_budget(): + """The shape that made this a real outage rather than a cosmetic leak.""" + _registry_with("researcher") + bus = AgentBus() + guard = SpawnGuard(BUDGET) + bus.set_spawn_guard(guard) + + for _ in range(3): + with pytest.raises(KeyError): + await bus.submit( + "t1", SubTask(question="q", role_id="nope"), + estimated_tokens=10_000, + ) + + _assert_guard_is_clean(guard) + # Budget still admits a job that needs the whole allowance. + await bus.submit( + "t1", SubTask(question="q", role_id="researcher"), + estimated_tokens=30_000, + ) + + +@pytest.mark.asyncio +async def test_successful_submit_keeps_the_reservation_until_the_job_ends( + monkeypatch, +): + """Guard against over-releasing: the fix must not free a live job's slot.""" + gate = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", _loop_factory(gate), + ) + _registry_with("researcher") + bus = AgentBus() + guard = SpawnGuard(BUDGET) + bus.set_spawn_guard(guard) + + await bus.submit( + "t1", SubTask(question="q", role_id="researcher"), estimated_tokens=10_000, + ) + await asyncio.sleep(0) + + assert guard.active_count == 1 + assert guard.tokens_reserved == 10_000 + + +# ── session paths ───────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_failed_free_dispatch_releases_the_reservation(): + _registry_with("researcher") + bus = AgentBus() + guard = SpawnGuard(BUDGET) + bus.set_spawn_guard(guard) + sid = await _session(bus) + + session = bus.get_session(sid) + assert session is not None + session.trimmer = MagicMock() + session.trimmer.trim.side_effect = RuntimeError("trim blew up") + + with pytest.raises(RuntimeError, match="trim blew up"): + await bus.submit_task_to_session(sid, "T1", estimated_tokens=10_000) + + _assert_guard_is_clean(guard) + + +@pytest.mark.asyncio +async def test_failed_queued_dispatch_releases_the_reservation(monkeypatch): + """``_drain_session_queue`` logs dispatch failures and moves on. + + So the queued task's reservation can only come back from inside + ``_dispatch_session_task`` itself — nothing upstream is still watching. + """ + gate = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", _loop_factory(gate), + ) + _registry_with("researcher") + bus = AgentBus() + guard = SpawnGuard(BUDGET) + bus.set_spawn_guard(guard) + sid = await _session(bus) + + await bus.submit_task_to_session(sid, "T1", estimated_tokens=10_000) + j2 = await bus.submit_task_to_session(sid, "T2", estimated_tokens=10_000) + + session = bus.get_session(sid) + assert session is not None + assert [p.job_id for p in session.pending_tasks] == [j2] + assert guard.tokens_reserved == 20_000 + + # T2 dispatches only once T1 finalises — break it in between. + session.trimmer = MagicMock() + session.trimmer.trim.side_effect = RuntimeError("trim blew up") + gate.set() + for _ in range(20): + await asyncio.sleep(0) + if not session.pending_tasks and guard.active_count == 0: + break + + _assert_guard_is_clean(guard) + + +@pytest.mark.asyncio +async def test_queued_task_keeps_its_reservation_while_it_waits(monkeypatch): + """Queued tasks reserve at submit time on purpose — don't undo that.""" + gate = asyncio.Event() + monkeypatch.setattr( + "agent_core.components.agent_bus.bus.run_agent_loop", _loop_factory(gate), + ) + _registry_with("researcher") + bus = AgentBus() + guard = SpawnGuard(BUDGET) + bus.set_spawn_guard(guard) + sid = await _session(bus) + + await bus.submit_task_to_session(sid, "T1", estimated_tokens=10_000) + await bus.submit_task_to_session(sid, "T2", estimated_tokens=10_000) + + assert guard.tokens_reserved == 20_000 + assert guard.active_count == 2 + + # Drain so the queued task is not garbage-collected mid-flight. + gate.set() + for _ in range(40): + await asyncio.sleep(0) + if guard.active_count == 0: + break + _assert_guard_is_clean(guard) diff --git a/tests/test_agent_comm_store_contract.py b/tests/test_agent_comm_store_contract.py new file mode 100644 index 0000000..e1504b1 --- /dev/null +++ b/tests/test_agent_comm_store_contract.py @@ -0,0 +1,371 @@ +"""The ``_AgentCommEventStore`` contract that ``AgentComm.consume`` rests on. + +``AgentComm`` is the durable half of inter-agent messaging: the cursor in +``consume`` is what makes result recovery exactly-once across a restart. That +cursor is an *integer* — ``EventReader.after_id`` — while ``KernelEvent.id`` is +an opaque ``EventId`` string. So the store owes AgentComm a total-order +ordinal, and this module pins that debt from both sides: + +* what a conforming store looks like, and the three ordering guarantees the + ``consume`` docstring promises callers (order, exactly-once under ``limit``, + per-``message_type`` cursor isolation); +* that a non-conforming store fails loudly and actionably instead of raising + ``ValueError`` from ``int()`` — which the bus used to swallow at debug level, + leaving durable recovery silently dead; +* that a plain ``EventSink`` is *not* an AgentComm store. ``EventSink.append`` + is flat ``(task_id, event_type, payload, agent_role)`` and cannot carry + ``from_agent`` / ``to_agent`` / ``message_type`` / ``correlation_id``, which + are exactly the fields the reader queries on. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from agent_core.components.agent_bus.agent_comm import ( + AgentComm, + DeliveryMode, + EventStoreContractError, + _AgentCommEventStore, + event_ordinal, +) +from agent_core.events import EventType +from agent_core.models.agent_message import AgentMessage +from agent_core.models.event import KernelEvent +from agent_core.protocols import EventSink +from agent_core.runtime.events import EventBus +from agent_core.types import EventId, TaskId, new_event_id + +RESULT = "subtask_result" +STATUS = "status_update" + + +# ────────────────────────────────────────────────────────────────────────── +# Reference implementations +# ────────────────────────────────────────────────────────────────────────── + + +class ReferenceStore: + """Minimal conforming store: stamps a monotonic ``seq`` on append.""" + + def __init__(self) -> None: + self.events: list[KernelEvent] = [] + self._next_seq = 1 + + async def append(self, event: KernelEvent) -> KernelEvent: + stamped = event.model_copy(update={ + "id": EventId(new_event_id()), + "seq": self._next_seq, + }) + self._next_seq += 1 + self.events.append(stamped) + return stamped + + def replay(self, task_id: str): # pragma: no cover — unused here + async def _empty(): + if False: + yield None + return _empty() + + async def get_events( + self, + task_id: Any, + event_type: Any = None, + after_id: int = 0, + limit: int | None = None, + ) -> list[KernelEvent]: + found = [ + e for e in self.events + if str(e.task_id) == str(task_id) and (e.seq or 0) > after_id + ] + return found[:limit] if limit is not None else found + + async def get_events_for_agent( + self, + to_agent: str, + after_id: int = 0, + limit: int = 50, + *, + task_id: str | Any | None = None, + message_type: str | None = None, + ) -> list[KernelEvent]: + found = [ + e for e in self.events + if e.to_agent == to_agent + and (e.seq or 0) > after_id + and (task_id is None or str(e.task_id) == str(task_id)) + and (message_type is None or e.message_type == message_type) + ] + return found[:limit] + + +class DecimalIdStore(ReferenceStore): + """Legacy shape: the integer primary key *is* the event id, no ``seq``.""" + + async def append(self, event: KernelEvent) -> KernelEvent: + stamped = event.model_copy(update={ + "id": EventId(str(self._next_seq)), "seq": None, + }) + self._next_seq += 1 + self.events.append(stamped) + return stamped + + async def get_events_for_agent( + self, to_agent, after_id=0, limit=50, *, task_id=None, message_type=None, + ) -> list[KernelEvent]: + found = [ + e for e in self.events + if e.to_agent == to_agent + and int(str(e.id)) > after_id + and (task_id is None or str(e.task_id) == str(task_id)) + and (message_type is None or e.message_type == message_type) + ] + return found[:limit] + + +class OpaqueIdStore(ReferenceStore): + """Non-conforming: keeps ``new_event_id()`` hex and stamps no ordinal.""" + + async def append(self, event: KernelEvent) -> KernelEvent: + stamped = event.model_copy(update={ + "id": EventId(new_event_id()), "seq": None, + }) + self.events.append(stamped) + return stamped + + async def get_events_for_agent( + self, to_agent, after_id=0, limit=50, *, task_id=None, message_type=None, + ) -> list[KernelEvent]: + return [ + e for e in self.events + if e.to_agent == to_agent + and (task_id is None or str(e.task_id) == str(task_id)) + and (message_type is None or e.message_type == message_type) + ] + + +class _FlatSink: + """A valid ``EventSink`` — and therefore *not* an AgentComm store.""" + + async def append( + self, + task_id: Any = "", + event_type: Any = None, + payload: dict[str, Any] | None = None, + agent_role: str = "system", + ) -> None: + return None + + def replay(self, task_id: str): # pragma: no cover + async def _empty(): + if False: + yield None + return _empty() + + +def _msg( + to: str, *, task: str = "t1", mtype: str = RESULT, job: str = "j1", +) -> AgentMessage: + return AgentMessage( + task_id=task, + from_agent="worker", + to_agent=to, + message_type=mtype, + content={"correlation_id": job, "result": {"final_content": job}}, + ) + + +# ────────────────────────────────────────────────────────────────────────── +# event_ordinal: what counts as a total order +# ────────────────────────────────────────────────────────────────────────── + + +def test_seq_is_the_ordinal_when_present(): + event = KernelEvent( + task_id=TaskId("t1"), event_type=EventType.AGENT_MESSAGE, + id=EventId("ae8dfb7969004eea"), seq=7, + ) + assert event_ordinal(event) == 7 + + +def test_decimal_id_is_accepted_as_a_fallback_ordinal(): + """Stores that expose their integer primary key as the id keep working.""" + event = KernelEvent( + task_id=TaskId("t1"), event_type=EventType.AGENT_MESSAGE, + id=EventId("42"), + ) + assert event_ordinal(event) == 42 + + +def test_seq_wins_over_a_decimal_id(): + event = KernelEvent( + task_id=TaskId("t1"), event_type=EventType.AGENT_MESSAGE, + id=EventId("42"), seq=7, + ) + assert event_ordinal(event) == 7 + + +@pytest.mark.parametrize("raw_id", ["", "ae8dfb7969004eea", "evt-1", "1.5", "-3"]) +def test_opaque_id_without_seq_is_a_contract_error(raw_id): + """Actionable error, not ``ValueError`` from ``int()``. + + ``new_event_id()`` hex is the trap: neither numeric nor monotonic, and it + is the default id generator in this package — so a host can build a + plausible-looking store that cannot order a cursor at all. + """ + event = KernelEvent( + task_id=TaskId("t1"), event_type=EventType.AGENT_MESSAGE, + id=EventId(raw_id), + ) + with pytest.raises(EventStoreContractError) as excinfo: + event_ordinal(event) + assert "KernelEvent.seq" in str(excinfo.value) + + +def test_default_event_id_generator_cannot_order_a_cursor(): + """Pins the reason ``seq`` exists at all.""" + event = KernelEvent( + task_id=TaskId("t1"), event_type=EventType.AGENT_MESSAGE, + id=new_event_id(), + ) + with pytest.raises(EventStoreContractError): + event_ordinal(event) + + +# ────────────────────────────────────────────────────────────────────────── +# The store protocol: which shape of ``append`` AgentComm requires +# ────────────────────────────────────────────────────────────────────────── + + +def test_reference_store_satisfies_the_agent_comm_contract(): + assert isinstance(ReferenceStore(), _AgentCommEventStore) + + +def test_flat_event_sink_is_not_an_agent_comm_store(): + """``EventSink`` is a telemetry sink, not a message store. + + It is structurally a valid ``EventSink`` and must stay one — the bus uses + that shape for session lifecycle events — but AgentComm needs a + whole-event ``append`` that can carry the routing fields. + """ + sink = _FlatSink() + assert isinstance(sink, EventSink) + assert not isinstance(sink, _AgentCommEventStore) + + +@pytest.mark.asyncio +async def test_send_round_trips_the_routing_fields(): + """The fields ``EventSink.append`` cannot express are the ones queried on.""" + store = ReferenceStore() + comm = AgentComm(store, EventBus()) + + persisted = await comm.send(_msg("critic"), DeliveryMode.QUEUE) + + assert persisted.to_agent == "critic" + assert persisted.from_agent == "worker" + assert persisted.message_type == RESULT + assert persisted.correlation_id == "j1" + assert event_ordinal(persisted) == 1 + + +# ────────────────────────────────────────────────────────────────────────── +# The three guarantees the consume() docstring makes to callers +# ────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +@pytest.mark.parametrize("store_cls", [ReferenceStore, DecimalIdStore]) +async def test_awaited_sends_come_back_in_send_order(store_cls): + comm = AgentComm(store_cls(), EventBus()) + for i in range(5): + await comm.send(_msg("critic", job=f"j{i}")) + + events = await comm.consume("critic", task_id="t1", message_type=RESULT) + + assert [e.payload["content"]["correlation_id"] for e in events] == [ + f"j{i}" for i in range(5) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("store_cls", [ReferenceStore, DecimalIdStore]) +async def test_limit_truncation_delivers_each_message_exactly_once(store_cls): + """The remainder arrives on the next call, still in order, never twice.""" + comm = AgentComm(store_cls(), EventBus()) + for i in range(5): + await comm.send(_msg("critic", job=f"j{i}")) + + seen: list[str] = [] + for _ in range(3): + batch = await comm.consume( + "critic", task_id="t1", limit=2, message_type=RESULT, + ) + seen += [e.payload["content"]["correlation_id"] for e in batch] + + assert seen == [f"j{i}" for i in range(5)] + assert await comm.consume( + "critic", task_id="t1", limit=2, message_type=RESULT, + ) == [] + + +@pytest.mark.asyncio +async def test_typed_cursors_do_not_carry_each_others_mail_past_unread(): + """A shared cursor would let the result reader consume status updates.""" + comm = AgentComm(ReferenceStore(), EventBus()) + await comm.send(_msg("critic", mtype=STATUS, job="s0")) + await comm.send(_msg("critic", mtype=RESULT, job="r0")) + await comm.send(_msg("critic", mtype=STATUS, job="s1")) + + results = await comm.consume("critic", task_id="t1", message_type=RESULT) + statuses = await comm.consume("critic", task_id="t1", message_type=STATUS) + + assert [e.payload["content"]["correlation_id"] for e in results] == ["r0"] + assert [e.payload["content"]["correlation_id"] for e in statuses] == ["s0", "s1"] + + +@pytest.mark.asyncio +async def test_reset_cursor_replays_the_selected_stream(): + comm = AgentComm(ReferenceStore(), EventBus()) + await comm.send(_msg("critic", job="r0")) + + first = await comm.consume("critic", task_id="t1", message_type=RESULT) + assert await comm.consume("critic", task_id="t1", message_type=RESULT) == [] + + comm.reset_cursor("critic", "t1", message_type=RESULT) + replayed = await comm.consume("critic", task_id="t1", message_type=RESULT) + + assert [e.id for e in replayed] == [e.id for e in first] + + +# ────────────────────────────────────────────────────────────────────────── +# Non-conforming stores: loud, and the cursor stays put +# ────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_consume_rejects_an_ordinal_less_store_actionably(): + """Previously a bare ``ValueError`` the bus swallowed at debug level.""" + comm = AgentComm(OpaqueIdStore(), EventBus()) + await comm.send(_msg("critic")) + + with pytest.raises(EventStoreContractError): + await comm.consume("critic", task_id="t1", message_type=RESULT) + + +@pytest.mark.asyncio +async def test_misordered_batch_leaves_the_cursor_untouched(): + """Advancing to max(ids) would permanently skip the lower ids.""" + store = ReferenceStore() + comm = AgentComm(store, EventBus()) + await comm.send(_msg("critic", job="r0")) + await comm.send(_msg("critic", job="r1")) + store.events.reverse() + + with pytest.raises(RuntimeError, match="ascending"): + await comm.consume("critic", task_id="t1", message_type=RESULT) + + store.events.reverse() + recovered = await comm.consume("critic", task_id="t1", message_type=RESULT) + assert [e.payload["content"]["correlation_id"] for e in recovered] == ["r0", "r1"] diff --git a/tests/test_finalization_budget.py b/tests/test_finalization_budget.py index 9ffc036..a5a19f9 100644 --- a/tests/test_finalization_budget.py +++ b/tests/test_finalization_budget.py @@ -58,6 +58,68 @@ def test_shortest_positive_ceiling_wins(monkeypatch): assert wall.hard_total_s == 3600 +def test_both_keys_keep_the_total_ceiling_enforced(monkeypatch): + """A research budget must not hide the total-task ceiling beside it. + + Selecting by key *presence* made ``wall_deadline_s`` unreadable whenever + ``research_wall_time_s`` existed, so ``hard_total_s`` came back 0 — + ``check_wall_feasibility`` then passed unconditionally and callers handed + ``deadline_monotonic_s=None`` to ``remaining_phase_budget_s``, leaving the + ceiling silently unenforced. + """ + monkeypatch.delenv("MIROHARNESS_TASK_WALL_TIME_S", raising=False) + monkeypatch.delenv("RESEARCH_WALL_TIME", raising=False) + + wall = resolve_research_wall( + {"research_wall_time_s": 600, "wall_deadline_s": 900}, + reserve_s=180, label_prefix="t", + ) + + assert wall.hard_total_s == 900 + # Research keeps its own (shorter) budget; the ceiling only caps it. + assert wall.research_deadline_s == 600 + + +def test_total_ceiling_can_shorten_a_longer_research_budget(monkeypatch): + monkeypatch.delenv("MIROHARNESS_TASK_WALL_TIME_S", raising=False) + monkeypatch.delenv("RESEARCH_WALL_TIME", raising=False) + + wall = resolve_research_wall( + {"research_wall_time_s": 9000, "wall_deadline_s": 900}, + reserve_s=180, label_prefix="t", + ) + + assert wall.hard_total_s == 900 + assert wall.research_deadline_s == soft_wall_deadline_s(900, 180) + + +@pytest.mark.parametrize("raw", [None, "", "0", "-1", "not-a-number"]) +def test_unusable_research_value_still_honours_wall_deadline_s(monkeypatch, raw): + """A YAML ``research_wall_time_s:`` with no value discarded both numbers.""" + monkeypatch.delenv("MIROHARNESS_TASK_WALL_TIME_S", raising=False) + monkeypatch.delenv("RESEARCH_WALL_TIME", raising=False) + + wall = resolve_research_wall( + {"research_wall_time_s": raw, "wall_deadline_s": 3600}, + reserve_s=600, label_prefix="t", + ) + + assert wall.hard_total_s == 3600 + assert wall.research_deadline_s == 3000 + + +def test_research_env_override_does_not_hide_the_profile_ceiling(monkeypatch): + monkeypatch.delenv("MIROHARNESS_TASK_WALL_TIME_S", raising=False) + monkeypatch.setenv("RESEARCH_WALL_TIME", "300") + + wall = resolve_research_wall( + {"wall_deadline_s": 900}, reserve_s=180, label_prefix="t", + ) + + assert wall.hard_total_s == 900 + assert wall.research_deadline_s == 300 + + def test_no_ceiling_anywhere(monkeypatch): monkeypatch.delenv("MIROHARNESS_TASK_WALL_TIME_S", raising=False) wall = resolve_research_wall({}, reserve_s=600, label_prefix="t")