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..afddf8e 100644 --- a/agent_core/components/agent_bus/__init__.py +++ b/agent_core/components/agent_bus/__init__.py @@ -1 +1,103 @@ -"""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, + EventStoreContractError, + event_ordinal, +) +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, + 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", + "EventStoreContractError", + "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", + "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", + "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..40411a8 --- /dev/null +++ b/agent_core/components/agent_bus/agent_comm.py @@ -0,0 +1,278 @@ +"""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 +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(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. + + 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/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) + # 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 = [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 + # 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..6768f9b --- /dev/null +++ b/agent_core/components/agent_bus/bus.py @@ -0,0 +1,2307 @@ +# 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] +_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( + 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 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, +) -> 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.""" + if _default_pause_check_factory is None: + return None + try: + return _default_pause_check_factory(task_id) + except (ImportError, RuntimeError): + 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, + session_activity: bool | 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 + 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. + + 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 + 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: + """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: + 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. + + 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 + + # 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: + await guard.pre_check( + job_id, dispatch_depth, estimated_tokens, + ) + + 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) + + # 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=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: + 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) + 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 {}), + ) + + # 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 + # 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. + + 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 + 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, + ) + + if self._session_activity: + 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(), + 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(), + ) + + 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, + EventStoreContractError, + ) + + 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 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", + 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(), + ) 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..c577024 --- /dev/null +++ b/agent_core/components/agent_bus/fan_in.py @@ -0,0 +1,436 @@ +# 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", + "thrash_no_progress", +}) + +_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" + ), + "thrash_no_progress": ( + "coordinator was spinning on sub-agent create/assign without new " + "information; stopped and synthesised from collected reports" + ), +} + + +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..f6866a2 --- /dev/null +++ b/agent_core/components/finalization/budget.py @@ -0,0 +1,267 @@ +"""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. + + 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_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_s = positive_seconds( + 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. + 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_total_s = positive_seconds(os.environ.get(resolved_env), label=resolved_env) + + # 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 + ] + candidates = [ + value + 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(total_candidates) if total_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..b531773 --- /dev/null +++ b/agent_core/models/event.py @@ -0,0 +1,238 @@ +"""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("")) + # 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 + 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..62cff1f 100644 --- a/agent_core/scheduling/__init__.py +++ b/agent_core/scheduling/__init__.py @@ -1 +1,27 @@ -"""Portable pipeline and topology registries.""" +"""Portable scheduling engine and pipeline/topology registries.""" + +from agent_core.scheduling.scheduler import ( + ProcessManager, + Scheduler, + SchedulerTask, + TaskWallTimeExceeded, + WallTimeResolver, + 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", + "WallTimeResolver", + "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..6feb62f --- /dev/null +++ b/agent_core/scheduling/scheduler.py @@ -0,0 +1,580 @@ +# 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, Callable +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: ... + + +WallTimeResolver = Callable[[int | None, str | None], int | None] + + +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, + wall_time_resolver: WallTimeResolver = resolve_wall_time_s, + ) -> 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._wall_time_resolver = wall_time_resolver + 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) + + 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, + 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 = self._wall_time_resolver(wall_time_s, 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_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. + + 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 + + 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/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..4b84235 --- /dev/null +++ b/docs/agent-bus-boundary.md @@ -0,0 +1,69 @@ +# 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. + +## 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_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_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_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_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_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_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 new file mode 100644 index 0000000..a5a19f9 --- /dev/null +++ b/tests/test_finalization_budget.py @@ -0,0 +1,308 @@ +"""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_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") + 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..56b5c3f --- /dev/null +++ b/tests/test_scheduler.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any + +import pytest + +from agent_core.components.agent_bus.fan_in import INCOMPLETE_STOP_REASONS +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) + + +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 + 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() + + +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"}, + ) + 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