From 4ba8bd4198ded295a29f54275ccf9387cffacfa6 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 7 Aug 2026 15:24:01 +0800 Subject: [PATCH 01/15] refactor(agent): introduce runtime contracts and step engine --- flocks/agent/runtime/__init__.py | 34 ++++++ flocks/agent/runtime/contracts.py | 167 ++++++++++++++++++++++++++ flocks/agent/runtime/events.py | 17 +++ flocks/agent/runtime/ports.py | 55 +++++++++ flocks/session/runner.py | 68 ++--------- flocks/session/session_loop.py | 29 +++-- flocks/session/step_engine.py | 35 ++++++ tests/agent/runtime/test_contracts.py | 45 +++++++ tests/session/test_step_engine.py | 53 ++++++++ 9 files changed, 435 insertions(+), 68 deletions(-) create mode 100644 flocks/agent/runtime/__init__.py create mode 100644 flocks/agent/runtime/contracts.py create mode 100644 flocks/agent/runtime/events.py create mode 100644 flocks/agent/runtime/ports.py create mode 100644 flocks/session/step_engine.py create mode 100644 tests/agent/runtime/test_contracts.py create mode 100644 tests/session/test_step_engine.py diff --git a/flocks/agent/runtime/__init__.py b/flocks/agent/runtime/__init__.py new file mode 100644 index 000000000..329cda955 --- /dev/null +++ b/flocks/agent/runtime/__init__.py @@ -0,0 +1,34 @@ +"""Reusable agent runtime contracts and control flow.""" + +from flocks.agent.runtime.contracts import ( + AgentRunOutcome, + AgentRunState, + AgentRunStatus, + AttemptEffects, + ContinuationDecision, + FailoverDecision, + ModelTurnSnapshot, + QueuedInputBatch, + RuntimeModel, + StepFailure, + StepResult, + ToolCall, +) +from flocks.agent.runtime.ports import RuntimeServices, StepEngine + +__all__ = [ + "AgentRunOutcome", + "AgentRunState", + "AgentRunStatus", + "AttemptEffects", + "ContinuationDecision", + "FailoverDecision", + "ModelTurnSnapshot", + "QueuedInputBatch", + "RuntimeModel", + "RuntimeServices", + "StepEngine", + "StepFailure", + "StepResult", + "ToolCall", +] diff --git a/flocks/agent/runtime/contracts.py b/flocks/agent/runtime/contracts.py new file mode 100644 index 000000000..895aa2289 --- /dev/null +++ b/flocks/agent/runtime/contracts.py @@ -0,0 +1,167 @@ +"""Data contracts shared by the agent core and session host. + +The contracts in this module intentionally avoid importing session storage, +server, CLI, provider, or tool-registry implementations. Session-specific +adapters may carry their native message objects through the generic message +type while the agent loop remains independent of those implementations. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from types import MappingProxyType +from typing import Any, Generic, Mapping, Optional, TypeVar + + +MessageT = TypeVar("MessageT") + + +@dataclass(frozen=True) +class RuntimeModel: + """Concrete provider/model selection for one model turn.""" + + provider_id: str + model_id: str + + +@dataclass +class AttemptEffects: + """Observable effects accumulated during one provider attempt.""" + + received_chunk: bool = False + observable_output_started: bool = False + tool_execution_started: bool = False + durable_side_effect_possible: bool = False + + @property + def replay_safe(self) -> bool: + """Return whether another provider may replay the logical request.""" + return not (self.observable_output_started or self.tool_execution_started or self.durable_side_effect_possible) + + +@dataclass(frozen=True) +class FailoverDecision: + """Classification used by the session host's recovery policy.""" + + eligible: bool + reason: str + + +@dataclass(frozen=True) +class ToolCall: + """Tool call emitted by a model response.""" + + id: str + name: str + arguments: dict[str, Any] + + +@dataclass +class StepFailure: + """Failure returned by a step when host finalization is deferred.""" + + message: str + error_data: dict[str, Any] + assistant_message_id: Optional[str] + reason: str + allow_fallback: bool + attempt_state: AttemptEffects + attempts: int = 0 + + +@dataclass +class StepResult: + """Result of one model turn, including any tool execution.""" + + action: str + content: str = "" + tool_calls: list[ToolCall] = field(default_factory=list) + error: Optional[str] = None + usage: Optional[dict[str, int]] = None + failure: Optional[StepFailure] = None + + +@dataclass +class AgentRunState(Generic[MessageT]): + """Agent-loop-owned mutable state for one resumable agent run.""" + + session_id: str + agent_name: str + active_model: RuntimeModel + messages: list[MessageT] = field(default_factory=list) + model_turn_index: int = 0 + trace_step_offset: int = 0 + consumed_input_cursor: Optional[str] = None + current_user_id: Optional[str] = None + metadata: dict[str, Any] = field(default_factory=dict) + + @property + def trace_step(self) -> int: + """Return the session-cumulative model-turn index.""" + return self.trace_step_offset + self.model_turn_index + + +@dataclass(frozen=True) +class ModelTurnSnapshot(Generic[MessageT]): + """Immutable input presented to a step engine for one model turn.""" + + session_id: str + agent_name: str + active_model: RuntimeModel + model_turn_index: int + trace_step: int + messages: tuple[MessageT, ...] + last_user: MessageT + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + """Defensively freeze caller-owned collections.""" + object.__setattr__(self, "messages", tuple(self.messages)) + object.__setattr__( + self, + "metadata", + MappingProxyType(dict(self.metadata)), + ) + + +@dataclass(frozen=True) +class QueuedInputBatch(Generic[MessageT]): + """New input made visible to the loop at a model-turn boundary.""" + + messages: tuple[MessageT, ...] = () + cursor: Optional[str] = None + + +@dataclass(frozen=True) +class ContinuationDecision(Generic[MessageT]): + """Host-owned continuation policy result consumed by the agent loop.""" + + messages: tuple[MessageT, ...] = () + reason: Optional[str] = None + + @property + def should_continue(self) -> bool: + """Return whether the loop should process another model turn.""" + return bool(self.messages) + + +class AgentRunStatus(str, Enum): + """Terminal states returned from the agent core to the session host.""" + + COMPLETED = "completed" + RETRYABLE_FAILURE = "retryable_failure" + CONTEXT_OVERFLOW = "context_overflow" + FATAL_FAILURE = "fatal_failure" + ABORTED = "aborted" + + +@dataclass(frozen=True) +class AgentRunOutcome(Generic[MessageT]): + """Structured terminal result for a resumable agent-loop invocation.""" + + status: AgentRunStatus + state: AgentRunState[MessageT] + last_message: Optional[MessageT] = None + error: Optional[str] = None + failure: Optional[StepFailure] = None diff --git a/flocks/agent/runtime/events.py b/flocks/agent/runtime/events.py new file mode 100644 index 000000000..efbd1501d --- /dev/null +++ b/flocks/agent/runtime/events.py @@ -0,0 +1,17 @@ +"""Runtime event contract emitted by the agent core.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Mapping, Optional + + +@dataclass(frozen=True) +class RuntimeEvent: + """Host-neutral event produced during an agent run.""" + + type: str + session_id: str + model_turn_index: int + payload: Mapping[str, Any] = field(default_factory=dict) + trace_step: Optional[int] = None diff --git a/flocks/agent/runtime/ports.py b/flocks/agent/runtime/ports.py new file mode 100644 index 000000000..c6c1f4ccb --- /dev/null +++ b/flocks/agent/runtime/ports.py @@ -0,0 +1,55 @@ +"""Ports implemented by session and infrastructure adapters.""" + +from __future__ import annotations + +from typing import Generic, Protocol, TypeVar + +from flocks.agent.runtime.contracts import ( + AgentRunState, + ContinuationDecision, + ModelTurnSnapshot, + QueuedInputBatch, + StepResult, +) +from flocks.agent.runtime.events import RuntimeEvent + + +MessageT = TypeVar("MessageT") + + +class StepEngine(Protocol, Generic[MessageT]): + """Execute one model turn from an immutable snapshot.""" + + async def run(self, snapshot: ModelTurnSnapshot[MessageT]) -> StepResult: + """Run one streamed model turn and return its result.""" + ... + + +class RuntimeServices(Protocol, Generic[MessageT]): + """Narrow session-host interface consumed by the agent loop.""" + + async def prepare_model_turn( + self, + state: AgentRunState[MessageT], + ) -> ModelTurnSnapshot[MessageT]: + """Prepare stable model, prompt, tool, and message inputs.""" + ... + + async def drain_queued_inputs( + self, + state: AgentRunState[MessageT], + ) -> QueuedInputBatch[MessageT]: + """Return inputs that arrived since the last consumed cursor.""" + ... + + async def resolve_continuation( + self, + state: AgentRunState[MessageT], + step_result: StepResult, + ) -> ContinuationDecision[MessageT]: + """Resolve goal and turn-finish-hook continuation policy.""" + ... + + async def emit_event(self, event: RuntimeEvent) -> None: + """Forward a runtime event to host-owned sinks.""" + ... diff --git a/flocks/session/runner.py b/flocks/session/runner.py index 85e4d333b..57c6fdfe5 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -18,11 +18,18 @@ import time from datetime import datetime from typing import Optional, Dict, Any, List, Callable, Awaitable, Tuple -from dataclasses import dataclass, field +from dataclasses import dataclass import httpcore import httpx +from flocks.agent.runtime.contracts import ( + AttemptEffects, + FailoverDecision, + StepFailure, + StepResult, + ToolCall, +) from flocks.utils.log import Log from flocks.utils.id import Identifier from flocks.session.session import Session, SessionInfo @@ -213,58 +220,7 @@ def _find_retryable_transport_exception(exception: Exception) -> Optional[Except return None -@dataclass -class ToolCall: - """Tool call from LLM response.""" - id: str - name: str - arguments: Dict[str, Any] - - -@dataclass -class LlmAttemptState: - """Observable side effects accumulated across retries for one model.""" - - received_chunk: bool = False - observable_output_started: bool = False - tool_execution_started: bool = False - - @property - def replay_safe(self) -> bool: - """Whether the same logical LLM call can safely run on another model.""" - return not self.observable_output_started and not self.tool_execution_started - - -@dataclass(frozen=True) -class FailoverDecision: - """Hermes-aligned retry/failover classification for a provider error.""" - - eligible: bool - reason: str - - -@dataclass -class StepFailure: - """Failure details returned to SessionLoop when finalization is deferred.""" - - message: str - error_data: Dict[str, Any] - assistant_message_id: Optional[str] - reason: str - allow_fallback: bool - attempt_state: LlmAttemptState - attempts: int = 0 - - -@dataclass -class StepResult: - """Result of a single processing step.""" - action: str # "stop", "continue", "compact" - content: str = "" - tool_calls: List[ToolCall] = field(default_factory=list) - error: Optional[str] = None - usage: Optional[Dict[str, int]] = None - failure: Optional[StepFailure] = None +LlmAttemptState = AttemptEffects @dataclass @@ -1386,7 +1342,7 @@ async def _process_step( "data": {"message": error}, }, assistant_message_id=None, - decision=FailoverDecision(True, "provider_unavailable", 0), + decision=FailoverDecision(True, "provider_unavailable"), attempts=0, ) error_dict = self._build_session_error_dict( @@ -1426,7 +1382,7 @@ async def _process_step( "data": {"message": error}, }, assistant_message_id=None, - decision=FailoverDecision(True, "provider_unavailable", 0), + decision=FailoverDecision(True, "provider_unavailable"), attempts=0, ) error_dict = self._build_session_error_dict( @@ -1742,7 +1698,7 @@ async def device_asset_prompt_factory() -> Optional[str]: message=empty_error_msg, error_data=empty_error_dict, assistant_message_id=assistant_msg.id, - decision=FailoverDecision(True, "empty_response", 3), + decision=FailoverDecision(True, "empty_response"), attempts=empty_attempt, ) if self.callbacks.on_error: diff --git a/flocks/session/session_loop.py b/flocks/session/session_loop.py index 7c098cd50..cbf8f9f11 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -20,6 +20,7 @@ from dataclasses import dataclass, field from datetime import datetime +from flocks.agent.runtime.contracts import ModelTurnSnapshot, RuntimeModel from flocks.utils.log import Log from flocks.utils.id import Identifier from flocks.session.session import ( @@ -56,14 +57,6 @@ CHAIN_EXHAUSTION_COOLDOWN_SECONDS = 5.0 -@dataclass(frozen=True) -class RuntimeModel: - """Concrete provider/model candidate used by Auto failover.""" - - provider_id: str - model_id: str - - @dataclass class AutoFailoverCooldown: """Process-local Hermes-style starting candidate cooldown.""" @@ -1297,6 +1290,7 @@ async def _process_step_with_failover( ) -> Any: """Run one logical step, moving across candidates without replaying output.""" from flocks.session.runner import RunnerCallbacks, SessionRunner + from flocks.session.step_engine import SessionStepEngine while True: runner_cbs = callbacks.runner_callbacks @@ -1323,10 +1317,21 @@ async def _process_step_with_failover( turn_additional_context=ctx.turn_additional_context, session_start_pending=ctx.session_start_pending, ) - runner._step = ctx.trace_step - - step_result = await runner._process_step(messages, last_user) - if runner._session_start_fired: + step_engine = SessionStepEngine(runner) + snapshot = ModelTurnSnapshot( + session_id=ctx.session.id, + agent_name=ctx.agent_name, + active_model=RuntimeModel( + provider_id=ctx.provider_id, + model_id=ctx.model_id, + ), + model_turn_index=ctx.step, + trace_step=ctx.trace_step, + messages=tuple(messages), + last_user=last_user, + ) + step_result = await step_engine.run(snapshot) + if step_engine.session_start_fired: ctx.session_start_pending = False failure = step_result.failure if not ctx.auto_failover or failure is None: diff --git a/flocks/session/step_engine.py b/flocks/session/step_engine.py new file mode 100644 index 000000000..c8d870184 --- /dev/null +++ b/flocks/session/step_engine.py @@ -0,0 +1,35 @@ +"""Session adapter for the agent runtime's step-engine port.""" + +from __future__ import annotations + +from flocks.agent.runtime.contracts import ModelTurnSnapshot, StepResult +from flocks.session.message import MessageInfo +from flocks.session.runner import SessionRunner + + +class SessionStepEngine: + """Run the existing session runner behind the StepEngine contract.""" + + def __init__(self, runner: SessionRunner): + self._runner = runner + + @property + def session_start_fired(self) -> bool: + """Return whether this engine fired the session-start hook.""" + return self._runner._session_start_fired + + @property + def attempt_effects(self): + """Return effects recorded by the latest provider attempt.""" + return self._runner._attempt_state + + async def run( + self, + snapshot: ModelTurnSnapshot[MessageInfo], + ) -> StepResult: + """Delegate one immutable model-turn snapshot to SessionRunner.""" + self._runner._step = snapshot.trace_step + return await self._runner._process_step( + list(snapshot.messages), + snapshot.last_user, + ) diff --git a/tests/agent/runtime/test_contracts.py b/tests/agent/runtime/test_contracts.py new file mode 100644 index 000000000..a12ed2a14 --- /dev/null +++ b/tests/agent/runtime/test_contracts.py @@ -0,0 +1,45 @@ +"""Tests for host-neutral agent runtime contracts.""" + +from flocks.agent.runtime.contracts import ( + AttemptEffects, + ModelTurnSnapshot, + RuntimeModel, +) + + +def test_attempt_effects_allow_replay_only_before_observable_effects() -> None: + effects = AttemptEffects(received_chunk=True) + + assert effects.replay_safe is True + + effects.observable_output_started = True + assert effects.replay_safe is False + + effects.observable_output_started = False + effects.tool_execution_started = True + assert effects.replay_safe is False + + effects.tool_execution_started = False + effects.durable_side_effect_possible = True + assert effects.replay_safe is False + + +def test_model_turn_snapshot_defensively_freezes_collections() -> None: + messages = ["user"] + metadata = {"tool_revision": 1} + snapshot = ModelTurnSnapshot( + session_id="session-1", + agent_name="rex", + active_model=RuntimeModel("provider", "model"), + model_turn_index=2, + trace_step=5, + messages=tuple(messages), + last_user="user", + metadata=metadata, + ) + + messages.append("new input") + metadata["tool_revision"] = 2 + + assert snapshot.messages == ("user",) + assert snapshot.metadata == {"tool_revision": 1} diff --git a/tests/session/test_step_engine.py b/tests/session/test_step_engine.py new file mode 100644 index 000000000..e940a54f8 --- /dev/null +++ b/tests/session/test_step_engine.py @@ -0,0 +1,53 @@ +"""Tests for the existing session runner's StepEngine adapter.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from flocks.agent.runtime.contracts import ( + AttemptEffects, + ModelTurnSnapshot, + RuntimeModel, + StepResult, +) +from flocks.session.runner import ( + LlmAttemptState, + StepResult as LegacyStepResult, +) +from flocks.session.step_engine import SessionStepEngine + + +@pytest.mark.asyncio +async def test_session_step_engine_delegates_immutable_snapshot() -> None: + last_user = SimpleNamespace(id="user-1") + expected = StepResult(action="stop", content="done") + runner = SimpleNamespace( + _process_step=AsyncMock(return_value=expected), + _session_start_fired=True, + _attempt_state=AttemptEffects(received_chunk=True), + _step=0, + ) + engine = SessionStepEngine(runner) + snapshot = ModelTurnSnapshot( + session_id="session-1", + agent_name="rex", + active_model=RuntimeModel("provider", "model"), + model_turn_index=1, + trace_step=8, + messages=(last_user,), + last_user=last_user, + ) + + result = await engine.run(snapshot) + + assert result is expected + assert runner._step == 8 + runner._process_step.assert_awaited_once_with([last_user], last_user) + assert engine.session_start_fired is True + assert engine.attempt_effects.received_chunk is True + + +def test_legacy_runner_contract_exports_remain_compatible() -> None: + assert LegacyStepResult is StepResult + assert LlmAttemptState is AttemptEffects From da0fc8e57ab5d5b97183c0947c13becf904a8214 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 7 Aug 2026 15:24:37 +0800 Subject: [PATCH 02/15] refactor(agent): add host-neutral agent loop --- flocks/agent/runtime/__init__.py | 25 +- flocks/agent/runtime/agent_loop.py | 173 +++++++++++ flocks/agent/runtime/contracts.py | 30 ++ flocks/agent/runtime/ports.py | 111 ++++++- tests/agent/runtime/test_agent_loop.py | 384 +++++++++++++++++++++++++ 5 files changed, 715 insertions(+), 8 deletions(-) create mode 100644 flocks/agent/runtime/agent_loop.py create mode 100644 tests/agent/runtime/test_agent_loop.py diff --git a/flocks/agent/runtime/__init__.py b/flocks/agent/runtime/__init__.py index 329cda955..d72b1b1b6 100644 --- a/flocks/agent/runtime/__init__.py +++ b/flocks/agent/runtime/__init__.py @@ -1,5 +1,6 @@ """Reusable agent runtime contracts and control flow.""" +from flocks.agent.runtime.agent_loop import AgentLoop from flocks.agent.runtime.contracts import ( AgentRunOutcome, AgentRunState, @@ -7,28 +8,50 @@ AttemptEffects, ContinuationDecision, FailoverDecision, + ModelTurnBoundary, + ModelTurnPreparation, ModelTurnSnapshot, QueuedInputBatch, RuntimeModel, StepFailure, StepResult, ToolCall, + TurnPreparationStatus, +) +from flocks.agent.runtime.ports import ( + ExternalRuntimePorts, + HookPort, + ModelPort, + PromptPort, + RuntimeEventSink, + RuntimeServices, + StepEngine, + ToolPort, ) -from flocks.agent.runtime.ports import RuntimeServices, StepEngine __all__ = [ "AgentRunOutcome", "AgentRunState", "AgentRunStatus", + "AgentLoop", "AttemptEffects", "ContinuationDecision", "FailoverDecision", + "ExternalRuntimePorts", + "HookPort", + "ModelPort", + "ModelTurnBoundary", + "ModelTurnPreparation", "ModelTurnSnapshot", + "PromptPort", "QueuedInputBatch", "RuntimeModel", + "RuntimeEventSink", "RuntimeServices", "StepEngine", "StepFailure", "StepResult", "ToolCall", + "ToolPort", + "TurnPreparationStatus", ] diff --git a/flocks/agent/runtime/agent_loop.py b/flocks/agent/runtime/agent_loop.py new file mode 100644 index 000000000..f55878b61 --- /dev/null +++ b/flocks/agent/runtime/agent_loop.py @@ -0,0 +1,173 @@ +"""Host-neutral model/tool/continuation control loop.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Generic, Optional, TypeVar + +from flocks.agent.runtime.contracts import ( + AgentRunOutcome, + AgentRunState, + AgentRunStatus, + TurnPreparationStatus, +) +from flocks.agent.runtime.ports import RuntimeServices, StepEngine + + +MessageT = TypeVar("MessageT") + + +class AgentLoop(Generic[MessageT]): + """Coordinate model turns without owning session policy or persistence.""" + + def __init__( + self, + step_engine: StepEngine[MessageT], + services: RuntimeServices[MessageT], + *, + abort_requested: Optional[Callable[[], bool]] = None, + ): + self._step_engine = step_engine + self._services = services + self._abort_requested = abort_requested or (lambda: False) + + async def run( + self, + state: AgentRunState[MessageT], + ) -> AgentRunOutcome[MessageT]: + """Run or resume an agent until it settles or needs host recovery.""" + last_message: Optional[MessageT] = None + + while not self._abort_requested(): + preparation = await self._services.prepare_model_turn(state) + if preparation.status == TurnPreparationStatus.CONTINUE: + continue + if preparation.status == TurnPreparationStatus.COMPLETE: + return AgentRunOutcome( + status=AgentRunStatus.COMPLETED, + state=state, + last_message=preparation.last_message or last_message, + ) + if preparation.status == TurnPreparationStatus.FATAL: + return AgentRunOutcome( + status=AgentRunStatus.FATAL_FAILURE, + state=state, + last_message=preparation.last_message or last_message, + error=preparation.error, + ) + + snapshot = preparation.snapshot + if snapshot is None: + return AgentRunOutcome( + status=AgentRunStatus.FATAL_FAILURE, + state=state, + last_message=last_message, + error="Runtime services returned READY without a model-turn snapshot", + ) + + state.active_model = snapshot.active_model + state.model_turn_index = snapshot.model_turn_index + state.messages = list(snapshot.messages) + + step_result = await self._step_engine.run(snapshot) + boundary = await self._services.complete_model_turn(state, step_result) + state.messages = list(boundary.messages) + last_message = boundary.last_message or last_message + if boundary.queued_inputs.cursor is not None: + state.consumed_input_cursor = boundary.queued_inputs.cursor + + if self._abort_requested(): + return AgentRunOutcome( + status=AgentRunStatus.ABORTED, + state=state, + last_message=last_message, + error=step_result.error, + ) + + if boundary.queued_inputs.messages: + self._append_new_messages( + state, + boundary.queued_inputs.messages, + ) + continue + + failure = step_result.failure + if failure is not None: + status = ( + AgentRunStatus.RETRYABLE_FAILURE + if failure.allow_fallback and failure.attempt_state.replay_safe + else AgentRunStatus.FATAL_FAILURE + ) + return AgentRunOutcome( + status=status, + state=state, + last_message=last_message, + error=failure.message, + failure=failure, + ) + + if step_result.action == "continue": + continue + if step_result.action == "compact": + return AgentRunOutcome( + status=AgentRunStatus.CONTEXT_OVERFLOW, + state=state, + last_message=last_message, + error=step_result.error, + ) + if step_result.action != "stop": + return AgentRunOutcome( + status=AgentRunStatus.FATAL_FAILURE, + state=state, + last_message=last_message, + error=f"Unknown step action: {step_result.action}", + ) + if step_result.error: + return AgentRunOutcome( + status=AgentRunStatus.FATAL_FAILURE, + state=state, + last_message=last_message, + error=step_result.error, + ) + + continuation = await self._services.resolve_continuation( + state, + step_result, + ) + if continuation.should_continue: + self._append_new_messages(state, continuation.messages) + continue + + return AgentRunOutcome( + status=AgentRunStatus.COMPLETED, + state=state, + last_message=last_message, + ) + + return AgentRunOutcome( + status=AgentRunStatus.ABORTED, + state=state, + last_message=last_message, + error="Aborted", + ) + + @staticmethod + def _append_new_messages( + state: AgentRunState[MessageT], + messages: tuple[MessageT, ...], + ) -> None: + """Append messages not already present in the current runtime view.""" + existing_ids = {AgentLoop._message_identity(message) for message in state.messages} + for message in messages: + identity = AgentLoop._message_identity(message) + if identity not in existing_ids: + state.messages.append(message) + existing_ids.add(identity) + + @staticmethod + def _message_identity(message: MessageT) -> tuple[str, object]: + """Return a stable identity for persisted or in-memory messages.""" + message_id = getattr(message, "id", None) + if message_id is not None: + return ("message_id", message_id) + return ("object_id", id(message)) diff --git a/flocks/agent/runtime/contracts.py b/flocks/agent/runtime/contracts.py index 895aa2289..2359515da 100644 --- a/flocks/agent/runtime/contracts.py +++ b/flocks/agent/runtime/contracts.py @@ -125,6 +125,25 @@ def __post_init__(self) -> None: ) +class TurnPreparationStatus(str, Enum): + """Host preparation result before the next model turn.""" + + READY = "ready" + CONTINUE = "continue" + COMPLETE = "complete" + FATAL = "fatal" + + +@dataclass(frozen=True) +class ModelTurnPreparation(Generic[MessageT]): + """Result of host-owned preparation at a model-turn boundary.""" + + status: TurnPreparationStatus + snapshot: Optional[ModelTurnSnapshot[MessageT]] = None + last_message: Optional[MessageT] = None + error: Optional[str] = None + + @dataclass(frozen=True) class QueuedInputBatch(Generic[MessageT]): """New input made visible to the loop at a model-turn boundary.""" @@ -133,6 +152,17 @@ class QueuedInputBatch(Generic[MessageT]): cursor: Optional[str] = None +@dataclass(frozen=True) +class ModelTurnBoundary(Generic[MessageT]): + """Committed session view after one model turn finishes.""" + + messages: tuple[MessageT, ...] + last_message: Optional[MessageT] = None + queued_inputs: QueuedInputBatch[MessageT] = field( + default_factory=QueuedInputBatch, + ) + + @dataclass(frozen=True) class ContinuationDecision(Generic[MessageT]): """Host-owned continuation policy result consumed by the agent loop.""" diff --git a/flocks/agent/runtime/ports.py b/flocks/agent/runtime/ports.py index c6c1f4ccb..9b476e9f4 100644 --- a/flocks/agent/runtime/ports.py +++ b/flocks/agent/runtime/ports.py @@ -2,13 +2,15 @@ from __future__ import annotations -from typing import Generic, Protocol, TypeVar +from dataclasses import dataclass +from typing import Any, Generic, Optional, Protocol, TypeVar from flocks.agent.runtime.contracts import ( AgentRunState, ContinuationDecision, + ModelTurnBoundary, + ModelTurnPreparation, ModelTurnSnapshot, - QueuedInputBatch, StepResult, ) from flocks.agent.runtime.events import RuntimeEvent @@ -31,15 +33,16 @@ class RuntimeServices(Protocol, Generic[MessageT]): async def prepare_model_turn( self, state: AgentRunState[MessageT], - ) -> ModelTurnSnapshot[MessageT]: - """Prepare stable model, prompt, tool, and message inputs.""" + ) -> ModelTurnPreparation[MessageT]: + """Prepare or defer the next stable model-turn input.""" ... - async def drain_queued_inputs( + async def complete_model_turn( self, state: AgentRunState[MessageT], - ) -> QueuedInputBatch[MessageT]: - """Return inputs that arrived since the last consumed cursor.""" + step_result: StepResult, + ) -> ModelTurnBoundary[MessageT]: + """Return the committed post-step view and queued inputs.""" ... async def resolve_continuation( @@ -53,3 +56,97 @@ async def resolve_continuation( async def emit_event(self, event: RuntimeEvent) -> None: """Forward a runtime event to host-owned sinks.""" ... + + +class PromptPort(Protocol): + """Build provider-ready system prompt sections.""" + + async def build_system_prompts(self, **kwargs: Any) -> list[str]: + """Build prompts from one stable model-turn configuration.""" + ... + + +class ToolPort(Protocol): + """Expose the tool registry without coupling the core to its storage.""" + + def revision(self) -> int: + """Return a cache revision for the visible tool set.""" + ... + + def list_tools(self) -> list[Any]: + """Return registered tool metadata entries.""" + ... + + def get(self, name: str) -> Optional[Any]: + """Resolve one executable tool by name.""" + ... + + +class ModelPort(Protocol): + """Resolve and configure provider/model adapters.""" + + def get_provider(self, provider_id: str) -> Optional[Any]: + """Return one configured provider adapter.""" + ... + + async def apply_config(self, provider_id: str) -> None: + """Apply persisted provider configuration before execution.""" + ... + + def resolve_model(self, provider_id: str, model_id: str) -> Optional[Any]: + """Return model capability metadata.""" + ... + + def resolve_model_info( + self, + provider_id: str, + model_id: str, + ) -> tuple[int, int, Optional[int]]: + """Return context, output, and input token limits.""" + ... + + +class HookPort(Protocol): + """Run existing Flocks hook stages at explicit runtime boundaries.""" + + async def run_session_start(self, data: dict[str, Any]) -> Any: + """Run SessionStart hooks.""" + ... + + async def has_stage_handlers( + self, + stage: Any, + metadata: dict[str, Any], + ) -> bool: + """Return whether a hook stage has eligible handlers.""" + ... + + async def run_llm_before(self, data: dict[str, Any]) -> Any: + """Run LLMBefore hooks.""" + ... + + async def run_llm_after( + self, + metadata: dict[str, Any], + result: dict[str, Any], + ) -> Any: + """Run LLMAfter hooks.""" + ... + + +class RuntimeEventSink(Protocol): + """Receive observable runtime events; not an event-sourcing store.""" + + async def emit(self, event: RuntimeEvent) -> None: + """Forward an event to UI, tracing, or audit subscribers.""" + ... + + +@dataclass(frozen=True) +class ExternalRuntimePorts: + """External interfaces captured once for a stable model attempt.""" + + prompts: PromptPort + tools: ToolPort + models: ModelPort + hooks: HookPort diff --git a/tests/agent/runtime/test_agent_loop.py b/tests/agent/runtime/test_agent_loop.py new file mode 100644 index 000000000..690f5b923 --- /dev/null +++ b/tests/agent/runtime/test_agent_loop.py @@ -0,0 +1,384 @@ +"""Tests for the host-neutral agent control loop.""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Callable +from dataclasses import dataclass + +import pytest + +from flocks.agent.runtime import ( + AgentLoop, + AgentRunState, + AgentRunStatus, + AttemptEffects, + ContinuationDecision, + ModelTurnBoundary, + ModelTurnPreparation, + ModelTurnSnapshot, + QueuedInputBatch, + RuntimeModel, + StepFailure, + StepResult, + TurnPreparationStatus, +) + + +@dataclass(frozen=True) +class Message: + """Minimal persisted-message stand-in with stable identity.""" + + id: str + content: str + + +PreparationFactory = Callable[ + [AgentRunState[Message]], + ModelTurnPreparation[Message], +] + + +class FakeStepEngine: + """Return deterministic step results and record immutable inputs.""" + + def __init__(self, results: list[StepResult]): + self._results = deque(results) + self.snapshots: list[ModelTurnSnapshot[Message]] = [] + + async def run(self, snapshot: ModelTurnSnapshot[Message]) -> StepResult: + self.snapshots.append(snapshot) + return self._results.popleft() + + +class FakeRuntimeServices: + """Expose scripted host-boundary decisions to the agent loop.""" + + def __init__( + self, + preparations: list[ModelTurnPreparation[Message] | PreparationFactory], + boundaries: list[ModelTurnBoundary[Message]], + continuations: list[ContinuationDecision[Message]] | None = None, + ): + self._preparations = deque(preparations) + self._boundaries = deque(boundaries) + self._continuations = deque(continuations or []) + self.prepared_messages: list[tuple[Message, ...]] = [] + self.events = [] + + async def prepare_model_turn( + self, + state: AgentRunState[Message], + ) -> ModelTurnPreparation[Message]: + self.prepared_messages.append(tuple(state.messages)) + preparation = self._preparations.popleft() + if callable(preparation): + return preparation(state) + return preparation + + async def complete_model_turn( + self, + state: AgentRunState[Message], + step_result: StepResult, + ) -> ModelTurnBoundary[Message]: + del state, step_result + return self._boundaries.popleft() + + async def resolve_continuation( + self, + state: AgentRunState[Message], + step_result: StepResult, + ) -> ContinuationDecision[Message]: + del state, step_result + if not self._continuations: + return ContinuationDecision() + return self._continuations.popleft() + + async def emit_event(self, event) -> None: + self.events.append(event) + + +def _state(messages: list[Message] | None = None) -> AgentRunState[Message]: + return AgentRunState( + session_id="session-1", + agent_name="rex", + active_model=RuntimeModel("provider-a", "model-a"), + messages=list(messages or [Message("user-1", "hello")]), + ) + + +def _ready( + state: AgentRunState[Message], + *, + turn: int = 0, +) -> ModelTurnPreparation[Message]: + messages = tuple(state.messages) + return ModelTurnPreparation( + status=TurnPreparationStatus.READY, + snapshot=ModelTurnSnapshot( + session_id=state.session_id, + agent_name=state.agent_name, + active_model=state.active_model, + model_turn_index=turn, + trace_step=turn, + messages=messages, + last_user=messages[-1], + ), + ) + + +@pytest.mark.asyncio +async def test_loop_honors_deferred_preparation_then_completes() -> None: + user = Message("user-1", "hello") + assistant = Message("assistant-1", "done") + state = _state([user]) + engine = FakeStepEngine([StepResult(action="stop")]) + services = FakeRuntimeServices( + preparations=[ + ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE), + _ready, + ], + boundaries=[ModelTurnBoundary(messages=(user, assistant), last_message=assistant)], + ) + + outcome = await AgentLoop(engine, services).run(state) + + assert outcome.status == AgentRunStatus.COMPLETED + assert outcome.last_message == assistant + assert len(engine.snapshots) == 1 + assert len(services.prepared_messages) == 2 + + +@pytest.mark.asyncio +async def test_loop_runs_another_turn_after_tool_continue() -> None: + user = Message("user-1", "hello") + tool_result = Message("tool-1", "tool result") + assistant = Message("assistant-1", "done") + state = _state([user]) + engine = FakeStepEngine( + [StepResult(action="continue"), StepResult(action="stop")], + ) + services = FakeRuntimeServices( + preparations=[_ready, lambda current: _ready(current, turn=1)], + boundaries=[ + ModelTurnBoundary(messages=(user, tool_result), last_message=tool_result), + ModelTurnBoundary( + messages=(user, tool_result, assistant), + last_message=assistant, + ), + ], + ) + + outcome = await AgentLoop(engine, services).run(state) + + assert outcome.status == AgentRunStatus.COMPLETED + assert len(engine.snapshots) == 2 + assert engine.snapshots[1].messages == (user, tool_result) + + +@pytest.mark.asyncio +async def test_queued_input_precedes_natural_stop() -> None: + user = Message("user-1", "hello") + assistant = Message("assistant-1", "first answer") + duplicate_assistant = Message("assistant-1", "reloaded answer") + queued_user = Message("user-2", "follow up") + final = Message("assistant-2", "second answer") + state = _state([user]) + engine = FakeStepEngine( + [StepResult(action="stop"), StepResult(action="stop")], + ) + services = FakeRuntimeServices( + preparations=[_ready, lambda current: _ready(current, turn=1)], + boundaries=[ + ModelTurnBoundary( + messages=(user, assistant), + last_message=assistant, + queued_inputs=QueuedInputBatch( + messages=(duplicate_assistant, queued_user), + cursor="input-2", + ), + ), + ModelTurnBoundary( + messages=(user, assistant, queued_user, final), + last_message=final, + ), + ], + ) + + outcome = await AgentLoop(engine, services).run(state) + + assert outcome.status == AgentRunStatus.COMPLETED + assert state.consumed_input_cursor == "input-2" + assert engine.snapshots[1].messages == (user, assistant, queued_user) + + +@pytest.mark.asyncio +async def test_queued_input_is_not_lost_after_final_step_failure() -> None: + user = Message("user-1", "hello") + failed = Message("assistant-1", "provider failed") + queued_user = Message("user-2", "try this instead") + final = Message("assistant-2", "done") + state = _state([user]) + failure = StepFailure( + message="provider failed", + error_data={}, + assistant_message_id=failed.id, + reason="provider_error", + allow_fallback=False, + attempt_state=AttemptEffects(observable_output_started=True), + ) + engine = FakeStepEngine( + [ + StepResult(action="stop", error=failure.message, failure=failure), + StepResult(action="stop"), + ] + ) + services = FakeRuntimeServices( + preparations=[_ready, lambda current: _ready(current, turn=1)], + boundaries=[ + ModelTurnBoundary( + messages=(user, failed, queued_user), + last_message=failed, + queued_inputs=QueuedInputBatch( + messages=(queued_user,), + cursor=queued_user.id, + ), + ), + ModelTurnBoundary( + messages=(user, failed, queued_user, final), + last_message=final, + ), + ], + ) + + outcome = await AgentLoop(engine, services).run(state) + + assert outcome.status == AgentRunStatus.COMPLETED + assert len(engine.snapshots) == 2 + assert engine.snapshots[1].messages == (user, failed, queued_user) + + +@pytest.mark.asyncio +async def test_host_continuation_starts_another_turn() -> None: + user = Message("user-1", "hello") + assistant = Message("assistant-1", "working") + continuation = Message("user-2", "continue goal") + final = Message("assistant-2", "done") + state = _state([user]) + engine = FakeStepEngine( + [StepResult(action="stop"), StepResult(action="stop")], + ) + services = FakeRuntimeServices( + preparations=[_ready, lambda current: _ready(current, turn=1)], + boundaries=[ + ModelTurnBoundary(messages=(user, assistant), last_message=assistant), + ModelTurnBoundary( + messages=(user, assistant, continuation, final), + last_message=final, + ), + ], + continuations=[ + ContinuationDecision(messages=(continuation,), reason="goal"), + ContinuationDecision(), + ], + ) + + outcome = await AgentLoop(engine, services).run(state) + + assert outcome.status == AgentRunStatus.COMPLETED + assert engine.snapshots[1].messages == (user, assistant, continuation) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("effects", "expected_status"), + [ + (AttemptEffects(received_chunk=True), AgentRunStatus.RETRYABLE_FAILURE), + ( + AttemptEffects(tool_execution_started=True), + AgentRunStatus.FATAL_FAILURE, + ), + ], +) +async def test_failure_is_retryable_only_before_observable_effects( + effects: AttemptEffects, + expected_status: AgentRunStatus, +) -> None: + user = Message("user-1", "hello") + state = _state([user]) + failure = StepFailure( + message="provider failed", + error_data={}, + assistant_message_id=None, + reason="provider_error", + allow_fallback=True, + attempt_state=effects, + ) + engine = FakeStepEngine( + [StepResult(action="stop", error=failure.message, failure=failure)], + ) + services = FakeRuntimeServices( + preparations=[_ready], + boundaries=[ModelTurnBoundary(messages=(user,))], + ) + + outcome = await AgentLoop(engine, services).run(state) + + assert outcome.status == expected_status + assert outcome.failure is failure + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("result", "expected_status"), + [ + (StepResult(action="compact", error="context overflow"), AgentRunStatus.CONTEXT_OVERFLOW), + (StepResult(action="unexpected"), AgentRunStatus.FATAL_FAILURE), + ], +) +async def test_loop_returns_structured_non_success_outcomes( + result: StepResult, + expected_status: AgentRunStatus, +) -> None: + user = Message("user-1", "hello") + engine = FakeStepEngine([result]) + services = FakeRuntimeServices( + preparations=[_ready], + boundaries=[ModelTurnBoundary(messages=(user,))], + ) + + outcome = await AgentLoop(engine, services).run(_state([user])) + + assert outcome.status == expected_status + + +@pytest.mark.asyncio +async def test_loop_aborts_after_current_step_boundary() -> None: + user = Message("user-1", "hello") + checks = iter([False, True]) + engine = FakeStepEngine([StepResult(action="stop")]) + services = FakeRuntimeServices( + preparations=[_ready], + boundaries=[ModelTurnBoundary(messages=(user,))], + ) + + outcome = await AgentLoop( + engine, + services, + abort_requested=lambda: next(checks), + ).run(_state([user])) + + assert outcome.status == AgentRunStatus.ABORTED + + +@pytest.mark.asyncio +async def test_ready_preparation_requires_snapshot() -> None: + services = FakeRuntimeServices( + preparations=[ModelTurnPreparation(status=TurnPreparationStatus.READY)], + boundaries=[], + ) + + outcome = await AgentLoop(FakeStepEngine([]), services).run(_state()) + + assert outcome.status == AgentRunStatus.FATAL_FAILURE + assert outcome.error == "Runtime services returned READY without a model-turn snapshot" From 79b9f1a03c72d564bffaf8bf333e2dbb584bfb11 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 7 Aug 2026 15:26:53 +0800 Subject: [PATCH 03/15] refactor(runtime): add external service ports --- flocks/session/runner.py | 77 ++++++++++++++------ flocks/session/runtime_adapters.py | 85 ++++++++++++++++++++++ tests/session/test_runtime_ports.py | 108 ++++++++++++++++++++++++++++ 3 files changed, 248 insertions(+), 22 deletions(-) create mode 100644 flocks/session/runtime_adapters.py create mode 100644 tests/session/test_runtime_ports.py diff --git a/flocks/session/runner.py b/flocks/session/runner.py index 57c6fdfe5..a71c10a94 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -30,6 +30,7 @@ StepResult, ToolCall, ) +from flocks.agent.runtime.ports import ExternalRuntimePorts from flocks.utils.log import Log from flocks.utils.id import Identifier from flocks.session.session import Session, SessionInfo @@ -270,6 +271,7 @@ def __init__( failover_available: bool = False, turn_additional_context: Optional[str] = None, session_start_pending: bool = False, + runtime_ports: Optional[ExternalRuntimePorts] = None, ): self.session = session from flocks.session.core.defaults import fallback_provider_id, fallback_model_id @@ -290,6 +292,26 @@ def __init__( self._session_start_pending = session_start_pending self._session_start_fired = False self._attempt_state = LlmAttemptState() + if runtime_ports is None: + from flocks.session.runtime_adapters import ( + create_default_runtime_ports, + ) + + runtime_ports = create_default_runtime_ports() + self._runtime_ports = runtime_ports + + @property + def _ports(self) -> ExternalRuntimePorts: + """Return injected ports, lazily supporting legacy test instances.""" + ports = getattr(self, "_runtime_ports", None) + if ports is None: + from flocks.session.runtime_adapters import ( + create_default_runtime_ports, + ) + + ports = create_default_runtime_ports() + self._runtime_ports = ports + return ports @staticmethod def _canonical_tool_signature(tool_name: str, arguments: Dict[str, Any]) -> str: @@ -333,9 +355,7 @@ async def _run_session_start_hook(self, agent: Any) -> None: return self._session_start_fired = True try: - from flocks.hooks.pipeline import HookPipeline - - await HookPipeline.run_session_start({ + await self._ports.hooks.run_session_start({ "sessionID": self.session.id, "workspace": self.session.directory, "agent": agent.name, @@ -553,7 +573,7 @@ async def _list_callable_tool_infos_for_turn( execution_mode == SessionExecutionMode.PLAN and all(tool_info.name != "plan_exit" for tool_info in tool_infos) ): - plan_exit = ToolRegistry.get("plan_exit") + plan_exit = self._ports.tools.get("plan_exit") if plan_exit is not None and getattr(plan_exit.info, "enabled", True): tool_infos.append(plan_exit.info) metadata = dict(result.metadata) @@ -625,7 +645,10 @@ def _log_perf(self, event: str, started_at: float, **extra: Any) -> None: def _provider_capability_key(self) -> str: interleaved = None try: - active_model = Provider.resolve_model(self.provider_id, self.model_id) + active_model = self._ports.models.resolve_model( + self.provider_id, + self.model_id, + ) if active_model and getattr(active_model, "capabilities", None): interleaved = getattr(active_model.capabilities, "interleaved", None) except Exception: @@ -644,7 +667,7 @@ def _tool_schema_cache_key( text_tool_call_mode: bool, ) -> Tuple[Any, ...]: return ( - ToolRegistry.revision(), + self._ports.tools.revision(), getattr(agent, "name", ""), tuple(sorted(getattr(agent, "tools", None) or ())), tuple(tool_info.name for tool_info in selected_tool_infos), @@ -804,9 +827,7 @@ def _model_supports_vision(self) -> bool: unknown configurations. """ try: - from flocks.provider.provider import Provider as _Provider - - provider = _Provider.get(self.provider_id) + provider = self._ports.models.get_provider(self.provider_id) if provider is not None: for model in getattr(provider, "_config_models", []) or []: if model.id == self.model_id: @@ -1330,7 +1351,7 @@ async def _process_step( is_last_step = self._step >= max_steps # Get provider - provider = Provider.get(self.provider_id) + provider = self._ports.models.get_provider(self.provider_id) if not provider: error = f"Provider {self.provider_id} not found" if self._defer_step_errors: @@ -1364,7 +1385,7 @@ async def _process_step( # Apply config-based provider options (api_key/base_url) try: - await Provider.apply_config(provider_id=self.provider_id) + await self._ports.models.apply_config(self.provider_id) except Exception as e: log.debug("runner.provider.apply_config.error", { "provider": self.provider_id, @@ -1425,7 +1446,7 @@ async def device_asset_prompt_factory() -> Optional[str]: current_device_revision = None prompts_started_at = time.perf_counter() - system_prompts = await SessionPrompt.build_system_prompts( + system_prompts = await self._ports.prompts.build_system_prompts( session_id=self.session.id, session_directory=self.session.directory, agent_name=agent.name, @@ -1438,7 +1459,7 @@ async def device_asset_prompt_factory() -> Optional[str]: plan_file=self._turn_plan_file, ), prompt_tool_names=prompt_tool_names, - tool_revision=ToolRegistry.revision(), + tool_revision=self._ports.tools.revision(), memory_bootstrap_data=self._memory_bootstrap_data, static_cache=self._static_cache, sandbox_prompt_factory=sandbox_prompt_factory, @@ -2011,7 +2032,7 @@ async def _build_device_asset_hint(self) -> Optional[str]: vendor_by_storage_key: Dict[str, str] = {} try: - for tool_info in ToolRegistry.list_tools(): + for tool_info in self._ports.tools.list_tools(): if getattr(tool_info, "source", None) != "device": continue storage_key = str(getattr(tool_info, "provider", "") or "").strip() @@ -2298,7 +2319,7 @@ async def _build_callable_tool_schema( def _agent_declares_tool(self, agent: AgentInfo, tool_name: str) -> bool: """Check if agent statically declares a tool.""" - tool = ToolRegistry.get(tool_name) + tool = self._ports.tools.get(tool_name) if tool is None: return False metadata = get_tool_catalog_metadata(tool_name, tool.info) @@ -2405,7 +2426,10 @@ def _exception_to_error_dict(self, exception: Exception) -> Dict[str, Any]: def _get_context_window_tokens(self) -> int: """Resolve the context window size for the current model.""" try: - ctx, _, _ = Provider.resolve_model_info(self.provider_id, self.model_id) + ctx, _, _ = self._ports.models.resolve_model_info( + self.provider_id, + self.model_id, + ) if ctx and ctx > 0: return ctx except Exception: @@ -2607,7 +2631,10 @@ async def _to_chat_messages( tool_result_refs: List[Dict[str, Any]] = [] turn_index = 0 queued_user_message_ids: set[str] = set(getattr(self, "_queued_user_message_ids", set()) or set()) - active_model = Provider.resolve_model(self.provider_id, self.model_id) + active_model = self._ports.models.resolve_model( + self.provider_id, + self.model_id, + ) active_interleaved = ( getattr(active_model.capabilities, "interleaved", None) if active_model and getattr(active_model, "capabilities", None) @@ -3248,14 +3275,18 @@ async def _on_tool_execution_start( llm_after_enabled = False self._llm_call_aborted = False try: - llm_before_enabled = await HookPipeline.has_stage_handlers( + llm_before_enabled = ( + await self._ports.hooks.has_stage_handlers( HookStage.LLM_BEFORE, llm_hook_metadata, ) - llm_after_enabled = await HookPipeline.has_stage_handlers( + ) + llm_after_enabled = ( + await self._ports.hooks.has_stage_handlers( HookStage.LLM_AFTER, llm_hook_metadata, ) + ) except Exception as exc: log.debug("runner.hook.stage_probe.error", {"error": str(exc)}) @@ -3273,7 +3304,9 @@ async def _on_tool_execution_start( } try: hook_started_at = time.perf_counter() - await HookPipeline.run_llm_before(llm_before_hook_input) + await self._ports.hooks.run_llm_before( + llm_before_hook_input, + ) self._log_perf( "runner.hook.llm_before.complete", hook_started_at, @@ -3463,7 +3496,7 @@ async def _on_tool_execution_start( ) if llm_after_enabled: try: - await HookPipeline.run_llm_after( + await self._ports.hooks.run_llm_after( llm_hook_metadata, { "durationMs": int((time.perf_counter() - llm_call_started_at) * 1000), @@ -3575,7 +3608,7 @@ async def _on_tool_execution_start( if llm_after_enabled: try: hook_started_at = time.perf_counter() - await HookPipeline.run_llm_after( + await self._ports.hooks.run_llm_after( llm_hook_metadata, { "durationMs": int((time.perf_counter() - llm_call_started_at) * 1000), diff --git a/flocks/session/runtime_adapters.py b/flocks/session/runtime_adapters.py new file mode 100644 index 000000000..def73f248 --- /dev/null +++ b/flocks/session/runtime_adapters.py @@ -0,0 +1,85 @@ +"""Default adapters from runtime ports to existing Flocks services.""" + +from __future__ import annotations + +from typing import Any, Optional + +from flocks.agent.runtime.ports import ExternalRuntimePorts +from flocks.hooks.pipeline import HookPipeline +from flocks.provider.provider import Provider +from flocks.session.prompt import SessionPrompt +from flocks.tool.registry import ToolRegistry + + +class FlocksPromptPort: + """Adapt the existing SessionPrompt builder.""" + + async def build_system_prompts(self, **kwargs: Any) -> list[str]: + return await SessionPrompt.build_system_prompts(**kwargs) + + +class FlocksToolPort: + """Adapt the process-wide ToolRegistry.""" + + def revision(self) -> int: + return ToolRegistry.revision() + + def list_tools(self) -> list[Any]: + return ToolRegistry.list_tools() + + def get(self, name: str) -> Optional[Any]: + return ToolRegistry.get(name) + + +class FlocksModelPort: + """Adapt provider configuration and model metadata lookup.""" + + def get_provider(self, provider_id: str) -> Optional[Any]: + return Provider.get(provider_id) + + async def apply_config(self, provider_id: str) -> None: + await Provider.apply_config(provider_id=provider_id) + + def resolve_model(self, provider_id: str, model_id: str) -> Optional[Any]: + return Provider.resolve_model(provider_id, model_id) + + def resolve_model_info( + self, + provider_id: str, + model_id: str, + ) -> tuple[int, int, Optional[int]]: + return Provider.resolve_model_info(provider_id, model_id) + + +class FlocksHookPort: + """Adapt HookPipeline while preserving hook names and payloads.""" + + async def run_session_start(self, data: dict[str, Any]) -> Any: + return await HookPipeline.run_session_start(data) + + async def has_stage_handlers( + self, + stage: Any, + metadata: dict[str, Any], + ) -> bool: + return await HookPipeline.has_stage_handlers(stage, metadata) + + async def run_llm_before(self, data: dict[str, Any]) -> Any: + return await HookPipeline.run_llm_before(data) + + async def run_llm_after( + self, + metadata: dict[str, Any], + result: dict[str, Any], + ) -> Any: + return await HookPipeline.run_llm_after(metadata, result) + + +def create_default_runtime_ports() -> ExternalRuntimePorts: + """Create adapters for one runner without changing public APIs.""" + return ExternalRuntimePorts( + prompts=FlocksPromptPort(), + tools=FlocksToolPort(), + models=FlocksModelPort(), + hooks=FlocksHookPort(), + ) diff --git a/tests/session/test_runtime_ports.py b/tests/session/test_runtime_ports.py new file mode 100644 index 000000000..1b76477c0 --- /dev/null +++ b/tests/session/test_runtime_ports.py @@ -0,0 +1,108 @@ +"""Tests for prompt, tool, model, and hook runtime ports.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from flocks.agent.runtime import ExternalRuntimePorts +from flocks.session.runner import SessionRunner +from flocks.session.runtime_adapters import ( + FlocksHookPort, + FlocksModelPort, + FlocksPromptPort, + FlocksToolPort, +) + + +@pytest.mark.asyncio +async def test_runner_uses_injected_ports_at_external_boundaries() -> None: + prompt_port = SimpleNamespace(build_system_prompts=AsyncMock(return_value=[])) + tool_port = SimpleNamespace( + revision=MagicMock(return_value=7), + list_tools=MagicMock(return_value=[]), + get=MagicMock(return_value=None), + ) + model_port = SimpleNamespace( + get_provider=MagicMock( + return_value=SimpleNamespace(_config_models=[]), + ), + apply_config=AsyncMock(), + resolve_model=MagicMock( + return_value=SimpleNamespace( + capabilities=SimpleNamespace(interleaved=True), + ) + ), + resolve_model_info=MagicMock(return_value=(100_000, 8_192, None)), + ) + hook_port = SimpleNamespace( + run_session_start=AsyncMock(), + has_stage_handlers=AsyncMock(return_value=False), + run_llm_before=AsyncMock(), + run_llm_after=AsyncMock(), + ) + ports = ExternalRuntimePorts( + prompts=prompt_port, + tools=tool_port, + models=model_port, + hooks=hook_port, + ) + runner = SessionRunner( + session=SimpleNamespace(id="session-1", directory="/tmp"), + provider_id="provider-a", + model_id="model-a", + runtime_ports=ports, + session_start_pending=True, + ) + + await runner._run_session_start_hook(SimpleNamespace(name="rex")) + capability_key = runner._provider_capability_key() + cache_key = runner._tool_schema_cache_key( + SimpleNamespace(name="rex", tools=[]), + [], + text_tool_call_mode=False, + ) + + hook_port.run_session_start.assert_awaited_once() + model_port.resolve_model.assert_called_once_with("provider-a", "model-a") + assert "interleaved=true" in capability_key + assert cache_key[0] == 7 + + +@pytest.mark.asyncio +async def test_default_adapters_preserve_existing_flocks_interfaces( + monkeypatch: pytest.MonkeyPatch, +) -> None: + build_prompts = AsyncMock(return_value=["system"]) + apply_config = AsyncMock() + session_start = AsyncMock(return_value="hook-result") + monkeypatch.setattr( + "flocks.session.runtime_adapters.SessionPrompt.build_system_prompts", + build_prompts, + ) + monkeypatch.setattr( + "flocks.session.runtime_adapters.ToolRegistry.revision", + MagicMock(return_value=11), + ) + monkeypatch.setattr( + "flocks.session.runtime_adapters.Provider.get", + MagicMock(return_value="provider"), + ) + monkeypatch.setattr( + "flocks.session.runtime_adapters.Provider.apply_config", + apply_config, + ) + monkeypatch.setattr( + "flocks.session.runtime_adapters.HookPipeline.run_session_start", + session_start, + ) + + assert await FlocksPromptPort().build_system_prompts(session_id="session-1") == ["system"] + assert FlocksToolPort().revision() == 11 + assert FlocksModelPort().get_provider("provider-a") == "provider" + await FlocksModelPort().apply_config("provider-a") + assert await FlocksHookPort().run_session_start({"sessionID": "session-1"}) == ("hook-result") + + build_prompts.assert_awaited_once_with(session_id="session-1") + apply_config.assert_awaited_once_with(provider_id="provider-a") + session_start.assert_awaited_once_with({"sessionID": "session-1"}) From 75ceb24aaf94c5a8babdf5f4e8d335f69c3e81f7 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 7 Aug 2026 15:27:45 +0800 Subject: [PATCH 04/15] refactor(session): introduce lifecycle session host --- flocks/session/runtime_services.py | 103 ++ flocks/session/session_host.py | 487 ++++++ flocks/session/session_loop.py | 2381 ++++++++++++++-------------- 3 files changed, 1775 insertions(+), 1196 deletions(-) create mode 100644 flocks/session/runtime_services.py create mode 100644 flocks/session/session_host.py diff --git a/flocks/session/runtime_services.py b/flocks/session/runtime_services.py new file mode 100644 index 000000000..f30e1748e --- /dev/null +++ b/flocks/session/runtime_services.py @@ -0,0 +1,103 @@ +"""Session adapters consumed by the host-neutral agent runtime.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +from flocks.agent.runtime.contracts import ( + AgentRunState, + ContinuationDecision, + ModelTurnBoundary, + ModelTurnPreparation, + ModelTurnSnapshot, + StepResult, +) +from flocks.agent.runtime.events import RuntimeEvent + + +class SessionStepCancelled(Exception): + """Internal signal raised when a hosted model turn is cancelled.""" + + +class SessionLoopStepEngine: + """Bind the session runner and cancellation state to one AgentLoop.""" + + def __init__(self, context: Any, callbacks: Any, policy: Any): + self._context = context + self._callbacks = callbacks + self._policy = policy + + async def run(self, snapshot: ModelTurnSnapshot[Any]) -> StepResult: + """Execute a model turn while exposing its task to session abort.""" + task = asyncio.create_task( + self._policy._process_model_step( + self._context, + self._callbacks, + snapshot, + ) + ) + self._context._current_step_task = task + started_at = asyncio.get_running_loop().time() + try: + return await task + except asyncio.CancelledError as exc: + raise SessionStepCancelled from exc + finally: + self._context._current_step_task = None + duration_ms = int((asyncio.get_running_loop().time() - started_at) * 1000) + self._policy._log_step_complete(self._context, duration_ms) + + +class SessionRuntimeServices: + """Adapt session persistence and policy to narrow runtime ports.""" + + def __init__(self, context: Any, callbacks: Any, policy: Any): + self._context = context + self._callbacks = callbacks + self._policy = policy + + async def prepare_model_turn( + self, + state: AgentRunState[Any], + ) -> ModelTurnPreparation[Any]: + """Prepare persisted session state for a stable model turn.""" + return await self._policy._prepare_model_turn( + self._context, + self._callbacks, + state, + ) + + async def complete_model_turn( + self, + state: AgentRunState[Any], + step_result: StepResult, + ) -> ModelTurnBoundary[Any]: + """Commit and expose the state written by a completed model turn.""" + return await self._policy._complete_model_turn( + self._context, + self._callbacks, + state, + step_result, + ) + + async def resolve_continuation( + self, + state: AgentRunState[Any], + step_result: StepResult, + ) -> ContinuationDecision[Any]: + """Resolve queued goal and TurnFinish continuation policy.""" + return await self._policy._resolve_continuation( + self._context, + self._callbacks, + state, + step_result, + ) + + async def emit_event(self, event: RuntimeEvent) -> None: + """Forward a host-neutral runtime event to session subscribers.""" + await self._policy._publish_runtime_event( + self._callbacks, + event.type, + dict(event.payload), + ) diff --git a/flocks/session/session_host.py b/flocks/session/session_host.py new file mode 100644 index 000000000..52792b112 --- /dev/null +++ b/flocks/session/session_host.py @@ -0,0 +1,487 @@ +"""Session lifecycle host for one resumable agent execution.""" + +from __future__ import annotations + +import time +from collections.abc import Awaitable, Callable, MutableMapping +from dataclasses import dataclass, replace +from typing import Any, Optional + +from flocks.agent.runtime.contracts import ( + ModelTurnSnapshot, + RuntimeModel, + StepResult, +) +from flocks.session.core.context import DefaultSessionContext +from flocks.session.core.status import ( + SessionStatus, + SessionStatusBusy, + SessionStatusIdle, +) +from flocks.session.core.turn_state import clear_turn_state +from flocks.session.message import Message +from flocks.session.runtime_adapters import create_default_runtime_ports +from flocks.session.session import Session, is_model_auto_session_category +from flocks.utils.log import Log + + +log = Log.create(service="session.host") + + +@dataclass(frozen=True) +class SessionHostDependencies: + """Compatibility adapters supplied by the public SessionLoop facade.""" + + create_context: Callable[..., Any] + create_callbacks: Callable[[], Any] + create_result: Callable[..., Any] + resolve_model: Callable[..., Awaitable[tuple[str, str]]] + run_logical_turn: Callable[[Any, Any], Awaitable[Any]] + publish_session_status: Callable[[Any, str, str], Awaitable[None]] + + +@dataclass(frozen=True) +class SessionLease: + """Process-local ownership record protected by Session.lifecycle_lock.""" + + session_id: str + context: Any + + +class SessionLeaseRegistry: + """Manage active session ownership without exposing lifecycle policy.""" + + def __init__(self, active_contexts: MutableMapping[str, Any]): + self._active_contexts = active_contexts + + def get(self, session_id: str) -> Optional[Any]: + """Return the active context, if this process owns the session.""" + return self._active_contexts.get(session_id) + + def acquire(self, session_id: str, context: Any) -> Optional[SessionLease]: + """Acquire process-local ownership; caller holds the lifecycle lock.""" + if session_id in self._active_contexts: + return None + self._active_contexts[session_id] = context + return SessionLease(session_id=session_id, context=context) + + def release(self, lease: SessionLease) -> None: + """Release ownership only when the stored context is still ours.""" + if self._active_contexts.get(lease.session_id) is lease.context: + self._active_contexts.pop(lease.session_id, None) + + +class SessionHostStepEngine: + """Apply host-owned cross-model recovery around one-candidate attempts.""" + + def __init__( + self, + *, + context: Any, + callbacks: Any, + attempt_engine: Any, + cooldowns: MutableMapping[str, Any], + cooldown_factory: Callable[..., Any], + select_candidate: Callable[[Any, int], None], + finalize_failure: Callable[[Any, Any, Any], Awaitable[None]], + publish_event: Callable[[Any, str, dict[str, Any]], Awaitable[None]], + rate_limit_cooldown_seconds: float, + chain_exhaustion_cooldown_seconds: float, + ): + self._context = context + self._callbacks = callbacks + self._attempt_engine = attempt_engine + self._cooldowns = cooldowns + self._cooldown_factory = cooldown_factory + self._select_candidate = select_candidate + self._finalize_failure = finalize_failure + self._publish_event = publish_event + self._rate_limit_cooldown_seconds = rate_limit_cooldown_seconds + self._chain_exhaustion_cooldown_seconds = chain_exhaustion_cooldown_seconds + + async def run(self, snapshot: ModelTurnSnapshot[Any]) -> StepResult: + """Retry a replay-safe snapshot across the configured model chain.""" + while True: + active_model = RuntimeModel( + self._context.provider_id, + self._context.model_id, + ) + result = await self._attempt_engine.run( + replace(snapshot, active_model=active_model), + ) + failure = result.failure + if not self._context.auto_failover or failure is None: + return result + + next_index = self._context.candidate_index + 1 + has_next = next_index < len(self._context.model_candidates) + if not failure.allow_fallback or not failure.attempt_state.replay_safe or not has_next: + self._record_chain_exhaustion(failure, has_next) + await self._finalize_failure( + self._context, + failure, + snapshot.last_user, + ) + return result + + if not await self._remove_failed_attempt(failure): + await self._finalize_failure( + self._context, + failure, + snapshot.last_user, + ) + return result + + await self._switch_candidate(next_index, failure.reason) + + def _record_chain_exhaustion(self, failure: Any, has_next: bool) -> None: + context = self._context + if not ( + context.model_candidate_policy == "automatic" + and failure.allow_fallback + and failure.attempt_state.replay_safe + and not has_next + and context.candidate_index > 0 + and failure.reason not in {"rate_limit", "billing"} + ): + return + + expires_at = time.monotonic() + self._chain_exhaustion_cooldown_seconds + existing = self._cooldowns.get(context.session.id) + if existing and existing.expires_at > expires_at: + return + self._cooldowns[context.session.id] = self._cooldown_factory( + model=context.model_candidates[context.candidate_index], + primary=context.model_candidates[0], + expires_at=expires_at, + reason="chain_exhausted", + ) + + async def _remove_failed_attempt(self, failure: Any) -> bool: + message_id = failure.assistant_message_id + if not message_id: + return True + try: + deleted = await Message.delete(self._context.session.id, message_id) + except Exception as exc: + deleted = False + log.error( + "session.model.fallback_cleanup_failed", + { + "session_id": self._context.session.id, + "message_id": message_id, + "error": str(exc), + }, + ) + if not deleted: + return False + await self._publish_event( + self._callbacks, + "message.removed", + { + "sessionID": self._context.session.id, + "messageID": message_id, + }, + ) + return True + + async def _switch_candidate(self, next_index: int, reason: str) -> None: + context = self._context + previous = context.model_candidates[context.candidate_index] + next_candidate = context.model_candidates[next_index] + if context.model_candidate_policy == "automatic": + if context.candidate_index == 0 and reason in { + "rate_limit", + "billing", + }: + self._cooldowns[context.session.id] = self._cooldown_factory( + model=next_candidate, + primary=context.model_candidates[0], + expires_at=(time.monotonic() + self._rate_limit_cooldown_seconds), + reason=reason, + ) + else: + cooldown = self._cooldowns.get(context.session.id) + if cooldown and cooldown.expires_at > time.monotonic(): + cooldown.model = next_candidate + + self._select_candidate(context, next_index) + payload = { + "sessionID": context.session.id, + "from": { + "providerID": previous.provider_id, + "modelID": previous.model_id, + }, + "to": { + "providerID": next_candidate.provider_id, + "modelID": next_candidate.model_id, + }, + "reason": reason, + "candidateIndex": next_index, + } + log.warn( + "session.model.fallback", + { + "from": payload["from"], + "to": payload["to"], + "reason": reason, + "candidateIndex": next_index, + }, + ) + await self._publish_event( + self._callbacks, + "session.model.fallback", + payload, + ) + + +class SessionHost: + """Own session acquisition, recovery, execution, and final cleanup.""" + + def __init__( + self, + dependencies: SessionHostDependencies, + active_contexts: MutableMapping[str, Any], + ): + self._dependencies = dependencies + self._leases = SessionLeaseRegistry(active_contexts) + + async def run( + self, + session_id: str, + provider_id: Optional[str] = None, + model_id: Optional[str] = None, + agent_name: Optional[str] = None, + callbacks: Optional[Any] = None, + working_directory: Optional[str] = None, + auto_failover: bool = False, + ) -> Any: + """Acquire and host one session until its logical turn settles.""" + active_context = self._leases.get(session_id) + if active_context is not None: + log.info("session.already_running", {"session_id": session_id}) + self._authorize_auto_failover(active_context, auto_failover) + return self._dependencies.create_result( + action="queued", + error="Loop already running", + ) + + session = await Session.get_by_id(session_id) + if session is None: + log.warning("session.not_found", {"session_id": session_id}) + return self._dependencies.create_result( + action="error", + error=f"Session {session_id} not found", + ) + if session.status != "active": + log.warning( + "session.not_active", + {"session_id": session_id, "status": session.status}, + ) + return self._dependencies.create_result( + action="error", + error=f"Session {session_id} is {session.status}", + ) + if working_directory: + session = session.model_copy(update={"directory": working_directory}) + + if not provider_id or not model_id: + resolved_provider, resolved_model = await self._dependencies.resolve_model( + session, + provider_id, + model_id, + ) + provider_id = provider_id or resolved_provider + model_id = model_id or resolved_model + + primary_model = RuntimeModel(provider_id=provider_id, model_id=model_id) + auto_failover = bool( + auto_failover + and is_model_auto_session_category( + getattr(session, "category", "user"), + ) + ) + session.provider = provider_id + session.model = model_id + + trace_offset = await self._load_trace_offset(session_id) + context = self._dependencies.create_context( + session=session, + provider_id=provider_id, + model_id=model_id, + agent_name=agent_name or session.agent or "rex", + session_ctx=DefaultSessionContext(session), + trace_step_offset=trace_offset, + auto_failover=auto_failover, + auto_failover_allowed=auto_failover, + model_candidates=[primary_model], + candidate_index=0, + session_start_pending=trace_offset == 0, + runtime_ports=create_default_runtime_ports(), + ) + lease_or_result = await self._acquire_lease(session_id, context) + if not isinstance(lease_or_result, SessionLease): + return lease_or_result + lease = lease_or_result + + runtime_callbacks = callbacks or self._dependencies.create_callbacks() + await self._mark_busy(session_id, runtime_callbacks) + await self._recover_orphan_tools(session_id) + + try: + return await self._dependencies.run_logical_turn( + context, + runtime_callbacks, + ) + except Exception as exc: + return await self._handle_execution_error( + context, + callbacks, + exc, + ) + finally: + await self._release_session( + lease, + session, + runtime_callbacks, + ) + + @staticmethod + def _authorize_auto_failover(context: Any, requested: bool) -> None: + if requested and is_model_auto_session_category( + getattr(context.session, "category", "user"), + ): + context.auto_failover_allowed = True + + @staticmethod + async def _load_trace_offset(session_id: str) -> int: + try: + messages = await Message.list(session_id) + return sum(1 for message in messages if message.role == "assistant") + except Exception as exc: + log.debug("session.trace_offset.error", {"error": str(exc)}) + return 0 + + async def _acquire_lease( + self, + session_id: str, + context: Any, + ) -> SessionLease | Any: + async with Session.lifecycle_lock(session_id): + latest_session = await Session.get_by_id(session_id) + if latest_session is None: + log.warning( + "session.not_found_before_lease", + {"session_id": session_id}, + ) + return self._dependencies.create_result( + action="error", + error=f"Session {session_id} not found", + ) + if latest_session.status != "active": + log.warning( + "session.not_active_before_lease", + { + "session_id": session_id, + "status": latest_session.status, + }, + ) + return self._dependencies.create_result( + action="error", + error=f"Session {session_id} is {latest_session.status}", + ) + if Session.is_lifecycle_transitioning(session_id): + return self._dependencies.create_result( + action="error", + error=f"Session {session_id} is changing lifecycle state", + ) + lease = self._leases.acquire(session_id, context) + if lease is None: + return self._dependencies.create_result( + action="queued", + error="Loop already running", + ) + return lease + + async def _mark_busy(self, session_id: str, callbacks: Any) -> None: + SessionStatus.set(session_id, SessionStatusBusy()) + await self._dependencies.publish_session_status( + callbacks, + session_id, + "busy", + ) + + @staticmethod + async def _recover_orphan_tools(session_id: str) -> None: + try: + from flocks.session.orphan_tools import abort_orphan_running_parts + + await abort_orphan_running_parts(session_id) + except Exception as exc: + log.warn( + "session.orphan_cleanup_failed", + {"session_id": session_id, "error": str(exc)}, + ) + + async def _handle_execution_error( + self, + context: Any, + callbacks: Optional[Any], + error: Exception, + ) -> Any: + session_id = context.session.id + log.error( + "session.execution_error", + {"session_id": session_id, "error": str(error)}, + ) + if callbacks and callbacks.on_error: + try: + await callbacks.on_error(str(error)) + except Exception as callback_error: + log.debug( + "session.error_callback_failed", + {"error": str(callback_error)}, + ) + try: + from flocks.bus.bus import Bus + from flocks.bus.events import SessionError + + await Bus.publish( + SessionError, + {"sessionID": session_id, "error": str(error)}, + ) + except Exception as publish_error: + log.warn( + "session.error_event_failed", + {"error": str(publish_error)}, + ) + return self._dependencies.create_result( + action="error", + error=str(error), + provider_id=context.provider_id, + model_id=context.model_id, + ) + + async def _release_session( + self, + lease: SessionLease, + session: Any, + callbacks: Any, + ) -> None: + self._leases.release(lease) + clear_turn_state(lease.session_id) + SessionStatus.set(lease.session_id, SessionStatusIdle()) + await self._dependencies.publish_session_status( + callbacks, + lease.session_id, + "idle", + ) + await Session.touch(session.project_id, lease.session_id) + + try: + from flocks.bus.bus import Bus + from flocks.bus.events import SessionIdle + + await Bus.publish(SessionIdle, {"sessionID": lease.session_id}) + except Exception as exc: + log.warn("session.idle_event_failed", {"error": str(exc)}) diff --git a/flocks/session/session_loop.py b/flocks/session/session_loop.py index cbf8f9f11..354ffa4a8 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -20,7 +20,20 @@ from dataclasses import dataclass, field from datetime import datetime -from flocks.agent.runtime.contracts import ModelTurnSnapshot, RuntimeModel +from flocks.agent.runtime.agent_loop import AgentLoop +from flocks.agent.runtime.contracts import ( + AgentRunState, + AgentRunStatus, + ContinuationDecision, + ModelTurnBoundary, + ModelTurnPreparation, + ModelTurnSnapshot, + QueuedInputBatch, + RuntimeModel, + StepResult, + TurnPreparationStatus, +) +from flocks.agent.runtime.ports import ExternalRuntimePorts from flocks.utils.log import Log from flocks.utils.id import Identifier from flocks.session.session import ( @@ -29,12 +42,11 @@ is_model_auto_session_category, ) from flocks.session.message import Message, MessageInfo, MessageRole -from flocks.session.core.status import SessionStatus, SessionStatusBusy, SessionStatusIdle +from flocks.session.core.status import SessionStatus, SessionStatusBusy from flocks.session.core.task_utils import fire_and_forget from flocks.session.core.turn_state import ( set_turn_state, set_context_state, - clear_turn_state, ) from flocks.session.lifecycle.compaction import ( SessionCompaction, @@ -110,6 +122,7 @@ class LoopContext: turn_additional_context: Optional[str] = None stop_hook_active: bool = False session_start_pending: bool = False + runtime_ports: Optional[ExternalRuntimePorts] = field(default=None, repr=False) @property def trace_step(self) -> int: @@ -576,224 +589,32 @@ async def run( working_directory: Optional[str] = None, auto_failover: bool = False, ) -> LoopResult: - """ - Run session loop - - Main entry point matching Flocks' SessionPrompt.loop() - - When provider_id/model_id are not provided, resolves from: - 1. Session's stored model (if set during creation) - 2. Global default LLM (default_models.llm -> config.model) - 3. Environment variables - 4. Hardcoded fallback - - Args: - session_id: Session ID to process - provider_id: Provider ID - model_id: Model ID - agent_name: Agent name (default: build) - callbacks: Loop callbacks - - Returns: - LoopResult with final state - """ - # Check if already running. - # Return action="queued" (not "error") so the route layer knows to skip - # creating a spurious empty assistant message. The new user message is - # already persisted in the DB; the active loop will pick it up on its - # next iteration once it finishes the current step. - if cls.is_running(session_id): - log.info("loop.already_running", {"session_id": session_id}) - if auto_failover: - active_ctx = cls._active_loops.get(session_id) - if ( - active_ctx is not None - and is_model_auto_session_category( - getattr(active_ctx.session, "category", "user") - ) - ): - active_ctx.auto_failover_allowed = True - return LoopResult( - action="queued", - error="Loop already running", - ) - - # Get session - session = await Session.get_by_id(session_id) - if not session: - log.warning("loop.session_not_found", {"session_id": session_id}) - return LoopResult( - action="error", - error=f"Session {session_id} not found", - ) - if session.status != "active": - log.warning("loop.session_not_active", { - "session_id": session_id, - "status": session.status, - }) - return LoopResult( - action="error", - error=f"Session {session_id} is {session.status}", - ) - if working_directory: - session = session.model_copy(update={"directory": working_directory}) - - # Resolve model when not explicitly provided - if not provider_id or not model_id: - resolved_provider, resolved_model = await cls._resolve_model( - session, provider_id, model_id - ) - provider_id = provider_id or resolved_provider - model_id = model_id or resolved_model - - primary_model = RuntimeModel( - provider_id=provider_id, - model_id=model_id, - ) - model_candidates = [primary_model] - candidate_index = 0 - auto_failover = bool( - auto_failover - and is_model_auto_session_category( - getattr(session, "category", "user") - ) + """Run one session through the lifecycle-owned SessionHost.""" + from flocks.session.session_host import ( + SessionHost, + SessionHostDependencies, ) - # Keep the in-memory session aligned with the runtime model so - # downstream helpers (title generation, compaction checks, etc.) see - # the model actually selected for this loop iteration. Unpinned - # sessions must not persist these values; otherwise switching the - # global default model would keep older sessions stuck on stale data. - if provider_id: - session.provider = provider_id - if model_id: - session.model = model_id - - # Create SessionContext interface for decoupled access - from flocks.session.core.context import DefaultSessionContext - session_ctx = DefaultSessionContext(session) - - # Compute trace step offset from existing assistant messages so - # observability step numbers are cumulative across the whole session. - trace_offset = 0 - try: - existing_msgs = await Message.list(session_id) - trace_offset = sum(1 for m in existing_msgs if m.role == "assistant") - except Exception as _trace_err: - log.debug("loop.trace_offset.error", {"error": str(_trace_err)}) - - # Create context - ctx = LoopContext( - session=session, + host = SessionHost( + dependencies=SessionHostDependencies( + create_context=LoopContext, + create_callbacks=LoopCallbacks, + create_result=LoopResult, + resolve_model=cls._resolve_model, + run_logical_turn=cls._run_loop, + publish_session_status=cls._publish_session_status, + ), + active_contexts=cls._active_loops, + ) + return await host.run( + session_id=session_id, provider_id=provider_id, model_id=model_id, - agent_name=agent_name or session.agent or "rex", - session_ctx=session_ctx, - trace_step_offset=trace_offset, + agent_name=agent_name, + callbacks=callbacks, + working_directory=working_directory, auto_failover=auto_failover, - auto_failover_allowed=auto_failover, - model_candidates=model_candidates, - candidate_index=candidate_index, - session_start_pending=trace_offset == 0, ) - - # Register under the same lock used by archive/delete. This closes the - # gap where archival could commit after the status check above but - # before the loop became visible to the lifecycle stop logic. - async with Session.lifecycle_lock(session_id): - latest_session = await Session.get_by_id(session_id) - if latest_session is None: - log.warning("loop.session_not_found_before_register", { - "session_id": session_id, - }) - return LoopResult( - action="error", - error=f"Session {session_id} not found", - ) - if latest_session.status != "active": - log.warning("loop.session_not_active_before_register", { - "session_id": session_id, - "status": latest_session.status, - }) - return LoopResult( - action="error", - error=f"Session {session_id} is {latest_session.status}", - ) - if Session.is_lifecycle_transitioning(session_id): - return LoopResult( - action="error", - error=f"Session {session_id} is changing lifecycle state", - ) - if cls.is_running(session_id): - return LoopResult( - action="queued", - error="Loop already running", - ) - cls._active_loops[session_id] = ctx - - # Set status to busy - SessionStatus.set(session_id, SessionStatusBusy()) - await cls._publish_session_status(callbacks or LoopCallbacks(), session_id, "busy") - - # Mark orphaned running tool parts as error (e.g. from server restart). - # Wrapped in try/except so cleanup failures never block the session loop. - try: - from flocks.session.orphan_tools import abort_orphan_running_parts - - await abort_orphan_running_parts(session_id) - except Exception as exc: - log.warn("loop.orphan_cleanup_failed", { - "session_id": session_id, - "error": str(exc), - }) - - try: - # Run loop iteration - result = await cls._run_loop(ctx, callbacks or LoopCallbacks()) - return result - except Exception as e: - log.error("loop.error", {"session_id": session_id, "error": str(e)}) - # Report error to callbacks so CLI/TUI can display it - if callbacks and callbacks.on_error: - try: - await callbacks.on_error(str(e)) - except Exception as _cb_err: - log.debug("loop.error.callback_failed", {"error": str(_cb_err)}) - try: - from flocks.bus.bus import Bus - from flocks.bus.events import SessionError - await Bus.publish(SessionError, { - "sessionID": session_id, - "error": str(e), - }) - except Exception as exc: - log.warn("loop.error.event_error", {"error": str(exc)}) - return LoopResult( - action="error", - error=str(e), - provider_id=ctx.provider_id, - model_id=ctx.model_id, - ) - finally: - # Clean up - if session_id in cls._active_loops: - del cls._active_loops[session_id] - clear_turn_state(session_id) - - # Set status to idle - SessionStatus.set(session_id, SessionStatusIdle()) - await cls._publish_session_status(callbacks or LoopCallbacks(), session_id, "idle") - - # Touch session (update timestamp) - await Session.touch(session.project_id, session_id) - - # Publish idle event - try: - from flocks.bus.bus import Bus - from flocks.bus.events import SessionIdle - await Bus.publish(SessionIdle, {"sessionID": session_id}) - except Exception as exc: - log.warn("loop.idle.event_error", {"error": str(exc)}) @staticmethod async def _resolve_model( @@ -1281,1033 +1102,1197 @@ async def _finalize_deferred_failure( ) @classmethod - async def _process_step_with_failover( + async def _process_model_step( cls, ctx: LoopContext, callbacks: LoopCallbacks, - messages: List[MessageInfo], - last_user: MessageInfo, - ) -> Any: - """Run one logical step, moving across candidates without replaying output.""" + snapshot: ModelTurnSnapshot[MessageInfo], + ) -> StepResult: + """Run one candidate attempt; provider-local retries stay in runner.""" from flocks.session.runner import RunnerCallbacks, SessionRunner from flocks.session.step_engine import SessionStepEngine - while True: - runner_cbs = callbacks.runner_callbacks - if runner_cbs is None: - runner_cbs = RunnerCallbacks() - if callbacks.event_publish_callback and not runner_cbs.event_publish_callback: - runner_cbs.event_publish_callback = callbacks.event_publish_callback - - runner = SessionRunner( - session=ctx.session, - provider_id=ctx.provider_id, - model_id=ctx.model_id, - agent_name=ctx.agent_name, - abort_event=ctx.abort_event, - callbacks=runner_cbs, - session_ctx=ctx.session_ctx, - memory_bootstrap_data=ctx.memory_bootstrap_data, - static_cache=ctx.runner_static_cache, - defer_step_errors=ctx.auto_failover, - failover_available=( - ctx.auto_failover - and ctx.candidate_index + 1 < len(ctx.model_candidates) - ), - turn_additional_context=ctx.turn_additional_context, - session_start_pending=ctx.session_start_pending, - ) - step_engine = SessionStepEngine(runner) - snapshot = ModelTurnSnapshot( - session_id=ctx.session.id, - agent_name=ctx.agent_name, - active_model=RuntimeModel( - provider_id=ctx.provider_id, - model_id=ctx.model_id, - ), - model_turn_index=ctx.step, - trace_step=ctx.trace_step, - messages=tuple(messages), - last_user=last_user, + runner_callbacks = callbacks.runner_callbacks + if runner_callbacks is None: + runner_callbacks = RunnerCallbacks() + if ( + callbacks.event_publish_callback + and not runner_callbacks.event_publish_callback + ): + runner_callbacks.event_publish_callback = ( + callbacks.event_publish_callback ) - step_result = await step_engine.run(snapshot) - if step_engine.session_start_fired: - ctx.session_start_pending = False - failure = step_result.failure - if not ctx.auto_failover or failure is None: - return step_result - - next_index = ctx.candidate_index + 1 - has_next = next_index < len(ctx.model_candidates) - if not failure.allow_fallback or not has_next: - if ( - ctx.model_candidate_policy == "automatic" - and failure.allow_fallback - and not has_next - and ctx.candidate_index > 0 - and failure.reason not in {"rate_limit", "billing"} - ): - expires_at = time.monotonic() + CHAIN_EXHAUSTION_COOLDOWN_SECONDS - existing_cooldown = cls._auto_failover_cooldowns.get(ctx.session.id) - if not ( - existing_cooldown - and existing_cooldown.expires_at > expires_at - ): - cls._auto_failover_cooldowns[ctx.session.id] = AutoFailoverCooldown( - model=ctx.model_candidates[ctx.candidate_index], - primary=ctx.model_candidates[0], - expires_at=expires_at, - reason="chain_exhausted", - ) - await cls._finalize_deferred_failure(ctx, failure, last_user) - return step_result - # A candidate may be removed only while its attempt is completely - # replay-safe. Failure to delete stops the switch to avoid leaving - # two assistant cards for one logical response. - if failure.assistant_message_id: - try: - deleted = await Message.delete( - ctx.session.id, - failure.assistant_message_id, - ) - except Exception as exc: - deleted = False - log.error("session.model.fallback_cleanup_failed", { - "session_id": ctx.session.id, - "message_id": failure.assistant_message_id, - "error": str(exc), - }) - if not deleted: - await cls._finalize_deferred_failure(ctx, failure, last_user) - return step_result - await cls._publish_runtime_event(callbacks, "message.removed", { - "sessionID": ctx.session.id, - "messageID": failure.assistant_message_id, - }) + runner = SessionRunner( + session=ctx.session, + provider_id=ctx.provider_id, + model_id=ctx.model_id, + agent_name=ctx.agent_name, + abort_event=ctx.abort_event, + callbacks=runner_callbacks, + session_ctx=ctx.session_ctx, + memory_bootstrap_data=ctx.memory_bootstrap_data, + static_cache=ctx.runner_static_cache, + defer_step_errors=ctx.auto_failover, + failover_available=( + ctx.auto_failover + and ctx.candidate_index + 1 < len(ctx.model_candidates) + ), + turn_additional_context=ctx.turn_additional_context, + session_start_pending=ctx.session_start_pending, + runtime_ports=ctx.runtime_ports, + ) + step_engine = SessionStepEngine(runner) + result = await step_engine.run(snapshot) + if step_engine.session_start_fired: + ctx.session_start_pending = False + return result + + @classmethod + def _create_host_step_engine( + cls, + ctx: LoopContext, + callbacks: LoopCallbacks, + ) -> Any: + """Compose one-candidate execution with host-owned model recovery.""" + from flocks.session.runtime_services import SessionLoopStepEngine + from flocks.session.session_host import SessionHostStepEngine + + return SessionHostStepEngine( + context=ctx, + callbacks=callbacks, + attempt_engine=SessionLoopStepEngine(ctx, callbacks, cls), + cooldowns=cls._auto_failover_cooldowns, + cooldown_factory=AutoFailoverCooldown, + select_candidate=cls._select_candidate, + finalize_failure=cls._finalize_deferred_failure, + publish_event=cls._publish_runtime_event, + rate_limit_cooldown_seconds=RATE_LIMIT_COOLDOWN_SECONDS, + chain_exhaustion_cooldown_seconds=( + CHAIN_EXHAUSTION_COOLDOWN_SECONDS + ), + ) - previous = ctx.model_candidates[ctx.candidate_index] - next_candidate = ctx.model_candidates[next_index] + @classmethod + async def _process_step_with_failover( + cls, + ctx: LoopContext, + callbacks: LoopCallbacks, + messages: List[MessageInfo], + last_user: MessageInfo, + ) -> StepResult: + """Compatibility entry point for host-owned cross-model recovery.""" + snapshot = ModelTurnSnapshot( + session_id=ctx.session.id, + agent_name=ctx.agent_name, + active_model=RuntimeModel(ctx.provider_id, ctx.model_id), + model_turn_index=ctx.step, + trace_step=ctx.trace_step, + messages=tuple(messages), + last_user=last_user, + ) + return await cls._create_host_step_engine(ctx, callbacks).run(snapshot) - if ctx.model_candidate_policy == "automatic": - if ctx.candidate_index == 0 and failure.reason in {"rate_limit", "billing"}: - cls._auto_failover_cooldowns[ctx.session.id] = AutoFailoverCooldown( - model=next_candidate, - primary=ctx.model_candidates[0], - expires_at=time.monotonic() + RATE_LIMIT_COOLDOWN_SECONDS, - reason=failure.reason, - ) - else: - cooldown = cls._auto_failover_cooldowns.get(ctx.session.id) - if cooldown and cooldown.expires_at > time.monotonic(): - cooldown.model = next_candidate + @staticmethod + def _log_step_complete(ctx: LoopContext, duration_ms: int) -> None: + """Record model-turn latency from the session StepEngine adapter.""" + log.debug( + "loop.step_complete", + { + "session_id": ctx.session.id, + "step": ctx.step, + "duration_ms": duration_ms, + }, + ) - cls._select_candidate(ctx, next_index) - event_payload = { - "sessionID": ctx.session.id, - "from": { - "providerID": previous.provider_id, - "modelID": previous.model_id, + @classmethod + async def _complete_model_turn( + cls, + ctx: LoopContext, + callbacks: LoopCallbacks, + state: AgentRunState[MessageInfo], + step_result: StepResult, + ) -> ModelTurnBoundary[MessageInfo]: + """Expose the persisted state written by one completed model turn.""" + if callbacks.on_step_end: + await callbacks.on_step_end(ctx.step) + + if step_result.error and callbacks.on_error: + await callbacks.on_error(step_result.error) + + if ctx.session_ctx: + post_messages = await ctx.session_ctx.get_messages() + else: + post_messages = await Message.list(ctx.session.id) + + last_user = state.metadata.get("last_user") + last_message = next( + ( + message + for message in reversed(post_messages) + if message.role == MessageRole.ASSISTANT + and ( + not ctx.auto_failover + or last_user is None + or getattr(message, "parentID", None) == last_user.id + ) + ), + None, + ) + state.metadata["last_message"] = last_message + + queued_user = None + if last_user is not None: + queued_user = await cls._detect_queued_user_message( + ctx.session.id, + post_messages, + last_user.id, + last_message, + ) + + if queued_user is not None: + turn_state = set_turn_state( + ctx.session.id, + step=ctx.step, + status="continued", + continue_reason="queued_message", + queued_message_detected=True, + ) + await cls._publish_runtime_event( + callbacks, + "turn.continued", + { + **turn_state.model_dump(by_alias=True), + "queuedUserMessageID": queued_user.id, }, - "to": { - "providerID": next_candidate.provider_id, - "modelID": next_candidate.model_id, + ) + log.info( + "loop.continuing_for_queued_message", + { + "session_id": ctx.session.id, + "queued_user_id": queued_user.id, + "last_assistant_id": ( + last_message.id if last_message else None + ), }, - "reason": failure.reason, - "candidateIndex": next_index, - } - log.warn("session.model.fallback", { - "from": event_payload["from"], - "to": event_payload["to"], - "reason": event_payload["reason"], - "candidateIndex": event_payload["candidateIndex"], - }) + ) + elif step_result.action == "continue": + turn_state = set_turn_state( + ctx.session.id, + step=ctx.step, + status="continued", + continue_reason="tool_calls", + queued_message_detected=False, + ) await cls._publish_runtime_event( callbacks, - "session.model.fallback", - event_payload, + "turn.continued", + turn_state.model_dump(by_alias=True), ) + elif step_result.error: + await cls._publish_turn_stopped( + callbacks, + ctx.session.id, + step=ctx.step, + stop_reason=step_result.error, + ) + + return ModelTurnBoundary( + messages=tuple(post_messages), + last_message=last_message, + queued_inputs=QueuedInputBatch( + messages=(queued_user,) if queued_user is not None else (), + cursor=queued_user.id if queued_user is not None else None, + ), + ) @classmethod - async def _run_loop( + async def _resolve_continuation( cls, ctx: LoopContext, callbacks: LoopCallbacks, - ) -> LoopResult: - """ - Main loop iteration logic - - 完全匹配 TUI SessionPrompt.loop() 的结构: - 1. Get messages and analyze (lastUser, lastAssistant, lastFinished) - 2. Check exit conditions - 3. Generate title on first step - 4. Check for pending tasks (subtask/compaction) - 5. Check context overflow (compaction before step) - 6. Process step (call LLM + tools) - 7. Loop until complete - """ - last_message: Optional[MessageInfo] = None - loop_error: Optional[str] = None - - while not ctx.should_abort(): - # Set status to busy - SessionStatus.set(ctx.session.id, SessionStatusBusy()) - - ctx.step += 1 - turn_state = set_turn_state( + state: AgentRunState[MessageInfo], + step_result: StepResult, + ) -> ContinuationDecision[MessageInfo]: + """Resolve goal and TurnFinish continuation after a natural stop.""" + last_user = state.metadata.get("last_user") + last_message = state.metadata.get("last_message") + if last_user is None or last_message is None: + await cls._publish_turn_stopped( + callbacks, ctx.session.id, step=ctx.step, - status="started", - queued_message_detected=False, + stop_reason="stop", ) - await cls._publish_runtime_event(callbacks, "turn.started", turn_state.model_dump(by_alias=True)) - log.info("loop.step", { - "session_id": ctx.session.id, - "step": ctx.step, - }) - - # Callback: step start - if callbacks.on_step_start: - await callbacks.on_step_start(ctx.step) - - # Get messages via SessionContext interface - messages_started_at = asyncio.get_event_loop().time() - if ctx.session_ctx: - messages = await ctx.session_ctx.get_messages() - else: - messages = await Message.list(ctx.session.id) - log.debug("loop.messages_loaded", { - "session_id": ctx.session.id, - "step": ctx.step, - "message_count": len(messages), - "duration_ms": int((asyncio.get_event_loop().time() - messages_started_at) * 1000), - }) - if not messages: - log.info("loop.no_messages", {"session_id": ctx.session.id}) - await cls._publish_turn_stopped( - callbacks, - ctx.session.id, - step=ctx.step, - stop_reason="no_messages", - ) - break - - # Analyze messages (matching TUI lines 277-292) - last_user: Optional[MessageInfo] = None - last_assistant: Optional[MessageInfo] = None - last_finished: Optional[MessageInfo] = None - tasks: List[tuple[str, Any]] = [] # (type, part) - compaction or subtask - - scan_started_at = asyncio.get_event_loop().time() - for msg in reversed(messages): - # Find lastUser - if not last_user and msg.role == MessageRole.USER: - last_user = msg - - # Find lastAssistant - if not last_assistant and msg.role == MessageRole.ASSISTANT: - last_assistant = msg - - # Find lastFinished - if not last_finished and msg.role == MessageRole.ASSISTANT and hasattr(msg, 'finish') and msg.finish: - last_finished = msg - - # Stop when we have both lastUser and lastFinished - if last_user and last_finished: - break - - # Collect pending tasks before lastFinished - if not last_finished: - parts = await Message.parts(msg.id, ctx.session.id) - for part in parts: - if part.type == "compaction": - tasks.append(("compaction", part)) - elif part.type == "subtask": - tasks.append(("subtask", part)) - log.debug("loop.message_scan_complete", { - "session_id": ctx.session.id, - "step": ctx.step, - "task_count": len(tasks), - "duration_ms": int((asyncio.get_event_loop().time() - scan_started_at) * 1000), - }) - - # Check if we have a user message - if not last_user: - log.info("loop.no_user_message", { - "session_id": ctx.session.id, - "message_count": len(messages), - "roles": [str(getattr(msg, "role", "")) for msg in messages[-5:]], - }) - await cls._publish_turn_stopped( - callbacks, - ctx.session.id, - step=ctx.step, - stop_reason="no_user_message", - ) - break + return ContinuationDecision() - last_assistant_parts = ( - await Message.parts(last_assistant.id, ctx.session.id) - if last_assistant - else [] + try: + content_result = Message.get_text_content(last_message) + last_response = ( + await content_result + if inspect.isawaitable(content_result) + else content_result ) - - # Check exit conditions (matching TUI lines 295-302) - if cls._should_exit(last_user, last_assistant, last_assistant_parts): - log.info("loop.exit_condition", { + except Exception as exc: + log.warn( + "goal.last_response.error", + { "session_id": ctx.session.id, - "last_user_id": last_user.id, - "last_assistant_id": last_assistant.id if last_assistant else None, - "finish": last_assistant.finish if last_assistant else None, - "has_tool_parts": any( - getattr(part, "type", None) == "tool" - for part in last_assistant_parts - ), - }) - last_message = last_assistant - break + "message_id": getattr(last_message, "id", None), + "error": str(exc), + }, + ) + last_response = getattr(last_message, "content", "") or "" - if await cls._prepare_auto_turn(ctx, last_user): - ctx.turn_additional_context = None - ctx.stop_hook_active = False - await cls._run_user_prompt_submit_hook(ctx, last_user) - - # Bootstrap memory on first step (once per loop, stored in ctx) - if ctx.step == 1 and ctx.session.memory_enabled and ctx.memory_bootstrap_data is None: - try: - from flocks.memory.bootstrap import MemoryBootstrap - ctx.memory_bootstrap_data = await MemoryBootstrap( - project_id=ctx.session.project_id, - ).bootstrap(load_daily=False) - log.info("loop.memory_bootstrap_done", { - "session_id": ctx.session.id, - "has_main": ctx.memory_bootstrap_data.get("main_memory") is not None, - }) - except Exception as e: - log.error("loop.memory_bootstrap_error", {"error": str(e)}) - - # Early title generation: fire concurrently with the first LLM call so - # the title is ready (or nearly so) by the time the response completes. - # This is an optimistic fast-path — CLISessionRunner._process_message() - # also calls generate_title_after_first_message() after the loop as a - # guaranteed safety net (handles single-run mode where asyncio cleanup - # may cancel this task before it finishes). - # generate_title_after_first_message is idempotent: if this task saves - # the title first, the safety-net call returns immediately. - if ctx.step == 1 and not ctx.auto_failover: - try: - from flocks.session.lifecycle.title import SessionTitle - # UserMessageInfo.model is Dict[str, str] {"providerID": ..., "modelID": ...} - user_model = getattr(last_user, 'model', None) if last_user else None - if isinstance(user_model, dict): - title_model_id = user_model.get("modelID", ctx.model_id) - title_provider_id = user_model.get("providerID", ctx.provider_id) - else: - title_model_id = ctx.model_id - title_provider_id = ctx.provider_id - fire_and_forget( - SessionTitle.ensure_title( - session_id=ctx.session.id, - model_id=title_model_id, - provider_id=title_provider_id, - messages=messages, - event_publish_callback=callbacks.event_publish_callback if callbacks else None, - ), - label="title_generation", - name=f"title:{ctx.session.id}", - ) - except Exception as e: - log.error("loop.title_generation.error", {"error": str(e)}) - - # Check for pending tasks (matching TUI lines 314-493) - if tasks: - task_type, task_part = tasks.pop() - - # Handle pending subtask (matching TUI lines 316-481) - if task_type == "subtask": - log.info("loop.subtask_detected", { - "session_id": ctx.session.id, - "step": ctx.step, - }) - - # Execute subtask using tool execution - await cls._execute_subtask(ctx, last_user, task_part) - - # Continue to next iteration - continue - - # Handle pending compaction (matching TUI lines 483-494) - elif task_type == "compaction": - log.info("loop.compaction_pending", { - "session_id": ctx.session.id, - "step": ctx.step, - "auto": getattr(task_part, 'auto', False), - }) - - # Callback: compaction - if callbacks.on_compaction: - await callbacks.on_compaction() - - # Build dynamic CompactionPolicy from model info - compaction_policy = cls._build_compaction_policy(ctx) - - # Auto-compaction also surfaces a "Compacting..." - # banner on the UI (driven by ``session.status`` → - # ``compacting``), so we wire the same SSE progress - # adapter as the manual ``/compact`` route. The - # closure captures ``ctx.session.id`` and the - # publish callback explicitly to keep behaviour - # identical between loop and route paths. - _publish = callbacks.event_publish_callback if callbacks else None - _session_id_for_progress = ctx.session.id - progress_callback = None - if _publish is not None: - async def progress_callback(stage: str, data: dict) -> None: - await _publish("session.compaction_progress", { - "sessionID": _session_id_for_progress, - "stage": stage, - "data": data, - }) - - # Process compaction - try: - compaction_result = await run_compaction( - ctx.session.id, - parent_message_id=last_user.id, - messages=messages, - provider_id=ctx.provider_id, - model_id=ctx.model_id, - auto=getattr(task_part, 'auto', False), - event_publish_callback=_publish, - status_after="busy", - policy=compaction_policy, - progress_callback=progress_callback, - ) - - if compaction_result == "stop": - log.error("loop.compaction_failed", {"session_id": ctx.session.id}) - if callbacks.on_error: - await callbacks.on_error("Compaction failed") - break - - if compaction_result == "skipped": - log.info("loop.manual_compaction_skipped", { - "session_id": ctx.session.id, - "step": ctx.step, - }) - - # Continue after compaction (whether compacted or skipped) - continue - - except Exception as e: - log.error("loop.compaction_error", {"error": str(e)}) - if callbacks.on_error: - await callbacks.on_error(f"Compaction error: {str(e)}") - break - - # ---------------------------------------------------------------- - # Context overflow detection & recovery - # - # Matches OpenClaw run.ts overflow recovery cascade: - # 1. Detect overflow - # 2. Try tool result truncation (once per run) - # 3. Full compaction (up to MAX_OVERFLOW_COMPACTION_ATTEMPTS) - # 4. Give up with error if still overflowing - # ---------------------------------------------------------------- - if last_finished and not getattr(last_finished, 'summary', False): - # Get model context limit from flocks.json / provider registry - model_context, model_output, model_input = Provider.resolve_model_info( - ctx.provider_id, ctx.model_id - ) - - # Check for overflow using dynamic CompactionPolicy - if model_context > 0: - compaction_policy = CompactionPolicy.from_model( - context_window=model_context, - max_output_tokens=model_output or 4096, - max_input_tokens=model_input, - ) - - # Build tokens_dict from last_finished.tokens if available. - # last_finished.tokens may be a TokenUsage Pydantic model (not a - # plain dict), so we normalise it to a dict here to ensure the - # provider-reported usage is actually read instead of silently - # falling through to the chars/4 estimation path. - tokens_dict = {} - if hasattr(last_finished, 'tokens') and last_finished.tokens: - raw_tok = last_finished.tokens - if isinstance(raw_tok, dict): - tokens_dict = raw_tok - elif hasattr(raw_tok, 'model_dump'): - tokens_dict = raw_tok.model_dump() - elif hasattr(raw_tok, '__dict__'): - tokens_dict = vars(raw_tok) - - # Check if provider returned actual usage data (not all zeros) - input_tokens = tokens_dict.get("input", 0) - _cache = tokens_dict.get("cache") or {} - cache_read = _cache.get("read", 0) if isinstance(_cache, dict) else 0 - output_tokens = tokens_dict.get("output", 0) - reported_total = input_tokens + cache_read + output_tokens - - # B3 — Observed-value-first token decision. Always prefer - # the provider's actual usage figure (input + cache_read) - # over our synthetic estimate, because that is what the - # next turn's prompt will be billed against. Cache it on - # the LoopContext so subsequent turns can reuse it as a - # baseline. Estimation only kicks in when the provider - # genuinely reports no usage (all zero) — we feed the - # compaction policy into ``estimate_full_context_tokens`` - # so it includes system-prompt + tool-schema overhead - # and applies the 1.2 safety margin. - if reported_total > 0: - ctx.last_observed_prompt_tokens = input_tokens + cache_read + output_tokens - log.info("loop.tokens_decision", { - "session_id": ctx.session.id, - "source": "observed", - "effective_tokens": input_tokens + cache_read, - "overflow_threshold": compaction_policy.overflow_threshold, - }) - else: - estimated_tokens = await SessionPrompt.estimate_full_context_tokens( - ctx.session.id, - messages, - policy=compaction_policy, - ) - tokens_dict = {"input": estimated_tokens, "output": 0, "cache": {"read": 0, "write": 0}} - log.info("loop.tokens_decision", { - "session_id": ctx.session.id, - "source": "estimated", - "effective_tokens": estimated_tokens, - "message_count": len(messages), - "overflow_threshold": compaction_policy.overflow_threshold, - }) - - try: - _tok_cache = tokens_dict.get("cache") or {} - current_input_tokens = ( - tokens_dict.get("input", 0) - + (_tok_cache.get("read", 0) if isinstance(_tok_cache, dict) else 0) - ) - recent_compaction = cls._has_recent_compaction_cooldown(ctx) - near_overflow = current_input_tokens >= compaction_policy.preemptive_threshold - - if near_overflow and ctx.last_cleanup_step != ctx.step: - try: - trunc_count = await SessionCompaction.truncate_oversized_tool_outputs( - ctx.session.id, - context_window_tokens=model_context, - ) - ctx.last_cleanup_step = ctx.step - if trunc_count > 0: - set_context_state( - ctx.session.id, - tool_results_compacted=True, - last_compaction_step=ctx.last_compaction_step, - last_compaction_reason="pre_compact_cleanup", - ) - await cls._publish_runtime_event(callbacks, "context.compacted", { - "sessionID": ctx.session.id, - "step": ctx.step, - "reason": "pre_compact_cleanup", - "truncatedToolResults": trunc_count, - "cooldownActive": recent_compaction, - }) - log.info("loop.pre_compact_cleanup_applied", { - "session_id": ctx.session.id, - "step": ctx.step, - "truncated": trunc_count, - "preemptive_threshold": compaction_policy.preemptive_threshold, - "input_tokens": current_input_tokens, - "cooldown_active": recent_compaction, - }) - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="continued", - continue_reason="pre_compact_cleanup", - queued_message_detected=False, - ) - await cls._publish_runtime_event( - callbacks, - "turn.continued", - turn_state.model_dump(by_alias=True), - ) - continue - except Exception as trunc_err: - log.warn("loop.pre_compact_cleanup_error", { - "session_id": ctx.session.id, - "error": str(trunc_err), - }) - - is_overflow = await SessionCompaction.is_overflow( - tokens=tokens_dict, - model_context=model_context, - policy=compaction_policy, - ) - - if is_overflow: - log.info("loop.context_overflow_detected", { - "session_id": ctx.session.id, - "step": ctx.step, - "tokens": tokens_dict, - "tier": compaction_policy.tier.value, - "overflow_compaction_attempts": ctx.overflow_compaction_attempts, - }) - - # Check if we've exhausted compaction attempts - # (matches OpenClaw MAX_OVERFLOW_COMPACTION_ATTEMPTS) - if ctx.overflow_compaction_attempts >= MAX_OVERFLOW_COMPACTION_ATTEMPTS: - # Distinguish "provider down / in cooldown" from - # "context genuinely too large" so users get - # actionable advice instead of a generic error. - compaction_hist = _get_compaction_history(ctx.session.id) - _provider_error = compaction_hist.summary_last_error - _in_cooldown = ( - compaction_hist.summary_cooldown_until > 0 - and compaction_hist.summary_cooldown_until - > time.monotonic() - ) - _cooldown_secs = max( - 0, - round(compaction_hist.summary_cooldown_until - - time.monotonic()), - ) - - if _in_cooldown or _provider_error: - # Provider-side issue: cooldown still active - # or last call recorded an error. Tell the - # user to wait / retry rather than open a new - # session (their context is fine). - _notice_msg = ( - "摘要模型暂时不可用，上下文压缩跳过了本轮压缩。" - + ( - f"冷却剩余约 {_cooldown_secs} 秒，" - if _in_cooldown else "" - ) - + "建议稍后继续，或切换到其他模型重试。" - ) - _error_msg = ( - "Compaction skipped: summary provider unavailable " - f"({_provider_error or 'cooldown active'})." - + ( - f" Cooldown expires in ~{_cooldown_secs}s." - if _in_cooldown else "" - ) - + " Wait for the provider to recover or switch models." - ) - else: - # Context is genuinely too large even after - # repeated compaction — advise reducing scope. - _notice_msg = ( - "当前任务上下文过重，已经多次 compact 仍接近上限。" - "建议收敛工具输出、缩小搜索范围，或开启新会话。" - ) - _error_msg = ( - "Context overflow: prompt too large for the model after " - f"{ctx.overflow_compaction_attempts} compaction attempts. " - "Try starting a new session or use a larger-context model." - ) - - await cls._publish_session_notice( - callbacks, - ctx.session.id, - level="warning", - message=_notice_msg, - details={ - "attempts": ctx.overflow_compaction_attempts, - "maxAttempts": MAX_OVERFLOW_COMPACTION_ATTEMPTS, - "tokens": tokens_dict, - "providerError": _provider_error or None, - "cooldownRemainingSeconds": ( - _cooldown_secs if _in_cooldown else 0 - ), - }, - ) - log.error("loop.overflow_compaction_exhausted", { - "session_id": ctx.session.id, - "attempts": ctx.overflow_compaction_attempts, - "max": MAX_OVERFLOW_COMPACTION_ATTEMPTS, - "tokens": tokens_dict, - "in_cooldown": _in_cooldown, - "provider_error": _provider_error or None, - }) - if callbacks.on_error: - await callbacks.on_error(_error_msg) - break - - # Recovery step 1: try truncating oversized tool - # results (once per run, matches OpenClaw - # toolResultTruncationAttempted) - if not ctx.tool_result_truncation_attempted: - ctx.tool_result_truncation_attempted = True - try: - trunc_count = await SessionCompaction.truncate_oversized_tool_outputs( - ctx.session.id, - context_window_tokens=model_context, - ) - if trunc_count > 0: - log.info("loop.oversized_tool_truncated", { - "session_id": ctx.session.id, - "truncated": trunc_count, - }) - # Re-check overflow after truncation - # (B3) Reuse the active v2 policy so the - # re-estimate includes the same overhead - # + safety margin as the first decision. - re_est = await SessionPrompt.estimate_full_context_tokens( - ctx.session.id, - messages, - policy=compaction_policy, - ) - re_tokens = {"input": re_est, "output": 0, "cache": {"read": 0, "write": 0}} - still_overflow = await SessionCompaction.is_overflow( - tokens=re_tokens, - model_context=model_context, - policy=compaction_policy, - ) - if not still_overflow: - log.info("loop.overflow_resolved_by_truncation", { - "session_id": ctx.session.id, - }) - # Do NOT reset overflow_compaction_attempts - # (matches OpenClaw OC-65) - continue - except Exception as trunc_err: - log.warn("loop.oversized_truncation_error", { - "session_id": ctx.session.id, - "error": str(trunc_err), - }) - - # Recovery step 2: full compaction - ctx.overflow_compaction_attempts += 1 - if ctx.overflow_compaction_attempts >= 2: - await cls._publish_session_notice( - callbacks, - ctx.session.id, - level="info", - message=( - "本轮上下文持续接近模型上限，系统将优先尝试压缩历史工具输出。" - ), - details={ - "attempt": ctx.overflow_compaction_attempts, - "threshold": compaction_policy.overflow_threshold, - "buffer": compaction_policy.overflow_buffer, - }, - ) - log.warn("loop.overflow_compaction_attempt", { - "session_id": ctx.session.id, - "attempt": ctx.overflow_compaction_attempts, - "max": MAX_OVERFLOW_COMPACTION_ATTEMPTS, - }) - - # --- Compaction start: notify all UIs --- - if callbacks.on_compaction: - await callbacks.on_compaction() - - # Prune first, then summarize - await SessionCompaction.prune( - ctx.session.id, - policy=compaction_policy, - ) - - # Same SSE progress adapter as the manual - # /compact route — mirrored here so the - # overflow-driven path also drives the - # multi-stage UI panel. - _publish_overflow = callbacks.event_publish_callback if callbacks else None - _session_id_overflow = ctx.session.id - progress_callback_overflow = None - if _publish_overflow is not None: - async def progress_callback_overflow(stage: str, data: dict) -> None: - await _publish_overflow("session.compaction_progress", { - "sessionID": _session_id_overflow, - "stage": stage, - "data": data, - }) - - # Trigger compaction (summarization + memory flush) - compaction_result = await run_compaction( - ctx.session.id, - parent_message_id=last_user.id, - messages=messages, - provider_id=ctx.provider_id, - model_id=ctx.model_id, - auto=True, - event_publish_callback=_publish_overflow, - status_after="busy", - policy=compaction_policy, - progress_callback=progress_callback_overflow, - ) + pending_user_input = False + try: + from flocks.server.routes.question import has_pending_questions - if compaction_result == "stop": - log.error("loop.compaction_failed", {"session_id": ctx.session.id}) - if callbacks.on_error: - await callbacks.on_error("Compaction failed") - break - - if compaction_result == "skipped": - # Anti-thrashing cooldown or summary-provider - # cooldown fired — nothing was archived, do NOT - # update last_compaction_step or publish the - # compacted event (would mislead cooldown logic - # and UI into thinking compaction succeeded). - log.info("loop.compaction_skipped", { - "session_id": ctx.session.id, - "step": ctx.step, - }) - else: - # compaction_result == "continue": real success - ctx.last_compaction_step = ctx.step - set_context_state( - ctx.session.id, - compaction_performed=True, - last_compaction_step=ctx.step, - last_compaction_reason="full_compaction", - ) - await cls._publish_runtime_event(callbacks, "context.compacted", { - "sessionID": ctx.session.id, - "step": ctx.step, - "reason": "full_compaction", - "attempt": ctx.overflow_compaction_attempts, - "cooldownUntilStep": ctx.step + POST_COMPACTION_COOLDOWN_STEPS, - }) - - # Continuation user message is now created inside - # SessionCompaction.process() (matching Flocks). - # Just continue — the new user message flips the - # ID ordering so _should_exit() won't trigger. - continue - except Exception as e: - log.error("loop.compaction_overflow_check_error", {"error": str(e)}) - - # Process single step — wrap in a Task so abort() can cancel it immediately - # rather than waiting for the current tool call to finish. - step_task = asyncio.create_task( - cls._process_step_with_failover( - ctx, - callbacks, - messages, - last_user, - ) + pending_user_input = has_pending_questions(ctx.session.id) + except Exception as exc: + log.warn( + "goal.pending_question_check.error", + {"session_id": ctx.session.id, "error": str(exc)}, ) - ctx._current_step_task = step_task - step_started_at = asyncio.get_event_loop().time() - try: - step_result = await step_task - except asyncio.CancelledError: - log.info("loop.step_cancelled", {"session_id": ctx.session.id, "step": ctx.step}) - break - finally: - ctx._current_step_task = None - log.debug("loop.step_complete", { - "session_id": ctx.session.id, - "step": ctx.step, - "duration_ms": int((asyncio.get_event_loop().time() - step_started_at) * 1000), - }) - - # Callback: step end - if callbacks.on_step_end: - await callbacks.on_step_end(ctx.step) - - # Handle result - if step_result.action == "stop": - loop_error = step_result.error - # Report error if step failed - if step_result.error and callbacks.on_error: - await callbacks.on_error(step_result.error) - - # Get last assistant message via SessionContext - if ctx.session_ctx: - post_messages = await ctx.session_ctx.get_messages() - else: - post_messages = await Message.list(ctx.session.id) - for msg in reversed(post_messages): - if ( - msg.role == MessageRole.ASSISTANT - and ( - not ctx.auto_failover - or getattr(msg, "parentID", None) == last_user.id - ) - ): - last_message = msg - break - queued_user = await cls._detect_queued_user_message( - ctx.session.id, - post_messages, - last_user.id, - last_message, - ) - if queued_user is not None: - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="continued", - continue_reason="queued_message", - queued_message_detected=True, - ) - await cls._publish_runtime_event(callbacks, "turn.continued", { - **turn_state.model_dump(by_alias=True), - "queuedUserMessageID": queued_user.id, - }) - log.info("loop.continuing_for_queued_message", { - "session_id": ctx.session.id, - "queued_user_id": queued_user.id, - "last_assistant_id": last_message.id if last_message else None, - }) - continue + goal_decision = await GoalManager.evaluate_after_turn( + ctx.session.id, + str(last_response or ""), + pending_user_input=pending_user_input, + provider_id=ctx.provider_id, + model_id=ctx.model_id, + ) + if ( + goal_decision.status in {"completed", "blocked", "paused"} + and goal_decision.objective + ): + await cls._publish_runtime_event( + callbacks, + "session.goal.updated", + { + "sessionID": ctx.session.id, + "status": goal_decision.status, + "objective": goal_decision.objective, + "reason": goal_decision.reason, + }, + ) + if goal_decision.should_continue and goal_decision.continuation_prompt: + goal_user = await Message.create( + session_id=ctx.session.id, + role=MessageRole.USER, + content=goal_decision.continuation_prompt, + agent=( + last_user.agent + if hasattr(last_user, "agent") + else ctx.agent_name + ), + model=( + last_user.model + if hasattr(last_user, "model") + else { + "providerID": ctx.provider_id, + "modelID": ctx.model_id, + } + ), + provider=( + last_user.provider + if hasattr(last_user, "provider") + else ctx.provider_id + ), + synthetic=True, + part_metadata={ + "goalContinuation": True, + "goalVerdict": goal_decision.verdict, + "goalReason": goal_decision.reason, + }, + ) + turn_state = set_turn_state( + ctx.session.id, + step=ctx.step, + status="continued", + continue_reason="goal", + queued_message_detected=False, + ) + await cls._publish_runtime_event( + callbacks, + "turn.continued", + { + **turn_state.model_dump(by_alias=True), + "goalMessageID": goal_user.id, + "goalVerdict": goal_decision.verdict, + }, + ) + log.info( + "loop.continuing_for_goal", + { + "session_id": ctx.session.id, + "goal_message_id": goal_user.id, + "reason": goal_decision.reason, + }, + ) + return ContinuationDecision( + messages=(goal_user,), + reason="goal", + ) + + if ( + not ctx.should_abort() + and getattr(last_message, "finish", None) == "stop" + and await cls._run_turn_finish_hook( + ctx, + callbacks, + last_user, + last_message, + ) + ): + if ctx.session_ctx: + post_hook_messages = await ctx.session_ctx.get_messages() + else: + post_hook_messages = await Message.list(ctx.session.id) + existing_ids = {message.id for message in state.messages} + new_messages = tuple( + message + for message in post_hook_messages + if message.id not in existing_ids + ) + return ContinuationDecision( + messages=new_messages, + reason="turn_finish_hook", + ) + + stop_reason = getattr(last_message, "finish", None) or "stop" + await cls._publish_turn_stopped( + callbacks, + ctx.session.id, + step=ctx.step, + stop_reason=stop_reason, + ) + return ContinuationDecision() + + @classmethod + async def _prepare_model_turn( + cls, + ctx: LoopContext, + callbacks: LoopCallbacks, + state: AgentRunState[MessageInfo], + ) -> ModelTurnPreparation[MessageInfo]: + """Prepare one immutable model-turn snapshot from session state.""" + SessionStatus.set(ctx.session.id, SessionStatusBusy()) + ctx.step += 1 + state.model_turn_index = ctx.step + turn_state = set_turn_state( + ctx.session.id, + step=ctx.step, + status="started", + queued_message_detected=False, + ) + await cls._publish_runtime_event( + callbacks, + "turn.started", + turn_state.model_dump(by_alias=True), + ) + log.info( + "loop.step", + {"session_id": ctx.session.id, "step": ctx.step}, + ) + if callbacks.on_step_start: + await callbacks.on_step_start(ctx.step) + + messages_started_at = asyncio.get_running_loop().time() + if ctx.session_ctx: + messages = await ctx.session_ctx.get_messages() + else: + messages = await Message.list(ctx.session.id) + log.debug( + "loop.messages_loaded", + { + "session_id": ctx.session.id, + "step": ctx.step, + "message_count": len(messages), + "duration_ms": int( + (asyncio.get_running_loop().time() - messages_started_at) + * 1000 + ), + }, + ) + if not messages: + log.info("loop.no_messages", {"session_id": ctx.session.id}) + await cls._publish_turn_stopped( + callbacks, + ctx.session.id, + step=ctx.step, + stop_reason="no_messages", + ) + return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) + + last_user: Optional[MessageInfo] = None + last_assistant: Optional[MessageInfo] = None + last_finished: Optional[MessageInfo] = None + tasks: List[tuple[str, Any]] = [] + scan_started_at = asyncio.get_running_loop().time() + for message in reversed(messages): + if last_user is None and message.role == MessageRole.USER: + last_user = message + if last_assistant is None and message.role == MessageRole.ASSISTANT: + last_assistant = message + if ( + last_finished is None + and message.role == MessageRole.ASSISTANT + and getattr(message, "finish", None) + ): + last_finished = message + if last_user is not None and last_finished is not None: + break + if last_finished is None: + for part in await Message.parts(message.id, ctx.session.id): + if part.type == "compaction": + tasks.append(("compaction", part)) + elif part.type == "subtask": + tasks.append(("subtask", part)) + log.debug( + "loop.message_scan_complete", + { + "session_id": ctx.session.id, + "step": ctx.step, + "task_count": len(tasks), + "duration_ms": int( + (asyncio.get_running_loop().time() - scan_started_at) * 1000 + ), + }, + ) + + if last_user is None: + log.info( + "loop.no_user_message", + { + "session_id": ctx.session.id, + "message_count": len(messages), + "roles": [ + str(getattr(message, "role", "")) + for message in messages[-5:] + ], + }, + ) + await cls._publish_turn_stopped( + callbacks, + ctx.session.id, + step=ctx.step, + stop_reason="no_user_message", + ) + return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) + + last_assistant_parts = ( + await Message.parts(last_assistant.id, ctx.session.id) + if last_assistant + else [] + ) + if cls._should_exit(last_user, last_assistant, last_assistant_parts): + log.info( + "loop.exit_condition", + { + "session_id": ctx.session.id, + "last_user_id": last_user.id, + "last_assistant_id": ( + last_assistant.id if last_assistant else None + ), + "finish": last_assistant.finish if last_assistant else None, + "has_tool_parts": any( + getattr(part, "type", None) == "tool" + for part in last_assistant_parts + ), + }, + ) + return ModelTurnPreparation( + status=TurnPreparationStatus.COMPLETE, + last_message=last_assistant, + ) + + if await cls._prepare_auto_turn(ctx, last_user): + ctx.turn_additional_context = None + ctx.stop_hook_active = False + await cls._run_user_prompt_submit_hook(ctx, last_user) + + state.current_user_id = last_user.id + state.metadata["last_user"] = last_user + await cls._prepare_memory(ctx) + cls._schedule_title_generation(ctx, callbacks, last_user, messages) + + if tasks: + task_preparation = await cls._prepare_pending_task( + ctx, + callbacks, + messages, + last_user, + tasks.pop(), + ) + if task_preparation is not None: + return task_preparation + + context_preparation = await cls._prepare_context_window( + ctx, + callbacks, + messages, + last_user, + last_finished, + ) + if context_preparation is not None: + return context_preparation + + active_model = RuntimeModel(ctx.provider_id, ctx.model_id) + state.active_model = active_model + state.messages = list(messages) + return ModelTurnPreparation( + status=TurnPreparationStatus.READY, + snapshot=ModelTurnSnapshot( + session_id=ctx.session.id, + agent_name=ctx.agent_name, + active_model=active_model, + model_turn_index=ctx.step, + trace_step=ctx.trace_step, + messages=tuple(messages), + last_user=last_user, + ), + ) + + @staticmethod + async def _prepare_memory(ctx: LoopContext) -> None: + """Load memory once before the first model turn.""" + if ( + ctx.step != 1 + or not ctx.session.memory_enabled + or ctx.memory_bootstrap_data is not None + ): + return + try: + from flocks.memory.bootstrap import MemoryBootstrap + + ctx.memory_bootstrap_data = await MemoryBootstrap( + project_id=ctx.session.project_id, + ).bootstrap(load_daily=False) + log.info( + "loop.memory_bootstrap_done", + { + "session_id": ctx.session.id, + "has_main": ( + ctx.memory_bootstrap_data.get("main_memory") is not None + ), + }, + ) + except Exception as exc: + log.error("loop.memory_bootstrap_error", {"error": str(exc)}) + + @staticmethod + def _schedule_title_generation( + ctx: LoopContext, + callbacks: LoopCallbacks, + last_user: MessageInfo, + messages: List[MessageInfo], + ) -> None: + """Start optimistic first-turn title generation without blocking.""" + if ctx.step != 1 or ctx.auto_failover: + return + try: + from flocks.session.lifecycle.title import SessionTitle + + user_model = getattr(last_user, "model", None) + if isinstance(user_model, dict): + title_model_id = user_model.get("modelID", ctx.model_id) + title_provider_id = user_model.get( + "providerID", + ctx.provider_id, + ) + else: + title_model_id = ctx.model_id + title_provider_id = ctx.provider_id + fire_and_forget( + SessionTitle.ensure_title( + session_id=ctx.session.id, + model_id=title_model_id, + provider_id=title_provider_id, + messages=messages, + event_publish_callback=callbacks.event_publish_callback, + ), + label="title_generation", + name=f"title:{ctx.session.id}", + ) + except Exception as exc: + log.error("loop.title_generation.error", {"error": str(exc)}) + + @classmethod + async def _prepare_pending_task( + cls, + ctx: LoopContext, + callbacks: LoopCallbacks, + messages: List[MessageInfo], + last_user: MessageInfo, + task: tuple[str, Any], + ) -> Optional[ModelTurnPreparation[MessageInfo]]: + """Finish persisted subtask or compaction work before the model turn.""" + task_type, task_part = task + if task_type == "subtask": + log.info( + "loop.subtask_detected", + {"session_id": ctx.session.id, "step": ctx.step}, + ) + await cls._execute_subtask(ctx, last_user, task_part) + return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) - if not step_result.error and last_message is not None: - try: - content_result = Message.get_text_content(last_message) - last_response = ( - await content_result - if inspect.isawaitable(content_result) - else content_result + log.info( + "loop.compaction_pending", + { + "session_id": ctx.session.id, + "step": ctx.step, + "auto": getattr(task_part, "auto", False), + }, + ) + if callbacks.on_compaction: + await callbacks.on_compaction() + + publish = callbacks.event_publish_callback + progress_callback = None + if publish is not None: + + async def progress_callback(stage: str, data: dict) -> None: + await publish( + "session.compaction_progress", + { + "sessionID": ctx.session.id, + "stage": stage, + "data": data, + }, + ) + + try: + compaction_result = await run_compaction( + ctx.session.id, + parent_message_id=last_user.id, + messages=messages, + provider_id=ctx.provider_id, + model_id=ctx.model_id, + auto=getattr(task_part, "auto", False), + event_publish_callback=publish, + status_after="busy", + policy=cls._build_compaction_policy(ctx), + progress_callback=progress_callback, + ) + if compaction_result == "stop": + log.error( + "loop.compaction_failed", + {"session_id": ctx.session.id}, + ) + if callbacks.on_error: + await callbacks.on_error("Compaction failed") + return ModelTurnPreparation( + status=TurnPreparationStatus.COMPLETE, + ) + if compaction_result == "skipped": + log.info( + "loop.manual_compaction_skipped", + {"session_id": ctx.session.id, "step": ctx.step}, + ) + return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) + except Exception as exc: + log.error("loop.compaction_error", {"error": str(exc)}) + if callbacks.on_error: + await callbacks.on_error(f"Compaction error: {exc}") + return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) + + @classmethod + async def _prepare_context_window( + cls, + ctx: LoopContext, + callbacks: LoopCallbacks, + messages: List[MessageInfo], + last_user: MessageInfo, + last_finished: Optional[MessageInfo], + ) -> Optional[ModelTurnPreparation[MessageInfo]]: + """Recover a near-overflow context before the next model turn.""" + if last_finished is None or getattr(last_finished, "summary", False): + return None + + model_context, model_output, model_input = Provider.resolve_model_info( + ctx.provider_id, + ctx.model_id, + ) + if model_context <= 0: + return None + + policy = CompactionPolicy.from_model( + context_window=model_context, + max_output_tokens=model_output or 4096, + max_input_tokens=model_input, + ) + tokens = cls._normalise_token_usage(last_finished) + input_tokens = tokens.get("input", 0) + cache = tokens.get("cache") or {} + cache_read = cache.get("read", 0) if isinstance(cache, dict) else 0 + output_tokens = tokens.get("output", 0) + reported_total = input_tokens + cache_read + output_tokens + if reported_total > 0: + ctx.last_observed_prompt_tokens = reported_total + log.info( + "loop.tokens_decision", + { + "session_id": ctx.session.id, + "source": "observed", + "effective_tokens": input_tokens + cache_read, + "overflow_threshold": policy.overflow_threshold, + }, + ) + else: + estimated_tokens = await SessionPrompt.estimate_full_context_tokens( + ctx.session.id, + messages, + policy=policy, + ) + tokens = { + "input": estimated_tokens, + "output": 0, + "cache": {"read": 0, "write": 0}, + } + log.info( + "loop.tokens_decision", + { + "session_id": ctx.session.id, + "source": "estimated", + "effective_tokens": estimated_tokens, + "message_count": len(messages), + "overflow_threshold": policy.overflow_threshold, + }, + ) + + try: + cache = tokens.get("cache") or {} + current_input_tokens = tokens.get("input", 0) + ( + cache.get("read", 0) if isinstance(cache, dict) else 0 + ) + recent_compaction = cls._has_recent_compaction_cooldown(ctx) + near_overflow = current_input_tokens >= policy.preemptive_threshold + if near_overflow and ctx.last_cleanup_step != ctx.step: + cleanup_result = await cls._prepare_tool_result_cleanup( + ctx, + callbacks, + model_context, + policy, + current_input_tokens, + recent_compaction, + ) + if cleanup_result is not None: + return cleanup_result + + is_overflow = await SessionCompaction.is_overflow( + tokens=tokens, + model_context=model_context, + policy=policy, + ) + if not is_overflow: + return None + + log.info( + "loop.context_overflow_detected", + { + "session_id": ctx.session.id, + "step": ctx.step, + "tokens": tokens, + "tier": policy.tier.value, + "overflow_compaction_attempts": ( + ctx.overflow_compaction_attempts + ), + }, + ) + if ( + ctx.overflow_compaction_attempts + >= MAX_OVERFLOW_COMPACTION_ATTEMPTS + ): + await cls._report_compaction_exhausted( + ctx, + callbacks, + tokens, + ) + return ModelTurnPreparation( + status=TurnPreparationStatus.COMPLETE, + ) + + if not ctx.tool_result_truncation_attempted: + ctx.tool_result_truncation_attempted = True + try: + truncation_count = ( + await SessionCompaction.truncate_oversized_tool_outputs( + ctx.session.id, + context_window_tokens=model_context, ) - except Exception as exc: - log.warn("goal.last_response.error", { - "session_id": ctx.session.id, - "message_id": getattr(last_message, "id", None), - "error": str(exc), - }) - last_response = getattr(last_message, "content", "") or "" - pending_user_input = False - try: - from flocks.server.routes.question import has_pending_questions - - pending_user_input = has_pending_questions(ctx.session.id) - except Exception as exc: - log.warn("goal.pending_question_check.error", { - "session_id": ctx.session.id, - "error": str(exc), - }) - goal_decision = await GoalManager.evaluate_after_turn( - ctx.session.id, - str(last_response or ""), - pending_user_input=pending_user_input, - provider_id=ctx.provider_id, - model_id=ctx.model_id, ) - if goal_decision.status in {"completed", "blocked", "paused"} and goal_decision.objective: - await cls._publish_runtime_event(callbacks, "session.goal.updated", { - "sessionID": ctx.session.id, - "status": goal_decision.status, - "objective": goal_decision.objective, - "reason": goal_decision.reason, - }) - if goal_decision.should_continue and goal_decision.continuation_prompt: - # Hermes-style goal continuation: append a user-role - # prompt to history so the model continues, while - # marking the part synthetic so UIs do not treat it as - # user-authored text. - goal_user = await Message.create( - session_id=ctx.session.id, - role=MessageRole.USER, - content=goal_decision.continuation_prompt, - agent=last_user.agent if hasattr(last_user, "agent") else ctx.agent_name, - model=last_user.model if hasattr(last_user, "model") else { - "providerID": ctx.provider_id, - "modelID": ctx.model_id, - }, - provider=last_user.provider if hasattr(last_user, "provider") else ctx.provider_id, - synthetic=True, - part_metadata={ - "goalContinuation": True, - "goalVerdict": goal_decision.verdict, - "goalReason": goal_decision.reason, + if truncation_count > 0: + log.info( + "loop.oversized_tool_truncated", + { + "session_id": ctx.session.id, + "truncated": truncation_count, }, ) - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="continued", - continue_reason="goal", - queued_message_detected=False, + estimated_tokens = ( + await SessionPrompt.estimate_full_context_tokens( + ctx.session.id, + messages, + policy=policy, + ) + ) + still_overflow = await SessionCompaction.is_overflow( + tokens={ + "input": estimated_tokens, + "output": 0, + "cache": {"read": 0, "write": 0}, + }, + model_context=model_context, + policy=policy, ) - await cls._publish_runtime_event(callbacks, "turn.continued", { - **turn_state.model_dump(by_alias=True), - "goalMessageID": goal_user.id, - "goalVerdict": goal_decision.verdict, - }) - log.info("loop.continuing_for_goal", { - "session_id": ctx.session.id, - "goal_message_id": goal_user.id, - "reason": goal_decision.reason, - }) - continue - - if ( - not step_result.error - and not ctx.should_abort() - and last_message is not None - and getattr(last_message, "finish", None) == "stop" - and await cls._run_turn_finish_hook( - ctx, - callbacks, - last_user, - last_message, + if not still_overflow: + log.info( + "loop.overflow_resolved_by_truncation", + {"session_id": ctx.session.id}, + ) + return ModelTurnPreparation( + status=TurnPreparationStatus.CONTINUE, + ) + except Exception as exc: + log.warn( + "loop.oversized_truncation_error", + {"session_id": ctx.session.id, "error": str(exc)}, ) - ): - continue - stop_reason = step_result.error or (getattr(last_message, "finish", None) if last_message else None) or "stop" - turn_state = set_turn_state( + return await cls._prepare_full_compaction( + ctx, + callbacks, + messages, + last_user, + policy, + ) + except Exception as exc: + log.error( + "loop.compaction_overflow_check_error", + {"error": str(exc)}, + ) + return None + + @staticmethod + def _normalise_token_usage(message: MessageInfo) -> Dict[str, Any]: + """Normalise provider token usage into the legacy mapping shape.""" + raw_tokens = getattr(message, "tokens", None) + if not raw_tokens: + return {} + if isinstance(raw_tokens, dict): + return raw_tokens + if hasattr(raw_tokens, "model_dump"): + return raw_tokens.model_dump() + if hasattr(raw_tokens, "__dict__"): + return vars(raw_tokens) + return {} + + @classmethod + async def _prepare_tool_result_cleanup( + cls, + ctx: LoopContext, + callbacks: LoopCallbacks, + model_context: int, + policy: CompactionPolicy, + current_input_tokens: int, + recent_compaction: bool, + ) -> Optional[ModelTurnPreparation[MessageInfo]]: + """Apply the cheap tool-result cleanup before full compaction.""" + try: + truncation_count = ( + await SessionCompaction.truncate_oversized_tool_outputs( ctx.session.id, - step=ctx.step, - status="stopped", - stop_reason=stop_reason, - queued_message_detected=False, + context_window_tokens=model_context, ) - await cls._publish_runtime_event(callbacks, "turn.stopped", turn_state.model_dump(by_alias=True)) + ) + ctx.last_cleanup_step = ctx.step + if truncation_count <= 0: + return None - break - - elif step_result.action == "continue": - if ctx.session_ctx: - post_messages = await ctx.session_ctx.get_messages() - else: - post_messages = await Message.list(ctx.session.id) - last_assistant_after_step = next( - ( - msg for msg in reversed(post_messages) - if msg.role == MessageRole.ASSISTANT - ), - None, + set_context_state( + ctx.session.id, + tool_results_compacted=True, + last_compaction_step=ctx.last_compaction_step, + last_compaction_reason="pre_compact_cleanup", + ) + await cls._publish_runtime_event( + callbacks, + "context.compacted", + { + "sessionID": ctx.session.id, + "step": ctx.step, + "reason": "pre_compact_cleanup", + "truncatedToolResults": truncation_count, + "cooldownActive": recent_compaction, + }, + ) + log.info( + "loop.pre_compact_cleanup_applied", + { + "session_id": ctx.session.id, + "step": ctx.step, + "truncated": truncation_count, + "preemptive_threshold": policy.preemptive_threshold, + "input_tokens": current_input_tokens, + "cooldown_active": recent_compaction, + }, + ) + turn_state = set_turn_state( + ctx.session.id, + step=ctx.step, + status="continued", + continue_reason="pre_compact_cleanup", + queued_message_detected=False, + ) + await cls._publish_runtime_event( + callbacks, + "turn.continued", + turn_state.model_dump(by_alias=True), + ) + return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) + except Exception as exc: + log.warn( + "loop.pre_compact_cleanup_error", + {"session_id": ctx.session.id, "error": str(exc)}, + ) + return None + + @classmethod + async def _report_compaction_exhausted( + cls, + ctx: LoopContext, + callbacks: LoopCallbacks, + tokens: Dict[str, Any], + ) -> None: + """Surface whether exhaustion came from context or provider health.""" + history = _get_compaction_history(ctx.session.id) + provider_error = history.summary_last_error + in_cooldown = ( + history.summary_cooldown_until > 0 + and history.summary_cooldown_until > time.monotonic() + ) + cooldown_seconds = max( + 0, + round(history.summary_cooldown_until - time.monotonic()), + ) + if in_cooldown or provider_error: + notice = ( + "摘要模型暂时不可用，上下文压缩跳过了本轮压缩。" + + ( + f"冷却剩余约 {cooldown_seconds} 秒，" + if in_cooldown + else "" ) - queued_user = await cls._detect_queued_user_message( - ctx.session.id, - post_messages, - last_user.id, - last_assistant_after_step, + + "建议稍后继续，或切换到其他模型重试。" + ) + error = ( + "Compaction skipped: summary provider unavailable " + f"({provider_error or 'cooldown active'})." + + ( + f" Cooldown expires in ~{cooldown_seconds}s." + if in_cooldown + else "" ) - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="continued", - continue_reason="queued_message" if queued_user is not None else "tool_calls", - queued_message_detected=queued_user is not None, + + " Wait for the provider to recover or switch models." + ) + else: + notice = ( + "当前任务上下文过重，已经多次 compact 仍接近上限。" + "建议收敛工具输出、缩小搜索范围，或开启新会话。" + ) + error = ( + "Context overflow: prompt too large for the model after " + f"{ctx.overflow_compaction_attempts} compaction attempts. " + "Try starting a new session or use a larger-context model." + ) + + await cls._publish_session_notice( + callbacks, + ctx.session.id, + level="warning", + message=notice, + details={ + "attempts": ctx.overflow_compaction_attempts, + "maxAttempts": MAX_OVERFLOW_COMPACTION_ATTEMPTS, + "tokens": tokens, + "providerError": provider_error or None, + "cooldownRemainingSeconds": ( + cooldown_seconds if in_cooldown else 0 + ), + }, + ) + log.error( + "loop.overflow_compaction_exhausted", + { + "session_id": ctx.session.id, + "attempts": ctx.overflow_compaction_attempts, + "max": MAX_OVERFLOW_COMPACTION_ATTEMPTS, + "tokens": tokens, + "in_cooldown": in_cooldown, + "provider_error": provider_error or None, + }, + ) + if callbacks.on_error: + await callbacks.on_error(error) + + @classmethod + async def _prepare_full_compaction( + cls, + ctx: LoopContext, + callbacks: LoopCallbacks, + messages: List[MessageInfo], + last_user: MessageInfo, + policy: CompactionPolicy, + ) -> ModelTurnPreparation[MessageInfo]: + """Run full compaction and request preparation to reload the session.""" + ctx.overflow_compaction_attempts += 1 + if ctx.overflow_compaction_attempts >= 2: + await cls._publish_session_notice( + callbacks, + ctx.session.id, + level="info", + message=( + "本轮上下文持续接近模型上限，系统将优先尝试压缩历史工具输出。" + ), + details={ + "attempt": ctx.overflow_compaction_attempts, + "threshold": policy.overflow_threshold, + "buffer": policy.overflow_buffer, + }, + ) + log.warn( + "loop.overflow_compaction_attempt", + { + "session_id": ctx.session.id, + "attempt": ctx.overflow_compaction_attempts, + "max": MAX_OVERFLOW_COMPACTION_ATTEMPTS, + }, + ) + if callbacks.on_compaction: + await callbacks.on_compaction() + await SessionCompaction.prune(ctx.session.id, policy=policy) + + publish = callbacks.event_publish_callback + progress_callback = None + if publish is not None: + + async def progress_callback(stage: str, data: dict) -> None: + await publish( + "session.compaction_progress", + { + "sessionID": ctx.session.id, + "stage": stage, + "data": data, + }, ) - payload = turn_state.model_dump(by_alias=True) - if queued_user is not None: - payload["queuedUserMessageID"] = queued_user.id - await cls._publish_runtime_event(callbacks, "turn.continued", payload) - # Continue to next iteration - continue - - else: - # Unknown action - log.warn("loop.unknown_action", { + + result = await run_compaction( + ctx.session.id, + parent_message_id=last_user.id, + messages=messages, + provider_id=ctx.provider_id, + model_id=ctx.model_id, + auto=True, + event_publish_callback=publish, + status_after="busy", + policy=policy, + progress_callback=progress_callback, + ) + if result == "stop": + log.error( + "loop.compaction_failed", + {"session_id": ctx.session.id}, + ) + if callbacks.on_error: + await callbacks.on_error("Compaction failed") + return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) + if result == "skipped": + log.info( + "loop.compaction_skipped", + {"session_id": ctx.session.id, "step": ctx.step}, + ) + else: + ctx.last_compaction_step = ctx.step + set_context_state( + ctx.session.id, + compaction_performed=True, + last_compaction_step=ctx.step, + last_compaction_reason="full_compaction", + ) + await cls._publish_runtime_event( + callbacks, + "context.compacted", + { + "sessionID": ctx.session.id, + "step": ctx.step, + "reason": "full_compaction", + "attempt": ctx.overflow_compaction_attempts, + "cooldownUntilStep": ( + ctx.step + POST_COMPACTION_COOLDOWN_STEPS + ), + }, + ) + return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) + + @classmethod + async def _run_loop( + cls, + ctx: LoopContext, + callbacks: LoopCallbacks, + ) -> LoopResult: + """Run the host-neutral AgentLoop against session-owned adapters.""" + from flocks.session.runtime_services import ( + SessionRuntimeServices, + SessionStepCancelled, + ) + + state = AgentRunState[MessageInfo]( + session_id=ctx.session.id, + agent_name=ctx.agent_name, + active_model=RuntimeModel(ctx.provider_id, ctx.model_id), + model_turn_index=ctx.step, + trace_step_offset=ctx.trace_step_offset, + current_user_id=ctx.turn_user_id, + ) + services = SessionRuntimeServices(ctx, callbacks, cls) + step_engine = cls._create_host_step_engine(ctx, callbacks) + try: + outcome = await AgentLoop( + step_engine, + services, + abort_requested=ctx.should_abort, + ).run(state) + except SessionStepCancelled: + log.info( + "loop.step_cancelled", + {"session_id": ctx.session.id, "step": ctx.step}, + ) + return LoopResult( + action="stop", + provider_id=ctx.provider_id, + model_id=ctx.model_id, + metadata={ + "steps": ctx.step, "session_id": ctx.session.id, - "action": step_result.action, - }) - break - - # Return result + "last_compaction_step": ctx.last_compaction_step, + "aborted": True, + }, + ) + + loop_error = ( + outcome.error + if outcome.status + in { + AgentRunStatus.RETRYABLE_FAILURE, + AgentRunStatus.FATAL_FAILURE, + AgentRunStatus.CONTEXT_OVERFLOW, + } + else None + ) return LoopResult( action="error" if ctx.auto_failover and loop_error else "stop", - last_message=last_message, + last_message=outcome.last_message, error=loop_error if ctx.auto_failover else None, provider_id=ctx.provider_id, model_id=ctx.model_id, @@ -2315,10 +2300,14 @@ async def progress_callback_overflow(stage: str, data: dict) -> None: "steps": ctx.step, "session_id": ctx.session.id, "last_compaction_step": ctx.last_compaction_step, - **({"aborted": True} if ctx.should_abort() else {}), + **( + {"aborted": True} + if outcome.status == AgentRunStatus.ABORTED + else {} + ), }, ) - + @classmethod def _build_compaction_policy(cls, ctx: LoopContext) -> CompactionPolicy: """ From 19fad5bdea8850d6a6ad914a6a1c853ba33b935d Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 7 Aug 2026 19:02:46 +0800 Subject: [PATCH 05/15] refactor(session): consolidate runtime execution loops --- flocks/agent/runtime/__init__.py | 57 - flocks/agent/runtime/events.py | 17 - flocks/agent/runtime/ports.py | 152 - flocks/channel/inbound/dispatcher.py | 15 +- flocks/cli/session_runner.py | 54 +- flocks/provider/options.py | 2 +- flocks/provider/sdk/google.py | 2 +- flocks/server/routes/session.py | 25 +- flocks/session/__init__.py | 13 - flocks/session/actions.py | 132 + flocks/session/features/activity_forwarder.py | 9 +- flocks/session/features/subtask.py | 4 +- flocks/session/prompt_strings.py | 2 +- flocks/session/runtime/__init__.py | 1 + .../{agent => session}/runtime/agent_loop.py | 135 +- flocks/session/runtime/continuation_policy.py | 464 +++ .../{agent => session}/runtime/contracts.py | 116 +- flocks/session/runtime/event_sink.py | 92 + flocks/session/runtime/model_policy.py | 397 +++ flocks/session/runtime/session_turn.py | 1272 +++++++ .../{runner.py => runtime/step_engine.py} | 1062 +++--- flocks/session/runtime_adapters.py | 85 - flocks/session/runtime_services.py | 103 - flocks/session/session.py | 2 - flocks/session/session_host.py | 487 --- flocks/session/session_loop.py | 2937 +++-------------- flocks/session/step_engine.py | 35 - flocks/session/utils/file_extractor.py | 2 +- flocks/task/background.py | 46 +- tests/agent/runtime/test_agent_loop.py | 384 --- tests/agent/runtime/test_contracts.py | 45 - tests/agent/test_unified_session_loop.py | 94 +- tests/channel/test_channel.py | 77 +- tests/integration/test_real_tool_calls.py | 7 +- .../test_langfuse_observability.py | 12 +- tests/session/runtime/test_agent_loop.py | 324 ++ tests/session/runtime/test_contracts.py | 87 + tests/session/runtime/test_session_loop.py | 354 ++ tests/session/runtime/test_step_engine.py | 112 + tests/session/test_actions.py | 71 + tests/session/test_auto_model_failover.py | 532 +-- tests/session/test_callable_state.py | 17 +- ...est_cli_session_runner_model_resolution.py | 6 +- tests/session/test_execution_mode.py | 6 +- tests/session/test_lifecycle_hooks.py | 283 +- tests/session/test_runner_chunk_handling.py | 18 +- tests/session/test_runner_device_hint.py | 8 +- .../session/test_runner_langfuse_payloads.py | 6 +- .../session/test_runner_llm_hook_payloads.py | 77 +- tests/session/test_runner_llm_hooks.py | 14 +- tests/session/test_runner_provider_version.py | 4 +- tests/session/test_runner_step.py | 174 +- tests/session/test_runtime_ports.py | 108 - tests/session/test_session_abort_inject.py | 185 +- tests/session/test_session_context.py | 44 +- .../test_session_loop_working_directory.py | 29 +- .../test_session_runner_tool_only_message.py | 12 +- tests/session/test_step_engine.py | 53 - tests/session_runtime_testkit.py | 35 + tests/task/test_task.py | 14 +- 60 files changed, 5602 insertions(+), 5310 deletions(-) delete mode 100644 flocks/agent/runtime/__init__.py delete mode 100644 flocks/agent/runtime/events.py delete mode 100644 flocks/agent/runtime/ports.py create mode 100644 flocks/session/actions.py create mode 100644 flocks/session/runtime/__init__.py rename flocks/{agent => session}/runtime/agent_loop.py (52%) create mode 100644 flocks/session/runtime/continuation_policy.py rename flocks/{agent => session}/runtime/contracts.py (56%) create mode 100644 flocks/session/runtime/event_sink.py create mode 100644 flocks/session/runtime/model_policy.py create mode 100644 flocks/session/runtime/session_turn.py rename flocks/session/{runner.py => runtime/step_engine.py} (88%) delete mode 100644 flocks/session/runtime_adapters.py delete mode 100644 flocks/session/runtime_services.py delete mode 100644 flocks/session/session_host.py delete mode 100644 flocks/session/step_engine.py delete mode 100644 tests/agent/runtime/test_agent_loop.py delete mode 100644 tests/agent/runtime/test_contracts.py create mode 100644 tests/session/runtime/test_agent_loop.py create mode 100644 tests/session/runtime/test_contracts.py create mode 100644 tests/session/runtime/test_session_loop.py create mode 100644 tests/session/runtime/test_step_engine.py create mode 100644 tests/session/test_actions.py delete mode 100644 tests/session/test_runtime_ports.py delete mode 100644 tests/session/test_step_engine.py create mode 100644 tests/session_runtime_testkit.py diff --git a/flocks/agent/runtime/__init__.py b/flocks/agent/runtime/__init__.py deleted file mode 100644 index d72b1b1b6..000000000 --- a/flocks/agent/runtime/__init__.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Reusable agent runtime contracts and control flow.""" - -from flocks.agent.runtime.agent_loop import AgentLoop -from flocks.agent.runtime.contracts import ( - AgentRunOutcome, - AgentRunState, - AgentRunStatus, - AttemptEffects, - ContinuationDecision, - FailoverDecision, - ModelTurnBoundary, - ModelTurnPreparation, - ModelTurnSnapshot, - QueuedInputBatch, - RuntimeModel, - StepFailure, - StepResult, - ToolCall, - TurnPreparationStatus, -) -from flocks.agent.runtime.ports import ( - ExternalRuntimePorts, - HookPort, - ModelPort, - PromptPort, - RuntimeEventSink, - RuntimeServices, - StepEngine, - ToolPort, -) - -__all__ = [ - "AgentRunOutcome", - "AgentRunState", - "AgentRunStatus", - "AgentLoop", - "AttemptEffects", - "ContinuationDecision", - "FailoverDecision", - "ExternalRuntimePorts", - "HookPort", - "ModelPort", - "ModelTurnBoundary", - "ModelTurnPreparation", - "ModelTurnSnapshot", - "PromptPort", - "QueuedInputBatch", - "RuntimeModel", - "RuntimeEventSink", - "RuntimeServices", - "StepEngine", - "StepFailure", - "StepResult", - "ToolCall", - "ToolPort", - "TurnPreparationStatus", -] diff --git a/flocks/agent/runtime/events.py b/flocks/agent/runtime/events.py deleted file mode 100644 index efbd1501d..000000000 --- a/flocks/agent/runtime/events.py +++ /dev/null @@ -1,17 +0,0 @@ -"""Runtime event contract emitted by the agent core.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Mapping, Optional - - -@dataclass(frozen=True) -class RuntimeEvent: - """Host-neutral event produced during an agent run.""" - - type: str - session_id: str - model_turn_index: int - payload: Mapping[str, Any] = field(default_factory=dict) - trace_step: Optional[int] = None diff --git a/flocks/agent/runtime/ports.py b/flocks/agent/runtime/ports.py deleted file mode 100644 index 9b476e9f4..000000000 --- a/flocks/agent/runtime/ports.py +++ /dev/null @@ -1,152 +0,0 @@ -"""Ports implemented by session and infrastructure adapters.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Generic, Optional, Protocol, TypeVar - -from flocks.agent.runtime.contracts import ( - AgentRunState, - ContinuationDecision, - ModelTurnBoundary, - ModelTurnPreparation, - ModelTurnSnapshot, - StepResult, -) -from flocks.agent.runtime.events import RuntimeEvent - - -MessageT = TypeVar("MessageT") - - -class StepEngine(Protocol, Generic[MessageT]): - """Execute one model turn from an immutable snapshot.""" - - async def run(self, snapshot: ModelTurnSnapshot[MessageT]) -> StepResult: - """Run one streamed model turn and return its result.""" - ... - - -class RuntimeServices(Protocol, Generic[MessageT]): - """Narrow session-host interface consumed by the agent loop.""" - - async def prepare_model_turn( - self, - state: AgentRunState[MessageT], - ) -> ModelTurnPreparation[MessageT]: - """Prepare or defer the next stable model-turn input.""" - ... - - async def complete_model_turn( - self, - state: AgentRunState[MessageT], - step_result: StepResult, - ) -> ModelTurnBoundary[MessageT]: - """Return the committed post-step view and queued inputs.""" - ... - - async def resolve_continuation( - self, - state: AgentRunState[MessageT], - step_result: StepResult, - ) -> ContinuationDecision[MessageT]: - """Resolve goal and turn-finish-hook continuation policy.""" - ... - - async def emit_event(self, event: RuntimeEvent) -> None: - """Forward a runtime event to host-owned sinks.""" - ... - - -class PromptPort(Protocol): - """Build provider-ready system prompt sections.""" - - async def build_system_prompts(self, **kwargs: Any) -> list[str]: - """Build prompts from one stable model-turn configuration.""" - ... - - -class ToolPort(Protocol): - """Expose the tool registry without coupling the core to its storage.""" - - def revision(self) -> int: - """Return a cache revision for the visible tool set.""" - ... - - def list_tools(self) -> list[Any]: - """Return registered tool metadata entries.""" - ... - - def get(self, name: str) -> Optional[Any]: - """Resolve one executable tool by name.""" - ... - - -class ModelPort(Protocol): - """Resolve and configure provider/model adapters.""" - - def get_provider(self, provider_id: str) -> Optional[Any]: - """Return one configured provider adapter.""" - ... - - async def apply_config(self, provider_id: str) -> None: - """Apply persisted provider configuration before execution.""" - ... - - def resolve_model(self, provider_id: str, model_id: str) -> Optional[Any]: - """Return model capability metadata.""" - ... - - def resolve_model_info( - self, - provider_id: str, - model_id: str, - ) -> tuple[int, int, Optional[int]]: - """Return context, output, and input token limits.""" - ... - - -class HookPort(Protocol): - """Run existing Flocks hook stages at explicit runtime boundaries.""" - - async def run_session_start(self, data: dict[str, Any]) -> Any: - """Run SessionStart hooks.""" - ... - - async def has_stage_handlers( - self, - stage: Any, - metadata: dict[str, Any], - ) -> bool: - """Return whether a hook stage has eligible handlers.""" - ... - - async def run_llm_before(self, data: dict[str, Any]) -> Any: - """Run LLMBefore hooks.""" - ... - - async def run_llm_after( - self, - metadata: dict[str, Any], - result: dict[str, Any], - ) -> Any: - """Run LLMAfter hooks.""" - ... - - -class RuntimeEventSink(Protocol): - """Receive observable runtime events; not an event-sourcing store.""" - - async def emit(self, event: RuntimeEvent) -> None: - """Forward an event to UI, tracing, or audit subscribers.""" - ... - - -@dataclass(frozen=True) -class ExternalRuntimePorts: - """External interfaces captured once for a stable model attempt.""" - - prompts: PromptPort - tools: ToolPort - models: ModelPort - hooks: HookPort diff --git a/flocks/channel/inbound/dispatcher.py b/flocks/channel/inbound/dispatcher.py index 910020e0d..f0e0ca1e9 100644 --- a/flocks/channel/inbound/dispatcher.py +++ b/flocks/channel/inbound/dispatcher.py @@ -234,14 +234,14 @@ async def deliver_text(self, text: str) -> None: session_id=self.session_id, ) - def to_loop_callbacks(self, runner_callbacks=None): + def to_loop_callbacks(self, *, on_text_delta=None): """Convert to a LoopCallbacks dataclass understood by SessionLoop.""" from flocks.session.session_loop import LoopCallbacks return LoopCallbacks( on_step_end=self.on_step_end, + on_text_delta=on_text_delta, on_error=self.on_error, event_publish_callback=self._publish_sse_event, - runner_callbacks=runner_callbacks, ) @staticmethod @@ -464,8 +464,8 @@ async def dispatch(self, msg: InboundMessage) -> None: # what _process_session_message does in the WebUI route. Storing # the resolved model on the user message keeps two things aligned # between WebUI and channel: - # - Title generation (``SessionLoop._run_loop`` reads - # ``last_user.model``). + # - Title generation and turn preparation read + # ``last_user.model``. # - The provider-specific base prompt template # (``SystemPrompt.provider``) selected on the next loop tick. # Without this, channel sessions ended up with the hardcoded @@ -1150,13 +1150,12 @@ async def _run_agent_with_streaming( return try: - from flocks.session.runner import RunnerCallbacks - async def _on_text_delta(delta: str) -> None: await card.append(delta) - runner_cbs = RunnerCallbacks(on_text_delta=_on_text_delta) - loop_callbacks = callbacks.to_loop_callbacks(runner_callbacks=runner_cbs) + loop_callbacks = callbacks.to_loop_callbacks( + on_text_delta=_on_text_delta, + ) result = await InboundDispatcher._run_session_loop( binding, diff --git a/flocks/cli/session_runner.py b/flocks/cli/session_runner.py index 5d5aa6766..9e92792d5 100644 --- a/flocks/cli/session_runner.py +++ b/flocks/cli/session_runner.py @@ -6,7 +6,7 @@ - Tool execution display - Streaming text display -Core logic is in session/runner.py +Core logic is exposed through SessionLoop. """ import asyncio @@ -26,11 +26,11 @@ from flocks.utils.log import Log from flocks.session.session import Session, SessionInfo -from flocks.session.runner import SessionRunner, RunnerCallbacks, ToolResult +from flocks.session.session_loop import LoopCallbacks from flocks.session.message import Message, MessageRole from flocks.agent.registry import Agent from flocks.provider.provider import Provider -from flocks.tool.registry import ToolRegistry +from flocks.tool.registry import ToolRegistry, ToolResult from flocks.project.project import Project from dotenv import load_dotenv @@ -38,21 +38,6 @@ log = Log.create(service="cli.runner") -# Module-level storage for CLI callbacks (used by SessionRunner during loop execution) -_CLI_CALLBACKS: Optional['RunnerCallbacks'] = None - - -def _set_cli_callbacks(callbacks: Optional['RunnerCallbacks']) -> None: - """Set CLI callbacks for current execution""" - global _CLI_CALLBACKS - _CLI_CALLBACKS = callbacks - - -def _get_cli_callbacks() -> Optional['RunnerCallbacks']: - """Get CLI callbacks for current execution""" - return _CLI_CALLBACKS - - # Tool display styles TOOL_STYLES: Dict[str, tuple] = { "todo": ("Todo", "yellow bold"), @@ -71,7 +56,7 @@ def _get_cli_callbacks() -> Optional['RunnerCallbacks']: class CLISessionRunner: """ - CLI wrapper for SessionRunner. + CLI wrapper for the public SessionLoop entry point. Handles all CLI-specific display logic. """ @@ -90,7 +75,6 @@ def __init__( self.agent_name = agent self.auto_confirm = auto_confirm self._session: Optional[SessionInfo] = None - self._runner: Optional[SessionRunner] = None self._live: Optional[Live] = None self._content_buffer: list[str] = [] @@ -299,8 +283,10 @@ async def _interactive_loop(self) -> None: except KeyboardInterrupt: self.console.print("\n[dim]Interrupted[/dim]") - if self._runner: - self._runner.abort() + if self._session: + from flocks.session.session_loop import SessionLoop + + SessionLoop.abort(self._session.id) break except EOFError: break @@ -348,7 +334,6 @@ async def _process_message( from flocks.input.dispatcher import dispatch_user_input from flocks.input.events import UserInputEvent from flocks.input.output import CliOutputSink - from flocks.session.message import Message event = UserInputEvent( source_type="cli", @@ -408,29 +393,21 @@ async def _clear_history() -> None: model={"providerID": provider_id, "modelID": model_id}, ) - # Import SessionLoop and LoopCallbacks - from flocks.session.session_loop import SessionLoop, LoopCallbacks - from flocks.session.runner import RunnerCallbacks + # Import the stable session execution entry point. + from flocks.session.session_loop import SessionLoop - # Create loop callbacks (wrapping runner callbacks) + # Pass one explicit callback set through the full runtime. loop_callbacks = LoopCallbacks( on_step_start=self._on_step_start, on_step_end=self._on_step_end, - on_error=self._on_error, - on_compaction=self._on_compaction, - ) - - # Store runner callbacks for tool events - # We need to hook into SessionRunner to get tool callbacks - # This is done by temporarily storing callbacks in a module-level variable - _set_cli_callbacks(RunnerCallbacks( on_text_delta=self._on_text_delta, on_reasoning_delta=self._on_reasoning_delta, on_tool_start=self._on_tool_start, on_tool_end=self._on_tool_end, on_permission_request=self._on_permission_request, on_error=self._on_error, - )) + on_compaction=self._on_compaction, + ) # Start streaming display self._content_buffer = [] @@ -486,9 +463,6 @@ async def _clear_history() -> None: live.update(Text("")) self._live = None - # Clear callbacks - _set_cli_callbacks(None) - # Print any remaining content not yet printed if self._content_buffer: self._flush_content() @@ -806,8 +780,6 @@ def _print_help(self) -> None: __all__ = [ "CLISessionRunner", "run_session", - "_get_cli_callbacks", - "_set_cli_callbacks", ] diff --git a/flocks/provider/options.py b/flocks/provider/options.py index 733093df5..41c6a7e3a 100644 --- a/flocks/provider/options.py +++ b/flocks/provider/options.py @@ -4,7 +4,7 @@ Centralises the logic for assembling thinking / reasoning / token-limit kwargs that get forwarded to each provider's ``chat_stream`` call. -Both ``SessionRunner`` (session/runner.py) and ``AgentExecutor`` +Both ``StepEngine`` and ``AgentExecutor`` (agent/runtime/executor.py) delegate to :func:`build_provider_options` so that provider rules are maintained in exactly one place. """ diff --git a/flocks/provider/sdk/google.py b/flocks/provider/sdk/google.py index 39df39669..4ac53fe0f 100644 --- a/flocks/provider/sdk/google.py +++ b/flocks/provider/sdk/google.py @@ -81,7 +81,7 @@ def _convert_messages( Rewrites history as text to bypass binary thought_signature requirements. ``session_id`` is forwarded by the runner via kwargs (see - ``SessionRunner._call_llm``). When provided, we attempt to reconstruct + ``StepEngine._call_llm``). When provided, we attempt to reconstruct the conversation directly from persisted session messages – including reasoning parts – which gives Gemini perfect context. As a defensive fallback we also honour ``messages[0].sessionID`` / ``session_id`` diff --git a/flocks/server/routes/session.py b/flocks/server/routes/session.py index a7c920c04..5fd1a6f3c 100644 --- a/flocks/server/routes/session.py +++ b/flocks/server/routes/session.py @@ -1733,30 +1733,24 @@ async def unshare_session_local(sessionID: str, http_request: Request) -> Sessio async def _abort_session_processing(sessionID: str) -> bool: """Abort active processing for a session and notify subscribers. - Aborts both the SessionLoop (sets abort_event so the next step check - stops the loop) and the SessionRunner (stops the current LLM stream). + Aborts SessionLoop, which owns the active StepEngine abort signal. Also auto-rejects any pending Question tool requests so the question handler polling loop unblocks immediately instead of timing out. Cascades abort to all child sub-agent sessions (synchronous subtasks and background tasks) so they stop together with the parent. """ - from flocks.session.runner import SessionRunner from flocks.session.session_loop import SessionLoop from flocks.server.routes.question import reject_session_questions # Abort the loop-level context (propagates to runner via shared abort_event) loop_aborted = SessionLoop.abort(sessionID) - # Also cancel through the runner's own path (sets status to idle) - SessionRunner.cancel(sessionID) - # Unblock any pending Question tool waiting for user input questions_rejected = await reject_session_questions(sessionID) # --- Cascade abort to child sub-agent sessions --- children_loops_aborted = SessionLoop.abort_children(sessionID) - children_runners_cancelled = SessionRunner.cancel_children(sessionID) # Cancel background sub-agent tasks spawned by this session bg_cancelled = 0 @@ -1771,7 +1765,6 @@ async def _abort_session_processing(sessionID: str) -> bool: "loop_aborted": loop_aborted, "questions_rejected": questions_rejected, "children_loops_aborted": children_loops_aborted, - "children_runners_cancelled": children_runners_cancelled, "bg_tasks_cancelled": bg_cancelled, }) @@ -1828,7 +1821,7 @@ class InitRequest(BaseModel): ) async def initialize_session(sessionID: str, request: InitRequest, http_request: Request) -> bool: """Initialize session""" - from flocks.session.runner import SessionRunner + from flocks.session.actions import render_session_command current_user = require_user(http_request) session = await _get_session_by_id_unfiltered(sessionID) @@ -1840,12 +1833,10 @@ async def initialize_session(sessionID: str, request: InitRequest, http_request: _require_session_write_access(session, current_user) # Execute INIT command - await SessionRunner.command( + await render_session_command( session_id=sessionID, command="init", arguments="", - message_id=request.messageID, - model=f"{request.providerID}/{request.modelID}", ) log.info("session.initialized", {"session_id": sessionID}) @@ -3460,7 +3451,6 @@ async def _process_session_message( from flocks.agent.registry import Agent from flocks.provider.provider import Provider from flocks.session.session_loop import SessionLoop, LoopCallbacks - from flocks.session.runner import RunnerCallbacks import time import os @@ -4903,7 +4893,7 @@ class ShellRequest(BaseModel): ) async def run_shell_command(sessionID: str, request: ShellRequest, http_request: Request): """Run shell command""" - from flocks.session.runner import SessionRunner + from flocks.session.actions import run_session_shell current_user = require_user(http_request) session = await _get_session_by_id_unfiltered(sessionID) @@ -4914,17 +4904,12 @@ async def run_shell_command(sessionID: str, request: ShellRequest, http_request: ) _require_session_write_access(session, current_user) - model = None - if request.model: - model = {"providerID": request.model.providerID, "modelID": request.model.modelID} - try: async with Session.active_operation(sessionID): - result = await SessionRunner.shell( + result = await run_session_shell( session_id=sessionID, agent=request.agent, command=request.command, - model=model, ) except SessionNotFoundError as exc: raise HTTPException( diff --git a/flocks/session/__init__.py b/flocks/session/__init__.py index 54840e2fb..2dc887f82 100644 --- a/flocks/session/__init__.py +++ b/flocks/session/__init__.py @@ -35,13 +35,6 @@ from flocks.session.prompt import SessionPrompt, SystemPrompt, ContextInfo from flocks.session.lifecycle.compaction import SessionCompaction, CompactionResult, CompactionPolicy, ContextTier from flocks.session.lifecycle.summary import SessionSummary, FileDiff -from flocks.session.runner import ( - SessionRunner, - RunnerCallbacks, - ToolCall, - StepResult, - run_session, -) from flocks.session.session_loop import ( SessionLoop, LoopContext, @@ -104,12 +97,6 @@ # Summary "SessionSummary", "FileDiff", - # Runner - "SessionRunner", - "RunnerCallbacks", - "ToolCall", - "StepResult", - "run_session", # Session Loop "SessionLoop", "LoopContext", diff --git a/flocks/session/actions.py b/flocks/session/actions.py new file mode 100644 index 000000000..71c0220b1 --- /dev/null +++ b/flocks/session/actions.py @@ -0,0 +1,132 @@ +"""Session actions that are independent of the agent execution loop.""" + +import asyncio +import os +from typing import Any, Optional + +from flocks.session.message import Message, MessageRole +from flocks.session.session import Session +from flocks.utils.id import Identifier +from flocks.utils.log import Log + + +log = Log.create(service="session.actions") + + +async def render_session_command( + session_id: str, + command: str, + arguments: str = "", +) -> dict[str, Any]: + """Resolve and render one slash-command template.""" + from flocks.command.command import Command + + command_info = Command.get(command) + if not command_info: + raise ValueError(f"Command '{command}' not found") + template = command_info.template.replace("$ARGUMENTS", arguments) + log.info( + "session.command", + { + "session_id": session_id, + "command": command, + "arguments": arguments[:50] if arguments else "", + }, + ) + return { + "command": command, + "arguments": arguments, + "template": template, + } + + +async def run_session_shell( + session_id: str, + agent: str, + command: str, +) -> dict[str, Any]: + """Execute one explicit user shell action and return its tool part.""" + session = await Session.get_by_id(session_id) + if not session: + raise ValueError(f"Session {session_id} not found") + + user_message = await Message.create( + session_id=session_id, + role=MessageRole.USER, + content="The following tool was executed by the user", + agent=agent, + ) + assistant_message = await Message.create( + session_id=session_id, + role=MessageRole.ASSISTANT, + content="", + agent=agent, + parent_id=user_message.id, + ) + + started_at = asyncio.get_event_loop().time() + process: Optional[asyncio.subprocess.Process] = None + try: + process = await asyncio.create_subprocess_shell( + command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=session.directory or os.getcwd(), + ) + stdout_bytes, stderr_bytes = await asyncio.wait_for( + process.communicate(), + timeout=300, + ) + output = ( + (stdout_bytes or b"").decode("utf-8", errors="replace") + + (stderr_bytes or b"").decode("utf-8", errors="replace") + ) + exit_code = process.returncode or 0 + except asyncio.TimeoutError: + output = "Command timed out after 300 seconds" + exit_code = -1 + if process is not None: + try: + process.kill() + except Exception as exc: + log.debug( + "session.shell.kill_failed", + {"error": str(exc)}, + ) + except Exception as exc: + output = f"Error executing command: {exc}" + exit_code = -1 + + log.info( + "session.shell", + { + "session_id": session_id, + "command": command[:50], + "exit_code": exit_code, + "duration_ms": int( + (asyncio.get_event_loop().time() - started_at) * 1000, + ), + }, + ) + return { + "info": { + "id": assistant_message.id, + "sessionID": session_id, + "role": "assistant", + "agent": agent, + }, + "parts": [ + { + "id": Identifier.create("part"), + "messageID": assistant_message.id, + "sessionID": session_id, + "type": "tool", + "tool": "bash", + "state": { + "status": "completed", + "input": {"command": command}, + "output": output, + }, + }, + ], + } diff --git a/flocks/session/features/activity_forwarder.py b/flocks/session/features/activity_forwarder.py index a4eeff11e..5898e229b 100644 --- a/flocks/session/features/activity_forwarder.py +++ b/flocks/session/features/activity_forwarder.py @@ -69,15 +69,12 @@ def build_callbacks(self, event_publish_callback=None): server layer. Avoids session → server reverse dependency. """ from flocks.session.session_loop import LoopCallbacks - from flocks.session.runner import RunnerCallbacks return LoopCallbacks( event_publish_callback=event_publish_callback, - runner_callbacks=RunnerCallbacks( - on_tool_start=self._on_tool_start, - on_tool_end=self._on_tool_end, - on_text_delta=self._on_text_delta, - ), + on_tool_start=self._on_tool_start, + on_tool_end=self._on_tool_end, + on_text_delta=self._on_text_delta, ) # ------------------------------------------------------------------ diff --git a/flocks/session/features/subtask.py b/flocks/session/features/subtask.py index f1231cea5..b7aac9ca2 100644 --- a/flocks/session/features/subtask.py +++ b/flocks/session/features/subtask.py @@ -44,14 +44,14 @@ class SubtaskResult: class SessionSubtask: """Subtask manager stub — business logic removed (was dead code). - The active execution path is SessionLoop._execute_subtask() in session_loop.py. + The active execution path is SessionTurn._execute_subtask(). """ @classmethod async def execute_subtask(cls, *args, **kwargs) -> SubtaskResult: raise NotImplementedError( "SessionSubtask.execute_subtask() is deprecated. " - "Subtask execution is handled by SessionLoop._execute_subtask()." + "Subtask execution is handled by SessionTurn._execute_subtask()." ) diff --git a/flocks/session/prompt_strings.py b/flocks/session/prompt_strings.py index 8ccb18b5d..73abada35 100644 --- a/flocks/session/prompt_strings.py +++ b/flocks/session/prompt_strings.py @@ -181,7 +181,7 @@ """ # ============================================================================= -# Runner prompt snippets (used by SessionRunner._process_step) +# Step prompt snippets (used by StepEngine._process_step) # ============================================================================= PROMPT_TOOL_RESULTS_AVAILABLE = ( diff --git a/flocks/session/runtime/__init__.py b/flocks/session/runtime/__init__.py new file mode 100644 index 000000000..f138134e9 --- /dev/null +++ b/flocks/session/runtime/__init__.py @@ -0,0 +1 @@ +"""Internal session turn, agent loop, and step execution package.""" diff --git a/flocks/agent/runtime/agent_loop.py b/flocks/session/runtime/agent_loop.py similarity index 52% rename from flocks/agent/runtime/agent_loop.py rename to flocks/session/runtime/agent_loop.py index f55878b61..9a54133ed 100644 --- a/flocks/agent/runtime/agent_loop.py +++ b/flocks/session/runtime/agent_loop.py @@ -1,45 +1,36 @@ -"""Host-neutral model/tool/continuation control loop.""" +"""The control loop for one logical user input.""" from __future__ import annotations -from collections.abc import Callable -from typing import Generic, Optional, TypeVar - -from flocks.agent.runtime.contracts import ( +from flocks.session.message import MessageInfo +from flocks.session.runtime.contracts import ( AgentRunOutcome, - AgentRunState, AgentRunStatus, + StepAction, TurnPreparationStatus, ) -from flocks.agent.runtime.ports import RuntimeServices, StepEngine - +from flocks.session.runtime.session_turn import SessionTurn +from flocks.session.runtime.step_engine import StepCancelled, StepEngine +from flocks.utils.log import Log -MessageT = TypeVar("MessageT") +log = Log.create(service="session.agent_loop") -class AgentLoop(Generic[MessageT]): - """Coordinate model turns without owning session policy or persistence.""" - def __init__( - self, - step_engine: StepEngine[MessageT], - services: RuntimeServices[MessageT], - *, - abort_requested: Optional[Callable[[], bool]] = None, - ): - self._step_engine = step_engine - self._services = services - self._abort_requested = abort_requested or (lambda: False) +class AgentLoop: + """Decide whether one logical user input needs another model step.""" async def run( self, - state: AgentRunState[MessageT], - ) -> AgentRunOutcome[MessageT]: - """Run or resume an agent until it settles or needs host recovery.""" - last_message: Optional[MessageT] = None - - while not self._abort_requested(): - preparation = await self._services.prepare_model_turn(state) + turn: SessionTurn, + engine: StepEngine, + ) -> AgentRunOutcome[MessageInfo]: + """Run the current logical input to a session-level boundary.""" + state = turn.state + last_message = None + + while not turn.aborted: + preparation = await turn.prepare_step() if preparation.status == TurnPreparationStatus.CONTINUE: continue if preparation.status == TurnPreparationStatus.COMPLETE: @@ -62,21 +53,41 @@ async def run( status=AgentRunStatus.FATAL_FAILURE, state=state, last_message=last_message, - error="Runtime services returned READY without a model-turn snapshot", + error=( + "SessionTurn returned READY without a model-turn " + "snapshot" + ), ) state.active_model = snapshot.active_model state.model_turn_index = snapshot.model_turn_index state.messages = list(snapshot.messages) - step_result = await self._step_engine.run(snapshot) - boundary = await self._services.complete_model_turn(state, step_result) + try: + step_result = await engine.run(snapshot) + except StepCancelled: + log.info( + "session.step.cancelled", + { + "session_id": turn.session.id, + "step": turn.step, + }, + ) + return AgentRunOutcome( + status=AgentRunStatus.ABORTED, + state=state, + last_message=last_message, + error="Aborted", + ) + + state.active_model = ( + step_result.effective_model or snapshot.active_model + ) + boundary = await turn.commit_step(step_result) state.messages = list(boundary.messages) last_message = boundary.last_message or last_message - if boundary.queued_inputs.cursor is not None: - state.consumed_input_cursor = boundary.queued_inputs.cursor - if self._abort_requested(): + if turn.aborted: return AgentRunOutcome( status=AgentRunStatus.ABORTED, state=state, @@ -85,17 +96,22 @@ async def run( ) if boundary.queued_inputs.messages: - self._append_new_messages( - state, - boundary.queued_inputs.messages, + return AgentRunOutcome( + status=AgentRunStatus.INPUT_AVAILABLE, + state=state, + last_message=last_message, + error=step_result.error, + step_result=step_result, ) - continue failure = step_result.failure if failure is not None: status = ( AgentRunStatus.RETRYABLE_FAILURE - if failure.allow_fallback and failure.attempt_state.replay_safe + if ( + failure.allow_fallback + and failure.attempt_state.replay_safe + ) else AgentRunStatus.FATAL_FAILURE ) return AgentRunOutcome( @@ -103,24 +119,26 @@ async def run( state=state, last_message=last_message, error=failure.message, - failure=failure, + step_result=step_result, ) - if step_result.action == "continue": + if step_result.action == StepAction.CONTINUE: continue - if step_result.action == "compact": + if step_result.action == StepAction.COMPACT: return AgentRunOutcome( status=AgentRunStatus.CONTEXT_OVERFLOW, state=state, last_message=last_message, error=step_result.error, + step_result=step_result, ) - if step_result.action != "stop": + if step_result.action != StepAction.STOP: return AgentRunOutcome( status=AgentRunStatus.FATAL_FAILURE, state=state, last_message=last_message, error=f"Unknown step action: {step_result.action}", + step_result=step_result, ) if step_result.error: return AgentRunOutcome( @@ -128,20 +146,14 @@ async def run( state=state, last_message=last_message, error=step_result.error, + step_result=step_result, ) - continuation = await self._services.resolve_continuation( - state, - step_result, - ) - if continuation.should_continue: - self._append_new_messages(state, continuation.messages) - continue - return AgentRunOutcome( status=AgentRunStatus.COMPLETED, state=state, last_message=last_message, + step_result=step_result, ) return AgentRunOutcome( @@ -150,24 +162,3 @@ async def run( last_message=last_message, error="Aborted", ) - - @staticmethod - def _append_new_messages( - state: AgentRunState[MessageT], - messages: tuple[MessageT, ...], - ) -> None: - """Append messages not already present in the current runtime view.""" - existing_ids = {AgentLoop._message_identity(message) for message in state.messages} - for message in messages: - identity = AgentLoop._message_identity(message) - if identity not in existing_ids: - state.messages.append(message) - existing_ids.add(identity) - - @staticmethod - def _message_identity(message: MessageT) -> tuple[str, object]: - """Return a stable identity for persisted or in-memory messages.""" - message_id = getattr(message, "id", None) - if message_id is not None: - return ("message_id", message_id) - return ("object_id", id(message)) diff --git a/flocks/session/runtime/continuation_policy.py b/flocks/session/runtime/continuation_policy.py new file mode 100644 index 000000000..5fbc1a86d --- /dev/null +++ b/flocks/session/runtime/continuation_policy.py @@ -0,0 +1,464 @@ +"""Session-level logical turn preparation and continuation policy.""" + +from __future__ import annotations + +import inspect +from typing import Any, Optional + +from flocks.hooks.pipeline import HookPipeline +from flocks.session.runtime.contracts import ( + AgentRunOutcome, + ContinuationDecision, +) +from flocks.session.core.turn_state import set_turn_state +from flocks.session.runtime.event_sink import SessionEventSink +from flocks.session.goal import GoalManager +from flocks.session.message import Message, MessageInfo, MessageRole +from flocks.session.runtime.model_policy import ( + DEFAULT_MODEL_ROUTING_POLICY, + ModelRoutingPolicy, +) +from flocks.utils.log import Log + + +log = Log.create(service="session.continuation_policy") + + +class ContinuationPolicy: + """Own boundaries between durable logical user turns.""" + + def __init__(self, model_policy: ModelRoutingPolicy) -> None: + self._model_policy = model_policy + + async def publish_turn_stopped( + self, + turn: Any, + *, + stop_reason: str, + ) -> None: + """Publish the terminal state of one logical turn.""" + await SessionEventSink.turn_stopped( + turn.callbacks, + turn.session.id, + step=turn.step, + stop_reason=stop_reason, + ) + + @staticmethod + async def detect_queued_user_message( + _session_id: str, + post_messages: list[MessageInfo], + current_user_id: str, + _last_message: Optional[MessageInfo], + ) -> Optional[MessageInfo]: + """Return the newest user message after the current logical input.""" + newest_user = next( + (message for message in reversed(post_messages) if message.role == MessageRole.USER), + None, + ) + if newest_user is None or newest_user.id <= current_user_id: + return None + return newest_user + + async def prepare_logical_turn(self, context: Any) -> None: + """Prepare model routing and UserPromptSubmit once per logical input.""" + if context.session_store: + messages = await context.session_store.get_messages() + else: + messages = await Message.list(context.session.id) + context.prepared_messages = list(messages) + last_user = next( + (message for message in reversed(messages) if message.role == MessageRole.USER), + None, + ) + if last_user is None or last_user.id == context.prepared_user_id: + return + + is_real_user_turn = await self._model_policy.prepare_turn( + context, + last_user, + ) + if is_real_user_turn: + context.turn_additional_context = None + context.stop_hook_active = False + await self.run_user_prompt_submit(context, last_user) + context.prepared_user_id = last_user.id + + @staticmethod + async def run_user_prompt_submit(context: Any, last_user: MessageInfo) -> None: + """Run UserPromptSubmit at the session logical-turn boundary.""" + try: + prompt = await Message.get_text_content(last_user) + hook_context = await HookPipeline.run_user_prompt_submit( + { + "sessionID": context.session.id, + "workspace": context.session.directory, + "agent": getattr(last_user, "agent", None) or context.agent_name, + "model": { + "providerID": context.provider_id, + "modelID": context.model_id, + }, + "messageID": last_user.id, + "prompt": prompt, + } + ) + additional_context = hook_context.output.get("additionalContext") + if isinstance(additional_context, str) and additional_context.strip(): + context.turn_additional_context = additional_context.strip() + except Exception as exc: + log.debug( + "session.hook.user_prompt_submit.error", + { + "session_id": context.session.id, + "message_id": last_user.id, + "error": str(exc), + }, + ) + + async def resolve( + self, + context: Any, + outcome: AgentRunOutcome[MessageInfo], + ) -> ContinuationDecision[MessageInfo]: + """Resolve goal and TurnFinish into a new logical turn.""" + last_user = outcome.state.metadata.get("last_user") + last_message = outcome.last_message + if last_user is None or last_message is None: + await self.publish_turn_stopped( + context, + stop_reason="stop", + ) + return ContinuationDecision() + + queued_decision = await self._materialize_continuation( + context, + last_user, + last_message, + ) + if queued_decision.should_continue: + return queued_decision + + try: + content_result = Message.get_text_content(last_message) + last_response = await content_result if inspect.isawaitable(content_result) else content_result + except Exception as exc: + log.warn( + "session.goal.last_response_error", + { + "session_id": context.session.id, + "message_id": getattr(last_message, "id", None), + "error": str(exc), + }, + ) + last_response = getattr(last_message, "content", "") or "" + + pending_user_input = False + try: + from flocks.server.routes.question import has_pending_questions + + pending_user_input = has_pending_questions(context.session.id) + except Exception as exc: + log.warn( + "session.goal.pending_question_check_error", + {"session_id": context.session.id, "error": str(exc)}, + ) + + goal_decision = await GoalManager.evaluate_after_turn( + context.session.id, + str(last_response or ""), + pending_user_input=pending_user_input, + provider_id=context.provider_id, + model_id=context.model_id, + ) + if goal_decision.status in {"completed", "blocked", "paused"} and goal_decision.objective: + await SessionEventSink.emit( + context.callbacks, + "session.goal.updated", + { + "sessionID": context.session.id, + "status": goal_decision.status, + "objective": goal_decision.objective, + "reason": goal_decision.reason, + }, + ) + if goal_decision.should_continue and goal_decision.continuation_prompt: + allow_synthetic = await self._synthetic_continuation_allowed( + context, + last_message, + ) + goal_continuation = await self._materialize_continuation( + context, + last_user, + last_message, + candidate_reason="goal", + content=goal_decision.continuation_prompt, + agent=( + last_user.agent + if hasattr(last_user, "agent") + else context.agent_name + ), + model=( + last_user.model + if hasattr(last_user, "model") + else { + "providerID": context.provider_id, + "modelID": context.model_id, + } + ), + provider=( + last_user.provider + if hasattr(last_user, "provider") + else context.provider_id + ), + part_metadata={ + "goalContinuation": True, + "goalVerdict": goal_decision.verdict, + "goalReason": goal_decision.reason, + }, + event_metadata={"goalVerdict": goal_decision.verdict}, + allow_synthetic=allow_synthetic, + ) + if goal_continuation.should_continue: + return goal_continuation + await self.publish_turn_stopped(context, stop_reason="stop") + return ContinuationDecision() + + queued_decision = await self._materialize_continuation( + context, + last_user, + last_message, + ) + if queued_decision.should_continue: + return queued_decision + + if not context.should_abort() and getattr(last_message, "finish", None) == "stop": + hook_decision = await self.run_turn_finish( + context, + last_user, + last_message, + ) + if hook_decision.should_continue: + return hook_decision + + stop_reason = getattr(last_message, "finish", None) or "stop" + await self.publish_turn_stopped( + context, + stop_reason=stop_reason, + ) + return ContinuationDecision() + + async def run_turn_finish( + self, + context: Any, + last_user: MessageInfo, + last_message: MessageInfo, + ) -> ContinuationDecision[MessageInfo]: + """Run TurnFinish and materialize a blocked-stop continuation.""" + try: + hook_user = last_user + if context.turn_user_id: + hook_user = await Message.get(context.session.id, context.turn_user_id) or last_user + user_text = await Message.get_text_content(hook_user) + assistant_text = await Message.get_text_content(last_message) + hook_context = await HookPipeline.run_turn_finish( + { + "sessionID": context.session.id, + "workspace": context.session.directory, + "agent": getattr(last_message, "agent", None) or context.agent_name, + "model": { + "providerID": context.provider_id, + "modelID": context.model_id, + }, + "step": context.trace_step, + "userMessage": { + "id": hook_user.id, + "content": user_text, + }, + "assistantMessage": { + "id": last_message.id, + "content": assistant_text, + }, + "finishReason": "stop", + "stopHookActive": context.stop_hook_active, + } + ) + except Exception as exc: + log.debug( + "session.hook.turn_finish_error", + { + "session_id": context.session.id, + "message_id": getattr(last_message, "id", None), + "error": str(exc), + }, + ) + return ContinuationDecision() + + decision = str(hook_context.output.get("decision") or "").strip().lower() + reason = str(hook_context.output.get("reason") or "").strip() + if decision != "block" or not reason or context.should_abort(): + return ContinuationDecision() + + allow_synthetic = await self._synthetic_continuation_allowed( + context, + last_message, + ) + continuation = await self._materialize_continuation( + context, + last_user, + last_message, + candidate_reason="turn_finish_hook", + content=reason, + agent=getattr(hook_user, "agent", None) or context.agent_name, + model={ + "providerID": context.provider_id, + "modelID": context.model_id, + }, + part_metadata={ + "turnFinishContinuation": True, + "stopHookActive": True, + "sourceAssistantMessageID": last_message.id, + }, + allow_synthetic=allow_synthetic, + ) + if continuation.reason == "turn_finish_hook": + context.stop_hook_active = True + return continuation + + @staticmethod + async def _synthetic_continuation_allowed( + context: Any, + last_message: MessageInfo, + ) -> bool: + """Protect all synthetic continuations with abort and step limits.""" + if context.should_abort(): + return False + + from flocks.agent.registry import Agent + from flocks.session.core.defaults import DEFAULT_MAX_TOOL_STEPS + + try: + agent = await Agent.get( + getattr(last_message, "agent", None) or context.agent_name + ) + except Exception as exc: + log.debug( + "session.continuation.agent_load_error", + {"session_id": context.session.id, "error": str(exc)}, + ) + agent = None + max_steps = ( + agent.steps + if agent is not None and getattr(agent, "steps", None) is not None + else DEFAULT_MAX_TOOL_STEPS + ) + return context.trace_step < max_steps + + async def _materialize_continuation( + self, + context: Any, + last_user: MessageInfo, + last_message: MessageInfo, + *, + candidate_reason: Optional[str] = None, + content: Optional[str] = None, + agent: Optional[str] = None, + model: Any = None, + provider: Optional[str] = None, + part_metadata: Optional[dict[str, Any]] = None, + event_metadata: Optional[dict[str, Any]] = None, + allow_synthetic: bool = True, + ) -> ContinuationDecision[MessageInfo]: + """Atomically let queued input preempt one synthetic candidate.""" + from flocks.session.session import Session + + try: + async with Session.lifecycle_lock(context.session.id): + if context.session_store: + messages = await context.session_store.get_messages() + else: + messages = await Message.list(context.session.id) + queued_user = await self.detect_queued_user_message( + context.session.id, + messages, + last_user.id, + last_message, + ) + if queued_user is not None: + selected = ContinuationDecision( + messages=(queued_user,), + reason="queued_message", + ) + elif ( + candidate_reason is None + or not content + or not allow_synthetic + or context.should_abort() + ): + selected = ContinuationDecision() + else: + create_kwargs = { + "session_id": context.session.id, + "role": MessageRole.USER, + "content": content, + "agent": agent or context.agent_name, + "model": model, + "synthetic": True, + "part_metadata": part_metadata or {}, + } + if provider is not None: + create_kwargs["provider"] = provider + continuation = await Message.create(**create_kwargs) + selected = ContinuationDecision( + messages=(continuation,), + reason=candidate_reason, + ) + except Exception as exc: + log.error( + "session.continuation.materialize_error", + {"session_id": context.session.id, "error": str(exc)}, + ) + return ContinuationDecision() + + if selected.should_continue: + await self._publish_continuation( + context, + selected, + event_metadata=event_metadata, + ) + return selected + + @staticmethod + async def _publish_continuation( + context: Any, + decision: ContinuationDecision[MessageInfo], + *, + event_metadata: Optional[dict[str, Any]] = None, + ) -> None: + """Publish the one continuation selected by the lifecycle boundary.""" + reason = decision.reason + message = decision.messages[0] + queued = reason == "queued_message" + turn_state = set_turn_state( + context.session.id, + step=context.step, + status="continued", + continue_reason=reason, + queued_message_detected=queued, + ) + message_id_key = { + "queued_message": "queuedUserMessageID", + "goal": "goalMessageID", + "turn_finish_hook": "turnFinishMessageID", + }.get(reason, "continuationMessageID") + await SessionEventSink.emit( + context.callbacks, + "turn.continued", + { + **turn_state.model_dump(by_alias=True), + message_id_key: message.id, + **(event_metadata or {}), + }, + ) + + +DEFAULT_CONTINUATION_POLICY = ContinuationPolicy(DEFAULT_MODEL_ROUTING_POLICY) diff --git a/flocks/agent/runtime/contracts.py b/flocks/session/runtime/contracts.py similarity index 56% rename from flocks/agent/runtime/contracts.py rename to flocks/session/runtime/contracts.py index 2359515da..fed280fda 100644 --- a/flocks/agent/runtime/contracts.py +++ b/flocks/session/runtime/contracts.py @@ -1,4 +1,4 @@ -"""Data contracts shared by the agent core and session host. +"""Data contracts shared by the agent loop and session runtime. The contracts in this module intentionally avoid importing session storage, server, CLI, provider, or tool-registry implementations. Session-specific @@ -8,6 +8,7 @@ from __future__ import annotations +import copy from dataclasses import dataclass, field from enum import Enum from types import MappingProxyType @@ -15,6 +16,47 @@ MessageT = TypeVar("MessageT") +ProviderMessageT = TypeVar("ProviderMessageT") + + +def _freeze(value: Any) -> Any: + """Recursively freeze request mappings and sequences.""" + if isinstance(value, Mapping): + return MappingProxyType( + {key: _freeze(item) for key, item in value.items()}, + ) + if isinstance(value, (list, tuple)): + return tuple(_freeze(item) for item in value) + if isinstance(value, set): + return frozenset(_freeze(item) for item in value) + if isinstance(value, (str, bytes, int, float, bool, type(None))): + return value + if hasattr(value, "model_copy"): + return value.model_copy(deep=True) + return copy.copy(value) + + +def _thaw(value: Any) -> Any: + """Return a provider-owned mutable copy of a frozen request value.""" + if isinstance(value, Mapping): + return {key: _thaw(item) for key, item in value.items()} + if isinstance(value, (tuple, list)): + return [_thaw(item) for item in value] + if isinstance(value, frozenset): + return {_thaw(item) for item in value} + if isinstance(value, (str, bytes, int, float, bool, type(None))): + return value + if hasattr(value, "model_copy"): + return value.model_copy(deep=True) + return copy.copy(value) + + +def _clone_provider_message(value: ProviderMessageT) -> ProviderMessageT: + if isinstance(value, (Mapping, list, tuple)): + return _thaw(_freeze(value)) + if hasattr(value, "model_copy"): + return value.model_copy(deep=True) + return copy.copy(value) @dataclass(frozen=True) @@ -25,24 +67,64 @@ class RuntimeModel: model_id: str +@dataclass(frozen=True) +class ModelRequest(Generic[ProviderMessageT]): + """Frozen provider request reused by retries of one model attempt.""" + + provider_id: str + model_id: str + messages: tuple[ProviderMessageT, ...] + tools: tuple[Mapping[str, Any], ...] + options: Mapping[str, Any] + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__( + self, + "messages", + tuple(_clone_provider_message(message) for message in self.messages), + ) + object.__setattr__( + self, + "tools", + tuple(_freeze(tool) for tool in self.tools), + ) + object.__setattr__(self, "options", _freeze(self.options)) + object.__setattr__(self, "metadata", _freeze(self.metadata)) + + def provider_messages(self) -> list[ProviderMessageT]: + """Return an isolated mutable copy for one provider invocation.""" + return [_clone_provider_message(message) for message in self.messages] + + def provider_tools(self) -> list[dict[str, Any]]: + """Return an isolated mutable tool-schema payload.""" + return [_thaw(tool) for tool in self.tools] + + def provider_options(self) -> dict[str, Any]: + """Return isolated provider options for one invocation.""" + return _thaw(self.options) + + @dataclass class AttemptEffects: """Observable effects accumulated during one provider attempt.""" + request_sent: bool = False received_chunk: bool = False observable_output_started: bool = False tool_execution_started: bool = False - durable_side_effect_possible: bool = False + tool_execution_completed: bool = False + externally_visible: bool = False @property def replay_safe(self) -> bool: """Return whether another provider may replay the logical request.""" - return not (self.observable_output_started or self.tool_execution_started or self.durable_side_effect_possible) + return not (self.observable_output_started or self.tool_execution_started) @dataclass(frozen=True) class FailoverDecision: - """Classification used by the session host's recovery policy.""" + """Classification used by the session runtime's recovery policy.""" eligible: bool reason: str @@ -57,9 +139,17 @@ class ToolCall: arguments: dict[str, Any] +class StepAction(str, Enum): + """Control-flow action produced by one model/tool step.""" + + CONTINUE = "continue" + STOP = "stop" + COMPACT = "compact" + + @dataclass class StepFailure: - """Failure returned by a step when host finalization is deferred.""" + """Failure returned by a step when runtime finalization is deferred.""" message: str error_data: dict[str, Any] @@ -74,12 +164,13 @@ class StepFailure: class StepResult: """Result of one model turn, including any tool execution.""" - action: str + action: StepAction | str content: str = "" tool_calls: list[ToolCall] = field(default_factory=list) error: Optional[str] = None usage: Optional[dict[str, int]] = None failure: Optional[StepFailure] = None + effective_model: Optional[RuntimeModel] = None @dataclass @@ -92,7 +183,6 @@ class AgentRunState(Generic[MessageT]): messages: list[MessageT] = field(default_factory=list) model_turn_index: int = 0 trace_step_offset: int = 0 - consumed_input_cursor: Optional[str] = None current_user_id: Optional[str] = None metadata: dict[str, Any] = field(default_factory=dict) @@ -126,7 +216,7 @@ def __post_init__(self) -> None: class TurnPreparationStatus(str, Enum): - """Host preparation result before the next model turn.""" + """Session preparation result before the next model turn.""" READY = "ready" CONTINUE = "continue" @@ -136,7 +226,7 @@ class TurnPreparationStatus(str, Enum): @dataclass(frozen=True) class ModelTurnPreparation(Generic[MessageT]): - """Result of host-owned preparation at a model-turn boundary.""" + """Result of session-owned preparation at a model-turn boundary.""" status: TurnPreparationStatus snapshot: Optional[ModelTurnSnapshot[MessageT]] = None @@ -149,7 +239,6 @@ class QueuedInputBatch(Generic[MessageT]): """New input made visible to the loop at a model-turn boundary.""" messages: tuple[MessageT, ...] = () - cursor: Optional[str] = None @dataclass(frozen=True) @@ -165,7 +254,7 @@ class ModelTurnBoundary(Generic[MessageT]): @dataclass(frozen=True) class ContinuationDecision(Generic[MessageT]): - """Host-owned continuation policy result consumed by the agent loop.""" + """Session-owned continuation policy result consumed by the outer loop.""" messages: tuple[MessageT, ...] = () reason: Optional[str] = None @@ -177,9 +266,10 @@ def should_continue(self) -> bool: class AgentRunStatus(str, Enum): - """Terminal states returned from the agent core to the session host.""" + """Terminal states returned from the agent core to the session runtime.""" COMPLETED = "completed" + INPUT_AVAILABLE = "input_available" RETRYABLE_FAILURE = "retryable_failure" CONTEXT_OVERFLOW = "context_overflow" FATAL_FAILURE = "fatal_failure" @@ -194,4 +284,4 @@ class AgentRunOutcome(Generic[MessageT]): state: AgentRunState[MessageT] last_message: Optional[MessageT] = None error: Optional[str] = None - failure: Optional[StepFailure] = None + step_result: Optional[StepResult] = None diff --git a/flocks/session/runtime/event_sink.py b/flocks/session/runtime/event_sink.py new file mode 100644 index 000000000..be16b9662 --- /dev/null +++ b/flocks/session/runtime/event_sink.py @@ -0,0 +1,92 @@ +"""Best-effort delivery of observable session runtime events.""" + +from __future__ import annotations + +from typing import Any, Optional + +from flocks.session.core.turn_state import set_turn_state +from flocks.utils.log import Log + + +log = Log.create(service="session.events") + + +class SessionEventSink: + """Forward runtime events without coupling control flow to observers.""" + + @staticmethod + async def emit( + callbacks: Any, + event_name: str, + payload: dict[str, Any], + ) -> None: + """Publish one event; observer failures never fail the agent run.""" + publish = getattr(callbacks, "event_publish_callback", None) + if publish is None: + return + try: + await publish(event_name, payload) + except Exception as exc: + log.debug( + "session.event.publish_failed", + {"event": event_name, "error": str(exc)}, + ) + + @classmethod + async def turn_stopped( + cls, + callbacks: Any, + session_id: str, + *, + step: int, + stop_reason: str, + ) -> None: + """Publish the terminal state of one logical turn.""" + turn_state = set_turn_state( + session_id, + step=step, + status="stopped", + stop_reason=stop_reason, + queued_message_detected=False, + ) + await cls.emit( + callbacks, + "turn.stopped", + turn_state.model_dump(by_alias=True), + ) + + @classmethod + async def session_status( + cls, + callbacks: Any, + session_id: str, + status: str, + ) -> None: + """Publish the current process-local session execution status.""" + await cls.emit( + callbacks, + "session.status", + {"sessionID": session_id, "status": {"type": status}}, + ) + + @classmethod + async def notice( + cls, + callbacks: Any, + session_id: str, + *, + level: str, + message: str, + details: Optional[dict[str, Any]] = None, + ) -> None: + """Publish a user-visible session notice.""" + await cls.emit( + callbacks, + "session.notice", + { + "sessionID": session_id, + "level": level, + "message": message, + "details": details or {}, + }, + ) diff --git a/flocks/session/runtime/model_policy.py b/flocks/session/runtime/model_policy.py new file mode 100644 index 000000000..10347af0f --- /dev/null +++ b/flocks/session/runtime/model_policy.py @@ -0,0 +1,397 @@ +"""Session-owned model routing and cross-model candidate policy.""" + +from __future__ import annotations + +import hashlib +import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Any, Optional + +from flocks.session.runtime.contracts import RuntimeModel +from flocks.provider.provider import Provider +from flocks.session.message import Message +from flocks.session.session import Session, is_model_auto_session_category +from flocks.utils.log import Log + + +log = Log.create(service="session.model_policy") + + +@dataclass +class AutoFailoverCooldown: + """Process-local starting candidate cooldown for automatic routing.""" + + model: RuntimeModel + primary: RuntimeModel + expires_at: float + reason: str + + +ModelValidator = Callable[..., Awaitable[tuple[bool, str]]] + + +class ModelRoutingPolicy: + """Own candidate discovery, per-turn routing, and failover cooldown state.""" + + def __init__(self) -> None: + self.cooldowns: dict[str, AutoFailoverCooldown] = {} + + def clear(self, session_id: str) -> None: + """Clear process-local routing state for one session.""" + self.cooldowns.pop(session_id, None) + + async def validate_runtime_model( + self, + provider_id: str, + model_id: str, + *, + config: Optional[Any] = None, + ) -> tuple[bool, str]: + """Validate a configured LLM candidate without a network health probe.""" + from flocks.config.config import Config + from flocks.provider.model_manager import get_model_manager + from flocks.provider.types import ModelType + + Provider._ensure_initialized() + config = config or await Config.get() + if provider_id in (getattr(config, "disabled_providers", None) or []): + return False, "provider_disabled" + enabled_providers = getattr(config, "enabled_providers", None) or [] + if enabled_providers and provider_id not in enabled_providers: + return False, "provider_disabled" + try: + await Provider.apply_config(config, provider_id=provider_id) + except Exception as exc: + log.warn( + "session.model.candidate_config_failed", + { + "provider_id": provider_id, + "model_id": model_id, + "error": str(exc), + }, + ) + return False, "provider_config_error" + + provider = Provider.get(provider_id) + if provider is None: + return False, "provider_not_found" + + model_manager = get_model_manager() + definition = model_manager.get_model(provider_id, model_id) + if definition is None: + return False, "model_not_found" + if getattr(definition, "model_type", None) != ModelType.LLM: + return False, "not_llm" + + setting = model_manager.get_setting(provider_id, model_id) + if setting is not None and not setting.enabled: + return False, "model_disabled" + if not provider.is_configured(): + return False, "provider_not_configured" + return True, "available" + + async def build_candidates( + self, + primary: RuntimeModel, + *, + route_seed: str, + preferred: Optional[RuntimeModel] = None, + config: Optional[Any] = None, + validate_model: Optional[ModelValidator] = None, + ) -> list[RuntimeModel]: + """Build a configured chain or stable automatic discovery chain.""" + from flocks.config.config import Config + from flocks.provider.model_manager import get_model_manager + from flocks.provider.types import ModelType + + validate_model = validate_model or self.validate_runtime_model + config = config or await Config.get() + await Provider.apply_config(config) + + configured_fallbacks = getattr(config, "fallback_providers", None) or [] + if configured_fallbacks: + candidates = [primary] + seen = {(primary.provider_id, primary.model_id)} + for index, raw in enumerate(configured_fallbacks): + provider_id = raw.get("provider_id") if isinstance(raw, dict) else raw.provider_id + model_id = raw.get("model_id") if isinstance(raw, dict) else raw.model_id + candidate = RuntimeModel(provider_id=provider_id, model_id=model_id) + identity = (candidate.provider_id, candidate.model_id) + if identity in seen: + continue + seen.add(identity) + + available, reason = await validate_model( + candidate.provider_id, + candidate.model_id, + config=config, + ) + if not available: + log.warn( + "session.model.fallback_skipped", + { + "provider_id": candidate.provider_id, + "model_id": candidate.model_id, + "configured_index": index, + "reason": reason, + }, + ) + continue + candidates.append(candidate) + return candidates + + definitions = get_model_manager().list_models( + model_type=ModelType.LLM, + enabled_only=True, + ) + discovered = {RuntimeModel(definition.provider_id, definition.id) for definition in definitions} + discovered.discard(primary) + + same_provider: list[RuntimeModel] = [] + other_providers: list[RuntimeModel] = [] + for candidate in sorted( + discovered, + key=lambda item: (item.provider_id, item.model_id), + ): + available, reason = await validate_model( + candidate.provider_id, + candidate.model_id, + config=config, + ) + if not available: + log.debug( + "session.model.fallback_skipped", + { + "provider_id": candidate.provider_id, + "model_id": candidate.model_id, + "reason": reason, + }, + ) + continue + if candidate.provider_id == primary.provider_id: + same_provider.append(candidate) + else: + other_providers.append(candidate) + + candidates = [primary] + for tier, pool in ( + ("same_provider", same_provider), + ("other_provider", other_providers), + ): + if not pool: + continue + selected = ( + preferred + if preferred is not None and preferred in pool + else self._stable_candidate_choice(pool, route_seed, tier) + ) + candidates.append(selected) + return candidates + + @staticmethod + def _stable_candidate_choice( + candidates: list[RuntimeModel], + route_seed: str, + tier: str, + ) -> RuntimeModel: + """Choose pseudo-randomly without process-randomized hash values.""" + ordered = sorted( + candidates, + key=lambda item: (item.provider_id, item.model_id), + ) + digest = hashlib.sha256(f"{route_seed}\0{tier}".encode("utf-8")).digest() + index = int.from_bytes(digest[:8], "big") % len(ordered) + return ordered[index] + + async def validate_auto_configuration(self) -> tuple[bool, str]: + """Validate that a newly selected Auto mode has a usable primary.""" + from flocks.config.config import Config + + default_llm = await Config.resolve_default_llm() + if not default_llm: + return False, "default_model_missing" + available, reason = await self.validate_runtime_model( + default_llm["provider_id"], + default_llm["model_id"], + ) + if not available: + return False, f"primary_{reason}" + return True, "available" + + def active_cooldown_model( + self, + session_id: str, + primary: RuntimeModel, + ) -> Optional[RuntimeModel]: + """Return a valid cooldown target for the current primary model.""" + cooldown = self.cooldowns.get(session_id) + if cooldown is None: + return None + if cooldown.expires_at <= time.monotonic() or cooldown.primary != primary: + self.cooldowns.pop(session_id, None) + return None + return cooldown.model + + def cooldown_candidate_index( + self, + session_id: str, + candidates: list[RuntimeModel], + ) -> int: + """Resolve the candidate index selected by an active cooldown.""" + if not candidates: + return 0 + cooldown_model = self.active_cooldown_model(session_id, candidates[0]) + if cooldown_model is None: + return 0 + try: + return candidates.index(cooldown_model) + except ValueError: + self.cooldowns.pop(session_id, None) + return 0 + + @staticmethod + def select_candidate(context: Any, index: int) -> None: + """Activate a candidate and invalidate model-specific runner caches.""" + candidate = context.model_candidates[index] + context.candidate_index = index + context.provider_id = candidate.provider_id + context.model_id = candidate.model_id + context.session.provider = candidate.provider_id + context.session.model = candidate.model_id + tool_loop_guard = context.step_static_cache.get("tool_loop_guard") + context.step_static_cache.clear() + if tool_loop_guard is not None: + context.step_static_cache["tool_loop_guard"] = tool_loop_guard + + async def reset_turn_candidates( + self, + context: Any, + primary: RuntimeModel, + user_message_id: str, + config: Any, + ) -> int: + """Rebuild and activate the model chain for one logical user turn.""" + configured = bool(getattr(config, "fallback_providers", None)) + if configured: + self.clear(context.session.id) + preferred = None + else: + preferred = self.active_cooldown_model(context.session.id, primary) + + context.model_candidates = await self.build_candidates( + primary, + route_seed=f"{context.session.id}:{user_message_id}", + preferred=preferred, + config=config, + ) + context.model_candidate_policy = "configured" if configured else "automatic" + context.auto_failover = True + next_index = ( + 0 + if configured + else self.cooldown_candidate_index( + context.session.id, + context.model_candidates, + ) + ) + self.select_candidate(context, next_index) + return next_index + + async def prepare_turn(self, context: Any, last_user: Any) -> bool: + """Synchronize model routing when a new real user turn begins.""" + if last_user.id == context.turn_user_id: + return False + + parts = await Message.parts(last_user.id, context.session.id) + if any(bool(getattr(part, "synthetic", False)) for part in parts): + return False + + if context.turn_user_id is None: + context.turn_user_id = last_user.id + if context.auto_failover and context.auto_failover_allowed: + from flocks.config.config import Config + + await self.reset_turn_candidates( + context, + context.model_candidates[0], + last_user.id, + config=await Config.get(), + ) + return True + + context.turn_user_id = last_user.id + persisted_session = await Session.get_by_id(context.session.id) + persisted_model_auto = bool( + persisted_session + and is_model_auto_session_category(getattr(persisted_session, "category", "user")) + and getattr(persisted_session, "model_auto", False) + ) + persisted_auto = persisted_model_auto and context.auto_failover_allowed + + user_model = getattr(last_user, "model", None) + user_provider_id = None + user_model_id = None + if isinstance(user_model, dict): + user_provider_id = user_model.get("providerID") or user_model.get("provider_id") + user_model_id = user_model.get("modelID") or user_model.get("model_id") + + if not persisted_auto: + context.auto_failover = False + if not persisted_model_auto: + self.clear(context.session.id) + context.auto_failover_allowed = False + provider_id = ( + getattr(persisted_session, "provider", None) + if Session.has_pinned_model(persisted_session) + else user_provider_id + ) or context.provider_id + model_id = ( + getattr(persisted_session, "model", None) + if Session.has_pinned_model(persisted_session) + else user_model_id + ) or context.model_id + context.model_candidates = [RuntimeModel(provider_id, model_id)] + context.model_candidate_policy = "fixed" + self.select_candidate(context, 0) + log.info( + "session.model.auto_disabled_for_turn", + { + "session_id": context.session.id, + "provider_id": provider_id, + "model_id": model_id, + }, + ) + return True + + from flocks.config.config import Config + + config = await Config.get() + previous = RuntimeModel(context.provider_id, context.model_id) + default_llm = await Config.resolve_default_llm() + primary = RuntimeModel( + provider_id=(default_llm or {}).get("provider_id") or user_provider_id or context.provider_id, + model_id=(default_llm or {}).get("model_id") or user_model_id or context.model_id, + ) + next_index = await self.reset_turn_candidates( + context, + primary, + last_user.id, + config=config, + ) + active = context.model_candidates[next_index] + log.info( + "session.model.auto_turn_reset", + { + "session_id": context.session.id, + "from_provider_id": previous.provider_id, + "from_model_id": previous.model_id, + "to_provider_id": active.provider_id, + "to_model_id": active.model_id, + "cooldown_active": next_index > 0, + }, + ) + return True + + +DEFAULT_MODEL_ROUTING_POLICY = ModelRoutingPolicy() diff --git a/flocks/session/runtime/session_turn.py b/flocks/session/runtime/session_turn.py new file mode 100644 index 000000000..666c9175e --- /dev/null +++ b/flocks/session/runtime/session_turn.py @@ -0,0 +1,1272 @@ +"""State and persistence boundary for one logical session turn. + +Implements model-turn preparation with support for: +- Message processing +- Tool execution +- Compaction +- Subtask handling +- Reminders +""" + +import asyncio +import time +from typing import Optional, List, Dict, Any, Callable, Awaitable, Literal +from dataclasses import dataclass, field +from datetime import datetime + +from flocks.session.runtime.contracts import ( + AgentRunState, + ModelTurnBoundary, + ModelTurnPreparation, + ModelTurnSnapshot, + QueuedInputBatch, + RuntimeModel, + StepAction, + StepResult, + TurnPreparationStatus, +) +from flocks.utils.log import Log +from flocks.utils.id import Identifier +from flocks.session.session import ( + Session, + SessionInfo, +) +from flocks.session.message import Message, MessageInfo, MessageRole +from flocks.session.runtime.event_sink import SessionEventSink +from flocks.session.core.status import SessionStatus, SessionStatusBusy +from flocks.session.core.task_utils import fire_and_forget +from flocks.session.core.turn_state import ( + set_turn_state, + set_context_state, +) +from flocks.session.lifecycle.compaction import ( + SessionCompaction, + CompactionPolicy, + build_compaction_policy, + run_compaction, +) +from flocks.session.lifecycle.compaction.compaction import _get_compaction_history +from flocks.session.prompt import SessionPrompt +from flocks.provider.provider import Provider + + +log = Log.create(service="session.loop") + + +MAX_OVERFLOW_COMPACTION_ATTEMPTS = 3 +POST_COMPACTION_COOLDOWN_STEPS = 2 + + +@dataclass +class LoopCallbacks: + """Callbacks for loop events""" + + on_step_start: Optional[Callable[[int], Awaitable[None]]] = None + on_step_end: Optional[Callable[[int], Awaitable[None]]] = None + on_text_delta: Optional[Callable[[str], Awaitable[None]]] = None + on_reasoning_delta: Optional[Callable[[str], Awaitable[None]]] = None + on_tool_start: Optional[ + Callable[[str, Dict[str, Any]], Awaitable[None]] + ] = None + on_tool_end: Optional[Callable[[str, Any], Awaitable[None]]] = None + on_permission_request: Optional[ + Callable[[Any], Awaitable[bool]] + ] = None + on_compaction: Optional[Callable[[], Awaitable[None]]] = None + on_error: Optional[Callable[[str], Awaitable[None]]] = None + on_reminder: Optional[Callable[[str], Awaitable[None]]] = None + # SSE event publishing callback (for TUI/WebUI real-time updates) + event_publish_callback: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None + + +@dataclass +class LoopResult: + """Result of loop execution""" + + action: str # "stop", "continue", "compact", "error", "queued" + last_message: Optional[MessageInfo] = None + error: Optional[str] = None + provider_id: Optional[str] = None + model_id: Optional[str] = None + metadata: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class SessionTurn: + """Own state and persistence boundaries for one logical user input. + + Supports: + - Message iteration + - Compaction triggers + - Subtask management + - Reminder injection + """ + + session: SessionInfo + provider_id: str + model_id: str + agent_name: str + callbacks: LoopCallbacks = field(default_factory=LoopCallbacks, repr=False) + step: int = 0 + abort_event: asyncio.Event = field(default_factory=asyncio.Event) + session_store: Optional[Any] = None + trace_step_offset: int = 0 + _current_step_task: Optional[asyncio.Task] = field(default=None, repr=False) + memory_bootstrap_data: Optional[Dict[str, Any]] = field(default=None, repr=False) + step_static_cache: Dict[str, Any] = field(default_factory=dict, repr=False) + overflow_compaction_attempts: int = 0 + tool_result_truncation_attempted: bool = False + last_compaction_step: Optional[int] = None + last_cleanup_step: Optional[int] = None + last_observed_prompt_tokens: int = 0 + auto_failover: bool = False + auto_failover_allowed: bool = False + model_candidates: List[RuntimeModel] = field(default_factory=list) + candidate_index: int = 0 + model_candidate_policy: Literal["fixed", "automatic", "configured"] = "automatic" + turn_user_id: Optional[str] = None + turn_additional_context: Optional[str] = None + stop_hook_active: bool = False + prepared_user_id: Optional[str] = None + prepared_messages: Optional[List[MessageInfo]] = field(default=None, repr=False) + session_start_pending: bool = False + model_policy: Optional[Any] = field(default=None, repr=False) + continuation_policy: Optional[Any] = field(default=None, repr=False) + state: AgentRunState[MessageInfo] = field(init=False, repr=False) + + def __post_init__(self) -> None: + self.reset() + + def reset(self) -> None: + """Start a fresh AgentLoop state for the current logical input.""" + self.state = AgentRunState[MessageInfo]( + session_id=self.session.id, + agent_name=self.agent_name, + active_model=RuntimeModel(self.provider_id, self.model_id), + model_turn_index=self.step, + trace_step_offset=self.trace_step_offset, + current_user_id=self.turn_user_id, + ) + + @property + def trace_step(self) -> int: + """Return the session-cumulative step number for observability.""" + return self.trace_step_offset + self.step + + @property + def aborted(self) -> bool: + """Return whether this turn was asked to stop.""" + return self.abort_event.is_set() + + def should_abort(self) -> bool: + """Keep the existing callable abort boundary for infrastructure.""" + return self.aborted + + def signal_abort(self) -> None: + """Stop the turn and cancel its active model step immediately.""" + self.abort_event.set() + task = self._current_step_task + if task is not None and not task.done(): + task.cancel() + + def _has_recent_compaction_cooldown(self) -> bool: + return ( + self.last_compaction_step is not None + and (self.step - self.last_compaction_step) <= POST_COMPACTION_COOLDOWN_STEPS + ) + + + async def finalize_failure( + self, + failure: Any, + last_user: MessageInfo, + ) -> None: + """Persist only the final Auto candidate failure.""" + if not failure.assistant_message_id: + assistant = await Message.create( + session_id=self.session.id, + role=MessageRole.ASSISTANT, + content="", + agent=getattr(last_user, "agent", None) or self.agent_name or "rex", + model_id=self.model_id, + provider_id=self.provider_id, + parent_id=last_user.id, + error=failure.error_data, + finish="error", + ) + failure.assistant_message_id = assistant.id + return + await Message.update( + self.session.id, + failure.assistant_message_id, + error=failure.error_data, + finish="error", + ) + + async def prepare_step( + self, + ) -> ModelTurnPreparation[MessageInfo]: + """Prepare one immutable model-turn snapshot from session state.""" + state = self.state + SessionStatus.set(self.session.id, SessionStatusBusy()) + self.step += 1 + state.model_turn_index = self.step + turn_state = set_turn_state( + self.session.id, + step=self.step, + status="started", + queued_message_detected=False, + ) + await SessionEventSink.emit( + self.callbacks, + "turn.started", + turn_state.model_dump(by_alias=True), + ) + log.info( + "loop.step", + {"session_id": self.session.id, "step": self.step}, + ) + if self.callbacks.on_step_start: + await self.callbacks.on_step_start(self.step) + + messages_started_at = asyncio.get_running_loop().time() + if self.prepared_messages is not None: + messages = self.prepared_messages + self.prepared_messages = None + elif self.session_store: + messages = await self.session_store.get_messages() + else: + messages = await Message.list(self.session.id) + log.debug( + "loop.messages_loaded", + { + "session_id": self.session.id, + "step": self.step, + "message_count": len(messages), + "duration_ms": int((asyncio.get_running_loop().time() - messages_started_at) * 1000), + }, + ) + if not messages: + log.info("loop.no_messages", {"session_id": self.session.id}) + await SessionEventSink.turn_stopped( + self.callbacks, + self.session.id, + step=self.step, + stop_reason="no_messages", + ) + return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) + + last_user: Optional[MessageInfo] = None + last_assistant: Optional[MessageInfo] = None + last_finished: Optional[MessageInfo] = None + tasks: List[tuple[str, Any]] = [] + scan_started_at = asyncio.get_running_loop().time() + for message in reversed(messages): + if last_user is None and message.role == MessageRole.USER: + last_user = message + if last_assistant is None and message.role == MessageRole.ASSISTANT: + last_assistant = message + if last_finished is None and message.role == MessageRole.ASSISTANT and getattr(message, "finish", None): + last_finished = message + if last_user is not None and last_finished is not None: + break + if last_finished is None: + for part in await Message.parts(message.id, self.session.id): + if part.type == "compaction": + tasks.append(("compaction", part)) + elif part.type == "subtask": + tasks.append(("subtask", part)) + log.debug( + "loop.message_scan_complete", + { + "session_id": self.session.id, + "step": self.step, + "task_count": len(tasks), + "duration_ms": int((asyncio.get_running_loop().time() - scan_started_at) * 1000), + }, + ) + + if last_user is None: + log.info( + "loop.no_user_message", + { + "session_id": self.session.id, + "message_count": len(messages), + "roles": [str(getattr(message, "role", "")) for message in messages[-5:]], + }, + ) + await SessionEventSink.turn_stopped( + self.callbacks, + self.session.id, + step=self.step, + stop_reason="no_user_message", + ) + return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) + + last_assistant_parts = await Message.parts(last_assistant.id, self.session.id) if last_assistant else [] + if self._should_exit(last_user, last_assistant, last_assistant_parts): + log.info( + "loop.exit_condition", + { + "session_id": self.session.id, + "last_user_id": last_user.id, + "last_assistant_id": (last_assistant.id if last_assistant else None), + "finish": last_assistant.finish if last_assistant else None, + "has_tool_parts": any(getattr(part, "type", None) == "tool" for part in last_assistant_parts), + }, + ) + return ModelTurnPreparation( + status=TurnPreparationStatus.COMPLETE, + last_message=last_assistant, + ) + + state.current_user_id = last_user.id + state.metadata["last_user"] = last_user + await self._prepare_memory() + self._schedule_title_generation(last_user, messages) + + if tasks: + task_preparation = await self._prepare_pending_task( + messages, + last_user, + tasks.pop(), + ) + if task_preparation is not None: + return task_preparation + + context_preparation = await self._prepare_context_window( + messages, + last_user, + last_finished, + ) + if context_preparation is not None: + return context_preparation + + active_model = RuntimeModel(self.provider_id, self.model_id) + state.active_model = active_model + state.messages = list(messages) + return ModelTurnPreparation( + status=TurnPreparationStatus.READY, + snapshot=ModelTurnSnapshot( + session_id=self.session.id, + agent_name=self.agent_name, + active_model=active_model, + model_turn_index=self.step, + trace_step=self.trace_step, + messages=tuple(messages), + last_user=last_user, + ), + ) + + async def commit_step( + self, + step_result: StepResult, + ) -> ModelTurnBoundary[MessageInfo]: + """Commit one executed step and expose its next control boundary.""" + if self.callbacks.on_step_end: + await self.callbacks.on_step_end(self.step) + if step_result.error and self.callbacks.on_error: + await self.callbacks.on_error(step_result.error) + + SessionStatus.set(self.session.id, SessionStatusBusy()) + if self.session_store: + post_messages = await self.session_store.get_messages() + else: + post_messages = await Message.list(self.session.id) + + last_user = self.state.metadata.get("last_user") + last_message = next( + ( + message + for message in reversed(post_messages) + if message.role == MessageRole.ASSISTANT + and ( + not self.auto_failover + or last_user is None + or getattr(message, "parentID", None) == last_user.id + ) + ), + None, + ) + + queued_user = None + if last_user is not None: + policy = self.continuation_policy + if policy is None: + from flocks.session.runtime.continuation_policy import ( + DEFAULT_CONTINUATION_POLICY, + ) + + policy = DEFAULT_CONTINUATION_POLICY + queued_user = await policy.detect_queued_user_message( + self.session.id, + post_messages, + last_user.id, + last_message, + ) + + if queued_user is not None: + turn_state = set_turn_state( + self.session.id, + step=self.step, + status="continued", + continue_reason="queued_message", + queued_message_detected=True, + ) + await SessionEventSink.emit( + self.callbacks, + "turn.continued", + { + **turn_state.model_dump(by_alias=True), + "queuedUserMessageID": queued_user.id, + }, + ) + log.info( + "session.turn.queued_input", + { + "session_id": self.session.id, + "queued_user_id": queued_user.id, + "last_assistant_id": ( + last_message.id if last_message else None + ), + }, + ) + elif step_result.action == StepAction.CONTINUE: + turn_state = set_turn_state( + self.session.id, + step=self.step, + status="continued", + continue_reason="tool_calls", + queued_message_detected=False, + ) + await SessionEventSink.emit( + self.callbacks, + "turn.continued", + turn_state.model_dump(by_alias=True), + ) + elif step_result.error: + await SessionEventSink.turn_stopped( + self.callbacks, + self.session.id, + step=self.step, + stop_reason=step_result.error, + ) + + return ModelTurnBoundary( + messages=tuple(post_messages), + last_message=last_message, + queued_inputs=QueuedInputBatch( + messages=(queued_user,) if queued_user is not None else (), + ), + ) + + async def has_late_input(self, processed_user_id: Optional[str]) -> bool: + """Return whether a newer persisted user input arrived before settle.""" + if processed_user_id is None: + return False + messages = await Message.list(self.session.id) + latest_user_id = next( + ( + message.id + for message in reversed(messages) + if getattr(message, "role", None) == "user" + ), + None, + ) + return latest_user_id is not None and latest_user_id != processed_user_id + + async def _prepare_memory(self) -> None: + """Load memory once before the first model turn.""" + if self.step != 1 or not self.session.memory_enabled or self.memory_bootstrap_data is not None: + return + try: + from flocks.memory.bootstrap import MemoryBootstrap + + self.memory_bootstrap_data = await MemoryBootstrap( + project_id=self.session.project_id, + ).bootstrap(load_daily=False) + log.info( + "loop.memory_bootstrap_done", + { + "session_id": self.session.id, + "has_main": (self.memory_bootstrap_data.get("main_memory") is not None), + }, + ) + except Exception as exc: + log.error("loop.memory_bootstrap_error", {"error": str(exc)}) + + def _schedule_title_generation( + self, + last_user: MessageInfo, + messages: List[MessageInfo], + ) -> None: + """Start optimistic first-turn title generation without blocking.""" + if self.step != 1 or self.auto_failover: + return + try: + from flocks.session.lifecycle.title import SessionTitle + + user_model = getattr(last_user, "model", None) + if isinstance(user_model, dict): + title_model_id = user_model.get("modelID", self.model_id) + title_provider_id = user_model.get( + "providerID", + self.provider_id, + ) + else: + title_model_id = self.model_id + title_provider_id = self.provider_id + fire_and_forget( + SessionTitle.ensure_title( + session_id=self.session.id, + model_id=title_model_id, + provider_id=title_provider_id, + messages=messages, + event_publish_callback=self.callbacks.event_publish_callback, + ), + label="title_generation", + name=f"title:{self.session.id}", + ) + except Exception as exc: + log.error("loop.title_generation.error", {"error": str(exc)}) + + async def _prepare_pending_task( + self, + messages: List[MessageInfo], + last_user: MessageInfo, + task: tuple[str, Any], + ) -> Optional[ModelTurnPreparation[MessageInfo]]: + """Finish persisted subtask or compaction work before the model turn.""" + task_type, task_part = task + if task_type == "subtask": + log.info( + "loop.subtask_detected", + {"session_id": self.session.id, "step": self.step}, + ) + await self._execute_subtask(last_user, task_part) + return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) + + log.info( + "loop.compaction_pending", + { + "session_id": self.session.id, + "step": self.step, + "auto": getattr(task_part, "auto", False), + }, + ) + if self.callbacks.on_compaction: + await self.callbacks.on_compaction() + + publish = self.callbacks.event_publish_callback + progress_callback = None + if publish is not None: + + async def progress_callback(stage: str, data: dict) -> None: + await publish( + "session.compaction_progress", + { + "sessionID": self.session.id, + "stage": stage, + "data": data, + }, + ) + + try: + compaction_result = await run_compaction( + self.session.id, + parent_message_id=last_user.id, + messages=messages, + provider_id=self.provider_id, + model_id=self.model_id, + auto=getattr(task_part, "auto", False), + event_publish_callback=publish, + status_after="busy", + policy=self._build_compaction_policy(), + progress_callback=progress_callback, + ) + if compaction_result == "stop": + log.error( + "loop.compaction_failed", + {"session_id": self.session.id}, + ) + if self.callbacks.on_error: + await self.callbacks.on_error("Compaction failed") + return ModelTurnPreparation( + status=TurnPreparationStatus.COMPLETE, + ) + if compaction_result == "skipped": + log.info( + "loop.manual_compaction_skipped", + {"session_id": self.session.id, "step": self.step}, + ) + return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) + except Exception as exc: + log.error("loop.compaction_error", {"error": str(exc)}) + if self.callbacks.on_error: + await self.callbacks.on_error(f"Compaction error: {exc}") + return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) + + async def _prepare_context_window( + self, + messages: List[MessageInfo], + last_user: MessageInfo, + last_finished: Optional[MessageInfo], + ) -> Optional[ModelTurnPreparation[MessageInfo]]: + """Recover a near-overflow context before the next model turn.""" + if last_finished is None or getattr(last_finished, "summary", False): + return None + + model_context, model_output, model_input = Provider.resolve_model_info( + self.provider_id, + self.model_id, + ) + if model_context <= 0: + return None + + policy = CompactionPolicy.from_model( + context_window=model_context, + max_output_tokens=model_output or 4096, + max_input_tokens=model_input, + ) + tokens = self._normalise_token_usage(last_finished) + input_tokens = tokens.get("input", 0) + cache = tokens.get("cache") or {} + cache_read = cache.get("read", 0) if isinstance(cache, dict) else 0 + output_tokens = tokens.get("output", 0) + reported_total = input_tokens + cache_read + output_tokens + if reported_total > 0: + self.last_observed_prompt_tokens = reported_total + log.info( + "loop.tokens_decision", + { + "session_id": self.session.id, + "source": "observed", + "effective_tokens": input_tokens + cache_read, + "overflow_threshold": policy.overflow_threshold, + }, + ) + else: + estimated_tokens = await SessionPrompt.estimate_full_context_tokens( + self.session.id, + messages, + policy=policy, + ) + tokens = { + "input": estimated_tokens, + "output": 0, + "cache": {"read": 0, "write": 0}, + } + log.info( + "loop.tokens_decision", + { + "session_id": self.session.id, + "source": "estimated", + "effective_tokens": estimated_tokens, + "message_count": len(messages), + "overflow_threshold": policy.overflow_threshold, + }, + ) + + try: + cache = tokens.get("cache") or {} + current_input_tokens = tokens.get("input", 0) + (cache.get("read", 0) if isinstance(cache, dict) else 0) + recent_compaction = self._has_recent_compaction_cooldown() + near_overflow = current_input_tokens >= policy.preemptive_threshold + if near_overflow and self.last_cleanup_step != self.step: + cleanup_result = await self._prepare_tool_result_cleanup( + model_context, + policy, + current_input_tokens, + recent_compaction, + ) + if cleanup_result is not None: + return cleanup_result + + is_overflow = await SessionCompaction.is_overflow( + tokens=tokens, + model_context=model_context, + policy=policy, + ) + if not is_overflow: + return None + + log.info( + "loop.context_overflow_detected", + { + "session_id": self.session.id, + "step": self.step, + "tokens": tokens, + "tier": policy.tier.value, + "overflow_compaction_attempts": (self.overflow_compaction_attempts), + }, + ) + if self.overflow_compaction_attempts >= MAX_OVERFLOW_COMPACTION_ATTEMPTS: + await self._report_compaction_exhausted( + tokens, + ) + return ModelTurnPreparation( + status=TurnPreparationStatus.COMPLETE, + ) + + if not self.tool_result_truncation_attempted: + self.tool_result_truncation_attempted = True + try: + truncation_count = await SessionCompaction.truncate_oversized_tool_outputs( + self.session.id, + context_window_tokens=model_context, + ) + if truncation_count > 0: + log.info( + "loop.oversized_tool_truncated", + { + "session_id": self.session.id, + "truncated": truncation_count, + }, + ) + estimated_tokens = await SessionPrompt.estimate_full_context_tokens( + self.session.id, + messages, + policy=policy, + ) + still_overflow = await SessionCompaction.is_overflow( + tokens={ + "input": estimated_tokens, + "output": 0, + "cache": {"read": 0, "write": 0}, + }, + model_context=model_context, + policy=policy, + ) + if not still_overflow: + log.info( + "loop.overflow_resolved_by_truncation", + {"session_id": self.session.id}, + ) + return ModelTurnPreparation( + status=TurnPreparationStatus.CONTINUE, + ) + except Exception as exc: + log.warn( + "loop.oversized_truncation_error", + {"session_id": self.session.id, "error": str(exc)}, + ) + + return await self._prepare_full_compaction( + messages, + last_user, + policy, + ) + except Exception as exc: + log.error( + "loop.compaction_overflow_check_error", + {"error": str(exc)}, + ) + return None + + @staticmethod + def _normalise_token_usage(message: MessageInfo) -> Dict[str, Any]: + """Normalise provider token usage into the legacy mapping shape.""" + raw_tokens = getattr(message, "tokens", None) + if not raw_tokens: + return {} + if isinstance(raw_tokens, dict): + return raw_tokens + if hasattr(raw_tokens, "model_dump"): + return raw_tokens.model_dump() + if hasattr(raw_tokens, "__dict__"): + return vars(raw_tokens) + return {} + + async def _prepare_tool_result_cleanup( + self, + model_context: int, + policy: CompactionPolicy, + current_input_tokens: int, + recent_compaction: bool, + ) -> Optional[ModelTurnPreparation[MessageInfo]]: + """Apply the cheap tool-result cleanup before full compaction.""" + try: + truncation_count = await SessionCompaction.truncate_oversized_tool_outputs( + self.session.id, + context_window_tokens=model_context, + ) + self.last_cleanup_step = self.step + if truncation_count <= 0: + return None + + set_context_state( + self.session.id, + tool_results_compacted=True, + last_compaction_step=self.last_compaction_step, + last_compaction_reason="pre_compact_cleanup", + ) + await SessionEventSink.emit( + self.callbacks, + "context.compacted", + { + "sessionID": self.session.id, + "step": self.step, + "reason": "pre_compact_cleanup", + "truncatedToolResults": truncation_count, + "cooldownActive": recent_compaction, + }, + ) + log.info( + "loop.pre_compact_cleanup_applied", + { + "session_id": self.session.id, + "step": self.step, + "truncated": truncation_count, + "preemptive_threshold": policy.preemptive_threshold, + "input_tokens": current_input_tokens, + "cooldown_active": recent_compaction, + }, + ) + turn_state = set_turn_state( + self.session.id, + step=self.step, + status="continued", + continue_reason="pre_compact_cleanup", + queued_message_detected=False, + ) + await SessionEventSink.emit( + self.callbacks, + "turn.continued", + turn_state.model_dump(by_alias=True), + ) + return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) + except Exception as exc: + log.warn( + "loop.pre_compact_cleanup_error", + {"session_id": self.session.id, "error": str(exc)}, + ) + return None + + async def _report_compaction_exhausted( + self, + tokens: Dict[str, Any], + ) -> None: + """Surface whether exhaustion came from context or provider health.""" + history = _get_compaction_history(self.session.id) + provider_error = history.summary_last_error + in_cooldown = history.summary_cooldown_until > 0 and history.summary_cooldown_until > time.monotonic() + cooldown_seconds = max( + 0, + round(history.summary_cooldown_until - time.monotonic()), + ) + if in_cooldown or provider_error: + notice = ( + "摘要模型暂时不可用，上下文压缩跳过了本轮压缩。" + + (f"冷却剩余约 {cooldown_seconds} 秒，" if in_cooldown else "") + + "建议稍后继续，或切换到其他模型重试。" + ) + error = ( + "Compaction skipped: summary provider unavailable " + f"({provider_error or 'cooldown active'})." + + (f" Cooldown expires in ~{cooldown_seconds}s." if in_cooldown else "") + + " Wait for the provider to recover or switch models." + ) + else: + notice = "当前任务上下文过重，已经多次 compact 仍接近上限。建议收敛工具输出、缩小搜索范围，或开启新会话。" + error = ( + "Context overflow: prompt too large for the model after " + f"{self.overflow_compaction_attempts} compaction attempts. " + "Try starting a new session or use a larger-context model." + ) + + await SessionEventSink.notice( + self.callbacks, + self.session.id, + level="warning", + message=notice, + details={ + "attempts": self.overflow_compaction_attempts, + "maxAttempts": MAX_OVERFLOW_COMPACTION_ATTEMPTS, + "tokens": tokens, + "providerError": provider_error or None, + "cooldownRemainingSeconds": (cooldown_seconds if in_cooldown else 0), + }, + ) + log.error( + "loop.overflow_compaction_exhausted", + { + "session_id": self.session.id, + "attempts": self.overflow_compaction_attempts, + "max": MAX_OVERFLOW_COMPACTION_ATTEMPTS, + "tokens": tokens, + "in_cooldown": in_cooldown, + "provider_error": provider_error or None, + }, + ) + if self.callbacks.on_error: + await self.callbacks.on_error(error) + + async def _prepare_full_compaction( + self, + messages: List[MessageInfo], + last_user: MessageInfo, + policy: CompactionPolicy, + ) -> ModelTurnPreparation[MessageInfo]: + """Run full compaction and request preparation to reload the session.""" + self.overflow_compaction_attempts += 1 + if self.overflow_compaction_attempts >= 2: + await SessionEventSink.notice( + self.callbacks, + self.session.id, + level="info", + message=("本轮上下文持续接近模型上限，系统将优先尝试压缩历史工具输出。"), + details={ + "attempt": self.overflow_compaction_attempts, + "threshold": policy.overflow_threshold, + "buffer": policy.overflow_buffer, + }, + ) + log.warn( + "loop.overflow_compaction_attempt", + { + "session_id": self.session.id, + "attempt": self.overflow_compaction_attempts, + "max": MAX_OVERFLOW_COMPACTION_ATTEMPTS, + }, + ) + if self.callbacks.on_compaction: + await self.callbacks.on_compaction() + await SessionCompaction.prune(self.session.id, policy=policy) + + publish = self.callbacks.event_publish_callback + progress_callback = None + if publish is not None: + + async def progress_callback(stage: str, data: dict) -> None: + await publish( + "session.compaction_progress", + { + "sessionID": self.session.id, + "stage": stage, + "data": data, + }, + ) + + result = await run_compaction( + self.session.id, + parent_message_id=last_user.id, + messages=messages, + provider_id=self.provider_id, + model_id=self.model_id, + auto=True, + event_publish_callback=publish, + status_after="busy", + policy=policy, + progress_callback=progress_callback, + ) + if result == "stop": + log.error( + "loop.compaction_failed", + {"session_id": self.session.id}, + ) + if self.callbacks.on_error: + await self.callbacks.on_error("Compaction failed") + return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) + if result == "skipped": + log.info( + "loop.compaction_skipped", + {"session_id": self.session.id, "step": self.step}, + ) + else: + self.last_compaction_step = self.step + set_context_state( + self.session.id, + compaction_performed=True, + last_compaction_step=self.step, + last_compaction_reason="full_compaction", + ) + await SessionEventSink.emit( + self.callbacks, + "context.compacted", + { + "sessionID": self.session.id, + "step": self.step, + "reason": "full_compaction", + "attempt": self.overflow_compaction_attempts, + "cooldownUntilStep": (self.step + POST_COMPACTION_COOLDOWN_STEPS), + }, + ) + return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) + + def _build_compaction_policy(self) -> CompactionPolicy: + """ + Construct a CompactionPolicy from the current model's info. + + Falls back to ``CompactionPolicy.default()`` when the model info + cannot be resolved (e.g. unknown provider or missing context_window). + """ + return build_compaction_policy(self.provider_id, self.model_id) + + @staticmethod + def _should_exit( + last_user: MessageInfo, + last_assistant: Optional[MessageInfo], + last_assistant_parts: Optional[List[Any]] = None, + ) -> bool: + """ + Check if loop should exit + + Ported from original exit logic: + - Exit if assistant has responded with finish != tool-calls + - Exit if assistant message is after user message + """ + if not last_assistant: + return False + + if any(getattr(part, "type", None) == "tool" for part in (last_assistant_parts or [])): + return False + + # Check finish reason + if last_assistant.finish: + if last_assistant.finish not in ("tool-calls", "unknown", "summary"): + # Assistant finished with stop/error/etc + if last_user.id < last_assistant.id: + # Assistant responded after user + return True + + return False + + async def _check_reminders( + self, + messages: List[MessageInfo], + ) -> None: + """ + Check and inject reminders (P1 feature) + + Reminders are system messages injected periodically to: + - Remind agent of task goals + - Prevent drift from original intent + - Nudge towards completion + """ + from flocks.session.features.reminders import SessionReminders, ReminderContext + + # Calculate elapsed time + if messages: + first_msg = messages[0] + if hasattr(first_msg, "time") and hasattr(first_msg.time, "created"): + first_time = first_msg.time.created + current_time = int(datetime.now().timestamp() * 1000) + elapsed_ms = current_time - first_time + else: + elapsed_ms = 0 + else: + elapsed_ms = 0 + + # Extract original task + original_task = await SessionReminders.extract_original_task(messages) + + # Create reminder context + reminder_ctx = ReminderContext( + session_id=self.session.id, + step_count=self.step, + message_count=len(messages), + elapsed_ms=elapsed_ms, + original_task=original_task, + ) + + # Check if reminder should be injected + if SessionReminders.should_remind(self.session.id, reminder_ctx): + # Create and inject reminder + reminder_msg = await SessionReminders.create_reminder( + self.session.id, + reminder_ctx, + ) + + if reminder_msg and self.callbacks.on_reminder: + await self.callbacks.on_reminder( + await Message.get_text_content(reminder_msg), + ) + + async def _execute_subtask( + self, + last_user: MessageInfo, + task_part: Any, + ) -> None: + """ + Execute subtask (matching TUI lines 316-481) + + 完全匹配 TUI 的 subtask 执行流程: + 1. 创建 assistant message + 2. 创建 tool part (Task tool) + 3. 执行 Task tool + 4. 更新 part 状态 + 5. 创建 synthetic user message + """ + from flocks.tool.registry import ToolRegistry + from flocks.agent.registry import Agent + + # Extract subtask information from part + agent_name = getattr(task_part, "agent", "hephaestus") + prompt = getattr(task_part, "prompt", "") + description = getattr(task_part, "description", "") + command = getattr(task_part, "command", None) + model_info = getattr(task_part, "model", None) + + # Get agent + agent = await Agent.get(agent_name) or await Agent.get("rex") + + # Determine model + if model_info: + provider_id = model_info.get("providerID", self.provider_id) + model_id = model_info.get("modelID", self.model_id) + else: + provider_id = self.provider_id + model_id = self.model_id + + # Create assistant message for subtask + assistant_msg = await Message.create( + session_id=self.session.id, + role=MessageRole.ASSISTANT, + content="", + agent=agent_name, + model=model_id, + provider=provider_id, + parent_id=last_user.id, + ) + + # Create tool part for Task + tool_call_id = Identifier.create("call") + from flocks.session.message import ToolPart, ToolStateRunning + + tool_part = ToolPart( + id=Identifier.ascending("part"), + sessionID=self.session.id, + messageID=assistant_msg.id, + type="tool", + callID=tool_call_id, + tool="task", + state=ToolStateRunning( + status="running", + input={ + "prompt": prompt, + "description": description, + "subagent_type": agent_name, + "command": command, + }, + time={"start": int(datetime.now().timestamp() * 1000)}, + ), + ) + + # Add part to message + await Message.add_part(self.session.id, assistant_msg.id, tool_part) + + # Get Task tool + task_tool = ToolRegistry.get("task") + if not task_tool: + log.error("loop.subtask.task_tool_not_found", {"session_id": self.session.id}) + return + + # Execute Task tool + task_args = { + "prompt": prompt, + "description": description, + "subagent_type": agent_name, + "command": command, + } + + # Create tool context + from flocks.tool.registry import ToolContext + + tool_ctx = ToolContext( + session_id=self.session.id, + message_id=assistant_msg.id, + agent=agent_name, + abort_event=self.abort_event, + ) + + execution_error: Optional[Exception] = None + result = None + + try: + result = await task_tool.execute(tool_ctx, **task_args) + except Exception as e: + execution_error = e + log.error( + "loop.subtask.execution_failed", + { + "error": str(e), + "agent": agent_name, + "description": description, + }, + ) + + # Update message finish + await Message.update(self.session.id, assistant_msg.id, finish="tool-calls") + + # Update tool part status + from flocks.session.message import ToolStateCompleted, ToolStateError + + if result: + # Create completed state + completed_state = ToolStateCompleted( + status="completed", + input={ + "prompt": prompt, + "description": description, + "subagent_type": agent_name, + "command": command, + }, + output=result.output if hasattr(result, "output") else str(result), + title=result.title if hasattr(result, "title") else None, + metadata=result.metadata if hasattr(result, "metadata") else {}, + time={ + "start": tool_part.state.time.get("start"), + "end": int(datetime.now().timestamp() * 1000), + }, + ) + await Message.update_part( + session_id=self.session.id, + message_id=assistant_msg.id, + part_id=tool_part.id, + state=completed_state, + ) + else: + # Create error state + error_msg = str(execution_error) if execution_error else "Tool execution failed" + error_state = ToolStateError( + status="error", + error=f"Tool execution failed: {error_msg}", + time={ + "start": tool_part.state.time.get("start"), + "end": int(datetime.now().timestamp() * 1000), + }, + metadata={}, + input={ + "prompt": prompt, + "description": description, + "subagent_type": agent_name, + "command": command, + }, + ) + await Message.update_part( + session_id=self.session.id, + message_id=assistant_msg.id, + part_id=tool_part.id, + state=error_state, + ) + + # Create synthetic user message (matching TUI lines 457-478) + # This prevents reasoning models from erroring due to missing user messages + synthetic_user_msg = await Message.create( + session_id=self.session.id, + role=MessageRole.USER, + content="Summarize the task tool output above and continue with your task.", + agent=last_user.agent if hasattr(last_user, "agent") else agent_name, + model=last_user.model if hasattr(last_user, "model") else model_id, + provider=last_user.provider if hasattr(last_user, "provider") else provider_id, + synthetic=True, + ) + + log.info( + "loop.subtask.completed", + { + "session_id": self.session.id, + "agent": agent_name, + "success": result is not None, + }, + ) diff --git a/flocks/session/runner.py b/flocks/session/runtime/step_engine.py similarity index 88% rename from flocks/session/runner.py rename to flocks/session/runtime/step_engine.py index a71c10a94..e58095e26 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runtime/step_engine.py @@ -1,13 +1,4 @@ -""" -Session runner module. - -Core session execution logic including: -- Session loop (message processing) -- Tool resolution and execution -- LLM interaction with tool support - -Implements session/prompt.ts SessionPrompt namespace pattern. -""" +"""Own one complete model/tool step from frozen input to StepResult.""" import asyncio import copy @@ -16,21 +7,31 @@ import re import sys import time +from collections.abc import Mapping from datetime import datetime -from typing import Optional, Dict, Any, List, Callable, Awaitable, Tuple -from dataclasses import dataclass +from typing import Optional, Dict, Any, List, Tuple +from dataclasses import replace import httpcore import httpx -from flocks.agent.runtime.contracts import ( +from flocks.session.runtime.contracts import ( AttemptEffects, FailoverDecision, + ModelRequest, + ModelTurnSnapshot, + RuntimeModel, StepFailure, StepResult, ToolCall, ) -from flocks.agent.runtime.ports import ExternalRuntimePorts +from flocks.session.runtime.event_sink import SessionEventSink +from flocks.session.runtime.model_policy import ( + DEFAULT_MODEL_ROUTING_POLICY, + AutoFailoverCooldown, + ModelRoutingPolicy, +) +from flocks.session.runtime.session_turn import LoopCallbacks from flocks.utils.log import Log from flocks.utils.id import Identifier from flocks.session.session import Session, SessionInfo @@ -92,7 +93,7 @@ from flocks.session.plan_file import session_plan_file -log = Log.create(service="session.runner") +log = Log.create(service="session.step_engine") TOOL_RESULT_CHAR_BUDGET_RATIO = 0.70 TOOL_RESULT_TURN_BUDGET_RATIO = 0.35 @@ -122,6 +123,8 @@ def _annotate_with_provider_version(tool_info: Any, description: Optional[str]) return f"{base.rstrip()}\n\n{note}" TOOL_RESULT_MIN_TURN_BUDGET = 4_000 TOOL_RESULT_PREVIEW_CHARS = 160 +RATE_LIMIT_COOLDOWN_SECONDS = 60.0 +CHAIN_EXHAUSTION_COOLDOWN_SECONDS = 5.0 # Maximum seconds to wait for the *first* chunk from the LLM stream. # If the model never starts responding, the stream times out and the session @@ -224,66 +227,40 @@ def _find_retryable_transport_exception(exception: Exception) -> Optional[Except LlmAttemptState = AttemptEffects -@dataclass -class RunnerCallbacks: - """Callbacks for runner events.""" - on_step_start: Optional[Callable[[int], Awaitable[None]]] = None - on_step_end: Optional[Callable[[int], Awaitable[None]]] = None - on_text_delta: Optional[Callable[[str], Awaitable[None]]] = None - on_reasoning_delta: Optional[Callable[[str], Awaitable[None]]] = None - on_tool_start: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None - on_tool_end: Optional[Callable[[str, ToolResult], Awaitable[None]]] = None - on_permission_request: Optional[Callable[[Any], Awaitable[bool]]] = None - on_error: Optional[Callable[[str], Awaitable[None]]] = None - # SSE event publishing callback (for TUI/WebUI real-time updates) - event_publish_callback: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None +class StepCancelled(Exception): + """Signal that the user cancelled the active session step.""" -class SessionRunner: - """ - Core session runner. - - Manages the session execution loop: - 1. Get messages from session - 2. Check if LLM response is needed - 3. Call LLM with tools - 4. Execute tool calls - 5. Loop until complete - - Implements SessionPrompt.loop() - """ - - # Class-level state for active sessions - _active_sessions: Dict[str, 'SessionRunner'] = {} - +class StepEngine: + """Own one complete model/tool step, including retries and failover.""" + def __init__( self, session: SessionInfo, provider_id: Optional[str] = None, model_id: Optional[str] = None, agent_name: Optional[str] = None, - callbacks: Optional[RunnerCallbacks] = None, + callbacks: Optional[LoopCallbacks] = None, abort_event: Optional[asyncio.Event] = None, - session_ctx: Optional[Any] = None, # SessionContext interface + session_store: Optional[Any] = None, memory_bootstrap_data: Optional[Dict[str, Any]] = None, static_cache: Optional[Dict[str, Any]] = None, defer_step_errors: bool = False, failover_available: bool = False, turn_additional_context: Optional[str] = None, session_start_pending: bool = False, - runtime_ports: Optional[ExternalRuntimePorts] = None, ): self.session = session from flocks.session.core.defaults import fallback_provider_id, fallback_model_id self.provider_id = provider_id or fallback_provider_id() self.model_id = model_id or fallback_model_id() self.agent_name = agent_name or "rex" - self.callbacks = callbacks or RunnerCallbacks() + self.callbacks = callbacks or LoopCallbacks() self._abort = asyncio.Event() self._external_abort = abort_event # External abort event (e.g. from SessionLoop) self._step = 0 self._recent_tool_calls: List[tuple[str, str]] = [] # Track recent (tool_name, args_json) for doom loop - self.session_ctx = session_ctx # SessionContext interface for decoupled access + self.session_store = session_store self._memory_bootstrap_data: Optional[Dict[str, Any]] = memory_bootstrap_data self._static_cache = static_cache if static_cache is not None else {} self._defer_step_errors = defer_step_errors @@ -292,26 +269,245 @@ def __init__( self._session_start_pending = session_start_pending self._session_start_fired = False self._attempt_state = LlmAttemptState() - if runtime_ports is None: - from flocks.session.runtime_adapters import ( - create_default_runtime_ports, + self._hooked_model_requests: Dict[ + str, + Tuple[ModelRequest[ChatMessage], bool], + ] = {} + self._turn: Optional[Any] = None + self._model_policy: ModelRoutingPolicy = DEFAULT_MODEL_ROUTING_POLICY + + @classmethod + def from_turn( + cls, + turn: Any, + model_policy: Optional[ModelRoutingPolicy] = None, + ) -> "StepEngine": + """Create the production engine for one stateful ``SessionTurn``.""" + engine = cls( + session=turn.session, + provider_id=turn.provider_id, + model_id=turn.model_id, + agent_name=turn.agent_name, + abort_event=turn.abort_event, + callbacks=turn.callbacks, + session_store=turn.session_store, + memory_bootstrap_data=turn.memory_bootstrap_data, + static_cache=turn.step_static_cache, + defer_step_errors=turn.auto_failover, + failover_available=( + turn.auto_failover + and turn.candidate_index + 1 < len(turn.model_candidates) + ), + turn_additional_context=turn.turn_additional_context, + session_start_pending=turn.session_start_pending, + ) + engine._turn = turn + engine._model_policy = ( + model_policy + or turn.model_policy + or DEFAULT_MODEL_ROUTING_POLICY + ) + return engine + + async def run( + self, + snapshot: ModelTurnSnapshot[MessageInfo], + ) -> StepResult: + """Execute a replay-safe snapshot across the active model chain.""" + turn = self._require_turn() + self._step_agent = None + self._frozen_tool_request = None + while True: + active_model = RuntimeModel(turn.provider_id, turn.model_id) + result = await self._run_candidate( + replace(snapshot, active_model=active_model), ) + result.effective_model = active_model + failure = result.failure + if not turn.auto_failover or failure is None: + return result - runtime_ports = create_default_runtime_ports() - self._runtime_ports = runtime_ports + next_index = turn.candidate_index + 1 + has_next = next_index < len(turn.model_candidates) + if ( + not failure.allow_fallback + or not failure.attempt_state.replay_safe + or not has_next + ): + self._record_chain_exhaustion(failure, has_next) + await turn.finalize_failure(failure, snapshot.last_user) + return result - @property - def _ports(self) -> ExternalRuntimePorts: - """Return injected ports, lazily supporting legacy test instances.""" - ports = getattr(self, "_runtime_ports", None) - if ports is None: - from flocks.session.runtime_adapters import ( - create_default_runtime_ports, + if not await self._remove_failed_attempt(failure): + await turn.finalize_failure(failure, snapshot.last_user) + return result + + await self._switch_candidate(next_index, failure.reason) + + def _require_turn(self) -> Any: + if self._turn is None: + raise RuntimeError( + "StepEngine.run() requires StepEngine.from_turn()", ) + return self._turn - ports = create_default_runtime_ports() - self._runtime_ports = ports - return ports + async def _run_candidate( + self, + snapshot: ModelTurnSnapshot[MessageInfo], + ) -> StepResult: + """Execute one candidate without introducing another runner object.""" + turn = self._require_turn() + self.provider_id = snapshot.active_model.provider_id + self.model_id = snapshot.active_model.model_id + self._step = snapshot.trace_step + self._defer_step_errors = turn.auto_failover + self._failover_available = ( + turn.auto_failover + and turn.candidate_index + 1 < len(turn.model_candidates) + ) + self._turn_additional_context = turn.turn_additional_context + self._session_start_pending = turn.session_start_pending + + task = asyncio.create_task( + self._process_step(list(snapshot.messages), snapshot.last_user), + ) + turn._current_step_task = task + started_at = asyncio.get_running_loop().time() + try: + result = await task + if self._session_start_fired: + turn.session_start_pending = False + return result + except asyncio.CancelledError as exc: + if turn.aborted: + raise StepCancelled from exc + raise + finally: + turn._current_step_task = None + log.debug( + "session.step.complete", + { + "session_id": turn.session.id, + "step": turn.step, + "duration_ms": int( + ( + asyncio.get_running_loop().time() + - started_at + ) + * 1000 + ), + }, + ) + + def _record_chain_exhaustion( + self, + failure: Any, + has_next: bool, + ) -> None: + turn = self._require_turn() + if not ( + turn.model_candidate_policy == "automatic" + and failure.allow_fallback + and failure.attempt_state.replay_safe + and not has_next + and turn.candidate_index > 0 + and failure.reason not in {"rate_limit", "billing"} + ): + return + + expires_at = time.monotonic() + CHAIN_EXHAUSTION_COOLDOWN_SECONDS + existing = self._model_policy.cooldowns.get(turn.session.id) + if existing and existing.expires_at > expires_at: + return + self._model_policy.cooldowns[turn.session.id] = AutoFailoverCooldown( + model=turn.model_candidates[turn.candidate_index], + primary=turn.model_candidates[0], + expires_at=expires_at, + reason="chain_exhausted", + ) + + async def _remove_failed_attempt(self, failure: Any) -> bool: + turn = self._require_turn() + message_id = failure.assistant_message_id + if not message_id: + return True + try: + deleted = await Message.delete(turn.session.id, message_id) + except Exception as exc: + deleted = False + log.error( + "session.model.fallback_cleanup_failed", + { + "session_id": turn.session.id, + "message_id": message_id, + "error": str(exc), + }, + ) + if not deleted: + return False + await SessionEventSink.emit( + turn.callbacks, + "message.removed", + { + "sessionID": turn.session.id, + "messageID": message_id, + }, + ) + return True + + async def _switch_candidate(self, next_index: int, reason: str) -> None: + turn = self._require_turn() + previous = turn.model_candidates[turn.candidate_index] + next_candidate = turn.model_candidates[next_index] + if turn.model_candidate_policy == "automatic": + if turn.candidate_index == 0 and reason in { + "rate_limit", + "billing", + }: + self._model_policy.cooldowns[turn.session.id] = ( + AutoFailoverCooldown( + model=next_candidate, + primary=turn.model_candidates[0], + expires_at=( + time.monotonic() + + RATE_LIMIT_COOLDOWN_SECONDS + ), + reason=reason, + ) + ) + else: + cooldown = self._model_policy.cooldowns.get(turn.session.id) + if cooldown and cooldown.expires_at > time.monotonic(): + cooldown.model = next_candidate + + self._model_policy.select_candidate(turn, next_index) + payload = { + "sessionID": turn.session.id, + "from": { + "providerID": previous.provider_id, + "modelID": previous.model_id, + }, + "to": { + "providerID": next_candidate.provider_id, + "modelID": next_candidate.model_id, + }, + "reason": reason, + "candidateIndex": next_index, + } + log.warn( + "session.model.fallback", + { + "from": payload["from"], + "to": payload["to"], + "reason": reason, + "candidateIndex": next_index, + }, + ) + await SessionEventSink.emit( + turn.callbacks, + "session.model.fallback", + payload, + ) @staticmethod def _canonical_tool_signature(tool_name: str, arguments: Dict[str, Any]) -> str: @@ -355,7 +551,7 @@ async def _run_session_start_hook(self, agent: Any) -> None: return self._session_start_fired = True try: - await self._ports.hooks.run_session_start({ + await HookPipeline.run_session_start({ "sessionID": self.session.id, "workspace": self.session.directory, "agent": agent.name, @@ -573,7 +769,7 @@ async def _list_callable_tool_infos_for_turn( execution_mode == SessionExecutionMode.PLAN and all(tool_info.name != "plan_exit" for tool_info in tool_infos) ): - plan_exit = self._ports.tools.get("plan_exit") + plan_exit = ToolRegistry.get("plan_exit") if plan_exit is not None and getattr(plan_exit.info, "enabled", True): tool_infos.append(plan_exit.info) metadata = dict(result.metadata) @@ -645,7 +841,7 @@ def _log_perf(self, event: str, started_at: float, **extra: Any) -> None: def _provider_capability_key(self) -> str: interleaved = None try: - active_model = self._ports.models.resolve_model( + active_model = Provider.resolve_model( self.provider_id, self.model_id, ) @@ -667,7 +863,7 @@ def _tool_schema_cache_key( text_tool_call_mode: bool, ) -> Tuple[Any, ...]: return ( - self._ports.tools.revision(), + ToolRegistry.revision(), getattr(agent, "name", ""), tuple(sorted(getattr(agent, "tools", None) or ())), tuple(tool_info.name for tool_info in selected_tool_infos), @@ -827,7 +1023,7 @@ def _model_supports_vision(self) -> bool: unknown configurations. """ try: - provider = self._ports.models.get_provider(self.provider_id) + provider = Provider.get(self.provider_id) if provider is not None: for model in getattr(provider, "_config_models", []) or []: if model.id == self.model_id: @@ -945,236 +1141,10 @@ def _append_file_content_block( placeholder = placeholder[:MAX_PLACEHOLDER_CHARS] + "…" text_fallbacks.append(placeholder) - @classmethod - async def loop(cls, session_id: str) -> Optional['MessageInfo']: - """ - Start or continue session processing loop. - - This is the main entry point for session execution, - matching Flocks' SessionPrompt.loop() behavior. - - Now delegates to SessionLoop for better separation of concerns. - - Args: - session_id: Session ID to process - - Returns: - Last assistant message with parts - """ - # Delegate to SessionLoop (new architecture) - from flocks.session.session_loop import SessionLoop - - result = await SessionLoop.run(session_id) - return result.last_message - - @classmethod - def cancel(cls, session_id: str) -> bool: - """ - Cancel a running session. - - Args: - session_id: Session ID to cancel - - Returns: - True if session was cancelled - """ - from flocks.session.core.status import SessionStatus - - runner = cls._active_sessions.get(session_id) - if runner: - runner.abort() - del cls._active_sessions[session_id] - log.info("runner.cancelled", {"session_id": session_id}) - - # Set status to idle (Flocks compatibility) - from flocks.session.core.status import SessionStatusIdle - SessionStatus.set(session_id, SessionStatusIdle()) - return True - - @classmethod - def cancel_children(cls, parent_session_id: str) -> int: - """Cancel all runners whose session.parent_id matches, recursively.""" - from flocks.session.core.status import SessionStatus, SessionStatusIdle - - cancelled = 0 - child_ids = [ - sid for sid, runner in list(cls._active_sessions.items()) - if getattr(runner.session, 'parent_id', None) == parent_session_id - ] - for sid in child_ids: - runner = cls._active_sessions.pop(sid, None) - if runner: - runner.abort() - SessionStatus.set(sid, SessionStatusIdle()) - cancelled += 1 - log.info("runner.child_cancelled", { - "session_id": sid, - "parent_session_id": parent_session_id, - }) - cancelled += cls.cancel_children(sid) - return cancelled - - @classmethod - async def command( - cls, - session_id: str, - command: str, - arguments: str = "", - message_id: Optional[str] = None, - agent: Optional[str] = None, - model: Optional[str] = None, - variant: Optional[str] = None, - ) -> Dict[str, Any]: - """ - Execute a slash command in a session. - - Args: - session_id: Session ID - command: Command name (e.g., "init", "help") - arguments: Command arguments - message_id: Optional message ID - agent: Optional agent name - model: Optional model string (provider/model) - variant: Optional model variant - - Returns: - Command execution result - """ - from flocks.command.command import Command - - # Get command definition - cmd = Command.get(command) - if not cmd: - raise ValueError(f"Command '{command}' not found") - - # Parse model if provided - provider_id, model_id = None, None - if model: - parts = model.split("/", 1) - if len(parts) == 2: - provider_id, model_id = parts - - # Execute command template - template = cmd.template - - # Replace placeholders - template = template.replace("$ARGUMENTS", arguments) - - # Create prompt request - parts = [{"type": "text", "text": template}] - - log.info("runner.command", { - "session_id": session_id, - "command": command, - "arguments": arguments[:50] if arguments else "", - }) - - return { - "command": command, - "arguments": arguments, - "template": template, - } - - @classmethod - async def shell( - cls, - session_id: str, - agent: str, - command: str, - model: Optional[Dict[str, str]] = None, - ) -> Dict[str, Any]: - """ - Execute a shell command in session context. - - Args: - session_id: Session ID - agent: Agent name - command: Shell command to execute - model: Optional model info - - Returns: - Shell execution result - """ - session = await Session.get_by_id(session_id) - if not session: - raise ValueError(f"Session {session_id} not found") - - cwd = session.directory or os.getcwd() - - user_msg = await Message.create( - session_id=session_id, - role=MessageRole.USER, - content="The following tool was executed by the user", - agent=agent, - ) - - assistant_msg = await Message.create( - session_id=session_id, - role=MessageRole.ASSISTANT, - content="", - agent=agent, - parent_id=user_msg.id, - ) - - start_time = asyncio.get_event_loop().time() - try: - proc = await asyncio.create_subprocess_shell( - command, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=cwd, - ) - stdout_bytes, stderr_bytes = await asyncio.wait_for( - proc.communicate(), timeout=300, - ) - output = (stdout_bytes or b"").decode("utf-8", errors="replace") + \ - (stderr_bytes or b"").decode("utf-8", errors="replace") - exit_code = proc.returncode or 0 - except asyncio.TimeoutError: - output = "Command timed out after 300 seconds" - exit_code = -1 - try: - proc.kill() - except Exception as _kill_err: - log.debug("runner.shell.kill_failed", {"error": str(_kill_err)}) - except Exception as e: - output = f"Error executing command: {str(e)}" - exit_code = -1 - - end_time = asyncio.get_event_loop().time() - - log.info("runner.shell", { - "session_id": session_id, - "command": command[:50], - "exit_code": exit_code, - "duration_ms": int((end_time - start_time) * 1000), - }) - - return { - "info": { - "id": assistant_msg.id, - "sessionID": session_id, - "role": "assistant", - "agent": agent, - }, - "parts": [{ - "id": Identifier.create("part"), - "messageID": assistant_msg.id, - "sessionID": session_id, - "type": "tool", - "tool": "bash", - "state": { - "status": "completed", - "input": {"command": command}, - "output": output, - }, - }], - } - def abort(self) -> None: """Signal abort to stop the loop.""" self._abort.set() - + @property def is_aborted(self) -> bool: """Check if abort was signaled (internal or external).""" @@ -1282,9 +1252,12 @@ def _deferred_failure_result( attempts: int, ) -> StepResult: state = LlmAttemptState( + request_sent=self._attempt_state.request_sent, received_chunk=self._attempt_state.received_chunk, observable_output_started=self._attempt_state.observable_output_started, tool_execution_started=self._attempt_state.tool_execution_started, + tool_execution_completed=self._attempt_state.tool_execution_completed, + externally_visible=self._attempt_state.externally_visible, ) return StepResult( action="stop", @@ -1299,7 +1272,7 @@ def _deferred_failure_result( attempts=attempts, ), ) - + async def _process_step( self, messages: List[MessageInfo], @@ -1317,27 +1290,13 @@ async def _process_step( self.session, worktree=Instance.get_worktree(), ) - # Check for CLI callbacks (if running in CLI mode) - # Only use CLI fallback if no callbacks were explicitly provided via constructor - has_explicit_callbacks = any([ - self.callbacks.on_text_delta, - self.callbacks.on_tool_start, - self.callbacks.on_tool_end, - self.callbacks.on_error, - self.callbacks.event_publish_callback, - ]) - if not has_explicit_callbacks: - try: - from flocks.cli.session_runner import _get_cli_callbacks - cli_callbacks = _get_cli_callbacks() - if cli_callbacks: - self.callbacks = cli_callbacks - except ImportError: - pass - # Resolve agent agent_name = last_user.agent or self.agent_name - agent = await Agent.get(agent_name) or await Agent.get("rex") + agent = getattr(self, "_step_agent", None) + if agent is None: + agent = await Agent.get(agent_name) or await Agent.get("rex") + if self._turn is not None: + self._step_agent = agent # Track session agent (Flocks compatibility) try: @@ -1345,13 +1304,13 @@ async def _process_step( set_session_agent(self.session.id, agent.name) except Exception as e: log.debug("runner.session_agent.error", {"error": str(e)}) - + # Check if we've reached max steps (matching Flocks logic) max_steps = agent.steps if hasattr(agent, 'steps') and agent.steps is not None else DEFAULT_MAX_TOOL_STEPS is_last_step = self._step >= max_steps - + # Get provider - provider = self._ports.models.get_provider(self.provider_id) + provider = Provider.get(self.provider_id) if not provider: error = f"Provider {self.provider_id} not found" if self._defer_step_errors: @@ -1385,13 +1344,13 @@ async def _process_step( # Apply config-based provider options (api_key/base_url) try: - await self._ports.models.apply_config(self.provider_id) + await Provider.apply_config(provider_id=self.provider_id) except Exception as e: log.debug("runner.provider.apply_config.error", { "provider": self.provider_id, "error": str(e), }) - + if not provider.is_configured(): error = f"Provider {self.provider_id} not configured" if self._defer_step_errors: @@ -1422,10 +1381,14 @@ async def _process_step( if self.callbacks.on_error: await self.callbacks.on_error(error_dict["data"]["displayMessage"]) return StepResult(action="stop", error=error_dict["data"]["displayMessage"]) - + # Build prompts and tools tools_started_at = time.perf_counter() - tools = await self._build_callable_tool_schema(agent, messages) + frozen_tool_request = getattr(self, "_frozen_tool_request", None) + if isinstance(frozen_tool_request, ModelRequest): + tools = frozen_tool_request.provider_tools() + else: + tools = await self._build_callable_tool_schema(agent, messages) self._log_perf("runner.process_step.tools_ready", tools_started_at, tool_count=len(tools)) prompt_tool_names = self._get_prompt_tool_names_from_schema(tools) @@ -1446,7 +1409,7 @@ async def device_asset_prompt_factory() -> Optional[str]: current_device_revision = None prompts_started_at = time.perf_counter() - system_prompts = await self._ports.prompts.build_system_prompts( + system_prompts = await SessionPrompt.build_system_prompts( session_id=self.session.id, session_directory=self.session.directory, agent_name=agent.name, @@ -1459,7 +1422,7 @@ async def device_asset_prompt_factory() -> Optional[str]: plan_file=self._turn_plan_file, ), prompt_tool_names=prompt_tool_names, - tool_revision=self._ports.tools.revision(), + tool_revision=ToolRegistry.revision(), memory_bootstrap_data=self._memory_bootstrap_data, static_cache=self._static_cache, sandbox_prompt_factory=sandbox_prompt_factory, @@ -1499,7 +1462,7 @@ async def device_asset_prompt_factory() -> Optional[str]: if has_tool_result and not has_text: from flocks.session.prompt_strings import PROMPT_TOOL_RESULTS_AVAILABLE system_prompts.append(PROMPT_TOOL_RESULTS_AVAILABLE) - + if has_tool_result and self._should_warn_about_tool_loop(last_user_id=last_user.id): state = self._get_tool_loop_guard_state(last_user_id=last_user.id) log.warn("runner.repeated_tool_calls_detected", { @@ -1539,7 +1502,7 @@ async def device_asset_prompt_factory() -> Optional[str]: "message_count": len(messages), }) raise - + # CRITICAL FIX: Ensure messages don't end with assistant role when tools are present # This prevents "assistant role in the final position when tools are used" API error # This commonly happens when: @@ -1561,7 +1524,7 @@ async def device_asset_prompt_factory() -> Optional[str]: "step": self._step, "session_id": self.session.id, }) - + # Add max steps warning if this is the last step (matching Flocks) if is_last_step: from flocks.session.prompt_strings import PROMPT_MAX_STEPS @@ -1569,17 +1532,26 @@ async def device_asset_prompt_factory() -> Optional[str]: role="assistant", content=PROMPT_MAX_STEPS, )) - + log.warn("runner.max_steps_reached", { "step": self._step, "max_steps": max_steps, "session_id": self.session.id, }) - + # Disable tools when max steps reached tools = [] - - # Create assistant message (will be reused across retries) + + request = self._build_model_request( + messages=chat_messages, + tools=tools, + agent=agent, + ) + if self._turn is not None and self._frozen_tool_request is None: + self._frozen_tool_request = request + self._active_model_request = request + + # Create the persisted assistant attempt after the request is frozen. assistant_msg = await Message.create( session_id=self.session.id, role=MessageRole.ASSISTANT, @@ -1589,7 +1561,8 @@ async def device_asset_prompt_factory() -> Optional[str]: provider_id=self.provider_id, parent_id=last_user.id, ) - + self._attempt_state.externally_visible = True + # Publish assistant message SSE event so frontends can show the message card if self.callbacks.event_publish_callback: import time as _time @@ -1607,7 +1580,7 @@ async def device_asset_prompt_factory() -> Optional[str]: "tokens": {"input": 0, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}}, } }) - + # Retry loop matching Flocks' SessionProcessor.process() # MAX_ERROR_RETRIES caps exception-based retries so a permanently-failing # model endpoint (e.g. repeated 500) cannot hold the session loop open @@ -1822,7 +1795,7 @@ async def device_asset_prompt_factory() -> Optional[str]: "reason": retry_message, "max_retries": retry_limit, }) - + # Set retry status SessionStatus.set( self.session.id, @@ -1832,10 +1805,10 @@ async def device_asset_prompt_factory() -> Optional[str]: next=next_retry_time, ) ) - + # Wait before retry await SessionRetry.sleep(delay_ms, self._abort) - + # Continue to next retry attempt continue else: @@ -1874,7 +1847,7 @@ async def device_asset_prompt_factory() -> Optional[str]: if self.callbacks.on_error: await self.callbacks.on_error(final_error_message) - + # Update assistant message with error (must be dict, not string) await Message.update( self.session.id, @@ -1890,9 +1863,9 @@ async def device_asset_prompt_factory() -> Optional[str]: error_dict=error_dict, text_part=text_part, ) - + return StepResult(action="stop", error=final_error_message) - + # Aborted return StepResult(action="stop", error="Aborted") @@ -2019,7 +1992,7 @@ async def _record_usage_if_available( "model_id": self.model_id, "error": str(exc), }) - + async def _build_device_asset_hint(self) -> Optional[str]: """Return concise device-aware tool guidance plus enabled device summary.""" try: @@ -2032,7 +2005,7 @@ async def _build_device_asset_hint(self) -> Optional[str]: vendor_by_storage_key: Dict[str, str] = {} try: - for tool_info in self._ports.tools.list_tools(): + for tool_info in ToolRegistry.list_tools(): if getattr(tool_info, "source", None) != "device": continue storage_key = str(getattr(tool_info, "provider", "") or "").strip() @@ -2254,7 +2227,7 @@ def _build_text_tool_call_catalog_prompt(self, tools: List[Dict[str, Any]]) -> O lines.append(f" - `{param_name}` ({param_type}, {required_suffix})") return "\n".join(lines) - + async def _build_callable_tool_schema( self, agent: AgentInfo, @@ -2316,19 +2289,19 @@ async def _build_callable_tool_schema( enabled=selection_metadata.get("enabledToolCount"), ) return tools - + def _agent_declares_tool(self, agent: AgentInfo, tool_name: str) -> bool: """Check if agent statically declares a tool.""" - tool = self._ports.tools.get(tool_name) + tool = ToolRegistry.get(tool_name) if tool is None: return False metadata = get_tool_catalog_metadata(tool_name, tool.info) return agent_declares_tool(agent, tool_name) or metadata.always_load - + def _exception_to_error_dict(self, exception: Exception) -> Dict[str, Any]: """ Convert exception to error dict for retry checking. - + Ported from original MessageV2.fromError() structure. """ error_dict = { @@ -2353,7 +2326,7 @@ def _exception_to_error_dict(self, exception: Exception) -> Dict[str, Any]: "transportExceptionType": transport_type, "transportExceptionModule": type(transport_exception).__module__, }) - + # Provider SDKs expose HTTP status through several shapes. Walk the # normal exception chain so lightweight wrapper errors do not hide it. status_code = None @@ -2392,11 +2365,11 @@ def _exception_to_error_dict(self, exception: Exception) -> Dict[str, Any]: if status_code is not None: error_dict["name"] = "APIError" error_dict["data"]["statusCode"] = status_code - + # Determine if retryable based on status code is_retryable = status_code in {429, 500, 502, 503, 504} error_dict["data"]["isRetryable"] = is_retryable - + # Extract response headers if available response = getattr(status_exception, "response", None) headers = getattr(response, "headers", None) @@ -2405,7 +2378,7 @@ def _exception_to_error_dict(self, exception: Exception) -> Dict[str, Any]: error_dict["data"]["responseHeaders"] = dict(headers) except (TypeError, ValueError): pass - + # Check for common retryable error patterns error_msg = str(exception).lower() if any(pattern in error_msg for pattern in [ @@ -2420,13 +2393,13 @@ def _exception_to_error_dict(self, exception: Exception) -> Dict[str, Any]: error_dict["name"] = "APIError" error_dict["data"]["isRetryable"] = True error_dict["data"]["displayMessage"] = CONNECTION_ERROR_DISPLAY_MESSAGE - + return error_dict - + def _get_context_window_tokens(self) -> int: """Resolve the context window size for the current model.""" try: - ctx, _, _ = self._ports.models.resolve_model_info( + ctx, _, _ = Provider.resolve_model_info( self.provider_id, self.model_id, ) @@ -2619,7 +2592,7 @@ async def _to_chat_messages( ) -> List[ChatMessage]: """ Convert messages to chat format with tool calls. - + Ported from original MessageV2.toModelMessage() logic: - Include text parts - Include tool calls and results @@ -2631,7 +2604,7 @@ async def _to_chat_messages( tool_result_refs: List[Dict[str, Any]] = [] turn_index = 0 queued_user_message_ids: set[str] = set(getattr(self, "_queued_user_message_ids", set()) or set()) - active_model = self._ports.models.resolve_model( + active_model = Provider.resolve_model( self.provider_id, self.model_id, ) @@ -2701,7 +2674,7 @@ async def _to_chat_messages( role="system", content=system_content, )) - + # Convert each message with parts for idx, msg in enumerate(messages): if idx < resume_message_index: @@ -2712,7 +2685,7 @@ async def _to_chat_messages( is_latest_user_turn = msg.id == last_user_msg_id # Get message parts parts = preloaded_parts[idx] - + if not parts: # Fallback: use text content only content = await Message.get_text_content(msg) @@ -2725,7 +2698,7 @@ async def _to_chat_messages( content=normalized_content, )) continue - + # Build message content from parts if msg.role == MessageRole.USER or (isinstance(msg.role, str) and msg.role == "user"): is_queued_user_turn = msg.id in queued_user_message_ids @@ -2801,7 +2774,7 @@ async def _to_chat_messages( role="user", content=user_text, )) - + elif msg.role == MessageRole.ASSISTANT or (isinstance(msg.role, str) and msg.role == "assistant"): # Skip messages with errors (matching Flocks logic) # Flocks: skip if error exists, UNLESS it's AbortedError with useful content @@ -2811,7 +2784,7 @@ async def _to_chat_messages( if isinstance(msg.error, dict): error_name = msg.error.get('name', '') is_aborted_error = error_name in ('MessageAbortedError', 'AbortedError') - + # If AbortedError, check if message has useful content if is_aborted_error: has_content = any( @@ -2825,7 +2798,7 @@ async def _to_chat_messages( else: # Non-AbortedError - skip continue - + assistant_content_parts = [] assistant_reasoning_parts = [] assistant_reasoning_content_parts = [] @@ -2836,11 +2809,11 @@ async def _to_chat_messages( structured_tool_calls: List[Dict[str, Any]] = [] # Corresponding tool-result messages (role="tool") pending_tool_results: List[ChatMessage] = [] - + for part in parts: if not hasattr(part, 'type'): continue - + # Text parts if part.type == "text" and hasattr(part, 'text'): if getattr(part, "ignored", False): @@ -2891,13 +2864,13 @@ async def _to_chat_messages( "type": "thinking", "thinking": part.text, }) - + # Tool parts - use structured OpenAI function-calling format elif part.type == "tool" and hasattr(part, 'state'): tool_name = getattr(part, 'tool', 'unknown') call_id = getattr(part, 'callID', None) or f"call_{id(part)}" tool_input = getattr(part.state, 'input', {}) - + if part.state.status == "completed": tool_output_str, was_dyn_truncated, persisted_placeholder = self._build_tool_output_text( part, @@ -2911,7 +2884,7 @@ async def _to_chat_messages( "context_window": ctx_window_tokens, "truncated_len": len(tool_output_str), }) - + # Build structured tool call for assistant message args_str = json.dumps(tool_input, ensure_ascii=False) if not isinstance(tool_input, str) else tool_input structured_tool_calls.append({ @@ -2939,7 +2912,7 @@ async def _to_chat_messages( "compacted": bool(persisted_placeholder), "dirty": False, }) - + log.debug("runner.to_chat_messages.tool_result_added", { "message_id": msg.id, "tool_name": tool_name, @@ -2947,7 +2920,7 @@ async def _to_chat_messages( "output_length": len(tool_output_str), "compacted": bool(persisted_placeholder), }) - + elif part.state.status == "error": tool_error = getattr(part.state, 'error', 'Unknown error') args_str = json.dumps(tool_input, ensure_ascii=False) if not isinstance(tool_input, str) else tool_input @@ -2965,7 +2938,7 @@ async def _to_chat_messages( tool_call_id=call_id, name=tool_name, )) - + elif part.state.status == "running": # Tool was interrupted (e.g., by user abort) before completing. # Include it in chat context so the LLM knows this tool call was @@ -2990,7 +2963,7 @@ async def _to_chat_messages( "tool_name": tool_name, "call_id": call_id, }) - + has_assistant_reasoning = bool( assistant_reasoning_parts or assistant_reasoning_content_parts @@ -3025,7 +2998,7 @@ async def _to_chat_messages( "parts_count": len(parts), "has_error": hasattr(msg, 'error') and bool(msg.error), }) - + budget_result = await self._apply_tool_result_budget(tool_result_refs, ctx_window_tokens) if budget_result.get("compacted"): log.info("runner.context_budget_enforced", { @@ -3054,9 +3027,133 @@ async def _to_chat_messages( source_message_count=len(messages), chat_message_count=len(chat_messages), ) - + return chat_messages - + + def _build_model_request( + self, + *, + messages: List[ChatMessage], + tools: List[Dict[str, Any]], + agent: AgentInfo, + ) -> ModelRequest[ChatMessage]: + """Freeze the exact provider-bound input for same-model retries.""" + from flocks.provider.options import build_provider_options + + provider_tools_enabled = not self._should_use_text_tool_call_mode() + return ModelRequest( + provider_id=self.provider_id, + model_id=self.model_id, + messages=tuple(messages), + tools=tuple(tools), + options=build_provider_options(self.provider_id, self.model_id), + metadata={ + "sessionID": self.session.id, + "workspace": self.session.directory, + "agent": agent.name, + "step": self._step, + "providerToolsEnabled": provider_tools_enabled, + }, + ) + + @staticmethod + def _serialize_model_message(message: ChatMessage) -> Dict[str, Any]: + payload = message.model_dump(exclude_none=True) + if not payload.get("custom_settings"): + payload.pop("custom_settings", None) + return payload + + async def _apply_before_model_hook( + self, + request: ModelRequest[ChatMessage], + hook_metadata: Dict[str, Any], + ) -> ModelRequest[ChatMessage]: + """Apply hook changes and freeze the request that will be sent.""" + request_payload = { + "providerID": request.provider_id, + "modelID": request.model_id, + "messageCount": len(request.messages), + "messages": [ + self._serialize_model_message(message) + for message in request.messages + ], + "toolCount": len(request.tools), + "tools": request.provider_tools(), + "providerOptions": request.provider_options(), + "providerToolsEnabled": bool( + request.metadata.get("providerToolsEnabled"), + ), + } + hook_input = {**hook_metadata, "request": request_payload} + started_at = time.perf_counter() + hook_context = await HookPipeline.run_llm_before(hook_input) + self._log_perf( + "runner.hook.llm_before.complete", + started_at, + message_count=len(request.messages), + tool_count=len(request.tools), + ) + + hook_output = getattr(hook_context, "output", {}) or {} + if hook_output.get("abort") or hook_output.get("blocked"): + reason = hook_output.get("reason") or "Model request blocked by hook" + raise RuntimeError(str(reason)) + effective_input = getattr(hook_context, "input", hook_input) + effective_payload = hook_output.get("request") + if not isinstance(effective_payload, Mapping): + effective_payload = effective_input.get("request", request_payload) + if not isinstance(effective_payload, Mapping): + raise TypeError("llm_before hook request must be a mapping") + + provider_id = str( + effective_payload.get("providerID", request.provider_id), + ) + model_id = str(effective_payload.get("modelID", request.model_id)) + if (provider_id, model_id) != (request.provider_id, request.model_id): + raise ValueError( + "llm_before hook cannot override ModelRoutingPolicy", + ) + + effective_messages: List[ChatMessage] = [] + for message in effective_payload.get("messages", request.messages): + if isinstance(message, ChatMessage): + effective_messages.append(message) + elif isinstance(message, Mapping): + effective_messages.append(ChatMessage(**dict(message))) + else: + raise TypeError( + "llm_before hook messages must be ChatMessage mappings", + ) + + effective_tools = effective_payload.get( + "tools", + request.provider_tools(), + ) + if not isinstance(effective_tools, (list, tuple)): + raise TypeError("llm_before hook tools must be a sequence") + effective_options = effective_payload.get( + "providerOptions", + request.provider_options(), + ) + if not isinstance(effective_options, Mapping): + raise TypeError("llm_before hook providerOptions must be a mapping") + + metadata = dict(request.metadata) + metadata["providerToolsEnabled"] = bool( + effective_payload.get( + "providerToolsEnabled", + metadata.get("providerToolsEnabled"), + ), + ) + return ModelRequest( + provider_id=request.provider_id, + model_id=request.model_id, + messages=tuple(effective_messages), + tools=tuple(dict(tool) for tool in effective_tools), + options=dict(effective_options), + metadata=metadata, + ) + async def _call_llm( self, provider: Any, @@ -3067,16 +3164,17 @@ async def _call_llm( ) -> StepResult: """ Call LLM and process response with event-driven streaming. - + Uses StreamProcessor to handle events and execute tools synchronously. Ported from Flocks' SessionProcessor.process() behavior. """ - def _serialize_message(message: ChatMessage) -> Dict[str, Any]: - payload = message.model_dump(exclude_none=True) - if not payload.get("custom_settings"): - payload.pop("custom_settings", None) - return payload - + request = getattr(self, "_active_model_request", None) + if not isinstance(request, ModelRequest): + request = self._build_model_request( + messages=messages, + tools=tools, + agent=agent, + ) def _build_llm_response_payload( *, content: str, @@ -3097,6 +3195,51 @@ def _build_llm_response_payload( ], } + llm_hook_metadata = { + "sessionID": self.session.id, + "messageID": assistant_msg.id, + "workspace": self.session.directory, + "agent": agent.name, + "step": self._step, + "model": { + "providerID": request.provider_id, + "modelID": request.model_id, + }, + } + cached_request = self._hooked_model_requests.get(assistant_msg.id) + if cached_request is None: + llm_before_enabled = False + llm_after_enabled = False + try: + llm_before_enabled = ( + await HookPipeline.has_stage_handlers( + HookStage.LLM_BEFORE, + llm_hook_metadata, + ) + ) + llm_after_enabled = ( + await HookPipeline.has_stage_handlers( + HookStage.LLM_AFTER, + llm_hook_metadata, + ) + ) + except Exception as exc: + log.debug("runner.hook.stage_probe.error", {"error": str(exc)}) + if llm_before_enabled: + request = await self._apply_before_model_hook( + request, + llm_hook_metadata, + ) + self._hooked_model_requests[assistant_msg.id] = ( + request, + llm_after_enabled, + ) + else: + request, llm_after_enabled = cached_request + self._active_model_request = request + messages = request.provider_messages() + tools = request.provider_tools() + # Create stream processor main_session_key = self.session.id try: @@ -3116,6 +3259,14 @@ async def _on_tool_execution_start( if self.callbacks.on_tool_start: await self.callbacks.on_tool_start(tool_name, tool_input) + async def _on_tool_execution_end( + tool_name: str, + result: ToolResult, + ) -> None: + self._attempt_state.tool_execution_completed = True + if self.callbacks.on_tool_end: + await self.callbacks.on_tool_end(tool_name, result) + turn_plan_file = getattr(self, "_turn_plan_file", None) if turn_plan_file is None: turn_plan_file = session_plan_file(self.session) @@ -3128,7 +3279,7 @@ async def _on_tool_execution_start( text_delta_callback=self.callbacks.on_text_delta, reasoning_delta_callback=self.callbacks.on_reasoning_delta, tool_start_callback=_on_tool_execution_start, - tool_end_callback=self.callbacks.on_tool_end, + tool_end_callback=_on_tool_execution_end, event_publish_callback=self.callbacks.event_publish_callback, session_key=self.session.id, main_session_key=main_session_key, @@ -3146,11 +3297,13 @@ async def _on_tool_execution_start( plan_relative_path=turn_plan_file.relative_path, plan_permission_path=turn_plan_file.permission_path, ) - - # Build provider options (thinking / reasoning / max_tokens) - from flocks.provider.options import build_provider_options - provider_options = build_provider_options(self.provider_id, self.model_id) - provider_tools = None if self._should_use_text_tool_call_mode() else (tools if tools else None) + + provider_options = request.provider_options() + provider_tools = ( + tools + if request.metadata.get("providerToolsEnabled") and tools + else None + ) # Clean up any leftover reasoning state from a previous (failed) call if hasattr(self, '_current_reasoning_id'): @@ -3187,7 +3340,7 @@ async def _on_tool_execution_start( provider_options=provider_options, ) trace_ctx = trace_scope( - name="SessionRunner.step", + name="StepEngine.step", session_id=self.session.id, tags=trace_tags, input=request_payload, @@ -3227,7 +3380,7 @@ async def _on_tool_execution_start( log.debug("runner.observability.init_failed", {"error": str(exc)}) trace_ctx = None generation_ctx = None - + # Validate messages - ensure we have at least one non-system message non_system_messages = [m for m in messages if m.role != "system"] if not non_system_messages: @@ -3237,20 +3390,20 @@ async def _on_tool_execution_start( }) self._end_observability(generation_ctx, trace_ctx, output="No valid messages", level="ERROR") return StepResult(action="stop", content="", error="No valid messages to send to LLM") - + log.debug("runner.call_llm.messages", { "total": len(messages), "non_system": len(non_system_messages), "roles": [m.role for m in messages], }) - + # Emit start event await processor.process_event(StartEvent()) - + # Lightweight counters instead of storing all chunks in memory chunk_counts = {"total": 0, "reasoning": 0, "text": 0, "tool": 0} stream_usage: Optional[Dict[str, int]] = None - + # Stream response and convert chunks to events if provider_tools is None and tools: log.info("runner.text_tool_call_mode.enabled", { @@ -3260,61 +3413,7 @@ async def _on_tool_execution_start( "tool_count": len(tools), }) - llm_hook_metadata = { - "sessionID": self.session.id, - "messageID": assistant_msg.id, - "workspace": self.session.directory, - "agent": agent.name, - "step": self._step, - "model": { - "providerID": self.provider_id, - "modelID": self.model_id, - }, - } - llm_before_enabled = False - llm_after_enabled = False self._llm_call_aborted = False - try: - llm_before_enabled = ( - await self._ports.hooks.has_stage_handlers( - HookStage.LLM_BEFORE, - llm_hook_metadata, - ) - ) - llm_after_enabled = ( - await self._ports.hooks.has_stage_handlers( - HookStage.LLM_AFTER, - llm_hook_metadata, - ) - ) - except Exception as exc: - log.debug("runner.hook.stage_probe.error", {"error": str(exc)}) - - if llm_before_enabled: - llm_before_hook_input = { - **llm_hook_metadata, - "request": { - "messageCount": len(messages), - "messages": [_serialize_message(message) for message in messages], - "toolCount": len(tools), - "tools": copy.deepcopy(tools), - "providerOptions": dict(provider_options), - "providerToolsEnabled": provider_tools is not None, - }, - } - try: - hook_started_at = time.perf_counter() - await self._ports.hooks.run_llm_before( - llm_before_hook_input, - ) - self._log_perf( - "runner.hook.llm_before.complete", - hook_started_at, - message_count=len(messages), - tool_count=len(tools), - ) - except Exception as exc: - log.debug("runner.hook.llm_before.error", {"error": str(exc)}) llm_call_started_at = time.perf_counter() first_chunk_logged = False @@ -3328,6 +3427,7 @@ async def _on_tool_execution_start( "local_endpoint": stream_timeouts.is_local, }) try: + self._attempt_state.request_sent = True async for chunk in _iter_with_chunk_timeout( provider.chat_stream( model_id=self.model_id, @@ -3496,7 +3596,7 @@ async def _on_tool_execution_start( ) if llm_after_enabled: try: - await self._ports.hooks.run_llm_after( + await HookPipeline.run_llm_after( llm_hook_metadata, { "durationMs": int((time.perf_counter() - llm_call_started_at) * 1000), @@ -3512,7 +3612,7 @@ async def _on_tool_execution_start( except Exception as hook_exc: log.debug("runner.hook.llm_after.error", {"error": str(hook_exc)}) raise - + log.debug("runner.stream.summary", { "total_chunks": chunk_counts["total"], "reasoning_chunks": chunk_counts["reasoning"], @@ -3524,11 +3624,11 @@ async def _on_tool_execution_start( }) await tool_accumulator.flush_remaining(stream_finish_reason) - + # End text block if started if text_started: await processor.process_event(TextEndEvent()) - + # End any remaining reasoning block if hasattr(self, '_current_reasoning_id'): reasoning_end_metadata = getattr(self, '_current_reasoning_metadata', {}) or {} @@ -3539,7 +3639,7 @@ async def _on_tool_execution_start( delattr(self, '_current_reasoning_id') if hasattr(self, '_current_reasoning_metadata'): delattr(self, '_current_reasoning_metadata') - + # Emit finish event await processor.process_event(FinishEvent( finish_reason=processor.get_finish_reason() @@ -3549,11 +3649,11 @@ async def _on_tool_execution_start( # streaming so sibling subagents can start in the same assistant turn. # Drain them here before exposing tool results to the next loop step. await processor.drain_parallel_tool_calls() - + # Get processed content content = processor.get_text_content() reasoning = processor.get_reasoning_content() - + # Update message tokens if provider reported usage tokens_update = self._build_tokens_update(stream_usage) if tokens_update: @@ -3570,7 +3670,7 @@ async def _on_tool_execution_start( }) except Exception as e: log.warn("runner.stream.usage_update_failed", {"error": str(e)}) - + # Log summary log.debug("runner.stream.complete", { "text_length": len(content), @@ -3578,7 +3678,7 @@ async def _on_tool_execution_start( "tool_calls": len(processor.tool_calls), "usage": stream_usage, }) - + # Update assistant message with content if content: await Message.update( @@ -3587,7 +3687,7 @@ async def _on_tool_execution_start( content=content, ) self._llm_call_aborted = aborted_during_stream - + # Note: Tools were already executed synchronously during streaming # Build tool call list for result tool_calls_for_result = [ @@ -3608,7 +3708,7 @@ async def _on_tool_execution_start( if llm_after_enabled: try: hook_started_at = time.perf_counter() - await self._ports.hooks.run_llm_after( + await HookPipeline.run_llm_after( llm_hook_metadata, { "durationMs": int((time.perf_counter() - llm_call_started_at) * 1000), @@ -3634,7 +3734,7 @@ async def _on_tool_execution_start( ) except Exception as exc: log.debug("runner.hook.llm_after.error", {"error": str(exc)}) - + if tool_calls_for_result: response_payload = self._build_langfuse_response_payload( action="continue", @@ -3660,7 +3760,7 @@ async def _on_tool_execution_start( tool_calls=tool_calls_for_result, usage=stream_usage, ) - + response_payload = self._build_langfuse_response_payload( action="stop", content=content, @@ -3680,7 +3780,7 @@ async def _on_tool_execution_start( trace_output=response_payload, ) return StepResult(action=result_action, content=content, usage=stream_usage) - + @staticmethod def _end_observability( generation_ctx: Any, @@ -3767,41 +3867,3 @@ async def _handle_permission(self, request) -> None: always=list(getattr(request, "always", None) or []), tool={"name": request.permission}, ) - - -async def run_session( - session: SessionInfo, - provider_id: Optional[str] = None, - model_id: Optional[str] = None, - agent_name: Optional[str] = None, - callbacks: Optional[RunnerCallbacks] = None, -) -> Optional[MessageInfo]: - """ - Run a session to completion. - - Delegates to SessionLoop which is the single authoritative execution path. - - Args: - session: Session to run - provider_id: Provider ID - model_id: Model ID - agent_name: Agent name - callbacks: RunnerCallbacks (wrapped into LoopCallbacks) - - Returns: - Last assistant message - """ - from flocks.session.session_loop import SessionLoop, LoopCallbacks - - loop_callbacks = LoopCallbacks( - runner_callbacks=callbacks, - event_publish_callback=callbacks.event_publish_callback if callbacks else None, - ) - result = await SessionLoop.run( - session_id=session.id, - provider_id=provider_id, - model_id=model_id, - agent_name=agent_name, - callbacks=loop_callbacks, - ) - return result.last_message diff --git a/flocks/session/runtime_adapters.py b/flocks/session/runtime_adapters.py deleted file mode 100644 index def73f248..000000000 --- a/flocks/session/runtime_adapters.py +++ /dev/null @@ -1,85 +0,0 @@ -"""Default adapters from runtime ports to existing Flocks services.""" - -from __future__ import annotations - -from typing import Any, Optional - -from flocks.agent.runtime.ports import ExternalRuntimePorts -from flocks.hooks.pipeline import HookPipeline -from flocks.provider.provider import Provider -from flocks.session.prompt import SessionPrompt -from flocks.tool.registry import ToolRegistry - - -class FlocksPromptPort: - """Adapt the existing SessionPrompt builder.""" - - async def build_system_prompts(self, **kwargs: Any) -> list[str]: - return await SessionPrompt.build_system_prompts(**kwargs) - - -class FlocksToolPort: - """Adapt the process-wide ToolRegistry.""" - - def revision(self) -> int: - return ToolRegistry.revision() - - def list_tools(self) -> list[Any]: - return ToolRegistry.list_tools() - - def get(self, name: str) -> Optional[Any]: - return ToolRegistry.get(name) - - -class FlocksModelPort: - """Adapt provider configuration and model metadata lookup.""" - - def get_provider(self, provider_id: str) -> Optional[Any]: - return Provider.get(provider_id) - - async def apply_config(self, provider_id: str) -> None: - await Provider.apply_config(provider_id=provider_id) - - def resolve_model(self, provider_id: str, model_id: str) -> Optional[Any]: - return Provider.resolve_model(provider_id, model_id) - - def resolve_model_info( - self, - provider_id: str, - model_id: str, - ) -> tuple[int, int, Optional[int]]: - return Provider.resolve_model_info(provider_id, model_id) - - -class FlocksHookPort: - """Adapt HookPipeline while preserving hook names and payloads.""" - - async def run_session_start(self, data: dict[str, Any]) -> Any: - return await HookPipeline.run_session_start(data) - - async def has_stage_handlers( - self, - stage: Any, - metadata: dict[str, Any], - ) -> bool: - return await HookPipeline.has_stage_handlers(stage, metadata) - - async def run_llm_before(self, data: dict[str, Any]) -> Any: - return await HookPipeline.run_llm_before(data) - - async def run_llm_after( - self, - metadata: dict[str, Any], - result: dict[str, Any], - ) -> Any: - return await HookPipeline.run_llm_after(metadata, result) - - -def create_default_runtime_ports() -> ExternalRuntimePorts: - """Create adapters for one runner without changing public APIs.""" - return ExternalRuntimePorts( - prompts=FlocksPromptPort(), - tools=FlocksToolPort(), - models=FlocksModelPort(), - hooks=FlocksHookPort(), - ) diff --git a/flocks/session/runtime_services.py b/flocks/session/runtime_services.py deleted file mode 100644 index f30e1748e..000000000 --- a/flocks/session/runtime_services.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Session adapters consumed by the host-neutral agent runtime.""" - -from __future__ import annotations - -import asyncio -from typing import Any - -from flocks.agent.runtime.contracts import ( - AgentRunState, - ContinuationDecision, - ModelTurnBoundary, - ModelTurnPreparation, - ModelTurnSnapshot, - StepResult, -) -from flocks.agent.runtime.events import RuntimeEvent - - -class SessionStepCancelled(Exception): - """Internal signal raised when a hosted model turn is cancelled.""" - - -class SessionLoopStepEngine: - """Bind the session runner and cancellation state to one AgentLoop.""" - - def __init__(self, context: Any, callbacks: Any, policy: Any): - self._context = context - self._callbacks = callbacks - self._policy = policy - - async def run(self, snapshot: ModelTurnSnapshot[Any]) -> StepResult: - """Execute a model turn while exposing its task to session abort.""" - task = asyncio.create_task( - self._policy._process_model_step( - self._context, - self._callbacks, - snapshot, - ) - ) - self._context._current_step_task = task - started_at = asyncio.get_running_loop().time() - try: - return await task - except asyncio.CancelledError as exc: - raise SessionStepCancelled from exc - finally: - self._context._current_step_task = None - duration_ms = int((asyncio.get_running_loop().time() - started_at) * 1000) - self._policy._log_step_complete(self._context, duration_ms) - - -class SessionRuntimeServices: - """Adapt session persistence and policy to narrow runtime ports.""" - - def __init__(self, context: Any, callbacks: Any, policy: Any): - self._context = context - self._callbacks = callbacks - self._policy = policy - - async def prepare_model_turn( - self, - state: AgentRunState[Any], - ) -> ModelTurnPreparation[Any]: - """Prepare persisted session state for a stable model turn.""" - return await self._policy._prepare_model_turn( - self._context, - self._callbacks, - state, - ) - - async def complete_model_turn( - self, - state: AgentRunState[Any], - step_result: StepResult, - ) -> ModelTurnBoundary[Any]: - """Commit and expose the state written by a completed model turn.""" - return await self._policy._complete_model_turn( - self._context, - self._callbacks, - state, - step_result, - ) - - async def resolve_continuation( - self, - state: AgentRunState[Any], - step_result: StepResult, - ) -> ContinuationDecision[Any]: - """Resolve queued goal and TurnFinish continuation policy.""" - return await self._policy._resolve_continuation( - self._context, - self._callbacks, - state, - step_result, - ) - - async def emit_event(self, event: RuntimeEvent) -> None: - """Forward a host-neutral runtime event to session subscribers.""" - await self._policy._publish_runtime_event( - self._callbacks, - event.type, - dict(event.payload), - ) diff --git a/flocks/session/session.py b/flocks/session/session.py index 67bfc0ddc..76f4fe18c 100644 --- a/flocks/session/session.py +++ b/flocks/session/session.py @@ -1056,13 +1056,11 @@ async def _stop_session_tree_for_archive( ) -> bool: """Stop persisted and in-memory work before committing archive state.""" from flocks.session.interaction_queue import InteractionQueue - from flocks.session.runner import SessionRunner from flocks.session.session_loop import SessionLoop session_ids = [session.id for session in sessions] for session_id in session_ids: SessionLoop.abort(session_id) - SessionRunner.cancel(session_id) if clear_prompt_queue: await InteractionQueue.clear(session_id) try: diff --git a/flocks/session/session_host.py b/flocks/session/session_host.py deleted file mode 100644 index 52792b112..000000000 --- a/flocks/session/session_host.py +++ /dev/null @@ -1,487 +0,0 @@ -"""Session lifecycle host for one resumable agent execution.""" - -from __future__ import annotations - -import time -from collections.abc import Awaitable, Callable, MutableMapping -from dataclasses import dataclass, replace -from typing import Any, Optional - -from flocks.agent.runtime.contracts import ( - ModelTurnSnapshot, - RuntimeModel, - StepResult, -) -from flocks.session.core.context import DefaultSessionContext -from flocks.session.core.status import ( - SessionStatus, - SessionStatusBusy, - SessionStatusIdle, -) -from flocks.session.core.turn_state import clear_turn_state -from flocks.session.message import Message -from flocks.session.runtime_adapters import create_default_runtime_ports -from flocks.session.session import Session, is_model_auto_session_category -from flocks.utils.log import Log - - -log = Log.create(service="session.host") - - -@dataclass(frozen=True) -class SessionHostDependencies: - """Compatibility adapters supplied by the public SessionLoop facade.""" - - create_context: Callable[..., Any] - create_callbacks: Callable[[], Any] - create_result: Callable[..., Any] - resolve_model: Callable[..., Awaitable[tuple[str, str]]] - run_logical_turn: Callable[[Any, Any], Awaitable[Any]] - publish_session_status: Callable[[Any, str, str], Awaitable[None]] - - -@dataclass(frozen=True) -class SessionLease: - """Process-local ownership record protected by Session.lifecycle_lock.""" - - session_id: str - context: Any - - -class SessionLeaseRegistry: - """Manage active session ownership without exposing lifecycle policy.""" - - def __init__(self, active_contexts: MutableMapping[str, Any]): - self._active_contexts = active_contexts - - def get(self, session_id: str) -> Optional[Any]: - """Return the active context, if this process owns the session.""" - return self._active_contexts.get(session_id) - - def acquire(self, session_id: str, context: Any) -> Optional[SessionLease]: - """Acquire process-local ownership; caller holds the lifecycle lock.""" - if session_id in self._active_contexts: - return None - self._active_contexts[session_id] = context - return SessionLease(session_id=session_id, context=context) - - def release(self, lease: SessionLease) -> None: - """Release ownership only when the stored context is still ours.""" - if self._active_contexts.get(lease.session_id) is lease.context: - self._active_contexts.pop(lease.session_id, None) - - -class SessionHostStepEngine: - """Apply host-owned cross-model recovery around one-candidate attempts.""" - - def __init__( - self, - *, - context: Any, - callbacks: Any, - attempt_engine: Any, - cooldowns: MutableMapping[str, Any], - cooldown_factory: Callable[..., Any], - select_candidate: Callable[[Any, int], None], - finalize_failure: Callable[[Any, Any, Any], Awaitable[None]], - publish_event: Callable[[Any, str, dict[str, Any]], Awaitable[None]], - rate_limit_cooldown_seconds: float, - chain_exhaustion_cooldown_seconds: float, - ): - self._context = context - self._callbacks = callbacks - self._attempt_engine = attempt_engine - self._cooldowns = cooldowns - self._cooldown_factory = cooldown_factory - self._select_candidate = select_candidate - self._finalize_failure = finalize_failure - self._publish_event = publish_event - self._rate_limit_cooldown_seconds = rate_limit_cooldown_seconds - self._chain_exhaustion_cooldown_seconds = chain_exhaustion_cooldown_seconds - - async def run(self, snapshot: ModelTurnSnapshot[Any]) -> StepResult: - """Retry a replay-safe snapshot across the configured model chain.""" - while True: - active_model = RuntimeModel( - self._context.provider_id, - self._context.model_id, - ) - result = await self._attempt_engine.run( - replace(snapshot, active_model=active_model), - ) - failure = result.failure - if not self._context.auto_failover or failure is None: - return result - - next_index = self._context.candidate_index + 1 - has_next = next_index < len(self._context.model_candidates) - if not failure.allow_fallback or not failure.attempt_state.replay_safe or not has_next: - self._record_chain_exhaustion(failure, has_next) - await self._finalize_failure( - self._context, - failure, - snapshot.last_user, - ) - return result - - if not await self._remove_failed_attempt(failure): - await self._finalize_failure( - self._context, - failure, - snapshot.last_user, - ) - return result - - await self._switch_candidate(next_index, failure.reason) - - def _record_chain_exhaustion(self, failure: Any, has_next: bool) -> None: - context = self._context - if not ( - context.model_candidate_policy == "automatic" - and failure.allow_fallback - and failure.attempt_state.replay_safe - and not has_next - and context.candidate_index > 0 - and failure.reason not in {"rate_limit", "billing"} - ): - return - - expires_at = time.monotonic() + self._chain_exhaustion_cooldown_seconds - existing = self._cooldowns.get(context.session.id) - if existing and existing.expires_at > expires_at: - return - self._cooldowns[context.session.id] = self._cooldown_factory( - model=context.model_candidates[context.candidate_index], - primary=context.model_candidates[0], - expires_at=expires_at, - reason="chain_exhausted", - ) - - async def _remove_failed_attempt(self, failure: Any) -> bool: - message_id = failure.assistant_message_id - if not message_id: - return True - try: - deleted = await Message.delete(self._context.session.id, message_id) - except Exception as exc: - deleted = False - log.error( - "session.model.fallback_cleanup_failed", - { - "session_id": self._context.session.id, - "message_id": message_id, - "error": str(exc), - }, - ) - if not deleted: - return False - await self._publish_event( - self._callbacks, - "message.removed", - { - "sessionID": self._context.session.id, - "messageID": message_id, - }, - ) - return True - - async def _switch_candidate(self, next_index: int, reason: str) -> None: - context = self._context - previous = context.model_candidates[context.candidate_index] - next_candidate = context.model_candidates[next_index] - if context.model_candidate_policy == "automatic": - if context.candidate_index == 0 and reason in { - "rate_limit", - "billing", - }: - self._cooldowns[context.session.id] = self._cooldown_factory( - model=next_candidate, - primary=context.model_candidates[0], - expires_at=(time.monotonic() + self._rate_limit_cooldown_seconds), - reason=reason, - ) - else: - cooldown = self._cooldowns.get(context.session.id) - if cooldown and cooldown.expires_at > time.monotonic(): - cooldown.model = next_candidate - - self._select_candidate(context, next_index) - payload = { - "sessionID": context.session.id, - "from": { - "providerID": previous.provider_id, - "modelID": previous.model_id, - }, - "to": { - "providerID": next_candidate.provider_id, - "modelID": next_candidate.model_id, - }, - "reason": reason, - "candidateIndex": next_index, - } - log.warn( - "session.model.fallback", - { - "from": payload["from"], - "to": payload["to"], - "reason": reason, - "candidateIndex": next_index, - }, - ) - await self._publish_event( - self._callbacks, - "session.model.fallback", - payload, - ) - - -class SessionHost: - """Own session acquisition, recovery, execution, and final cleanup.""" - - def __init__( - self, - dependencies: SessionHostDependencies, - active_contexts: MutableMapping[str, Any], - ): - self._dependencies = dependencies - self._leases = SessionLeaseRegistry(active_contexts) - - async def run( - self, - session_id: str, - provider_id: Optional[str] = None, - model_id: Optional[str] = None, - agent_name: Optional[str] = None, - callbacks: Optional[Any] = None, - working_directory: Optional[str] = None, - auto_failover: bool = False, - ) -> Any: - """Acquire and host one session until its logical turn settles.""" - active_context = self._leases.get(session_id) - if active_context is not None: - log.info("session.already_running", {"session_id": session_id}) - self._authorize_auto_failover(active_context, auto_failover) - return self._dependencies.create_result( - action="queued", - error="Loop already running", - ) - - session = await Session.get_by_id(session_id) - if session is None: - log.warning("session.not_found", {"session_id": session_id}) - return self._dependencies.create_result( - action="error", - error=f"Session {session_id} not found", - ) - if session.status != "active": - log.warning( - "session.not_active", - {"session_id": session_id, "status": session.status}, - ) - return self._dependencies.create_result( - action="error", - error=f"Session {session_id} is {session.status}", - ) - if working_directory: - session = session.model_copy(update={"directory": working_directory}) - - if not provider_id or not model_id: - resolved_provider, resolved_model = await self._dependencies.resolve_model( - session, - provider_id, - model_id, - ) - provider_id = provider_id or resolved_provider - model_id = model_id or resolved_model - - primary_model = RuntimeModel(provider_id=provider_id, model_id=model_id) - auto_failover = bool( - auto_failover - and is_model_auto_session_category( - getattr(session, "category", "user"), - ) - ) - session.provider = provider_id - session.model = model_id - - trace_offset = await self._load_trace_offset(session_id) - context = self._dependencies.create_context( - session=session, - provider_id=provider_id, - model_id=model_id, - agent_name=agent_name or session.agent or "rex", - session_ctx=DefaultSessionContext(session), - trace_step_offset=trace_offset, - auto_failover=auto_failover, - auto_failover_allowed=auto_failover, - model_candidates=[primary_model], - candidate_index=0, - session_start_pending=trace_offset == 0, - runtime_ports=create_default_runtime_ports(), - ) - lease_or_result = await self._acquire_lease(session_id, context) - if not isinstance(lease_or_result, SessionLease): - return lease_or_result - lease = lease_or_result - - runtime_callbacks = callbacks or self._dependencies.create_callbacks() - await self._mark_busy(session_id, runtime_callbacks) - await self._recover_orphan_tools(session_id) - - try: - return await self._dependencies.run_logical_turn( - context, - runtime_callbacks, - ) - except Exception as exc: - return await self._handle_execution_error( - context, - callbacks, - exc, - ) - finally: - await self._release_session( - lease, - session, - runtime_callbacks, - ) - - @staticmethod - def _authorize_auto_failover(context: Any, requested: bool) -> None: - if requested and is_model_auto_session_category( - getattr(context.session, "category", "user"), - ): - context.auto_failover_allowed = True - - @staticmethod - async def _load_trace_offset(session_id: str) -> int: - try: - messages = await Message.list(session_id) - return sum(1 for message in messages if message.role == "assistant") - except Exception as exc: - log.debug("session.trace_offset.error", {"error": str(exc)}) - return 0 - - async def _acquire_lease( - self, - session_id: str, - context: Any, - ) -> SessionLease | Any: - async with Session.lifecycle_lock(session_id): - latest_session = await Session.get_by_id(session_id) - if latest_session is None: - log.warning( - "session.not_found_before_lease", - {"session_id": session_id}, - ) - return self._dependencies.create_result( - action="error", - error=f"Session {session_id} not found", - ) - if latest_session.status != "active": - log.warning( - "session.not_active_before_lease", - { - "session_id": session_id, - "status": latest_session.status, - }, - ) - return self._dependencies.create_result( - action="error", - error=f"Session {session_id} is {latest_session.status}", - ) - if Session.is_lifecycle_transitioning(session_id): - return self._dependencies.create_result( - action="error", - error=f"Session {session_id} is changing lifecycle state", - ) - lease = self._leases.acquire(session_id, context) - if lease is None: - return self._dependencies.create_result( - action="queued", - error="Loop already running", - ) - return lease - - async def _mark_busy(self, session_id: str, callbacks: Any) -> None: - SessionStatus.set(session_id, SessionStatusBusy()) - await self._dependencies.publish_session_status( - callbacks, - session_id, - "busy", - ) - - @staticmethod - async def _recover_orphan_tools(session_id: str) -> None: - try: - from flocks.session.orphan_tools import abort_orphan_running_parts - - await abort_orphan_running_parts(session_id) - except Exception as exc: - log.warn( - "session.orphan_cleanup_failed", - {"session_id": session_id, "error": str(exc)}, - ) - - async def _handle_execution_error( - self, - context: Any, - callbacks: Optional[Any], - error: Exception, - ) -> Any: - session_id = context.session.id - log.error( - "session.execution_error", - {"session_id": session_id, "error": str(error)}, - ) - if callbacks and callbacks.on_error: - try: - await callbacks.on_error(str(error)) - except Exception as callback_error: - log.debug( - "session.error_callback_failed", - {"error": str(callback_error)}, - ) - try: - from flocks.bus.bus import Bus - from flocks.bus.events import SessionError - - await Bus.publish( - SessionError, - {"sessionID": session_id, "error": str(error)}, - ) - except Exception as publish_error: - log.warn( - "session.error_event_failed", - {"error": str(publish_error)}, - ) - return self._dependencies.create_result( - action="error", - error=str(error), - provider_id=context.provider_id, - model_id=context.model_id, - ) - - async def _release_session( - self, - lease: SessionLease, - session: Any, - callbacks: Any, - ) -> None: - self._leases.release(lease) - clear_turn_state(lease.session_id) - SessionStatus.set(lease.session_id, SessionStatusIdle()) - await self._dependencies.publish_session_status( - callbacks, - lease.session_id, - "idle", - ) - await Session.touch(session.project_id, lease.session_id) - - try: - from flocks.bus.bus import Bus - from flocks.bus.events import SessionIdle - - await Bus.publish(SessionIdle, {"sessionID": lease.session_id}) - except Exception as exc: - log.warn("session.idle_event_failed", {"error": str(exc)}) diff --git a/flocks/session/session_loop.py b/flocks/session/session_loop.py index 354ffa4a8..ae96f7f7a 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -1,621 +1,334 @@ -""" -Session Loop Module - -Core session execution loop logic extracted from runner.py. -Implements the main session processing loop with support for: -- Message processing -- Tool execution -- Compaction -- Subtask handling -- Reminders - -Ported from original SessionPrompt.loop() pattern. -""" - -import asyncio -import hashlib -import inspect -import time -from typing import Optional, List, Dict, Any, Callable, Awaitable, Literal -from dataclasses import dataclass, field -from datetime import datetime - -from flocks.agent.runtime.agent_loop import AgentLoop -from flocks.agent.runtime.contracts import ( +"""Public entry point and lifecycle owner for session execution.""" + +from __future__ import annotations + +from collections.abc import MutableMapping +from dataclasses import dataclass +from typing import Any, ClassVar, Optional + +from flocks.session.core.context import DefaultSessionContext +from flocks.session.core.status import ( + SessionStatus, + SessionStatusBusy, + SessionStatusIdle, +) +from flocks.session.core.turn_state import clear_turn_state +from flocks.session.message import Message +from flocks.session.runtime.agent_loop import AgentLoop +from flocks.session.runtime.continuation_policy import ( + DEFAULT_CONTINUATION_POLICY, + ContinuationPolicy, +) +from flocks.session.runtime.contracts import ( + AgentRunOutcome, AgentRunState, AgentRunStatus, - ContinuationDecision, - ModelTurnBoundary, - ModelTurnPreparation, - ModelTurnSnapshot, - QueuedInputBatch, RuntimeModel, - StepResult, - TurnPreparationStatus, ) -from flocks.agent.runtime.ports import ExternalRuntimePorts -from flocks.utils.log import Log -from flocks.utils.id import Identifier +from flocks.session.runtime.event_sink import SessionEventSink +from flocks.session.runtime.model_policy import ( + DEFAULT_MODEL_ROUTING_POLICY, + ModelRoutingPolicy, +) +from flocks.session.runtime.session_turn import ( + LoopCallbacks, + LoopResult, + SessionTurn, +) +from flocks.session.runtime.step_engine import StepEngine from flocks.session.session import ( Session, - SessionInfo, is_model_auto_session_category, ) -from flocks.session.message import Message, MessageInfo, MessageRole -from flocks.session.core.status import SessionStatus, SessionStatusBusy -from flocks.session.core.task_utils import fire_and_forget -from flocks.session.core.turn_state import ( - set_turn_state, - set_context_state, -) -from flocks.session.lifecycle.compaction import ( - SessionCompaction, - CompactionPolicy, - build_compaction_policy, - run_compaction, -) -from flocks.session.lifecycle.compaction.compaction import _get_compaction_history -from flocks.session.prompt import SessionPrompt -from flocks.provider.provider import Provider -from flocks.session.goal import GoalManager +from flocks.utils.log import Log log = Log.create(service="session.loop") +# Public compatibility name. There is only one state implementation. +LoopContext = SessionTurn + + +@dataclass(frozen=True) +class _SessionLease: + """One process-local ownership record.""" + + session_id: str + turn: SessionTurn + + +class _SessionLeaseRegistry: + """Keep lease bookkeeping out of the SessionLoop control flow.""" -MAX_OVERFLOW_COMPACTION_ATTEMPTS = 3 -POST_COMPACTION_COOLDOWN_STEPS = 2 -RATE_LIMIT_COOLDOWN_SECONDS = 60.0 -CHAIN_EXHAUSTION_COOLDOWN_SECONDS = 5.0 - - -@dataclass -class AutoFailoverCooldown: - """Process-local Hermes-style starting candidate cooldown.""" - - model: RuntimeModel - primary: RuntimeModel - expires_at: float - reason: str - - -@dataclass -class LoopContext: - """Context for session loop execution""" - session: SessionInfo - provider_id: str - model_id: str - agent_name: str - step: int = 0 - abort_event: asyncio.Event = field(default_factory=asyncio.Event) - # SessionContext interface for decoupled session access - session_ctx: Optional[Any] = None # Type: Optional[SessionContext] - # Offset so observability step numbers are cumulative across the session - trace_step_offset: int = 0 - # Track current step asyncio.Task so abort() can cancel it immediately - _current_step_task: Optional[asyncio.Task] = field(default=None, repr=False) - # Memory bootstrap data loaded once on step 1; passed to each SessionRunner - memory_bootstrap_data: Optional[Dict[str, Any]] = field(default=None, repr=False) - # Reusable runner artifacts that stay stable across steps in the same loop. - runner_static_cache: Dict[str, Any] = field(default_factory=dict, repr=False) - # Overflow compaction attempt counter (matches OpenClaw MAX_OVERFLOW_COMPACTION_ATTEMPTS) - overflow_compaction_attempts: int = 0 - # Tool result truncation attempted once per run (matches OpenClaw toolResultTruncationAttempted) - tool_result_truncation_attempted: bool = False - # Cooldown window to prefer cheap cleanup over repeated full compaction. - last_compaction_step: Optional[int] = None - last_cleanup_step: Optional[int] = None - # ``input + cache.read + output`` reported by the provider on the most - # recent finished assistant turn. When non-zero this beats the - # synthetic estimate from ``estimate_full_context_tokens`` because it - # is what the upstream will actually bill us for on the next turn - # (matches the "observed value wins" rule from docs/design/context-compaction-v2.md §B3). - last_observed_prompt_tokens: int = 0 - auto_failover: bool = False - # Entrypoint authorization is separate from persisted model_auto. Only a - # WebUI message route may set this bit; non-WebUI entrypoints use the default. - auto_failover_allowed: bool = False - model_candidates: List[RuntimeModel] = field(default_factory=list) - candidate_index: int = 0 - model_candidate_policy: Literal["fixed", "automatic", "configured"] = "automatic" - turn_user_id: Optional[str] = None - turn_additional_context: Optional[str] = None - stop_hook_active: bool = False - session_start_pending: bool = False - runtime_ports: Optional[ExternalRuntimePorts] = field(default=None, repr=False) - - @property - def trace_step(self) -> int: - """Session-cumulative step number for observability.""" - return self.trace_step_offset + self.step - - def should_abort(self) -> bool: - """Check if loop should abort""" - return self.abort_event.is_set() - - def signal_abort(self) -> None: - """Signal abort to stop loop, and cancel the current step task if running.""" - self.abort_event.set() - task = self._current_step_task - if task is not None and not task.done(): - task.cancel() - - -@dataclass -class LoopCallbacks: - """Callbacks for loop events""" - on_step_start: Optional[Callable[[int], Awaitable[None]]] = None - on_step_end: Optional[Callable[[int], Awaitable[None]]] = None - on_compaction: Optional[Callable[[], Awaitable[None]]] = None - on_error: Optional[Callable[[str], Awaitable[None]]] = None - on_reminder: Optional[Callable[[str], Awaitable[None]]] = None - # SSE event publishing callback (for TUI/WebUI real-time updates) - event_publish_callback: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None - # Runner-level callbacks (text delta, tool events, permissions, etc.) - # Type: Optional[RunnerCallbacks] - using Any to avoid circular import - runner_callbacks: Optional[Any] = None - - -@dataclass -class LoopResult: - """Result of loop execution""" - action: str # "stop", "continue", "compact", "error", "queued" - last_message: Optional[MessageInfo] = None - error: Optional[str] = None - provider_id: Optional[str] = None - model_id: Optional[str] = None - metadata: Dict[str, Any] = field(default_factory=dict) + def __init__(self, active_turns: MutableMapping[str, SessionTurn]): + self._active_turns = active_turns + + def get(self, session_id: str) -> Optional[SessionTurn]: + return self._active_turns.get(session_id) + + def acquire( + self, + session_id: str, + turn: SessionTurn, + ) -> Optional[_SessionLease]: + if session_id in self._active_turns: + return None + self._active_turns[session_id] = turn + return _SessionLease(session_id=session_id, turn=turn) + + def release(self, lease: _SessionLease) -> None: + if self._active_turns.get(lease.session_id) is lease.turn: + self._active_turns.pop(lease.session_id, None) + + def owns(self, lease: _SessionLease) -> bool: + return self._active_turns.get(lease.session_id) is lease.turn class SessionLoop: - """ - Session loop manager - - Handles the main session execution loop with support for: - - Message iteration - - Compaction triggers - - Subtask management - - Reminder injection - - Loop control (abort, pause, resume) - """ - - # Active loop contexts by session ID - _active_loops: Dict[str, LoopContext] = {} - _auto_failover_cooldowns: Dict[str, AutoFailoverCooldown] = {} + """Decide whether a persistent session should continue or settle.""" - @classmethod - def clear_auto_failover_state(cls, session_id: str) -> None: - """Clear process-local routing state when WebUI Auto is disabled.""" - cls._auto_failover_cooldowns.pop(session_id, None) + _active_turns: ClassVar[dict[str, SessionTurn]] = {} + _leases: ClassVar[_SessionLeaseRegistry] = _SessionLeaseRegistry( + _active_turns, + ) + _model_policy: ClassVar[ModelRoutingPolicy] = DEFAULT_MODEL_ROUTING_POLICY + _continuation_policy: ClassVar[ContinuationPolicy] = ( + DEFAULT_CONTINUATION_POLICY + ) @classmethod - async def validate_runtime_model( + async def run( cls, - provider_id: str, - model_id: str, - *, - config: Optional[Any] = None, - ) -> tuple[bool, str]: - """Validate a configured LLM candidate without a network health probe.""" - from flocks.config.config import Config - from flocks.provider.model_manager import get_model_manager - from flocks.provider.types import ModelType - - Provider._ensure_initialized() - config = config or await Config.get() - if provider_id in (getattr(config, "disabled_providers", None) or []): - return False, "provider_disabled" - enabled_providers = getattr(config, "enabled_providers", None) or [] - if enabled_providers and provider_id not in enabled_providers: - return False, "provider_disabled" - try: - await Provider.apply_config(config, provider_id=provider_id) - except Exception as exc: - log.warn("session.model.candidate_config_failed", { - "provider_id": provider_id, - "model_id": model_id, - "error": str(exc), - }) - return False, "provider_config_error" - - provider = Provider.get(provider_id) - if provider is None: - return False, "provider_not_found" - - definition = get_model_manager().get_model(provider_id, model_id) - if definition is None: - return False, "model_not_found" - if getattr(definition, "model_type", None) != ModelType.LLM: - return False, "not_llm" - - setting = get_model_manager().get_setting(provider_id, model_id) - if setting is not None and not setting.enabled: - return False, "model_disabled" - if not provider.is_configured(): - return False, "provider_not_configured" - return True, "available" + session_id: str, + provider_id: Optional[str] = None, + model_id: Optional[str] = None, + agent_name: Optional[str] = None, + callbacks: Optional[LoopCallbacks] = None, + working_directory: Optional[str] = None, + auto_failover: bool = False, + ) -> LoopResult: + """Run one session until queued and synthetic continuations settle.""" + active_turn = cls._leases.get(session_id) + if active_turn is not None: + log.info("session.already_running", {"session_id": session_id}) + cls._authorize_auto_failover(active_turn, auto_failover) + return LoopResult( + action="queued", + error="Loop already running", + ) - @classmethod - async def _build_model_candidates( - cls, - primary: RuntimeModel, - *, - route_seed: str, - preferred: Optional[RuntimeModel] = None, - config: Optional[Any] = None, - ) -> List[RuntimeModel]: - """Build a configured chain or the stable automatic discovery chain.""" - from flocks.config.config import Config - from flocks.provider.model_manager import get_model_manager - from flocks.provider.types import ModelType - - config = config or await Config.get() - await Provider.apply_config(config) - - configured_fallbacks = getattr(config, "fallback_providers", None) or [] - if configured_fallbacks: - candidates = [primary] - seen = {(primary.provider_id, primary.model_id)} - for index, raw in enumerate(configured_fallbacks): - provider_id = ( - raw.get("provider_id") - if isinstance(raw, dict) - else raw.provider_id - ) - model_id = ( - raw.get("model_id") - if isinstance(raw, dict) - else raw.model_id - ) - candidate = RuntimeModel( - provider_id=provider_id, - model_id=model_id, - ) - identity = (candidate.provider_id, candidate.model_id) - if identity in seen: - continue - seen.add(identity) + session = await Session.get_by_id(session_id) + if session is None: + log.warning("session.not_found", {"session_id": session_id}) + return LoopResult( + action="error", + error=f"Session {session_id} not found", + ) + if session.status != "active": + log.warning( + "session.not_active", + {"session_id": session_id, "status": session.status}, + ) + return LoopResult( + action="error", + error=f"Session {session_id} is {session.status}", + ) + if working_directory: + session = session.model_copy( + update={"directory": working_directory}, + ) - available, reason = await cls.validate_runtime_model( - candidate.provider_id, - candidate.model_id, - config=config, - ) - if not available: - log.warn("session.model.fallback_skipped", { - "provider_id": candidate.provider_id, - "model_id": candidate.model_id, - "configured_index": index, - "reason": reason, - }) - continue - candidates.append(candidate) - return candidates + if not provider_id or not model_id: + resolved_provider, resolved_model = await cls._resolve_model( + session, + provider_id, + model_id, + ) + provider_id = provider_id or resolved_provider + model_id = model_id or resolved_model - definitions = get_model_manager().list_models( - model_type=ModelType.LLM, - enabled_only=True, + primary_model = RuntimeModel( + provider_id=provider_id, + model_id=model_id, ) - discovered = { - RuntimeModel(definition.provider_id, definition.id) - for definition in definitions - } - discovered.discard(primary) - - same_provider: List[RuntimeModel] = [] - other_providers: List[RuntimeModel] = [] - for candidate in sorted( - discovered, - key=lambda item: (item.provider_id, item.model_id), - ): - available, reason = await cls.validate_runtime_model( - candidate.provider_id, - candidate.model_id, - config=config, - ) - if not available: - log.debug("session.model.fallback_skipped", { - "provider_id": candidate.provider_id, - "model_id": candidate.model_id, - "reason": reason, - }) - continue - - if candidate.provider_id == primary.provider_id: - same_provider.append(candidate) - else: - other_providers.append(candidate) - - candidates = [primary] - for tier, pool in ( - ("same_provider", same_provider), - ("other_provider", other_providers), - ): - if not pool: - continue - selected = ( - preferred - if preferred is not None and preferred in pool - else cls._stable_candidate_choice(pool, route_seed, tier) + auto_failover = bool( + auto_failover + and is_model_auto_session_category( + getattr(session, "category", "user"), ) - candidates.append(selected) - return candidates - - @staticmethod - def _stable_candidate_choice( - candidates: List[RuntimeModel], - route_seed: str, - tier: str, - ) -> RuntimeModel: - """Choose pseudo-randomly without Python's process-randomized hash().""" - ordered = sorted( - candidates, - key=lambda item: (item.provider_id, item.model_id), ) - digest = hashlib.sha256( - f"{route_seed}\0{tier}".encode("utf-8") - ).digest() - index = int.from_bytes(digest[:8], "big") % len(ordered) - return ordered[index] + session.provider = provider_id + session.model = model_id - @classmethod - async def validate_auto_configuration(cls) -> tuple[bool, str]: - """Validate that a newly selected Auto mode has a usable chain.""" - from flocks.config.config import Config + trace_offset = await cls._load_trace_offset(session_id) + runtime_callbacks = callbacks or LoopCallbacks() + turn = SessionTurn( + session=session, + provider_id=provider_id, + model_id=model_id, + agent_name=agent_name or session.agent or "rex", + callbacks=runtime_callbacks, + session_store=DefaultSessionContext(session), + trace_step_offset=trace_offset, + auto_failover=auto_failover, + auto_failover_allowed=auto_failover, + model_candidates=[primary_model], + candidate_index=0, + session_start_pending=trace_offset == 0, + model_policy=cls._model_policy, + continuation_policy=cls._continuation_policy, + ) + lease_or_result = await cls._acquire_lease(session_id, turn) + if not isinstance(lease_or_result, _SessionLease): + return lease_or_result + lease = lease_or_result + + settled = False + processed_user_id: Optional[str] = None + try: + await cls._mark_busy(session_id, runtime_callbacks) + await cls._recover_orphan_tools(session_id) - default_llm = await Config.resolve_default_llm() - if not default_llm: - return False, "default_model_missing" - primary = RuntimeModel( - default_llm["provider_id"], - default_llm["model_id"], - ) - available, reason = await cls.validate_runtime_model( - primary.provider_id, - primary.model_id, - ) - if not available: - return False, f"primary_{reason}" - return True, "available" + while True: + continuation_policy = ( + turn.continuation_policy or cls._continuation_policy + ) + try: + await continuation_policy.prepare_logical_turn(turn) + processed_user_id = ( + turn.prepared_user_id or processed_user_id + ) + outcome = await cls._run_logical_input(turn) + processed_user_id = ( + outcome.state.current_user_id or processed_user_id + ) + if await cls._should_continue( + turn, + continuation_policy, + outcome, + ): + continue + except Exception as exc: + outcome = await cls._handle_execution_error( + turn, + exc, + processed_user_id, + ) - @classmethod - def _active_cooldown_model( - cls, - session_id: str, - primary: RuntimeModel, - ) -> Optional[RuntimeModel]: - """Return a still-valid cooldown target for the current primary.""" - cooldown = cls._auto_failover_cooldowns.get(session_id) - if cooldown is None: - return None - if cooldown.expires_at <= time.monotonic() or cooldown.primary != primary: - cls._auto_failover_cooldowns.pop(session_id, None) - return None - return cooldown.model + if await cls._settle_or_continue( + lease, + runtime_callbacks, + processed_user_id, + ): + continue + + settled = True + return cls._to_loop_result(turn, outcome) + finally: + if not settled and cls._leases.owns(lease): + await cls._release_session(lease, runtime_callbacks) @classmethod - def _cooldown_candidate_index( + async def _run_logical_input( cls, - session_id: str, - candidates: List[RuntimeModel], - ) -> int: - if not candidates: - return 0 - cooldown_model = cls._active_cooldown_model(session_id, candidates[0]) - if cooldown_model is None: - return 0 - try: - return candidates.index(cooldown_model) - except ValueError: - cls._auto_failover_cooldowns.pop(session_id, None) - return 0 + turn: SessionTurn, + ) -> AgentRunOutcome[Any]: + """Execute one prepared logical input through AgentLoop.""" + turn.reset() + return await AgentLoop().run( + turn, + StepEngine.from_turn(turn), + ) + + @staticmethod + async def _should_continue( + turn: SessionTurn, + continuation_policy: ContinuationPolicy, + outcome: AgentRunOutcome[Any], + ) -> bool: + """Resolve queued input, goal, and TurnFinish continuation.""" + if outcome.status == AgentRunStatus.INPUT_AVAILABLE: + return True + if ( + outcome.status == AgentRunStatus.COMPLETED + and outcome.step_result is not None + ): + continuation = await continuation_policy.resolve(turn, outcome) + return continuation.should_continue + return False - @classmethod - def _select_candidate(cls, ctx: LoopContext, index: int) -> None: - candidate = ctx.model_candidates[index] - ctx.candidate_index = index - ctx.provider_id = candidate.provider_id - ctx.model_id = candidate.model_id - ctx.session.provider = candidate.provider_id - ctx.session.model = candidate.model_id - # Prompt and model-capability caches are keyed in most places, but a - # fresh dict makes the runtime rebuild guarantee explicit. The tool - # loop guard is turn state rather than model state, so it must survive - # a provider switch to keep repeated-tool protection effective. - tool_loop_guard = ctx.runner_static_cache.get("tool_loop_guard") - ctx.runner_static_cache.clear() - if tool_loop_guard is not None: - ctx.runner_static_cache["tool_loop_guard"] = tool_loop_guard - @classmethod def is_running(cls, session_id: str) -> bool: - """Check if loop is running for session""" - return session_id in cls._active_loops - + """Return whether this process owns the session.""" + return session_id in cls._active_turns + @classmethod - def get_context(cls, session_id: str) -> Optional[LoopContext]: - """Get loop context for session""" - return cls._active_loops.get(session_id) - + def get_context(cls, session_id: str) -> Optional[SessionTurn]: + """Return the active turn used by public session controls.""" + return cls._active_turns.get(session_id) + @classmethod def abort(cls, session_id: str) -> bool: - """Abort running loop""" - ctx = cls._active_loops.get(session_id) - if ctx: - ctx.signal_abort() - return True - return False - + """Abort one active session run.""" + turn = cls._active_turns.get(session_id) + if turn is None: + return False + turn.signal_abort() + return True + @classmethod def abort_children(cls, parent_session_id: str) -> int: - """Abort all child loops whose session.parent_id matches, recursively.""" + """Abort all active descendants of one parent session.""" aborted = 0 child_ids = [ - sid for sid, ctx in list(cls._active_loops.items()) - if getattr(ctx.session, 'parent_id', None) == parent_session_id + session_id + for session_id, turn in list(cls._active_turns.items()) + if getattr(turn.session, "parent_id", None) == parent_session_id ] - for sid in child_ids: - ctx = cls._active_loops.get(sid) - if ctx and not ctx.should_abort(): - ctx.signal_abort() + for session_id in child_ids: + turn = cls._active_turns.get(session_id) + if turn is not None and not turn.aborted: + turn.signal_abort() aborted += 1 - aborted += cls.abort_children(sid) + aborted += cls.abort_children(session_id) return aborted @classmethod - async def _publish_runtime_event( - cls, - callbacks: "LoopCallbacks", - event_name: str, - payload: Dict[str, Any], - ) -> None: - if not callbacks.event_publish_callback: - return - try: - await callbacks.event_publish_callback(event_name, payload) - except Exception as exc: - log.debug("loop.runtime_event.publish_failed", { - "event": event_name, - "error": str(exc), - }) - - @classmethod - async def _publish_turn_stopped( - cls, - callbacks: "LoopCallbacks", - session_id: str, - *, - step: int, - stop_reason: str, - ) -> None: - turn_state = set_turn_state( - session_id, - step=step, - status="stopped", - stop_reason=stop_reason, - queued_message_detected=False, - ) - await cls._publish_runtime_event( - callbacks, - "turn.stopped", - turn_state.model_dump(by_alias=True), - ) - - @classmethod - async def _publish_session_status( - cls, - callbacks: "LoopCallbacks", - session_id: str, - status: str, - ) -> None: - if not callbacks.event_publish_callback: - return - try: - await callbacks.event_publish_callback("session.status", { - "sessionID": session_id, - "status": {"type": status}, - }) - except Exception as exc: - log.debug("loop.session_status.publish_failed", { - "session_id": session_id, - "status": status, - "error": str(exc), - }) + def clear_auto_failover_state(cls, session_id: str) -> None: + """Clear model-routing cooldown state for one session.""" + cls._model_policy.clear(session_id) @classmethod - async def _publish_session_notice( + async def validate_runtime_model( cls, - callbacks: "LoopCallbacks", - session_id: str, + provider_id: str, + model_id: str, *, - level: str, - message: str, - details: Optional[Dict[str, Any]] = None, - ) -> None: - if not callbacks.event_publish_callback: - return - try: - await callbacks.event_publish_callback("session.notice", { - "sessionID": session_id, - "level": level, - "message": message, - "details": details or {}, - }) - except Exception as exc: - log.debug("loop.session_notice.publish_failed", {"error": str(exc)}) - - @classmethod - def _has_recent_compaction_cooldown(cls, ctx: LoopContext) -> bool: - return ( - ctx.last_compaction_step is not None - and (ctx.step - ctx.last_compaction_step) <= POST_COMPACTION_COOLDOWN_STEPS + config: Optional[Any] = None, + ) -> tuple[bool, str]: + """Validate one provider/model candidate.""" + return await cls._model_policy.validate_runtime_model( + provider_id, + model_id, + config=config, ) @classmethod - async def _detect_queued_user_message( - cls, - _session_id: str, - post_messages: List[MessageInfo], - current_user_id: str, - _last_message: Optional[MessageInfo], - ) -> Optional[MessageInfo]: - if not post_messages: - return None - - newest_user = None - for msg in reversed(post_messages): - if msg.role == MessageRole.USER: - newest_user = msg - break + async def validate_auto_configuration(cls) -> tuple[bool, str]: + """Validate that Auto mode has an available primary model.""" + from flocks.config.config import Config - if newest_user is None: - return None - if newest_user.id <= current_user_id: - return None - # A fallback assistant is created after a user message that arrived - # while the primary model was running. Its newer ID must not make that - # user message look handled; the current turn's user ID is the stable - # boundary for queued work. - return newest_user - - @classmethod - async def run( - cls, - session_id: str, - provider_id: Optional[str] = None, - model_id: Optional[str] = None, - agent_name: Optional[str] = None, - callbacks: Optional[LoopCallbacks] = None, - working_directory: Optional[str] = None, - auto_failover: bool = False, - ) -> LoopResult: - """Run one session through the lifecycle-owned SessionHost.""" - from flocks.session.session_host import ( - SessionHost, - SessionHostDependencies, + default_llm = await Config.resolve_default_llm() + if not default_llm: + return False, "default_model_missing" + available, reason = await cls.validate_runtime_model( + default_llm["provider_id"], + default_llm["model_id"], ) + if not available: + return False, f"primary_{reason}" + return True, "available" - host = SessionHost( - dependencies=SessionHostDependencies( - create_context=LoopContext, - create_callbacks=LoopCallbacks, - create_result=LoopResult, - resolve_model=cls._resolve_model, - run_logical_turn=cls._run_loop, - publish_session_status=cls._publish_session_status, - ), - active_contexts=cls._active_loops, - ) - return await host.run( - session_id=session_id, - provider_id=provider_id, - model_id=model_id, - agent_name=agent_name, - callbacks=callbacks, - working_directory=working_directory, - auto_failover=auto_failover, - ) - @staticmethod async def _resolve_model( session: Any, @@ -624,105 +337,122 @@ async def _resolve_model( *, include_source: bool = False, ) -> tuple: - """ - Resolve provider_id and model_id for session execution. - - Priority: - 1. Explicitly passed provider_id / model_id (already handled by caller) - 2. Session's stored model/provider (set during Session.create) - 3. Agent model override from Storage (set via WebUI) - 4. Agent-specific model from AgentInfo.model (agent.yaml / config) - 5. Parent session's model/provider (inherits from parent — TUI/CLI default) - 6. Global default LLM (default_models.llm -> config.model) - 7. Environment variables - 8. Hardcoded fallback - - Returns: - (provider_id, model_id) tuple - """ + """Resolve the concrete model used to open a session turn.""" import os - + resolved_provider = provider_id resolved_model = model_id source = "explicit" if provider_id and model_id else "unknown" - - # Priority 2: Session's stored model/provider - if (not resolved_provider or not resolved_model) and Session.has_pinned_model(session): + + if ( + (not resolved_provider or not resolved_model) + and Session.has_pinned_model(session) + ): resolved_provider = resolved_provider or session.provider resolved_model = resolved_model or session.model if resolved_provider and resolved_model: source = "session" - - # Priority 3: Agent model override from Storage (set via WebUI) + if not resolved_provider or not resolved_model: - agent_name = getattr(session, 'agent', None) + agent_name = getattr(session, "agent", None) if agent_name: try: from flocks.storage.storage import Storage + overrides = await Storage.read("agent/model_overrides") if isinstance(overrides, dict) and agent_name in overrides: override = overrides[agent_name] - override_provider = override.get('providerID') - override_model = override.get('modelID') + override_provider = override.get("providerID") + override_model = override.get("modelID") if override_provider and override_model: resolved_provider = override_provider resolved_model = override_model source = "agent_override" - except Exception as _e: - log.debug("loop.resolve_model.storage_override_failed", {"error": str(_e)}) - - # Priority 4: Agent-specific model from AgentInfo + except Exception as exc: + log.debug( + "loop.resolve_model.storage_override_failed", + {"error": str(exc)}, + ) + if not resolved_provider or not resolved_model: - agent_name = getattr(session, 'agent', None) + agent_name = getattr(session, "agent", None) if agent_name: try: from flocks.agent.registry import Agent + agent_info = await Agent.get(agent_name) if agent_info and agent_info.model: - resolved_provider = resolved_provider or agent_info.model.provider_id - resolved_model = resolved_model or agent_info.model.model_id + resolved_provider = ( + resolved_provider + or agent_info.model.provider_id + ) + resolved_model = ( + resolved_model or agent_info.model.model_id + ) if resolved_provider and resolved_model: source = "agent" - except Exception as _e: - log.debug("loop.resolve_model.agent_model_failed", {"error": str(_e)}) - - # Priority 5: Parent session's model/provider (inherit from Rex etc.) + except Exception as exc: + log.debug( + "loop.resolve_model.agent_model_failed", + {"error": str(exc)}, + ) + if not resolved_provider or not resolved_model: - parent_id = getattr(session, 'parent_id', None) + parent_id = getattr(session, "parent_id", None) if parent_id: try: parent = await Session.get_by_id(parent_id) if Session.has_pinned_model(parent): - resolved_provider = resolved_provider or getattr(parent, 'provider', None) - resolved_model = resolved_model or getattr(parent, 'model', None) + resolved_provider = resolved_provider or getattr( + parent, + "provider", + None, + ) + resolved_model = resolved_model or getattr( + parent, + "model", + None, + ) if resolved_provider and resolved_model: source = "parent_session" - except Exception as _e: - log.debug("loop.resolve_model.parent_failed", {"error": str(_e)}) - - # Priority 6: Global default LLM (default_models.llm -> config.model) + except Exception as exc: + log.debug( + "loop.resolve_model.parent_failed", + {"error": str(exc)}, + ) + if not resolved_provider or not resolved_model: try: from flocks.config.config import Config + default_llm = await Config.resolve_default_llm() if default_llm: - resolved_provider = resolved_provider or default_llm["provider_id"] - resolved_model = resolved_model or default_llm["model_id"] + resolved_provider = ( + resolved_provider or default_llm["provider_id"] + ) + resolved_model = ( + resolved_model or default_llm["model_id"] + ) if resolved_provider and resolved_model: source = "config" - except Exception as _e: - log.debug("loop.resolve_model.config_default_failed", {"error": str(_e)}) - - # Priority 7: Environment variables + except Exception as exc: + log.debug( + "loop.resolve_model.config_default_failed", + {"error": str(exc)}, + ) + if not resolved_provider: resolved_provider = os.environ.get("LLM_PROVIDER") if not resolved_model: resolved_model = os.environ.get("LLM_MODEL") if resolved_provider and resolved_model and source == "unknown": source = "env_default" - - # Priority 8: Hardcoded fallback - from flocks.session.core.defaults import fallback_provider_id, fallback_model_id + + from flocks.session.core.defaults import ( + fallback_model_id, + fallback_provider_id, + ) + resolved_provider = resolved_provider or fallback_provider_id() resolved_model = resolved_model or fallback_model_id() if source == "unknown": @@ -733,1864 +463,235 @@ async def _resolve_model( return resolved_provider, resolved_model @classmethod - async def _reset_auto_turn_candidates( + def _to_loop_result( cls, - ctx: LoopContext, - primary: RuntimeModel, - user_message_id: str, - config: Any, - ) -> int: - """Rebuild and activate the configured or automatic chain for one turn.""" - configured = bool(getattr(config, "fallback_providers", None)) - if configured: - cls.clear_auto_failover_state(ctx.session.id) - preferred = None - else: - preferred = cls._active_cooldown_model(ctx.session.id, primary) - - ctx.model_candidates = await cls._build_model_candidates( - primary, - route_seed=f"{ctx.session.id}:{user_message_id}", - preferred=preferred, - config=config, + turn: SessionTurn, + outcome: AgentRunOutcome[Any], + ) -> LoopResult: + loop_error = ( + outcome.error + if outcome.status + in { + AgentRunStatus.RETRYABLE_FAILURE, + AgentRunStatus.FATAL_FAILURE, + AgentRunStatus.CONTEXT_OVERFLOW, + } + else None ) - ctx.model_candidate_policy = ( - "configured" if configured else "automatic" + unhandled_runtime_error = bool( + outcome.state.metadata.get("unhandled_runtime_error"), ) - ctx.auto_failover = True - next_index = ( - 0 - if configured - else cls._cooldown_candidate_index( - ctx.session.id, - ctx.model_candidates, - ) + return LoopResult( + action=( + "error" + if ( + unhandled_runtime_error + or (turn.auto_failover and loop_error) + ) + else "stop" + ), + last_message=outcome.last_message, + error=( + loop_error + if unhandled_runtime_error or turn.auto_failover + else None + ), + provider_id=turn.provider_id, + model_id=turn.model_id, + metadata={ + "steps": turn.step, + "session_id": turn.session.id, + "last_compaction_step": turn.last_compaction_step, + **( + {"aborted": True} + if outcome.status == AgentRunStatus.ABORTED + else {} + ), + }, ) - cls._select_candidate(ctx, next_index) - return next_index - - @classmethod - async def _prepare_auto_turn( - cls, - ctx: LoopContext, - last_user: MessageInfo, - ) -> bool: - """Synchronize routing when the loop advances to a real WebUI turn. - Returns: - True when ``last_user`` starts a new non-synthetic user turn. - """ - if last_user.id == ctx.turn_user_id: - return False + @staticmethod + def _authorize_auto_failover( + turn: SessionTurn, + requested: bool, + ) -> None: + if requested and is_model_auto_session_category( + getattr(turn.session, "category", "user"), + ): + turn.auto_failover_allowed = True - parts = await Message.parts(last_user.id, ctx.session.id) - if any(bool(getattr(part, "synthetic", False)) for part in parts): - return False - - if ctx.turn_user_id is None: - ctx.turn_user_id = last_user.id - if ctx.auto_failover and ctx.auto_failover_allowed: - from flocks.config.config import Config - - primary = ctx.model_candidates[0] - config = await Config.get() - await cls._reset_auto_turn_candidates( - ctx, - primary, - last_user.id, - config=config, - ) - return True - - ctx.turn_user_id = last_user.id - persisted_session = await Session.get_by_id(ctx.session.id) - persisted_model_auto = bool( - persisted_session - and is_model_auto_session_category( - getattr(persisted_session, "category", "user") - ) - and getattr(persisted_session, "model_auto", False) - ) - persisted_auto = persisted_model_auto and ctx.auto_failover_allowed - - user_model = getattr(last_user, "model", None) - user_provider_id = None - user_model_id = None - if isinstance(user_model, dict): - user_provider_id = user_model.get("providerID") or user_model.get("provider_id") - user_model_id = user_model.get("modelID") or user_model.get("model_id") - - if not persisted_auto: - ctx.auto_failover = False - if not persisted_model_auto: - cls.clear_auto_failover_state(ctx.session.id) - ctx.auto_failover_allowed = False - provider_id = ( - getattr(persisted_session, "provider", None) - if Session.has_pinned_model(persisted_session) - else user_provider_id - ) or ctx.provider_id - model_id = ( - getattr(persisted_session, "model", None) - if Session.has_pinned_model(persisted_session) - else user_model_id - ) or ctx.model_id - ctx.model_candidates = [RuntimeModel(provider_id, model_id)] - ctx.model_candidate_policy = "fixed" - cls._select_candidate(ctx, 0) - log.info("session.model.auto_disabled_for_turn", { - "session_id": ctx.session.id, - "provider_id": provider_id, - "model_id": model_id, - }) - return True - - from flocks.config.config import Config - - config = await Config.get() - previous = RuntimeModel(ctx.provider_id, ctx.model_id) - default_llm = await Config.resolve_default_llm() - primary = RuntimeModel( - provider_id=(default_llm or {}).get("provider_id") or user_provider_id or ctx.provider_id, - model_id=(default_llm or {}).get("model_id") or user_model_id or ctx.model_id, - ) - next_index = await cls._reset_auto_turn_candidates( - ctx, - primary, - last_user.id, - config=config, - ) - active = ctx.model_candidates[next_index] - log.info("session.model.auto_turn_reset", { - "session_id": ctx.session.id, - "from_provider_id": previous.provider_id, - "from_model_id": previous.model_id, - "to_provider_id": active.provider_id, - "to_model_id": active.model_id, - "cooldown_active": next_index > 0, - }) - return True - - @classmethod - async def _run_user_prompt_submit_hook( - cls, - ctx: LoopContext, - last_user: MessageInfo, - ) -> None: - """Run UserPromptSubmit once for a newly observed real user turn.""" - try: - from flocks.hooks.pipeline import HookPipeline - - prompt = await Message.get_text_content(last_user) - hook_ctx = await HookPipeline.run_user_prompt_submit({ - "sessionID": ctx.session.id, - "workspace": ctx.session.directory, - "agent": getattr(last_user, "agent", None) or ctx.agent_name, - "model": { - "providerID": ctx.provider_id, - "modelID": ctx.model_id, - }, - "messageID": last_user.id, - "prompt": prompt, - }) - additional_context = hook_ctx.output.get("additionalContext") - if isinstance(additional_context, str) and additional_context.strip(): - ctx.turn_additional_context = additional_context.strip() - except Exception as exc: - log.debug("loop.hook.user_prompt_submit.error", { - "session_id": ctx.session.id, - "message_id": last_user.id, - "error": str(exc), - }) + @staticmethod + async def _load_trace_offset(session_id: str) -> int: + try: + messages = await Message.list(session_id) + return sum( + 1 for message in messages if message.role == "assistant" + ) + except Exception as exc: + log.debug("session.trace_offset.error", {"error": str(exc)}) + return 0 @classmethod - async def _run_turn_finish_hook( + async def _acquire_lease( cls, - ctx: LoopContext, - callbacks: LoopCallbacks, - last_user: MessageInfo, - last_message: MessageInfo, - ) -> bool: - """Run TurnFinish and continue the loop when the hook blocks stopping.""" - try: - from flocks.hooks.pipeline import HookPipeline - - hook_user = last_user - if ctx.turn_user_id: - hook_user = ( - await Message.get(ctx.session.id, ctx.turn_user_id) - or last_user + session_id: str, + turn: SessionTurn, + ) -> _SessionLease | LoopResult: + async with Session.lifecycle_lock(session_id): + latest_session = await Session.get_by_id(session_id) + if latest_session is None: + return LoopResult( + action="error", + error=f"Session {session_id} not found", ) - user_text = await Message.get_text_content(hook_user) - assistant_text = await Message.get_text_content(last_message) - hook_ctx = await HookPipeline.run_turn_finish({ - "sessionID": ctx.session.id, - "workspace": ctx.session.directory, - "agent": getattr(last_message, "agent", None) or ctx.agent_name, - "model": { - "providerID": ctx.provider_id, - "modelID": ctx.model_id, - }, - "step": ctx.trace_step, - "userMessage": { - "id": hook_user.id, - "content": user_text, - }, - "assistantMessage": { - "id": last_message.id, - "content": assistant_text, - }, - "finishReason": "stop", - "stopHookActive": ctx.stop_hook_active, - }) - except Exception as exc: - log.debug("loop.hook.turn_finish.error", { - "session_id": ctx.session.id, - "message_id": getattr(last_message, "id", None), - "error": str(exc), - }) - return False - - decision = str(hook_ctx.output.get("decision") or "").strip().lower() - reason = str(hook_ctx.output.get("reason") or "").strip() - if decision != "block": - return False - if not reason: - log.warn("loop.hook.turn_finish.missing_reason", { - "session_id": ctx.session.id, - "message_id": last_message.id, - }) - return False - if ctx.should_abort(): - log.info("loop.hook.turn_finish.ignored_after_abort", { - "session_id": ctx.session.id, - "message_id": last_message.id, - }) - return False - - try: - if ctx.session_ctx: - post_hook_messages = await ctx.session_ctx.get_messages() - else: - post_hook_messages = await Message.list(ctx.session.id) - queued_user = await cls._detect_queued_user_message( - ctx.session.id, - post_hook_messages, - last_user.id, - last_message, - ) - except Exception as exc: - queued_user = None - log.debug("loop.hook.turn_finish.queued_recheck_error", { - "session_id": ctx.session.id, - "error": str(exc), - }) - if queued_user is not None: - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="continued", - continue_reason="queued_message", - queued_message_detected=True, - ) - await cls._publish_runtime_event(callbacks, "turn.continued", { - **turn_state.model_dump(by_alias=True), - "queuedUserMessageID": queued_user.id, - }) - log.info("loop.hook.turn_finish.queued_message_won", { - "session_id": ctx.session.id, - "queued_user_id": queued_user.id, - "source_assistant_message_id": last_message.id, - }) - return True + if latest_session.status != "active": + return LoopResult( + action="error", + error=f"Session {session_id} is {latest_session.status}", + ) + if Session.is_lifecycle_transitioning(session_id): + return LoopResult( + action="error", + error=f"Session {session_id} is changing lifecycle state", + ) + lease = cls._leases.acquire(session_id, turn) + if lease is None: + return LoopResult( + action="queued", + error="Loop already running", + ) + return lease - from flocks.agent.registry import Agent - from flocks.session.core.defaults import DEFAULT_MAX_TOOL_STEPS + @staticmethod + async def _mark_busy( + session_id: str, + callbacks: LoopCallbacks, + ) -> None: + SessionStatus.set(session_id, SessionStatusBusy()) + await SessionEventSink.session_status(callbacks, session_id, "busy") + @staticmethod + async def _recover_orphan_tools(session_id: str) -> None: try: - agent = await Agent.get( - getattr(last_message, "agent", None) or ctx.agent_name - ) - except Exception as exc: - log.debug("loop.hook.turn_finish.agent_load_error", { - "session_id": ctx.session.id, - "error": str(exc), - }) - agent = None - max_steps = ( - agent.steps - if agent is not None and getattr(agent, "steps", None) is not None - else DEFAULT_MAX_TOOL_STEPS - ) - if ctx.trace_step >= max_steps: - log.warn("loop.hook.turn_finish.step_limit", { - "session_id": ctx.session.id, - "step": ctx.trace_step, - "max_steps": max_steps, - }) - return False + from flocks.session.orphan_tools import abort_orphan_running_parts - try: - continuation = await Message.create( - session_id=ctx.session.id, - role=MessageRole.USER, - content=reason, - agent=getattr(hook_user, "agent", None) or ctx.agent_name, - model={ - "providerID": ctx.provider_id, - "modelID": ctx.model_id, - }, - synthetic=True, - part_metadata={ - "turnFinishContinuation": True, - "stopHookActive": True, - "sourceAssistantMessageID": last_message.id, - }, - ) + await abort_orphan_running_parts(session_id) except Exception as exc: - log.error("loop.hook.turn_finish.continuation_error", { - "session_id": ctx.session.id, - "error": str(exc), - }) - return False - ctx.stop_hook_active = True - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="continued", - continue_reason="turn_finish_hook", - queued_message_detected=False, - ) - await cls._publish_runtime_event(callbacks, "turn.continued", { - **turn_state.model_dump(by_alias=True), - "turnFinishMessageID": continuation.id, - }) - log.info("loop.continuing_for_turn_finish_hook", { - "session_id": ctx.session.id, - "continuation_message_id": continuation.id, - "source_assistant_message_id": last_message.id, - }) - return True - - @classmethod - async def _finalize_deferred_failure( - cls, - ctx: LoopContext, - failure: Any, - last_user: MessageInfo, - ) -> None: - """Persist only the final Auto candidate failure.""" - if not failure.assistant_message_id: - assistant = await Message.create( - session_id=ctx.session.id, - role=MessageRole.ASSISTANT, - content="", - agent=getattr(last_user, "agent", None) or ctx.agent_name or "rex", - model_id=ctx.model_id, - provider_id=ctx.provider_id, - parent_id=last_user.id, - error=failure.error_data, - finish="error", - ) - failure.assistant_message_id = assistant.id - return - await Message.update( - ctx.session.id, - failure.assistant_message_id, - error=failure.error_data, - finish="error", - ) - - @classmethod - async def _process_model_step( - cls, - ctx: LoopContext, - callbacks: LoopCallbacks, - snapshot: ModelTurnSnapshot[MessageInfo], - ) -> StepResult: - """Run one candidate attempt; provider-local retries stay in runner.""" - from flocks.session.runner import RunnerCallbacks, SessionRunner - from flocks.session.step_engine import SessionStepEngine - - runner_callbacks = callbacks.runner_callbacks - if runner_callbacks is None: - runner_callbacks = RunnerCallbacks() - if ( - callbacks.event_publish_callback - and not runner_callbacks.event_publish_callback - ): - runner_callbacks.event_publish_callback = ( - callbacks.event_publish_callback + log.warn( + "session.orphan_cleanup_failed", + {"session_id": session_id, "error": str(exc)}, ) - runner = SessionRunner( - session=ctx.session, - provider_id=ctx.provider_id, - model_id=ctx.model_id, - agent_name=ctx.agent_name, - abort_event=ctx.abort_event, - callbacks=runner_callbacks, - session_ctx=ctx.session_ctx, - memory_bootstrap_data=ctx.memory_bootstrap_data, - static_cache=ctx.runner_static_cache, - defer_step_errors=ctx.auto_failover, - failover_available=( - ctx.auto_failover - and ctx.candidate_index + 1 < len(ctx.model_candidates) - ), - turn_additional_context=ctx.turn_additional_context, - session_start_pending=ctx.session_start_pending, - runtime_ports=ctx.runtime_ports, - ) - step_engine = SessionStepEngine(runner) - result = await step_engine.run(snapshot) - if step_engine.session_start_fired: - ctx.session_start_pending = False - return result - - @classmethod - def _create_host_step_engine( - cls, - ctx: LoopContext, - callbacks: LoopCallbacks, - ) -> Any: - """Compose one-candidate execution with host-owned model recovery.""" - from flocks.session.runtime_services import SessionLoopStepEngine - from flocks.session.session_host import SessionHostStepEngine - - return SessionHostStepEngine( - context=ctx, - callbacks=callbacks, - attempt_engine=SessionLoopStepEngine(ctx, callbacks, cls), - cooldowns=cls._auto_failover_cooldowns, - cooldown_factory=AutoFailoverCooldown, - select_candidate=cls._select_candidate, - finalize_failure=cls._finalize_deferred_failure, - publish_event=cls._publish_runtime_event, - rate_limit_cooldown_seconds=RATE_LIMIT_COOLDOWN_SECONDS, - chain_exhaustion_cooldown_seconds=( - CHAIN_EXHAUSTION_COOLDOWN_SECONDS - ), - ) - - @classmethod - async def _process_step_with_failover( - cls, - ctx: LoopContext, - callbacks: LoopCallbacks, - messages: List[MessageInfo], - last_user: MessageInfo, - ) -> StepResult: - """Compatibility entry point for host-owned cross-model recovery.""" - snapshot = ModelTurnSnapshot( - session_id=ctx.session.id, - agent_name=ctx.agent_name, - active_model=RuntimeModel(ctx.provider_id, ctx.model_id), - model_turn_index=ctx.step, - trace_step=ctx.trace_step, - messages=tuple(messages), - last_user=last_user, - ) - return await cls._create_host_step_engine(ctx, callbacks).run(snapshot) - @staticmethod - def _log_step_complete(ctx: LoopContext, duration_ms: int) -> None: - """Record model-turn latency from the session StepEngine adapter.""" - log.debug( - "loop.step_complete", - { - "session_id": ctx.session.id, - "step": ctx.step, - "duration_ms": duration_ms, - }, + async def _handle_execution_error( + turn: SessionTurn, + error: Exception, + processed_user_id: Optional[str], + ) -> AgentRunOutcome[Any]: + session_id = turn.session.id + log.error( + "session.execution_error", + {"session_id": session_id, "error": str(error)}, ) - - @classmethod - async def _complete_model_turn( - cls, - ctx: LoopContext, - callbacks: LoopCallbacks, - state: AgentRunState[MessageInfo], - step_result: StepResult, - ) -> ModelTurnBoundary[MessageInfo]: - """Expose the persisted state written by one completed model turn.""" - if callbacks.on_step_end: - await callbacks.on_step_end(ctx.step) - - if step_result.error and callbacks.on_error: - await callbacks.on_error(step_result.error) - - if ctx.session_ctx: - post_messages = await ctx.session_ctx.get_messages() - else: - post_messages = await Message.list(ctx.session.id) - - last_user = state.metadata.get("last_user") - last_message = next( - ( - message - for message in reversed(post_messages) - if message.role == MessageRole.ASSISTANT - and ( - not ctx.auto_failover - or last_user is None - or getattr(message, "parentID", None) == last_user.id + if turn.callbacks.on_error: + try: + await turn.callbacks.on_error(str(error)) + except Exception as callback_error: + log.debug( + "session.error_callback_failed", + {"error": str(callback_error)}, ) - ), - None, - ) - state.metadata["last_message"] = last_message - - queued_user = None - if last_user is not None: - queued_user = await cls._detect_queued_user_message( - ctx.session.id, - post_messages, - last_user.id, - last_message, - ) - - if queued_user is not None: - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="continued", - continue_reason="queued_message", - queued_message_detected=True, - ) - await cls._publish_runtime_event( - callbacks, - "turn.continued", - { - **turn_state.model_dump(by_alias=True), - "queuedUserMessageID": queued_user.id, - }, - ) - log.info( - "loop.continuing_for_queued_message", - { - "session_id": ctx.session.id, - "queued_user_id": queued_user.id, - "last_assistant_id": ( - last_message.id if last_message else None - ), - }, - ) - elif step_result.action == "continue": - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="continued", - continue_reason="tool_calls", - queued_message_detected=False, - ) - await cls._publish_runtime_event( - callbacks, - "turn.continued", - turn_state.model_dump(by_alias=True), - ) - elif step_result.error: - await cls._publish_turn_stopped( - callbacks, - ctx.session.id, - step=ctx.step, - stop_reason=step_result.error, - ) - - return ModelTurnBoundary( - messages=tuple(post_messages), - last_message=last_message, - queued_inputs=QueuedInputBatch( - messages=(queued_user,) if queued_user is not None else (), - cursor=queued_user.id if queued_user is not None else None, - ), - ) - - @classmethod - async def _resolve_continuation( - cls, - ctx: LoopContext, - callbacks: LoopCallbacks, - state: AgentRunState[MessageInfo], - step_result: StepResult, - ) -> ContinuationDecision[MessageInfo]: - """Resolve goal and TurnFinish continuation after a natural stop.""" - last_user = state.metadata.get("last_user") - last_message = state.metadata.get("last_message") - if last_user is None or last_message is None: - await cls._publish_turn_stopped( - callbacks, - ctx.session.id, - step=ctx.step, - stop_reason="stop", - ) - return ContinuationDecision() - - try: - content_result = Message.get_text_content(last_message) - last_response = ( - await content_result - if inspect.isawaitable(content_result) - else content_result - ) - except Exception as exc: - log.warn( - "goal.last_response.error", - { - "session_id": ctx.session.id, - "message_id": getattr(last_message, "id", None), - "error": str(exc), - }, - ) - last_response = getattr(last_message, "content", "") or "" - - pending_user_input = False try: - from flocks.server.routes.question import has_pending_questions + from flocks.bus.bus import Bus + from flocks.bus.events import SessionError - pending_user_input = has_pending_questions(ctx.session.id) - except Exception as exc: + await Bus.publish( + SessionError, + {"sessionID": session_id, "error": str(error)}, + ) + except Exception as publish_error: log.warn( - "goal.pending_question_check.error", - {"session_id": ctx.session.id, "error": str(exc)}, + "session.error_event_failed", + {"error": str(publish_error)}, ) - - goal_decision = await GoalManager.evaluate_after_turn( - ctx.session.id, - str(last_response or ""), - pending_user_input=pending_user_input, - provider_id=ctx.provider_id, - model_id=ctx.model_id, + state = AgentRunState[Any]( + session_id=session_id, + agent_name=turn.agent_name, + active_model=RuntimeModel(turn.provider_id, turn.model_id), + model_turn_index=turn.step, + trace_step_offset=turn.trace_step_offset, + current_user_id=processed_user_id, + metadata={"unhandled_runtime_error": True}, ) - if ( - goal_decision.status in {"completed", "blocked", "paused"} - and goal_decision.objective - ): - await cls._publish_runtime_event( - callbacks, - "session.goal.updated", - { - "sessionID": ctx.session.id, - "status": goal_decision.status, - "objective": goal_decision.objective, - "reason": goal_decision.reason, - }, - ) - if goal_decision.should_continue and goal_decision.continuation_prompt: - goal_user = await Message.create( - session_id=ctx.session.id, - role=MessageRole.USER, - content=goal_decision.continuation_prompt, - agent=( - last_user.agent - if hasattr(last_user, "agent") - else ctx.agent_name - ), - model=( - last_user.model - if hasattr(last_user, "model") - else { - "providerID": ctx.provider_id, - "modelID": ctx.model_id, - } - ), - provider=( - last_user.provider - if hasattr(last_user, "provider") - else ctx.provider_id - ), - synthetic=True, - part_metadata={ - "goalContinuation": True, - "goalVerdict": goal_decision.verdict, - "goalReason": goal_decision.reason, - }, - ) - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="continued", - continue_reason="goal", - queued_message_detected=False, - ) - await cls._publish_runtime_event( - callbacks, - "turn.continued", - { - **turn_state.model_dump(by_alias=True), - "goalMessageID": goal_user.id, - "goalVerdict": goal_decision.verdict, - }, - ) - log.info( - "loop.continuing_for_goal", - { - "session_id": ctx.session.id, - "goal_message_id": goal_user.id, - "reason": goal_decision.reason, - }, - ) - return ContinuationDecision( - messages=(goal_user,), - reason="goal", - ) - - if ( - not ctx.should_abort() - and getattr(last_message, "finish", None) == "stop" - and await cls._run_turn_finish_hook( - ctx, - callbacks, - last_user, - last_message, - ) - ): - if ctx.session_ctx: - post_hook_messages = await ctx.session_ctx.get_messages() - else: - post_hook_messages = await Message.list(ctx.session.id) - existing_ids = {message.id for message in state.messages} - new_messages = tuple( - message - for message in post_hook_messages - if message.id not in existing_ids - ) - return ContinuationDecision( - messages=new_messages, - reason="turn_finish_hook", - ) - - stop_reason = getattr(last_message, "finish", None) or "stop" - await cls._publish_turn_stopped( - callbacks, - ctx.session.id, - step=ctx.step, - stop_reason=stop_reason, + return AgentRunOutcome( + status=AgentRunStatus.FATAL_FAILURE, + state=state, + error=str(error), ) - return ContinuationDecision() @classmethod - async def _prepare_model_turn( + async def _release_session( cls, - ctx: LoopContext, - callbacks: LoopCallbacks, - state: AgentRunState[MessageInfo], - ) -> ModelTurnPreparation[MessageInfo]: - """Prepare one immutable model-turn snapshot from session state.""" - SessionStatus.set(ctx.session.id, SessionStatusBusy()) - ctx.step += 1 - state.model_turn_index = ctx.step - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="started", - queued_message_detected=False, - ) - await cls._publish_runtime_event( - callbacks, - "turn.started", - turn_state.model_dump(by_alias=True), - ) - log.info( - "loop.step", - {"session_id": ctx.session.id, "step": ctx.step}, - ) - if callbacks.on_step_start: - await callbacks.on_step_start(ctx.step) - - messages_started_at = asyncio.get_running_loop().time() - if ctx.session_ctx: - messages = await ctx.session_ctx.get_messages() - else: - messages = await Message.list(ctx.session.id) - log.debug( - "loop.messages_loaded", - { - "session_id": ctx.session.id, - "step": ctx.step, - "message_count": len(messages), - "duration_ms": int( - (asyncio.get_running_loop().time() - messages_started_at) - * 1000 - ), - }, - ) - if not messages: - log.info("loop.no_messages", {"session_id": ctx.session.id}) - await cls._publish_turn_stopped( - callbacks, - ctx.session.id, - step=ctx.step, - stop_reason="no_messages", - ) - return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) - - last_user: Optional[MessageInfo] = None - last_assistant: Optional[MessageInfo] = None - last_finished: Optional[MessageInfo] = None - tasks: List[tuple[str, Any]] = [] - scan_started_at = asyncio.get_running_loop().time() - for message in reversed(messages): - if last_user is None and message.role == MessageRole.USER: - last_user = message - if last_assistant is None and message.role == MessageRole.ASSISTANT: - last_assistant = message - if ( - last_finished is None - and message.role == MessageRole.ASSISTANT - and getattr(message, "finish", None) - ): - last_finished = message - if last_user is not None and last_finished is not None: - break - if last_finished is None: - for part in await Message.parts(message.id, ctx.session.id): - if part.type == "compaction": - tasks.append(("compaction", part)) - elif part.type == "subtask": - tasks.append(("subtask", part)) - log.debug( - "loop.message_scan_complete", - { - "session_id": ctx.session.id, - "step": ctx.step, - "task_count": len(tasks), - "duration_ms": int( - (asyncio.get_running_loop().time() - scan_started_at) * 1000 - ), - }, - ) - - if last_user is None: - log.info( - "loop.no_user_message", - { - "session_id": ctx.session.id, - "message_count": len(messages), - "roles": [ - str(getattr(message, "role", "")) - for message in messages[-5:] - ], - }, - ) - await cls._publish_turn_stopped( - callbacks, - ctx.session.id, - step=ctx.step, - stop_reason="no_user_message", - ) - return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) - - last_assistant_parts = ( - await Message.parts(last_assistant.id, ctx.session.id) - if last_assistant - else [] - ) - if cls._should_exit(last_user, last_assistant, last_assistant_parts): - log.info( - "loop.exit_condition", - { - "session_id": ctx.session.id, - "last_user_id": last_user.id, - "last_assistant_id": ( - last_assistant.id if last_assistant else None - ), - "finish": last_assistant.finish if last_assistant else None, - "has_tool_parts": any( - getattr(part, "type", None) == "tool" - for part in last_assistant_parts - ), - }, - ) - return ModelTurnPreparation( - status=TurnPreparationStatus.COMPLETE, - last_message=last_assistant, - ) - - if await cls._prepare_auto_turn(ctx, last_user): - ctx.turn_additional_context = None - ctx.stop_hook_active = False - await cls._run_user_prompt_submit_hook(ctx, last_user) - - state.current_user_id = last_user.id - state.metadata["last_user"] = last_user - await cls._prepare_memory(ctx) - cls._schedule_title_generation(ctx, callbacks, last_user, messages) - - if tasks: - task_preparation = await cls._prepare_pending_task( - ctx, - callbacks, - messages, - last_user, - tasks.pop(), - ) - if task_preparation is not None: - return task_preparation - - context_preparation = await cls._prepare_context_window( - ctx, - callbacks, - messages, - last_user, - last_finished, - ) - if context_preparation is not None: - return context_preparation - - active_model = RuntimeModel(ctx.provider_id, ctx.model_id) - state.active_model = active_model - state.messages = list(messages) - return ModelTurnPreparation( - status=TurnPreparationStatus.READY, - snapshot=ModelTurnSnapshot( - session_id=ctx.session.id, - agent_name=ctx.agent_name, - active_model=active_model, - model_turn_index=ctx.step, - trace_step=ctx.trace_step, - messages=tuple(messages), - last_user=last_user, - ), - ) - - @staticmethod - async def _prepare_memory(ctx: LoopContext) -> None: - """Load memory once before the first model turn.""" - if ( - ctx.step != 1 - or not ctx.session.memory_enabled - or ctx.memory_bootstrap_data is not None - ): - return - try: - from flocks.memory.bootstrap import MemoryBootstrap - - ctx.memory_bootstrap_data = await MemoryBootstrap( - project_id=ctx.session.project_id, - ).bootstrap(load_daily=False) - log.info( - "loop.memory_bootstrap_done", - { - "session_id": ctx.session.id, - "has_main": ( - ctx.memory_bootstrap_data.get("main_memory") is not None - ), - }, - ) - except Exception as exc: - log.error("loop.memory_bootstrap_error", {"error": str(exc)}) - - @staticmethod - def _schedule_title_generation( - ctx: LoopContext, + lease: _SessionLease, callbacks: LoopCallbacks, - last_user: MessageInfo, - messages: List[MessageInfo], ) -> None: - """Start optimistic first-turn title generation without blocking.""" - if ctx.step != 1 or ctx.auto_failover: - return - try: - from flocks.session.lifecycle.title import SessionTitle - - user_model = getattr(last_user, "model", None) - if isinstance(user_model, dict): - title_model_id = user_model.get("modelID", ctx.model_id) - title_provider_id = user_model.get( - "providerID", - ctx.provider_id, - ) - else: - title_model_id = ctx.model_id - title_provider_id = ctx.provider_id - fire_and_forget( - SessionTitle.ensure_title( - session_id=ctx.session.id, - model_id=title_model_id, - provider_id=title_provider_id, - messages=messages, - event_publish_callback=callbacks.event_publish_callback, - ), - label="title_generation", - name=f"title:{ctx.session.id}", - ) - except Exception as exc: - log.error("loop.title_generation.error", {"error": str(exc)}) + async with Session.lifecycle_lock(lease.session_id): + cls._finalize_release_state_locked(lease) + await cls._publish_released(lease.turn, callbacks) @classmethod - async def _prepare_pending_task( + async def _settle_or_continue( cls, - ctx: LoopContext, + lease: _SessionLease, callbacks: LoopCallbacks, - messages: List[MessageInfo], - last_user: MessageInfo, - task: tuple[str, Any], - ) -> Optional[ModelTurnPreparation[MessageInfo]]: - """Finish persisted subtask or compaction work before the model turn.""" - task_type, task_part = task - if task_type == "subtask": - log.info( - "loop.subtask_detected", - {"session_id": ctx.session.id, "step": ctx.step}, - ) - await cls._execute_subtask(ctx, last_user, task_part) - return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) - - log.info( - "loop.compaction_pending", - { - "session_id": ctx.session.id, - "step": ctx.step, - "auto": getattr(task_part, "auto", False), - }, - ) - if callbacks.on_compaction: - await callbacks.on_compaction() - - publish = callbacks.event_publish_callback - progress_callback = None - if publish is not None: - - async def progress_callback(stage: str, data: dict) -> None: - await publish( - "session.compaction_progress", + processed_user_id: Optional[str], + ) -> bool: + """Atomically keep ownership for late input or settle idle.""" + async with Session.lifecycle_lock(lease.session_id): + if await lease.turn.has_late_input(processed_user_id): + log.info( + "session.continuing_for_late_input", { - "sessionID": ctx.session.id, - "stage": stage, - "data": data, + "session_id": lease.session_id, + "processed_user_id": processed_user_id, }, ) + return True + cls._finalize_release_state_locked(lease) - try: - compaction_result = await run_compaction( - ctx.session.id, - parent_message_id=last_user.id, - messages=messages, - provider_id=ctx.provider_id, - model_id=ctx.model_id, - auto=getattr(task_part, "auto", False), - event_publish_callback=publish, - status_after="busy", - policy=cls._build_compaction_policy(ctx), - progress_callback=progress_callback, - ) - if compaction_result == "stop": - log.error( - "loop.compaction_failed", - {"session_id": ctx.session.id}, - ) - if callbacks.on_error: - await callbacks.on_error("Compaction failed") - return ModelTurnPreparation( - status=TurnPreparationStatus.COMPLETE, - ) - if compaction_result == "skipped": - log.info( - "loop.manual_compaction_skipped", - {"session_id": ctx.session.id, "step": ctx.step}, - ) - return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) - except Exception as exc: - log.error("loop.compaction_error", {"error": str(exc)}) - if callbacks.on_error: - await callbacks.on_error(f"Compaction error: {exc}") - return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) + await cls._publish_released(lease.turn, callbacks) + return False @classmethod - async def _prepare_context_window( - cls, - ctx: LoopContext, - callbacks: LoopCallbacks, - messages: List[MessageInfo], - last_user: MessageInfo, - last_finished: Optional[MessageInfo], - ) -> Optional[ModelTurnPreparation[MessageInfo]]: - """Recover a near-overflow context before the next model turn.""" - if last_finished is None or getattr(last_finished, "summary", False): - return None - - model_context, model_output, model_input = Provider.resolve_model_info( - ctx.provider_id, - ctx.model_id, - ) - if model_context <= 0: - return None - - policy = CompactionPolicy.from_model( - context_window=model_context, - max_output_tokens=model_output or 4096, - max_input_tokens=model_input, - ) - tokens = cls._normalise_token_usage(last_finished) - input_tokens = tokens.get("input", 0) - cache = tokens.get("cache") or {} - cache_read = cache.get("read", 0) if isinstance(cache, dict) else 0 - output_tokens = tokens.get("output", 0) - reported_total = input_tokens + cache_read + output_tokens - if reported_total > 0: - ctx.last_observed_prompt_tokens = reported_total - log.info( - "loop.tokens_decision", - { - "session_id": ctx.session.id, - "source": "observed", - "effective_tokens": input_tokens + cache_read, - "overflow_threshold": policy.overflow_threshold, - }, - ) - else: - estimated_tokens = await SessionPrompt.estimate_full_context_tokens( - ctx.session.id, - messages, - policy=policy, - ) - tokens = { - "input": estimated_tokens, - "output": 0, - "cache": {"read": 0, "write": 0}, - } - log.info( - "loop.tokens_decision", - { - "session_id": ctx.session.id, - "source": "estimated", - "effective_tokens": estimated_tokens, - "message_count": len(messages), - "overflow_threshold": policy.overflow_threshold, - }, - ) - - try: - cache = tokens.get("cache") or {} - current_input_tokens = tokens.get("input", 0) + ( - cache.get("read", 0) if isinstance(cache, dict) else 0 - ) - recent_compaction = cls._has_recent_compaction_cooldown(ctx) - near_overflow = current_input_tokens >= policy.preemptive_threshold - if near_overflow and ctx.last_cleanup_step != ctx.step: - cleanup_result = await cls._prepare_tool_result_cleanup( - ctx, - callbacks, - model_context, - policy, - current_input_tokens, - recent_compaction, - ) - if cleanup_result is not None: - return cleanup_result - - is_overflow = await SessionCompaction.is_overflow( - tokens=tokens, - model_context=model_context, - policy=policy, - ) - if not is_overflow: - return None - - log.info( - "loop.context_overflow_detected", - { - "session_id": ctx.session.id, - "step": ctx.step, - "tokens": tokens, - "tier": policy.tier.value, - "overflow_compaction_attempts": ( - ctx.overflow_compaction_attempts - ), - }, - ) - if ( - ctx.overflow_compaction_attempts - >= MAX_OVERFLOW_COMPACTION_ATTEMPTS - ): - await cls._report_compaction_exhausted( - ctx, - callbacks, - tokens, - ) - return ModelTurnPreparation( - status=TurnPreparationStatus.COMPLETE, - ) - - if not ctx.tool_result_truncation_attempted: - ctx.tool_result_truncation_attempted = True - try: - truncation_count = ( - await SessionCompaction.truncate_oversized_tool_outputs( - ctx.session.id, - context_window_tokens=model_context, - ) - ) - if truncation_count > 0: - log.info( - "loop.oversized_tool_truncated", - { - "session_id": ctx.session.id, - "truncated": truncation_count, - }, - ) - estimated_tokens = ( - await SessionPrompt.estimate_full_context_tokens( - ctx.session.id, - messages, - policy=policy, - ) - ) - still_overflow = await SessionCompaction.is_overflow( - tokens={ - "input": estimated_tokens, - "output": 0, - "cache": {"read": 0, "write": 0}, - }, - model_context=model_context, - policy=policy, - ) - if not still_overflow: - log.info( - "loop.overflow_resolved_by_truncation", - {"session_id": ctx.session.id}, - ) - return ModelTurnPreparation( - status=TurnPreparationStatus.CONTINUE, - ) - except Exception as exc: - log.warn( - "loop.oversized_truncation_error", - {"session_id": ctx.session.id, "error": str(exc)}, - ) - - return await cls._prepare_full_compaction( - ctx, - callbacks, - messages, - last_user, - policy, - ) - except Exception as exc: - log.error( - "loop.compaction_overflow_check_error", - {"error": str(exc)}, - ) - return None + def _finalize_release_state_locked(cls, lease: _SessionLease) -> None: + clear_turn_state(lease.session_id) + SessionStatus.set(lease.session_id, SessionStatusIdle()) + cls._leases.release(lease) @staticmethod - def _normalise_token_usage(message: MessageInfo) -> Dict[str, Any]: - """Normalise provider token usage into the legacy mapping shape.""" - raw_tokens = getattr(message, "tokens", None) - if not raw_tokens: - return {} - if isinstance(raw_tokens, dict): - return raw_tokens - if hasattr(raw_tokens, "model_dump"): - return raw_tokens.model_dump() - if hasattr(raw_tokens, "__dict__"): - return vars(raw_tokens) - return {} - - @classmethod - async def _prepare_tool_result_cleanup( - cls, - ctx: LoopContext, + async def _publish_released( + turn: SessionTurn, callbacks: LoopCallbacks, - model_context: int, - policy: CompactionPolicy, - current_input_tokens: int, - recent_compaction: bool, - ) -> Optional[ModelTurnPreparation[MessageInfo]]: - """Apply the cheap tool-result cleanup before full compaction.""" + ) -> None: + session_id = turn.session.id + await SessionEventSink.session_status(callbacks, session_id, "idle") try: - truncation_count = ( - await SessionCompaction.truncate_oversized_tool_outputs( - ctx.session.id, - context_window_tokens=model_context, - ) - ) - ctx.last_cleanup_step = ctx.step - if truncation_count <= 0: - return None - - set_context_state( - ctx.session.id, - tool_results_compacted=True, - last_compaction_step=ctx.last_compaction_step, - last_compaction_reason="pre_compact_cleanup", - ) - await cls._publish_runtime_event( - callbacks, - "context.compacted", - { - "sessionID": ctx.session.id, - "step": ctx.step, - "reason": "pre_compact_cleanup", - "truncatedToolResults": truncation_count, - "cooldownActive": recent_compaction, - }, - ) - log.info( - "loop.pre_compact_cleanup_applied", - { - "session_id": ctx.session.id, - "step": ctx.step, - "truncated": truncation_count, - "preemptive_threshold": policy.preemptive_threshold, - "input_tokens": current_input_tokens, - "cooldown_active": recent_compaction, - }, - ) - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="continued", - continue_reason="pre_compact_cleanup", - queued_message_detected=False, - ) - await cls._publish_runtime_event( - callbacks, - "turn.continued", - turn_state.model_dump(by_alias=True), - ) - return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) + await Session.touch(turn.session.project_id, session_id) except Exception as exc: log.warn( - "loop.pre_compact_cleanup_error", - {"session_id": ctx.session.id, "error": str(exc)}, - ) - return None - - @classmethod - async def _report_compaction_exhausted( - cls, - ctx: LoopContext, - callbacks: LoopCallbacks, - tokens: Dict[str, Any], - ) -> None: - """Surface whether exhaustion came from context or provider health.""" - history = _get_compaction_history(ctx.session.id) - provider_error = history.summary_last_error - in_cooldown = ( - history.summary_cooldown_until > 0 - and history.summary_cooldown_until > time.monotonic() - ) - cooldown_seconds = max( - 0, - round(history.summary_cooldown_until - time.monotonic()), - ) - if in_cooldown or provider_error: - notice = ( - "摘要模型暂时不可用，上下文压缩跳过了本轮压缩。" - + ( - f"冷却剩余约 {cooldown_seconds} 秒，" - if in_cooldown - else "" - ) - + "建议稍后继续，或切换到其他模型重试。" - ) - error = ( - "Compaction skipped: summary provider unavailable " - f"({provider_error or 'cooldown active'})." - + ( - f" Cooldown expires in ~{cooldown_seconds}s." - if in_cooldown - else "" - ) - + " Wait for the provider to recover or switch models." - ) - else: - notice = ( - "当前任务上下文过重，已经多次 compact 仍接近上限。" - "建议收敛工具输出、缩小搜索范围，或开启新会话。" - ) - error = ( - "Context overflow: prompt too large for the model after " - f"{ctx.overflow_compaction_attempts} compaction attempts. " - "Try starting a new session or use a larger-context model." - ) - - await cls._publish_session_notice( - callbacks, - ctx.session.id, - level="warning", - message=notice, - details={ - "attempts": ctx.overflow_compaction_attempts, - "maxAttempts": MAX_OVERFLOW_COMPACTION_ATTEMPTS, - "tokens": tokens, - "providerError": provider_error or None, - "cooldownRemainingSeconds": ( - cooldown_seconds if in_cooldown else 0 - ), - }, - ) - log.error( - "loop.overflow_compaction_exhausted", - { - "session_id": ctx.session.id, - "attempts": ctx.overflow_compaction_attempts, - "max": MAX_OVERFLOW_COMPACTION_ATTEMPTS, - "tokens": tokens, - "in_cooldown": in_cooldown, - "provider_error": provider_error or None, - }, - ) - if callbacks.on_error: - await callbacks.on_error(error) - - @classmethod - async def _prepare_full_compaction( - cls, - ctx: LoopContext, - callbacks: LoopCallbacks, - messages: List[MessageInfo], - last_user: MessageInfo, - policy: CompactionPolicy, - ) -> ModelTurnPreparation[MessageInfo]: - """Run full compaction and request preparation to reload the session.""" - ctx.overflow_compaction_attempts += 1 - if ctx.overflow_compaction_attempts >= 2: - await cls._publish_session_notice( - callbacks, - ctx.session.id, - level="info", - message=( - "本轮上下文持续接近模型上限，系统将优先尝试压缩历史工具输出。" - ), - details={ - "attempt": ctx.overflow_compaction_attempts, - "threshold": policy.overflow_threshold, - "buffer": policy.overflow_buffer, - }, - ) - log.warn( - "loop.overflow_compaction_attempt", - { - "session_id": ctx.session.id, - "attempt": ctx.overflow_compaction_attempts, - "max": MAX_OVERFLOW_COMPACTION_ATTEMPTS, - }, - ) - if callbacks.on_compaction: - await callbacks.on_compaction() - await SessionCompaction.prune(ctx.session.id, policy=policy) - - publish = callbacks.event_publish_callback - progress_callback = None - if publish is not None: - - async def progress_callback(stage: str, data: dict) -> None: - await publish( - "session.compaction_progress", - { - "sessionID": ctx.session.id, - "stage": stage, - "data": data, - }, - ) - - result = await run_compaction( - ctx.session.id, - parent_message_id=last_user.id, - messages=messages, - provider_id=ctx.provider_id, - model_id=ctx.model_id, - auto=True, - event_publish_callback=publish, - status_after="busy", - policy=policy, - progress_callback=progress_callback, - ) - if result == "stop": - log.error( - "loop.compaction_failed", - {"session_id": ctx.session.id}, + "session.touch_failed", + {"session_id": session_id, "error": str(exc)}, ) - if callbacks.on_error: - await callbacks.on_error("Compaction failed") - return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) - if result == "skipped": - log.info( - "loop.compaction_skipped", - {"session_id": ctx.session.id, "step": ctx.step}, - ) - else: - ctx.last_compaction_step = ctx.step - set_context_state( - ctx.session.id, - compaction_performed=True, - last_compaction_step=ctx.step, - last_compaction_reason="full_compaction", - ) - await cls._publish_runtime_event( - callbacks, - "context.compacted", - { - "sessionID": ctx.session.id, - "step": ctx.step, - "reason": "full_compaction", - "attempt": ctx.overflow_compaction_attempts, - "cooldownUntilStep": ( - ctx.step + POST_COMPACTION_COOLDOWN_STEPS - ), - }, - ) - return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) - @classmethod - async def _run_loop( - cls, - ctx: LoopContext, - callbacks: LoopCallbacks, - ) -> LoopResult: - """Run the host-neutral AgentLoop against session-owned adapters.""" - from flocks.session.runtime_services import ( - SessionRuntimeServices, - SessionStepCancelled, - ) - - state = AgentRunState[MessageInfo]( - session_id=ctx.session.id, - agent_name=ctx.agent_name, - active_model=RuntimeModel(ctx.provider_id, ctx.model_id), - model_turn_index=ctx.step, - trace_step_offset=ctx.trace_step_offset, - current_user_id=ctx.turn_user_id, - ) - services = SessionRuntimeServices(ctx, callbacks, cls) - step_engine = cls._create_host_step_engine(ctx, callbacks) try: - outcome = await AgentLoop( - step_engine, - services, - abort_requested=ctx.should_abort, - ).run(state) - except SessionStepCancelled: - log.info( - "loop.step_cancelled", - {"session_id": ctx.session.id, "step": ctx.step}, - ) - return LoopResult( - action="stop", - provider_id=ctx.provider_id, - model_id=ctx.model_id, - metadata={ - "steps": ctx.step, - "session_id": ctx.session.id, - "last_compaction_step": ctx.last_compaction_step, - "aborted": True, - }, - ) - - loop_error = ( - outcome.error - if outcome.status - in { - AgentRunStatus.RETRYABLE_FAILURE, - AgentRunStatus.FATAL_FAILURE, - AgentRunStatus.CONTEXT_OVERFLOW, - } - else None - ) - return LoopResult( - action="error" if ctx.auto_failover and loop_error else "stop", - last_message=outcome.last_message, - error=loop_error if ctx.auto_failover else None, - provider_id=ctx.provider_id, - model_id=ctx.model_id, - metadata={ - "steps": ctx.step, - "session_id": ctx.session.id, - "last_compaction_step": ctx.last_compaction_step, - **( - {"aborted": True} - if outcome.status == AgentRunStatus.ABORTED - else {} - ), - }, - ) - - @classmethod - def _build_compaction_policy(cls, ctx: LoopContext) -> CompactionPolicy: - """ - Construct a CompactionPolicy from the current model's info. - - Falls back to ``CompactionPolicy.default()`` when the model info - cannot be resolved (e.g. unknown provider or missing context_window). - """ - return build_compaction_policy(ctx.provider_id, ctx.model_id) - - @classmethod - def _should_exit( - cls, - last_user: MessageInfo, - last_assistant: Optional[MessageInfo], - last_assistant_parts: Optional[List[Any]] = None, - ) -> bool: - """ - Check if loop should exit - - Ported from original exit logic: - - Exit if assistant has responded with finish != tool-calls - - Exit if assistant message is after user message - """ - if not last_assistant: - return False + from flocks.bus.bus import Bus + from flocks.bus.events import SessionIdle - if any( - getattr(part, "type", None) == "tool" - for part in (last_assistant_parts or []) - ): - return False - - # Check finish reason - if last_assistant.finish: - if last_assistant.finish not in ("tool-calls", "unknown", "summary"): - # Assistant finished with stop/error/etc - if last_user.id < last_assistant.id: - # Assistant responded after user - return True - - return False - - @classmethod - async def _check_reminders( - cls, - ctx: LoopContext, - messages: List[MessageInfo], - callbacks: LoopCallbacks, - ) -> None: - """ - Check and inject reminders (P1 feature) - - Reminders are system messages injected periodically to: - - Remind agent of task goals - - Prevent drift from original intent - - Nudge towards completion - """ - from flocks.session.features.reminders import SessionReminders, ReminderContext, ReminderConfig - - # Calculate elapsed time - if messages: - first_msg = messages[0] - if hasattr(first_msg, 'time') and hasattr(first_msg.time, 'created'): - first_time = first_msg.time.created - current_time = int(datetime.now().timestamp() * 1000) - elapsed_ms = current_time - first_time - else: - elapsed_ms = 0 - else: - elapsed_ms = 0 - - # Extract original task - original_task = await SessionReminders.extract_original_task(messages) - - # Create reminder context - reminder_ctx = ReminderContext( - session_id=ctx.session.id, - step_count=ctx.step, - message_count=len(messages), - elapsed_ms=elapsed_ms, - original_task=original_task, - ) - - # Check if reminder should be injected - if SessionReminders.should_remind(ctx.session.id, reminder_ctx): - # Create and inject reminder - reminder_msg = await SessionReminders.create_reminder( - ctx.session.id, - reminder_ctx, - ) - - if reminder_msg and callbacks.on_reminder: - await callbacks.on_reminder(await Message.get_text_content(reminder_msg)) - - @classmethod - async def _execute_subtask( - cls, - ctx: LoopContext, - last_user: MessageInfo, - task_part: Any, - ) -> None: - """ - Execute subtask (matching TUI lines 316-481) - - 完全匹配 TUI 的 subtask 执行流程: - 1. 创建 assistant message - 2. 创建 tool part (Task tool) - 3. 执行 Task tool - 4. 更新 part 状态 - 5. 创建 synthetic user message - """ - from flocks.tool.registry import ToolRegistry - from flocks.agent.registry import Agent - - # Extract subtask information from part - agent_name = getattr(task_part, 'agent', 'hephaestus') - prompt = getattr(task_part, 'prompt', '') - description = getattr(task_part, 'description', '') - command = getattr(task_part, 'command', None) - model_info = getattr(task_part, 'model', None) - - # Get agent - agent = await Agent.get(agent_name) or await Agent.get("rex") - - # Determine model - if model_info: - provider_id = model_info.get('providerID', ctx.provider_id) - model_id = model_info.get('modelID', ctx.model_id) - else: - provider_id = ctx.provider_id - model_id = ctx.model_id - - # Create assistant message for subtask - assistant_msg = await Message.create( - session_id=ctx.session.id, - role=MessageRole.ASSISTANT, - content="", - agent=agent_name, - model=model_id, - provider=provider_id, - parent_id=last_user.id, - ) - - # Create tool part for Task - tool_call_id = Identifier.create("call") - from flocks.session.message import ToolPart, ToolStateRunning - - tool_part = ToolPart( - id=Identifier.ascending("part"), - sessionID=ctx.session.id, - messageID=assistant_msg.id, - type="tool", - callID=tool_call_id, - tool="task", - state=ToolStateRunning( - status="running", - input={ - "prompt": prompt, - "description": description, - "subagent_type": agent_name, - "command": command, - }, - time={"start": int(datetime.now().timestamp() * 1000)}, - ), - ) - - # Add part to message - await Message.add_part(ctx.session.id, assistant_msg.id, tool_part) - - # Get Task tool - task_tool = ToolRegistry.get("task") - if not task_tool: - log.error("loop.subtask.task_tool_not_found", {"session_id": ctx.session.id}) - return - - # Execute Task tool - task_args = { - "prompt": prompt, - "description": description, - "subagent_type": agent_name, - "command": command, - } - - # Create tool context - from flocks.tool.registry import ToolContext - - tool_ctx = ToolContext( - session_id=ctx.session.id, - message_id=assistant_msg.id, - agent=agent_name, - abort_event=ctx.abort_event, - ) - - execution_error: Optional[Exception] = None - result = None - - try: - result = await task_tool.execute(tool_ctx, **task_args) - except Exception as e: - execution_error = e - log.error("loop.subtask.execution_failed", { - "error": str(e), - "agent": agent_name, - "description": description, - }) - - # Update message finish - await Message.update(ctx.session.id, assistant_msg.id, finish="tool-calls") - - # Update tool part status - from flocks.session.message import ToolStateCompleted, ToolStateError - - if result: - # Create completed state - completed_state = ToolStateCompleted( - status="completed", - input={ - "prompt": prompt, - "description": description, - "subagent_type": agent_name, - "command": command, - }, - output=result.output if hasattr(result, 'output') else str(result), - title=result.title if hasattr(result, 'title') else None, - metadata=result.metadata if hasattr(result, 'metadata') else {}, - time={ - "start": tool_part.state.time.get("start"), - "end": int(datetime.now().timestamp() * 1000), - }, - ) - await Message.update_part( - session_id=ctx.session.id, - message_id=assistant_msg.id, - part_id=tool_part.id, - state=completed_state, - ) - else: - # Create error state - error_msg = str(execution_error) if execution_error else "Tool execution failed" - error_state = ToolStateError( - status="error", - error=f"Tool execution failed: {error_msg}", - time={ - "start": tool_part.state.time.get("start"), - "end": int(datetime.now().timestamp() * 1000), - }, - metadata={}, - input={ - "prompt": prompt, - "description": description, - "subagent_type": agent_name, - "command": command, - }, - ) - await Message.update_part( - session_id=ctx.session.id, - message_id=assistant_msg.id, - part_id=tool_part.id, - state=error_state, - ) - - # Create synthetic user message (matching TUI lines 457-478) - # This prevents reasoning models from erroring due to missing user messages - synthetic_user_msg = await Message.create( - session_id=ctx.session.id, - role=MessageRole.USER, - content="Summarize the task tool output above and continue with your task.", - agent=last_user.agent if hasattr(last_user, 'agent') else agent_name, - model=last_user.model if hasattr(last_user, 'model') else model_id, - provider=last_user.provider if hasattr(last_user, 'provider') else provider_id, - synthetic=True, - ) - - log.info("loop.subtask.completed", { - "session_id": ctx.session.id, - "agent": agent_name, - "success": result is not None, - }) - + await Bus.publish(SessionIdle, {"sessionID": session_id}) + except Exception as exc: + log.warn("session.idle_event_failed", {"error": str(exc)}) -# Export __all__ = [ "SessionLoop", "LoopContext", diff --git a/flocks/session/step_engine.py b/flocks/session/step_engine.py deleted file mode 100644 index c8d870184..000000000 --- a/flocks/session/step_engine.py +++ /dev/null @@ -1,35 +0,0 @@ -"""Session adapter for the agent runtime's step-engine port.""" - -from __future__ import annotations - -from flocks.agent.runtime.contracts import ModelTurnSnapshot, StepResult -from flocks.session.message import MessageInfo -from flocks.session.runner import SessionRunner - - -class SessionStepEngine: - """Run the existing session runner behind the StepEngine contract.""" - - def __init__(self, runner: SessionRunner): - self._runner = runner - - @property - def session_start_fired(self) -> bool: - """Return whether this engine fired the session-start hook.""" - return self._runner._session_start_fired - - @property - def attempt_effects(self): - """Return effects recorded by the latest provider attempt.""" - return self._runner._attempt_state - - async def run( - self, - snapshot: ModelTurnSnapshot[MessageInfo], - ) -> StepResult: - """Delegate one immutable model-turn snapshot to SessionRunner.""" - self._runner._step = snapshot.trace_step - return await self._runner._process_step( - list(snapshot.messages), - snapshot.last_user, - ) diff --git a/flocks/session/utils/file_extractor.py b/flocks/session/utils/file_extractor.py index 1c77aa2c2..f5a887a2c 100644 --- a/flocks/session/utils/file_extractor.py +++ b/flocks/session/utils/file_extractor.py @@ -1,7 +1,7 @@ """ File content extraction utilities for session message processing. -Extracted from SessionRunner to keep file-handling concerns separate +Extracted from StepEngine to keep file-handling concerns separate. from session execution logic. """ diff --git a/flocks/task/background.py b/flocks/task/background.py index e5a62d0ef..72781eec4 100644 --- a/flocks/task/background.py +++ b/flocks/task/background.py @@ -237,21 +237,27 @@ async def _inject_parent_completion(self, task: BackgroundTask) -> None: "" ) try: - await Message.create( - session_id=task.parent_session_id, - role=MessageRole.USER, - content=content, - agent=task.parent_agent or "rex", - model=task.parent_model, - synthetic=True, - part_metadata={ - "kind": "background_task_result", - "task_id": task.id, - "session_id": task.session_id, - "status": state, - }, + async def _persist_completion() -> None: + await Message.create( + session_id=task.parent_session_id, + role=MessageRole.USER, + content=content, + agent=task.parent_agent or "rex", + model=task.parent_model, + synthetic=True, + part_metadata={ + "kind": "background_task_result", + "task_id": task.id, + "session_id": task.session_id, + "status": state, + }, + ) + await self._update_parent_tool_part(task) + + await Session.run_active_write( + task.parent_session_id, + _persist_completion, ) - await self._update_parent_tool_part(task) task.completion_injected = True self._schedule_parent_resume(task) except Exception as exc: @@ -265,14 +271,6 @@ def _schedule_parent_resume(self, task: BackgroundTask) -> None: """Kick the parent session so Rex consumes injected background results.""" if not task.parent_session_id: return - if task.status not in ("completed", "error"): - return - if SessionLoop.is_running(task.parent_session_id): - log.info("background.parent_resume.already_running", { - "task_id": task.id, - "parent_session_id": task.parent_session_id, - }) - return async def _run_parent() -> None: try: @@ -423,7 +421,6 @@ def cancel_by_parent_session_id(self, parent_session_id: str) -> int: def _build_activity_callbacks(self, task: BackgroundTask): """构建带活跃时间更新的 LoopCallbacks，用于不活跃超时检测。""" from flocks.session.session_loop import LoopCallbacks - from flocks.session.runner import RunnerCallbacks from flocks.server.routes.event import publish_event def _touch() -> None: @@ -435,10 +432,9 @@ async def _on_step_start(_step: int) -> None: async def _on_text_delta(_text: str) -> None: _touch() - runner_cbs = RunnerCallbacks(on_text_delta=_on_text_delta) return LoopCallbacks( on_step_start=_on_step_start, - runner_callbacks=runner_cbs, + on_text_delta=_on_text_delta, event_publish_callback=publish_event, ) diff --git a/tests/agent/runtime/test_agent_loop.py b/tests/agent/runtime/test_agent_loop.py deleted file mode 100644 index 690f5b923..000000000 --- a/tests/agent/runtime/test_agent_loop.py +++ /dev/null @@ -1,384 +0,0 @@ -"""Tests for the host-neutral agent control loop.""" - -from __future__ import annotations - -from collections import deque -from collections.abc import Callable -from dataclasses import dataclass - -import pytest - -from flocks.agent.runtime import ( - AgentLoop, - AgentRunState, - AgentRunStatus, - AttemptEffects, - ContinuationDecision, - ModelTurnBoundary, - ModelTurnPreparation, - ModelTurnSnapshot, - QueuedInputBatch, - RuntimeModel, - StepFailure, - StepResult, - TurnPreparationStatus, -) - - -@dataclass(frozen=True) -class Message: - """Minimal persisted-message stand-in with stable identity.""" - - id: str - content: str - - -PreparationFactory = Callable[ - [AgentRunState[Message]], - ModelTurnPreparation[Message], -] - - -class FakeStepEngine: - """Return deterministic step results and record immutable inputs.""" - - def __init__(self, results: list[StepResult]): - self._results = deque(results) - self.snapshots: list[ModelTurnSnapshot[Message]] = [] - - async def run(self, snapshot: ModelTurnSnapshot[Message]) -> StepResult: - self.snapshots.append(snapshot) - return self._results.popleft() - - -class FakeRuntimeServices: - """Expose scripted host-boundary decisions to the agent loop.""" - - def __init__( - self, - preparations: list[ModelTurnPreparation[Message] | PreparationFactory], - boundaries: list[ModelTurnBoundary[Message]], - continuations: list[ContinuationDecision[Message]] | None = None, - ): - self._preparations = deque(preparations) - self._boundaries = deque(boundaries) - self._continuations = deque(continuations or []) - self.prepared_messages: list[tuple[Message, ...]] = [] - self.events = [] - - async def prepare_model_turn( - self, - state: AgentRunState[Message], - ) -> ModelTurnPreparation[Message]: - self.prepared_messages.append(tuple(state.messages)) - preparation = self._preparations.popleft() - if callable(preparation): - return preparation(state) - return preparation - - async def complete_model_turn( - self, - state: AgentRunState[Message], - step_result: StepResult, - ) -> ModelTurnBoundary[Message]: - del state, step_result - return self._boundaries.popleft() - - async def resolve_continuation( - self, - state: AgentRunState[Message], - step_result: StepResult, - ) -> ContinuationDecision[Message]: - del state, step_result - if not self._continuations: - return ContinuationDecision() - return self._continuations.popleft() - - async def emit_event(self, event) -> None: - self.events.append(event) - - -def _state(messages: list[Message] | None = None) -> AgentRunState[Message]: - return AgentRunState( - session_id="session-1", - agent_name="rex", - active_model=RuntimeModel("provider-a", "model-a"), - messages=list(messages or [Message("user-1", "hello")]), - ) - - -def _ready( - state: AgentRunState[Message], - *, - turn: int = 0, -) -> ModelTurnPreparation[Message]: - messages = tuple(state.messages) - return ModelTurnPreparation( - status=TurnPreparationStatus.READY, - snapshot=ModelTurnSnapshot( - session_id=state.session_id, - agent_name=state.agent_name, - active_model=state.active_model, - model_turn_index=turn, - trace_step=turn, - messages=messages, - last_user=messages[-1], - ), - ) - - -@pytest.mark.asyncio -async def test_loop_honors_deferred_preparation_then_completes() -> None: - user = Message("user-1", "hello") - assistant = Message("assistant-1", "done") - state = _state([user]) - engine = FakeStepEngine([StepResult(action="stop")]) - services = FakeRuntimeServices( - preparations=[ - ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE), - _ready, - ], - boundaries=[ModelTurnBoundary(messages=(user, assistant), last_message=assistant)], - ) - - outcome = await AgentLoop(engine, services).run(state) - - assert outcome.status == AgentRunStatus.COMPLETED - assert outcome.last_message == assistant - assert len(engine.snapshots) == 1 - assert len(services.prepared_messages) == 2 - - -@pytest.mark.asyncio -async def test_loop_runs_another_turn_after_tool_continue() -> None: - user = Message("user-1", "hello") - tool_result = Message("tool-1", "tool result") - assistant = Message("assistant-1", "done") - state = _state([user]) - engine = FakeStepEngine( - [StepResult(action="continue"), StepResult(action="stop")], - ) - services = FakeRuntimeServices( - preparations=[_ready, lambda current: _ready(current, turn=1)], - boundaries=[ - ModelTurnBoundary(messages=(user, tool_result), last_message=tool_result), - ModelTurnBoundary( - messages=(user, tool_result, assistant), - last_message=assistant, - ), - ], - ) - - outcome = await AgentLoop(engine, services).run(state) - - assert outcome.status == AgentRunStatus.COMPLETED - assert len(engine.snapshots) == 2 - assert engine.snapshots[1].messages == (user, tool_result) - - -@pytest.mark.asyncio -async def test_queued_input_precedes_natural_stop() -> None: - user = Message("user-1", "hello") - assistant = Message("assistant-1", "first answer") - duplicate_assistant = Message("assistant-1", "reloaded answer") - queued_user = Message("user-2", "follow up") - final = Message("assistant-2", "second answer") - state = _state([user]) - engine = FakeStepEngine( - [StepResult(action="stop"), StepResult(action="stop")], - ) - services = FakeRuntimeServices( - preparations=[_ready, lambda current: _ready(current, turn=1)], - boundaries=[ - ModelTurnBoundary( - messages=(user, assistant), - last_message=assistant, - queued_inputs=QueuedInputBatch( - messages=(duplicate_assistant, queued_user), - cursor="input-2", - ), - ), - ModelTurnBoundary( - messages=(user, assistant, queued_user, final), - last_message=final, - ), - ], - ) - - outcome = await AgentLoop(engine, services).run(state) - - assert outcome.status == AgentRunStatus.COMPLETED - assert state.consumed_input_cursor == "input-2" - assert engine.snapshots[1].messages == (user, assistant, queued_user) - - -@pytest.mark.asyncio -async def test_queued_input_is_not_lost_after_final_step_failure() -> None: - user = Message("user-1", "hello") - failed = Message("assistant-1", "provider failed") - queued_user = Message("user-2", "try this instead") - final = Message("assistant-2", "done") - state = _state([user]) - failure = StepFailure( - message="provider failed", - error_data={}, - assistant_message_id=failed.id, - reason="provider_error", - allow_fallback=False, - attempt_state=AttemptEffects(observable_output_started=True), - ) - engine = FakeStepEngine( - [ - StepResult(action="stop", error=failure.message, failure=failure), - StepResult(action="stop"), - ] - ) - services = FakeRuntimeServices( - preparations=[_ready, lambda current: _ready(current, turn=1)], - boundaries=[ - ModelTurnBoundary( - messages=(user, failed, queued_user), - last_message=failed, - queued_inputs=QueuedInputBatch( - messages=(queued_user,), - cursor=queued_user.id, - ), - ), - ModelTurnBoundary( - messages=(user, failed, queued_user, final), - last_message=final, - ), - ], - ) - - outcome = await AgentLoop(engine, services).run(state) - - assert outcome.status == AgentRunStatus.COMPLETED - assert len(engine.snapshots) == 2 - assert engine.snapshots[1].messages == (user, failed, queued_user) - - -@pytest.mark.asyncio -async def test_host_continuation_starts_another_turn() -> None: - user = Message("user-1", "hello") - assistant = Message("assistant-1", "working") - continuation = Message("user-2", "continue goal") - final = Message("assistant-2", "done") - state = _state([user]) - engine = FakeStepEngine( - [StepResult(action="stop"), StepResult(action="stop")], - ) - services = FakeRuntimeServices( - preparations=[_ready, lambda current: _ready(current, turn=1)], - boundaries=[ - ModelTurnBoundary(messages=(user, assistant), last_message=assistant), - ModelTurnBoundary( - messages=(user, assistant, continuation, final), - last_message=final, - ), - ], - continuations=[ - ContinuationDecision(messages=(continuation,), reason="goal"), - ContinuationDecision(), - ], - ) - - outcome = await AgentLoop(engine, services).run(state) - - assert outcome.status == AgentRunStatus.COMPLETED - assert engine.snapshots[1].messages == (user, assistant, continuation) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("effects", "expected_status"), - [ - (AttemptEffects(received_chunk=True), AgentRunStatus.RETRYABLE_FAILURE), - ( - AttemptEffects(tool_execution_started=True), - AgentRunStatus.FATAL_FAILURE, - ), - ], -) -async def test_failure_is_retryable_only_before_observable_effects( - effects: AttemptEffects, - expected_status: AgentRunStatus, -) -> None: - user = Message("user-1", "hello") - state = _state([user]) - failure = StepFailure( - message="provider failed", - error_data={}, - assistant_message_id=None, - reason="provider_error", - allow_fallback=True, - attempt_state=effects, - ) - engine = FakeStepEngine( - [StepResult(action="stop", error=failure.message, failure=failure)], - ) - services = FakeRuntimeServices( - preparations=[_ready], - boundaries=[ModelTurnBoundary(messages=(user,))], - ) - - outcome = await AgentLoop(engine, services).run(state) - - assert outcome.status == expected_status - assert outcome.failure is failure - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("result", "expected_status"), - [ - (StepResult(action="compact", error="context overflow"), AgentRunStatus.CONTEXT_OVERFLOW), - (StepResult(action="unexpected"), AgentRunStatus.FATAL_FAILURE), - ], -) -async def test_loop_returns_structured_non_success_outcomes( - result: StepResult, - expected_status: AgentRunStatus, -) -> None: - user = Message("user-1", "hello") - engine = FakeStepEngine([result]) - services = FakeRuntimeServices( - preparations=[_ready], - boundaries=[ModelTurnBoundary(messages=(user,))], - ) - - outcome = await AgentLoop(engine, services).run(_state([user])) - - assert outcome.status == expected_status - - -@pytest.mark.asyncio -async def test_loop_aborts_after_current_step_boundary() -> None: - user = Message("user-1", "hello") - checks = iter([False, True]) - engine = FakeStepEngine([StepResult(action="stop")]) - services = FakeRuntimeServices( - preparations=[_ready], - boundaries=[ModelTurnBoundary(messages=(user,))], - ) - - outcome = await AgentLoop( - engine, - services, - abort_requested=lambda: next(checks), - ).run(_state([user])) - - assert outcome.status == AgentRunStatus.ABORTED - - -@pytest.mark.asyncio -async def test_ready_preparation_requires_snapshot() -> None: - services = FakeRuntimeServices( - preparations=[ModelTurnPreparation(status=TurnPreparationStatus.READY)], - boundaries=[], - ) - - outcome = await AgentLoop(FakeStepEngine([]), services).run(_state()) - - assert outcome.status == AgentRunStatus.FATAL_FAILURE - assert outcome.error == "Runtime services returned READY without a model-turn snapshot" diff --git a/tests/agent/runtime/test_contracts.py b/tests/agent/runtime/test_contracts.py deleted file mode 100644 index a12ed2a14..000000000 --- a/tests/agent/runtime/test_contracts.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Tests for host-neutral agent runtime contracts.""" - -from flocks.agent.runtime.contracts import ( - AttemptEffects, - ModelTurnSnapshot, - RuntimeModel, -) - - -def test_attempt_effects_allow_replay_only_before_observable_effects() -> None: - effects = AttemptEffects(received_chunk=True) - - assert effects.replay_safe is True - - effects.observable_output_started = True - assert effects.replay_safe is False - - effects.observable_output_started = False - effects.tool_execution_started = True - assert effects.replay_safe is False - - effects.tool_execution_started = False - effects.durable_side_effect_possible = True - assert effects.replay_safe is False - - -def test_model_turn_snapshot_defensively_freezes_collections() -> None: - messages = ["user"] - metadata = {"tool_revision": 1} - snapshot = ModelTurnSnapshot( - session_id="session-1", - agent_name="rex", - active_model=RuntimeModel("provider", "model"), - model_turn_index=2, - trace_step=5, - messages=tuple(messages), - last_user="user", - metadata=metadata, - ) - - messages.append("new input") - metadata["tool_revision"] = 2 - - assert snapshot.messages == ("user",) - assert snapshot.metadata == {"tool_revision": 1} diff --git a/tests/agent/test_unified_session_loop.py b/tests/agent/test_unified_session_loop.py index 44aaa8fa0..8aa4bb81c 100644 --- a/tests/agent/test_unified_session_loop.py +++ b/tests/agent/test_unified_session_loop.py @@ -2,9 +2,9 @@ Tests for Phase 1: Unified UI entry via SessionLoop. Verifies that: -1. RunnerCallbacks.event_publish_callback is passed through to StreamProcessor -2. LoopCallbacks carries runner_callbacks and event_publish_callback -3. SessionRunner uses explicit callbacks (doesn't override with CLI fallback) +1. LoopCallbacks carries model, tool, and event callbacks directly +2. StepEngine receives the same explicit callback object +3. The runtime has no reverse dependency on CLI callback globals 4. _resolve_model implements 5-level priority correctly """ @@ -14,78 +14,64 @@ from unittest.mock import AsyncMock, MagicMock, patch from dataclasses import dataclass -from flocks.session.runner import RunnerCallbacks from flocks.session.session_loop import LoopCallbacks -class TestRunnerCallbacksEventPublish: - """RunnerCallbacks should carry event_publish_callback.""" +class TestLoopCallbacksFields: + """LoopCallbacks should carry all runtime callbacks directly.""" def test_event_publish_callback_field_exists(self): - cb = RunnerCallbacks() + cb = LoopCallbacks() assert hasattr(cb, 'event_publish_callback') assert cb.event_publish_callback is None def test_event_publish_callback_can_be_set(self): publish = AsyncMock() - cb = RunnerCallbacks(event_publish_callback=publish) + cb = LoopCallbacks(event_publish_callback=publish) assert cb.event_publish_callback is publish - -class TestLoopCallbacksFields: - """LoopCallbacks should carry event_publish_callback and runner_callbacks.""" - - def test_event_publish_callback_field(self): - cb = LoopCallbacks() - assert hasattr(cb, 'event_publish_callback') - assert cb.event_publish_callback is None - - def test_runner_callbacks_field(self): - cb = LoopCallbacks() - assert hasattr(cb, 'runner_callbacks') - assert cb.runner_callbacks is None - - def test_pass_runner_callbacks(self): - runner_cb = RunnerCallbacks(on_error=AsyncMock()) - loop_cb = LoopCallbacks(runner_callbacks=runner_cb) - assert loop_cb.runner_callbacks is runner_cb - assert loop_cb.runner_callbacks.on_error is not None + def test_runtime_callbacks_are_flat(self): + on_text_delta = AsyncMock() + on_tool_start = AsyncMock() + callbacks = LoopCallbacks( + on_text_delta=on_text_delta, + on_tool_start=on_tool_start, + ) + assert callbacks.on_text_delta is on_text_delta + assert callbacks.on_tool_start is on_tool_start -class TestCallbackPrecedence: - """SessionRunner should not override explicit callbacks with CLI fallback.""" +class TestCallbackIdentity: + """The runtime should use the callbacks explicitly injected by callers.""" - def test_explicit_callbacks_not_overridden(self): - """When event_publish_callback is set, CLI fallback should NOT be used.""" + def test_explicit_callbacks_are_complete(self): publish = AsyncMock() - cb = RunnerCallbacks(event_publish_callback=publish) - - # Verify the check that _process_step uses - has_explicit = any([ - cb.on_text_delta, - cb.on_tool_start, - cb.on_tool_end, - cb.on_error, - cb.event_publish_callback, - ]) - assert has_explicit is True - - def test_empty_callbacks_allows_cli_fallback(self): - """When no callbacks are set, CLI fallback should be used.""" - cb = RunnerCallbacks() - has_explicit = any([ - cb.on_text_delta, - cb.on_tool_start, - cb.on_tool_end, - cb.on_error, - cb.event_publish_callback, - ]) - assert has_explicit is False + cb = LoopCallbacks(event_publish_callback=publish) + assert cb.event_publish_callback is publish class TestResolveModel: """Test the _resolve_model 5-level priority.""" + @pytest.fixture(autouse=True) + def _active_write_passthrough(self, monkeypatch): + """Persist mocked route messages without requiring stored sessions.""" + from flocks.session.session import Session + + async def run_active_write( + _cls, + _session_id, + operation, + **_kwargs, + ): + return await operation() + + monkeypatch.setattr( + Session, + "run_active_write", + classmethod(run_active_write), + ) + @pytest.mark.asyncio async def test_priority_1_request_model(self): """Request model takes highest priority.""" diff --git a/tests/channel/test_channel.py b/tests/channel/test_channel.py index ecf2e380c..148d44ba2 100644 --- a/tests/channel/test_channel.py +++ b/tests/channel/test_channel.py @@ -39,6 +39,30 @@ from flocks.utils.rate_limiter import AsyncTokenBucket +@pytest.fixture +def active_write_passthrough(monkeypatch): + """Execute channel writes while recording the lifecycle boundary.""" + from flocks.session.session import Session + + session_ids: list[str] = [] + + async def run_active_write( + _cls, + session_id, + operation, + **_kwargs, + ): + session_ids.append(session_id) + return await operation() + + monkeypatch.setattr( + Session, + "run_active_write", + classmethod(run_active_write), + ) + return session_ids + + # ===================================================================== # Helpers — minimal concrete ChannelPlugin for testing # ===================================================================== @@ -1033,7 +1057,11 @@ async def fake_deliver(ctx, session_id=None): assert delivered == ["已清空当前会话历史，共删除 3 条消息。"] @pytest.mark.asyncio - async def test_append_user_message_stores_feishu_media_part(self, monkeypatch): + async def test_append_user_message_stores_feishu_media_part( + self, + monkeypatch, + active_write_passthrough, + ): from flocks.channel.inbound.dispatcher import InboundDispatcher from flocks.config.config import ChannelConfig @@ -1081,9 +1109,14 @@ async def test_append_user_message_stores_feishu_media_part(self, monkeypatch): assert stored_part.filename == "diagram.png" assert stored_part.mime == "image/png" assert stored_part.url == "file:///tmp/diagram.png" + assert active_write_passthrough == ["session_1"] @pytest.mark.asyncio - async def test_append_user_message_accepts_windows_file_uri(self, monkeypatch): + async def test_append_user_message_accepts_windows_file_uri( + self, + monkeypatch, + active_write_passthrough, + ): from flocks.channel.inbound.dispatcher import InboundDispatcher from flocks.config.config import ChannelConfig @@ -1130,24 +1163,25 @@ def fake_isfile(path: str) -> bool: assert stored_part.type == "file" assert stored_part.filename == "channel image.png" assert stored_part.mime == "image/png" + assert active_write_passthrough == ["session_1"] class TestMultimodalInput: @pytest.mark.asyncio async def test_runner_builds_multimodal_user_message_for_image_parts(self, tmp_path, monkeypatch): from flocks.session.message import FilePart, MessageRole, TextPart - from flocks.session.runner import SessionRunner + from flocks.session.runtime.step_engine import StepEngine image_path = tmp_path / "sample.png" image_path.write_bytes(b"image-bytes") - runner = SessionRunner( + runner = StepEngine( session=SimpleNamespace(id="session_1"), provider_id="anthropic", ) monkeypatch.setattr( - "flocks.session.runner.Message.parts", + "flocks.session.runtime.step_engine.Message.parts", AsyncMock( return_value=[ TextPart( @@ -1211,18 +1245,18 @@ def test_anthropic_provider_formats_image_blocks(self): @pytest.mark.asyncio async def test_runner_extracts_plain_text_file_content(self, tmp_path, monkeypatch): from flocks.session.message import FilePart, MessageRole - from flocks.session.runner import SessionRunner + from flocks.session.runtime.step_engine import StepEngine text_path = tmp_path / "notes.txt" text_path.write_text("line 1\nline 2", encoding="utf-8") - runner = SessionRunner( + runner = StepEngine( session=SimpleNamespace(id="session_1"), provider_id="anthropic", ) monkeypatch.setattr( - "flocks.session.runner.Message.parts", + "flocks.session.runtime.step_engine.Message.parts", AsyncMock( return_value=[ FilePart( @@ -1250,18 +1284,18 @@ async def test_runner_extracts_plain_text_file_content(self, tmp_path, monkeypat @pytest.mark.asyncio async def test_runner_extracts_pdf_content(self, tmp_path, monkeypatch): from flocks.session.message import FilePart, MessageRole - from flocks.session.runner import SessionRunner + from flocks.session.runtime.step_engine import StepEngine pdf_path = tmp_path / "report.pdf" pdf_path.write_bytes(b"%PDF-test") - runner = SessionRunner( + runner = StepEngine( session=SimpleNamespace(id="session_1"), provider_id="anthropic", ) monkeypatch.setattr( - "flocks.session.runner.Message.parts", + "flocks.session.runtime.step_engine.Message.parts", AsyncMock( return_value=[ FilePart( @@ -2001,7 +2035,11 @@ async def fake_download(msg, config): assert store_part.await_args.args[2].type == "file" @pytest.mark.asyncio - async def test_wecom_pipeline_stores_file_part(self, monkeypatch): + async def test_wecom_pipeline_stores_file_part( + self, + monkeypatch, + active_write_passthrough, + ): from flocks.channel.inbound.dispatcher import InboundDispatcher from flocks.config.config import ChannelConfig @@ -2042,9 +2080,14 @@ async def fake_download(msg, config): assert stored_part.filename == "report.pdf" assert stored_part.mime == "application/pdf" assert stored_part.url == "file:///tmp/report.pdf" + assert active_write_passthrough == ["s1"] @pytest.mark.asyncio - async def test_dingtalk_pipeline_stores_file_part(self, monkeypatch): + async def test_dingtalk_pipeline_stores_file_part( + self, + monkeypatch, + active_write_passthrough, + ): from flocks.channel.inbound.dispatcher import InboundDispatcher created_message = SimpleNamespace(id="m1") @@ -2082,9 +2125,14 @@ async def fake_download(msg, config): stored_part = store_part.await_args_list[0].args[2] assert stored_part.type == "file" assert stored_part.filename == "image.png" + assert active_write_passthrough == ["s1"] @pytest.mark.asyncio - async def test_telegram_pipeline_stores_file_part(self, monkeypatch): + async def test_telegram_pipeline_stores_file_part( + self, + monkeypatch, + active_write_passthrough, + ): from flocks.channel.inbound.dispatcher import InboundDispatcher created_message = SimpleNamespace(id="m1") @@ -2122,3 +2170,4 @@ async def fake_download(msg, config): stored_part = store_part.await_args_list[0].args[2] assert stored_part.type == "file" assert stored_part.filename == "photo.jpg" + assert active_write_passthrough == ["s1"] diff --git a/tests/integration/test_real_tool_calls.py b/tests/integration/test_real_tool_calls.py index 86adb3f5f..7b88006b2 100644 --- a/tests/integration/test_real_tool_calls.py +++ b/tests/integration/test_real_tool_calls.py @@ -191,16 +191,11 @@ async def on_tool_start(tool_name, args): async def on_tool_end(tool_name, result): tool_ends.append((tool_name, result)) - from flocks.session.runner import RunnerCallbacks - runner_callbacks = RunnerCallbacks( + callbacks = LoopCallbacks( on_tool_start=on_tool_start, on_tool_end=on_tool_end, ) - callbacks = LoopCallbacks( - runner_callbacks=runner_callbacks - ) - # Mock LLM with patch('flocks.provider.provider.Provider.chat') as mock_chat: first_response = MagicMock() diff --git a/tests/observability/test_langfuse_observability.py b/tests/observability/test_langfuse_observability.py index 8dcc02ef8..fa476e53a 100644 --- a/tests/observability/test_langfuse_observability.py +++ b/tests/observability/test_langfuse_observability.py @@ -125,7 +125,7 @@ def test_create_trace_forwards_tags(monkeypatch): monkeypatch.setattr(lf, "_get_client", lambda: client) obs = lf.create_trace( - name="SessionRunner.step", + name="StepEngine.step", session_id="s1", tags=["session:s1", "step:2", "session_step:s1:2"], input={"step": 2}, @@ -141,7 +141,7 @@ def test_create_trace_uses_start_observation_for_new_sdk(monkeypatch): monkeypatch.setattr(lf, "_get_client", lambda: client) obs = lf.create_trace( - name="SessionRunner.step", + name="StepEngine.step", session_id="s1", user_id="u1", tags=["session:s1", "step:2"], @@ -157,7 +157,7 @@ def test_create_trace_uses_start_observation_for_new_sdk(monkeypatch): assert client.start_observation_payload["metadata"]["session_id"] == "s1" assert client.start_observation_payload["metadata"]["user_id"] == "u1" assert client.start_observation_payload["metadata"]["tags"] == ["session:s1", "step:2"] - assert obs._otel_span.attributes["langfuse.trace.name"] == "SessionRunner.step" + assert obs._otel_span.attributes["langfuse.trace.name"] == "StepEngine.step" assert obs._otel_span.attributes["session.id"] == "s1" assert obs._otel_span.attributes["user.id"] == "u1" assert obs._otel_span.attributes["langfuse.trace.tags"] == ["session:s1", "step:2"] @@ -167,7 +167,7 @@ def test_generation_and_span_inherit_trace_dimensions_from_parent(monkeypatch): monkeypatch.setattr(lf, "_get_client", lambda: object()) parent = _TrackingObservation("trace", {"name": "trace"}) - parent._otel_span.attributes["langfuse.trace.name"] = "SessionRunner.step" + parent._otel_span.attributes["langfuse.trace.name"] = "StepEngine.step" parent._otel_span.attributes["session.id"] = "s1" parent._otel_span.attributes["user.id"] = "u1" parent._otel_span.attributes["langfuse.trace.tags"] = ["session:s1", "step:2"] @@ -175,11 +175,11 @@ def test_generation_and_span_inherit_trace_dimensions_from_parent(monkeypatch): gen = lf.create_generation(parent=parent, name="LLM.generate", model="gpt-5", input={"x": 1}) span = lf.create_span(parent=parent, name="Tool.execute.read", input={"path": "/tmp/a"}) - assert gen._otel_span.attributes["langfuse.trace.name"] == "SessionRunner.step" + assert gen._otel_span.attributes["langfuse.trace.name"] == "StepEngine.step" assert gen._otel_span.attributes["session.id"] == "s1" assert gen._otel_span.attributes["user.id"] == "u1" assert gen._otel_span.attributes["langfuse.trace.tags"] == ["session:s1", "step:2"] - assert span._otel_span.attributes["langfuse.trace.name"] == "SessionRunner.step" + assert span._otel_span.attributes["langfuse.trace.name"] == "StepEngine.step" assert span._otel_span.attributes["session.id"] == "s1" assert span._otel_span.attributes["user.id"] == "u1" assert span._otel_span.attributes["langfuse.trace.tags"] == ["session:s1", "step:2"] diff --git a/tests/session/runtime/test_agent_loop.py b/tests/session/runtime/test_agent_loop.py new file mode 100644 index 000000000..a8410f4af --- /dev/null +++ b/tests/session/runtime/test_agent_loop.py @@ -0,0 +1,324 @@ +"""Tests for the logical-input AgentLoop.""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Callable +from dataclasses import dataclass +from types import SimpleNamespace + +import pytest + +from flocks.session.runtime.agent_loop import AgentLoop +from flocks.session.runtime.contracts import ( + AgentRunState, + AgentRunStatus, + AttemptEffects, + ModelTurnBoundary, + ModelTurnPreparation, + ModelTurnSnapshot, + QueuedInputBatch, + RuntimeModel, + StepFailure, + StepResult, + TurnPreparationStatus, +) + + +@dataclass(frozen=True) +class Message: + id: str + content: str + + +PreparationFactory = Callable[ + [AgentRunState[Message]], + ModelTurnPreparation[Message], +] + + +class FakeStepEngine: + def __init__(self, results: list[StepResult]): + self._results = deque(results) + self.snapshots: list[ModelTurnSnapshot[Message]] = [] + + async def run(self, snapshot: ModelTurnSnapshot[Message]) -> StepResult: + self.snapshots.append(snapshot) + return self._results.popleft() + + +class FakeTurn: + """Script the two boundaries AgentLoop is allowed to call.""" + + def __init__( + self, + state: AgentRunState[Message], + preparations: list[ModelTurnPreparation[Message] | PreparationFactory], + boundaries: list[ModelTurnBoundary[Message]], + *, + abort_after_commit: bool = False, + ) -> None: + self.state = state + self._preparations = deque(preparations) + self._boundaries = deque(boundaries) + self.prepared_messages: list[tuple[Message, ...]] = [] + self.aborted = False + self._abort_after_commit = abort_after_commit + self.session = SimpleNamespace(id=state.session_id) + self.step = 0 + + async def prepare_step(self) -> ModelTurnPreparation[Message]: + self.prepared_messages.append(tuple(self.state.messages)) + preparation = self._preparations.popleft() + if callable(preparation): + return preparation(self.state) + return preparation + + async def commit_step( + self, + _step_result: StepResult, + ) -> ModelTurnBoundary[Message]: + boundary = self._boundaries.popleft() + if self._abort_after_commit: + self.aborted = True + return boundary + + +def _state(messages: list[Message] | None = None) -> AgentRunState[Message]: + return AgentRunState( + session_id="session-1", + agent_name="rex", + active_model=RuntimeModel("provider-a", "model-a"), + messages=list(messages or [Message("user-1", "hello")]), + ) + + +def _ready( + state: AgentRunState[Message], + *, + turn: int = 0, +) -> ModelTurnPreparation[Message]: + messages = tuple(state.messages) + return ModelTurnPreparation( + status=TurnPreparationStatus.READY, + snapshot=ModelTurnSnapshot( + session_id=state.session_id, + agent_name=state.agent_name, + active_model=state.active_model, + model_turn_index=turn, + trace_step=turn, + messages=messages, + last_user=messages[-1], + ), + ) + + +async def _run( + state: AgentRunState[Message], + engine: FakeStepEngine, + preparations, + boundaries, + *, + abort_after_commit: bool = False, +): + turn = FakeTurn( + state, + preparations, + boundaries, + abort_after_commit=abort_after_commit, + ) + outcome = await AgentLoop().run(turn, engine) + return outcome, turn + + +@pytest.mark.asyncio +async def test_loop_honors_deferred_preparation_then_completes() -> None: + user = Message("user-1", "hello") + assistant = Message("assistant-1", "done") + engine = FakeStepEngine([StepResult(action="stop")]) + outcome, turn = await _run( + _state([user]), + engine, + [ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE), _ready], + [ModelTurnBoundary(messages=(user, assistant), last_message=assistant)], + ) + + assert outcome.status == AgentRunStatus.COMPLETED + assert outcome.last_message == assistant + assert len(engine.snapshots) == 1 + assert len(turn.prepared_messages) == 2 + + +@pytest.mark.asyncio +async def test_loop_records_model_actually_used_by_step_engine() -> None: + user = Message("user-1", "hello") + assistant = Message("assistant-1", "done") + fallback = RuntimeModel("provider-b", "model-b") + outcome, _ = await _run( + _state([user]), + FakeStepEngine([StepResult(action="stop", effective_model=fallback)]), + [_ready], + [ModelTurnBoundary(messages=(user, assistant), last_message=assistant)], + ) + assert outcome.state.active_model == fallback + + +@pytest.mark.asyncio +async def test_loop_runs_another_step_after_tool_continue() -> None: + user = Message("user-1", "hello") + tool_result = Message("tool-1", "tool result") + assistant = Message("assistant-1", "done") + engine = FakeStepEngine( + [StepResult(action="continue"), StepResult(action="stop")], + ) + outcome, _ = await _run( + _state([user]), + engine, + [_ready, lambda current: _ready(current, turn=1)], + [ + ModelTurnBoundary(messages=(user, tool_result), last_message=tool_result), + ModelTurnBoundary( + messages=(user, tool_result, assistant), + last_message=assistant, + ), + ], + ) + assert outcome.status == AgentRunStatus.COMPLETED + assert len(engine.snapshots) == 2 + assert engine.snapshots[1].messages == (user, tool_result) + + +@pytest.mark.asyncio +async def test_queued_input_yields_to_session_loop() -> None: + user = Message("user-1", "hello") + assistant = Message("assistant-1", "first answer") + queued_user = Message("user-2", "follow up") + outcome, _ = await _run( + _state([user]), + FakeStepEngine([StepResult(action="stop")]), + [_ready], + [ + ModelTurnBoundary( + messages=(user, assistant), + last_message=assistant, + queued_inputs=QueuedInputBatch(messages=(queued_user,)), + ), + ], + ) + assert outcome.status == AgentRunStatus.INPUT_AVAILABLE + + +@pytest.mark.asyncio +async def test_queued_input_precedes_final_step_failure() -> None: + user = Message("user-1", "hello") + failed = Message("assistant-1", "provider failed") + queued_user = Message("user-2", "try this instead") + failure = StepFailure( + message="provider failed", + error_data={}, + assistant_message_id=failed.id, + reason="provider_error", + allow_fallback=False, + attempt_state=AttemptEffects(observable_output_started=True), + ) + outcome, _ = await _run( + _state([user]), + FakeStepEngine( + [StepResult(action="stop", error=failure.message, failure=failure)], + ), + [_ready], + [ + ModelTurnBoundary( + messages=(user, failed, queued_user), + last_message=failed, + queued_inputs=QueuedInputBatch(messages=(queued_user,)), + ), + ], + ) + assert outcome.status == AgentRunStatus.INPUT_AVAILABLE + assert outcome.step_result.failure is failure + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("effects", "expected_status"), + [ + (AttemptEffects(received_chunk=True), AgentRunStatus.RETRYABLE_FAILURE), + ( + AttemptEffects(tool_execution_started=True), + AgentRunStatus.FATAL_FAILURE, + ), + ], +) +async def test_failure_is_retryable_only_before_observable_effects( + effects: AttemptEffects, + expected_status: AgentRunStatus, +) -> None: + user = Message("user-1", "hello") + failure = StepFailure( + message="provider failed", + error_data={}, + assistant_message_id=None, + reason="provider_error", + allow_fallback=True, + attempt_state=effects, + ) + outcome, _ = await _run( + _state([user]), + FakeStepEngine( + [StepResult(action="stop", error=failure.message, failure=failure)], + ), + [_ready], + [ModelTurnBoundary(messages=(user,))], + ) + assert outcome.status == expected_status + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("result", "expected_status"), + [ + ( + StepResult(action="compact", error="context overflow"), + AgentRunStatus.CONTEXT_OVERFLOW, + ), + (StepResult(action="unexpected"), AgentRunStatus.FATAL_FAILURE), + ], +) +async def test_loop_returns_structured_non_success_outcomes( + result: StepResult, + expected_status: AgentRunStatus, +) -> None: + user = Message("user-1", "hello") + outcome, _ = await _run( + _state([user]), + FakeStepEngine([result]), + [_ready], + [ModelTurnBoundary(messages=(user,))], + ) + assert outcome.status == expected_status + + +@pytest.mark.asyncio +async def test_loop_aborts_after_current_step_boundary() -> None: + user = Message("user-1", "hello") + outcome, _ = await _run( + _state([user]), + FakeStepEngine([StepResult(action="stop")]), + [_ready], + [ModelTurnBoundary(messages=(user,))], + abort_after_commit=True, + ) + assert outcome.status == AgentRunStatus.ABORTED + + +@pytest.mark.asyncio +async def test_ready_preparation_requires_snapshot() -> None: + outcome, _ = await _run( + _state(), + FakeStepEngine([]), + [ModelTurnPreparation(status=TurnPreparationStatus.READY)], + [], + ) + assert outcome.status == AgentRunStatus.FATAL_FAILURE + assert outcome.error == "SessionTurn returned READY without a model-turn snapshot" diff --git a/tests/session/runtime/test_contracts.py b/tests/session/runtime/test_contracts.py new file mode 100644 index 000000000..0808ca940 --- /dev/null +++ b/tests/session/runtime/test_contracts.py @@ -0,0 +1,87 @@ +"""Tests for session-neutral agent runtime contracts.""" + +import dataclasses + +import pytest + +from flocks.session.runtime.contracts import ( + AttemptEffects, + ModelRequest, + ModelTurnSnapshot, + RuntimeModel, + StepAction, + StepResult, +) + + +def test_attempt_effects_allow_replay_only_before_observable_effects() -> None: + effects = AttemptEffects(received_chunk=True) + + assert effects.replay_safe is True + + effects.observable_output_started = True + assert effects.replay_safe is False + + effects.observable_output_started = False + effects.tool_execution_started = True + assert effects.replay_safe is False + + +def test_model_turn_snapshot_defensively_freezes_collections() -> None: + messages = ["user"] + metadata = {"tool_revision": 1} + snapshot = ModelTurnSnapshot( + session_id="session-1", + agent_name="rex", + active_model=RuntimeModel("provider", "model"), + model_turn_index=2, + trace_step=5, + messages=tuple(messages), + last_user="user", + metadata=metadata, + ) + + messages.append("new input") + metadata["tool_revision"] = 2 + + assert snapshot.messages == ("user",) + assert snapshot.metadata == {"tool_revision": 1} + + +def test_model_request_freezes_and_isolates_provider_payloads() -> None: + message = {"role": "user", "content": ["hello"]} + tool = {"type": "function", "function": {"name": "read"}} + options = {"reasoning": {"effort": "high"}} + request = ModelRequest( + provider_id="provider", + model_id="model", + messages=(message,), + tools=(tool,), + options=options, + ) + + tool["function"]["name"] = "write" + options["reasoning"]["effort"] = "low" + first_tools = request.provider_tools() + first_tools[0]["function"]["name"] = "mutated" + + assert request.provider_tools()[0]["function"]["name"] == "read" + assert request.provider_options()["reasoning"]["effort"] == "high" + + +def test_model_request_identity_is_frozen() -> None: + request = ModelRequest( + provider_id="provider", + model_id="model", + messages=(), + tools=(), + options={}, + ) + + with pytest.raises(dataclasses.FrozenInstanceError): + request.model_id = "other" + + +def test_step_actions_are_explicit_but_unknown_adapter_values_remain_reportable() -> None: + assert StepResult(action=StepAction.CONTINUE).action == "continue" + assert StepResult(action="unexpected").action == "unexpected" diff --git a/tests/session/runtime/test_session_loop.py b/tests/session/runtime/test_session_loop.py new file mode 100644 index 000000000..540dba233 --- /dev/null +++ b/tests/session/runtime/test_session_loop.py @@ -0,0 +1,354 @@ +"""SessionLoop lifecycle and logical-turn ownership tests.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from flocks.session.core.status import SessionStatus +from flocks.session.message import MessageRole +from flocks.session.runtime.agent_loop import AgentLoop +from flocks.session.runtime.contracts import ( + AgentRunOutcome, + AgentRunState, + AgentRunStatus, + ContinuationDecision, + RuntimeModel, + StepResult, +) +from flocks.session.session import Session, SessionInfo +from flocks.session.session_loop import ( + SessionLoop, + _SessionLeaseRegistry, +) + + +def _session() -> SessionInfo: + return SessionInfo.model_construct( + id="ses_runtime", + projectID="project", + directory="/tmp/project", + agent="rex", + provider="provider", + model="model", + category="user", + status="active", + ) + + +def _message(message_id: str) -> SimpleNamespace: + return SimpleNamespace(id=message_id, role=MessageRole.USER) + + +def _outcome( + turn, + user_id: str, + label: str, + *, + status: AgentRunStatus = AgentRunStatus.COMPLETED, +) -> AgentRunOutcome: + state = AgentRunState( + session_id=turn.session.id, + agent_name=turn.agent_name, + active_model=RuntimeModel(turn.provider_id, turn.model_id), + current_user_id=user_id, + ) + return AgentRunOutcome( + status=status, + state=state, + last_message=SimpleNamespace(label=label), + step_result=( + StepResult(action="stop") + if status == AgentRunStatus.COMPLETED + else None + ), + ) + + +@pytest.fixture +def loop_io(monkeypatch): + session = _session() + active: dict[str, object] = {} + monkeypatch.setattr(SessionLoop, "_active_turns", active) + monkeypatch.setattr( + SessionLoop, + "_leases", + _SessionLeaseRegistry(active), + ) + monkeypatch.setattr( + Session, + "get_by_id", + AsyncMock(return_value=session), + ) + monkeypatch.setattr( + "flocks.session.orphan_tools.abort_orphan_running_parts", + AsyncMock(), + ) + monkeypatch.setattr(Session, "touch", AsyncMock()) + monkeypatch.setattr("flocks.bus.bus.Bus.publish", AsyncMock()) + return session, active + + +@pytest.mark.asyncio +async def test_late_input_keeps_one_lease_and_runs_next_logical_turn( + monkeypatch, + loop_io, +) -> None: + session, active = loop_io + first_user = _message("msg_001") + second_user = _message("msg_002") + prepare = AsyncMock() + + async def prepare_turn(turn): + turn.prepared_user_id = ( + first_user.id if prepare.await_count == 1 else second_user.id + ) + + prepare.side_effect = prepare_turn + continuation = SimpleNamespace( + prepare_logical_turn=prepare, + resolve=AsyncMock(return_value=ContinuationDecision()), + ) + monkeypatch.setattr(SessionLoop, "_continuation_policy", continuation) + run = AsyncMock() + lease_ids: list[int] = [] + + async def run_turn(turn, _engine): + lease_ids.append(id(active[session.id])) + return _outcome( + turn, + first_user.id if run.await_count == 1 else second_user.id, + "first" if run.await_count == 1 else "second", + ) + + run.side_effect = run_turn + monkeypatch.setattr(AgentLoop, "run", run) + monkeypatch.setattr( + "flocks.session.session_loop.Message.list", + AsyncMock( + side_effect=[ + [], + [first_user, second_user], + [first_user, second_user], + ], + ), + ) + + result = await SessionLoop.run( + session.id, + provider_id="provider", + model_id="model", + ) + + assert result.last_message.label == "second" + assert run.await_count == 2 + assert prepare.await_count == 2 + assert continuation.resolve.await_count == 2 + assert len(set(lease_ids)) == 1 + assert active == {} + assert SessionStatus.get(session.id).type == "idle" + + +@pytest.mark.asyncio +async def test_agent_turn_error_settles_without_replaying_current_input( + monkeypatch, + loop_io, +) -> None: + session, active = loop_io + user = _message("msg_001") + + async def prepare(turn): + turn.prepared_user_id = user.id + + continuation = SimpleNamespace( + prepare_logical_turn=AsyncMock(side_effect=prepare), + resolve=AsyncMock(), + ) + monkeypatch.setattr(SessionLoop, "_continuation_policy", continuation) + run = AsyncMock(side_effect=RuntimeError("turn failed")) + monkeypatch.setattr(AgentLoop, "run", run) + monkeypatch.setattr( + "flocks.session.session_loop.Message.list", + AsyncMock(side_effect=[[], [user]]), + ) + + result = await SessionLoop.run( + session.id, + provider_id="provider", + model_id="model", + ) + + assert result.action == "error" + assert result.error == "turn failed" + assert run.await_count == 1 + assert active == {} + + +@pytest.mark.asyncio +async def test_input_available_skips_terminal_continuation_resolution( + monkeypatch, + loop_io, +) -> None: + session, _ = loop_io + user = _message("msg_001") + continuation = SimpleNamespace( + prepare_logical_turn=AsyncMock(), + resolve=AsyncMock(return_value=ContinuationDecision()), + ) + monkeypatch.setattr(SessionLoop, "_continuation_policy", continuation) + run = AsyncMock() + + async def run_turn(turn, _engine): + return _outcome( + turn, + user.id, + "queued" if run.await_count == 1 else "done", + status=( + AgentRunStatus.INPUT_AVAILABLE + if run.await_count == 1 + else AgentRunStatus.COMPLETED + ), + ) + + run.side_effect = run_turn + monkeypatch.setattr(AgentLoop, "run", run) + monkeypatch.setattr( + "flocks.session.session_loop.Message.list", + AsyncMock(side_effect=[[], [user]]), + ) + + result = await SessionLoop.run( + session.id, + provider_id="provider", + model_id="model", + ) + + assert result.last_message.label == "done" + assert run.await_count == 2 + continuation.resolve.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_continuation_runs_under_same_lease( + monkeypatch, + loop_io, +) -> None: + session, active = loop_io + user = _message("msg_001") + continuation = SimpleNamespace( + prepare_logical_turn=AsyncMock(), + resolve=AsyncMock( + side_effect=[ + ContinuationDecision(messages=(user,), reason="goal"), + ContinuationDecision(), + ], + ), + ) + monkeypatch.setattr(SessionLoop, "_continuation_policy", continuation) + run = AsyncMock() + lease_ids: list[int] = [] + + async def run_turn(turn, _engine): + lease_ids.append(id(active[session.id])) + return _outcome( + turn, + user.id, + "first" if run.await_count == 1 else "second", + ) + + run.side_effect = run_turn + monkeypatch.setattr(AgentLoop, "run", run) + monkeypatch.setattr( + "flocks.session.session_loop.Message.list", + AsyncMock(side_effect=[[], [user]]), + ) + + result = await SessionLoop.run( + session.id, + provider_id="provider", + model_id="model", + ) + + assert result.last_message.label == "second" + assert len(set(lease_ids)) == 1 + assert continuation.resolve.await_count == 2 + + +@pytest.mark.asyncio +async def test_idle_is_visible_before_lease_release( + monkeypatch, + loop_io, +) -> None: + session, _ = loop_io + user = _message("msg_001") + continuation = SimpleNamespace( + prepare_logical_turn=AsyncMock(), + resolve=AsyncMock(return_value=ContinuationDecision()), + ) + monkeypatch.setattr(SessionLoop, "_continuation_policy", continuation) + monkeypatch.setattr( + AgentLoop, + "run", + AsyncMock(side_effect=lambda turn, _engine: _outcome( + turn, + user.id, + "done", + )), + ) + release_statuses: list[str] = [] + release = SessionLoop._leases.release + + def record_release(lease): + release_statuses.append(SessionStatus.get(session.id).type) + release(lease) + + monkeypatch.setattr(SessionLoop._leases, "release", record_release) + + async def touch_outside_lock(_project_id, session_id): + assert not Session.lifecycle_lock(session_id).locked() + + monkeypatch.setattr(Session, "touch", AsyncMock(side_effect=touch_outside_lock)) + monkeypatch.setattr( + "flocks.session.session_loop.Message.list", + AsyncMock(side_effect=[[], [user]]), + ) + + await SessionLoop.run( + session.id, + provider_id="provider", + model_id="model", + ) + + assert release_statuses == ["idle"] + + +@pytest.mark.asyncio +async def test_failed_busy_transition_releases_lease( + monkeypatch, + loop_io, +) -> None: + session, active = loop_io + monkeypatch.setattr( + "flocks.session.session_loop.Message.list", + AsyncMock(return_value=[]), + ) + monkeypatch.setattr( + SessionLoop, + "_mark_busy", + AsyncMock(side_effect=RuntimeError("busy failed")), + ) + run = AsyncMock() + monkeypatch.setattr(AgentLoop, "run", run) + + with pytest.raises(RuntimeError, match="busy failed"): + await SessionLoop.run( + session.id, + provider_id="provider", + model_id="model", + ) + + assert active == {} + assert SessionStatus.get(session.id).type == "idle" + run.assert_not_awaited() diff --git a/tests/session/runtime/test_step_engine.py b/tests/session/runtime/test_step_engine.py new file mode 100644 index 000000000..39302b85c --- /dev/null +++ b/tests/session/runtime/test_step_engine.py @@ -0,0 +1,112 @@ +"""Tests for the concrete session StepEngine.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest + +from flocks.session.runtime.step_engine import ( + LlmAttemptState, + StepResult as LegacyStepResult, +) +from flocks.session.runtime.contracts import ( + AttemptEffects, + ModelTurnSnapshot, + RuntimeModel, + StepResult, +) +from flocks.session.runtime.session_turn import LoopCallbacks, SessionTurn +from flocks.session.runtime.step_engine import StepCancelled, StepEngine +from flocks.session.session import SessionInfo + + +def _turn(*, aborted: bool = False) -> SessionTurn: + abort_event = asyncio.Event() + if aborted: + abort_event.set() + return SessionTurn( + session=SessionInfo.model_construct( + id="session-1", + projectID="project", + directory="/tmp/project", + agent="rex", + status="active", + ), + provider_id="provider", + model_id="model", + agent_name="rex", + callbacks=LoopCallbacks(event_publish_callback=None), + abort_event=abort_event, + model_candidates=[RuntimeModel("provider", "model")], + session_start_pending=True, + ) + + +def _snapshot(last_user) -> ModelTurnSnapshot: + return ModelTurnSnapshot( + session_id="session-1", + agent_name="rex", + active_model=RuntimeModel("provider", "model"), + model_turn_index=1, + trace_step=8, + messages=(last_user,), + last_user=last_user, + ) + + +@pytest.mark.asyncio +async def test_step_engine_executes_one_immutable_snapshot() -> None: + last_user = SimpleNamespace(id="user-1") + expected = StepResult(action="stop", content="done") + turn = _turn() + engine = StepEngine.from_turn(turn) + + async def execute(messages, user): + engine._session_start_fired = True + return expected + + process_step = AsyncMock(side_effect=execute) + + with patch.object(StepEngine, "_process_step", process_step): + result = await engine.run(_snapshot(last_user)) + + assert result is expected + assert engine._step == 8 + process_step.assert_awaited_once_with([last_user], last_user) + assert turn.session_start_pending is False + assert turn._current_step_task is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("aborted", "expected_error"), + [ + (False, asyncio.CancelledError), + (True, StepCancelled), + ], +) +async def test_step_engine_only_translates_user_abort( + aborted, + expected_error, +) -> None: + last_user = SimpleNamespace(id="user-1") + turn = _turn(aborted=aborted) + engine = StepEngine.from_turn(turn) + + with ( + patch.object( + StepEngine, + "_process_step", + AsyncMock(side_effect=asyncio.CancelledError), + ), + pytest.raises(expected_error), + ): + await engine.run(_snapshot(last_user)) + + assert turn._current_step_task is None + + +def test_legacy_runner_contract_exports_remain_compatible() -> None: + assert LegacyStepResult is StepResult + assert LlmAttemptState is AttemptEffects diff --git a/tests/session/test_actions.py b/tests/session/test_actions.py new file mode 100644 index 000000000..b6106edcb --- /dev/null +++ b/tests/session/test_actions.py @@ -0,0 +1,71 @@ +"""Tests for non-agent session actions.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from flocks.command.command import Command +from flocks.session import actions + + +@pytest.mark.asyncio +async def test_render_session_command_resolves_template(monkeypatch) -> None: + monkeypatch.setattr( + Command, + "get", + lambda _name: SimpleNamespace(template="Do $ARGUMENTS"), + ) + + result = await actions.render_session_command( + "session-1", + "test", + "the work", + ) + + assert result["template"] == "Do the work" + + +@pytest.mark.asyncio +async def test_run_session_shell_preserves_legacy_response(monkeypatch) -> None: + monkeypatch.setattr( + actions.Session, + "get_by_id", + AsyncMock( + return_value=SimpleNamespace(directory="/tmp/project"), + ), + ) + messages = [ + SimpleNamespace(id="user-1"), + SimpleNamespace(id="assistant-1"), + ] + monkeypatch.setattr( + actions.Message, + "create", + AsyncMock(side_effect=messages), + ) + process = SimpleNamespace( + communicate=AsyncMock(return_value=(b"done", b"")), + returncode=0, + ) + create_process = AsyncMock(return_value=process) + monkeypatch.setattr( + actions.asyncio, + "create_subprocess_shell", + create_process, + ) + + result = await actions.run_session_shell( + "session-1", + "rex", + "echo done", + ) + + create_process.assert_awaited_once_with( + "echo done", + stdout=actions.asyncio.subprocess.PIPE, + stderr=actions.asyncio.subprocess.PIPE, + cwd="/tmp/project", + ) + assert result["info"]["id"] == "assistant-1" + assert result["parts"][0]["state"]["output"] == "done" diff --git a/tests/session/test_auto_model_failover.py b/tests/session/test_auto_model_failover.py index b1d655b1f..d1e8d1a60 100644 --- a/tests/session/test_auto_model_failover.py +++ b/tests/session/test_auto_model_failover.py @@ -6,22 +6,28 @@ import pytest +from flocks.session.runtime.continuation_policy import DEFAULT_CONTINUATION_POLICY +from flocks.session.runtime.agent_loop import AgentLoop +from flocks.session.runtime.contracts import ModelTurnSnapshot, RuntimeModel from flocks.session.message import Message, MessageRole -from flocks.session.runner import ( +from flocks.session.runtime.model_policy import ( + DEFAULT_MODEL_ROUTING_POLICY, + AutoFailoverCooldown, +) +from flocks.session.runtime.step_engine import ( LlmAttemptState, - SessionRunner, + StepEngine, StepFailure, StepResult, ) from flocks.session.session import Session, SessionInfo from flocks.session.session_loop import ( - AutoFailoverCooldown, LoopCallbacks, LoopContext, LoopResult, - RuntimeModel, SessionLoop, ) +from tests.session_runtime_testkit import run_logical_turns def _session(**updates) -> SessionInfo: @@ -66,6 +72,22 @@ def _ctx( ) +async def _build_model_candidates( + primary: RuntimeModel, + *, + route_seed: str, + preferred: RuntimeModel | None = None, + config=None, +): + return await DEFAULT_MODEL_ROUTING_POLICY.build_candidates( + primary, + route_seed=route_seed, + preferred=preferred, + config=config, + validate_model=SessionLoop.validate_runtime_model, + ) + + def _failure( *, assistant_id: str, @@ -89,11 +111,31 @@ def _failure( ) +async def _process_step_with_failover( + turn: LoopContext, + callbacks: LoopCallbacks, + messages, + last_user, +) -> StepResult: + turn.callbacks = callbacks + return await StepEngine.from_turn(turn).run( + ModelTurnSnapshot( + session_id=turn.session.id, + agent_name=turn.agent_name, + active_model=RuntimeModel(turn.provider_id, turn.model_id), + model_turn_index=turn.step, + trace_step=turn.trace_step, + messages=tuple(messages), + last_user=last_user, + ), + ) + + @pytest.fixture(autouse=True) def _clear_cooldowns(): - SessionLoop._auto_failover_cooldowns.clear() + DEFAULT_MODEL_ROUTING_POLICY.cooldowns.clear() yield - SessionLoop._auto_failover_cooldowns.clear() + DEFAULT_MODEL_ROUTING_POLICY.cooldowns.clear() @pytest.mark.parametrize( @@ -116,10 +158,12 @@ def test_failover_classifier( message: str, reason: str, ): - decision = SessionRunner.classify_failover_error({ - "name": "APIError", - "data": {"message": message, "statusCode": status_code}, - }) + decision = StepEngine.classify_failover_error( + { + "name": "APIError", + "data": {"message": message, "statusCode": status_code}, + } + ) assert decision.eligible is True assert decision.reason == reason @@ -144,7 +188,7 @@ async def test_auto_runner_uses_standard_retry_policy( status_code: int, expected_calls: int, ): - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="primary", model_id="primary-model", @@ -160,19 +204,21 @@ async def test_auto_runner_uses_standard_retry_policy( call_llm = AsyncMock(side_effect=failure) monkeypatch.setattr( - "flocks.session.runner.Agent.get", - AsyncMock(return_value=SimpleNamespace( - name="rex", - steps=None, - mode="primary", - prompt="", - tools=[], - )), - ) - monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) - monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + "flocks.session.runtime.step_engine.Agent.get", + AsyncMock( + return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + ) + ), + ) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.apply_config", AsyncMock()) monkeypatch.setattr( - "flocks.session.runner.SessionPrompt.build_system_prompts", + "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompts", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -186,7 +232,7 @@ async def test_auto_runner_uses_standard_retry_policy( monkeypatch.setattr(Message, "create", AsyncMock(return_value=assistant)) monkeypatch.setattr(Message, "update", AsyncMock()) monkeypatch.setattr(runner, "_call_llm", call_llm) - monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", AsyncMock()) + monkeypatch.setattr("flocks.session.runtime.step_engine.SessionRetry.sleep", AsyncMock()) result = await runner._process_step([last_user], last_user) @@ -208,7 +254,7 @@ async def test_last_auto_candidate_uses_standard_retry_policy( expected_calls: int, ): """The last candidate uses the same retry policy as every other mode.""" - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="fallback", model_id="fallback-model", @@ -224,19 +270,21 @@ async def test_last_auto_candidate_uses_standard_retry_policy( call_llm = AsyncMock(side_effect=failure) monkeypatch.setattr( - "flocks.session.runner.Agent.get", - AsyncMock(return_value=SimpleNamespace( - name="rex", - steps=None, - mode="primary", - prompt="", - tools=[], - )), - ) - monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) - monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + "flocks.session.runtime.step_engine.Agent.get", + AsyncMock( + return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + ) + ), + ) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.apply_config", AsyncMock()) monkeypatch.setattr( - "flocks.session.runner.SessionPrompt.build_system_prompts", + "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompts", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -250,7 +298,7 @@ async def test_last_auto_candidate_uses_standard_retry_policy( monkeypatch.setattr(Message, "create", AsyncMock(return_value=assistant)) monkeypatch.setattr(Message, "update", AsyncMock()) monkeypatch.setattr(runner, "_call_llm", call_llm) - monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", AsyncMock()) + monkeypatch.setattr("flocks.session.runtime.step_engine.SessionRetry.sleep", AsyncMock()) result = await runner._process_step([last_user], last_user) @@ -262,9 +310,7 @@ async def test_last_auto_candidate_uses_standard_retry_policy( ("exception", "status_code", "reason"), [ ( - type("GoogleSdkError", (RuntimeError,), {"code": 429})( - "Resource exhausted" - ), + type("GoogleSdkError", (RuntimeError,), {"code": 429})("Resource exhausted"), 429, "rate_limit", ), @@ -289,7 +335,7 @@ def test_exception_status_is_normalized_from_sdk_shapes( status_code: int, reason: str, ): - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="primary", model_id="primary-model", @@ -298,16 +344,14 @@ def test_exception_status_is_normalized_from_sdk_shapes( error = runner._exception_to_error_dict(exception) assert error["data"]["statusCode"] == status_code - assert SessionRunner.classify_failover_error(error).reason == reason + assert StepEngine.classify_failover_error(error).reason == reason def test_exception_status_is_normalized_from_cause_chain(): - inner = type("GoogleSdkError", (RuntimeError,), {"code": 401})( - "Unauthenticated" - ) + inner = type("GoogleSdkError", (RuntimeError,), {"code": 401})("Unauthenticated") outer = RuntimeError("Provider wrapper failed") outer.__cause__ = inner - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="primary", model_id="primary-model", @@ -316,34 +360,40 @@ def test_exception_status_is_normalized_from_cause_chain(): error = runner._exception_to_error_dict(outer) assert error["data"]["statusCode"] == 401 - assert SessionRunner.classify_failover_error(error).reason == "auth" + assert StepEngine.classify_failover_error(error).reason == "auth" def test_local_validation_error_never_fails_over(): - decision = SessionRunner.classify_failover_error({ - "name": "ValidationError", - "data": {"message": "Local prompt schema validation failed"}, - }) + decision = StepEngine.classify_failover_error( + { + "name": "ValidationError", + "data": {"message": "Local prompt schema validation failed"}, + } + ) assert decision.eligible is False assert decision.reason == "local_error" def test_model_not_found_without_status_fails_over(): - decision = SessionRunner.classify_failover_error({ - "name": "ValueError", - "data": {"message": "Model acme-v2 not found for provider custom"}, - }) + decision = StepEngine.classify_failover_error( + { + "name": "ValueError", + "data": {"message": "Model acme-v2 not found for provider custom"}, + } + ) assert decision.eligible is True assert decision.reason == "model_not_found" def test_content_filter_error_fails_over_immediately(): - decision = SessionRunner.classify_failover_error({ - "name": "BadRequestError", - "data": {"message": "Response blocked by content_filter"}, - }) + decision = StepEngine.classify_failover_error( + { + "name": "BadRequestError", + "data": {"message": "Response blocked by content_filter"}, + } + ) assert decision.eligible is True assert decision.reason == "content_policy" @@ -356,22 +406,24 @@ def test_candidate_switch_keeps_tool_loop_guard_only(): "signature": "same-tool-call", "count": 2, } - ctx.runner_static_cache.update({ - "tool_loop_guard": tool_loop_guard, - "tool_schema_cache": {"primary": "schema"}, - "chat_context_cache": {"primary": "context"}, - "system_prompt": "primary prompt", - }) + ctx.step_static_cache.update( + { + "tool_loop_guard": tool_loop_guard, + "tool_schema_cache": {"primary": "schema"}, + "chat_context_cache": {"primary": "context"}, + "system_prompt": "primary prompt", + } + ) - SessionLoop._select_candidate(ctx, 1) + DEFAULT_MODEL_ROUTING_POLICY.select_candidate(ctx, 1) - assert ctx.runner_static_cache == {"tool_loop_guard": tool_loop_guard} - assert ctx.runner_static_cache["tool_loop_guard"] is tool_loop_guard + assert ctx.step_static_cache == {"tool_loop_guard": tool_loop_guard} + assert ctx.step_static_cache["tool_loop_guard"] is tool_loop_guard @pytest.mark.asyncio async def test_reasoning_only_empty_response_is_not_replayed(monkeypatch): - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="primary", model_id="primary-model", @@ -391,19 +443,21 @@ async def call_llm(*_args, **_kwargs): return StepResult(action="stop", content="") monkeypatch.setattr( - "flocks.session.runner.Agent.get", - AsyncMock(return_value=SimpleNamespace( - name="rex", - steps=None, - mode="primary", - prompt="", - tools=[], - )), - ) - monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) - monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + "flocks.session.runtime.step_engine.Agent.get", + AsyncMock( + return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + ) + ), + ) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.apply_config", AsyncMock()) monkeypatch.setattr( - "flocks.session.runner.SessionPrompt.build_system_prompts", + "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompts", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -418,7 +472,7 @@ async def call_llm(*_args, **_kwargs): monkeypatch.setattr(Message, "update", AsyncMock()) monkeypatch.setattr(runner, "_call_llm", call_llm) sleep = AsyncMock() - monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", sleep) + monkeypatch.setattr("flocks.session.runtime.step_engine.SessionRetry.sleep", sleep) result = await runner._process_step([last_user], last_user) @@ -490,11 +544,13 @@ def get_reasoning_content(self): reasoning=None, event_type=None, metadata={}, - tool_calls=[{ - "index": 0, - "id": "call_1", - "function": {"name": "example_tool", "arguments": "{}"}, - }], + tool_calls=[ + { + "index": 0, + "id": "call_1", + "function": {"name": "example_tool", "arguments": "{}"}, + } + ], finish_reason=None, usage=None, ) @@ -518,7 +574,7 @@ async def stream(): return stream() provider = FailingStreamProvider() - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="primary", model_id="primary-model", @@ -529,19 +585,21 @@ async def stream(): assistant = SimpleNamespace(id="msg_assistant") monkeypatch.setattr( - "flocks.session.runner.Agent.get", - AsyncMock(return_value=SimpleNamespace( - name="rex", - steps=None, - mode="primary", - prompt="", - tools=[], - )), - ) - monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) - monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + "flocks.session.runtime.step_engine.Agent.get", + AsyncMock( + return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + ) + ), + ) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.apply_config", AsyncMock()) monkeypatch.setattr( - "flocks.session.runner.SessionPrompt.build_system_prompts", + "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompts", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -555,18 +613,18 @@ async def stream(): monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr(Message, "create", AsyncMock(return_value=assistant)) monkeypatch.setattr(Message, "update", AsyncMock()) - monkeypatch.setattr("flocks.session.runner.StreamProcessor", FakeStreamProcessor) + monkeypatch.setattr("flocks.session.runtime.step_engine.StreamProcessor", FakeStreamProcessor) monkeypatch.setattr( - "flocks.session.runner.HookPipeline.has_stage_handlers", + "flocks.session.runtime.step_engine.HookPipeline.has_stage_handlers", AsyncMock(return_value=False), ) - monkeypatch.setattr("flocks.session.runner.langfuse_is_active", lambda: False) + monkeypatch.setattr("flocks.session.runtime.step_engine.langfuse_is_active", lambda: False) monkeypatch.setattr( "flocks.provider.options.build_provider_options", lambda _provider_id, _model_id: {}, ) sleep = AsyncMock() - monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", sleep) + monkeypatch.setattr("flocks.session.runtime.step_engine.SessionRetry.sleep", sleep) result = await runner._process_step([last_user], last_user) @@ -575,9 +633,7 @@ async def stream(): assert result.failure.allow_fallback is False assert result.failure.attempt_state.received_chunk is True assert result.failure.attempt_state.observable_output_started is True - assert result.failure.attempt_state.tool_execution_started is ( - chunk_kind == "tool" - ) + assert result.failure.attempt_state.tool_execution_started is (chunk_kind == "tool") sleep.assert_not_awaited() @@ -592,14 +648,14 @@ async def process_step(runner, _messages, _last_user): return _failure(assistant_id="msg_failed") return StepResult(action="stop", content="recovered") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) delete = AsyncMock(return_value=True) monkeypatch.setattr(Message, "delete", delete) async def publish(event, payload): events.append((event, payload)) - result = await SessionLoop._process_step_with_failover( + result = await _process_step_with_failover( ctx, LoopCallbacks(event_publish_callback=publish), [last_user], @@ -622,7 +678,7 @@ async def test_queued_user_is_detected_before_replacement_assistant(): role=MessageRole.ASSISTANT, ) - detected = await SessionLoop._detect_queued_user_message( + detected = await DEFAULT_CONTINUATION_POLICY.detect_queued_user_message( "ses_auto", [current_user, queued_user, replacement_assistant], current_user.id, @@ -662,10 +718,10 @@ async def preflight_failure(_runner, _messages, _last_user): final_assistant = SimpleNamespace(id="msg_final_error") create = AsyncMock(return_value=final_assistant) - monkeypatch.setattr(SessionRunner, "_process_step", preflight_failure) + monkeypatch.setattr(StepEngine, "_process_step", preflight_failure) monkeypatch.setattr(Message, "create", create) - result = await SessionLoop._process_step_with_failover( + result = await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], @@ -696,7 +752,7 @@ async def test_failed_blank_message_deletion_stops_switch(monkeypatch): ctx = _ctx() last_user = SimpleNamespace(id="msg_user", agent="rex") monkeypatch.setattr( - SessionRunner, + StepEngine, "_process_step", AsyncMock(return_value=_failure(assistant_id="msg_failed")), ) @@ -704,7 +760,7 @@ async def test_failed_blank_message_deletion_stops_switch(monkeypatch): update = AsyncMock() monkeypatch.setattr(Message, "update", update) - result = await SessionLoop._process_step_with_failover( + result = await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], @@ -738,11 +794,11 @@ async def process_step(runner, _messages, _last_user): return _failure(assistant_id=f"msg_{runner.provider_id}") return StepResult(action="stop", content="recovered") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) delete = AsyncMock(return_value=True) monkeypatch.setattr(Message, "delete", delete) - result = await SessionLoop._process_step_with_failover( + result = await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], @@ -771,13 +827,13 @@ async def test_chain_exhaustion_finalizes_only_last_candidate(monkeypatch): async def process_step(runner, _messages, _last_user): return _failure(assistant_id=f"msg_{runner.provider_id}") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) delete = AsyncMock(return_value=True) update = AsyncMock() monkeypatch.setattr(Message, "delete", delete) monkeypatch.setattr(Message, "update", update) - result = await SessionLoop._process_step_with_failover( + result = await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], @@ -788,7 +844,7 @@ async def process_step(runner, _messages, _last_user): assert delete.await_count == 2 update.assert_awaited_once() assert update.await_args.args[1] == "msg_fallback-2" - cooldown = SessionLoop._auto_failover_cooldowns[ctx.session.id] + cooldown = DEFAULT_MODEL_ROUTING_POLICY.cooldowns[ctx.session.id] assert cooldown.model == RuntimeModel("fallback-2", "model-2") assert cooldown.reason == "chain_exhausted" @@ -811,11 +867,13 @@ async def test_full_loop_reports_chain_exhaustion_once(monkeypatch): parentID=user.id, finish="error", ) - ctx.session_ctx = SimpleNamespace( - get_messages=AsyncMock(side_effect=[ - [user], - [user, final_assistant], - ]) + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user], + [user, final_assistant], + ] + ) ) attempts = [] @@ -823,14 +881,14 @@ async def process_step(runner, _messages, _last_user): attempts.append((runner.provider_id, runner.model_id)) return _failure(assistant_id=f"msg_{runner.provider_id}") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) update = AsyncMock() monkeypatch.setattr(Message, "update", update) on_error = AsyncMock() - result = await SessionLoop._run_loop( + result = await run_logical_turns( ctx, LoopCallbacks( on_error=on_error, @@ -859,7 +917,7 @@ async def test_observable_failure_is_finalized_without_replay(monkeypatch): ctx = _ctx() last_user = SimpleNamespace(id="msg_user", agent="rex") monkeypatch.setattr( - SessionRunner, + StepEngine, "_process_step", AsyncMock(return_value=_failure(assistant_id="msg_partial", safe=False)), ) @@ -868,7 +926,7 @@ async def test_observable_failure_is_finalized_without_replay(monkeypatch): monkeypatch.setattr(Message, "delete", delete) monkeypatch.setattr(Message, "update", update) - result = await SessionLoop._process_step_with_failover( + result = await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], @@ -891,23 +949,26 @@ async def process_step(runner, _messages, _last_user): return _failure(assistant_id="msg_rate", reason="rate_limit") return StepResult(action="stop", content="recovered") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) - await SessionLoop._process_step_with_failover( + await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], last_user, ) - cooldown = SessionLoop._auto_failover_cooldowns[ctx.session.id] + cooldown = DEFAULT_MODEL_ROUTING_POLICY.cooldowns[ctx.session.id] assert cooldown.model == RuntimeModel("fallback", "fallback-model") assert cooldown.reason == "rate_limit" - assert SessionLoop._cooldown_candidate_index( - ctx.session.id, - ctx.model_candidates, - ) == 1 + assert ( + DEFAULT_MODEL_ROUTING_POLICY.cooldown_candidate_index( + ctx.session.id, + ctx.model_candidates, + ) + == 1 + ) @pytest.mark.asyncio @@ -921,10 +982,10 @@ async def process_step(runner, _messages, _last_user): return _failure(assistant_id="msg_rate", reason="rate_limit") return StepResult(action="stop", content="recovered") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) - await SessionLoop._process_step_with_failover( + await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], @@ -932,15 +993,17 @@ async def process_step(runner, _messages, _last_user): ) assert (ctx.provider_id, ctx.model_id) == ("fallback", "fallback-model") - assert ctx.session.id not in SessionLoop._auto_failover_cooldowns + assert ctx.session.id not in DEFAULT_MODEL_ROUTING_POLICY.cooldowns @pytest.mark.asyncio async def test_403_quota_failure_sets_primary_cooldown(monkeypatch): - decision = SessionRunner.classify_failover_error({ - "name": "APIError", - "data": {"message": "Quota exceeded", "statusCode": 403}, - }) + decision = StepEngine.classify_failover_error( + { + "name": "APIError", + "data": {"message": "Quota exceeded", "statusCode": 403}, + } + ) ctx = _ctx() last_user = SimpleNamespace(id="msg_user", agent="rex") @@ -952,17 +1015,17 @@ async def process_step(runner, _messages, _last_user): ) return StepResult(action="stop", content="recovered") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) - await SessionLoop._process_step_with_failover( + await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], last_user, ) - cooldown = SessionLoop._auto_failover_cooldowns[ctx.session.id] + cooldown = DEFAULT_MODEL_ROUTING_POLICY.cooldowns[ctx.session.id] assert cooldown.reason == "rate_limit" assert cooldown.model == RuntimeModel("fallback", "fallback-model") assert cooldown.expires_at > time.monotonic() + 50 @@ -980,18 +1043,18 @@ async def process_step(runner, _messages, _last_user): reason=reason, ) - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) monkeypatch.setattr(Message, "update", AsyncMock()) - await SessionLoop._process_step_with_failover( + await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], last_user, ) - cooldown = SessionLoop._auto_failover_cooldowns[ctx.session.id] + cooldown = DEFAULT_MODEL_ROUTING_POLICY.cooldowns[ctx.session.id] assert cooldown.reason == "rate_limit" assert cooldown.model == RuntimeModel("fallback", "fallback-model") # A 5s anti-replay window must not replace the primary's 60s cooldown. @@ -1032,11 +1095,11 @@ async def validate(provider_id, _model_id, **_kwargs): monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) primary = RuntimeModel("primary", "primary-model") - first = await SessionLoop._build_model_candidates( + first = await _build_model_candidates( primary, route_seed="ses_auto:msg_1", ) - repeated = await SessionLoop._build_model_candidates( + repeated = await _build_model_candidates( primary, route_seed="ses_auto:msg_1", ) @@ -1050,10 +1113,12 @@ async def validate(provider_id, _model_id, **_kwargs): assert all(candidate.provider_id != "missing" for candidate in first) selections = { - tuple(await SessionLoop._build_model_candidates( - primary, - route_seed=f"ses_auto:msg_{index}", - )) + tuple( + await _build_model_candidates( + primary, + route_seed=f"ses_auto:msg_{index}", + ) + ) for index in range(12) } assert len(selections) > 1 @@ -1092,7 +1157,7 @@ async def test_candidate_builder_keeps_active_cooldown_model_in_its_tier( primary = RuntimeModel("primary", "primary-model") cooldown_model = RuntimeModel("other", "other-b") - candidates = await SessionLoop._build_model_candidates( + candidates = await _build_model_candidates( primary, route_seed="ses_auto:new-turn", preferred=cooldown_model, @@ -1107,21 +1172,19 @@ async def test_candidate_builder_keeps_active_cooldown_model_in_its_tier( async def test_auto_configuration_only_requires_available_primary(monkeypatch): monkeypatch.setattr( "flocks.config.config.Config.resolve_default_llm", - AsyncMock(return_value={ - "provider_id": "primary", - "model_id": "primary-model", - }), + AsyncMock( + return_value={ + "provider_id": "primary", + "model_id": "primary-model", + } + ), ) monkeypatch.setattr( SessionLoop, "validate_runtime_model", AsyncMock(return_value=(True, "available")), ) - build_candidates = AsyncMock() - monkeypatch.setattr(SessionLoop, "_build_model_candidates", build_candidates) - assert await SessionLoop.validate_auto_configuration() == (True, "available") - build_candidates.assert_not_awaited() @pytest.mark.asyncio @@ -1145,7 +1208,7 @@ async def test_candidate_builder_allows_primary_only_chain(monkeypatch): primary = RuntimeModel("primary", "primary-model") - assert await SessionLoop._build_model_candidates( + assert await _build_model_candidates( primary, route_seed="ses_auto:msg_primary_only", ) == [primary] @@ -1155,11 +1218,13 @@ async def test_candidate_builder_allows_primary_only_chain(monkeypatch): async def test_candidate_builder_uses_configured_order_without_discovery( monkeypatch, ): - config = SimpleNamespace(fallback_providers=[ - SimpleNamespace(provider_id="other", model_id="model-b"), - SimpleNamespace(provider_id="primary", model_id="model-a"), - SimpleNamespace(provider_id="missing", model_id="missing-model"), - ]) + config = SimpleNamespace( + fallback_providers=[ + SimpleNamespace(provider_id="other", model_id="model-b"), + SimpleNamespace(provider_id="primary", model_id="model-a"), + SimpleNamespace(provider_id="missing", model_id="missing-model"), + ] + ) model_manager = MagicMock() monkeypatch.setattr( "flocks.provider.provider.Provider.apply_config", @@ -1177,7 +1242,7 @@ async def validate(provider_id, _model_id, **_kwargs): monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) primary = RuntimeModel("primary", "primary-model") - candidates = await SessionLoop._build_model_candidates( + candidates = await _build_model_candidates( primary, route_seed="unused-for-configured", preferred=RuntimeModel("other", "model-b"), @@ -1196,9 +1261,11 @@ async def validate(provider_id, _model_id, **_kwargs): async def test_configured_chain_with_no_available_fallbacks_keeps_primary_only( monkeypatch, ): - config = SimpleNamespace(fallback_providers=[ - SimpleNamespace(provider_id="missing", model_id="missing-model"), - ]) + config = SimpleNamespace( + fallback_providers=[ + SimpleNamespace(provider_id="missing", model_id="missing-model"), + ] + ) monkeypatch.setattr( "flocks.provider.provider.Provider.apply_config", AsyncMock(), @@ -1210,7 +1277,7 @@ async def test_configured_chain_with_no_available_fallbacks_keeps_primary_only( ) primary = RuntimeModel("primary", "primary-model") - assert await SessionLoop._build_model_candidates( + assert await _build_model_candidates( primary, route_seed="unused-for-configured", config=config, @@ -1222,15 +1289,15 @@ def test_cooldown_is_cleared_when_primary_changes(): RuntimeModel("new-primary", "new-model"), RuntimeModel("fallback", "fallback-model"), ] - SessionLoop._auto_failover_cooldowns["ses_auto"] = AutoFailoverCooldown( + DEFAULT_MODEL_ROUTING_POLICY.cooldowns["ses_auto"] = AutoFailoverCooldown( model=RuntimeModel("fallback", "fallback-model"), primary=RuntimeModel("old-primary", "old-model"), expires_at=float("inf"), reason="rate_limit", ) - assert SessionLoop._cooldown_candidate_index("ses_auto", candidates) == 0 - assert "ses_auto" not in SessionLoop._auto_failover_cooldowns + assert DEFAULT_MODEL_ROUTING_POLICY.cooldown_candidate_index("ses_auto", candidates) == 0 + assert "ses_auto" not in DEFAULT_MODEL_ROUTING_POLICY.cooldowns @pytest.mark.asyncio @@ -1248,7 +1315,7 @@ async def test_synthetic_subtask_continuation_keeps_fallback(monkeypatch): AsyncMock(return_value=[SimpleNamespace(synthetic=True)]), ) - await SessionLoop._prepare_auto_turn(ctx, synthetic_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, synthetic_user) assert ctx.auto_failover is True assert ctx.turn_user_id == "msg_real" @@ -1275,9 +1342,9 @@ async def test_first_real_turn_builds_stable_chain_from_user_id(monkeypatch): AsyncMock(return_value=config), ) build = AsyncMock(return_value=rebuilt) - monkeypatch.setattr(SessionLoop, "_build_model_candidates", build) + monkeypatch.setattr(DEFAULT_MODEL_ROUTING_POLICY, "build_candidates", build) - await SessionLoop._prepare_auto_turn(ctx, first_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, first_user) assert ctx.turn_user_id == "msg_first" assert ctx.model_candidates == rebuilt @@ -1299,14 +1366,16 @@ async def test_configured_first_real_turn_ignores_cooldown_and_starts_primary( id="msg_first", model={"providerID": "primary", "modelID": "primary-model"}, ) - config = SimpleNamespace(fallback_providers=[ - SimpleNamespace(provider_id="fallback", model_id="fallback-model"), - ]) + config = SimpleNamespace( + fallback_providers=[ + SimpleNamespace(provider_id="fallback", model_id="fallback-model"), + ] + ) rebuilt = [ RuntimeModel("primary", "primary-model"), RuntimeModel("fallback", "fallback-model"), ] - SessionLoop._auto_failover_cooldowns[ctx.session.id] = AutoFailoverCooldown( + DEFAULT_MODEL_ROUTING_POLICY.cooldowns[ctx.session.id] = AutoFailoverCooldown( model=rebuilt[1], primary=rebuilt[0], expires_at=float("inf"), @@ -1317,17 +1386,17 @@ async def test_configured_first_real_turn_ignores_cooldown_and_starts_primary( AsyncMock(return_value=config), ) monkeypatch.setattr( - SessionLoop, - "_build_model_candidates", + DEFAULT_MODEL_ROUTING_POLICY, + "build_candidates", AsyncMock(return_value=rebuilt), ) - await SessionLoop._prepare_auto_turn(ctx, first_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, first_user) assert ctx.model_candidate_policy == "configured" assert ctx.candidate_index == 0 assert (ctx.provider_id, ctx.model_id) == ("primary", "primary-model") - assert ctx.session.id not in SessionLoop._auto_failover_cooldowns + assert ctx.session.id not in DEFAULT_MODEL_ROUTING_POLICY.cooldowns @pytest.mark.asyncio @@ -1350,7 +1419,7 @@ async def test_queued_explicit_model_disables_auto(monkeypatch): AsyncMock(return_value=persisted), ) - await SessionLoop._prepare_auto_turn(ctx, queued_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, queued_user) assert ctx.auto_failover is False assert ctx.model_candidates == [RuntimeModel("explicit", "explicit-model")] @@ -1372,7 +1441,7 @@ async def test_non_webui_loop_cannot_activate_persisted_auto(monkeypatch): AsyncMock(return_value=_session(model_auto=True)), ) - await SessionLoop._prepare_auto_turn(ctx, queued_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, queued_user) assert ctx.auto_failover is False assert ctx.auto_failover_allowed is False @@ -1399,10 +1468,12 @@ async def test_queued_webui_turn_rebuilds_auto_chain(monkeypatch): ) monkeypatch.setattr( "flocks.config.config.Config.resolve_default_llm", - AsyncMock(return_value={ - "provider_id": "primary", - "model_id": "primary-model", - }), + AsyncMock( + return_value={ + "provider_id": "primary", + "model_id": "primary-model", + } + ), ) config = SimpleNamespace(fallback_providers=None) monkeypatch.setattr( @@ -1410,9 +1481,9 @@ async def test_queued_webui_turn_rebuilds_auto_chain(monkeypatch): AsyncMock(return_value=config), ) build = AsyncMock(return_value=rebuilt) - monkeypatch.setattr(SessionLoop, "_build_model_candidates", build) + monkeypatch.setattr(DEFAULT_MODEL_ROUTING_POLICY, "build_candidates", build) - await SessionLoop._prepare_auto_turn(ctx, queued_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, queued_user) assert ctx.auto_failover is True assert ctx.model_candidates == rebuilt @@ -1438,9 +1509,11 @@ async def test_queued_configured_turn_restarts_from_primary(monkeypatch): RuntimeModel("primary", "primary-model"), RuntimeModel("fallback", "fallback-model"), ] - config = SimpleNamespace(fallback_providers=[ - SimpleNamespace(provider_id="fallback", model_id="fallback-model"), - ]) + config = SimpleNamespace( + fallback_providers=[ + SimpleNamespace(provider_id="fallback", model_id="fallback-model"), + ] + ) monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr( "flocks.session.session.Session.get_by_id", @@ -1452,18 +1525,20 @@ async def test_queued_configured_turn_restarts_from_primary(monkeypatch): ) monkeypatch.setattr( "flocks.config.config.Config.resolve_default_llm", - AsyncMock(return_value={ - "provider_id": "primary", - "model_id": "primary-model", - }), + AsyncMock( + return_value={ + "provider_id": "primary", + "model_id": "primary-model", + } + ), ) monkeypatch.setattr( - SessionLoop, - "_build_model_candidates", + DEFAULT_MODEL_ROUTING_POLICY, + "build_candidates", AsyncMock(return_value=rebuilt), ) - await SessionLoop._prepare_auto_turn(ctx, queued_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, queued_user) assert ctx.model_candidate_policy == "configured" assert ctx.candidate_index == 0 @@ -1474,11 +1549,11 @@ async def test_queued_configured_turn_restarts_from_primary(monkeypatch): @pytest.mark.parametrize("category", ["user", "entity-config", "workflow"]) async def test_queued_webui_auto_authorizes_active_loop(category): ctx = _ctx(auto=False, category=category) - SessionLoop._active_loops[ctx.session.id] = ctx + SessionLoop._active_turns[ctx.session.id] = ctx try: result = await SessionLoop.run(ctx.session.id, auto_failover=True) finally: - SessionLoop._active_loops.pop(ctx.session.id, None) + SessionLoop._active_turns.pop(ctx.session.id, None) assert result.action == "queued" assert ctx.auto_failover_allowed is True @@ -1491,19 +1566,36 @@ async def test_unsupported_session_loop_ignores_auto_authorization( task_session = _session(category="task") captured_ctx = None - async def run_loop(ctx, _callbacks): + async def run_turn(_loop, ctx, _engine): + from flocks.session.runtime.contracts import ( + AgentRunOutcome, + AgentRunState, + AgentRunStatus, + ) + nonlocal captured_ctx captured_ctx = ctx - return LoopResult(action="stop") + state = AgentRunState( + session_id=ctx.session.id, + agent_name=ctx.agent_name, + active_model=RuntimeModel(ctx.provider_id, ctx.model_id), + ) + return AgentRunOutcome( + status=AgentRunStatus.ABORTED, + state=state, + ) build_candidates = AsyncMock() monkeypatch.setattr( "flocks.session.session.Session.get_by_id", AsyncMock(return_value=task_session), ) - monkeypatch.setattr(SessionLoop, "_build_model_candidates", build_candidates) - monkeypatch.setattr(SessionLoop, "_run_loop", run_loop) - monkeypatch.setattr(SessionLoop, "_publish_session_status", AsyncMock()) + monkeypatch.setattr( + DEFAULT_MODEL_ROUTING_POLICY, + "build_candidates", + build_candidates, + ) + monkeypatch.setattr(AgentLoop, "run", run_turn) monkeypatch.setattr(Message, "list", AsyncMock(return_value=[])) monkeypatch.setattr( "flocks.session.orphan_tools.abort_orphan_running_parts", @@ -1525,20 +1617,18 @@ async def run_loop(ctx, _callbacks): assert captured_ctx is not None assert captured_ctx.auto_failover is False assert captured_ctx.auto_failover_allowed is False - assert captured_ctx.model_candidates == [ - RuntimeModel("primary", "primary-model") - ] + assert captured_ctx.model_candidates == [RuntimeModel("primary", "primary-model")] build_candidates.assert_not_awaited() @pytest.mark.asyncio async def test_active_unsupported_loop_rejects_auto_authorization(): ctx = _ctx(auto=False, category="task") - SessionLoop._active_loops[ctx.session.id] = ctx + SessionLoop._active_turns[ctx.session.id] = ctx try: result = await SessionLoop.run(ctx.session.id, auto_failover=True) finally: - SessionLoop._active_loops.pop(ctx.session.id, None) + SessionLoop._active_turns.pop(ctx.session.id, None) assert result.action == "queued" assert ctx.auto_failover_allowed is False @@ -1547,7 +1637,7 @@ async def test_active_unsupported_loop_rejects_auto_authorization(): @pytest.mark.asyncio async def test_session_delete_clears_auto_failover_cooldown(monkeypatch): session = _session() - SessionLoop._auto_failover_cooldowns[session.id] = AutoFailoverCooldown( + DEFAULT_MODEL_ROUTING_POLICY.cooldowns[session.id] = AutoFailoverCooldown( model=RuntimeModel("fallback", "fallback-model"), primary=RuntimeModel("primary", "primary-model"), expires_at=float("inf"), @@ -1564,4 +1654,4 @@ async def test_session_delete_clears_auto_failover_cooldown(monkeypatch): monkeypatch.setattr("flocks.bus.bus.Bus.publish", AsyncMock()) assert await Session.delete("project", session.id) is True - assert session.id not in SessionLoop._auto_failover_cooldowns + assert session.id not in DEFAULT_MODEL_ROUTING_POLICY.cooldowns diff --git a/tests/session/test_callable_state.py b/tests/session/test_callable_state.py index 385e7cf1e..f20121253 100644 --- a/tests/session/test_callable_state.py +++ b/tests/session/test_callable_state.py @@ -1,6 +1,3 @@ -from pathlib import Path -import tempfile - import pytest from flocks.storage.storage import Storage @@ -10,18 +7,8 @@ get_session_callable_tools, ) - -@pytest.fixture -async def callable_storage(): - with tempfile.TemporaryDirectory() as tmpdir: - db_path = Path(tmpdir) / "test_session_callable.db" - await Storage.init(db_path) - yield - await Storage.clear() - - @pytest.mark.asyncio -async def test_session_callable_persists_unique_sorted_tools(callable_storage) -> None: +async def test_session_callable_persists_unique_sorted_tools() -> None: await add_session_callable_tools("session-callable", ["websearch", "task", "websearch"]) result = await get_session_callable_tools("session-callable") @@ -32,7 +19,7 @@ async def test_session_callable_persists_unique_sorted_tools(callable_storage) - @pytest.mark.asyncio -async def test_session_callable_clear_removes_cache_and_storage(callable_storage) -> None: +async def test_session_callable_clear_removes_cache_and_storage() -> None: await add_session_callable_tools("session-callable-clear", ["websearch"]) await clear_session_callable_tools("session-callable-clear") diff --git a/tests/session/test_cli_session_runner_model_resolution.py b/tests/session/test_cli_session_runner_model_resolution.py index 4c9625f82..f858022d7 100644 --- a/tests/session/test_cli_session_runner_model_resolution.py +++ b/tests/session/test_cli_session_runner_model_resolution.py @@ -73,8 +73,7 @@ async def test_reads_config_model_when_no_cli_flag(self): patch("flocks.agent.registry.Agent.default_agent", new_callable=AsyncMock, return_value="rex"), \ patch("flocks.agent.registry.Agent.get", new_callable=AsyncMock) as mock_agent_get, \ patch("flocks.session.message.Message.create", new_callable=AsyncMock) as mock_msg_create, \ - patch("flocks.session.session_loop.SessionLoop.run", new_callable=AsyncMock) as mock_loop_run, \ - patch("flocks.cli.session_runner._set_cli_callbacks"): + patch("flocks.session.session_loop.SessionLoop.run", new_callable=AsyncMock) as mock_loop_run: mock_agent = MagicMock() mock_agent.name = "rex" @@ -104,8 +103,7 @@ async def test_cli_flag_overrides_config(self): with patch("flocks.agent.registry.Agent.default_agent", new_callable=AsyncMock, return_value="rex"), \ patch("flocks.agent.registry.Agent.get", new_callable=AsyncMock) as mock_agent_get, \ patch("flocks.session.message.Message.create", new_callable=AsyncMock), \ - patch("flocks.session.session_loop.SessionLoop.run", new_callable=AsyncMock) as mock_loop_run, \ - patch("flocks.cli.session_runner._set_cli_callbacks"): + patch("flocks.session.session_loop.SessionLoop.run", new_callable=AsyncMock) as mock_loop_run: mock_agent = MagicMock() mock_agent.name = "rex" diff --git a/tests/session/test_execution_mode.py b/tests/session/test_execution_mode.py index 0ef33c1de..fceb95dc1 100644 --- a/tests/session/test_execution_mode.py +++ b/tests/session/test_execution_mode.py @@ -399,9 +399,9 @@ async def test_read_only_sandbox_allows_only_plan_artifact_write(tmp_path) -> No @pytest.mark.asyncio async def test_runner_filters_tools_with_message_mode(monkeypatch) -> None: - from flocks.session.runner import SessionRunner + from flocks.session.runtime.step_engine import StepEngine - runner = object.__new__(SessionRunner) + runner = object.__new__(StepEngine) runner.session = SimpleNamespace(id="session-1") runner._step = 1 runner.callbacks = SimpleNamespace(event_publish_callback=None) @@ -434,7 +434,7 @@ async def list_tools(**_kwargs): return result monkeypatch.setattr( - "flocks.session.runner.list_session_callable_tool_infos", + "flocks.session.runtime.step_engine.list_session_callable_tool_infos", list_tools, ) monkeypatch.setattr( diff --git a/tests/session/test_lifecycle_hooks.py b/tests/session/test_lifecycle_hooks.py index b08c3581e..b3412b0a1 100644 --- a/tests/session/test_lifecycle_hooks.py +++ b/tests/session/test_lifecycle_hooks.py @@ -4,15 +4,22 @@ import asyncio from types import SimpleNamespace -from unittest.mock import ANY, AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from flocks.session.runtime.contracts import ContinuationDecision from flocks.hooks.pipeline import HookContext, HookStage +from flocks.session.runtime.continuation_policy import DEFAULT_CONTINUATION_POLICY from flocks.session.goal import GoalDecision -from flocks.session.runner import SessionRunner, StepResult +from flocks.session.runtime.model_policy import DEFAULT_MODEL_ROUTING_POLICY +from flocks.session.runtime.step_engine import StepEngine, StepResult from flocks.session.session import SessionInfo -from flocks.session.session_loop import LoopCallbacks, LoopContext, SessionLoop +from flocks.session.session_loop import ( + LoopCallbacks, + LoopContext, +) +from tests.session_runtime_testkit import run_logical_turns def _session(session_id: str = "ses_lifecycle_hooks") -> SessionInfo: @@ -27,12 +34,13 @@ def _session(session_id: str = "ses_lifecycle_hooks") -> SessionInfo: def _loop_context(session_id: str = "ses_lifecycle_hooks") -> LoopContext: - return LoopContext( + context = LoopContext( session=_session(session_id), provider_id="test-provider", model_id="test-model", agent_name="rex", ) + return context @pytest.mark.asyncio @@ -43,7 +51,7 @@ async def test_real_user_turn_is_detected_once_and_synthetic_is_ignored() -> Non with ( patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock( side_effect=[ [], @@ -52,14 +60,14 @@ async def test_real_user_turn_is_detected_once_and_synthetic_is_ignored() -> Non ), ), patch.object( - SessionLoop, - "_run_user_prompt_submit_hook", + DEFAULT_CONTINUATION_POLICY, + "run_user_prompt_submit", AsyncMock(), ) as submit_hook, ): for user in (first_user, first_user, synthetic_user): - if await SessionLoop._prepare_auto_turn(ctx, user): - await SessionLoop._run_user_prompt_submit_hook(ctx, user) + if await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, user): + await DEFAULT_CONTINUATION_POLICY.run_user_prompt_submit(ctx, user) assert ctx.turn_user_id == first_user.id submit_hook.assert_awaited_once_with(ctx, first_user) @@ -79,7 +87,7 @@ async def test_user_prompt_submit_adds_ephemeral_turn_context() -> None: with ( patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", AsyncMock(return_value="implement hooks"), ), patch( @@ -87,7 +95,7 @@ async def test_user_prompt_submit_adds_ephemeral_turn_context() -> None: run_hook, ), ): - await SessionLoop._run_user_prompt_submit_hook(ctx, user) + await DEFAULT_CONTINUATION_POLICY.run_user_prompt_submit(ctx, user) assert ctx.turn_additional_context == "current sprint context" payload = run_hook.await_args.args[0] @@ -101,7 +109,7 @@ async def test_user_prompt_submit_adds_ephemeral_turn_context() -> None: @pytest.mark.asyncio async def test_session_start_runs_only_when_pending() -> None: - runner = SessionRunner( + runner = StepEngine( session=_session("ses_session_start"), provider_id="test-provider", model_id="test-model", @@ -110,7 +118,7 @@ async def test_session_start_runs_only_when_pending() -> None: run_hook = AsyncMock() with patch( - "flocks.session.runner.HookPipeline.run_session_start", + "flocks.session.runtime.step_engine.HookPipeline.run_session_start", run_hook, ): await runner._run_session_start_hook(SimpleNamespace(name="rex")) @@ -137,6 +145,7 @@ async def test_turn_finish_block_creates_synthetic_continuation() -> None: ) continuation = SimpleNamespace(id="msg_continuation") callbacks = LoopCallbacks(event_publish_callback=AsyncMock()) + ctx.callbacks = callbacks create_message = AsyncMock(return_value=continuation) run_hook = AsyncMock( return_value=HookContext( @@ -151,15 +160,15 @@ async def test_turn_finish_block_creates_synthetic_continuation() -> None: with ( patch( - "flocks.session.session_loop.Message.get", + "flocks.session.runtime.session_turn.Message.get", AsyncMock(return_value=user), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", AsyncMock(side_effect=["implement hooks", "implementation complete"]), ), patch( - "flocks.session.session_loop.Message.create", + "flocks.session.runtime.session_turn.Message.create", create_message, ), patch( @@ -171,12 +180,12 @@ async def test_turn_finish_block_creates_synthetic_continuation() -> None: AsyncMock(return_value=SimpleNamespace(steps=10)), ), ): - continued = await SessionLoop._run_turn_finish_hook( + decision = await DEFAULT_CONTINUATION_POLICY.run_turn_finish( ctx, - callbacks, user, assistant, ) + continued = decision.should_continue assert continued is True assert ctx.stop_hook_active is True @@ -206,25 +215,26 @@ async def test_queued_prompt_arriving_during_turn_finish_wins() -> None: finish="stop", ) queued_user = SimpleNamespace(id="msg_003", agent="rex", role="user") - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock( return_value=[user, assistant, queued_user], ) ) callbacks = LoopCallbacks(event_publish_callback=AsyncMock()) + ctx.callbacks = callbacks create_message = AsyncMock() with ( patch( - "flocks.session.session_loop.Message.get", + "flocks.session.runtime.session_turn.Message.get", AsyncMock(return_value=user), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", AsyncMock(side_effect=["prompt", "response"]), ), patch( - "flocks.session.session_loop.Message.create", + "flocks.session.runtime.session_turn.Message.create", create_message, ), patch( @@ -238,12 +248,12 @@ async def test_queued_prompt_arriving_during_turn_finish_wins() -> None: ), ), ): - continued = await SessionLoop._run_turn_finish_hook( + decision = await DEFAULT_CONTINUATION_POLICY.run_turn_finish( ctx, - callbacks, user, assistant, ) + continued = decision.should_continue assert continued is True create_message.assert_not_awaited() @@ -253,6 +263,74 @@ async def test_queued_prompt_arriving_during_turn_finish_wins() -> None: assert payload["queuedUserMessageID"] == queued_user.id +@pytest.mark.asyncio +async def test_real_user_arriving_during_goal_evaluation_wins() -> None: + ctx = _loop_context("ses_goal_queue_race") + user = SimpleNamespace( + id="msg_001", + agent="rex", + role="user", + model={"providerID": "test-provider", "modelID": "test-model"}, + provider="test-provider", + ) + assistant = SimpleNamespace( + id="msg_002", + agent="rex", + role="assistant", + finish="stop", + ) + queued_user = SimpleNamespace( + id="msg_003", + agent="rex", + role="user", + ) + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user, assistant], + [user, assistant, queued_user], + ] + ) + ) + ctx.callbacks = LoopCallbacks(event_publish_callback=AsyncMock()) + create_message = AsyncMock() + outcome = SimpleNamespace( + state=SimpleNamespace(metadata={"last_user": user}), + last_message=assistant, + ) + + with ( + patch( + "flocks.session.runtime.continuation_policy.Message.get_text_content", + AsyncMock(return_value="one failure remains"), + ), + patch( + "flocks.session.runtime.continuation_policy.Message.create", + create_message, + ), + patch( + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", + AsyncMock( + return_value=GoalDecision( + status="active", + verdict="continue", + should_continue=True, + continuation_prompt="continue fixing failures", + ) + ), + ), + patch( + "flocks.agent.registry.Agent.get", + AsyncMock(return_value=SimpleNamespace(steps=10)), + ), + ): + decision = await DEFAULT_CONTINUATION_POLICY.resolve(ctx, outcome) + + assert decision.reason == "queued_message" + assert decision.messages == (queued_user,) + create_message.assert_not_awaited() + + @pytest.mark.asyncio async def test_turn_finish_block_is_ignored_at_agent_step_limit() -> None: ctx = _loop_context("ses_turn_finish_limit") @@ -269,15 +347,15 @@ async def test_turn_finish_block_is_ignored_at_agent_step_limit() -> None: with ( patch( - "flocks.session.session_loop.Message.get", + "flocks.session.runtime.session_turn.Message.get", AsyncMock(return_value=user), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", AsyncMock(side_effect=["prompt", "response"]), ), patch( - "flocks.session.session_loop.Message.create", + "flocks.session.runtime.session_turn.Message.create", create_message, ), patch( @@ -295,12 +373,12 @@ async def test_turn_finish_block_is_ignored_at_agent_step_limit() -> None: AsyncMock(return_value=SimpleNamespace(steps=3)), ), ): - continued = await SessionLoop._run_turn_finish_hook( + decision = await DEFAULT_CONTINUATION_POLICY.run_turn_finish( ctx, - LoopCallbacks(), user, assistant, ) + continued = decision.should_continue assert continued is False create_message.assert_not_awaited() @@ -328,7 +406,7 @@ async def test_turn_finish_runs_only_after_persisted_stop() -> None: ctx = _loop_context("ses_turn_finish_integration") user = _message("msg_001", "user") assistant = _message("msg_002", "assistant", finish="stop") - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock( side_effect=[ [user], @@ -336,23 +414,23 @@ async def test_turn_finish_runs_only_after_persisted_stop() -> None: ] ) ) - run_turn_finish = AsyncMock(return_value=False) + run_turn_finish = AsyncMock(return_value=ContinuationDecision()) with ( patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", AsyncMock(return_value="final response"), ), patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.GoalManager.evaluate_after_turn", + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", AsyncMock( return_value=GoalDecision( status="inactive", @@ -361,15 +439,16 @@ async def test_turn_finish_runs_only_after_persisted_stop() -> None: ), ), patch( - "flocks.session.session_loop.SessionLoop._run_user_prompt_submit_hook", + "flocks.session.runtime.continuation_policy.ContinuationPolicy.run_user_prompt_submit", AsyncMock(), ), - patch( - "flocks.session.session_loop.SessionLoop._run_turn_finish_hook", + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_finish", run_turn_finish, ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(return_value=StepResult(action="stop")), ), patch( @@ -377,16 +456,15 @@ async def test_turn_finish_runs_only_after_persisted_stop() -> None: MagicMock(return_value=None), ), patch( - "flocks.session.session_loop.fire_and_forget", + "flocks.session.runtime.session_turn.fire_and_forget", MagicMock(), ), ): - result = await SessionLoop._run_loop(ctx, LoopCallbacks()) + result = await run_logical_turns(ctx, LoopCallbacks()) assert result.action == "stop" run_turn_finish.assert_awaited_once_with( ctx, - ANY, user, assistant, ) @@ -407,7 +485,7 @@ async def test_turn_finish_skips_errors_and_tool_calls( ctx = _loop_context(f"ses_turn_finish_{assistant_finish}") user = _message("msg_001", "user") assistant = _message("msg_002", "assistant", finish=assistant_finish) - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock( side_effect=[ [user], @@ -415,7 +493,7 @@ async def test_turn_finish_skips_errors_and_tool_calls( ] ) ) - run_turn_finish = AsyncMock(return_value=False) + run_turn_finish = AsyncMock(return_value=ContinuationDecision()) async def process_step(*_args, **_kwargs): if step_result.action == "continue": @@ -424,23 +502,24 @@ async def process_step(*_args, **_kwargs): with ( patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.SessionLoop._run_user_prompt_submit_hook", + "flocks.session.runtime.continuation_policy.ContinuationPolicy.run_user_prompt_submit", AsyncMock(), ), - patch( - "flocks.session.session_loop.SessionLoop._run_turn_finish_hook", + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_finish", run_turn_finish, ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(side_effect=process_step), ), patch( @@ -448,11 +527,11 @@ async def process_step(*_args, **_kwargs): MagicMock(return_value=None), ), patch( - "flocks.session.session_loop.fire_and_forget", + "flocks.session.runtime.session_turn.fire_and_forget", MagicMock(), ), ): - result = await SessionLoop._run_loop(ctx, LoopCallbacks()) + result = await run_logical_turns(ctx, LoopCallbacks()) assert result.action == "stop" run_turn_finish.assert_not_awaited() @@ -464,7 +543,7 @@ async def test_queued_user_message_takes_priority_over_turn_finish() -> None: user = _message("msg_001", "user") assistant = _message("msg_002", "assistant", finish="stop") queued_user = _message("msg_003", "user") - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock( side_effect=[ [user], @@ -472,7 +551,7 @@ async def test_queued_user_message_takes_priority_over_turn_finish() -> None: ] ) ) - run_turn_finish = AsyncMock(return_value=False) + run_turn_finish = AsyncMock(return_value=ContinuationDecision()) async def process_step(*_args, **_kwargs): ctx.signal_abort() @@ -480,23 +559,24 @@ async def process_step(*_args, **_kwargs): with ( patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.SessionLoop._run_user_prompt_submit_hook", + "flocks.session.runtime.continuation_policy.ContinuationPolicy.run_user_prompt_submit", AsyncMock(), ), - patch( - "flocks.session.session_loop.SessionLoop._run_turn_finish_hook", + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_finish", run_turn_finish, ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(side_effect=process_step), ), patch( @@ -504,11 +584,11 @@ async def process_step(*_args, **_kwargs): MagicMock(return_value=None), ), patch( - "flocks.session.session_loop.fire_and_forget", + "flocks.session.runtime.session_turn.fire_and_forget", MagicMock(), ), ): - await SessionLoop._run_loop(ctx, LoopCallbacks()) + await run_logical_turns(ctx, LoopCallbacks()) run_turn_finish.assert_not_awaited() @@ -519,7 +599,7 @@ async def test_goal_continuation_takes_priority_over_turn_finish() -> None: user = _message("msg_001", "user") assistant = _message("msg_002", "assistant", finish="stop") goal_user = _message("msg_003", "user") - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock( side_effect=[ [user], @@ -527,7 +607,7 @@ async def test_goal_continuation_takes_priority_over_turn_finish() -> None: ] ) ) - run_turn_finish = AsyncMock(return_value=False) + run_turn_finish = AsyncMock(return_value=ContinuationDecision()) async def process_step(*_args, **_kwargs): ctx.signal_abort() @@ -535,23 +615,23 @@ async def process_step(*_args, **_kwargs): with ( patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", AsyncMock(return_value="not done"), ), patch( - "flocks.session.session_loop.Message.create", + "flocks.session.runtime.session_turn.Message.create", AsyncMock(return_value=goal_user), ), patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.GoalManager.evaluate_after_turn", + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", AsyncMock( return_value=GoalDecision( status="active", @@ -562,15 +642,16 @@ async def process_step(*_args, **_kwargs): ), ), patch( - "flocks.session.session_loop.SessionLoop._run_user_prompt_submit_hook", + "flocks.session.runtime.continuation_policy.ContinuationPolicy.run_user_prompt_submit", AsyncMock(), ), - patch( - "flocks.session.session_loop.SessionLoop._run_turn_finish_hook", + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_finish", run_turn_finish, ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(side_effect=process_step), ), patch( @@ -578,11 +659,11 @@ async def process_step(*_args, **_kwargs): MagicMock(return_value=None), ), patch( - "flocks.session.session_loop.fire_and_forget", + "flocks.session.runtime.session_turn.fire_and_forget", MagicMock(), ), ): - await SessionLoop._run_loop(ctx, LoopCallbacks()) + await run_logical_turns(ctx, LoopCallbacks()) run_turn_finish.assert_not_awaited() @@ -591,40 +672,45 @@ async def process_step(*_args, **_kwargs): async def test_abort_does_not_trigger_turn_finish() -> None: ctx = _loop_context("ses_turn_finish_abort") user = _message("msg_001", "user") - ctx.session_ctx = SimpleNamespace(get_messages=AsyncMock(return_value=[user])) - run_turn_finish = AsyncMock(return_value=False) + ctx.session_store = SimpleNamespace(get_messages=AsyncMock(return_value=[user])) + run_turn_finish = AsyncMock(return_value=ContinuationDecision()) + + async def cancel_for_user_abort(*_args, **_kwargs): + ctx.signal_abort() + raise asyncio.CancelledError with ( patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.SessionLoop._run_user_prompt_submit_hook", + "flocks.session.runtime.continuation_policy.ContinuationPolicy.run_user_prompt_submit", AsyncMock(), ), - patch( - "flocks.session.session_loop.SessionLoop._run_turn_finish_hook", + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_finish", run_turn_finish, ), patch( - "flocks.session.runner.SessionRunner._process_step", - AsyncMock(side_effect=asyncio.CancelledError()), + "flocks.session.runtime.step_engine.StepEngine._process_step", + AsyncMock(side_effect=cancel_for_user_abort), ), patch( "flocks.session.lifecycle.title.SessionTitle.ensure_title", MagicMock(return_value=None), ), patch( - "flocks.session.session_loop.fire_and_forget", + "flocks.session.runtime.session_turn.fire_and_forget", MagicMock(), ), ): - await SessionLoop._run_loop(ctx, LoopCallbacks()) + await run_logical_turns(ctx, LoopCallbacks()) run_turn_finish.assert_not_awaited() @@ -634,7 +720,7 @@ async def test_late_abort_after_step_completion_skips_turn_finish() -> None: ctx = _loop_context("ses_turn_finish_late_abort") user = _message("msg_001", "user") assistant = _message("msg_002", "assistant", finish="stop") - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock( side_effect=[ [user], @@ -642,26 +728,26 @@ async def test_late_abort_after_step_completion_skips_turn_finish() -> None: ] ) ) - run_turn_finish = AsyncMock(return_value=False) + run_turn_finish = AsyncMock(return_value=ContinuationDecision()) async def abort_after_step(_step: int) -> None: ctx.abort_event.set() with ( patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", AsyncMock(return_value="final response"), ), patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.GoalManager.evaluate_after_turn", + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", AsyncMock( return_value=GoalDecision( status="inactive", @@ -670,15 +756,16 @@ async def abort_after_step(_step: int) -> None: ), ), patch( - "flocks.session.session_loop.SessionLoop._run_user_prompt_submit_hook", + "flocks.session.runtime.continuation_policy.ContinuationPolicy.run_user_prompt_submit", AsyncMock(), ), - patch( - "flocks.session.session_loop.SessionLoop._run_turn_finish_hook", + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_finish", run_turn_finish, ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(return_value=StepResult(action="stop")), ), patch( @@ -686,11 +773,11 @@ async def abort_after_step(_step: int) -> None: MagicMock(return_value=None), ), patch( - "flocks.session.session_loop.fire_and_forget", + "flocks.session.runtime.session_turn.fire_and_forget", MagicMock(), ), ): - result = await SessionLoop._run_loop( + result = await run_logical_turns( ctx, LoopCallbacks(on_step_end=abort_after_step), ) diff --git a/tests/session/test_runner_chunk_handling.py b/tests/session/test_runner_chunk_handling.py index 61c22aa2f..fc3a2a677 100644 --- a/tests/session/test_runner_chunk_handling.py +++ b/tests/session/test_runner_chunk_handling.py @@ -1,6 +1,6 @@ """ Regression tests for the chunk-handling logic in -``SessionRunner._call_llm`` (Issue #1 of PR review for Gemini 3 support). +``StepEngine._call_llm`` (Issue #1 of PR review for Gemini 3 support). The previous implementation treated any ``StreamChunk`` carrying ``reasoning`` as reasoning-only and immediately ``continue``d, silently dropping ``delta`` / @@ -8,10 +8,9 @@ fixed loop consumes all three event types out of a single mixed chunk and correctly opens / closes the reasoning block around interleaved text. -We exercise the loop in isolation by replicating the exact runner code so the -test pins the contract; the same loop is used in -``flocks/session/runner.py``. Drift is unlikely because the loop is small and -documented, but a follow-up could refactor the runner to call this helper +We exercise the loop in isolation by replicating the exact step-engine code so +the test pins the contract. Drift is unlikely because the loop is small and +documented, but a follow-up could refactor the engine to call this helper directly. """ @@ -20,9 +19,11 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Optional +import pytest + # --------------------------------------------------------------------------- -# Minimal stand-ins for runner imports so the test stays self-contained. +# Minimal stand-ins for step-engine imports so the test stays self-contained. # --------------------------------------------------------------------------- @@ -113,7 +114,7 @@ async def feed_chunk(self, tc): # --------------------------------------------------------------------------- # The function under test: a faithful copy of the consumer loop in -# SessionRunner._call_llm (kept in sync via comments + cross-references). +# StepEngine._call_llm (kept in sync via comments + cross-references). # --------------------------------------------------------------------------- @@ -213,9 +214,6 @@ async def consume_chunks(chunks, processor, tool_accumulator) -> Dict[str, int]: # --------------------------------------------------------------------------- -import pytest - - class TestBundledChunks: """Bundled (reasoning + text + tool_calls) chunks must not lose data.""" diff --git a/tests/session/test_runner_device_hint.py b/tests/session/test_runner_device_hint.py index 01773730f..3e53f7b79 100644 --- a/tests/session/test_runner_device_hint.py +++ b/tests/session/test_runner_device_hint.py @@ -3,7 +3,7 @@ import pytest -from flocks.session.runner import SessionRunner +from flocks.session.runtime.step_engine import StepEngine from flocks.tool.registry import ToolCategory, ToolInfo @@ -26,7 +26,7 @@ async def test_device_asset_hint_stays_short_and_strategy_only() -> None: ]), ) monkeypatch.setattr( - "flocks.session.runner.ToolRegistry.list_tools", + "flocks.session.runtime.step_engine.ToolRegistry.list_tools", lambda: [ ToolInfo( name="tdp_event_list", @@ -49,8 +49,8 @@ async def test_device_asset_hint_stays_short_and_strategy_only() -> None: ], ) - runner = SessionRunner.__new__(SessionRunner) - hint = await SessionRunner._build_device_asset_hint(runner) + runner = StepEngine.__new__(StepEngine) + hint = await StepEngine._build_device_asset_hint(runner) monkeypatch.undo() assert hint is not None diff --git a/tests/session/test_runner_langfuse_payloads.py b/tests/session/test_runner_langfuse_payloads.py index 129ee8b3a..697d5cefc 100644 --- a/tests/session/test_runner_langfuse_payloads.py +++ b/tests/session/test_runner_langfuse_payloads.py @@ -1,5 +1,5 @@ from flocks.provider.provider import ChatMessage -from flocks.session.runner import SessionRunner, ToolCall +from flocks.session.runtime.step_engine import StepEngine, ToolCall def test_build_langfuse_request_payload_keeps_full_messages_and_system_prompt() -> None: @@ -31,7 +31,7 @@ def test_build_langfuse_request_payload_keeps_full_messages_and_system_prompt() ), ] - payload = SessionRunner._build_langfuse_request_payload( + payload = StepEngine._build_langfuse_request_payload( step=3, messages=messages, request_tools=tools, @@ -60,7 +60,7 @@ def test_build_langfuse_response_payload_keeps_full_content_reasoning_and_tool_a ) ] - payload = SessionRunner._build_langfuse_response_payload( + payload = StepEngine._build_langfuse_response_payload( action="continue", content=full_content, reasoning=full_reasoning, diff --git a/tests/session/test_runner_llm_hook_payloads.py b/tests/session/test_runner_llm_hook_payloads.py index 1b4f8f201..27f0e972e 100644 --- a/tests/session/test_runner_llm_hook_payloads.py +++ b/tests/session/test_runner_llm_hook_payloads.py @@ -5,10 +5,10 @@ from flocks.agent.agent import AgentInfo from flocks.config.config import Config, ConfigInfo -from flocks.hooks.pipeline import HookPipeline +from flocks.hooks.pipeline import HookPipeline, HookStage from flocks.provider.provider import ChatMessage, StreamChunk from flocks.session.message import Message, MessageRole -from flocks.session.runner import SessionRunner +from flocks.session.runtime.step_engine import StepEngine from flocks.session.session import Session @@ -41,7 +41,7 @@ async def _run_call_llm_with_hooks( agent="rex", ) - runner = SessionRunner( + runner = StepEngine( session=session, provider_id="test-provider", model_id="test-model", @@ -113,6 +113,8 @@ async def test_call_llm_uses_full_hook_payloads_by_default( assert before_input["request"]["tools"][0]["function"]["name"] == "read" assert before_input["request"]["messageCount"] == 2 assert before_input["request"]["toolCount"] == 1 + assert before_input["request"]["providerID"] == "test-provider" + assert before_input["request"]["modelID"] == "test-model" assert "messageSummaries" not in before_input["request"] assert "toolSummaries" not in before_input["request"] @@ -126,3 +128,72 @@ async def test_call_llm_uses_full_hook_payloads_by_default( "model", } assert "request" not in after_input + + +@pytest.mark.asyncio +async def test_before_model_hook_changes_the_real_provider_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = await Session.create( + project_id="test_project_hook_request", + directory="/test/hooks", + ) + user_msg = await Message.create( + session_id=session.id, + role=MessageRole.USER, + content="hello", + ) + assistant_msg = await Message.create( + session_id=session.id, + role=MessageRole.ASSISTANT, + content="", + parentID=user_msg.id, + modelID="test-model", + providerID="test-provider", + agent="rex", + ) + runner = StepEngine( + session=session, + provider_id="test-provider", + model_id="test-model", + agent_name="rex", + ) + provider_calls: list[dict] = [] + + class ProviderStub: + async def chat_stream(self, **kwargs): # noqa: ANN003 + provider_calls.append(kwargs) + yield StreamChunk(delta="modified", finish_reason="stop") + + async def before_model(input_data, output_data=None): # noqa: ANN001, ANN202 + del output_data + modified = dict(input_data["request"]) + modified["messages"] = [ + {"role": "user", "content": "rewritten by hook"}, + ] + modified["tools"] = [] + modified["providerOptions"] = {"temperature": 0.7} + return SimpleNamespace( + input=input_data, + output={"request": modified}, + ) + + async def has_handlers(stage, _metadata): # noqa: ANN001, ANN202 + return stage == HookStage.LLM_BEFORE + + monkeypatch.setattr(HookPipeline, "has_stage_handlers", has_handlers) + monkeypatch.setattr(HookPipeline, "run_llm_before", before_model) + + result = await runner._call_llm( + provider=ProviderStub(), + messages=[ChatMessage(role="user", content="original")], + tools=[{"type": "function", "function": {"name": "read"}}], + agent=AgentInfo(name="rex"), + assistant_msg=assistant_msg, + ) + + assert result.content == "modified" + assert len(provider_calls) == 1 + assert provider_calls[0]["messages"][0].content == "rewritten by hook" + assert provider_calls[0]["tools"] is None + assert provider_calls[0]["temperature"] == 0.7 diff --git a/tests/session/test_runner_llm_hooks.py b/tests/session/test_runner_llm_hooks.py index a2c7ba864..95c93d73d 100644 --- a/tests/session/test_runner_llm_hooks.py +++ b/tests/session/test_runner_llm_hooks.py @@ -1,4 +1,4 @@ -"""Tests for LLM lifecycle hooks in SessionRunner and HookPipeline.""" +"""Tests for LLM lifecycle hooks in StepEngine and HookPipeline.""" from __future__ import annotations @@ -8,11 +8,11 @@ import pytest -import flocks.session.runner as runner_mod +import flocks.session.runtime.step_engine as runner_mod from flocks.hooks.pipeline import HookBase, HookPipeline from flocks.provider.provider import ChatMessage from flocks.session.streaming.stream_processor import StreamProcessor -from flocks.session.runner import SessionRunner +from flocks.session.runtime.step_engine import StepEngine from flocks.session.session import SessionInfo from flocks.tool.registry import ToolResult @@ -27,8 +27,8 @@ def _make_session(session_id: str = "ses_runner_llm_hooks") -> SessionInfo: ) -def _make_runner(session_id: str = "ses_runner_llm_hooks") -> SessionRunner: - return SessionRunner( +def _make_runner(session_id: str = "ses_runner_llm_hooks") -> StepEngine: + return StepEngine( session=_make_session(session_id), provider_id="anthropic", model_id="claude-sonnet", @@ -179,7 +179,7 @@ async def _after(payload, result): AsyncMock(side_effect=_after), ) monkeypatch.setattr( - runner_mod.SessionRunner, + runner_mod.StepEngine, "_end_observability", staticmethod(lambda *args, **kwargs: None), ) @@ -274,7 +274,7 @@ async def _after(payload, result): AsyncMock(side_effect=_after), ) monkeypatch.setattr( - runner_mod.SessionRunner, + runner_mod.StepEngine, "_end_observability", staticmethod(lambda *args, **kwargs: None), ) diff --git a/tests/session/test_runner_provider_version.py b/tests/session/test_runner_provider_version.py index 0131c5038..fed64ba9b 100644 --- a/tests/session/test_runner_provider_version.py +++ b/tests/session/test_runner_provider_version.py @@ -1,5 +1,5 @@ """ -Tests for ``flocks.session.runner._annotate_with_provider_version``. +Tests for ``StepEngine`` provider-version annotations. Ensures that when a tool's ``ToolInfo`` carries a ``provider_version`` (sourced from ``_provider.yaml``), the description handed to the LLM in the function @@ -15,7 +15,7 @@ from dataclasses import dataclass from typing import Optional -from flocks.session.runner import _annotate_with_provider_version +from flocks.session.runtime.step_engine import _annotate_with_provider_version @dataclass diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index 723f0e874..5bfbb666f 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -1,13 +1,13 @@ """ -Tests for SessionRunner internals in flocks/session/runner.py +Tests for StepEngine internals in flocks/session/runtime/step_engine.py. Covers: - _agent_declares_tool(): tool declaration filtering - _exception_to_error_dict(): exception to error dict conversion - _build_callable_tool_schema(): excluded tools filter -- RunnerCallbacks dataclass +- LoopCallbacks dataclass - ToolCall / StepResult dataclasses -- SessionRunner construction and abort behavior (from existing tests) +- StepEngine construction and abort behavior (from existing tests) """ import httpcore @@ -16,7 +16,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch, AsyncMock -import flocks.session.runner as runner_mod +import flocks.session.runtime.step_engine as runner_mod from flocks.provider.sdk.anthropic import AnthropicProvider from flocks.session.message import ( Message, @@ -27,12 +27,12 @@ ToolStateRunning, UserMessageInfo, ) -from flocks.session.runner import ( - RunnerCallbacks, - SessionRunner, +from flocks.session.runtime.step_engine import ( + StepEngine, StepResult, ToolCall, ) +from flocks.session.runtime.session_turn import LoopCallbacks from flocks.session.prompt import SessionPrompt from flocks.session.core.defaults import DEFAULT_MAX_TOOL_STEPS from flocks.session.session import Session, SessionInfo @@ -62,7 +62,7 @@ def _make_agent(name="rex", tools=None): def _make_runner(session_id="ses_runner_test"): session = _make_session(session_id) - return SessionRunner(session=session) + return StepEngine(session=session) def _make_callable_schema_result(*tool_names): @@ -167,12 +167,12 @@ def test_resets_after_text_response(self): # --------------------------------------------------------------------------- -# RunnerCallbacks dataclass +# LoopCallbacks dataclass # --------------------------------------------------------------------------- class TestRunnerCallbacks: def test_all_defaults_none(self): - cb = RunnerCallbacks() + cb = LoopCallbacks() assert cb.on_step_start is None assert cb.on_step_end is None assert cb.on_text_delta is None @@ -187,7 +187,7 @@ def test_set_callbacks(self): async def my_callback(x): pass - cb = RunnerCallbacks(on_text_delta=my_callback, on_error=my_callback) + cb = LoopCallbacks(on_text_delta=my_callback, on_error=my_callback) assert cb.on_text_delta is my_callback assert cb.on_error is my_callback assert cb.on_step_start is None @@ -382,7 +382,7 @@ async def test_excludes_invalid_tool(self): ) with patch( - "flocks.session.runner.ToolRegistry.list_tools", + "flocks.session.runtime.step_engine.ToolRegistry.list_tools", return_value=[invalid_tool, bash_tool], ): tools = await runner._build_callable_tool_schema(agent) @@ -447,7 +447,7 @@ async def test_excludes_noop_tool(self): ) with patch( - "flocks.session.runner.ToolRegistry.list_tools", + "flocks.session.runtime.step_engine.ToolRegistry.list_tools", return_value=[noop_tool, real_tool], ): tools = await runner._build_callable_tool_schema(agent) @@ -469,7 +469,7 @@ async def test_disabled_tools_excluded(self): ) with patch( - "flocks.session.runner.ToolRegistry.list_tools", + "flocks.session.runtime.step_engine.ToolRegistry.list_tools", return_value=[disabled_tool], ): tools = await runner._build_callable_tool_schema(agent) @@ -490,7 +490,7 @@ async def test_tool_format_is_function_type(self): ) with patch( - "flocks.session.runner.SessionRunner._list_callable_tool_infos_for_turn", + "flocks.session.runtime.step_engine.StepEngine._list_callable_tool_infos_for_turn", AsyncMock(return_value=([tool_info], {"enabledToolCount": 1})), ): tools = await runner._build_callable_tool_schema(agent) @@ -524,7 +524,7 @@ async def test_build_tools_reflects_latest_selector_result(self): ([tool_v1], {"enabledToolCount": 3}), ([tool_v2], {"enabledToolCount": 3}), ]) - with patch.object(SessionRunner, "_list_callable_tool_infos_for_turn", selector_mock): + with patch.object(StepEngine, "_list_callable_tool_infos_for_turn", selector_mock): tools1 = await runner._build_callable_tool_schema(agent, []) tools2 = await runner._build_callable_tool_schema(agent, []) @@ -549,8 +549,8 @@ def test_prompt_tool_names_from_schema_uses_loaded_tool_names(self): async def test_build_tools_calls_selector_for_each_runner_instance(self): shared_cache = {} session = _make_session("ses_tools_runner_instances") - runner1 = SessionRunner(session=session, static_cache=shared_cache) - runner2 = SessionRunner(session=session, static_cache=shared_cache) + runner1 = StepEngine(session=session, static_cache=shared_cache) + runner2 = StepEngine(session=session, static_cache=shared_cache) agent = _make_agent(name="rex") selected_tool = ToolInfo( @@ -562,7 +562,7 @@ async def test_build_tools_calls_selector_for_each_runner_instance(self): ) selector_mock = AsyncMock(return_value=([selected_tool], {"enabledToolCount": 3})) - with patch.object(SessionRunner, "_list_callable_tool_infos_for_turn", selector_mock): + with patch.object(StepEngine, "_list_callable_tool_infos_for_turn", selector_mock): tools1 = await runner1._build_callable_tool_schema(agent, []) tools2 = await runner2._build_callable_tool_schema(agent, []) @@ -585,7 +585,7 @@ async def test_build_tools_uses_selector_results_and_emits_event(self): ) with patch.object( - SessionRunner, + StepEngine, "_list_callable_tool_infos_for_turn", AsyncMock(return_value=( [selected_tool], @@ -612,7 +612,7 @@ async def test_build_tools_refreshes_skill_description_from_enabled_skills(self) ) with patch.object( - SessionRunner, + StepEngine, "_list_callable_tool_infos_for_turn", AsyncMock(return_value=([skill_tool], {"enabledToolCount": 3})), ), patch( @@ -633,8 +633,8 @@ class TestBuildSystemPrompts: async def test_build_system_prompts_reuses_loop_static_cache(self): shared_cache = {} session = _make_session("ses_prompts_cache") - runner1 = SessionRunner(session=session, static_cache=shared_cache) - runner2 = SessionRunner(session=session, static_cache=shared_cache) + runner1 = StepEngine(session=session, static_cache=shared_cache) + runner2 = StepEngine(session=session, static_cache=shared_cache) agent = _make_agent(name="rex") agent.prompt = "agent prompt" @@ -693,7 +693,7 @@ async def test_build_system_prompts_reuses_loop_static_cache(self): @pytest.mark.asyncio async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(self): session = _make_session("ses_prompts_order") - runner = SessionRunner(session=session) + runner = StepEngine(session=session) agent = _make_agent(name="rex") agent.prompt = "agent prompt" memory_bootstrap_data = { @@ -756,7 +756,7 @@ async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(sel async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): shared_cache = {} session = _make_session("ses_prompts_revision") - runner = SessionRunner(session=session, static_cache=shared_cache) + runner = StepEngine(session=session, static_cache=shared_cache) agent = _make_agent(name="rex") agent.prompt = "agent prompt v1" @@ -822,7 +822,7 @@ async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): async def test_build_system_prompts_reuses_static_device_hint_cache(self): shared_cache = {} session = _make_session("ses_prompts_static_device_hint") - runner = SessionRunner(session=session, static_cache=shared_cache) + runner = StepEngine(session=session, static_cache=shared_cache) agent = _make_agent(name="rex") agent.prompt = "agent prompt" @@ -881,7 +881,7 @@ async def test_build_system_prompts_reuses_static_device_hint_cache(self): async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): shared_cache = {} session = _make_session("ses_prompts_device_revision") - runner = SessionRunner(session=session, static_cache=shared_cache) + runner = StepEngine(session=session, static_cache=shared_cache) agent = _make_agent(name="rex") agent.prompt = "agent prompt" @@ -942,7 +942,7 @@ async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): shared_cache = {} session = _make_session("ses_prompts_agent_prompt") - runner = SessionRunner(session=session, static_cache=shared_cache) + runner = StepEngine(session=session, static_cache=shared_cache) agent = _make_agent(name="rex") agent.prompt = "agent prompt v1" @@ -988,7 +988,7 @@ async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): @pytest.mark.asyncio async def test_build_system_prompts_includes_filesystem_memory_guidance(self): session = _make_session("ses_prompts_memory_guidance") - runner = SessionRunner( + runner = StepEngine( session=session, memory_bootstrap_data={ "instructions": "memory guidance", @@ -1032,7 +1032,7 @@ async def test_build_system_prompts_includes_filesystem_memory_guidance(self): @pytest.mark.asyncio async def test_build_system_prompts_does_not_add_bash_guidance_prompt_when_bash_loaded(self): session = _make_session("ses_prompts_no_bash_guidance") - runner = SessionRunner(session=session) + runner = StepEngine(session=session) agent = _make_agent(name="rex") agent.prompt = "agent prompt" @@ -1057,7 +1057,7 @@ async def test_build_system_prompts_does_not_add_bash_guidance_prompt_when_bash_ @pytest.mark.asyncio async def test_build_system_prompts_skips_memory_guidance_without_management_tools(self): session = _make_session("ses_prompts_no_memory_guidance") - runner = SessionRunner( + runner = StepEngine( session=session, memory_bootstrap_data={ "instructions": "memory guidance", @@ -1093,7 +1093,7 @@ async def test_build_system_prompts_skips_memory_guidance_without_management_too async def test_filesystem_memory_guidance_depends_on_tool_names(self): shared_cache = {} session = _make_session("ses_prompts_tool_names") - runner = SessionRunner( + runner = StepEngine( session=session, static_cache=shared_cache, memory_bootstrap_data={ @@ -1157,7 +1157,7 @@ def test_build_tool_catalog_prompt_for_rex(self): agent.mode = "primary" with patch( - "flocks.session.runner.SessionRunner._list_catalog_tool_infos", + "flocks.session.runtime.step_engine.StepEngine._list_catalog_tool_infos", return_value=[ToolInfo( name="plugin_memory", description="Access project memory", @@ -1169,7 +1169,7 @@ def test_build_tool_catalog_prompt_for_rex(self): "flocks.agent.toolset.get_all_enabled_builtin_tool_names", return_value=["read", "bash"], ), patch( - "flocks.session.runner.get_always_load_tool_names", + "flocks.session.runtime.step_engine.get_always_load_tool_names", return_value={"question", "tool_search"}, ), patch( "flocks.command.direct.format_tools_catalog_summary", @@ -1204,13 +1204,13 @@ def test_build_tool_catalog_prompt_for_rex_excludes_builtin_and_always_load_tool ] with patch( - "flocks.session.runner.SessionRunner._list_catalog_tool_infos", + "flocks.session.runtime.step_engine.StepEngine._list_catalog_tool_infos", return_value=catalog_tools, ), patch( "flocks.agent.toolset.get_all_enabled_builtin_tool_names", return_value=["bash", "read"], ), patch( - "flocks.session.runner.get_always_load_tool_names", + "flocks.session.runtime.step_engine.get_always_load_tool_names", return_value={"question", "tool_search"}, ), patch( "flocks.command.direct.format_tools_catalog_summary", @@ -1248,13 +1248,13 @@ def test_build_tool_catalog_prompt_for_rex_excludes_device_tools(self): ] with patch( - "flocks.session.runner.SessionRunner._list_catalog_tool_infos", + "flocks.session.runtime.step_engine.StepEngine._list_catalog_tool_infos", return_value=catalog_tools, ), patch( "flocks.agent.toolset.get_all_enabled_builtin_tool_names", return_value=["bash", "read"], ), patch( - "flocks.session.runner.get_always_load_tool_names", + "flocks.session.runtime.step_engine.get_always_load_tool_names", return_value={"question", "tool_search"}, ), patch( "flocks.command.direct.format_tools_catalog_summary", @@ -1288,7 +1288,7 @@ def test_list_catalog_tool_infos_returns_full_catalog_for_rex(self): ) with patch( - "flocks.session.runner.list_tool_catalog_infos", + "flocks.session.runtime.step_engine.list_tool_catalog_infos", return_value=[shell_tool, helper_tool], ): infos = runner._list_catalog_tool_infos(agent) @@ -1306,7 +1306,7 @@ def test_list_catalog_tool_infos_filters_subagent_boundaries(self): ToolInfo(name="websearch", description="Search web", category=ToolCategory.BROWSER, native=True, enabled=True), ] - with patch("flocks.session.runner.list_tool_catalog_infos", return_value=tool_infos): + with patch("flocks.session.runtime.step_engine.list_tool_catalog_infos", return_value=tool_infos): infos = runner._list_catalog_tool_infos(agent) assert [tool.name for tool in infos] == ["read"] @@ -1323,7 +1323,7 @@ def test_list_catalog_tool_infos_keeps_always_load_tools_for_subagent(self): ToolInfo(name="bash", description="Run commands", category=ToolCategory.CODE, native=True, enabled=True), ] - with patch("flocks.session.runner.list_tool_catalog_infos", return_value=tool_infos): + with patch("flocks.session.runtime.step_engine.list_tool_catalog_infos", return_value=tool_infos): infos = runner._list_catalog_tool_infos(agent) assert [tool.name for tool in infos] == ["read", "question", "tool_search"] @@ -1339,7 +1339,7 @@ def test_list_catalog_tool_infos_does_not_fall_back_to_full_catalog_when_tools_m ToolInfo(name="bash", description="Run commands", category=ToolCategory.CODE, native=True, enabled=True), ] - with patch("flocks.session.runner.list_tool_catalog_infos", return_value=tool_infos): + with patch("flocks.session.runtime.step_engine.list_tool_catalog_infos", return_value=tool_infos): infos = runner._list_catalog_tool_infos(agent) assert [tool.name for tool in infos] == ["question", "tool_search"] @@ -1348,7 +1348,7 @@ def test_list_catalog_tool_infos_does_not_fall_back_to_full_catalog_when_tools_m class TestMiniMaxTextToolMode: def test_disabled_for_custom_threatbook_minimax(self): session = _make_session("ses_minimax_mode") - runner = SessionRunner( + runner = StepEngine( session=session, provider_id="custom-threatbook-internal", model_id="minimax:MiniMax-M2.5", @@ -1357,7 +1357,7 @@ def test_disabled_for_custom_threatbook_minimax(self): def test_disabled_for_custom_tb_inner_minimax(self): session = _make_session("ses_minimax_mode_tb_inner") - runner = SessionRunner( + runner = StepEngine( session=session, provider_id="custom-tb-inner", model_id="minimax:MiniMax-M2.7", @@ -1366,7 +1366,7 @@ def test_disabled_for_custom_tb_inner_minimax(self): def test_disabled_for_threatbook_cn_llm_minimax(self): session = _make_session("ses_minimax_threatbook_cn_llm") - runner = SessionRunner( + runner = StepEngine( session=session, provider_id="threatbook-cn-llm", model_id="minimax-m2.7", @@ -1375,7 +1375,7 @@ def test_disabled_for_threatbook_cn_llm_minimax(self): def test_disabled_for_threatbook_cn_llm_minimax_case_insensitive(self): session = _make_session("ses_minimax_threatbook_cn_llm_case") - runner = SessionRunner( + runner = StepEngine( session=session, provider_id="ThreatBook-CN-LLM", model_id="MiniMax-M2.5", @@ -1386,7 +1386,7 @@ def test_disabled_for_threatbook_cn_llm_non_minimax(self): # Other models routed through the same gateway (e.g. qwen, GLM) keep # the standard OpenAI native function-calling path. session = _make_session("ses_threatbook_cn_llm_qwen") - runner = SessionRunner( + runner = StepEngine( session=session, provider_id="threatbook-cn-llm", model_id="qwen3.6-plus", @@ -1395,7 +1395,7 @@ def test_disabled_for_threatbook_cn_llm_non_minimax(self): def test_disabled_for_other_models(self): session = _make_session("ses_normal_mode") - runner = SessionRunner( + runner = StepEngine( session=session, provider_id="anthropic", model_id="claude-sonnet-4-5-20250929", @@ -1405,7 +1405,7 @@ def test_disabled_for_other_models(self): @pytest.mark.asyncio async def test_system_prompts_add_minimax_native_tool_guidance(self): session = _make_session("ses_minimax_prompt") - runner = SessionRunner( + runner = StepEngine( session=session, provider_id="custom-tb-inner", model_id="minimax:MiniMax-M2.5", @@ -1432,7 +1432,7 @@ async def test_system_prompts_add_minimax_native_tool_guidance(self): def test_build_text_tool_call_catalog_prompt(self): session = _make_session("ses_minimax_catalog") - runner = SessionRunner( + runner = StepEngine( session=session, provider_id="custom-threatbook-internal", model_id="minimax:MiniMax-M2.5", @@ -1465,7 +1465,7 @@ def test_build_text_tool_call_catalog_prompt(self): @pytest.mark.asyncio async def test_to_chat_messages_uses_structured_anthropic_system_blocks(monkeypatch): - runner = SessionRunner( + runner = StepEngine( session=_make_session("ses_anthropic_system_blocks"), provider_id="anthropic", model_id="claude-sonnet", @@ -1488,7 +1488,7 @@ async def test_to_chat_messages_uses_structured_anthropic_system_blocks(monkeypa @pytest.mark.asyncio async def test_to_chat_messages_keeps_joined_system_prompt_for_openai(monkeypatch): - runner = SessionRunner( + runner = StepEngine( session=_make_session("ses_openai_system_blocks"), provider_id="openai", model_id="gpt-5", @@ -1518,7 +1518,7 @@ async def test_to_chat_messages_invalidates_shared_cache_when_message_parts_chan role=MessageRole.ASSISTANT, content="starting", ) - runner = SessionRunner(session=session, static_cache={}) + runner = StepEngine(session=session, static_cache={}) first_messages = await runner._to_chat_messages([assistant_message], []) @@ -1566,7 +1566,7 @@ async def test_to_chat_messages_excludes_ignored_assistant_text(): modelID="command", ignored=True, ) - runner = SessionRunner(session=session, static_cache={}) + runner = StepEngine(session=session, static_cache={}) chat_messages = await runner._to_chat_messages([assistant_message], []) @@ -1584,7 +1584,7 @@ async def test_to_chat_messages_preserves_assistant_reasoning_for_replay(): role=MessageRole.ASSISTANT, content="", ) - runner = SessionRunner(session=session, static_cache={}) + runner = StepEngine(session=session, static_cache={}) await Message.add_part( session.id, @@ -1633,7 +1633,7 @@ async def test_to_chat_messages_restores_provider_reasoning_fields_from_metadata role=MessageRole.ASSISTANT, content="", ) - runner = SessionRunner(session=session, static_cache={}) + runner = StepEngine(session=session, static_cache={}) runner.provider_id = "alibaba" runner.model_id = "qwen3-max" @@ -1700,7 +1700,7 @@ async def test_to_chat_messages_restores_redacted_anthropic_thinking_blocks(monk role=MessageRole.ASSISTANT, content="", ) - runner = SessionRunner(session=session, static_cache={}) + runner = StepEngine(session=session, static_cache={}) runner.provider_id = "anthropic" runner.model_id = "claude-sonnet-4-6" @@ -1762,7 +1762,7 @@ async def test_to_chat_messages_restores_signed_anthropic_thinking_blocks(monkey role=MessageRole.ASSISTANT, content="", ) - runner = SessionRunner(session=session, static_cache={}) + runner = StepEngine(session=session, static_cache={}) runner.provider_id = "anthropic" runner.model_id = "claude-sonnet-4-6" @@ -1828,7 +1828,7 @@ async def test_to_chat_messages_restores_unsigned_anthropic_thinking_blocks(monk role=MessageRole.ASSISTANT, content="", ) - runner = SessionRunner(session=session, static_cache={}) + runner = StepEngine(session=session, static_cache={}) runner.provider_id = "anthropic" runner.model_id = "claude-sonnet-4-6" @@ -1892,7 +1892,7 @@ async def test_runner_history_round_trip_formats_anthropic_payload(monkeypatch): role=MessageRole.ASSISTANT, content="Done", ) - runner = SessionRunner(session=session, static_cache={}) + runner = StepEngine(session=session, static_cache={}) runner.provider_id = "anthropic" runner.model_id = "claude-sonnet-4-6" @@ -1957,7 +1957,7 @@ async def test_to_chat_messages_prefers_provider_specific_interleaved_resolution role=MessageRole.ASSISTANT, content="", ) - runner = SessionRunner(session=session, static_cache={}) + runner = StepEngine(session=session, static_cache={}) runner.provider_id = "deepseek" runner.model_id = "shared-model" @@ -2037,7 +2037,7 @@ async def test_to_chat_messages_keeps_reasoning_only_assistant_message(monkeypat role=MessageRole.ASSISTANT, content="", ) - runner = SessionRunner(session=session, static_cache={}) + runner = StepEngine(session=session, static_cache={}) runner.provider_id = "alibaba" runner.model_id = "qwen3-max" @@ -2112,7 +2112,7 @@ async def test_to_chat_messages_wraps_only_queued_user_messages(): content="What version is installed?", ) - runner = SessionRunner(session=session, static_cache={}) + runner = StepEngine(session=session, static_cache={}) runner._step = 3 runner._queued_user_message_ids = {queued_user.id} @@ -2193,7 +2193,7 @@ def test_provider_capability_key_includes_interleaved_policy(monkeypatch): runner.provider_id = "deepseek" runner.model_id = "deepseek-v4-pro" - monkeypatch.setattr(SessionRunner, "_model_supports_vision", lambda self: False) + monkeypatch.setattr(StepEngine, "_model_supports_vision", lambda self: False) monkeypatch.setattr( runner_mod.Provider, "resolve_model", @@ -2219,7 +2219,7 @@ def test_provider_capability_key_includes_interleaved_policy(monkeypatch): @pytest.mark.asyncio async def test_process_step_creates_assistant_message_with_provider_and_model(monkeypatch): runner = _make_runner("ses_runner_provider_model") - runner.callbacks = RunnerCallbacks( + runner.callbacks = LoopCallbacks( on_text_delta=AsyncMock(), on_error=AsyncMock(), ) @@ -2273,7 +2273,7 @@ async def test_process_step_invalidates_chat_cache_for_queued_messages(monkeypat "chat_messages": [{"role": "user", "content": "stale"}], } } - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) root_user = SimpleNamespace(id="msg_200", role="user") last_user = UserMessageInfo( @@ -2324,7 +2324,7 @@ async def fake_to_chat_messages(_messages, _system_prompts): # noqa: ANN001 @pytest.mark.asyncio async def test_process_step_limits_connection_error_retries(monkeypatch): runner = _make_runner("ses_runner_connection_error") - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) last_user = UserMessageInfo( id="msg_user_connection_error", @@ -2397,7 +2397,7 @@ async def fake_call_llm(*_args, **_kwargs): @pytest.mark.asyncio async def test_process_step_marks_aborted_llm_message_as_error(monkeypatch): runner = _make_runner("ses_runner_aborted_result") - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) last_user = UserMessageInfo( id="msg_user_aborted_result", @@ -2460,7 +2460,7 @@ async def fake_call_llm(*_args, **_kwargs): @pytest.mark.asyncio async def test_call_llm_skips_observability_when_langfuse_inactive(monkeypatch): runner = _make_runner("ses_runner_langfuse_inactive") - runner.callbacks = RunnerCallbacks() + runner.callbacks = LoopCallbacks() agent = SimpleNamespace(name="rex") assistant_msg = SimpleNamespace(id="msg_assistant_langfuse") @@ -2490,7 +2490,7 @@ async def test_call_llm_skips_observability_when_langfuse_inactive(monkeypatch): @pytest.mark.asyncio async def test_call_llm_skips_llm_hook_payload_preparation_without_handlers(monkeypatch): runner = _make_runner("ses_runner_no_llm_hooks") - runner.callbacks = RunnerCallbacks() + runner.callbacks = LoopCallbacks() class _ProviderStub: async def chat_stream(self, **kwargs): # noqa: ANN003 @@ -2554,7 +2554,7 @@ async def on_error(error): runner.provider_id = "missing-provider" runner.model_id = "missing-model" - runner.callbacks = RunnerCallbacks( + runner.callbacks = LoopCallbacks( on_error=on_error, event_publish_callback=publish_event, ) @@ -2608,7 +2608,7 @@ async def on_error(error): runner.provider_id = "unconfigured-provider" runner.model_id = "unconfigured-model" - runner.callbacks = RunnerCallbacks( + runner.callbacks = LoopCallbacks( on_error=on_error, event_publish_callback=publish_event, ) @@ -2661,12 +2661,12 @@ async def publish_event(event_name, payload): monkeypatch.setattr(runner_mod.Provider, "get", staticmethod(lambda _provider_id: EmptyProvider())) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) - monkeypatch.setattr(SessionRunner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr(StepEngine, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.SessionRetry, "sleep", AsyncMock(return_value=None)) runner.provider_id = "empty-provider" runner.model_id = "empty-model" - runner.callbacks = RunnerCallbacks(event_publish_callback=publish_event) + runner.callbacks = LoopCallbacks(event_publish_callback=publish_event) result = await runner._process_step(messages, user) messages_with_parts = await Message.list_with_parts(runner.session.id) @@ -2692,7 +2692,7 @@ async def publish_event(event_name, payload): @pytest.mark.asyncio async def test_process_step_uses_loaded_tool_schema_names_for_prompt_guidance(monkeypatch): runner = _make_runner("ses_runner_prompt_guidance_tool_names") - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) last_user = UserMessageInfo( id="msg_user_prompt_guidance", @@ -2743,7 +2743,7 @@ async def test_process_step_uses_loaded_tool_schema_names_for_prompt_guidance(mo @pytest.mark.asyncio async def test_process_step_records_usage_after_success(monkeypatch): runner = _make_runner("ses_runner_usage_success") - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) last_user = UserMessageInfo( id="msg_user_usage_success", @@ -2792,7 +2792,7 @@ async def test_process_step_records_usage_after_success(monkeypatch): @pytest.mark.asyncio async def test_process_step_passes_device_hint_factory_into_build_system_prompts(monkeypatch): runner = _make_runner("ses_runner_device_hint_order") - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) last_user = UserMessageInfo( id="msg_user_device_hint_order", @@ -2847,7 +2847,7 @@ async def test_process_step_empty_retry_records_usage_per_attempt(monkeypatch): """Each empty-response attempt records its own usage so that provider charges are not lost when the model returns tokens but no content.""" runner = _make_runner("ses_runner_usage_retry") - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) last_user = UserMessageInfo( id="msg_user_usage_retry", @@ -2907,7 +2907,7 @@ async def test_process_step_empty_retry_records_usage_per_attempt(monkeypatch): @pytest.mark.asyncio async def test_process_step_retries_empty_transport_exception(monkeypatch): runner = _make_runner("ses_runner_transport_retry") - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) last_user = UserMessageInfo( id="msg_user_transport_retry", @@ -2956,7 +2956,7 @@ async def test_process_step_retries_empty_transport_exception(monkeypatch): @pytest.mark.asyncio async def test_process_step_does_not_retry_after_tool_execution_started(monkeypatch): runner = _make_runner("ses_runner_tool_side_effect_no_retry") - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) last_user = UserMessageInfo( id="msg_user_tool_side_effect_no_retry", @@ -3010,7 +3010,7 @@ async def _call_llm(*_args, **_kwargs): @pytest.mark.asyncio async def test_process_step_uses_default_max_steps_when_agent_steps_missing(monkeypatch): runner = _make_runner("ses_runner_default_max_steps") - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) runner._step = DEFAULT_MAX_TOOL_STEPS last_user = UserMessageInfo( @@ -3048,7 +3048,7 @@ async def fake_call_llm(self, provider, messages, tools, agent, assistant_msg): captured["tools"] = tools return StepResult(action="stop", content="done") - monkeypatch.setattr(SessionRunner, "_call_llm", fake_call_llm) + monkeypatch.setattr(StepEngine, "_call_llm", fake_call_llm) result = await runner._process_step([last_user], last_user) @@ -3059,7 +3059,7 @@ async def fake_call_llm(self, provider, messages, tools, agent, assistant_msg): @pytest.mark.asyncio async def test_process_step_respects_explicit_agent_steps_over_default(monkeypatch): runner = _make_runner("ses_runner_explicit_max_steps") - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) runner._step = DEFAULT_MAX_TOOL_STEPS last_user = UserMessageInfo( @@ -3097,7 +3097,7 @@ async def fake_call_llm(self, provider, messages, tools, agent, assistant_msg): captured["tools"] = tools return StepResult(action="stop", content="done") - monkeypatch.setattr(SessionRunner, "_call_llm", fake_call_llm) + monkeypatch.setattr(StepEngine, "_call_llm", fake_call_llm) result = await runner._process_step([last_user], last_user) @@ -3140,7 +3140,7 @@ async def fake_call_llm(self, provider, messages, tools, agent, assistant_msg): monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.Message, "create", create_mock) monkeypatch.setattr(runner_mod.Message, "update", update_mock) - monkeypatch.setattr(SessionRunner, "_call_llm", fake_call_llm) + monkeypatch.setattr(StepEngine, "_call_llm", fake_call_llm) last_user = UserMessageInfo( id="msg_user_tool_loop_guard", @@ -3152,8 +3152,8 @@ async def fake_call_llm(self, provider, messages, tools, agent, assistant_msg): ) for idx in range(1, 4): - runner = SessionRunner(session=_make_session("ses_runner_tool_loop_guard"), static_cache=shared_cache) - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + runner = StepEngine(session=_make_session("ses_runner_tool_loop_guard"), static_cache=shared_cache) + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[ {"type": "function", "function": {"name": "echo_tool", "description": "", "parameters": {}}} ])) diff --git a/tests/session/test_runtime_ports.py b/tests/session/test_runtime_ports.py deleted file mode 100644 index 1b76477c0..000000000 --- a/tests/session/test_runtime_ports.py +++ /dev/null @@ -1,108 +0,0 @@ -"""Tests for prompt, tool, model, and hook runtime ports.""" - -from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock - -import pytest - -from flocks.agent.runtime import ExternalRuntimePorts -from flocks.session.runner import SessionRunner -from flocks.session.runtime_adapters import ( - FlocksHookPort, - FlocksModelPort, - FlocksPromptPort, - FlocksToolPort, -) - - -@pytest.mark.asyncio -async def test_runner_uses_injected_ports_at_external_boundaries() -> None: - prompt_port = SimpleNamespace(build_system_prompts=AsyncMock(return_value=[])) - tool_port = SimpleNamespace( - revision=MagicMock(return_value=7), - list_tools=MagicMock(return_value=[]), - get=MagicMock(return_value=None), - ) - model_port = SimpleNamespace( - get_provider=MagicMock( - return_value=SimpleNamespace(_config_models=[]), - ), - apply_config=AsyncMock(), - resolve_model=MagicMock( - return_value=SimpleNamespace( - capabilities=SimpleNamespace(interleaved=True), - ) - ), - resolve_model_info=MagicMock(return_value=(100_000, 8_192, None)), - ) - hook_port = SimpleNamespace( - run_session_start=AsyncMock(), - has_stage_handlers=AsyncMock(return_value=False), - run_llm_before=AsyncMock(), - run_llm_after=AsyncMock(), - ) - ports = ExternalRuntimePorts( - prompts=prompt_port, - tools=tool_port, - models=model_port, - hooks=hook_port, - ) - runner = SessionRunner( - session=SimpleNamespace(id="session-1", directory="/tmp"), - provider_id="provider-a", - model_id="model-a", - runtime_ports=ports, - session_start_pending=True, - ) - - await runner._run_session_start_hook(SimpleNamespace(name="rex")) - capability_key = runner._provider_capability_key() - cache_key = runner._tool_schema_cache_key( - SimpleNamespace(name="rex", tools=[]), - [], - text_tool_call_mode=False, - ) - - hook_port.run_session_start.assert_awaited_once() - model_port.resolve_model.assert_called_once_with("provider-a", "model-a") - assert "interleaved=true" in capability_key - assert cache_key[0] == 7 - - -@pytest.mark.asyncio -async def test_default_adapters_preserve_existing_flocks_interfaces( - monkeypatch: pytest.MonkeyPatch, -) -> None: - build_prompts = AsyncMock(return_value=["system"]) - apply_config = AsyncMock() - session_start = AsyncMock(return_value="hook-result") - monkeypatch.setattr( - "flocks.session.runtime_adapters.SessionPrompt.build_system_prompts", - build_prompts, - ) - monkeypatch.setattr( - "flocks.session.runtime_adapters.ToolRegistry.revision", - MagicMock(return_value=11), - ) - monkeypatch.setattr( - "flocks.session.runtime_adapters.Provider.get", - MagicMock(return_value="provider"), - ) - monkeypatch.setattr( - "flocks.session.runtime_adapters.Provider.apply_config", - apply_config, - ) - monkeypatch.setattr( - "flocks.session.runtime_adapters.HookPipeline.run_session_start", - session_start, - ) - - assert await FlocksPromptPort().build_system_prompts(session_id="session-1") == ["system"] - assert FlocksToolPort().revision() == 11 - assert FlocksModelPort().get_provider("provider-a") == "provider" - await FlocksModelPort().apply_config("provider-a") - assert await FlocksHookPort().run_session_start({"sessionID": "session-1"}) == ("hook-result") - - build_prompts.assert_awaited_once_with(session_id="session-1") - apply_config.assert_awaited_once_with(provider_id="provider-a") - session_start.assert_awaited_once_with({"sessionID": "session-1"}) diff --git a/tests/session/test_session_abort_inject.py b/tests/session/test_session_abort_inject.py index e690ea944..36154e2d8 100644 --- a/tests/session/test_session_abort_inject.py +++ b/tests/session/test_session_abort_inject.py @@ -2,7 +2,7 @@ Tests for session abort and inject functionality. Tests cover: -- SessionRunner external abort_event propagation +- StepEngine external abort_event propagation - SessionLoop abort mechanism - Inject endpoint logic (message creation without starting new loop) - _should_exit behavior with injected messages @@ -14,12 +14,19 @@ import pytest +from flocks.session.runtime.continuation_policy import DEFAULT_CONTINUATION_POLICY +from flocks.session.runtime.session_turn import SessionTurn from flocks.session.message import ToolPart, ToolStateCompleted from flocks.session.goal import GoalDecision -from flocks.session.session_loop import SessionLoop, LoopCallbacks, LoopContext, LoopResult -from flocks.session.runner import SessionRunner, StepResult +from flocks.session.session_loop import ( + LoopCallbacks, + LoopContext, + SessionLoop, +) +from flocks.session.runtime.step_engine import StepEngine, StepResult from flocks.session.session import SessionInfo from flocks.server.routes import session as session_routes +from tests.session_runtime_testkit import run_logical_turns def _make_session_info(session_id: str = "test_session") -> SessionInfo: @@ -55,14 +62,14 @@ def _make_completed_tool_part(message_id: str) -> ToolPart: # --------------------------------------------------------------------------- class TestAbortPropagation: - """Test that abort_event propagates from SessionLoop to SessionRunner.""" + """Test that abort_event propagates from SessionLoop to StepEngine.""" def test_runner_accepts_external_abort_event(self): - """SessionRunner should accept an optional external abort_event.""" + """StepEngine should accept an optional external abort_event.""" external_event = asyncio.Event() session_info = _make_session_info() - runner = SessionRunner( + runner = StepEngine( session=session_info, abort_event=external_event, ) @@ -75,9 +82,9 @@ def test_runner_accepts_external_abort_event(self): assert runner.is_aborted is True def test_runner_internal_abort_still_works(self): - """SessionRunner's own abort() method should still work.""" + """StepEngine's own abort() method should still work.""" session_info = _make_session_info() - runner = SessionRunner(session=session_info) + runner = StepEngine(session=session_info) assert runner.is_aborted is False runner.abort() @@ -88,7 +95,7 @@ def test_runner_either_abort_triggers(self): external_event = asyncio.Event() session_info = _make_session_info() - runner = SessionRunner( + runner = StepEngine( session=session_info, abort_event=external_event, ) @@ -111,7 +118,7 @@ def test_runner_either_abort_triggers(self): def test_runner_without_external_event(self): """Runner created without abort_event should still work normally.""" session_info = _make_session_info() - runner = SessionRunner(session=session_info) + runner = StepEngine(session=session_info) assert runner._external_abort is None assert runner.is_aborted is False @@ -130,12 +137,12 @@ async def test_session_loop_run_publishes_busy_and_idle_status_events(self): ), patch( "flocks.session.session_loop.Message.list", AsyncMock(return_value=[]), + ), patch( + "flocks.session.runtime.session_turn.Message.list", + AsyncMock(return_value=[]), ), patch( "flocks.session.orphan_tools.abort_orphan_running_parts", AsyncMock(return_value=0), - ), patch( - "flocks.session.session_loop.SessionLoop._run_loop", - AsyncMock(return_value=LoopResult(action="stop")), ), patch( "flocks.session.session_loop.Session.touch", AsyncMock(), @@ -186,7 +193,7 @@ def test_abort_running_session(self): ) # Register the context - SessionLoop._active_loops["test_loop_abort"] = ctx + SessionLoop._active_turns["test_loop_abort"] = ctx try: assert ctx.should_abort() is False @@ -195,10 +202,10 @@ def test_abort_running_session(self): assert ctx.should_abort() is True finally: # Clean up - SessionLoop._active_loops.pop("test_loop_abort", None) + SessionLoop._active_turns.pop("test_loop_abort", None) def test_is_running(self): - """is_running should reflect _active_loops state.""" + """is_running should reflect the runtime lease registry.""" assert SessionLoop.is_running("not_there") is False session_info = _make_session_info("running_test") @@ -208,12 +215,12 @@ def test_is_running(self): model_id="test", agent_name="test", ) - SessionLoop._active_loops["running_test"] = ctx + SessionLoop._active_turns["running_test"] = ctx try: assert SessionLoop.is_running("running_test") is True finally: - SessionLoop._active_loops.pop("running_test", None) + SessionLoop._active_turns.pop("running_test", None) def test_get_context(self): """get_context should return the LoopContext for a running session.""" @@ -224,14 +231,14 @@ def test_get_context(self): model_id="test", agent_name="test", ) - SessionLoop._active_loops["ctx_get_test"] = ctx + SessionLoop._active_turns["ctx_get_test"] = ctx try: retrieved = SessionLoop.get_context("ctx_get_test") assert retrieved is ctx assert SessionLoop.get_context("nonexistent") is None finally: - SessionLoop._active_loops.pop("ctx_get_test", None) + SessionLoop._active_turns.pop("ctx_get_test", None) # --------------------------------------------------------------------------- @@ -256,7 +263,7 @@ def test_exit_when_assistant_after_user_and_finished(self): last_assistant = self._make_msg("msg_002", "assistant", finish="stop") # assistant.id > user.id → user.id < assistant.id → True → should exit - assert SessionLoop._should_exit(last_user, last_assistant) is True + assert SessionTurn._should_exit(last_user, last_assistant) is True def test_no_exit_when_user_injected_after_assistant(self): """Should NOT exit when a new user message appears after the assistant. @@ -268,34 +275,34 @@ def test_no_exit_when_user_injected_after_assistant(self): last_assistant = self._make_msg("msg_002", "assistant", finish="stop") # user.id > assistant.id → user.id < assistant.id → False → don't exit - assert SessionLoop._should_exit(last_user, last_assistant) is False + assert SessionTurn._should_exit(last_user, last_assistant) is False def test_no_exit_when_assistant_has_tool_calls(self): """Should NOT exit when assistant finish is 'tool-calls'.""" last_user = self._make_msg("msg_001", "user") last_assistant = self._make_msg("msg_002", "assistant", finish="tool-calls") - assert SessionLoop._should_exit(last_user, last_assistant) is False + assert SessionTurn._should_exit(last_user, last_assistant) is False def test_no_exit_when_no_assistant(self): """Should NOT exit when there is no assistant message yet.""" last_user = self._make_msg("msg_001", "user") - assert SessionLoop._should_exit(last_user, None) is False + assert SessionTurn._should_exit(last_user, None) is False def test_no_exit_when_assistant_finish_is_unknown(self): """Should NOT exit when finish reason is 'unknown'.""" last_user = self._make_msg("msg_001", "user") last_assistant = self._make_msg("msg_002", "assistant", finish="unknown") - assert SessionLoop._should_exit(last_user, last_assistant) is False + assert SessionTurn._should_exit(last_user, last_assistant) is False def test_no_exit_when_assistant_not_finished(self): """Should NOT exit when assistant has no finish status.""" last_user = self._make_msg("msg_001", "user") last_assistant = self._make_msg("msg_002", "assistant", finish=None) - assert SessionLoop._should_exit(last_user, last_assistant) is False + assert SessionTurn._should_exit(last_user, last_assistant) is False def test_no_exit_when_assistant_has_completed_tool_parts(self): """Should continue so completed tool results can be fed back to the model.""" @@ -303,7 +310,7 @@ def test_no_exit_when_assistant_has_completed_tool_parts(self): last_assistant = self._make_msg("msg_002", "assistant", finish="stop") last_assistant_parts = [_make_completed_tool_part(last_assistant.id)] - assert SessionLoop._should_exit( + assert SessionTurn._should_exit( last_user, last_assistant, last_assistant_parts, @@ -322,7 +329,7 @@ def _make_msg(msg_id: str, role: str): async def test_does_not_treat_current_user_as_queued_when_no_assistant_exists(self): current_user = self._make_msg("msg_001", "user") - queued = await SessionLoop._detect_queued_user_message( + queued = await DEFAULT_CONTINUATION_POLICY.detect_queued_user_message( "session-1", [current_user], current_user.id, @@ -336,7 +343,7 @@ async def test_detects_newer_user_when_step_failed_before_assistant_created(self current_user = self._make_msg("msg_001", "user") newer_user = self._make_msg("msg_002", "user") - queued = await SessionLoop._detect_queued_user_message( + queued = await DEFAULT_CONTINUATION_POLICY.detect_queued_user_message( "session-1", [current_user, newer_user], current_user.id, @@ -371,11 +378,11 @@ async def test_run_loop_stops_turn_when_messages_are_empty(self): model_id="test-model", agent_name="rex", ) - ctx.session_ctx = SimpleNamespace(get_messages=AsyncMock(return_value=[])) + ctx.session_store = SimpleNamespace(get_messages=AsyncMock(return_value=[])) event_callback = AsyncMock() callbacks = LoopCallbacks(event_publish_callback=event_callback) - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" event_names = [call.args[0] for call in event_callback.await_args_list] @@ -398,11 +405,11 @@ async def test_run_loop_stops_turn_when_no_user_message_exists(self): agent_name="rex", ) assistant = self._make_msg("msg_001", "assistant", finish="stop") - ctx.session_ctx = SimpleNamespace(get_messages=AsyncMock(return_value=[assistant])) + ctx.session_store = SimpleNamespace(get_messages=AsyncMock(return_value=[assistant])) event_callback = AsyncMock() callbacks = LoopCallbacks(event_publish_callback=event_callback) - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" event_names = [call.args[0] for call in event_callback.await_args_list] @@ -428,12 +435,16 @@ async def test_run_loop_continues_for_active_goal_after_stop(self): assistant = self._make_msg("msg_002", "assistant", finish="stop") goal_user = self._make_msg("msg_003", "user") assistant_after_goal = self._make_msg("msg_004", "assistant", finish="stop") - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[ [user], [user, assistant], + [user, assistant], + [user, assistant], [user, assistant, goal_user], [user, assistant, goal_user, assistant_after_goal], + [user, assistant, goal_user, assistant_after_goal], + [user, assistant, goal_user, assistant_after_goal], ]) ) event_callback = AsyncMock() @@ -450,25 +461,25 @@ async def test_run_loop_continues_for_active_goal_after_stop(self): ] with patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", MagicMock(return_value="still working"), ), patch( - "flocks.session.session_loop.Message.create", + "flocks.session.runtime.session_turn.Message.create", AsyncMock(return_value=goal_user), ) as create_message, patch( - "flocks.session.session_loop.GoalManager.evaluate_after_turn", + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", AsyncMock(side_effect=goal_decisions), ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(side_effect=[StepResult(action="stop"), StepResult(action="stop")]), ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" create_message.assert_awaited_once() @@ -502,7 +513,7 @@ async def test_run_loop_waits_for_user_input_after_goal_clarification(self): ) user = self._make_msg("msg_001", "user") assistant = self._make_msg("msg_002", "assistant", finish="stop") - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[ [user], [user, assistant], @@ -512,19 +523,19 @@ async def test_run_loop_waits_for_user_input_after_goal_clarification(self): callbacks = LoopCallbacks(event_publish_callback=event_callback) with patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", MagicMock(return_value="Please clarify what tests to write."), ), patch( - "flocks.session.session_loop.Message.create", + "flocks.session.runtime.session_turn.Message.create", AsyncMock(), ) as create_message, patch( - "flocks.session.session_loop.GoalManager.evaluate_after_turn", + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", AsyncMock(return_value=GoalDecision( status="active", verdict="waiting", @@ -532,10 +543,10 @@ async def test_run_loop_waits_for_user_input_after_goal_clarification(self): reason="waiting for user clarification", )), ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(return_value=StepResult(action="stop")), ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" create_message.assert_not_awaited() @@ -558,7 +569,7 @@ async def test_run_loop_passes_pending_question_to_goal_judge(self): ) user = self._make_msg("msg_001", "user") assistant = self._make_msg("msg_002", "assistant", finish="stop") - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[ [user], [user, assistant], @@ -574,28 +585,28 @@ async def test_run_loop_passes_pending_question_to_goal_judge(self): )) with patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", MagicMock(return_value="Please provide the input."), ), patch( "flocks.server.routes.question.has_pending_questions", MagicMock(return_value=True), ), patch( - "flocks.session.session_loop.Message.create", + "flocks.session.runtime.session_turn.Message.create", AsyncMock(), ) as create_message, patch( - "flocks.session.session_loop.GoalManager.evaluate_after_turn", + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", evaluate_goal, ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(return_value=StepResult(action="stop")), ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" create_message.assert_not_awaited() @@ -621,23 +632,23 @@ async def test_run_loop_publishes_goal_terminal_status(self): self._make_msg("msg_001", "user"), self._make_msg("msg_002", "assistant", finish="stop"), ] - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[[messages[0]], messages]) ) event_callback = AsyncMock() callbacks = LoopCallbacks(event_publish_callback=event_callback) with patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", MagicMock(return_value="Goal complete: done"), ), patch( - "flocks.session.session_loop.GoalManager.evaluate_after_turn", + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", AsyncMock(return_value=GoalDecision( status="completed", verdict="complete", @@ -645,10 +656,10 @@ async def test_run_loop_publishes_goal_terminal_status(self): objective="finish work", )), ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(return_value=StepResult(action="stop")), ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" event_names = [call.args[0] for call in event_callback.await_args_list] @@ -693,26 +704,26 @@ async def test_pre_compact_cleanup_emits_turn_continued_before_next_iteration(se tokens={"input": 0, "output": 0, "cache": {"read": 0, "write": 0}}, ), ] - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[overflow_messages, normal_messages, normal_messages]) ) event_callback = AsyncMock() callbacks = LoopCallbacks(event_publish_callback=event_callback) with patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(20000, 1024, None), ), patch( - "flocks.session.session_loop.SessionCompaction.truncate_oversized_tool_outputs", + "flocks.session.runtime.session_turn.SessionCompaction.truncate_oversized_tool_outputs", AsyncMock(return_value=1), ), patch( - "flocks.session.session_loop.SessionPrompt.estimate_full_context_tokens", + "flocks.session.runtime.session_turn.SessionPrompt.estimate_full_context_tokens", AsyncMock(return_value=0), ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(return_value=StepResult(action="stop")), ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" event_names = [call.args[0] for call in event_callback.await_args_list] @@ -745,7 +756,7 @@ async def test_run_loop_skips_exit_condition_when_assistant_has_tool_parts(self) self._make_msg("msg_001", "user"), self._make_msg("msg_002", "assistant", finish="stop"), ] - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[messages, messages]) ) event_callback = AsyncMock() @@ -754,25 +765,25 @@ async def test_run_loop_skips_exit_condition_when_assistant_has_tool_parts(self) log_info = MagicMock() with patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[_make_completed_tool_part("msg_002")]), ), patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( "flocks.session.lifecycle.title.SessionTitle.ensure_title", MagicMock(return_value=None), ), patch( - "flocks.session.session_loop.fire_and_forget", + "flocks.session.runtime.session_turn.fire_and_forget", MagicMock(), ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", process_step, ), patch( - "flocks.session.session_loop.log.info", + "flocks.session.runtime.session_turn.log.info", log_info, ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" assert result.last_message is messages[1] @@ -799,7 +810,7 @@ async def test_run_loop_breaks_on_exit_condition_without_tool_parts(self): self._make_msg("msg_001", "user"), self._make_msg("msg_002", "assistant", finish="stop"), ] - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(return_value=messages) ) event_callback = AsyncMock() @@ -808,16 +819,16 @@ async def test_run_loop_breaks_on_exit_condition_without_tool_parts(self): log_info = MagicMock() with patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.log.info", + "flocks.session.runtime.session_turn.log.info", log_info, ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", process_step, ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" assert result.last_message is messages[1] @@ -863,11 +874,11 @@ async def test_execute_subtask_passes_tool_context_first(self): with patch("flocks.agent.registry.Agent.get", AsyncMock(return_value=SimpleNamespace(name="helper"))), \ patch("flocks.tool.registry.ToolRegistry.get", return_value=task_tool), \ - patch("flocks.session.session_loop.Message.create", AsyncMock(side_effect=[assistant_msg, synthetic_msg])), \ - patch("flocks.session.session_loop.Message.add_part", AsyncMock()), \ - patch("flocks.session.session_loop.Message.update", AsyncMock()), \ - patch("flocks.session.session_loop.Message.update_part", AsyncMock()): - await SessionLoop._execute_subtask(ctx, last_user, task_part) + patch("flocks.session.runtime.session_turn.Message.create", AsyncMock(side_effect=[assistant_msg, synthetic_msg])), \ + patch("flocks.session.runtime.session_turn.Message.add_part", AsyncMock()), \ + patch("flocks.session.runtime.session_turn.Message.update", AsyncMock()), \ + patch("flocks.session.runtime.session_turn.Message.update_part", AsyncMock()): + await ctx._execute_subtask(last_user, task_part) task_tool.execute.assert_awaited_once() tool_ctx = task_tool.execute.await_args.args[0] diff --git a/tests/session/test_session_context.py b/tests/session/test_session_context.py index 77388f586..6216267d4 100644 --- a/tests/session/test_session_context.py +++ b/tests/session/test_session_context.py @@ -5,8 +5,8 @@ 1. SessionContext protocol is properly defined 2. DefaultSessionContext implements all methods 3. DefaultSessionContext delegates to underlying session modules -4. LoopContext carries session_ctx -5. SessionRunner accepts session_ctx +4. LoopContext carries session_store +5. StepEngine accepts session_store """ import pytest @@ -14,7 +14,7 @@ from flocks.session.core.context import SessionContext, DefaultSessionContext from flocks.session.session_loop import LoopContext -from flocks.session.runner import SessionRunner +from flocks.session.runtime.step_engine import StepEngine class TestSessionContextProtocol: @@ -137,10 +137,10 @@ async def test_touch_delegates_to_session(self): mock_touch.assert_called_once_with("proj-1", "ses-123") -class TestLoopContextSessionCtx: - """LoopContext should carry session_ctx.""" +class TestLoopContextSessionStore: + """LoopContext should carry session_store.""" - def test_loop_context_has_session_ctx_field(self): + def test_loop_context_has_session_store_field(self): import asyncio session = MagicMock() session.id = "test" @@ -153,24 +153,24 @@ def test_loop_context_has_session_ctx_field(self): model_id="claude-sonnet-4", agent_name="rex", ) - assert ctx.session_ctx is None + assert ctx.session_store is None - def test_loop_context_with_session_ctx(self): + def test_loop_context_with_session_store(self): session = MagicMock() session.id = "test" session.directory = "/test" session.project_id = "proj" - session_ctx = DefaultSessionContext(session) + session_store = DefaultSessionContext(session) ctx = LoopContext( session=session, provider_id="anthropic", model_id="claude-sonnet-4", agent_name="rex", - session_ctx=session_ctx, + session_store=session_store, ) - assert ctx.session_ctx is session_ctx - assert ctx.session_ctx.session_id == "test" + assert ctx.session_store is session_store + assert ctx.session_store.session_id == "test" def test_loop_context_tracks_observed_prompt_tokens(self): # B3 — LoopContext must expose ``last_observed_prompt_tokens`` so @@ -192,27 +192,27 @@ def test_loop_context_tracks_observed_prompt_tokens(self): assert ctx.last_observed_prompt_tokens == 123_456 -class TestRunnerSessionCtx: - """SessionRunner should accept session_ctx.""" +class TestStepEngineSessionStore: + """StepEngine should accept session_store.""" - def test_runner_accepts_session_ctx(self): + def test_step_engine_accepts_session_store(self): session = MagicMock() session.id = "test" session.directory = "/test" session.project_id = "proj" - session_ctx = DefaultSessionContext(session) - runner = SessionRunner( + session_store = DefaultSessionContext(session) + runner = StepEngine( session=session, - session_ctx=session_ctx, + session_store=session_store, ) - assert runner.session_ctx is session_ctx + assert runner.session_store is session_store - def test_runner_session_ctx_defaults_to_none(self): + def test_step_engine_session_store_defaults_to_none(self): session = MagicMock() session.id = "test" session.directory = "/test" session.project_id = "proj" - runner = SessionRunner(session=session) - assert runner.session_ctx is None + runner = StepEngine(session=session) + assert runner.session_store is None diff --git a/tests/session/test_session_loop_working_directory.py b/tests/session/test_session_loop_working_directory.py index 4c2cc5f3a..dbda31ed0 100644 --- a/tests/session/test_session_loop_working_directory.py +++ b/tests/session/test_session_loop_working_directory.py @@ -3,9 +3,16 @@ import pytest from flocks.bus.bus import Bus +from flocks.session.runtime.agent_loop import AgentLoop +from flocks.session.runtime.contracts import ( + AgentRunOutcome, + AgentRunState, + AgentRunStatus, + RuntimeModel, +) from flocks.session.message import Message from flocks.session.session import Session, SessionInfo -from flocks.session.session_loop import LoopResult, SessionLoop +from flocks.session.session_loop import SessionLoop @pytest.mark.asyncio @@ -16,12 +23,24 @@ async def test_run_uses_runtime_working_directory(monkeypatch: pytest.MonkeyPatc directory="/missing/original", title="Legacy session", ) - run_loop = AsyncMock(return_value=LoopResult(action="stop")) + + async def run_agent_turn(context, _engine): + state = AgentRunState( + session_id=context.session.id, + agent_name=context.agent_name, + active_model=RuntimeModel(context.provider_id, context.model_id), + ) + return AgentRunOutcome( + status=AgentRunStatus.ABORTED, + state=state, + ) + + run_turn = AsyncMock(side_effect=run_agent_turn) monkeypatch.setattr(Session, "get_by_id", AsyncMock(return_value=session)) monkeypatch.setattr(Session, "touch", AsyncMock()) monkeypatch.setattr(Message, "list", AsyncMock(return_value=[])) - monkeypatch.setattr(SessionLoop, "_run_loop", run_loop) + monkeypatch.setattr(AgentLoop, "run", run_turn) monkeypatch.setattr( "flocks.session.orphan_tools.abort_orphan_running_parts", AsyncMock(), @@ -36,7 +55,7 @@ async def test_run_uses_runtime_working_directory(monkeypatch: pytest.MonkeyPatc ) assert result.action == "stop" - loop_context = run_loop.await_args.args[0] + loop_context = run_turn.await_args.args[0] assert loop_context.session.directory == "/available/default" - assert loop_context.session_ctx.directory == "/available/default" + assert loop_context.session_store.directory == "/available/default" assert session.directory == "/missing/original" diff --git a/tests/session/test_session_runner_tool_only_message.py b/tests/session/test_session_runner_tool_only_message.py index 28e879e4f..f8624de42 100644 --- a/tests/session/test_session_runner_tool_only_message.py +++ b/tests/session/test_session_runner_tool_only_message.py @@ -5,7 +5,7 @@ from flocks.provider.provider import ChatMessage, Provider from flocks.session.message import Message, MessageRole, ToolPart, ToolStateCompleted from flocks.session.prompt import SessionPrompt -from flocks.session.runner import SessionRunner, StepResult +from flocks.session.runtime.step_engine import StepEngine, StepResult from flocks.session.session import Session from flocks.utils.id import Identifier @@ -99,13 +99,13 @@ async def fake_call_llm(self, provider, messages, tools, agent, assistant_msg): monkeypatch.setattr(Provider, "get", lambda _provider_id: DummyProvider()) monkeypatch.setattr(Provider, "apply_config", fake_apply_config) monkeypatch.setattr(Agent, "get", fake_agent_get) - monkeypatch.setattr(SessionRunner, "_get_prompt_tool_names", fake_get_prompt_tool_names) + monkeypatch.setattr(StepEngine, "_get_prompt_tool_names", fake_get_prompt_tool_names) monkeypatch.setattr(SessionPrompt, "build_system_prompts", fake_build_system_prompts) - monkeypatch.setattr(SessionRunner, "_build_callable_tool_schema", fake_build_callable_tool_schema) - monkeypatch.setattr(SessionRunner, "_to_chat_messages", fake_to_chat_messages) - monkeypatch.setattr(SessionRunner, "_call_llm", fake_call_llm) + monkeypatch.setattr(StepEngine, "_build_callable_tool_schema", fake_build_callable_tool_schema) + monkeypatch.setattr(StepEngine, "_to_chat_messages", fake_to_chat_messages) + monkeypatch.setattr(StepEngine, "_call_llm", fake_call_llm) - runner = SessionRunner(session=session, provider_id="test-provider", model_id="test-model", agent_name="rex") + runner = StepEngine(session=session, provider_id="test-provider", model_id="test-model", agent_name="rex") runner._step = 2 # ensure reminder wrapping branch doesn't break assumptions result = await runner._process_step(messages=messages, last_user=user_2) diff --git a/tests/session/test_step_engine.py b/tests/session/test_step_engine.py deleted file mode 100644 index e940a54f8..000000000 --- a/tests/session/test_step_engine.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Tests for the existing session runner's StepEngine adapter.""" - -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import pytest - -from flocks.agent.runtime.contracts import ( - AttemptEffects, - ModelTurnSnapshot, - RuntimeModel, - StepResult, -) -from flocks.session.runner import ( - LlmAttemptState, - StepResult as LegacyStepResult, -) -from flocks.session.step_engine import SessionStepEngine - - -@pytest.mark.asyncio -async def test_session_step_engine_delegates_immutable_snapshot() -> None: - last_user = SimpleNamespace(id="user-1") - expected = StepResult(action="stop", content="done") - runner = SimpleNamespace( - _process_step=AsyncMock(return_value=expected), - _session_start_fired=True, - _attempt_state=AttemptEffects(received_chunk=True), - _step=0, - ) - engine = SessionStepEngine(runner) - snapshot = ModelTurnSnapshot( - session_id="session-1", - agent_name="rex", - active_model=RuntimeModel("provider", "model"), - model_turn_index=1, - trace_step=8, - messages=(last_user,), - last_user=last_user, - ) - - result = await engine.run(snapshot) - - assert result is expected - assert runner._step == 8 - runner._process_step.assert_awaited_once_with([last_user], last_user) - assert engine.session_start_fired is True - assert engine.attempt_effects.received_chunk is True - - -def test_legacy_runner_contract_exports_remain_compatible() -> None: - assert LegacyStepResult is StepResult - assert LlmAttemptState is AttemptEffects diff --git a/tests/session_runtime_testkit.py b/tests/session_runtime_testkit.py new file mode 100644 index 000000000..15cf0f4e4 --- /dev/null +++ b/tests/session_runtime_testkit.py @@ -0,0 +1,35 @@ +"""Test-only helpers for exercising SessionLoop logical-turn control.""" + +from flocks.session.session_loop import ( + LoopCallbacks, + LoopContext, + LoopResult, + SessionLoop, +) + + +async def run_logical_turns( + turn: LoopContext, + callbacks: LoopCallbacks, +) -> LoopResult: + """Run logical turns without acquiring a persisted session lease.""" + turn.callbacks = callbacks + policy = turn.continuation_policy or SessionLoop._continuation_policy + processed_user_id = None + while True: + try: + await policy.prepare_logical_turn(turn) + processed_user_id = turn.prepared_user_id or processed_user_id + outcome = await SessionLoop._run_logical_input(turn) + processed_user_id = ( + outcome.state.current_user_id or processed_user_id + ) + if await SessionLoop._should_continue(turn, policy, outcome): + continue + except Exception as exc: + outcome = await SessionLoop._handle_execution_error( + turn, + exc, + processed_user_id, + ) + return SessionLoop._to_loop_result(turn, outcome) diff --git a/tests/task/test_task.py b/tests/task/test_task.py index f2b31aad7..f12d1a49e 100644 --- a/tests/task/test_task.py +++ b/tests/task/test_task.py @@ -624,6 +624,12 @@ async def test_background_task_completion_injects_parent_context( monkeypatch.setattr(background_module.Message, "update_part", update_part) monkeypatch.setattr(background_module.SessionLoop, "run", parent_loop_run) + async def run_active_write(_session_id, operation, **_kwargs): + return await operation() + + active_write = AsyncMock(side_effect=run_active_write) + monkeypatch.setattr(background_module.Session, "run_active_write", active_write) + manager = BackgroundManager() task = BackgroundTask( id="bg_parent_inject", @@ -642,6 +648,8 @@ async def test_background_task_completion_injects_parent_context( await manager._inject_parent_completion(task) + active_write.assert_awaited_once() + assert active_write.await_args.args[0] == "ses-parent" create_message.assert_awaited_once() kwargs = create_message.await_args.kwargs assert kwargs["session_id"] == "ses-parent" @@ -662,12 +670,11 @@ async def test_background_task_completion_injects_parent_context( @pytest.mark.asyncio -async def test_background_task_completion_does_not_resume_running_parent( +async def test_background_task_completion_always_attempts_parent_resume( monkeypatch: pytest.MonkeyPatch, ): parent_loop_run = AsyncMock(return_value=SimpleNamespace(action="stop")) monkeypatch.setattr(background_module.SessionLoop, "run", parent_loop_run) - monkeypatch.setattr(background_module.SessionLoop, "is_running", lambda _session_id: True) manager = BackgroundManager() task = BackgroundTask( @@ -685,7 +692,8 @@ async def test_background_task_completion_does_not_resume_running_parent( manager._schedule_parent_resume(task) await asyncio.sleep(0) - parent_loop_run.assert_not_awaited() + parent_loop_run.assert_awaited_once() + assert parent_loop_run.await_args.kwargs["session_id"] == "ses-parent" @pytest.mark.asyncio From 5870662e6a72c192a2b8f26088096be8bbe3cf06 Mon Sep 17 00:00:00 2001 From: xiami762 Date: Fri, 7 Aug 2026 21:16:57 +0800 Subject: [PATCH 06/15] refactor: remove legacy task and subtask systems --- .../agent/agents/hephaestus/prompt_builder.py | 43 +--- flocks/agent/agents/rex/prompt_builder.py | 57 ++--- flocks/cli/commands/import_.py | 15 +- flocks/cli/commands/task.py | 50 ++-- flocks/command/command.py | 1 - flocks/config/config.py | 1 - flocks/server/app.py | 22 +- flocks/server/routes/misc.py | 3 - flocks/server/routes/session.py | 9 - flocks/server/routes/skill.py | 2 - flocks/server/routes/stats.py | 4 +- flocks/server/routes/task_entities.py | 93 ++++--- flocks/session/__init__.py | 11 - flocks/session/context_usage.py | 8 +- flocks/session/execution_mode.py | 2 +- flocks/session/features/subtask.py | 62 ----- flocks/session/message.py | 33 +-- flocks/session/prompt/anthropic-20250930.txt | 4 +- flocks/session/prompt/anthropic.txt | 10 +- flocks/session/runtime/session_turn.py | 229 +---------------- flocks/session/runtime/step_engine.py | 6 - flocks/session/streaming/stream_processor.py | 2 +- flocks/task/__init__.py | 4 +- flocks/task/plugin_sync.py | 4 +- .../{manager.py => schedule_task_manager.py} | 12 +- flocks/task/scheduler.py | 4 +- flocks/tool/agent/delegate_task.py | 9 +- flocks/tool/agent/task.py | 106 -------- flocks/tool/catalog.py | 1 - flocks/tool/registry.py | 2 +- flocks/tool/task/schedule_task_center.py | 52 ++-- flocks/utils/id.py | 2 - flocks/workflow/tool_context.py | 2 +- .../test_task_queue_integration.py | 86 +++---- .../test_task_scheduler_context_route.py | 8 +- tests/server/test_lifespan.py | 4 +- tests/server/test_server.py | 43 ++-- tests/session/test_auto_model_failover.py | 4 +- tests/session/test_context_usage.py | 10 +- tests/session/test_execution_mode.py | 16 +- tests/session/test_message_parts.py | 41 ++-- tests/session/test_runner_step.py | 28 +-- tests/session/test_session_abort_inject.py | 54 ---- .../storage/test_sqlite_connection_config.py | 6 +- tests/task/test_task.py | 202 +++++++-------- tests/tool/test_builtin_management_tools.py | 4 +- tests/tool/test_delegate_task_compat.py | 16 +- tests/tool/test_task_center_compat.py | 44 ++-- tests/tool/test_task_list_routing.py | 2 +- tests/tool/test_task_model_pinning.py | 61 ----- tests/tool/test_tools.py | 12 +- tests/utils/test_id_compatibility.py | 1 - .../cli/cmd/tui/routes/session/index.tsx | 10 +- tui/flocks/command/index.ts | 4 +- tui/flocks/config/config.ts | 1 - tui/flocks/session/message-v2.ts | 22 -- tui/flocks/session/prompt.ts | 231 +----------------- .../session/prompt/anthropic-20250930.txt | 4 +- tui/flocks/session/prompt/anthropic.txt | 10 +- tui/flocks/tool/{task.ts => delegate-task.ts} | 5 +- .../tool/{task.txt => delegate-task.txt} | 18 +- tui/flocks/tool/registry.ts | 4 +- tui/sdk/gen/types.gen.ts | 23 +- tui/sdk/v2/gen/sdk.gen.ts | 5 +- tui/sdk/v2/gen/types.gen.ts | 33 +-- webui/src/api/skill.ts | 1 - 66 files changed, 500 insertions(+), 1378 deletions(-) delete mode 100644 flocks/session/features/subtask.py rename flocks/task/{manager.py => schedule_task_manager.py} (99%) delete mode 100644 flocks/tool/agent/task.py delete mode 100644 tests/tool/test_task_model_pinning.py rename tui/flocks/tool/{task.ts => delegate-task.ts} (97%) rename tui/flocks/tool/{task.txt => delegate-task.txt} (79%) diff --git a/flocks/agent/agents/hephaestus/prompt_builder.py b/flocks/agent/agents/hephaestus/prompt_builder.py index 9617affcb..3d220164b 100644 --- a/flocks/agent/agents/hephaestus/prompt_builder.py +++ b/flocks/agent/agents/hephaestus/prompt_builder.py @@ -29,7 +29,6 @@ def inject( available_agents=available_agents, available_tools=tools, available_skills=skills, - use_task_system=False, ) @@ -37,7 +36,6 @@ def build_hephaestus_prompt( available_agents: List["AvailableAgent"], available_tools: List["AvailableTool"], available_skills: List["AvailableSkill"], - use_task_system: bool = False, ) -> str: from flocks.agent.prompt_utils import ( build_agent_selection_table, @@ -62,7 +60,7 @@ def build_hephaestus_prompt( oracle_section = build_oracle_section(available_agents) hard_blocks = build_hard_blocks_section() anti_patterns = build_anti_patterns_section() - todo_discipline = _todo_discipline_section(use_task_system) + todo_discipline = _todo_discipline_section() template = """You are Hephaestus, an autonomous deep worker for software engineering. @@ -245,44 +243,7 @@ def build_hephaestus_prompt( return prompt -def _todo_discipline_section(use_task_system: bool) -> str: - if use_task_system: - return """## Task Discipline (NON-NEGOTIABLE) - -**Track ALL multi-step work with tasks. This is your execution backbone.** - -### When to Create Tasks (MANDATORY) - -| Trigger | Action | -|---------|--------| -| 2+ step task | `TaskCreate` FIRST, atomic breakdown | -| Uncertain scope | `TaskCreate` to clarify thinking | -| Complex single task | Break down into trackable steps | - -### Workflow (STRICT) - -1. **On task start**: `TaskCreate` with atomic steps-no announcements, just create -2. **Before each step**: `TaskUpdate(status="in_progress")` (ONE at a time) -3. **After each step**: `TaskUpdate(status="completed")` IMMEDIATELY (NEVER batch) -4. **Scope changes**: Update tasks BEFORE proceeding - -### Why This Matters - -- **Execution anchor**: Tasks prevent drift from original request -- **Recovery**: If interrupted, tasks enable seamless continuation -- **Accountability**: Each task = explicit commitment to deliver - -### Anti-Patterns (BLOCKING) - -| Violation | Why It Fails | -|-----------|--------------| -| Skipping tasks on multi-step work | Steps get forgotten, user has no visibility | -| Batch-completing multiple tasks | Defeats real-time tracking purpose | -| Proceeding without `in_progress` | No indication of current work | -| Finishing without completing tasks | Task appears incomplete | - -**NO TASKS ON MULTI-STEP WORK = INCOMPLETE WORK.**""" - +def _todo_discipline_section() -> str: return """## Todo Discipline (NON-NEGOTIABLE) **Track ALL multi-step work with todos. This is your execution backbone.** diff --git a/flocks/agent/agents/rex/prompt_builder.py b/flocks/agent/agents/rex/prompt_builder.py index 1e74df7ab..e705fb862 100644 --- a/flocks/agent/agents/rex/prompt_builder.py +++ b/flocks/agent/agents/rex/prompt_builder.py @@ -30,7 +30,6 @@ def inject( available_tools=tools, available_skills=skills, available_workflows=workflows or [], - use_task_system=False, ) @@ -39,7 +38,6 @@ def build_dynamic_rex_prompt( available_tools: List["AvailableTool"], available_skills: List["AvailableSkill"], available_workflows: Optional[List["AvailableWorkflow"]] = None, - use_task_system: bool = False, ) -> str: from flocks.agent.prompt_utils import ( build_agent_selection_table, @@ -58,12 +56,8 @@ def build_dynamic_rex_prompt( im_send_section = _build_im_send_pointer_section() anti_patterns = _build_rex_anti_patterns_section() command_guidance_section = _build_command_guidance_section() - task_management_section = _task_management_section(use_task_system) - todo_hook_note = ( - "YOUR TASK CREATION WOULD BE TRACKED BY HOOK([SYSTEM REMINDER - TASK CONTINUATION])" - if use_task_system - else "YOUR TODO CREATION WOULD BE TRACKED BY HOOK([SYSTEM REMINDER - TODO CONTINUATION])" - ) + task_management_section = _task_management_section() + todo_hook_note = "YOUR TODO CREATION WOULD BE TRACKED BY HOOK([SYSTEM REMINDER - TODO CONTINUATION])" template = """ You are "Rex" - Powerful AI orchestrator for security operations. @@ -143,7 +137,7 @@ def build_dynamic_rex_prompt( - Match existing codebase patterns when editing. - Fix bugs minimally; do not refactor during a bugfix unless required. - Keep search bounded: stop when you have enough context, when results repeat, or when direct evidence already answers the question. -- For independent parallel branches whose results are needed this turn, emit multiple foreground `delegate_task` / `task` tool calls in the same assistant turn. The runtime executes those sibling tool calls concurrently and returns all tool results before you continue. +- For independent parallel branches whose results are needed this turn, emit multiple foreground `delegate_task` tool calls in the same assistant turn. The runtime executes those sibling tool calls concurrently and returns all tool results before you continue. - Do not use `run_in_background=true`; background subagent execution is disabled. ## 5. Verify @@ -278,55 +272,42 @@ def _build_clarification_protocol() -> str: ```""" -def _task_management_section(use_task_system: bool) -> str: - title = "Task Management" if use_task_system else "Todo Management" - unit = "tasks" if use_task_system else "todos" - create_action = "`TaskCreate`" if use_task_system else '`todo(action="write")`' - progress_action = ( - '`TaskUpdate(status="in_progress")`' - if use_task_system - else "mark `in_progress`" - ) - complete_action = ( - '`TaskUpdate(status="completed")`' - if use_task_system - else "mark `completed`" - ) +def _task_management_section() -> str: clarification_protocol = _build_clarification_protocol() - return f""" -## {title} + return f""" +## Todo Management -Use {unit} as the primary coordination mechanism for non-trivial execution work. +Use todos as the primary coordination mechanism for non-trivial execution work. ### When They Are Mandatory | Trigger | Action | |---------|--------| -| Multi-step work (2+ steps) | Create {unit} first | -| Uncertain scope | Create {unit} to structure the work | -| User request with multiple items | Create {unit} first | -| Complex single task | Break it into {unit} | +| Multi-step work (2+ steps) | Create todos first | +| Uncertain scope | Create todos to structure the work | +| User request with multiple items | Create todos first | +| Complex single task | Break it into todos | ### Operating Rules -1. Start with {create_action} before implementation work begins. -2. ONLY add {unit} when the user wants execution, not when they only want analysis or planning. -3. Before each step, {progress_action}. Keep only one item in progress. -4. After each step, {complete_action} immediately. Never batch updates. -5. If scope changes, update the {unit} before continuing. +1. Start with `todo(action="write")` before implementation work begins. +2. ONLY add todos when the user wants execution, not when they only want analysis or planning. +3. Before each step, mark it `in_progress`. Keep only one item in progress. +4. After each step, mark it `completed` immediately. Never batch updates. +5. If scope changes, update the todos before continuing. ### Failure Modes | Violation | Why It Breaks the Workflow | |-----------|----------------------------| -| Skipping {unit} on non-trivial work | The user loses progress visibility and steps get dropped | -| Batch-completing multiple {unit} | Real-time tracking becomes meaningless | +| Skipping todos on non-trivial work | The user loses progress visibility and steps get dropped | +| Batch-completing multiple todos | Real-time tracking becomes meaningless | | Proceeding without an in-progress item | It is unclear what is being worked on | | Finishing without closing items | The work appears incomplete | {clarification_protocol} -""" +""" def _build_security_priority_section(available_agents: List["AvailableAgent"]) -> str: diff --git a/flocks/cli/commands/import_.py b/flocks/cli/commands/import_.py index fe04b4845..2dd701c8a 100644 --- a/flocks/cli/commands/import_.py +++ b/flocks/cli/commands/import_.py @@ -98,6 +98,17 @@ def _normalize_part_data( metadata = normalized.get("metadata") metadata_dict = metadata if isinstance(metadata, dict) else {} + if part_type == "subtask": + normalized["type"] = "text" + normalized["text"] = "" + normalized["ignored"] = True + normalized["metadata"] = { + **metadata_dict, + "legacyPartType": "subtask", + } + part_type = "text" + metadata_dict = normalized["metadata"] + if "content" in normalized and "text" not in normalized: normalized["text"] = normalized.get("content", "") @@ -135,10 +146,6 @@ def _normalize_part_data( ) elif part_type == "agent": normalized.setdefault("name", metadata_dict.get("name") or normalized.get("content") or "agent") - elif part_type == "subtask": - normalized.setdefault("prompt", metadata_dict.get("prompt") or normalized.get("content", "")) - normalized.setdefault("description", metadata_dict.get("description") or "") - normalized.setdefault("agent", metadata_dict.get("agent") or "agent") elif part_type == "retry": normalized.setdefault("attempt", metadata_dict.get("attempt") or 1) normalized.setdefault("error", metadata_dict.get("error") or {}) diff --git a/flocks/cli/commands/task.py b/flocks/cli/commands/task.py index 25ec387c1..cf7024b9c 100644 --- a/flocks/cli/commands/task.py +++ b/flocks/cli/commands/task.py @@ -45,11 +45,11 @@ def task_dashboard(): async def _dashboard(): await Storage.init() - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.store import TaskStore await TaskStore.init() - counts = await TaskManager.dashboard() + counts = await ScheduleTaskManager.dashboard() panel_lines = [ f"🟢 Running: {counts.get('running', 0)}", @@ -63,7 +63,7 @@ async def _dashboard(): console.print(Panel("\n".join(panel_lines), title="📋 Task Center", border_style="cyan")) - unviewed = await TaskManager.get_unviewed_results() + unviewed = await ScheduleTaskManager.get_unviewed_results() if unviewed: console.print() console.print("[bold]Unviewed completed tasks:[/bold]") @@ -88,7 +88,7 @@ def task_list( async def _list_tasks(status_val, type_val, limit, fmt): await Storage.init() - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.models import SchedulerStatus, TaskStatus from flocks.task.store import TaskStore await TaskStore.init() @@ -99,7 +99,7 @@ async def _list_tasks(status_val, type_val, limit, fmt): scheduler_status = SchedulerStatus.ACTIVE elif status_val in ("paused", "disabled"): scheduler_status = SchedulerStatus.DISABLED - tasks, total = await TaskManager.list_schedulers(status=scheduler_status, limit=limit) + tasks, total = await ScheduleTaskManager.list_schedulers(status=scheduler_status, limit=limit) else: task_status = None if status_val: @@ -108,7 +108,7 @@ async def _list_tasks(status_val, type_val, limit, fmt): task_status = TaskStatus(mapped_status) except ValueError as exc: raise typer.BadParameter(f"Invalid execution status: {status_val}") from exc - tasks, total = await TaskManager.list_executions( + tasks, total = await ScheduleTaskManager.list_executions( status=task_status, limit=limit, ) @@ -160,13 +160,13 @@ def task_show(task_id: str = typer.Argument(..., help="Task ID")): async def _show_task(task_id: str): await Storage.init() - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.store import TaskStore await TaskStore.init() - task = await TaskManager.get_execution(task_id) + task = await ScheduleTaskManager.get_execution(task_id) if task is None: - task = await TaskManager.get_scheduler(task_id) + task = await ScheduleTaskManager.get_scheduler(task_id) if not task: console.print(f"[red]Task {task_id} not found[/red]") raise typer.Exit(1) @@ -243,7 +243,7 @@ def task_create( async def _create_task(title, description, task_type, priority, mode, agent, workflow, skills, cron, prompt): await Storage.init() - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.models import ( ExecutionMode, SchedulerMode, @@ -270,7 +270,7 @@ async def _create_task(title, description, task_type, priority, mode, agent, wor source = TaskSource(user_prompt=prompt) if prompt else None - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title=title, description=description, mode=scheduler_mode, @@ -284,7 +284,7 @@ async def _create_task(title, description, task_type, priority, mode, agent, wor ) console.print(f"[green]✅ Created scheduler:[/green] {scheduler.id} {scheduler.title}") if trigger.run_immediately: - executions, _ = await TaskManager.list_scheduler_executions(scheduler.id, limit=1) + executions, _ = await ScheduleTaskManager.list_scheduler_executions(scheduler.id, limit=1) if executions: console.print( f"[green]↳ execution:[/green] {executions[0].id} ({executions[0].status.value})" @@ -306,20 +306,20 @@ def task_queue( async def _queue(pause: bool, resume: bool): await Storage.init() - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.store import TaskStore await TaskStore.init() if pause: - TaskManager.pause_queue() + ScheduleTaskManager.pause_queue() console.print("[yellow]Queue paused[/yellow]") return if resume: - TaskManager.resume_queue() + ScheduleTaskManager.resume_queue() console.print("[green]Queue resumed[/green]") return - qs = await TaskManager.queue_status() + qs = await ScheduleTaskManager.queue_status() console.print(Panel( f"Paused: {qs['paused']}\n" f"Max concurrent: {qs['max_concurrent']}\n" @@ -342,11 +342,11 @@ def task_scheduled(): async def _scheduled(): await Storage.init() - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.store import TaskStore await TaskStore.init() - tasks, _ = await TaskManager.list_schedulers(scheduled_only=True, limit=100) + tasks, _ = await ScheduleTaskManager.list_schedulers(scheduled_only=True, limit=100) if not tasks: console.print("[dim]No scheduled tasks[/dim]") return @@ -384,11 +384,11 @@ def task_cancel(task_id: str = typer.Argument(..., help="Task ID")): async def _cancel(task_id: str): await Storage.init() - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.store import TaskStore await TaskStore.init() - task = await TaskManager.cancel_execution(task_id) + task = await ScheduleTaskManager.cancel_execution(task_id) if not task: console.print(f"[red]Task {task_id} not found[/red]") raise typer.Exit(1) @@ -403,11 +403,11 @@ def task_retry(task_id: str = typer.Argument(..., help="Task ID")): async def _retry(task_id: str): await Storage.init() - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.store import TaskStore await TaskStore.init() - task = await TaskManager.retry_execution(task_id) + task = await ScheduleTaskManager.retry_execution(task_id) if not task: console.print(f"[red]Task {task_id} not found or not failed[/red]") raise typer.Exit(1) @@ -422,13 +422,13 @@ def task_rerun(task_id: str = typer.Argument(..., help="Task ID")): async def _rerun(task_id: str): await Storage.init() - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.store import TaskStore await TaskStore.init() - task = await TaskManager.rerun_execution(task_id) + task = await ScheduleTaskManager.rerun_execution(task_id) if task is None: - task = await TaskManager.rerun_scheduler(task_id) + task = await ScheduleTaskManager.rerun_scheduler(task_id) if not task: console.print(f"[red]Task {task_id} not found[/red]") raise typer.Exit(1) diff --git a/flocks/command/command.py b/flocks/command/command.py index c9322ba72..d38545dd3 100644 --- a/flocks/command/command.py +++ b/flocks/command/command.py @@ -24,7 +24,6 @@ class CommandDef: template: str agent: Optional[str] = None model: Optional[str] = None - subtask: Optional[bool] = None hidden: bool = False aliases: Tuple[str, ...] = field(default_factory=tuple) visible_surfaces: Tuple[CommandSurface, ...] = ("webui", "tui", "acp", "cli") diff --git a/flocks/config/config.py b/flocks/config/config.py index 684156edf..6024de221 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -158,7 +158,6 @@ class CommandConfig(BaseModel): description: Optional[str] = None agent: Optional[str] = None model: Optional[str] = None - subtask: Optional[bool] = None # ==================== Provider Configuration ==================== diff --git a/flocks/server/app.py b/flocks/server/app.py index 5a1635023..9d47761ec 100644 --- a/flocks/server/app.py +++ b/flocks/server/app.py @@ -350,17 +350,17 @@ async def _sync_workflows_phase() -> None: # Start Task Center (scheduler + queue executor) try: - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager await _run_startup_phase( log, - "task_manager.start", - TaskManager.start, + "schedule_task_manager.start", + ScheduleTaskManager.start, ) - log.info("task_manager.started") + log.info("schedule_task_manager.started") except Exception as e: - from flocks.task.manager import TaskManager - TaskManager.mark_start_failed(e) - log.warning("task_manager.start.failed", {"error": str(e)}) + from flocks.task.schedule_task_manager import ScheduleTaskManager + ScheduleTaskManager.mark_start_failed(e) + log.warning("schedule_task_manager.start.failed", {"error": str(e)}) # Seed built-in scheduled tasks from .flocks/plugins/tasks/*.json (idempotent) try: @@ -538,13 +538,13 @@ async def _delayed_trigger_runtime_start() -> None: # Stop Task Center try: - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.store import TaskStore - await TaskManager.stop() + await ScheduleTaskManager.stop() await TaskStore.close() - log.info("task_manager.stopped") + log.info("schedule_task_manager.stopped") except Exception as e: - log.warning("task_manager.stop.failed", {"error": str(e)}) + log.warning("schedule_task_manager.stop.failed", {"error": str(e)}) # Stop Skill file watcher try: diff --git a/flocks/server/routes/misc.py b/flocks/server/routes/misc.py index 2605ced40..230e0af2d 100644 --- a/flocks/server/routes/misc.py +++ b/flocks/server/routes/misc.py @@ -151,7 +151,6 @@ async def list_commands() -> List[Dict[str, Any]]: "template": cmd.template, "agent": cmd.agent, "model": cmd.model, - "subtask": cmd.subtask, "hidden": cmd.hidden, "aliases": list(cmd.aliases), "visible_surfaces": list(cmd.visible_surfaces), @@ -192,7 +191,6 @@ async def get_command(name: str) -> Dict[str, Any]: "template": cmd.template, "agent": cmd.agent, "model": cmd.model, - "subtask": cmd.subtask, "hidden": cmd.hidden, "aliases": list(cmd.aliases), "visible_surfaces": list(cmd.visible_surfaces), @@ -241,4 +239,3 @@ async def list_experimental_resources() -> Dict[str, Any]: # Return empty dict - resources are not implemented yet return {} - diff --git a/flocks/server/routes/session.py b/flocks/server/routes/session.py index 5fd1a6f3c..246222b6c 100644 --- a/flocks/server/routes/session.py +++ b/flocks/server/routes/session.py @@ -2075,15 +2075,6 @@ class AgentPartInput(BaseModel): name: str = Field(..., description="Agent name") -class SubtaskPartInput(BaseModel): - """Subtask part input for API compatibility""" - type: Literal["subtask"] = "subtask" - id: Optional[str] = Field(None, description="Part ID") - agent: str = Field(..., description="Agent name") - prompt: str = Field(..., description="Subtask prompt") - description: Optional[str] = Field(None, description="Subtask description") - - class PromptRequest(BaseModel): """ Request to send a prompt/message diff --git a/flocks/server/routes/skill.py b/flocks/server/routes/skill.py index 28d0e3d7f..4d16ee70c 100644 --- a/flocks/server/routes/skill.py +++ b/flocks/server/routes/skill.py @@ -164,7 +164,6 @@ class CommandResponse(BaseModel): template: str = Field(..., description="Command template") agent: Optional[str] = Field(None, description="Preferred agent") model: Optional[str] = Field(None, description="Preferred model") - subtask: Optional[bool] = Field(None, description="Run as subtask") hidden: bool = Field(False, description="Hidden from UI") aliases: List[str] = Field(default_factory=list, description="Alternate slash aliases") visible_surfaces: List[str] = Field(default_factory=list, description="Surfaces where the command is visible") @@ -182,7 +181,6 @@ def _command_to_response(cmd: CommandInfo) -> CommandResponse: template=cmd.template, agent=cmd.agent, model=cmd.model, - subtask=cmd.subtask, hidden=cmd.hidden, aliases=list(cmd.aliases), visible_surfaces=list(cmd.visible_surfaces), diff --git a/flocks/server/routes/stats.py b/flocks/server/routes/stats.py index af2e74766..205ce1be9 100644 --- a/flocks/server/routes/stats.py +++ b/flocks/server/routes/stats.py @@ -19,7 +19,7 @@ from flocks.server.routes.provider import list_providers from flocks.server.routes.workflow import _list_workflows_from_fs, _migrate_storage_to_filesystem from flocks.skill.skill import Skill -from flocks.task.manager import TaskManager +from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.tool.registry import ToolRegistry from flocks.utils.log import Log @@ -70,7 +70,7 @@ def _should_count_agent(agent: Any) -> bool: async def _task_dashboard() -> dict[str, Any]: - return await TaskManager.dashboard() + return await ScheduleTaskManager.dashboard() async def _safe_dashboard(failures: list[str]) -> dict[str, Any]: diff --git a/flocks/server/routes/task_entities.py b/flocks/server/routes/task_entities.py index 69170c4dd..18bbf35d6 100644 --- a/flocks/server/routes/task_entities.py +++ b/flocks/server/routes/task_entities.py @@ -145,10 +145,10 @@ def _parse_task_type(task_type: str) -> str: @router.get("/task-system/notice") async def get_task_system_notice(): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager started_at = time.perf_counter() - notice = await TaskManager.get_task_page_notice() + notice = await ScheduleTaskManager.get_task_page_notice() log_route_timing(log, "task.notice.complete", started_at=started_at, extra={ "has_notice": bool(notice), }) @@ -157,10 +157,10 @@ async def get_task_system_notice(): @router.get("/task-system/dashboard") async def task_dashboard(): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager started_at = time.perf_counter() - payload = await TaskManager.dashboard() + payload = await ScheduleTaskManager.dashboard() log_route_timing(log, "task.dashboard.complete", started_at=started_at, extra={ "running": payload.get("running"), "queued": payload.get("queued"), @@ -171,10 +171,10 @@ async def task_dashboard(): @router.get("/task-system/queue/status") async def task_queue_status(): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager started_at = time.perf_counter() - payload = await TaskManager.queue_status() + payload = await ScheduleTaskManager.queue_status() log_route_timing(log, "task.queue_status.complete", started_at=started_at, extra={ "queued": payload.get("queued"), "running": payload.get("running"), @@ -185,17 +185,17 @@ async def task_queue_status(): @router.post("/task-system/queue/pause") async def pause_task_queue(): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - TaskManager.pause_queue() + ScheduleTaskManager.pause_queue() return {"paused": True} @router.post("/task-system/queue/resume") async def resume_task_queue(): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - TaskManager.resume_queue() + ScheduleTaskManager.resume_queue() return {"paused": False} @@ -209,9 +209,9 @@ async def list_schedulers( offset: int = Query(0, ge=0), limit: int = Query(20, ge=1, le=100), ): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - items, total = await TaskManager.list_schedulers( + items, total = await ScheduleTaskManager.list_schedulers( status=_parse_scheduler_status_filter(status_filter), priority=_parse_priority(priority), scheduled_only=scheduled_only, @@ -230,7 +230,7 @@ async def list_schedulers( @router.post("/task-schedulers", status_code=status.HTTP_201_CREATED) async def create_scheduler(req: SchedulerCreateRequest): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.models import ( SchedulerMode, TaskSource, @@ -268,7 +268,7 @@ async def create_scheduler(req: SchedulerCreateRequest): detail=str(exc), ) from exc - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title=req.title, description=req.description, mode=mode, @@ -289,9 +289,9 @@ async def create_scheduler(req: SchedulerCreateRequest): @router.get("/task-schedulers/{scheduler_id}") async def get_scheduler(scheduler_id: str): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - scheduler = await TaskManager.get_scheduler(scheduler_id) + scheduler = await ScheduleTaskManager.get_scheduler(scheduler_id) if not scheduler: raise HTTPException(404, "Task scheduler not found") return scheduler.model_dump(mode="json", by_alias=True) @@ -299,7 +299,7 @@ async def get_scheduler(scheduler_id: str): @router.put("/task-schedulers/{scheduler_id}") async def update_scheduler(scheduler_id: str, req: SchedulerUpdateRequest): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager fields = {k: v for k, v in req.model_dump(exclude_none=True).items()} if "priority" in fields: @@ -313,7 +313,7 @@ async def update_scheduler(scheduler_id: str, req: SchedulerUpdateRequest): run_at = fields.pop("run_at", None) user_prompt = fields.pop("user_prompt", None) try: - scheduler = await TaskManager.update_scheduler_with_trigger( + scheduler = await ScheduleTaskManager.update_scheduler_with_trigger( scheduler_id, fields=fields, cron=cron, @@ -335,18 +335,18 @@ async def update_scheduler(scheduler_id: str, req: SchedulerUpdateRequest): @router.delete("/task-schedulers/{scheduler_id}") async def delete_scheduler(scheduler_id: str): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - if not await TaskManager.delete_scheduler(scheduler_id): + if not await ScheduleTaskManager.delete_scheduler(scheduler_id): raise HTTPException(404, "Task scheduler not found") return {"ok": True} @router.post("/task-schedulers/{scheduler_id}/enable") async def enable_scheduler(scheduler_id: str): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - scheduler = await TaskManager.enable_scheduler(scheduler_id) + scheduler = await ScheduleTaskManager.enable_scheduler(scheduler_id) if not scheduler: raise HTTPException(404, "Task scheduler not found") return scheduler.model_dump(mode="json", by_alias=True) @@ -354,9 +354,9 @@ async def enable_scheduler(scheduler_id: str): @router.post("/task-schedulers/{scheduler_id}/disable") async def disable_scheduler(scheduler_id: str): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - scheduler = await TaskManager.disable_scheduler(scheduler_id) + scheduler = await ScheduleTaskManager.disable_scheduler(scheduler_id) if not scheduler: raise HTTPException(404, "Task scheduler not found") return scheduler.model_dump(mode="json", by_alias=True) @@ -368,9 +368,9 @@ async def list_scheduler_executions( offset: int = Query(0, ge=0), limit: int = Query(20, ge=1, le=100), ): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - items, total = await TaskManager.list_scheduler_executions( + items, total = await ScheduleTaskManager.list_scheduler_executions( scheduler_id, offset=offset, limit=limit ) return PaginatedResponse( @@ -383,9 +383,9 @@ async def list_scheduler_executions( @router.post("/task-schedulers/{scheduler_id}/run") async def run_scheduler(scheduler_id: str): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - execution = await TaskManager.rerun_scheduler(scheduler_id) + execution = await ScheduleTaskManager.rerun_scheduler(scheduler_id) if not execution: raise HTTPException(404, "Task scheduler not found") return execution.model_dump(mode="json", by_alias=True) @@ -402,9 +402,9 @@ async def list_executions( offset: int = Query(0, ge=0), limit: int = Query(20, ge=1, le=100), ): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - items, total = await TaskManager.list_executions( + items, total = await ScheduleTaskManager.list_executions( scheduler_id=scheduler_id, status=_parse_execution_status_filter(status_filter), priority=_parse_priority(priority), @@ -424,23 +424,23 @@ async def list_executions( @router.post("/task-executions/batch/cancel") async def batch_cancel(req: BatchRequest): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - return {"cancelled": await TaskManager.batch_cancel(req.execution_ids)} + return {"cancelled": await ScheduleTaskManager.batch_cancel(req.execution_ids)} @router.post("/task-executions/batch/delete") async def batch_delete(req: BatchRequest): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - return {"deleted": await TaskManager.batch_delete(req.execution_ids)} + return {"deleted": await ScheduleTaskManager.batch_delete(req.execution_ids)} @router.get("/task-executions/{execution_id}") async def get_execution(execution_id: str): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - execution = await TaskManager.get_execution(execution_id) + execution = await ScheduleTaskManager.get_execution(execution_id) if not execution: raise HTTPException(404, "Task execution not found") return execution.model_dump(mode="json", by_alias=True) @@ -448,9 +448,9 @@ async def get_execution(execution_id: str): @router.post("/task-executions/{execution_id}/viewed") async def mark_execution_viewed(execution_id: str): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - execution = await TaskManager.mark_viewed(execution_id) + execution = await ScheduleTaskManager.mark_viewed(execution_id) if not execution: raise HTTPException(404, "Task execution not found") return execution.model_dump(mode="json", by_alias=True) @@ -458,18 +458,18 @@ async def mark_execution_viewed(execution_id: str): @router.post("/task-executions/{execution_id}/cancel") async def cancel_execution(execution_id: str): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - execution = await TaskManager.cancel_execution(execution_id) + execution = await ScheduleTaskManager.cancel_execution(execution_id) if not execution: raise HTTPException(404, "Task execution not found") return execution.model_dump(mode="json", by_alias=True) @router.post("/task-executions/{execution_id}/retry") async def retry_execution(execution_id: str): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - execution = await TaskManager.retry_execution(execution_id) + execution = await ScheduleTaskManager.retry_execution(execution_id) if not execution: raise HTTPException(404, "Task execution not found") return execution.model_dump(mode="json", by_alias=True) @@ -477,9 +477,9 @@ async def retry_execution(execution_id: str): @router.post("/task-executions/{execution_id}/rerun") async def rerun_execution(execution_id: str): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - execution = await TaskManager.rerun_execution(execution_id) + execution = await ScheduleTaskManager.rerun_execution(execution_id) if not execution: raise HTTPException(404, "Task execution not found") return execution.model_dump(mode="json", by_alias=True) @@ -487,10 +487,9 @@ async def rerun_execution(execution_id: str): @router.delete("/task-executions/{execution_id}") async def delete_execution(execution_id: str): - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - if not await TaskManager.delete_execution(execution_id): + if not await ScheduleTaskManager.delete_execution(execution_id): raise HTTPException(404, "Task execution not found") return {"ok": True} - diff --git a/flocks/session/__init__.py b/flocks/session/__init__.py index 2dc887f82..086208d66 100644 --- a/flocks/session/__init__.py +++ b/flocks/session/__init__.py @@ -30,7 +30,6 @@ ReasoningPart, PatchPart, AgentPart, - SubtaskPart, ) from flocks.session.prompt import SessionPrompt, SystemPrompt, ContextInfo from flocks.session.lifecycle.compaction import SessionCompaction, CompactionResult, CompactionPolicy, ContextTier @@ -46,11 +45,6 @@ ReminderConfig, ReminderContext, ) -from flocks.session.features.subtask import ( - SessionSubtask, - SubtaskInfo, - SubtaskResult, -) from flocks.session.lifecycle.revert import ( SessionRevertManager, RevertInput, @@ -86,7 +80,6 @@ "ReasoningPart", "PatchPart", "AgentPart", - "SubtaskPart", # Prompt "SessionPrompt", "SystemPrompt", @@ -106,10 +99,6 @@ "SessionReminders", "ReminderConfig", "ReminderContext", - # Subtask - "SessionSubtask", - "SubtaskInfo", - "SubtaskResult", # Revert "SessionRevertManager", "RevertInput", diff --git a/flocks/session/context_usage.py b/flocks/session/context_usage.py index ca35e09d1..376b41044 100644 --- a/flocks/session/context_usage.py +++ b/flocks/session/context_usage.py @@ -23,7 +23,7 @@ log = Log.create(service="context-usage") UsageSource = Literal["observed", "estimated"] -DELEGATION_TOOLS = {"delegate_task", "task"} +DELEGATION_TOOLS = {"delegate_task"} ZERO_VISIBLE_SEGMENTS = {"agentDelegation"} @@ -441,8 +441,8 @@ async def _estimate_message_breakdown(session_id: str, messages: List[Any]) -> t if part_type in {"reasoning", "thinking"}: tokens_by_key["reasoning"] += SessionPrompt.count_tokens(_field_value(part, "text", "") or "") continue - if part_type in {"agent", "subtask"}: - tokens_by_key["agentDelegation"] += _estimate_subtask_part_tokens(part) + if part_type == "agent": + tokens_by_key["agentDelegation"] += _estimate_agent_part_tokens(part) continue if part_type != "tool": continue @@ -499,7 +499,7 @@ def _context_key_for_tool(tool_name: str) -> str: return "tools" -def _estimate_subtask_part_tokens(part: Any) -> int: +def _estimate_agent_part_tokens(part: Any) -> int: total = 0 for field in ("prompt", "description", "name"): value = _field_value(part, field, "") diff --git a/flocks/session/execution_mode.py b/flocks/session/execution_mode.py index 13d12b154..00952ff5c 100644 --- a/flocks/session/execution_mode.py +++ b/flocks/session/execution_mode.py @@ -28,7 +28,7 @@ class SessionExecutionMode(str, Enum): "run_slash_command", } ) -PLAN_DELEGATION_TOOL_NAMES = frozenset({"delegate_task", "task"}) +PLAN_DELEGATION_TOOL_NAMES = frozenset({"delegate_task"}) PLAN_DELEGATABLE_AGENT_NAMES = frozenset({"explore", "librarian"}) PLAN_PATH_SCOPED_TOOL_NAMES = frozenset({"apply_patch", "edit", "write"}) diff --git a/flocks/session/features/subtask.py b/flocks/session/features/subtask.py deleted file mode 100644 index b7aac9ca2..000000000 --- a/flocks/session/features/subtask.py +++ /dev/null @@ -1,62 +0,0 @@ -""" -Session Subtask data models. - -Note: The SessionSubtask business logic has been removed as it was dead code. -The live subtask execution path is session_loop.py::_execute_subtask(), which -handles the full lifecycle inline without using this module. - -These data classes are kept because they are exported from session/__init__.py -and may be referenced by external consumers. -""" - -from dataclasses import dataclass, field -from datetime import datetime -from typing import Any, Dict, Optional - - -@dataclass -class SubtaskInfo: - """Information about a subtask""" - id: str - parent_session_id: str - child_session_id: Optional[str] = None - task_description: str = "" - agent: Optional[str] = None - model: Optional[str] = None - status: str = "pending" # pending, running, completed, error - result: Optional[str] = None - error: Optional[str] = None - created_at: int = field(default_factory=lambda: int(datetime.now().timestamp() * 1000)) - completed_at: Optional[int] = None - - -@dataclass -class SubtaskResult: - """Result of subtask execution""" - subtask_id: str - success: bool - output: str - error: Optional[str] = None - metadata: Dict[str, Any] = field(default_factory=dict) - - -# Minimal stub so imports of SessionSubtask don't break existing code. -class SessionSubtask: - """Subtask manager stub — business logic removed (was dead code). - - The active execution path is SessionTurn._execute_subtask(). - """ - - @classmethod - async def execute_subtask(cls, *args, **kwargs) -> SubtaskResult: - raise NotImplementedError( - "SessionSubtask.execute_subtask() is deprecated. " - "Subtask execution is handled by SessionTurn._execute_subtask()." - ) - - -__all__ = [ - "SessionSubtask", - "SubtaskInfo", - "SubtaskResult", -] diff --git a/flocks/session/message.py b/flocks/session/message.py index a8c4b3e20..f3ac17bef 100644 --- a/flocks/session/message.py +++ b/flocks/session/message.py @@ -259,21 +259,6 @@ class AgentPart(BaseModel): source: Optional[Dict[str, Any]] = Field(None, description="Source information") -class SubtaskPart(BaseModel): - """Subtask/subagent part - Flocks compatible""" - model_config = ConfigDict(populate_by_name=True, by_alias=True) - - id: str = Field(default_factory=lambda: Identifier.ascending("part")) - sessionID: str = Field(..., description="Session ID") - messageID: str = Field(..., description="Message ID") - type: Literal["subtask"] = "subtask" - prompt: str = Field(..., description="Task prompt") - description: str = Field(..., description="Task description") - agent: str = Field(..., description="Agent name") - model: Optional[Dict[str, str]] = Field(None, description="Model configuration") - command: Optional[str] = Field(None, description="Command to execute") - - class RetryPart(BaseModel): """Retry part - Flocks compatible""" model_config = ConfigDict(populate_by_name=True, by_alias=True) @@ -301,7 +286,6 @@ class CompactionPart(BaseModel): # Union type for all parts - matches Flocks MessageV2.Part PartType = Union[ TextPart, - SubtaskPart, ReasoningPart, FilePart, ToolPart, @@ -1122,6 +1106,18 @@ def _normalize_part_data( metadata = normalized.get("metadata") metadata_dict = metadata if isinstance(metadata, dict) else {} + if part_type == "subtask": + normalized["type"] = "text" + normalized["text"] = "" + normalized["ignored"] = True + normalized["metadata"] = { + **metadata_dict, + "legacyPartType": "subtask", + } + part_type = "text" + metadata = normalized["metadata"] + metadata_dict = metadata + if "content" in normalized and "text" not in normalized: normalized["text"] = normalized.get("content", "") @@ -1174,10 +1170,6 @@ def _normalize_part_data( normalized.setdefault("tokens", metadata_dict.get("tokens") or cls._default_token_usage()) elif part_type == "agent": normalized.setdefault("name", metadata_dict.get("name") or normalized.get("content") or "agent") - elif part_type == "subtask": - normalized.setdefault("prompt", metadata_dict.get("prompt") or normalized.get("content", "")) - normalized.setdefault("description", metadata_dict.get("description") or "") - normalized.setdefault("agent", metadata_dict.get("agent") or "agent") elif part_type == "retry": normalized.setdefault("attempt", metadata_dict.get("attempt") or 1) normalized.setdefault("error", metadata_dict.get("error") or {}) @@ -1255,7 +1247,6 @@ def deserialize_part( 'step-start': StepStartPart, 'step-finish': StepFinishPart, 'agent': AgentPart, - 'subtask': SubtaskPart, 'retry': RetryPart, 'compaction': CompactionPart, } diff --git a/flocks/session/prompt/anthropic-20250930.txt b/flocks/session/prompt/anthropic-20250930.txt index a8ada5ede..ec4e65c94 100644 --- a/flocks/session/prompt/anthropic-20250930.txt +++ b/flocks/session/prompt/anthropic-20250930.txt @@ -122,10 +122,10 @@ I've found existing rules. Let me mark the first todo as in_progress and start d Users may configure 'hooks', shell commands that execute in response to events like tool calls, in settings. Treat feedback from hooks, including , as coming from the user. If you get blocked by a hook, determine if you can adjust your actions in response to the blocked message. If not, ask the user to check their hooks configuration. # Tool usage policy -- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - Tool results and user messages may include tags. tags contain useful information and reminders. They are automatically added by the system, and bear no direct relation to the specific tool results or user messages in which they appear. - When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. diff --git a/flocks/session/prompt/anthropic.txt b/flocks/session/prompt/anthropic.txt index 655551871..2455ef584 100644 --- a/flocks/session/prompt/anthropic.txt +++ b/flocks/session/prompt/anthropic.txt @@ -66,22 +66,22 @@ I've found existing SIGMA rules. Let me mark the first todo as in_progress and s # Tool usage policy -- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - Tool results and user messages may include tags. tags contain useful information and reminders. They are automatically added by the system, and bear no direct relation to the specific tool results or user messages in which they appear. - When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. - You can call multiple tools in a single response. If you intend to call multiple tools and there are no dependencies between them, make all independent tool calls in parallel. Maximize use of parallel tool calls where possible to increase efficiency. However, if some tool calls depend on previous calls to inform dependent values, do NOT call these tools in parallel and instead call them sequentially. For instance, if one operation must complete before another starts, run these operations sequentially instead. Never use placeholders or guess missing parameters in tool calls. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. -- VERY IMPORTANT: When exploring security logs, configurations, or investigating incidents that require context gathering, use the Task tool for complex searches instead of running multiple search commands directly. +- VERY IMPORTANT: When exploring security logs, configurations, or investigating incidents that require context gathering, use `delegate_task` for complex searches instead of running multiple search commands directly. - IMPORTANT: Always respond in the same language as the user. user: Where are authentication failures logged in our application? -assistant: [Uses the Task tool to find authentication logging locations instead of using Glob or Grep directly] +assistant: [Uses `delegate_task` to find authentication logging locations instead of using Glob or Grep directly] user: Find all places where user input is processed without validation -assistant: [Uses the Task tool to comprehensively search for input validation gaps] +assistant: [Uses `delegate_task` to comprehensively search for input validation gaps] IMPORTANT: Always use `todo(action="write")` to plan and track tasks throughout the conversation. diff --git a/flocks/session/runtime/session_turn.py b/flocks/session/runtime/session_turn.py index 666c9175e..f970a5e22 100644 --- a/flocks/session/runtime/session_turn.py +++ b/flocks/session/runtime/session_turn.py @@ -4,7 +4,6 @@ - Message processing - Tool execution - Compaction -- Subtask handling - Reminders """ @@ -26,7 +25,6 @@ TurnPreparationStatus, ) from flocks.utils.log import Log -from flocks.utils.id import Identifier from flocks.session.session import ( Session, SessionInfo, @@ -98,7 +96,6 @@ class SessionTurn: Supports: - Message iteration - Compaction triggers - - Subtask management - Reminder injection """ @@ -259,7 +256,7 @@ async def prepare_step( last_user: Optional[MessageInfo] = None last_assistant: Optional[MessageInfo] = None last_finished: Optional[MessageInfo] = None - tasks: List[tuple[str, Any]] = [] + pending_compactions: List[Any] = [] scan_started_at = asyncio.get_running_loop().time() for message in reversed(messages): if last_user is None and message.role == MessageRole.USER: @@ -273,15 +270,13 @@ async def prepare_step( if last_finished is None: for part in await Message.parts(message.id, self.session.id): if part.type == "compaction": - tasks.append(("compaction", part)) - elif part.type == "subtask": - tasks.append(("subtask", part)) + pending_compactions.append(part) log.debug( "loop.message_scan_complete", { "session_id": self.session.id, "step": self.step, - "task_count": len(tasks), + "compaction_count": len(pending_compactions), "duration_ms": int((asyncio.get_running_loop().time() - scan_started_at) * 1000), }, ) @@ -325,14 +320,14 @@ async def prepare_step( await self._prepare_memory() self._schedule_title_generation(last_user, messages) - if tasks: - task_preparation = await self._prepare_pending_task( + if pending_compactions: + compaction_preparation = await self._prepare_pending_compaction( messages, last_user, - tasks.pop(), + pending_compactions.pop(), ) - if task_preparation is not None: - return task_preparation + if compaction_preparation is not None: + return compaction_preparation context_preparation = await self._prepare_context_window( messages, @@ -530,28 +525,19 @@ def _schedule_title_generation( except Exception as exc: log.error("loop.title_generation.error", {"error": str(exc)}) - async def _prepare_pending_task( + async def _prepare_pending_compaction( self, messages: List[MessageInfo], last_user: MessageInfo, - task: tuple[str, Any], + compaction_part: Any, ) -> Optional[ModelTurnPreparation[MessageInfo]]: - """Finish persisted subtask or compaction work before the model turn.""" - task_type, task_part = task - if task_type == "subtask": - log.info( - "loop.subtask_detected", - {"session_id": self.session.id, "step": self.step}, - ) - await self._execute_subtask(last_user, task_part) - return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) - + """Finish persisted compaction work before the model turn.""" log.info( "loop.compaction_pending", { "session_id": self.session.id, "step": self.step, - "auto": getattr(task_part, "auto", False), + "auto": getattr(compaction_part, "auto", False), }, ) if self.callbacks.on_compaction: @@ -578,7 +564,7 @@ async def progress_callback(stage: str, data: dict) -> None: messages=messages, provider_id=self.provider_id, model_id=self.model_id, - auto=getattr(task_part, "auto", False), + auto=getattr(compaction_part, "auto", False), event_publish_callback=publish, status_after="busy", policy=self._build_compaction_policy(), @@ -1081,192 +1067,3 @@ async def _check_reminders( await self.callbacks.on_reminder( await Message.get_text_content(reminder_msg), ) - - async def _execute_subtask( - self, - last_user: MessageInfo, - task_part: Any, - ) -> None: - """ - Execute subtask (matching TUI lines 316-481) - - 完全匹配 TUI 的 subtask 执行流程: - 1. 创建 assistant message - 2. 创建 tool part (Task tool) - 3. 执行 Task tool - 4. 更新 part 状态 - 5. 创建 synthetic user message - """ - from flocks.tool.registry import ToolRegistry - from flocks.agent.registry import Agent - - # Extract subtask information from part - agent_name = getattr(task_part, "agent", "hephaestus") - prompt = getattr(task_part, "prompt", "") - description = getattr(task_part, "description", "") - command = getattr(task_part, "command", None) - model_info = getattr(task_part, "model", None) - - # Get agent - agent = await Agent.get(agent_name) or await Agent.get("rex") - - # Determine model - if model_info: - provider_id = model_info.get("providerID", self.provider_id) - model_id = model_info.get("modelID", self.model_id) - else: - provider_id = self.provider_id - model_id = self.model_id - - # Create assistant message for subtask - assistant_msg = await Message.create( - session_id=self.session.id, - role=MessageRole.ASSISTANT, - content="", - agent=agent_name, - model=model_id, - provider=provider_id, - parent_id=last_user.id, - ) - - # Create tool part for Task - tool_call_id = Identifier.create("call") - from flocks.session.message import ToolPart, ToolStateRunning - - tool_part = ToolPart( - id=Identifier.ascending("part"), - sessionID=self.session.id, - messageID=assistant_msg.id, - type="tool", - callID=tool_call_id, - tool="task", - state=ToolStateRunning( - status="running", - input={ - "prompt": prompt, - "description": description, - "subagent_type": agent_name, - "command": command, - }, - time={"start": int(datetime.now().timestamp() * 1000)}, - ), - ) - - # Add part to message - await Message.add_part(self.session.id, assistant_msg.id, tool_part) - - # Get Task tool - task_tool = ToolRegistry.get("task") - if not task_tool: - log.error("loop.subtask.task_tool_not_found", {"session_id": self.session.id}) - return - - # Execute Task tool - task_args = { - "prompt": prompt, - "description": description, - "subagent_type": agent_name, - "command": command, - } - - # Create tool context - from flocks.tool.registry import ToolContext - - tool_ctx = ToolContext( - session_id=self.session.id, - message_id=assistant_msg.id, - agent=agent_name, - abort_event=self.abort_event, - ) - - execution_error: Optional[Exception] = None - result = None - - try: - result = await task_tool.execute(tool_ctx, **task_args) - except Exception as e: - execution_error = e - log.error( - "loop.subtask.execution_failed", - { - "error": str(e), - "agent": agent_name, - "description": description, - }, - ) - - # Update message finish - await Message.update(self.session.id, assistant_msg.id, finish="tool-calls") - - # Update tool part status - from flocks.session.message import ToolStateCompleted, ToolStateError - - if result: - # Create completed state - completed_state = ToolStateCompleted( - status="completed", - input={ - "prompt": prompt, - "description": description, - "subagent_type": agent_name, - "command": command, - }, - output=result.output if hasattr(result, "output") else str(result), - title=result.title if hasattr(result, "title") else None, - metadata=result.metadata if hasattr(result, "metadata") else {}, - time={ - "start": tool_part.state.time.get("start"), - "end": int(datetime.now().timestamp() * 1000), - }, - ) - await Message.update_part( - session_id=self.session.id, - message_id=assistant_msg.id, - part_id=tool_part.id, - state=completed_state, - ) - else: - # Create error state - error_msg = str(execution_error) if execution_error else "Tool execution failed" - error_state = ToolStateError( - status="error", - error=f"Tool execution failed: {error_msg}", - time={ - "start": tool_part.state.time.get("start"), - "end": int(datetime.now().timestamp() * 1000), - }, - metadata={}, - input={ - "prompt": prompt, - "description": description, - "subagent_type": agent_name, - "command": command, - }, - ) - await Message.update_part( - session_id=self.session.id, - message_id=assistant_msg.id, - part_id=tool_part.id, - state=error_state, - ) - - # Create synthetic user message (matching TUI lines 457-478) - # This prevents reasoning models from erroring due to missing user messages - synthetic_user_msg = await Message.create( - session_id=self.session.id, - role=MessageRole.USER, - content="Summarize the task tool output above and continue with your task.", - agent=last_user.agent if hasattr(last_user, "agent") else agent_name, - model=last_user.model if hasattr(last_user, "model") else model_id, - provider=last_user.provider if hasattr(last_user, "provider") else provider_id, - synthetic=True, - ) - - log.info( - "loop.subtask.completed", - { - "session_id": self.session.id, - "agent": agent_name, - "success": result is not None, - }, - ) diff --git a/flocks/session/runtime/step_engine.py b/flocks/session/runtime/step_engine.py index e58095e26..caec9f5da 100644 --- a/flocks/session/runtime/step_engine.py +++ b/flocks/session/runtime/step_engine.py @@ -2749,12 +2749,6 @@ async def _to_chat_messages( "type": "text", "text": "What did we do so far?", }) - elif part.type == "subtask": - user_content_parts.append("The following tool was executed by the user") - user_content_blocks.append({ - "type": "text", - "text": "The following tool was executed by the user", - }) if user_content_blocks and any( block.get("type") == "image" diff --git a/flocks/session/streaming/stream_processor.py b/flocks/session/streaming/stream_processor.py index 05397ab4a..aa80766db 100644 --- a/flocks/session/streaming/stream_processor.py +++ b/flocks/session/streaming/stream_processor.py @@ -1549,7 +1549,7 @@ async def _handle_text_end(self, event: TextEndEvent) -> None: def _should_run_tool_call_parallel(self, event: ToolCallEvent) -> bool: """Return true for independent foreground subagent tool-calls.""" - if event.tool_name not in {"delegate_task", "task"}: + if event.tool_name != "delegate_task": return False tool_input = event.input if isinstance(event.input, dict) else {} if tool_input.get("run_in_background") is True: diff --git a/flocks/task/__init__.py b/flocks/task/__init__.py index 180792a9d..28b314739 100644 --- a/flocks/task/__init__.py +++ b/flocks/task/__init__.py @@ -20,7 +20,7 @@ SchedulerStatus, build_schedule, ) -from .manager import TaskManager +from .schedule_task_manager import ScheduleTaskManager from .store import TaskStore __all__ = [ @@ -30,7 +30,7 @@ "RetryConfig", "TaskExecution", "TaskExecutionQueueRef", - "TaskManager", + "ScheduleTaskManager", "TaskPriority", "TaskScheduler", "TaskTrigger", diff --git a/flocks/task/plugin_sync.py b/flocks/task/plugin_sync.py index 1c27a582a..1483ad538 100644 --- a/flocks/task/plugin_sync.py +++ b/flocks/task/plugin_sync.py @@ -12,7 +12,7 @@ async def upsert_task_specs(specs: Sequence[TaskSpec]) -> int: - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.models import ( ExecutionMode, SchedulerMode, @@ -89,7 +89,7 @@ async def upsert_task_specs(specs: Sequence[TaskSpec]) -> int: log.warn("task.plugin.missing_cron", {"dedup_key": spec.dedup_key}) continue - await TaskManager.create_scheduler( + await ScheduleTaskManager.create_scheduler( title=spec.title, description=spec.description, mode=SchedulerMode.CRON, diff --git a/flocks/task/manager.py b/flocks/task/schedule_task_manager.py similarity index 99% rename from flocks/task/manager.py rename to flocks/task/schedule_task_manager.py index de038a25a..c769f4e6f 100644 --- a/flocks/task/manager.py +++ b/flocks/task/schedule_task_manager.py @@ -1,4 +1,4 @@ -"""Task Manager for scheduler/execution domain.""" +"""Schedule task manager for the scheduler/execution domain.""" import asyncio import json @@ -32,7 +32,7 @@ from .scheduler import TaskScheduler as SchedulerLoop from .store import TaskStore -log = Log.create(service="task.manager") +log = Log.create(service="task.schedule_manager") _TASK_EXPIRY_HOURS: int = 24 _CLEANUP_INTERVAL_S: int = 3600 @@ -49,8 +49,8 @@ class _TaskEventProps(_BaseModel): title: str -class TaskManager: - _instance: Optional["TaskManager"] = None +class ScheduleTaskManager: + _instance: Optional["ScheduleTaskManager"] = None _startup_error: Optional[str] = None def __init__( @@ -83,7 +83,7 @@ async def start( max_concurrent: int = 4, poll_interval: int = 5, scheduler_interval: int = 30, - ) -> "TaskManager": + ) -> "ScheduleTaskManager": if cls._instance and cls._instance._running: return cls._instance await TaskStore.init() @@ -128,7 +128,7 @@ async def stop(cls) -> None: log.info("manager.stopped") @classmethod - def get(cls) -> Optional["TaskManager"]: + def get(cls) -> Optional["ScheduleTaskManager"]: return cls._instance @classmethod diff --git a/flocks/task/scheduler.py b/flocks/task/scheduler.py index b8035dead..75acf6619 100644 --- a/flocks/task/scheduler.py +++ b/flocks/task/scheduler.py @@ -60,7 +60,7 @@ async def _loop(self) -> None: await asyncio.sleep(self._check_interval) async def _tick(self) -> None: - from .manager import TaskManager + from .schedule_task_manager import ScheduleTaskManager now = datetime.now(timezone.utc) schedulers = await TaskStore.list_due_schedulers() @@ -77,7 +77,7 @@ async def _tick(self) -> None: if scheduler.mode == SchedulerMode.ONCE else ExecutionTriggerType.SCHEDULED ) - await TaskManager.create_execution_from_scheduler( + await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=trigger_type, enqueue=True, diff --git a/flocks/tool/agent/delegate_task.py b/flocks/tool/agent/delegate_task.py index 0ee80c8df..b31cef7a4 100644 --- a/flocks/tool/agent/delegate_task.py +++ b/flocks/tool/agent/delegate_task.py @@ -329,7 +329,7 @@ def _derive_task_description( - Background subagent execution is disabled. Do not set run_in_background=true. - Foreground execution is always used: the tool waits for completion and returns results inline. - For independent parallel work needed this turn, emit multiple sibling - foreground delegate_task/task tool calls in the same assistant response. + foreground delegate_task tool calls in the same assistant response. The runtime executes them concurrently and the webui renders each as its own DelegateTaskCard. @@ -395,9 +395,8 @@ async def delegate_task_tool( load_skills: Optional[List[str]] = None, description: Optional[str] = None, # Internal-only: not exposed in the public schema. The registry rejects - # `run_in_background=True` at the schema layer for any caller, but legacy - # in-process call paths (e.g. `task.py` alias) may still pass it through. - # This guard is the second line of defense. + # `run_in_background=True` at the schema layer for any caller. This guard + # also protects direct in-process callers that bypass the registry. run_in_background: bool = False, subagent_type: Optional[str] = None, session_id: Optional[str] = None, @@ -409,7 +408,7 @@ async def delegate_task_tool( success=False, error=( "Background subagent execution is disabled. " - "Use foreground delegate_task/task calls; emit multiple sibling calls " + "Use foreground delegate_task calls; emit multiple sibling calls " "in the same assistant turn for parallel work." ), ) diff --git a/flocks/tool/agent/task.py b/flocks/tool/agent/task.py deleted file mode 100644 index 891231729..000000000 --- a/flocks/tool/agent/task.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Compatibility alias for delegate_task. - -The runtime keeps ``task`` as a registered tool name for workflow/backward -compatibility, but all scheduling behavior lives in ``delegate_task``. -Background subagent execution is disabled; run synchronously and emit -multiple sibling tool calls in one assistant turn for parallel work. -""" - -from __future__ import annotations - -from typing import Optional - -from flocks.tool.agent.delegate_task import delegate_task_tool -from flocks.tool.registry import ( - ParameterType, - ToolCategory, - ToolContext, - ToolParameter, - ToolRegistry, - ToolResult, -) - - -DESCRIPTION = """Compatibility alias for delegate_task. - -Use delegate_task directly for new prompts. Workflows may continue using task; -it accepts the same single-subagent shape and forwards it to delegate_task. -Background subagent execution is disabled; run synchronously and emit multiple -sibling tool calls in one assistant turn for parallel work. -""" - - -@ToolRegistry.register_function( - name="task", - description=DESCRIPTION, - category=ToolCategory.SYSTEM, - native=False, - parameters=[ - ToolParameter( - name="description", - type=ParameterType.STRING, - description="Optional short task description (3-5 words)", - required=False, - ), - ToolParameter( - name="prompt", - type=ParameterType.STRING, - description="Detailed prompt for the subagent.", - required=True, - ), - ToolParameter( - name="subagent_type", - type=ParameterType.STRING, - description="Delegatable agent name. Required for new tasks; omit when continuing with session_id.", - required=False, - ), - ToolParameter( - name="load_skills", - type=ParameterType.ARRAY, - description="Optional skill names to inject into the delegated agent", - required=False, - default=[], - ), - ToolParameter( - name="session_id", - type=ParameterType.STRING, - description="Existing subagent session to continue", - required=False, - ), - ToolParameter( - name="command", - type=ParameterType.STRING, - description="Deprecated command name retained for caller compatibility", - required=False, - ), - ToolParameter( - name="model", - type=ParameterType.STRING, - description="Optional model override (provider/model or model)", - required=False, - ), - ], -) -async def task_tool( - ctx: ToolContext, - description: Optional[str] = None, - prompt: Optional[str] = None, - subagent_type: Optional[str] = None, - load_skills: Optional[list] = None, - run_in_background: bool = False, - session_id: Optional[str] = None, - command: Optional[str] = None, - model: Optional[str] = None, -) -> ToolResult: - """Forward legacy task calls to delegate_task.""" - return await delegate_task_tool( - ctx=ctx, - prompt=prompt, - load_skills=load_skills, - description=description, - run_in_background=run_in_background, - subagent_type=subagent_type, - session_id=session_id, - command=command, - model=model, - ) diff --git a/flocks/tool/catalog.py b/flocks/tool/catalog.py index c37384cae..785bc2ada 100644 --- a/flocks/tool/catalog.py +++ b/flocks/tool/catalog.py @@ -37,7 +37,6 @@ class ToolCatalogMetadata(BaseModel): "webfetch": ["web", "http-fetch"], "websearch": ["web", "research"], "delegate_task": ["agent", "delegation"], - "task": ["agent", "delegation"], "schedule_task_create": ["scheduled-task", "task-management"], "schedule_task_list": ["scheduled-task", "task-management"], "schedule_task_status": ["scheduled-task", "task-management"], diff --git a/flocks/tool/registry.py b/flocks/tool/registry.py index 35d1c3cd7..17f05f904 100644 --- a/flocks/tool/registry.py +++ b/flocks/tool/registry.py @@ -1652,7 +1652,7 @@ def _register_builtin_tools(cls) -> None: # web/ — internet access ("flocks.tool.web", ["webfetch", "websearch"]), # agent/ — agent delegation/coordination - ("flocks.tool.agent", ["delegate_task", "task"]), + ("flocks.tool.agent", ["delegate_task"]), # task/ — task/workflow ("flocks.tool.task", [ "schedule_task_center", diff --git a/flocks/tool/task/schedule_task_center.py b/flocks/tool/task/schedule_task_center.py index aa1c6775e..1d27a2f64 100644 --- a/flocks/tool/task/schedule_task_center.py +++ b/flocks/tool/task/schedule_task_center.py @@ -312,7 +312,7 @@ async def schedule_task_create( enabled: Optional[bool] = None, action: Optional[str] = None, ) -> ToolResult: - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.models import ( SchedulerMode, TaskPriority, @@ -362,7 +362,7 @@ async def schedule_task_create( user_prompt=user_prompt, ) - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title=title, description=description, mode=mode, @@ -371,7 +371,7 @@ async def schedule_task_create( trigger=trigger, ) if enabled is False: - scheduler = await TaskManager.disable_scheduler(scheduler.id) or scheduler + scheduler = await ScheduleTaskManager.disable_scheduler(scheduler.id) or scheduler display_tz = resolve_task_timezone_name(scheduler) output_lines = [ @@ -382,7 +382,7 @@ async def schedule_task_create( f"Priority: {scheduler.priority.value}", ] if scheduler.trigger.run_immediately: - executions, _ = await TaskManager.list_scheduler_executions( + executions, _ = await ScheduleTaskManager.list_scheduler_executions( scheduler.id, limit=1, ) @@ -486,7 +486,7 @@ async def schedule_task_list( type: Optional[str] = None, limit: int = 10, ) -> ToolResult: - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.models import SchedulerStatus, TaskStatus if type is not None and type not in _VALID_TYPES: @@ -534,7 +534,7 @@ async def schedule_task_list( scheduler_status = SchedulerStatus.ACTIVE elif status in ("disabled", "paused"): scheduler_status = SchedulerStatus.DISABLED - tasks, total = await TaskManager.list_schedulers( + tasks, total = await ScheduleTaskManager.list_schedulers( status=scheduler_status, scheduled_only=True, limit=limit, @@ -552,7 +552,7 @@ async def schedule_task_list( f"Valid values: {', '.join(s.value for s in TaskStatus)}." ), ) - tasks, total = await TaskManager.list_executions( + tasks, total = await ScheduleTaskManager.list_executions( status=task_status, limit=limit, ) @@ -583,13 +583,13 @@ async def schedule_task_list( ], ) async def schedule_task_status(ctx: ToolContext, task_id: str) -> ToolResult: - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - task = await TaskManager.get_execution(task_id) + task = await ScheduleTaskManager.get_execution(task_id) if task and task.delivery_status.value == "unread": - await TaskManager.mark_notified(task_id) + await ScheduleTaskManager.mark_notified(task_id) if task is None: - task = await TaskManager.get_scheduler(task_id) + task = await ScheduleTaskManager.get_scheduler(task_id) if task is None: return ToolResult(success=False, error=f"Task {task_id} not found") @@ -738,7 +738,7 @@ async def schedule_task_update( user_prompt: Optional[str] = None, enabled: Optional[bool] = None, ) -> ToolResult: - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.models import TaskPriority normalized_action = (action or "update").lower() @@ -748,13 +748,13 @@ async def schedule_task_update( normalized_action = "enable" if normalized_action == "cancel": - task = await TaskManager.cancel_execution(task_id) + task = await ScheduleTaskManager.cancel_execution(task_id) elif normalized_action == "retry": - task = await TaskManager.retry_execution(task_id) + task = await ScheduleTaskManager.retry_execution(task_id) elif normalized_action == "disable": - task = await TaskManager.disable_scheduler(task_id) + task = await ScheduleTaskManager.disable_scheduler(task_id) elif normalized_action == "enable": - task = await TaskManager.enable_scheduler(task_id) + task = await ScheduleTaskManager.enable_scheduler(task_id) elif normalized_action == "update": fields = {} if priority: @@ -764,7 +764,7 @@ async def schedule_task_update( if description is not None: fields["description"] = description try: - task = await TaskManager.update_scheduler_with_trigger( + task = await ScheduleTaskManager.update_scheduler_with_trigger( task_id, fields=fields, cron=cron, @@ -777,10 +777,10 @@ async def schedule_task_update( except ValueError as exc: return ToolResult(success=False, error=str(exc)) if enabled is False: - task = await TaskManager.disable_scheduler(task_id) or task + task = await ScheduleTaskManager.disable_scheduler(task_id) or task normalized_action = "disable" elif enabled is True: - task = await TaskManager.enable_scheduler(task_id) or task + task = await ScheduleTaskManager.enable_scheduler(task_id) or task normalized_action = "enable" else: return ToolResult(success=False, error=f"Unknown action: {action}") @@ -813,13 +813,13 @@ async def schedule_task_update( ], ) async def schedule_task_delete(ctx: ToolContext, task_id: str) -> ToolResult: - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - execution = await TaskManager.get_execution(task_id) + execution = await ScheduleTaskManager.get_execution(task_id) if execution is not None: - ok = await TaskManager.delete_execution(task_id) + ok = await ScheduleTaskManager.delete_execution(task_id) else: - ok = await TaskManager.delete_scheduler(task_id) + ok = await ScheduleTaskManager.delete_scheduler(task_id) if not ok: return ToolResult(success=False, error=f"Task {task_id} not found") return ToolResult(success=True, output=f"Task {task_id} deleted.") @@ -843,11 +843,11 @@ async def schedule_task_delete(ctx: ToolContext, task_id: str) -> ToolResult: ], ) async def schedule_task_rerun(ctx: ToolContext, task_id: str) -> ToolResult: - from flocks.task.manager import TaskManager + from flocks.task.schedule_task_manager import ScheduleTaskManager - task = await TaskManager.rerun_execution(task_id) + task = await ScheduleTaskManager.rerun_execution(task_id) if task is None: - task = await TaskManager.rerun_scheduler(task_id) + task = await ScheduleTaskManager.rerun_scheduler(task_id) if not task: return ToolResult(success=False, error=f"Task {task_id} not found") diff --git a/flocks/utils/id.py b/flocks/utils/id.py index 7257ac0db..a7d682582 100644 --- a/flocks/utils/id.py +++ b/flocks/utils/id.py @@ -25,7 +25,6 @@ "call", # cal "step", # stp "agent", # agt - "subtask", # stk "event", # evt "tqref", # tqr "chbind", # chb (channel session binding) @@ -54,7 +53,6 @@ class Identifier: "call": "cal", "step": "stp", "agent": "agt", - "subtask": "stk", "event": "evt", "tqref": "tqr", "task": "tsk", diff --git a/flocks/workflow/tool_context.py b/flocks/workflow/tool_context.py index b803f9810..cb1f975ea 100644 --- a/flocks/workflow/tool_context.py +++ b/flocks/workflow/tool_context.py @@ -30,7 +30,7 @@ async def build_workflow_tool_context( Prefer the caller-provided session/message. When absent, create a temporary parent session and synthetic user message so workflow-internal tools such as - ``task`` / ``delegate_task`` can resolve a valid parent session. + ``delegate_task`` can resolve a valid parent session. """ effective_session_id = str(session_id or "").strip() diff --git a/tests/integration/test_task_queue_integration.py b/tests/integration/test_task_queue_integration.py index bde3ab7b0..32c20d961 100644 --- a/tests/integration/test_task_queue_integration.py +++ b/tests/integration/test_task_queue_integration.py @@ -8,12 +8,12 @@ import pytest -import flocks.task.manager as task_manager_module +import flocks.task.schedule_task_manager as schedule_task_manager_module from flocks.task.models import ExecutionMode from flocks.config.config import Config from flocks.storage.storage import Storage from flocks.task.executor import TaskExecutor -from flocks.task.manager import TaskManager +from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.models import ( ExecutionTriggerType, SchedulerMode, @@ -34,8 +34,8 @@ async def isolated_task_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): Config._cached_config = None Storage._db_path = None Storage._initialized = False - TaskManager._instance = None - TaskManager._startup_error = None + ScheduleTaskManager._instance = None + ScheduleTaskManager._startup_error = None TaskStore._initialized = False TaskStore._conn = None @@ -44,14 +44,14 @@ async def isolated_task_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): yield - await TaskManager.stop() + await ScheduleTaskManager.stop() await TaskStore.close() Config._global_config = None Config._cached_config = None Storage._db_path = None Storage._initialized = False - TaskManager._instance = None - TaskManager._startup_error = None + ScheduleTaskManager._instance = None + ScheduleTaskManager._startup_error = None TaskStore._initialized = False TaskStore._conn = None @@ -59,7 +59,7 @@ async def isolated_task_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): async def _wait_for_execution(execution_id: str, *, status: TaskStatus, timeout: float = 2.0): deadline = asyncio.get_running_loop().time() + timeout while asyncio.get_running_loop().time() < deadline: - execution = await TaskManager.get_execution(execution_id) + execution = await ScheduleTaskManager.get_execution(execution_id) if execution is not None and execution.status == status: return execution await asyncio.sleep(0.02) @@ -77,15 +77,15 @@ async def fake_dispatch(execution, scheduler): return await TaskStore.update_execution(execution) monkeypatch.setattr(TaskExecutor, "dispatch", fake_dispatch) - await TaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) + await ScheduleTaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="立即执行链路", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace"), ) - executions, total = await TaskManager.list_scheduler_executions(scheduler.id) + executions, total = await ScheduleTaskManager.list_scheduler_executions(scheduler.id) assert total == 1 completed = await _wait_for_execution(executions[0].id, status=TaskStatus.COMPLETED) @@ -108,18 +108,18 @@ async def fake_dispatch(execution, scheduler): return await TaskStore.update_execution(execution) monkeypatch.setattr(TaskExecutor, "dispatch", fake_dispatch) - await TaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) + await ScheduleTaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="重复执行", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace"), ) - first = (await TaskManager.list_scheduler_executions(scheduler.id))[0][0] + first = (await ScheduleTaskManager.list_scheduler_executions(scheduler.id))[0][0] first = await _wait_for_execution(first.id, status=TaskStatus.COMPLETED) - rerun = await TaskManager.rerun_execution(first.id) + rerun = await ScheduleTaskManager.rerun_execution(first.id) assert rerun is not None rerun = await _wait_for_execution(rerun.id, status=TaskStatus.COMPLETED) @@ -151,24 +151,24 @@ async def fake_dispatch(execution, scheduler): return await TaskStore.update_execution(execution) monkeypatch.setattr(TaskExecutor, "dispatch", fake_dispatch) - monkeypatch.setattr(task_manager_module, "_DISPATCH_GUARD_TIMEOUT_S", 0.05) - await TaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) + monkeypatch.setattr(schedule_task_manager_module, "_DISPATCH_GUARD_TIMEOUT_S", 0.05) + await ScheduleTaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) - first_scheduler = await TaskManager.create_scheduler( + first_scheduler = await ScheduleTaskManager.create_scheduler( title="超时任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace-1"), ) - second_scheduler = await TaskManager.create_scheduler( + second_scheduler = await ScheduleTaskManager.create_scheduler( title="后续任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace-2"), ) - first = (await TaskManager.list_scheduler_executions(first_scheduler.id))[0][0] - second = (await TaskManager.list_scheduler_executions(second_scheduler.id))[0][0] + first = (await ScheduleTaskManager.list_scheduler_executions(first_scheduler.id))[0][0] + second = (await ScheduleTaskManager.list_scheduler_executions(second_scheduler.id))[0][0] failed = await _wait_for_execution(first.id, status=TaskStatus.FAILED, timeout=1.0) completed = await _wait_for_execution(second.id, status=TaskStatus.COMPLETED, timeout=1.0) @@ -179,21 +179,21 @@ async def fake_dispatch(execution, scheduler): @pytest.mark.asyncio async def test_delete_scheduler_releases_claimed_queue_slot(tmp_path: Path): - await TaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) + await ScheduleTaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="删除释放槽位", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace-1"), ) - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) - manager = TaskManager.get() + manager = ScheduleTaskManager.get() assert manager is not None claimed = await manager.queue.dequeue() @@ -202,20 +202,20 @@ async def test_delete_scheduler_releases_claimed_queue_slot(tmp_path: Path): assert claimed.id == execution.id assert execution.id in manager.queue._running_ids - deleted = await TaskManager.delete_scheduler(scheduler.id) + deleted = await ScheduleTaskManager.delete_scheduler(scheduler.id) assert deleted is True assert execution.id not in manager.queue._running_ids - assert await TaskManager.get_execution(execution.id) is None + assert await ScheduleTaskManager.get_execution(execution.id) is None assert await TaskStore.get_queue_ref(execution.id) is None - next_scheduler = await TaskManager.create_scheduler( + next_scheduler = await ScheduleTaskManager.create_scheduler( title="后续可领取", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace-2"), ) - next_execution = await TaskManager.create_execution_from_scheduler( + next_execution = await ScheduleTaskManager.create_execution_from_scheduler( next_scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, @@ -236,21 +236,21 @@ async def fake_dispatch(_execution, _scheduler): raise asyncio.CancelledError() monkeypatch.setattr(TaskExecutor, "dispatch", fake_dispatch) - await TaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) + await ScheduleTaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="取消后回队", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) - manager = TaskManager.get() + manager = ScheduleTaskManager.get() assert manager is not None claimed = await manager.queue.dequeue() @@ -260,7 +260,7 @@ async def fake_dispatch(_execution, _scheduler): with pytest.raises(asyncio.CancelledError): await manager._run_execution(claimed) - refreshed = await TaskManager.get_execution(execution.id) + refreshed = await ScheduleTaskManager.get_execution(execution.id) queue_ref = await TaskStore.get_queue_ref(execution.id) assert refreshed is not None @@ -309,10 +309,10 @@ async def test_workflow_timeout_signals_cancel_and_stops_before_next_node( "read_workflow_from_fs", lambda workflow_id: {"workflowJson": workflow_json} if workflow_id == "wf_slow_cancel" else None, ) - monkeypatch.setattr(task_manager_module, "_DISPATCH_GUARD_TIMEOUT_S", 0.01) - await TaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) + monkeypatch.setattr(schedule_task_manager_module, "_DISPATCH_GUARD_TIMEOUT_S", 0.01) + await ScheduleTaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="可取消 workflow", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), @@ -321,7 +321,7 @@ async def test_workflow_timeout_signals_cancel_and_stops_before_next_node( workspace_directory=str(tmp_path / "workspace"), ) - execution = (await TaskManager.list_scheduler_executions(scheduler.id))[0][0] + execution = (await ScheduleTaskManager.list_scheduler_executions(scheduler.id))[0][0] failed = await _wait_for_execution(execution.id, status=TaskStatus.FAILED, timeout=1.0) await asyncio.sleep(0.15) @@ -466,8 +466,8 @@ async def test_standalone_legacy_migration_script_migrates_existing_tables(tmp_p stdout, stderr = await proc.communicate() assert proc.returncode == 0, (stdout or b"").decode() + (stderr or b"").decode() - scheduler = await TaskManager.get_scheduler("task_legacy_1") - execution = await TaskManager.get_execution("texec_legacy_1") + scheduler = await ScheduleTaskManager.get_scheduler("task_legacy_1") + execution = await ScheduleTaskManager.get_execution("texec_legacy_1") assert scheduler is not None assert scheduler.mode == SchedulerMode.ONCE @@ -475,7 +475,7 @@ async def test_standalone_legacy_migration_script_migrates_existing_tables(tmp_p assert execution.scheduler_id == "task_legacy_1" assert execution.status == TaskStatus.COMPLETED assert execution.session_id == "ses_123" - assert TaskManager._legacy_tables_exist() is False + assert ScheduleTaskManager._legacy_tables_exist() is False assert state_path.exists() is False @@ -581,8 +581,8 @@ async def test_standalone_legacy_migration_preserves_paused_scheduled_task_histo stdout, stderr = await proc.communicate() assert proc.returncode == 0, (stdout or b"").decode() + (stderr or b"").decode() - scheduler = await TaskManager.get_scheduler("task_paused_sched") - execution = await TaskManager.get_execution("legacy_exec_task_paused_sched") + scheduler = await ScheduleTaskManager.get_scheduler("task_paused_sched") + execution = await ScheduleTaskManager.get_execution("legacy_exec_task_paused_sched") assert scheduler is not None assert scheduler.status in (SchedulerStatus.ACTIVE, SchedulerStatus.DISABLED) diff --git a/tests/server/routes/test_task_scheduler_context_route.py b/tests/server/routes/test_task_scheduler_context_route.py index 8df880d77..e31c85214 100644 --- a/tests/server/routes/test_task_scheduler_context_route.py +++ b/tests/server/routes/test_task_scheduler_context_route.py @@ -2,13 +2,13 @@ import pytest -from flocks.task.manager import TaskManager +from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.models import ExecutionMode, ExecutionTriggerType, SchedulerMode, TaskTrigger @pytest.mark.asyncio async def test_update_scheduler_accepts_context_for_workflow_inputs(client): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="工作流定时任务", mode=SchedulerMode.CRON, trigger=TaskTrigger(cron="0 9 * * *", timezone="Asia/Shanghai"), @@ -28,11 +28,11 @@ async def test_update_scheduler_accepts_context_for_workflow_inputs(client): assert response.status_code == 200 assert response.json()["context"] == {"keyword": "after", "limit": 5} - updated = await TaskManager.get_scheduler(scheduler.id) + updated = await ScheduleTaskManager.get_scheduler(scheduler.id) assert updated is not None assert updated.context == {"keyword": "after", "limit": 5} - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( updated, trigger_type=ExecutionTriggerType.SCHEDULED, enqueue=False, diff --git a/tests/server/test_lifespan.py b/tests/server/test_lifespan.py index ab2b4fa60..9e34b9fc4 100644 --- a/tests/server/test_lifespan.py +++ b/tests/server/test_lifespan.py @@ -93,8 +93,8 @@ async def fake_async_noop(*_args, **_kwargs) -> None: ) monkeypatch.setitem( sys.modules, - "flocks.task.manager", - types.SimpleNamespace(TaskManager=types.SimpleNamespace(start=fake_async_noop, stop=fake_async_noop)), + "flocks.task.schedule_task_manager", + types.SimpleNamespace(ScheduleTaskManager=types.SimpleNamespace(start=fake_async_noop, stop=fake_async_noop)), ) monkeypatch.setitem( sys.modules, diff --git a/tests/server/test_server.py b/tests/server/test_server.py index d326a7b69..355849931 100644 --- a/tests/server/test_server.py +++ b/tests/server/test_server.py @@ -9,6 +9,7 @@ from fastapi import status from flocks.server.app import app +from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.store import TaskStore from flocks.task.models import ( DeliveryStatus, @@ -132,12 +133,12 @@ async def test_list_executions_invalid_priority_returns_422(client): @pytest.mark.asyncio async def test_task_schedulers_scheduled_only_excludes_immediate_queue_templates(client): - await TaskManager.create_scheduler( + await ScheduleTaskManager.create_scheduler( title="立即任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), ) - scheduled = await TaskManager.create_scheduler( + scheduled = await ScheduleTaskManager.create_scheduler( title="单次计划", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), @@ -154,12 +155,12 @@ async def test_task_schedulers_scheduled_only_excludes_immediate_queue_templates @pytest.mark.asyncio async def test_task_scheduler_list_accepts_legacy_paused_status_query(client): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="兼容旧 paused 调度查询", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), ) - await TaskManager.disable_scheduler(scheduler.id) + await ScheduleTaskManager.disable_scheduler(scheduler.id) response = await client.get("/api/task-schedulers", params={"status": "paused"}) @@ -170,13 +171,13 @@ async def test_task_scheduler_list_accepts_legacy_paused_status_query(client): @pytest.mark.asyncio async def test_task_schedulers_list_excludes_archived_builtin_after_delete(client): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="内置计划任务", mode=SchedulerMode.CRON, trigger=TaskTrigger(cron="*/5 * * * *", timezone="Asia/Shanghai"), dedup_key="builtin:test-scheduled-task", ) - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.SCHEDULED, enqueue=True, @@ -191,8 +192,8 @@ async def test_task_schedulers_list_excludes_archived_builtin_after_delete(clien ids = {item["id"] for item in data["items"]} assert scheduler.id not in ids - archived = await TaskManager.get_scheduler(scheduler.id) - cancelled_execution = await TaskManager.get_execution(execution.id) + archived = await ScheduleTaskManager.get_scheduler(scheduler.id) + cancelled_execution = await ScheduleTaskManager.get_execution(execution.id) assert archived is not None assert archived.status == SchedulerStatus.ARCHIVED assert cancelled_execution is not None @@ -201,13 +202,13 @@ async def test_task_schedulers_list_excludes_archived_builtin_after_delete(clien @pytest.mark.asyncio async def test_delete_scheduler_cleans_queue_state_for_non_builtin(client): - await TaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await TaskManager.create_scheduler( + await ScheduleTaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) + scheduler = await ScheduleTaskManager.create_scheduler( title="普通计划任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), ) - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, @@ -233,12 +234,12 @@ async def test_delete_scheduler_cleans_queue_state_for_non_builtin(client): @pytest.mark.asyncio async def test_task_execution_list_accepts_legacy_paused_status_query(client): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="兼容旧 paused 执行查询", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), ) - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -256,17 +257,17 @@ async def test_task_execution_list_accepts_legacy_paused_status_query(client): @pytest.mark.asyncio async def test_batch_cancel_endpoint_cancels_selected_executions(client): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="批量取消接口", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), ) - cancellable = await TaskManager.create_execution_from_scheduler( + cancellable = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) - completed = await TaskManager.create_execution_from_scheduler( + completed = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -283,19 +284,19 @@ async def test_batch_cancel_endpoint_cancels_selected_executions(client): assert response.status_code == status.HTTP_200_OK assert response.json()["cancelled"] == 1 - cancelled_execution = await TaskManager.get_execution(cancellable.id) + cancelled_execution = await ScheduleTaskManager.get_execution(cancellable.id) assert cancelled_execution is not None assert cancelled_execution.status == TaskStatus.CANCELLED @pytest.mark.asyncio async def test_execution_pause_and_resume_endpoints_are_removed(client): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="旧暂停接口", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), ) - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, @@ -310,12 +311,12 @@ async def test_execution_pause_and_resume_endpoints_are_removed(client): @pytest.mark.asyncio async def test_mark_execution_viewed_endpoint_updates_delivery_status(client): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="标记已读", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), ) - execution = (await TaskManager.list_scheduler_executions(scheduler.id, limit=1))[0][0] + execution = (await ScheduleTaskManager.list_scheduler_executions(scheduler.id, limit=1))[0][0] execution.status = TaskStatus.COMPLETED execution.delivery_status = DeliveryStatus.UNREAD await TaskStore.update_execution(execution) diff --git a/tests/session/test_auto_model_failover.py b/tests/session/test_auto_model_failover.py index d1e8d1a60..388b10c18 100644 --- a/tests/session/test_auto_model_failover.py +++ b/tests/session/test_auto_model_failover.py @@ -1301,12 +1301,12 @@ def test_cooldown_is_cleared_when_primary_changes(): @pytest.mark.asyncio -async def test_synthetic_subtask_continuation_keeps_fallback(monkeypatch): +async def test_synthetic_continuation_keeps_fallback(monkeypatch): ctx = _ctx(index=1) ctx.model_candidate_policy = "configured" ctx.turn_user_id = "msg_real" synthetic_user = SimpleNamespace( - id="msg_subtask_continue", + id="msg_synthetic_continue", model={"providerID": "primary", "modelID": "primary-model"}, ) monkeypatch.setattr( diff --git a/tests/session/test_context_usage.py b/tests/session/test_context_usage.py index 2bc9db71a..303dcf88a 100644 --- a/tests/session/test_context_usage.py +++ b/tests/session/test_context_usage.py @@ -275,7 +275,7 @@ async def test_context_usage_splits_skill_and_delegation_tools(context_usage_moc ), SimpleNamespace( type="tool", - tool="task", + tool="delegate_task", state=SimpleNamespace(input={}, output="t" * 80, time={"start": 3}), ), SimpleNamespace( @@ -288,20 +288,16 @@ async def test_context_usage_splits_skill_and_delegation_tools(context_usage_moc metadata={"tool": "skill_load"}, state=SimpleNamespace(input={}, output="m" * 40, time={"start": 5}), ), - SimpleNamespace( - type="subtask", - prompt="p" * 40, - description="q" * 40, - ), ] } snapshot = await context_usage.build_context_usage_snapshot("sess-1") assert [(segment.key, segment.tokens) for segment in snapshot.segments] == [ + ("conversation", 20), ("tools", 30), ("skillLoad", 30), - ("agentDelegation", 50), + ("agentDelegation", 30), ] tools_segment = next(segment for segment in snapshot.segments if segment.key == "tools") assert tools_segment.tokens == 30 diff --git a/tests/session/test_execution_mode.py b/tests/session/test_execution_mode.py index fceb95dc1..637d3cfcd 100644 --- a/tests/session/test_execution_mode.py +++ b/tests/session/test_execution_mode.py @@ -89,7 +89,6 @@ def test_plan_uses_read_only_permission_rules() -> None: assert is_tool_allowed(SessionExecutionMode.PLAN, "edit") assert is_tool_allowed(SessionExecutionMode.PLAN, "write") assert is_tool_allowed(SessionExecutionMode.PLAN, "unknown_plugin_tool") - assert is_tool_allowed(SessionExecutionMode.PLAN, "task") assert is_tool_allowed(SessionExecutionMode.PLAN, "delegate_task") assert not is_tool_allowed(SessionExecutionMode.PLAN, "run_slash_command") @@ -102,14 +101,13 @@ def test_plan_uses_read_only_permission_rules() -> None: assert execution_mode_prompt("build") == "" -@pytest.mark.parametrize("tool_name", ["task", "delegate_task"]) -def test_plan_delegation_only_allows_explore_and_librarian(tool_name) -> None: +def test_plan_delegation_only_allows_explore_and_librarian() -> None: ctx = ToolContext(session_id="session-1", message_id="message-1") for subagent_type in ("explore", "librarian"): assert tool_call_denial_reason( SessionExecutionMode.PLAN, - tool_name, + "delegate_task", {"subagent_type": subagent_type}, ctx, ) is None @@ -121,7 +119,7 @@ def test_plan_delegation_only_allows_explore_and_librarian(tool_name) -> None: ): reason = tool_call_denial_reason( SessionExecutionMode.PLAN, - tool_name, + "delegate_task", arguments, ctx, ) @@ -245,7 +243,7 @@ async def handler(_ctx, **_kwargs): tool = Tool( info=ToolInfo( - name="task", + name="delegate_task", description="Delegation test tool", category=ToolCategory.FILE, ), @@ -258,7 +256,7 @@ async def handler(_ctx, **_kwargs): ) explore = await ToolRegistry.execute( - "task", + "delegate_task", ctx=ToolContext( session_id="session-1", message_id="message-1", @@ -411,7 +409,6 @@ async def test_runner_filters_tools_with_message_mode(monkeypatch) -> None: "bash", "write", "edit", - "task", "delegate_task", "run_slash_command", ] @@ -423,7 +420,6 @@ async def test_runner_filters_tools_with_message_mode(monkeypatch) -> None: SimpleNamespace(name="bash"), SimpleNamespace(name="write"), SimpleNamespace(name="edit"), - SimpleNamespace(name="task"), SimpleNamespace(name="delegate_task"), SimpleNamespace(name="run_slash_command"), ], @@ -465,7 +461,6 @@ async def list_tools(**_kwargs): "bash", "write", "edit", - "task", "delegate_task", "plan_exit", ] @@ -476,6 +471,5 @@ async def list_tools(**_kwargs): "edit", "plan_exit", "read", - "task", "write", ] diff --git a/tests/session/test_message_parts.py b/tests/session/test_message_parts.py index 89504ba62..c28a7d1c0 100644 --- a/tests/session/test_message_parts.py +++ b/tests/session/test_message_parts.py @@ -27,7 +27,6 @@ SnapshotPart, StepFinishPart, StepStartPart, - SubtaskPart, TextPart, TokenCache, TokenUsage, @@ -313,21 +312,9 @@ def test_creation(self): # --------------------------------------------------------------------------- -# SubtaskPart / AgentPart +# AgentPart # --------------------------------------------------------------------------- -class TestSubtaskPart: - def test_creation(self): - part = SubtaskPart( - sessionID=SID, - messageID=MID, - prompt="Summarize findings", - description="Summarize", - agent="rex", - ) - assert part.type == "subtask" - assert part.agent == "rex" - class TestAgentPart: def test_creation(self): @@ -416,6 +403,24 @@ def test_deserialize_reasoning_part(self): assert deserialized is not None assert deserialized.type == "reasoning" + def test_deserialize_legacy_subtask_as_ignored_text(self): + deserialized = Message.deserialize_part( + { + "id": "part_legacy_subtask", + "sessionID": SID, + "messageID": MID, + "type": "subtask", + "prompt": "old delegated command", + "description": "legacy", + "agent": "rex", + } + ) + + assert deserialized.type == "text" + assert deserialized.text == "" + assert deserialized.ignored is True + assert deserialized.metadata == {"legacyPartType": "subtask"} + def test_deserialize_unknown_type_falls_back_to_text(self): # Unknown type falls back to TextPart; missing required fields raise exception with pytest.raises(Exception): @@ -537,11 +542,11 @@ async def test_store_part_does_not_downgrade_terminal_tool_state(self, monkeypat sessionID=sid, messageID=msg.id, callID="call_terminal_guard", - tool="task", + tool="delegate_task", state=ToolStateCompleted( input={"prompt": "run"}, output="done", - title="task", + title="delegate_task", metadata={"sessionId": "ses_child_done"}, time={"start": 1000, "end": 2000}, ), @@ -551,10 +556,10 @@ async def test_store_part_does_not_downgrade_terminal_tool_state(self, monkeypat sessionID=sid, messageID=msg.id, callID="call_terminal_guard", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "run"}, - title="task", + title="delegate_task", metadata={"sessionId": "ses_child_done", "status": "running"}, time={"start": 1000}, ), diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index 5bfbb666f..f23c04d48 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -1533,7 +1533,7 @@ async def test_to_chat_messages_invalidates_shared_cache_when_message_parts_chan sessionID=session.id, messageID=assistant_message.id, callID="call_cache_fix", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1546,7 +1546,7 @@ async def test_to_chat_messages_invalidates_shared_cache_when_message_parts_chan assert len(second_messages) == 2 assert second_messages[0].role == "assistant" assert second_messages[0].tool_calls is not None - assert second_messages[0].tool_calls[0]["function"]["name"] == "task" + assert second_messages[0].tool_calls[0]["function"]["name"] == "delegate_task" assert second_messages[1].role == "tool" assert second_messages[1].tool_call_id == "call_cache_fix" assert second_messages[1].content == "Error: Tool execution was interrupted" @@ -1603,7 +1603,7 @@ async def test_to_chat_messages_preserves_assistant_reasoning_for_replay(): sessionID=session.id, messageID=assistant_message.id, callID="call_reasoning_replay", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1617,7 +1617,7 @@ async def test_to_chat_messages_preserves_assistant_reasoning_for_replay(): assert chat_messages[0].role == "assistant" assert chat_messages[0].reasoning == "Need to call the tool first." assert chat_messages[0].tool_calls is not None - assert chat_messages[0].tool_calls[0]["function"]["name"] == "task" + assert chat_messages[0].tool_calls[0]["function"]["name"] == "delegate_task" assert chat_messages[1].role == "tool" assert chat_messages[1].tool_call_id == "call_reasoning_replay" @@ -1672,7 +1672,7 @@ async def test_to_chat_messages_restores_provider_reasoning_fields_from_metadata sessionID=session.id, messageID=assistant_message.id, callID="call_reasoning_metadata", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1686,7 +1686,7 @@ async def test_to_chat_messages_restores_provider_reasoning_fields_from_metadata assert chat_messages[0].reasoning == "Need to call the tool first." assert chat_messages[0].reasoning_content == "Need to call the tool first." assert chat_messages[0].reasoning_source == "native_reasoning_content" - assert chat_messages[0].tool_calls[0]["function"]["name"] == "task" + assert chat_messages[0].tool_calls[0]["function"]["name"] == "delegate_task" @pytest.mark.asyncio @@ -1734,7 +1734,7 @@ async def test_to_chat_messages_restores_redacted_anthropic_thinking_blocks(monk sessionID=session.id, messageID=assistant_message.id, callID="call_redacted_reasoning", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1748,7 +1748,7 @@ async def test_to_chat_messages_restores_redacted_anthropic_thinking_blocks(monk assert chat_messages[0].custom_settings["anthropic_thinking_blocks"] == [ {"type": "redacted_thinking", "data": "opaque_blob"} ] - assert chat_messages[0].tool_calls[0]["function"]["name"] == "task" + assert chat_messages[0].tool_calls[0]["function"]["name"] == "delegate_task" @pytest.mark.asyncio @@ -1796,7 +1796,7 @@ async def test_to_chat_messages_restores_signed_anthropic_thinking_blocks(monkey sessionID=session.id, messageID=assistant_message.id, callID="call_signed_reasoning", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1814,7 +1814,7 @@ async def test_to_chat_messages_restores_signed_anthropic_thinking_blocks(monkey "signature": "sig123", } ] - assert chat_messages[0].tool_calls[0]["function"]["name"] == "task" + assert chat_messages[0].tool_calls[0]["function"]["name"] == "delegate_task" @pytest.mark.asyncio @@ -1861,7 +1861,7 @@ async def test_to_chat_messages_restores_unsigned_anthropic_thinking_blocks(monk sessionID=session.id, messageID=assistant_message.id, callID="call_unsigned_reasoning", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1878,7 +1878,7 @@ async def test_to_chat_messages_restores_unsigned_anthropic_thinking_blocks(monk "thinking": "Unsigned plan before tool use.", } ] - assert chat_messages[0].tool_calls[0]["function"]["name"] == "task" + assert chat_messages[0].tool_calls[0]["function"]["name"] == "delegate_task" @pytest.mark.asyncio @@ -1926,7 +1926,7 @@ async def test_runner_history_round_trip_formats_anthropic_payload(monkeypatch): sessionID=session.id, messageID=assistant_message.id, callID="call_signed_reasoning", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -2010,7 +2010,7 @@ async def test_to_chat_messages_prefers_provider_specific_interleaved_resolution sessionID=session.id, messageID=assistant_message.id, callID="call_provider_specific_interleaved", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, diff --git a/tests/session/test_session_abort_inject.py b/tests/session/test_session_abort_inject.py index 36154e2d8..e55ef6df4 100644 --- a/tests/session/test_session_abort_inject.py +++ b/tests/session/test_session_abort_inject.py @@ -838,60 +838,6 @@ async def test_run_loop_breaks_on_exit_condition_without_tool_parts(self): assert event_names == ["turn.started"] -class TestExecuteSubtask: - @pytest.mark.asyncio - async def test_execute_subtask_passes_tool_context_first(self): - session_info = _make_session_info("subtask_exec_test") - ctx = LoopContext( - session=session_info, - provider_id="test-provider", - model_id="test-model", - agent_name="rex", - ) - last_user = SimpleNamespace( - id="msg_parent", - agent="rex", - model={"providerID": "test-provider", "modelID": "test-model"}, - provider="test-provider", - ) - task_part = SimpleNamespace( - agent="helper", - prompt="do the thing", - description="test task", - command=None, - model=None, - ) - - task_tool = MagicMock() - task_tool.execute = AsyncMock(return_value=SimpleNamespace( - output="done", - title="task complete", - metadata={"sessionId": "child-session"}, - )) - - assistant_msg = SimpleNamespace(id="msg_assistant") - synthetic_msg = SimpleNamespace(id="msg_synthetic") - - with patch("flocks.agent.registry.Agent.get", AsyncMock(return_value=SimpleNamespace(name="helper"))), \ - patch("flocks.tool.registry.ToolRegistry.get", return_value=task_tool), \ - patch("flocks.session.runtime.session_turn.Message.create", AsyncMock(side_effect=[assistant_msg, synthetic_msg])), \ - patch("flocks.session.runtime.session_turn.Message.add_part", AsyncMock()), \ - patch("flocks.session.runtime.session_turn.Message.update", AsyncMock()), \ - patch("flocks.session.runtime.session_turn.Message.update_part", AsyncMock()): - await ctx._execute_subtask(last_user, task_part) - - task_tool.execute.assert_awaited_once() - tool_ctx = task_tool.execute.await_args.args[0] - assert tool_ctx.session_id == session_info.id - assert tool_ctx.message_id == assistant_msg.id - assert task_tool.execute.await_args.kwargs == { - "prompt": "do the thing", - "description": "test task", - "subagent_type": "helper", - "command": None, - } - - # --------------------------------------------------------------------------- # LoopContext tests # --------------------------------------------------------------------------- diff --git a/tests/storage/test_sqlite_connection_config.py b/tests/storage/test_sqlite_connection_config.py index 19276d496..eb1aa1725 100644 --- a/tests/storage/test_sqlite_connection_config.py +++ b/tests/storage/test_sqlite_connection_config.py @@ -5,7 +5,7 @@ from flocks.config.config import Config from flocks.storage.storage import Storage -from flocks.task.manager import TaskManager +from flocks.task.schedule_task_manager import ScheduleTaskManager @pytest.fixture(autouse=True) @@ -73,10 +73,10 @@ def test_storage_connect_sync_applies_runtime_sqlite_pragmas() -> None: @pytest.mark.asyncio -async def test_task_manager_sync_connection_uses_storage_sqlite_contract() -> None: +async def test_schedule_task_manager_sync_connection_uses_storage_sqlite_contract() -> None: await Storage.init() - with TaskManager._with_db_connection() as db: + with ScheduleTaskManager._with_db_connection() as db: row = db.execute( "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'storage'" ).fetchone() diff --git a/tests/task/test_task.py b/tests/task/test_task.py index f12d1a49e..b1cb53520 100644 --- a/tests/task/test_task.py +++ b/tests/task/test_task.py @@ -10,7 +10,7 @@ from flocks.cli.commands import task as task_cli_commands import flocks.task.background as background_module -import flocks.task.manager as task_manager_module +import flocks.task.schedule_task_manager as schedule_task_manager_module import flocks.task.plugin_sync as plugin_sync_module from flocks.server.routes import question as question_routes from flocks.config.config import Config @@ -19,7 +19,7 @@ from flocks.task.background import BackgroundManager, BackgroundTask, LaunchInput from flocks.task.executor import TaskExecutor from flocks.task.formatting import format_task_datetime -from flocks.task.manager import TaskManager +from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.models import ( DeliveryStatus, ExecutionMode, @@ -49,8 +49,8 @@ async def isolated_task_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): Config._cached_config = None Storage._db_path = None Storage._initialized = False - TaskManager._instance = None - TaskManager._startup_error = None + ScheduleTaskManager._instance = None + ScheduleTaskManager._startup_error = None TaskStore._initialized = False TaskStore._conn = None @@ -59,21 +59,21 @@ async def isolated_task_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): yield - await TaskManager.stop() + await ScheduleTaskManager.stop() await TaskStore.close() Config._global_config = None Config._cached_config = None Storage._db_path = None Storage._initialized = False - TaskManager._instance = None - TaskManager._startup_error = None + ScheduleTaskManager._instance = None + ScheduleTaskManager._startup_error = None TaskStore._initialized = False TaskStore._conn = None @pytest.mark.asyncio async def test_immediate_scheduler_creates_single_queued_execution(tmp_path: Path): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="立即执行", description="创建后立刻入队", mode=SchedulerMode.ONCE, @@ -82,7 +82,7 @@ async def test_immediate_scheduler_creates_single_queued_execution(tmp_path: Pat workspace_directory=str(tmp_path / "workspace"), ) - executions, total = await TaskManager.list_scheduler_executions(scheduler.id) + executions, total = await ScheduleTaskManager.list_scheduler_executions(scheduler.id) assert total == 1 execution = executions[0] @@ -99,7 +99,7 @@ async def test_list_executions_tolerates_legacy_non_utf8_description( tmp_path: Path, ): legacy_description = "扫描 Windows 主机 192.168.254.1" - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="legacy encoding", description="placeholder", mode=SchedulerMode.ONCE, @@ -107,7 +107,7 @@ async def test_list_executions_tolerates_legacy_non_utf8_description( trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace"), ) - executions, _ = await TaskManager.list_scheduler_executions(scheduler.id) + executions, _ = await ScheduleTaskManager.list_scheduler_executions(scheduler.id) execution_id = executions[0].id db = await TaskStore.raw_db() @@ -121,7 +121,7 @@ async def test_list_executions_tolerates_legacy_non_utf8_description( ) await db.commit() - executions, total = await TaskManager.list_executions(limit=20) + executions, total = await ScheduleTaskManager.list_executions(limit=20) assert total == 1 assert executions[0].description == legacy_description @@ -129,7 +129,7 @@ async def test_list_executions_tolerates_legacy_non_utf8_description( @pytest.mark.asyncio async def test_once_scheduler_tick_creates_execution_and_disables_scheduler(tmp_path: Path): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="单次定时", mode=SchedulerMode.ONCE, trigger=TaskTrigger( @@ -142,8 +142,8 @@ async def test_once_scheduler_tick_creates_execution_and_disables_scheduler(tmp_ loop = SchedulerLoop() await loop._tick() - updated = await TaskManager.get_scheduler(scheduler.id) - executions, total = await TaskManager.list_scheduler_executions(scheduler.id) + updated = await ScheduleTaskManager.get_scheduler(scheduler.id) + executions, total = await ScheduleTaskManager.list_scheduler_executions(scheduler.id) assert updated is not None assert updated.status == SchedulerStatus.DISABLED @@ -154,7 +154,7 @@ async def test_once_scheduler_tick_creates_execution_and_disables_scheduler(tmp_ @pytest.mark.asyncio async def test_cron_scheduler_does_not_spawn_second_active_execution(tmp_path: Path): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="循环任务", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -172,7 +172,7 @@ async def test_cron_scheduler_does_not_spawn_second_active_execution(tmp_path: P await loop._tick() await loop._tick() - executions, total = await TaskManager.list_scheduler_executions(scheduler.id, limit=10) + executions, total = await ScheduleTaskManager.list_scheduler_executions(scheduler.id, limit=10) assert total == 1 assert executions[0].status == TaskStatus.QUEUED @@ -182,7 +182,7 @@ async def test_cron_scheduler_does_not_spawn_second_active_execution(tmp_path: P @pytest.mark.asyncio async def test_create_scheduler_rejects_six_field_cron(tmp_path: Path): with pytest.raises(ValueError, match="only 5-field cron is supported"): - await TaskManager.create_scheduler( + await ScheduleTaskManager.create_scheduler( title="非法 cron", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -196,7 +196,7 @@ async def test_create_scheduler_rejects_six_field_cron(tmp_path: Path): @pytest.mark.asyncio async def test_create_scheduler_rejects_out_of_range_five_field_cron(tmp_path: Path): with pytest.raises(ValueError, match="not a valid 5-field cron"): - await TaskManager.create_scheduler( + await ScheduleTaskManager.create_scheduler( title="越界 cron", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -214,7 +214,7 @@ async def test_create_scheduler_does_not_mutate_trigger_argument(tmp_path: Path) timezone="Asia/Shanghai", ) - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="不修改调用方 trigger", mode=SchedulerMode.CRON, trigger=trigger, @@ -276,7 +276,7 @@ async def test_plugin_sync_skips_new_scheduler_with_invalid_six_field_cron( async def test_plugin_sync_does_not_overwrite_existing_scheduler_with_invalid_six_field_cron( tmp_path: Path, ): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="原始内置任务", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -299,7 +299,7 @@ async def test_plugin_sync_does_not_overwrite_existing_scheduler_with_invalid_si ] ) - unchanged = await TaskManager.get_scheduler(scheduler.id) + unchanged = await ScheduleTaskManager.get_scheduler(scheduler.id) assert created == 0 assert unchanged is not None @@ -312,7 +312,7 @@ async def test_plugin_sync_does_not_overwrite_existing_scheduler_with_invalid_si async def test_update_scheduler_with_trigger_run_once_preserves_existing_cron_when_omitted( tmp_path: Path, ): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="切成单次任务时保留旧 cron", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -322,7 +322,7 @@ async def test_update_scheduler_with_trigger_run_once_preserves_existing_cron_wh workspace_directory=str(tmp_path / "workspace"), ) - updated = await TaskManager.update_scheduler_with_trigger( + updated = await ScheduleTaskManager.update_scheduler_with_trigger( scheduler.id, fields={}, run_once=True, @@ -338,8 +338,8 @@ async def test_update_scheduler_with_trigger_run_once_preserves_existing_cron_wh @pytest.mark.asyncio async def test_recover_stale_running_execution_unblocks_scheduler(tmp_path: Path): - manager = TaskManager(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await TaskManager.create_scheduler( + manager = ScheduleTaskManager(max_concurrent=1, poll_interval=999, scheduler_interval=999) + scheduler = await ScheduleTaskManager.create_scheduler( title="恢复阻塞循环任务", mode=SchedulerMode.CRON, trigger=TaskTrigger(cron="*/5 * * * *", timezone="Asia/Shanghai"), @@ -348,13 +348,13 @@ async def test_recover_stale_running_execution_unblocks_scheduler(tmp_path: Path scheduler.trigger.next_run = datetime.now(timezone.utc) - timedelta(seconds=1) await TaskStore.update_scheduler(scheduler) - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.SCHEDULED, enqueue=False, ) started_at = datetime.now(timezone.utc) - timedelta( - seconds=task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 5 + seconds=schedule_task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 5 ) execution.status = TaskStatus.RUNNING execution.started_at = started_at @@ -366,7 +366,7 @@ async def test_recover_stale_running_execution_unblocks_scheduler(tmp_path: Path recovered = await manager._recover_stale_active_executions() assert recovered == 1 - failed = await TaskManager.get_execution(execution.id) + failed = await ScheduleTaskManager.get_execution(execution.id) assert failed is not None assert failed.status == TaskStatus.FAILED assert "recovery threshold" in (failed.error or "") @@ -374,7 +374,7 @@ async def test_recover_stale_running_execution_unblocks_scheduler(tmp_path: Path loop = SchedulerLoop() await loop._tick() - executions, total = await TaskManager.list_scheduler_executions(scheduler.id, limit=10) + executions, total = await ScheduleTaskManager.list_scheduler_executions(scheduler.id, limit=10) assert total == 2 assert any(item.id == execution.id and item.status == TaskStatus.FAILED for item in executions) @@ -383,29 +383,29 @@ async def test_recover_stale_running_execution_unblocks_scheduler(tmp_path: Path @pytest.mark.asyncio async def test_recover_orphaned_queued_execution_restores_queue_ref(tmp_path: Path): - await TaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) + await ScheduleTaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) # Pause the execution loop so it cannot race the test by claiming the # execution before we manually simulate the orphan state (queued row # with no queue ref). - TaskManager.pause_queue() - scheduler = await TaskManager.create_scheduler( + ScheduleTaskManager.pause_queue() + scheduler = await ScheduleTaskManager.create_scheduler( title="恢复孤儿排队任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) await TaskStore.finish_queue_ref(execution.id) - manager = TaskManager.get() + manager = ScheduleTaskManager.get() assert manager is not None recovered = await manager._recover_orphaned_queued_executions() - refreshed = await TaskManager.get_execution(execution.id) + refreshed = await ScheduleTaskManager.get_execution(execution.id) assert recovered == 1 assert refreshed is not None @@ -420,54 +420,54 @@ async def test_recover_orphaned_queued_execution_restores_queue_ref(tmp_path: Pa @pytest.mark.asyncio async def test_queue_status_reports_stale_running_execution(tmp_path: Path): - await TaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await TaskManager.create_scheduler( + await ScheduleTaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) + scheduler = await ScheduleTaskManager.create_scheduler( title="阻塞诊断", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, ) started_at = datetime.now(timezone.utc) - timedelta( - seconds=task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 15 + seconds=schedule_task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 15 ) execution.status = TaskStatus.RUNNING execution.started_at = started_at execution.queued_at = started_at await TaskStore.update_execution(execution) - status = await TaskManager.queue_status() + status = await ScheduleTaskManager.queue_status() assert status["stale_running"] == 1 assert isinstance(status["oldest_running_seconds"], int) - assert status["oldest_running_seconds"] >= task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + assert status["oldest_running_seconds"] >= schedule_task_manager_module._RUNNING_RECOVERY_TIMEOUT_S @pytest.mark.asyncio async def test_queue_status_uses_largest_elapsed_running_time(tmp_path: Path): - await TaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await TaskManager.create_scheduler( + await ScheduleTaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) + scheduler = await ScheduleTaskManager.create_scheduler( title="阻塞诊断-多任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - older_execution = await TaskManager.create_execution_from_scheduler( + older_execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, ) - newer_execution = await TaskManager.create_execution_from_scheduler( + newer_execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, ) older_started_at = datetime.now(timezone.utc) - timedelta( - seconds=task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 15 + seconds=schedule_task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 15 ) newer_started_at = datetime.now(timezone.utc) - timedelta(seconds=10) older_execution.status = TaskStatus.RUNNING @@ -479,11 +479,11 @@ async def test_queue_status_uses_largest_elapsed_running_time(tmp_path: Path): await TaskStore.update_execution(older_execution) await TaskStore.update_execution(newer_execution) - status = await TaskManager.queue_status() + status = await ScheduleTaskManager.queue_status() assert status["stale_running"] == 1 assert isinstance(status["oldest_running_seconds"], int) - assert status["oldest_running_seconds"] >= task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 15 + assert status["oldest_running_seconds"] >= schedule_task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 15 @pytest.mark.asyncio @@ -747,15 +747,15 @@ def fake_run_workflow(*, workflow, inputs, **_kwargs): @pytest.mark.asyncio async def test_retry_queue_requeues_failed_execution(tmp_path: Path): - manager = TaskManager(max_concurrent=1, poll_interval=999, scheduler_interval=999) + manager = ScheduleTaskManager(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="失败重试", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -769,7 +769,7 @@ async def test_retry_queue_requeues_failed_execution(tmp_path: Path): await manager._process_retry_queue() - reloaded = await TaskManager.get_execution(execution.id) + reloaded = await ScheduleTaskManager.get_execution(execution.id) assert reloaded is not None assert reloaded.status == TaskStatus.QUEUED assert reloaded.retry.retry_after is None @@ -781,18 +781,18 @@ async def test_retry_queue_requeues_failed_execution(tmp_path: Path): @pytest.mark.asyncio async def test_queue_dequeue_respects_claimed_slots_before_running_status(tmp_path: Path): queue = TaskQueue(max_concurrent=1) - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="并发控制", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - first = await TaskManager.create_execution_from_scheduler( + first = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) - second = await TaskManager.create_execution_from_scheduler( + second = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, @@ -814,14 +814,14 @@ async def test_queue_dequeue_respects_claimed_slots_before_running_status(tmp_pa @pytest.mark.asyncio async def test_immediate_scheduler_dedup_does_not_create_duplicate_execution(tmp_path: Path): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="去重立即执行", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace"), dedup_key="dup-immediate", ) - duplicate = await TaskManager.create_scheduler( + duplicate = await ScheduleTaskManager.create_scheduler( title="去重立即执行", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), @@ -829,7 +829,7 @@ async def test_immediate_scheduler_dedup_does_not_create_duplicate_execution(tmp dedup_key="dup-immediate", ) - executions, total = await TaskManager.list_scheduler_executions(scheduler.id, limit=10) + executions, total = await ScheduleTaskManager.list_scheduler_executions(scheduler.id, limit=10) assert duplicate.id == scheduler.id assert total == 1 @@ -838,18 +838,18 @@ async def test_immediate_scheduler_dedup_does_not_create_duplicate_execution(tmp @pytest.mark.asyncio async def test_batch_cancel_counts_only_actual_cancellations(tmp_path: Path): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="批量取消", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - cancellable = await TaskManager.create_execution_from_scheduler( + cancellable = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) - completed = await TaskManager.create_execution_from_scheduler( + completed = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -858,7 +858,7 @@ async def test_batch_cancel_counts_only_actual_cancellations(tmp_path: Path): completed.completed_at = datetime.now(timezone.utc) await TaskStore.update_execution(completed) - cancelled = await TaskManager.batch_cancel([cancellable.id, completed.id]) + cancelled = await ScheduleTaskManager.batch_cancel([cancellable.id, completed.id]) assert cancelled == 1 @@ -874,28 +874,28 @@ async def fake_cancel_runtime(_cls, execution): cancelled_runtime_ids.append(execution.id) monkeypatch.setattr( - TaskManager, + ScheduleTaskManager, "_cancel_execution_runtime", classmethod(fake_cancel_runtime), ) - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="删除前清理普通计划", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - pending = await TaskManager.create_execution_from_scheduler( + pending = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, ) - queued = await TaskManager.create_execution_from_scheduler( + queued = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) - running = await TaskManager.create_execution_from_scheduler( + running = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -906,7 +906,7 @@ async def fake_cancel_runtime(_cls, execution): running.session_id = "ses_delete_running" await TaskStore.update_execution(running) await TaskStore.enqueue_execution_ref(running.id) - completed = await TaskManager.create_execution_from_scheduler( + completed = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -915,7 +915,7 @@ async def fake_cancel_runtime(_cls, execution): completed.completed_at = datetime.now(timezone.utc) await TaskStore.update_execution(completed) - deleted = await TaskManager.delete_scheduler(scheduler.id) + deleted = await ScheduleTaskManager.delete_scheduler(scheduler.id) assert deleted is True assert set(cancelled_runtime_ids) == { @@ -924,7 +924,7 @@ async def fake_cancel_runtime(_cls, execution): running.id, } for execution_id in (pending.id, queued.id, running.id, completed.id): - assert await TaskManager.get_execution(execution_id) is None + assert await ScheduleTaskManager.get_execution(execution_id) is None assert await TaskStore.get_queue_ref(queued.id) is None assert await TaskStore.get_queue_ref(running.id) is None @@ -946,13 +946,13 @@ async def fake_cancel_runtime(_cls, execution): @pytest.mark.asyncio async def test_store_init_normalizes_legacy_paused_execution(tmp_path: Path): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="兼容旧 paused 状态", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, @@ -974,7 +974,7 @@ async def test_store_init_normalizes_legacy_paused_execution(tmp_path: Path): await TaskStore.close() await TaskStore.init() - normalized = await TaskManager.get_execution(execution.id) + normalized = await ScheduleTaskManager.get_execution(execution.id) assert normalized is not None assert normalized.status == TaskStatus.CANCELLED @@ -985,12 +985,12 @@ async def test_store_init_normalizes_legacy_paused_execution(tmp_path: Path): @pytest.mark.asyncio async def test_cli_list_tasks_accepts_legacy_paused_status(monkeypatch: pytest.MonkeyPatch): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="CLI 兼容 paused 查询", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), ) - execution = await TaskManager.create_execution_from_scheduler( + execution = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -998,7 +998,7 @@ async def test_cli_list_tasks_accepts_legacy_paused_status(monkeypatch: pytest.M execution.status = TaskStatus.CANCELLED execution.completed_at = datetime.now(timezone.utc) await TaskStore.update_execution(execution) - await TaskManager.disable_scheduler(scheduler.id) + await ScheduleTaskManager.disable_scheduler(scheduler.id) printed: list[object] = [] monkeypatch.setattr(task_cli_commands.console, "print", lambda *args, **kwargs: printed.append(args)) @@ -1057,24 +1057,24 @@ async def fake_cancel_runtime(_cls, execution): cancelled_runtime_ids.append(execution.id) monkeypatch.setattr( - TaskManager, + ScheduleTaskManager, "_cancel_execution_runtime", classmethod(fake_cancel_runtime), ) - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="删除前清理内置计划", mode=SchedulerMode.CRON, trigger=TaskTrigger(cron="*/5 * * * *", timezone="Asia/Shanghai"), workspace_directory=str(tmp_path / "workspace"), dedup_key="builtin:test-delete-cleanup", ) - queued = await TaskManager.create_execution_from_scheduler( + queued = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.SCHEDULED, enqueue=True, ) - running = await TaskManager.create_execution_from_scheduler( + running = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.SCHEDULED, enqueue=False, @@ -1085,7 +1085,7 @@ async def fake_cancel_runtime(_cls, execution): running.session_id = "ses_builtin_running" await TaskStore.update_execution(running) await TaskStore.enqueue_execution_ref(running.id) - completed = await TaskManager.create_execution_from_scheduler( + completed = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.SCHEDULED, enqueue=False, @@ -1094,11 +1094,11 @@ async def fake_cancel_runtime(_cls, execution): completed.completed_at = datetime.now(timezone.utc) await TaskStore.update_execution(completed) - deleted = await TaskManager.delete_scheduler(scheduler.id) - archived = await TaskManager.get_scheduler(scheduler.id) - queued_execution = await TaskManager.get_execution(queued.id) - running_execution = await TaskManager.get_execution(running.id) - completed_execution = await TaskManager.get_execution(completed.id) + deleted = await ScheduleTaskManager.delete_scheduler(scheduler.id) + archived = await ScheduleTaskManager.get_scheduler(scheduler.id) + queued_execution = await ScheduleTaskManager.get_execution(queued.id) + running_execution = await ScheduleTaskManager.get_execution(running.id) + completed_execution = await ScheduleTaskManager.get_execution(completed.id) assert deleted is True assert archived is not None @@ -1116,39 +1116,39 @@ async def fake_cancel_runtime(_cls, execution): @pytest.mark.asyncio async def test_dashboard_counts_exclude_immediate_once_schedulers(tmp_path: Path): - await TaskManager.create_scheduler( + await ScheduleTaskManager.create_scheduler( title="队列任务模板", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace-1"), ) - await TaskManager.create_scheduler( + await ScheduleTaskManager.create_scheduler( title="单次计划任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False, run_at=datetime.now(timezone.utc) + timedelta(hours=1)), workspace_directory=str(tmp_path / "workspace-2"), ) - await TaskManager.create_scheduler( + await ScheduleTaskManager.create_scheduler( title="循环计划任务", mode=SchedulerMode.CRON, trigger=TaskTrigger(cron="*/5 * * * *", timezone="Asia/Shanghai"), workspace_directory=str(tmp_path / "workspace-3"), ) - counts = await TaskManager.dashboard() + counts = await ScheduleTaskManager.dashboard() assert counts["scheduled_active"] == 2 @pytest.mark.asyncio async def test_get_unviewed_results_includes_unread_and_notified_only(tmp_path: Path): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="未读结果", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - unread = await TaskManager.create_execution_from_scheduler( + unread = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -1158,7 +1158,7 @@ async def test_get_unviewed_results_includes_unread_and_notified_only(tmp_path: unread.delivery_status = DeliveryStatus.UNREAD await TaskStore.update_execution(unread) - notified = await TaskManager.create_execution_from_scheduler( + notified = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -1168,7 +1168,7 @@ async def test_get_unviewed_results_includes_unread_and_notified_only(tmp_path: notified.delivery_status = DeliveryStatus.NOTIFIED await TaskStore.update_execution(notified) - viewed = await TaskManager.create_execution_from_scheduler( + viewed = await ScheduleTaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -1178,7 +1178,7 @@ async def test_get_unviewed_results_includes_unread_and_notified_only(tmp_path: viewed.delivery_status = DeliveryStatus.VIEWED await TaskStore.update_execution(viewed) - results = await TaskManager.get_unviewed_results() + results = await ScheduleTaskManager.get_unviewed_results() result_ids = {item.id for item in results} assert unread.id in result_ids @@ -1191,11 +1191,11 @@ async def test_task_page_notice_drops_legacy_tables_after_third_display(): db = await TaskStore.raw_db() await db.execute("CREATE TABLE tasks (id TEXT PRIMARY KEY)") await db.commit() - TaskManager._write_migration_state({"failed": True, "notice_count": 0}) + ScheduleTaskManager._write_migration_state({"failed": True, "notice_count": 0}) - first = await TaskManager.get_task_page_notice() - second = await TaskManager.get_task_page_notice() - third = await TaskManager.get_task_page_notice() + first = await ScheduleTaskManager.get_task_page_notice() + second = await ScheduleTaskManager.get_task_page_notice() + third = await ScheduleTaskManager.get_task_page_notice() assert first == { "message": "系统更新了任务表的存储，旧表自动迁移失败，请手动重建任务 scheduler", @@ -1203,4 +1203,4 @@ async def test_task_page_notice_drops_legacy_tables_after_third_display(): } assert second is not None and second["displayCount"] == 2 assert third is not None and third["displayCount"] == 3 - assert TaskManager._legacy_tables_exist() is False + assert ScheduleTaskManager._legacy_tables_exist() is False diff --git a/tests/tool/test_builtin_management_tools.py b/tests/tool/test_builtin_management_tools.py index b1a61d86d..d5564f6d0 100644 --- a/tests/tool/test_builtin_management_tools.py +++ b/tests/tool/test_builtin_management_tools.py @@ -30,10 +30,10 @@ def test_lsp_remains_non_native_by_default() -> None: assert tool.info.native is False -def test_task_remains_non_native_when_declared() -> None: +def test_delegate_task_remains_non_native_when_declared() -> None: ToolRegistry.init() - tool = ToolRegistry.get("task") + tool = ToolRegistry.get("delegate_task") assert tool is not None assert tool.info.native is False diff --git a/tests/tool/test_delegate_task_compat.py b/tests/tool/test_delegate_task_compat.py index 644718c10..43ceeea6b 100644 --- a/tests/tool/test_delegate_task_compat.py +++ b/tests/tool/test_delegate_task_compat.py @@ -20,9 +20,8 @@ def test_delegate_description_requires_material_delegation_benefit(self): assert "Do not delegate trivial edits" in DESCRIPTION assert "The task requires multiple steps or research" not in DESCRIPTION - @pytest.mark.parametrize("tool_name", ["delegate_task", "task"]) - def test_delegate_schema_exposes_only_subagent_routing(self, tool_name): - schema = ToolRegistry.get_schema(tool_name) + def test_delegate_schema_exposes_only_subagent_routing(self): + schema = ToolRegistry.get_schema("delegate_task") assert schema is not None assert "prompt" in schema.required assert "subagent_type" in schema.properties @@ -75,7 +74,6 @@ async def test_delegate_task_derives_description_and_ignores_blank_skills(self): permissions = create_session.await_args.kwargs["permission"] denied_permissions = {rule.permission for rule in permissions if rule.action == "deny"} assert "delegate_task" not in denied_permissions - assert "task" not in denied_permissions @pytest.mark.asyncio async def test_delegate_task_explicit_model_override_is_pinned(self): @@ -115,10 +113,9 @@ async def test_delegate_task_explicit_model_override_is_pinned(self): assert loop_run.await_args.kwargs["model_id"] == "claude-haiku-4-5" @pytest.mark.asyncio - @pytest.mark.parametrize("tool_name", ["delegate_task", "task"]) - async def test_delegate_tools_reject_removed_category_parameter(self, tool_name): + async def test_delegate_task_rejects_removed_category_parameter(self): result = await ToolRegistry.execute( - tool_name, + "delegate_task", ctx=_make_ctx(), category="quick", prompt="Summarize the diff", @@ -128,10 +125,9 @@ async def test_delegate_tools_reject_removed_category_parameter(self, tool_name) assert "unknown parameters: category" in (result.error or "") @pytest.mark.asyncio - @pytest.mark.parametrize("tool_name", ["delegate_task", "task"]) - async def test_delegate_tools_accept_deprecated_command_parameter(self, tool_name): + async def test_delegate_task_accepts_deprecated_command_parameter(self): result = await ToolRegistry.execute( - tool_name, + "delegate_task", ctx=_make_ctx(), command="legacy-tracking-command", prompt="Summarize the diff", diff --git a/tests/tool/test_task_center_compat.py b/tests/tool/test_task_center_compat.py index afd9dd9ea..5f804f971 100644 --- a/tests/tool/test_task_center_compat.py +++ b/tests/tool/test_task_center_compat.py @@ -7,7 +7,7 @@ import flocks.tool.task.schedule_task_center # noqa: F401 from flocks.config.config import Config from flocks.storage.storage import Storage -from flocks.task.manager import TaskManager +from flocks.task.schedule_task_manager import ScheduleTaskManager from flocks.task.models import SchedulerMode, TaskPriority, TaskScheduler, TaskTrigger from flocks.task.store import TaskStore from flocks.tool.registry import ToolContext, ToolRegistry @@ -27,8 +27,8 @@ async def isolated_task_env(tmp_path: pytest.TempPathFactory, monkeypatch: pytes Config._cached_config = None Storage._db_path = None Storage._initialized = False - TaskManager._instance = None - TaskManager._startup_error = None + ScheduleTaskManager._instance = None + ScheduleTaskManager._startup_error = None TaskStore._initialized = False TaskStore._conn = None @@ -37,14 +37,14 @@ async def isolated_task_env(tmp_path: pytest.TempPathFactory, monkeypatch: pytes yield - await TaskManager.stop() + await ScheduleTaskManager.stop() await TaskStore.close() Config._global_config = None Config._cached_config = None Storage._db_path = None Storage._initialized = False - TaskManager._instance = None - TaskManager._startup_error = None + ScheduleTaskManager._instance = None + ScheduleTaskManager._startup_error = None TaskStore._initialized = False TaskStore._conn = None @@ -102,7 +102,7 @@ async def test_task_create_accepts_legacy_schedule_type_alias(self): assert result.success is True - schedulers, total = await TaskManager.list_schedulers(limit=10) + schedulers, total = await ScheduleTaskManager.list_schedulers(limit=10) assert total == 1 scheduler = schedulers[0] assert scheduler.mode == SchedulerMode.CRON @@ -123,7 +123,7 @@ async def test_task_create_infers_scheduled_type_from_cron(self): assert result.success is True - schedulers, total = await TaskManager.list_schedulers(limit=10) + schedulers, total = await ScheduleTaskManager.list_schedulers(limit=10) assert total == 1 scheduler = schedulers[0] assert scheduler.mode == SchedulerMode.CRON @@ -144,7 +144,7 @@ async def test_task_create_accepts_legacy_schedule_action_and_enabled_fields(sel assert result.success is True - schedulers, total = await TaskManager.list_schedulers(limit=10) + schedulers, total = await ScheduleTaskManager.list_schedulers(limit=10) assert total == 1 scheduler = schedulers[0] assert scheduler.mode == SchedulerMode.CRON @@ -153,7 +153,7 @@ async def test_task_create_accepts_legacy_schedule_action_and_enabled_fields(sel @pytest.mark.asyncio async def test_task_update_defaults_to_update_and_accepts_schedule_fields(self): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="原始任务", description="原始描述", mode=SchedulerMode.ONCE, @@ -177,7 +177,7 @@ async def test_task_update_defaults_to_update_and_accepts_schedule_fields(self): assert result.success is True - updated = await TaskManager.get_scheduler(scheduler.id) + updated = await ScheduleTaskManager.get_scheduler(scheduler.id) assert updated is not None assert updated.mode == SchedulerMode.CRON assert updated.description == "更新后的描述" @@ -206,7 +206,7 @@ async def test_task_create_rejects_run_once_without_time_instead_of_immediate(se assert result.error is not None assert "run_at" in result.error or "cron" in result.error - _, total = await TaskManager.list_schedulers(limit=10) + _, total = await ScheduleTaskManager.list_schedulers(limit=10) assert total == 0 @pytest.mark.asyncio @@ -224,7 +224,7 @@ async def test_task_create_schedule_json_accepts_string_boolean_run_once(self): assert result.success is True - schedulers, total = await TaskManager.list_schedulers(limit=10) + schedulers, total = await ScheduleTaskManager.list_schedulers(limit=10) assert total == 1 scheduler = schedulers[0] assert scheduler.mode == SchedulerMode.CRON @@ -234,7 +234,7 @@ async def test_task_create_schedule_json_accepts_string_boolean_run_once(self): @pytest.mark.asyncio async def test_task_create_rejects_six_field_cron_with_hint(self): result = await ToolRegistry.execute( - "task_create", + "schedule_task_create", ctx=_make_ctx(), title="每天早上 6 点执行", description="误传了 6 段 Quartz cron", @@ -247,12 +247,12 @@ async def test_task_create_rejects_six_field_cron_with_hint(self): assert "only 5-field cron is supported" in result.error assert "`0 6 * * *`" in result.error - _, total = await TaskManager.list_schedulers(limit=10) + _, total = await ScheduleTaskManager.list_schedulers(limit=10) assert total == 0 @pytest.mark.asyncio async def test_task_update_rejects_six_field_cron_with_hint(self): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="待更新的定时任务", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -262,7 +262,7 @@ async def test_task_update_rejects_six_field_cron_with_hint(self): ) result = await ToolRegistry.execute( - "task_update", + "schedule_task_update", ctx=_make_ctx(), task_id=scheduler.id, cron="0 0 6 * * *", @@ -274,7 +274,7 @@ async def test_task_update_rejects_six_field_cron_with_hint(self): assert "only 5-field cron is supported" in result.error assert "`0 6 * * *`" in result.error - unchanged = await TaskManager.get_scheduler(scheduler.id) + unchanged = await ScheduleTaskManager.get_scheduler(scheduler.id) assert unchanged is not None assert unchanged.trigger.cron == "*/5 * * * *" @@ -294,7 +294,7 @@ async def test_task_status_formats_scheduler_times_in_schedule_timezone(self): await TaskStore.create_scheduler(scheduler) result = await ToolRegistry.execute( - "task_status", + "schedule_task_status", ctx=_make_ctx(), task_id=scheduler.id, ) @@ -307,7 +307,7 @@ async def test_task_status_formats_scheduler_times_in_schedule_timezone(self): @pytest.mark.asyncio async def test_task_update_can_disable_and_enable_scheduled_task(self): - scheduler = await TaskManager.create_scheduler( + scheduler = await ScheduleTaskManager.create_scheduler( title="可停止的定时任务", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -324,7 +324,7 @@ async def test_task_update_can_disable_and_enable_scheduled_task(self): ) assert disable_result.success is True - disabled = await TaskManager.get_scheduler(scheduler.id) + disabled = await ScheduleTaskManager.get_scheduler(scheduler.id) assert disabled is not None assert disabled.status.value == "disabled" @@ -336,6 +336,6 @@ async def test_task_update_can_disable_and_enable_scheduled_task(self): ) assert enable_result.success is True - enabled = await TaskManager.get_scheduler(scheduler.id) + enabled = await ScheduleTaskManager.get_scheduler(scheduler.id) assert enabled is not None assert enabled.status.value == "active" diff --git a/tests/tool/test_task_list_routing.py b/tests/tool/test_task_list_routing.py index 8c867043a..0ac28a96e 100644 --- a/tests/tool/test_task_list_routing.py +++ b/tests/tool/test_task_list_routing.py @@ -15,7 +15,7 @@ from flocks.tool.registry import ToolContext, ToolResult from flocks.tool.task.schedule_task_center import schedule_task_list -_TM_PATH = "flocks.task.manager.TaskManager" +_TM_PATH = "flocks.task.schedule_task_manager.ScheduleTaskManager" def _ctx() -> ToolContext: diff --git a/tests/tool/test_task_model_pinning.py b/tests/tool/test_task_model_pinning.py deleted file mode 100644 index 4f75e10ab..000000000 --- a/tests/tool/test_task_model_pinning.py +++ /dev/null @@ -1,61 +0,0 @@ -from unittest.mock import AsyncMock, patch - -import pytest - -from flocks.tool.agent.task import task_tool -from flocks.tool.registry import ToolContext, ToolRegistry, ToolResult - - -def _make_ctx() -> ToolContext: - return ToolContext(session_id="test-session", message_id="test-message", agent="rex") - - -class TestTaskCompatibilityAlias: - def test_task_schema_does_not_expose_background_execution(self): - schema = ToolRegistry.get_schema("task") - assert schema is not None - assert "run_in_background" not in schema.properties - # Legacy batch shape is gone. - assert "tasks" not in schema.properties - - @pytest.mark.asyncio - async def test_task_tool_rejects_background_execution_when_called_directly(self): - result = await task_tool( - _make_ctx(), - description="delegate explore", - prompt="Inspect the repository", - subagent_type="explore", - run_in_background=True, - ) - - assert result.success is False - assert "Background subagent execution is disabled" in (result.error or "") - - @pytest.mark.asyncio - async def test_task_tool_forwards_single_call_to_delegate_task(self): - delegate_result = ToolResult( - success=True, - output="ok", - metadata={"sessionId": "ses-child"}, - ) - - with patch( - "flocks.tool.agent.task.delegate_task_tool", - AsyncMock(return_value=delegate_result), - ) as delegate: - result = await task_tool( - _make_ctx(), - description="delegate explore", - prompt="Inspect the repository", - subagent_type="explore", - model="openai/gpt-5", - ) - - assert result is delegate_result - delegate.assert_awaited_once() - kwargs = delegate.await_args.kwargs - assert kwargs["description"] == "delegate explore" - assert kwargs["prompt"] == "Inspect the repository" - assert kwargs["subagent_type"] == "explore" - assert kwargs["run_in_background"] is False - assert kwargs["model"] == "openai/gpt-5" diff --git a/tests/tool/test_tools.py b/tests/tool/test_tools.py index dfba07484..4ce0a37ee 100644 --- a/tests/tool/test_tools.py +++ b/tests/tool/test_tools.py @@ -156,7 +156,7 @@ def test_expected_tools_registered(self): # P1 tools "webfetch", "todo", "question", # P2 tools - "task", "lsp", "skill_load", + "delegate_task", "lsp", "skill_load", # P3 tools (2) "websearch", "apply_patch", ] @@ -801,13 +801,13 @@ async def test_webfetch_schema(self): # P2 Tools Tests # ============================================================================= -class TestTaskTool: - """Test the task tool""" +class TestDelegateTaskTool: + """Test the delegate_task tool""" @pytest.mark.asyncio - async def test_task_exists(self): - """Test that task tool is registered""" - tool = ToolRegistry.get("task") + async def test_delegate_task_exists(self): + """Test that delegate_task tool is registered""" + tool = ToolRegistry.get("delegate_task") assert tool is not None diff --git a/tests/utils/test_id_compatibility.py b/tests/utils/test_id_compatibility.py index 22b1691a2..c26d9c509 100644 --- a/tests/utils/test_id_compatibility.py +++ b/tests/utils/test_id_compatibility.py @@ -26,7 +26,6 @@ def test_prefix_mappings(self): "call": "cal", "step": "stp", "agent": "agt", - "subtask": "stk", "event": "evt", "tqref": "tqr", "task": "tsk", diff --git a/tui/flocks/cli/cmd/tui/routes/session/index.tsx b/tui/flocks/cli/cmd/tui/routes/session/index.tsx index bade13c9d..899a1e6b1 100644 --- a/tui/flocks/cli/cmd/tui/routes/session/index.tsx +++ b/tui/flocks/cli/cmd/tui/routes/session/index.tsx @@ -40,7 +40,7 @@ import type { GrepTool } from "@/tool/grep" import type { EditTool } from "@/tool/edit" import type { ApplyPatchTool } from "@/tool/apply_patch" import type { WebFetchTool } from "@/tool/webfetch" -import type { TaskTool } from "@/tool/task" +import type { DelegateTaskTool } from "@/tool/delegate-task" import type { QuestionTool } from "@/tool/question" import { useKeyboard, useRenderer, useTerminalDimensions, type JSX } from "@opentui/solid" import { useSDK } from "@tui/context/sdk" @@ -1861,7 +1861,7 @@ function SubagentActivity(props: { ) } -function Task(props: ToolProps) { +function Task(props: ToolProps) { const { theme } = useTheme() const keybind = useKeybind() const { navigate } = useRoute() @@ -1977,7 +1977,7 @@ function DelegateTask(props: ToolProps) { navigate({ type: "session", sessionID: sessionId()! }) : undefined} part={props.part} > @@ -1997,7 +1997,7 @@ function DelegateTask(props: ToolProps) { > - {delegateInput.description || "subtask"} + {delegateInput.description || "delegated task"} {isBackground() ? " (background)" : ""} {statusText()} @@ -2016,7 +2016,7 @@ function DelegateTask(props: ToolProps) { part={props.part} > {agentName()}{" "} - "{delegateInput.description || "subtask"}" + "{delegateInput.description || "delegated task"}" {isBackground() ? " (bg)" : ""} diff --git a/tui/flocks/command/index.ts b/tui/flocks/command/index.ts index 976f1cd51..56120ce30 100644 --- a/tui/flocks/command/index.ts +++ b/tui/flocks/command/index.ts @@ -30,7 +30,6 @@ export namespace Command { // workaround for zod not supporting async functions natively so we use getters // https://zod.dev/v4/changelog?id=zfunction template: z.promise(z.string()).or(z.string()), - subtask: z.boolean().optional(), hints: z.array(z.string()), }) .meta({ @@ -70,10 +69,10 @@ export namespace Command { [Default.REVIEW]: { name: Default.REVIEW, description: "review changes [commit|branch|pr], defaults to uncommitted", + agent: "oracle", get template() { return PROMPT_REVIEW.replace("${path}", Instance.worktree) }, - subtask: true, hints: hints(PROMPT_REVIEW), }, } @@ -87,7 +86,6 @@ export namespace Command { get template() { return command.template }, - subtask: command.subtask, hints: hints(command.template), } } diff --git a/tui/flocks/config/config.ts b/tui/flocks/config/config.ts index 5085330ed..10a8e4327 100644 --- a/tui/flocks/config/config.ts +++ b/tui/flocks/config/config.ts @@ -559,7 +559,6 @@ export namespace Config { description: z.string().optional(), agent: z.string().optional(), model: z.string().optional(), - subtask: z.boolean().optional(), }) export type Command = z.infer diff --git a/tui/flocks/session/message-v2.ts b/tui/flocks/session/message-v2.ts index f2f3331e5..ff596603b 100644 --- a/tui/flocks/session/message-v2.ts +++ b/tui/flocks/session/message-v2.ts @@ -163,21 +163,6 @@ export namespace MessageV2 { }) export type CompactionPart = z.infer - export const SubtaskPart = PartBase.extend({ - type: z.literal("subtask"), - prompt: z.string(), - description: z.string(), - agent: z.string(), - model: z - .object({ - providerID: z.string(), - modelID: z.string(), - }) - .optional(), - command: z.string().optional(), - }) - export type SubtaskPart = z.infer - export const RetryPart = PartBase.extend({ type: z.literal("retry"), attempt: z.number(), @@ -329,7 +314,6 @@ export namespace MessageV2 { export const Part = z .discriminatedUnion("type", [ TextPart, - SubtaskPart, ReasoningPart, FilePart, ToolPart, @@ -466,12 +450,6 @@ export namespace MessageV2 { text: "What did we do so far?", }) } - if (part.type === "subtask") { - userMessage.parts.push({ - type: "text", - text: "The following tool was executed by the user", - }) - } } } diff --git a/tui/flocks/session/prompt.ts b/tui/flocks/session/prompt.ts index befddd369..c0caae740 100644 --- a/tui/flocks/session/prompt.ts +++ b/tui/flocks/session/prompt.ts @@ -34,7 +34,6 @@ import { SessionSummary } from "./summary" import { NamedError } from "@flocks-ai/util/error" import { fn } from "@/util/fn" import { SessionProcessor } from "./processor" -import { TaskTool } from "@/tool/task" import { Tool } from "@/tool/tool" import { PermissionNext } from "@/permission/next" import { SessionStatus } from "./status" @@ -200,16 +199,6 @@ export namespace SessionPrompt { .meta({ ref: "AgentPartInput", }), - MessageV2.SubtaskPart.omit({ - messageID: true, - sessionID: true, - }) - .partial({ - id: true, - }) - .meta({ - ref: "SubtaskPartInput", - }), ]), ), }) @@ -344,7 +333,7 @@ export namespace SessionPrompt { let lastUser: MessageV2.User | undefined let lastAssistant: MessageV2.Assistant | undefined let lastFinished: MessageV2.Assistant | undefined - let tasks: (MessageV2.CompactionPart | MessageV2.SubtaskPart)[] = [] + const pendingCompactions: MessageV2.CompactionPart[] = [] for (let i = msgs.length - 1; i >= 0; i--) { const msg = msgs[i] if (!lastUser && msg.info.role === "user") lastUser = msg.info as MessageV2.User @@ -352,9 +341,9 @@ export namespace SessionPrompt { if (!lastFinished && msg.info.role === "assistant" && msg.info.finish) lastFinished = msg.info as MessageV2.Assistant if (lastUser && lastFinished) break - const task = msg.parts.filter((part) => part.type === "compaction" || part.type === "subtask") - if (task && !lastFinished) { - tasks.push(...task) + const compactions = msg.parts.filter((part) => part.type === "compaction") + if (compactions.length > 0 && !lastFinished) { + pendingCompactions.push(...compactions) } } @@ -378,183 +367,16 @@ export namespace SessionPrompt { }) const model = await Provider.getModel(lastUser.model.providerID, lastUser.model.modelID) - const task = tasks.pop() - - // pending subtask - // TODO: centralize "invoke tool" logic - if (task?.type === "subtask") { - const taskTool = await TaskTool.init() - const taskModel = task.model ? await Provider.getModel(task.model.providerID, task.model.modelID) : model - const assistantMessage = (await Session.updateMessage({ - id: Identifier.ascending("message"), - role: "assistant", - parentID: lastUser.id, - sessionID, - mode: task.agent, - agent: task.agent, - path: { - cwd: Instance.directory, - root: Instance.worktree, - }, - cost: 0, - tokens: { - input: 0, - output: 0, - reasoning: 0, - cache: { read: 0, write: 0 }, - }, - modelID: taskModel.id, - providerID: taskModel.providerID, - time: { - created: Date.now(), - }, - })) as MessageV2.Assistant - let part = (await Session.updatePart({ - id: Identifier.ascending("part"), - messageID: assistantMessage.id, - sessionID: assistantMessage.sessionID, - type: "tool", - callID: ulid(), - tool: TaskTool.id, - state: { - status: "running", - input: { - prompt: task.prompt, - description: task.description, - subagent_type: task.agent, - command: task.command, - }, - time: { - start: Date.now(), - }, - }, - })) as MessageV2.ToolPart - const taskArgs = { - prompt: task.prompt, - description: task.description, - subagent_type: task.agent, - command: task.command, - } - await Plugin.trigger( - "tool.execute.before", - { - tool: "task", - sessionID, - callID: part.id, - }, - { args: taskArgs }, - ) - let executionError: Error | undefined - const taskAgent = await Agent.get(task.agent) - const taskCtx: Tool.Context = { - agent: task.agent, - messageID: assistantMessage.id, - sessionID: sessionID, - abort, - callID: part.callID, - extra: { bypassAgentCheck: true }, - async metadata(input) { - await Session.updatePart({ - ...part, - type: "tool", - state: { - ...part.state, - ...input, - }, - } satisfies MessageV2.ToolPart) - }, - async ask(req) { - await PermissionNext.ask({ - ...req, - sessionID: sessionID, - ruleset: PermissionNext.merge(taskAgent.permission, session.permission ?? []), - }) - }, - } - const result = await taskTool.execute(taskArgs, taskCtx).catch((error) => { - executionError = error - log.error("subtask execution failed", { error, agent: task.agent, description: task.description }) - return undefined - }) - await Plugin.trigger( - "tool.execute.after", - { - tool: "task", - sessionID, - callID: part.id, - }, - result, - ) - assistantMessage.finish = "tool-calls" - assistantMessage.time.completed = Date.now() - await Session.updateMessage(assistantMessage) - if (result && part.state.status === "running") { - await Session.updatePart({ - ...part, - state: { - status: "completed", - input: part.state.input, - title: result.title, - metadata: result.metadata, - output: result.output, - attachments: result.attachments, - time: { - ...part.state.time, - end: Date.now(), - }, - }, - } satisfies MessageV2.ToolPart) - } - if (!result) { - await Session.updatePart({ - ...part, - state: { - status: "error", - error: executionError ? `Tool execution failed: ${executionError.message}` : "Tool execution failed", - time: { - start: part.state.status === "running" ? part.state.time.start : Date.now(), - end: Date.now(), - }, - metadata: part.metadata, - input: part.state.input, - }, - } satisfies MessageV2.ToolPart) - } - - // Add synthetic user message to prevent certain reasoning models from erroring - // If we create assistant messages w/ out user ones following mid loop thinking signatures - // will be missing and it can cause errors for models like gemini for example - const summaryUserMsg: MessageV2.User = { - id: Identifier.ascending("message"), - sessionID, - role: "user", - time: { - created: Date.now(), - }, - agent: lastUser.agent, - model: lastUser.model, - } - await Session.updateMessage(summaryUserMsg) - await Session.updatePart({ - id: Identifier.ascending("part"), - messageID: summaryUserMsg.id, - sessionID, - type: "text", - text: "Summarize the task tool output above and continue with your task.", - synthetic: true, - } satisfies MessageV2.TextPart) - - continue - } + const pendingCompaction = pendingCompactions.pop() // pending compaction - if (task?.type === "compaction") { + if (pendingCompaction) { const result = await SessionCompaction.process({ messages: msgs, parentID: lastUser.id, abort, sessionID, - auto: task.auto, + auto: pendingCompaction.auto, }) if (result === "stop") break continue @@ -1742,30 +1564,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the } const templateParts = await resolvePromptParts(template) - const isSubtask = (agent.mode === "subagent" && command.subtask !== false) || command.subtask === true - const parts = isSubtask - ? [ - { - type: "subtask" as const, - agent: agent.name, - description: command.description ?? "", - command: input.command, - model: { - providerID: taskModel.providerID, - modelID: taskModel.modelID, - }, - // TODO: how can we make task tool accept a more complex input? - prompt: templateParts.find((y) => y.type === "text")?.text ?? "", - }, - ] - : [...templateParts, ...(input.parts ?? [])] - - const userAgent = isSubtask ? (input.agent ?? (await Agent.defaultAgent())) : agentName - const userModel = isSubtask - ? input.model - ? Provider.parseModel(input.model) - : await lastModel(input.sessionID) - : taskModel + const parts = [...templateParts, ...(input.parts ?? [])] await Plugin.trigger( "command.execute.before", @@ -1780,8 +1579,8 @@ NOTE: At any point in time through this workflow you should feel free to ask the const result = (await prompt({ sessionID: input.sessionID, messageID: input.messageID, - model: userModel, - agent: userAgent, + model: taskModel, + agent: agentName, parts, variant: input.variant, })) as MessageV2.WithParts @@ -1817,15 +1616,9 @@ NOTE: At any point in time through this workflow you should feel free to ask the if (!isFirst) return // Gather all messages up to and including the first real user message for context - // This includes any shell/subtask executions that preceded the user's first prompt const contextMessages = input.history.slice(0, firstRealUserIdx + 1) const firstRealUser = contextMessages[firstRealUserIdx] - // For subtask-only messages (from command invocations), extract the prompt directly - // since toModelMessage converts subtask parts to generic "The following tool was executed by the user" - const subtaskParts = firstRealUser.parts.filter((p) => p.type === "subtask") as MessageV2.SubtaskPart[] - const hasOnlySubtaskParts = subtaskParts.length > 0 && firstRealUser.parts.every((p) => p.type === "subtask") - const agent = await Agent.get("title") if (!agent) return const result = await LLM.stream({ @@ -1848,9 +1641,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the role: "user", content: "Generate a title for this conversation:\n", }, - ...(hasOnlySubtaskParts - ? [{ role: "user" as const, content: subtaskParts.map((p) => p.prompt).join("\n") }] - : MessageV2.toModelMessage(contextMessages)), + ...MessageV2.toModelMessage(contextMessages), ], }) const text = await result.text.catch((err) => log.error("failed to generate title", { error: err })) diff --git a/tui/flocks/session/prompt/anthropic-20250930.txt b/tui/flocks/session/prompt/anthropic-20250930.txt index 676c4d8dc..01ac8e5b5 100644 --- a/tui/flocks/session/prompt/anthropic-20250930.txt +++ b/tui/flocks/session/prompt/anthropic-20250930.txt @@ -129,10 +129,10 @@ The user will primarily request you perform software engineering tasks. This inc # Tool usage policy -- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. diff --git a/tui/flocks/session/prompt/anthropic.txt b/tui/flocks/session/prompt/anthropic.txt index 7a0e5fd5c..0709a5c30 100644 --- a/tui/flocks/session/prompt/anthropic.txt +++ b/tui/flocks/session/prompt/anthropic.txt @@ -73,20 +73,20 @@ The user will primarily request you perform SecOps tasks. This includes security # Tool usage policy -- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. - You can call multiple tools in a single response. If you intend to call multiple tools and there are no dependencies between them, make all independent tool calls in parallel. Maximize use of parallel tool calls where possible to increase efficiency. However, if some tool calls depend on previous calls to inform dependent values, do NOT call these tools in parallel and instead call them sequentially. For instance, if one operation must complete before another starts, run these operations sequentially instead. Never use placeholders or guess missing parameters in tool calls. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. -- VERY IMPORTANT: When exploring the codebase to gather context or to answer a question that is not a needle query for a specific file/class/function, it is CRITICAL that you use the Task tool instead of running search commands directly. +- VERY IMPORTANT: When exploring the codebase to gather context or to answer a question that is not a needle query for a specific file/class/function, it is CRITICAL that you use `delegate_task` instead of running search commands directly. user: Where are errors from the client handled? -assistant: [Uses the Task tool to find the files that handle client errors instead of using Glob or Grep directly] +assistant: [Uses `delegate_task` to find the files that handle client errors instead of using Glob or Grep directly] user: What is the codebase structure? -assistant: [Uses the Task tool] +assistant: [Uses `delegate_task`] IMPORTANT: Always use `todo(action="write")` to plan and track tasks throughout the conversation. diff --git a/tui/flocks/tool/task.ts b/tui/flocks/tool/delegate-task.ts similarity index 97% rename from tui/flocks/tool/task.ts rename to tui/flocks/tool/delegate-task.ts index f98316b39..e2b9c5d92 100644 --- a/tui/flocks/tool/task.ts +++ b/tui/flocks/tool/delegate-task.ts @@ -1,5 +1,5 @@ import { Tool } from "./tool" -import DESCRIPTION from "./task.txt" +import DESCRIPTION from "./delegate-task.txt" import z from "zod" import { Session } from "../session" import { Bus } from "../bus" @@ -20,7 +20,7 @@ const parameters = z.object({ command: z.string().describe("The command that triggered this task").optional(), }) -export const TaskTool = Tool.define("task", async (ctx) => { +export const DelegateTaskTool = Tool.define("delegate_task", async (ctx) => { const agents = await Agent.list().then((x) => x.filter((a) => a.mode !== "primary")) // Filter agents by permissions if agent provided @@ -41,7 +41,6 @@ export const TaskTool = Tool.define("task", async (ctx) => { async execute(params: z.infer, ctx) { const config = await Config.get() - // Skip permission check when user explicitly invoked via @ or command subtask if (!ctx.extra?.bypassAgentCheck) { await ctx.ask({ permission: "task", diff --git a/tui/flocks/tool/task.txt b/tui/flocks/tool/delegate-task.txt similarity index 79% rename from tui/flocks/tool/task.txt rename to tui/flocks/tool/delegate-task.txt index 7af2a6f60..21258b793 100644 --- a/tui/flocks/tool/task.txt +++ b/tui/flocks/tool/delegate-task.txt @@ -1,17 +1,17 @@ -Launch a new agent to handle complex, multistep tasks autonomously. +Delegate a complex, multistep task to another agent. Available agent types and the tools they have access to: {agents} -When using the Task tool, you must specify a subagent_type parameter to select which agent type to use. +When using the delegate_task tool, you must specify a subagent_type parameter to select which agent type to use. -When to use the Task tool: -- When you are instructed to execute custom slash commands. Use the Task tool with the slash command invocation as the entire prompt. The slash command can take arguments. For example: Task(description="Check the file", prompt="/check-file path/to/file.py") +When to use the delegate_task tool: +- When you are instructed to execute custom slash commands. Use delegate_task with the slash command invocation as the entire prompt. The slash command can take arguments. For example: delegate_task(description="Check the file", prompt="/check-file path/to/file.py") -When NOT to use the Task tool: -- If you want to read a specific file path, use the Read or Glob tool instead of the Task tool, to find the match more quickly +When NOT to use the delegate_task tool: +- If you want to read a specific file path, use the Read or Glob tool instead of delegate_task, to find the match more quickly - If you are searching for a specific class definition like "class Foo", use the Glob tool instead, to find the match more quickly -- If you are searching for code within a specific file or set of 2-3 files, use the Read tool instead of the Task tool, to find the match more quickly +- If you are searching for code within a specific file or set of 2-3 files, use the Read tool instead of delegate_task, to find the match more quickly - Other tasks that are not related to the agent descriptions above @@ -48,7 +48,7 @@ function isPrime(n) { Since a significant piece of code was written and the task was completed, now use the code-reviewer agent to review the code assistant: Now let me use the code-reviewer agent to review the code -assistant: Uses the Task tool to launch the code-reviewer agent +assistant: Uses delegate_task to launch the code-reviewer agent @@ -56,5 +56,5 @@ user: "Hello" Since the user is greeting, use the greeting-responder agent to respond with a friendly joke -assistant: "I'm going to use the Task tool to launch the with the greeting-responder agent" +assistant: "I'm going to use delegate_task to launch the greeting-responder agent" diff --git a/tui/flocks/tool/registry.ts b/tui/flocks/tool/registry.ts index 4f4eed7a0..0714a55c7 100644 --- a/tui/flocks/tool/registry.ts +++ b/tui/flocks/tool/registry.ts @@ -4,7 +4,7 @@ import { EditTool } from "./edit" import { GlobTool } from "./glob" import { GrepTool } from "./grep" import { ReadTool } from "./read" -import { TaskTool } from "./task" +import { DelegateTaskTool } from "./delegate-task" import { TodoTool } from "./todo" import { WebFetchTool } from "./webfetch" import { WriteTool } from "./write" @@ -99,7 +99,7 @@ export namespace ToolRegistry { GrepTool, EditTool, WriteTool, - TaskTool, + DelegateTaskTool, WebFetchTool, TodoTool, WebSearchTool, diff --git a/tui/sdk/gen/types.gen.ts b/tui/sdk/gen/types.gen.ts index 8ac5c7342..ed4ba1843 100644 --- a/tui/sdk/gen/types.gen.ts +++ b/tui/sdk/gen/types.gen.ts @@ -383,15 +383,6 @@ export type CompactionPart = { export type Part = | TextPart - | { - id: string - sessionID: string - messageID: string - type: "subtask" - prompt: string - description: string - agent: string - } | ReasoningPart | FilePart | ToolPart @@ -1217,7 +1208,6 @@ export type Config = { description?: string agent?: string model?: string - subtask?: boolean } } watcher?: { @@ -1432,21 +1422,12 @@ export type AgentPartInput = { } } -export type SubtaskPartInput = { - id?: string - type: "subtask" - prompt: string - description: string - agent: string -} - export type Command = { name: string description?: string agent?: string model?: string template: string - subtask?: boolean } export type Model = { @@ -2591,7 +2572,7 @@ export type SessionPromptData = { tools?: { [key: string]: boolean } - parts: Array + parts: Array } path: { /** @@ -2686,7 +2667,7 @@ export type SessionPromptAsyncData = { tools?: { [key: string]: boolean } - parts: Array + parts: Array } path: { /** diff --git a/tui/sdk/v2/gen/sdk.gen.ts b/tui/sdk/v2/gen/sdk.gen.ts index a84c70938..033e235b6 100644 --- a/tui/sdk/v2/gen/sdk.gen.ts +++ b/tui/sdk/v2/gen/sdk.gen.ts @@ -134,7 +134,6 @@ import type { SessionUnshareResponses, SessionUpdateErrors, SessionUpdateResponses, - SubtaskPartInput, TextPartInput, ToolIdsErrors, ToolIdsResponses, @@ -1364,7 +1363,7 @@ export class Session extends HeyApiClient { } system?: string variant?: string - parts?: Array + parts?: Array }, options?: Options, ) { @@ -1452,7 +1451,7 @@ export class Session extends HeyApiClient { } system?: string variant?: string - parts?: Array + parts?: Array }, options?: Options, ) { diff --git a/tui/sdk/v2/gen/types.gen.ts b/tui/sdk/v2/gen/types.gen.ts index 77f869dc6..0109b8e50 100644 --- a/tui/sdk/v2/gen/types.gen.ts +++ b/tui/sdk/v2/gen/types.gen.ts @@ -429,20 +429,6 @@ export type CompactionPart = { export type Part = | TextPart - | { - id: string - sessionID: string - messageID: string - type: "subtask" - prompt: string - description: string - agent: string - model?: { - providerID: string - modelID: string - } - command?: string - } | ReasoningPart | FilePart | ToolPart @@ -1617,7 +1603,6 @@ export type Config = { description?: string agent?: string model?: string - subtask?: boolean } } watcher?: { @@ -1953,19 +1938,6 @@ export type AgentPartInput = { } } -export type SubtaskPartInput = { - id?: string - type: "subtask" - prompt: string - description: string - agent: string - model?: { - providerID: string - modelID: string - } - command?: string -} - export type ProviderAuthMethod = { type: "oauth" | "api" label: string @@ -2071,7 +2043,6 @@ export type Command = { model?: string mcp?: boolean template: string - subtask?: boolean hints: Array } @@ -3226,7 +3197,7 @@ export type SessionPromptData = { } system?: string variant?: string - parts: Array + parts: Array } path: { /** @@ -3413,7 +3384,7 @@ export type SessionPromptAsyncData = { } system?: string variant?: string - parts: Array + parts: Array } path: { /** diff --git a/webui/src/api/skill.ts b/webui/src/api/skill.ts index 254520cec..00c2f9dad 100644 --- a/webui/src/api/skill.ts +++ b/webui/src/api/skill.ts @@ -39,7 +39,6 @@ export interface Command { template: string; agent?: string; model?: string; - subtask?: boolean; hidden: boolean; aliases: string[]; visible_surfaces: string[]; From 942b6bec412eb6e941c5932614cc2d4dd51cde3b Mon Sep 17 00:00:00 2001 From: xiami762 Date: Fri, 7 Aug 2026 21:41:01 +0800 Subject: [PATCH 07/15] refactor: clarify runtime loop context --- flocks/session/runtime/agent_loop.py | 6 ++-- flocks/session/runtime/session_turn.py | 6 ++-- flocks/session/runtime/step_engine.py | 2 +- flocks/session/session_loop.py | 34 ++++++++++------------ tests/session/runtime/test_agent_loop.py | 2 +- tests/session/runtime/test_step_engine.py | 6 ++-- tests/session/test_runner_llm_hooks.py | 10 +++++++ tests/session/test_session_abort_inject.py | 16 +++++----- 8 files changed, 44 insertions(+), 38 deletions(-) diff --git a/flocks/session/runtime/agent_loop.py b/flocks/session/runtime/agent_loop.py index 9a54133ed..da57a86d2 100644 --- a/flocks/session/runtime/agent_loop.py +++ b/flocks/session/runtime/agent_loop.py @@ -9,7 +9,7 @@ StepAction, TurnPreparationStatus, ) -from flocks.session.runtime.session_turn import SessionTurn +from flocks.session.runtime.session_turn import LoopContext from flocks.session.runtime.step_engine import StepCancelled, StepEngine from flocks.utils.log import Log @@ -22,7 +22,7 @@ class AgentLoop: async def run( self, - turn: SessionTurn, + turn: LoopContext, engine: StepEngine, ) -> AgentRunOutcome[MessageInfo]: """Run the current logical input to a session-level boundary.""" @@ -54,7 +54,7 @@ async def run( state=state, last_message=last_message, error=( - "SessionTurn returned READY without a model-turn " + "LoopContext returned READY without a model-turn " "snapshot" ), ) diff --git a/flocks/session/runtime/session_turn.py b/flocks/session/runtime/session_turn.py index f970a5e22..a358175a7 100644 --- a/flocks/session/runtime/session_turn.py +++ b/flocks/session/runtime/session_turn.py @@ -90,11 +90,11 @@ class LoopResult: @dataclass -class SessionTurn: - """Own state and persistence boundaries for one logical user input. +class LoopContext: + """Own state and persistence boundaries for one session loop run. Supports: - - Message iteration + - Logical user-turn and message iteration - Compaction triggers - Reminder injection """ diff --git a/flocks/session/runtime/step_engine.py b/flocks/session/runtime/step_engine.py index caec9f5da..18ea14d79 100644 --- a/flocks/session/runtime/step_engine.py +++ b/flocks/session/runtime/step_engine.py @@ -282,7 +282,7 @@ def from_turn( turn: Any, model_policy: Optional[ModelRoutingPolicy] = None, ) -> "StepEngine": - """Create the production engine for one stateful ``SessionTurn``.""" + """Create the production engine for one stateful ``LoopContext``.""" engine = cls( session=turn.session, provider_id=turn.provider_id, diff --git a/flocks/session/session_loop.py b/flocks/session/session_loop.py index ae96f7f7a..5ef9cf31e 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -31,9 +31,9 @@ ModelRoutingPolicy, ) from flocks.session.runtime.session_turn import ( + LoopContext, LoopCallbacks, LoopResult, - SessionTurn, ) from flocks.session.runtime.step_engine import StepEngine from flocks.session.session import ( @@ -45,31 +45,27 @@ log = Log.create(service="session.loop") -# Public compatibility name. There is only one state implementation. -LoopContext = SessionTurn - - @dataclass(frozen=True) class _SessionLease: """One process-local ownership record.""" session_id: str - turn: SessionTurn + turn: LoopContext class _SessionLeaseRegistry: """Keep lease bookkeeping out of the SessionLoop control flow.""" - def __init__(self, active_turns: MutableMapping[str, SessionTurn]): + def __init__(self, active_turns: MutableMapping[str, LoopContext]): self._active_turns = active_turns - def get(self, session_id: str) -> Optional[SessionTurn]: + def get(self, session_id: str) -> Optional[LoopContext]: return self._active_turns.get(session_id) def acquire( self, session_id: str, - turn: SessionTurn, + turn: LoopContext, ) -> Optional[_SessionLease]: if session_id in self._active_turns: return None @@ -87,7 +83,7 @@ def owns(self, lease: _SessionLease) -> bool: class SessionLoop: """Decide whether a persistent session should continue or settle.""" - _active_turns: ClassVar[dict[str, SessionTurn]] = {} + _active_turns: ClassVar[dict[str, LoopContext]] = {} _leases: ClassVar[_SessionLeaseRegistry] = _SessionLeaseRegistry( _active_turns, ) @@ -162,7 +158,7 @@ async def run( trace_offset = await cls._load_trace_offset(session_id) runtime_callbacks = callbacks or LoopCallbacks() - turn = SessionTurn( + turn = LoopContext( session=session, provider_id=provider_id, model_id=model_id, @@ -231,7 +227,7 @@ async def run( @classmethod async def _run_logical_input( cls, - turn: SessionTurn, + turn: LoopContext, ) -> AgentRunOutcome[Any]: """Execute one prepared logical input through AgentLoop.""" turn.reset() @@ -242,7 +238,7 @@ async def _run_logical_input( @staticmethod async def _should_continue( - turn: SessionTurn, + turn: LoopContext, continuation_policy: ContinuationPolicy, outcome: AgentRunOutcome[Any], ) -> bool: @@ -263,7 +259,7 @@ def is_running(cls, session_id: str) -> bool: return session_id in cls._active_turns @classmethod - def get_context(cls, session_id: str) -> Optional[SessionTurn]: + def get_context(cls, session_id: str) -> Optional[LoopContext]: """Return the active turn used by public session controls.""" return cls._active_turns.get(session_id) @@ -465,7 +461,7 @@ async def _resolve_model( @classmethod def _to_loop_result( cls, - turn: SessionTurn, + turn: LoopContext, outcome: AgentRunOutcome[Any], ) -> LoopResult: loop_error = ( @@ -512,7 +508,7 @@ def _to_loop_result( @staticmethod def _authorize_auto_failover( - turn: SessionTurn, + turn: LoopContext, requested: bool, ) -> None: if requested and is_model_auto_session_category( @@ -535,7 +531,7 @@ async def _load_trace_offset(session_id: str) -> int: async def _acquire_lease( cls, session_id: str, - turn: SessionTurn, + turn: LoopContext, ) -> _SessionLease | LoopResult: async with Session.lifecycle_lock(session_id): latest_session = await Session.get_by_id(session_id) @@ -584,7 +580,7 @@ async def _recover_orphan_tools(session_id: str) -> None: @staticmethod async def _handle_execution_error( - turn: SessionTurn, + turn: LoopContext, error: Exception, processed_user_id: Optional[str], ) -> AgentRunOutcome[Any]: @@ -670,7 +666,7 @@ def _finalize_release_state_locked(cls, lease: _SessionLease) -> None: @staticmethod async def _publish_released( - turn: SessionTurn, + turn: LoopContext, callbacks: LoopCallbacks, ) -> None: session_id = turn.session.id diff --git a/tests/session/runtime/test_agent_loop.py b/tests/session/runtime/test_agent_loop.py index a8410f4af..84acc85d5 100644 --- a/tests/session/runtime/test_agent_loop.py +++ b/tests/session/runtime/test_agent_loop.py @@ -321,4 +321,4 @@ async def test_ready_preparation_requires_snapshot() -> None: [], ) assert outcome.status == AgentRunStatus.FATAL_FAILURE - assert outcome.error == "SessionTurn returned READY without a model-turn snapshot" + assert outcome.error == "LoopContext returned READY without a model-turn snapshot" diff --git a/tests/session/runtime/test_step_engine.py b/tests/session/runtime/test_step_engine.py index 39302b85c..77dcd152d 100644 --- a/tests/session/runtime/test_step_engine.py +++ b/tests/session/runtime/test_step_engine.py @@ -16,16 +16,16 @@ RuntimeModel, StepResult, ) -from flocks.session.runtime.session_turn import LoopCallbacks, SessionTurn +from flocks.session.runtime.session_turn import LoopCallbacks, LoopContext from flocks.session.runtime.step_engine import StepCancelled, StepEngine from flocks.session.session import SessionInfo -def _turn(*, aborted: bool = False) -> SessionTurn: +def _turn(*, aborted: bool = False) -> LoopContext: abort_event = asyncio.Event() if aborted: abort_event.set() - return SessionTurn( + return LoopContext( session=SessionInfo.model_construct( id="session-1", projectID="project", diff --git a/tests/session/test_runner_llm_hooks.py b/tests/session/test_runner_llm_hooks.py index 95c93d73d..001b5babb 100644 --- a/tests/session/test_runner_llm_hooks.py +++ b/tests/session/test_runner_llm_hooks.py @@ -178,6 +178,11 @@ async def _after(payload, result): "run_llm_after", AsyncMock(side_effect=_after), ) + monkeypatch.setattr( + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock(return_value=True), + ) monkeypatch.setattr( runner_mod.StepEngine, "_end_observability", @@ -273,6 +278,11 @@ async def _after(payload, result): "run_llm_after", AsyncMock(side_effect=_after), ) + monkeypatch.setattr( + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock(return_value=True), + ) monkeypatch.setattr( runner_mod.StepEngine, "_end_observability", diff --git a/tests/session/test_session_abort_inject.py b/tests/session/test_session_abort_inject.py index e55ef6df4..bc7693b67 100644 --- a/tests/session/test_session_abort_inject.py +++ b/tests/session/test_session_abort_inject.py @@ -15,7 +15,7 @@ import pytest from flocks.session.runtime.continuation_policy import DEFAULT_CONTINUATION_POLICY -from flocks.session.runtime.session_turn import SessionTurn +from flocks.session.runtime.session_turn import LoopContext as RuntimeLoopContext from flocks.session.message import ToolPart, ToolStateCompleted from flocks.session.goal import GoalDecision from flocks.session.session_loop import ( @@ -263,7 +263,7 @@ def test_exit_when_assistant_after_user_and_finished(self): last_assistant = self._make_msg("msg_002", "assistant", finish="stop") # assistant.id > user.id → user.id < assistant.id → True → should exit - assert SessionTurn._should_exit(last_user, last_assistant) is True + assert RuntimeLoopContext._should_exit(last_user, last_assistant) is True def test_no_exit_when_user_injected_after_assistant(self): """Should NOT exit when a new user message appears after the assistant. @@ -275,34 +275,34 @@ def test_no_exit_when_user_injected_after_assistant(self): last_assistant = self._make_msg("msg_002", "assistant", finish="stop") # user.id > assistant.id → user.id < assistant.id → False → don't exit - assert SessionTurn._should_exit(last_user, last_assistant) is False + assert RuntimeLoopContext._should_exit(last_user, last_assistant) is False def test_no_exit_when_assistant_has_tool_calls(self): """Should NOT exit when assistant finish is 'tool-calls'.""" last_user = self._make_msg("msg_001", "user") last_assistant = self._make_msg("msg_002", "assistant", finish="tool-calls") - assert SessionTurn._should_exit(last_user, last_assistant) is False + assert RuntimeLoopContext._should_exit(last_user, last_assistant) is False def test_no_exit_when_no_assistant(self): """Should NOT exit when there is no assistant message yet.""" last_user = self._make_msg("msg_001", "user") - assert SessionTurn._should_exit(last_user, None) is False + assert RuntimeLoopContext._should_exit(last_user, None) is False def test_no_exit_when_assistant_finish_is_unknown(self): """Should NOT exit when finish reason is 'unknown'.""" last_user = self._make_msg("msg_001", "user") last_assistant = self._make_msg("msg_002", "assistant", finish="unknown") - assert SessionTurn._should_exit(last_user, last_assistant) is False + assert RuntimeLoopContext._should_exit(last_user, last_assistant) is False def test_no_exit_when_assistant_not_finished(self): """Should NOT exit when assistant has no finish status.""" last_user = self._make_msg("msg_001", "user") last_assistant = self._make_msg("msg_002", "assistant", finish=None) - assert SessionTurn._should_exit(last_user, last_assistant) is False + assert RuntimeLoopContext._should_exit(last_user, last_assistant) is False def test_no_exit_when_assistant_has_completed_tool_parts(self): """Should continue so completed tool results can be fed back to the model.""" @@ -310,7 +310,7 @@ def test_no_exit_when_assistant_has_completed_tool_parts(self): last_assistant = self._make_msg("msg_002", "assistant", finish="stop") last_assistant_parts = [_make_completed_tool_part(last_assistant.id)] - assert SessionTurn._should_exit( + assert RuntimeLoopContext._should_exit( last_user, last_assistant, last_assistant_parts, From eec04fab5dddcf637f5c2b66d7cc80239e430593 Mon Sep 17 00:00:00 2001 From: xiami762 Date: Fri, 7 Aug 2026 21:58:30 +0800 Subject: [PATCH 08/15] fix(runtime): preserve memory and deduplicate errors --- flocks/session/runtime/step_engine.py | 10 +--------- tests/session/runtime/test_step_engine.py | 19 +++++++++++++++++++ tests/session/test_runner_step.py | 21 +++++++++++++++++---- 3 files changed, 37 insertions(+), 13 deletions(-) diff --git a/flocks/session/runtime/step_engine.py b/flocks/session/runtime/step_engine.py index 18ea14d79..c2b359421 100644 --- a/flocks/session/runtime/step_engine.py +++ b/flocks/session/runtime/step_engine.py @@ -367,6 +367,7 @@ async def _run_candidate( ) self._turn_additional_context = turn.turn_additional_context self._session_start_pending = turn.session_start_pending + self._memory_bootstrap_data = turn.memory_bootstrap_data task = asyncio.create_task( self._process_step(list(snapshot.messages), snapshot.last_user), @@ -1338,8 +1339,6 @@ async def _process_step( error_dict=error_dict, visible_text=error_dict["data"]["displayMessage"], ) - if self.callbacks.on_error: - await self.callbacks.on_error(error_dict["data"]["displayMessage"]) return StepResult(action="stop", error=error_dict["data"]["displayMessage"]) # Apply config-based provider options (api_key/base_url) @@ -1378,8 +1377,6 @@ async def _process_step( error_dict=error_dict, visible_text=error_dict["data"]["displayMessage"], ) - if self.callbacks.on_error: - await self.callbacks.on_error(error_dict["data"]["displayMessage"]) return StepResult(action="stop", error=error_dict["data"]["displayMessage"]) # Build prompts and tools @@ -1695,8 +1692,6 @@ async def device_asset_prompt_factory() -> Optional[str]: decision=FailoverDecision(True, "empty_response"), attempts=empty_attempt, ) - if self.callbacks.on_error: - await self.callbacks.on_error(empty_error_msg) await Message.update( self.session.id, assistant_msg.id, @@ -1845,9 +1840,6 @@ async def device_asset_prompt_factory() -> Optional[str]: attempts=error_attempt, ) - if self.callbacks.on_error: - await self.callbacks.on_error(final_error_message) - # Update assistant message with error (must be dict, not string) await Message.update( self.session.id, diff --git a/tests/session/runtime/test_step_engine.py b/tests/session/runtime/test_step_engine.py index 77dcd152d..684701058 100644 --- a/tests/session/runtime/test_step_engine.py +++ b/tests/session/runtime/test_step_engine.py @@ -78,6 +78,25 @@ async def execute(messages, user): assert turn._current_step_task is None +@pytest.mark.asyncio +async def test_step_engine_refreshes_memory_loaded_after_construction() -> None: + last_user = SimpleNamespace(id="user-1") + turn = _turn() + engine = StepEngine.from_turn(turn) + loaded_memory = { + "instructions": "remember this", + "main_memory": {"content": "project context", "inject": True}, + } + turn.memory_bootstrap_data = loaded_memory + + async def execute(_messages, _user): + assert engine._memory_bootstrap_data is loaded_memory + return StepResult(action="stop", content="done") + + with patch.object(StepEngine, "_process_step", side_effect=execute): + await engine.run(_snapshot(last_user)) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("aborted", "expected_error"), diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index f23c04d48..25091d6c2 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -32,7 +32,7 @@ StepResult, ToolCall, ) -from flocks.session.runtime.session_turn import LoopCallbacks +from flocks.session.runtime.session_turn import LoopCallbacks, LoopContext from flocks.session.prompt import SessionPrompt from flocks.session.core.defaults import DEFAULT_MAX_TOOL_STEPS from flocks.session.session import Session, SessionInfo @@ -2371,7 +2371,7 @@ async def fake_call_llm(*_args, **_kwargs): assert call_count == 6 assert result.action == "stop" assert result.error == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE - runner.callbacks.on_error.assert_awaited_with(runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE) + runner.callbacks.on_error.assert_not_awaited() final_update = update_mock.await_args_list[-1].kwargs assert final_update["finish"] == "error" @@ -2529,7 +2529,7 @@ async def chat_stream(self, **kwargs): # noqa: ANN003 @pytest.mark.asyncio -async def test_process_step_persists_visible_error_when_provider_missing(monkeypatch): +async def test_step_boundary_reports_provider_missing_once(monkeypatch): runner = _make_runner("ses_runner_missing_provider_error") user = await Message.create( runner.session.id, @@ -2560,6 +2560,19 @@ async def on_error(error): ) result = await runner._process_step(messages, user) + turn = LoopContext( + session=runner.session, + provider_id=runner.provider_id, + model_id=runner.model_id, + agent_name="rex", + callbacks=runner.callbacks, + session_store=SimpleNamespace( + get_messages=AsyncMock( + return_value=await Message.list(runner.session.id), + ), + ), + ) + await turn.commit_step(result) messages_with_parts = await Message.list_with_parts(runner.session.id) assistant = next(item for item in messages_with_parts if item.info.role == MessageRole.ASSISTANT) visible_text_parts = [ @@ -2623,7 +2636,7 @@ async def on_error(error): assert result.action == "stop" assert result.error == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE - assert callback_errors == [runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE] + assert callback_errors == [] assert assistant.info.finish == "error" assert assistant.info.error["name"] == "ProviderConfigurationError" assert visible_text_parts From 614e0f311907190b7fd385bb863170c8dca52974 Mon Sep 17 00:00:00 2001 From: xiami762 Date: Fri, 7 Aug 2026 23:42:43 +0800 Subject: [PATCH 09/15] refactor(runtime): centralize prompt context assembly Collect runtime prompt inputs before deterministic block assembly, preserve provider-specific cache boundaries, and keep context usage estimation aligned with execution. Prevent Bash tool errors from being rendered twice in the session UI. --- flocks/session/context_usage.py | 24 +- flocks/session/prompt.py | 371 ++++++++------ flocks/session/runtime/step_engine.py | 297 ++++++++--- tests/session/test_prompt_tokens.py | 37 +- tests/session/test_runner_step.py | 478 +++++++++++------- .../test_session_runner_tool_only_message.py | 8 +- .../src/components/common/SessionChat.test.ts | 23 + webui/src/components/common/SessionChat.tsx | 2 +- 8 files changed, 828 insertions(+), 412 deletions(-) diff --git a/flocks/session/context_usage.py b/flocks/session/context_usage.py index 376b41044..422249a77 100644 --- a/flocks/session/context_usage.py +++ b/flocks/session/context_usage.py @@ -15,7 +15,7 @@ from flocks.provider.provider import Provider from flocks.session.message import Message -from flocks.session.prompt import SessionPrompt +from flocks.session.prompt import SessionPrompt, TurnPromptContext from flocks.session.session import SessionInfo from flocks.utils.log import Log @@ -309,7 +309,16 @@ async def _estimate_system_prompt_tokens( if agent is None: agent = await Agent.get("rex") - prompts = await SessionPrompt.build_system_prompts( + from flocks.config import Config + from flocks.project.instance import Instance + + try: + config = await Config.get() + config_instructions = tuple(config.instructions or ()) + except Exception: + config_instructions = () + + prompt_blocks = await SessionPrompt.build_system_prompt_blocks( session_id=session_id, session_directory=getattr(session, "directory", None) if session is not None else None, agent_name=getattr(agent, "name", agent_name) if agent is not None else agent_name, @@ -317,9 +326,16 @@ async def _estimate_system_prompt_tokens( provider_id=provider_id, model_id=model_id, prompt_tool_names=prompt_tool_names, - tool_revision=ToolRegistry.revision(), + turn_context=TurnPromptContext( + worktree=Instance.get_worktree(), + config_instructions=config_instructions, + tool_revision=ToolRegistry.revision(), + ), + ) + return sum( + SessionPrompt.count_tokens(block.content) + for block in prompt_blocks ) - return sum(SessionPrompt.count_tokens(prompt) for prompt in prompts) except Exception as exc: log.debug("context_usage.system_prompt_estimate_failed", { "session_id": session_id, diff --git a/flocks/session/prompt.py b/flocks/session/prompt.py index 12e17ef64..2d9f44592 100644 --- a/flocks/session/prompt.py +++ b/flocks/session/prompt.py @@ -30,8 +30,7 @@ # Output token maximum OUTPUT_TOKEN_MAX = int(os.getenv("FLOCKS_OUTPUT_TOKEN_MAX", "32000")) SystemPromptCache = Dict[str, Any] -AsyncPromptFactory = Callable[[], Awaitable[Optional[str]]] -StringPromptFactory = Callable[[], Optional[str]] +AsyncPromptLoader = Callable[[], Awaitable[Optional[str]]] # Prompt template directory (same structure as Flocks) @@ -150,6 +149,25 @@ class SystemPromptBlock: cache_key: str +@dataclass(frozen=True) +class TurnPromptContext: + """Runtime prompt values collected once before deterministic assembly.""" + + tool_catalog: Optional[str] = None + device_asset_hint: Optional[str] = None + sandbox_context: Optional[str] = None + channel_context: Optional[str] = None + additional_context: Optional[str] = None + text_tool_catalog: Optional[str] = None + tool_results_reminder: Optional[str] = None + repeated_tool_calls_reminder: Optional[str] = None + worktree: Optional[str] = None + config_instructions: tuple[str, ...] = () + tool_revision: Optional[int] = None + device_revision: Optional[int] = None + minimal_prompt: Optional[bool] = None + + class SystemPrompt: """ System Prompt generation namespace @@ -763,49 +781,6 @@ def _layer_cache_key( """Build a layer cache key for one prompt block.""" return f"system_prompt_block:{name}:{cls._system_prompt_cache_digest(digest_inputs)}" - @classmethod - def _system_prompt_cache_key( - cls, - *, - session_id: str, - agent_name: str, - provider_id: str, - model_id: str, - block_keys: Iterable[str], - ) -> str: - """Build the cache key for the composed system prompt snapshot.""" - cache_digest = cls._system_prompt_cache_digest({ - "block_keys": tuple(block_keys), - }) - return f"system_prompts:{session_id}:{agent_name}:{provider_id}:{model_id}:{cache_digest}" - - @classmethod - def _read_system_prompt_cache( - cls, - static_cache: Optional[SystemPromptCache], - cache_key: Optional[str], - ) -> Optional[List[str]]: - """Return a defensive copy of cached prompt blocks when available.""" - if static_cache is None or cache_key is None: - return None - - cached = static_cache.get(cache_key) - if cached is None: - return None - return list(cached) - - @classmethod - def _write_system_prompt_cache( - cls, - static_cache: Optional[SystemPromptCache], - cache_key: Optional[str], - prompts: List[str], - ) -> None: - """Store a defensive copy of prompt blocks in the session cache.""" - if static_cache is None or cache_key is None: - return - static_cache[cache_key] = list(prompts) - @classmethod def _read_cached_prompt_block( cls, @@ -880,14 +855,15 @@ async def _build_cached_async_prompt_block( name: str, cache_scope: str, digest_inputs: Dict[str, Any], - builder: AsyncPromptFactory, + loader: AsyncPromptLoader, ) -> Optional[SystemPromptBlock]: """Build or reuse a cached async prompt block.""" cache_key = cls._layer_cache_key(name=name, digest_inputs=digest_inputs) content = cls._read_cached_prompt_block(static_cache, cache_key) if content is None: - content = cls._normalize_prompt_text(await builder()) - cls._write_cached_prompt_block(static_cache, cache_key, content) + content = cls._normalize_prompt_text(await loader()) + if content: + cls._write_cached_prompt_block(static_cache, cache_key, content) if not content: return None return SystemPromptBlock( @@ -986,26 +962,6 @@ def _prompt_blocks_to_list( if block is not None and block.content.strip() ] - @classmethod - async def _build_optional_async_prompt( - cls, - prompt_factory: Optional[AsyncPromptFactory], - ) -> Optional[str]: - """Run an optional async prompt factory.""" - if not prompt_factory: - return None - return await prompt_factory() - - @classmethod - def _build_optional_prompt( - cls, - prompt_factory: Optional[StringPromptFactory], - ) -> Optional[str]: - """Run an optional synchronous prompt factory.""" - if not prompt_factory: - return None - return prompt_factory() - @classmethod def _print_system_prompts_for_debug( cls, @@ -1014,7 +970,7 @@ def _print_system_prompts_for_debug( agent_name: str, provider_id: str, model_id: str, - prompts: List[str], + blocks: Iterable[SystemPromptBlock], ) -> None: """Print prompt blocks when FLOCKS_PRINT_SYSTEM_PROMPT is enabled.""" if os.getenv("FLOCKS_PRINT_SYSTEM_PROMPT", "").lower() not in ("1", "true", "yes"): @@ -1025,8 +981,12 @@ def _print_system_prompts_for_debug( f"agent={agent_name} model={provider_id}/{model_id} ===" ) print(header, file=sys.stderr) - for idx, prompt in enumerate(prompts): - print(f"\n--- prompt[{idx}] ---\n{prompt}\n", file=sys.stderr) + for idx, block in enumerate(blocks): + print( + f"\n--- prompt[{idx}] {block.name} scope={block.cache_scope} " + f"---\n{block.content}\n", + file=sys.stderr, + ) print("=== end system_prompt ===\n", file=sys.stderr) @classmethod @@ -1083,21 +1043,81 @@ async def _is_builtin_system_subagent_session( return False @classmethod - async def _build_subagent_minimal_prompts( + def _append_turn_tail_blocks( + cls, + *, + blocks: List[SystemPromptBlock], + turn_context: TurnPromptContext, + static_cache: Optional[SystemPromptCache], + session_id: str, + ) -> None: + """Append per-turn values in the exact order sent to the model.""" + tail_values = [ + ("turn_additional_context", turn_context.additional_context), + ("text_tool_catalog", turn_context.text_tool_catalog), + ("tool_results_reminder", turn_context.tool_results_reminder), + ( + "repeated_tool_calls_reminder", + turn_context.repeated_tool_calls_reminder, + ), + ] + for name, content in tail_values: + block = cls._build_cached_prompt_block( + static_cache=static_cache, + name=name, + cache_scope="runtime_tail", + digest_inputs={"session_id": session_id, "content": content or ""}, + builder=lambda value=content: cls._normalize_prompt_text(value), + ) + if block is not None: + blocks.append(block) + + @classmethod + def _build_subagent_minimal_blocks( cls, *, + session_id: str, session_directory: Optional[str], agent_prompt: Optional[str], - ) -> List[str]: - """Build minimal system prompts for built-in system subagents.""" - prompts = [ - cls._normalize_prompt_text(agent_prompt), - cls._build_minimal_environment(session_directory), - ] - return [prompt for prompt in prompts if prompt] + turn_context: TurnPromptContext, + static_cache: Optional[SystemPromptCache], + ) -> List[SystemPromptBlock]: + """Build minimal prompt blocks for built-in system subagents.""" + blocks: List[SystemPromptBlock] = [] + agent_block = cls._build_cached_prompt_block( + static_cache=static_cache, + name="agent_identity", + cache_scope="agent", + digest_inputs={"agent_prompt": agent_prompt or ""}, + builder=lambda: cls._normalize_prompt_text(agent_prompt), + ) + if agent_block is not None: + blocks.append(agent_block) + + environment_block = cls._build_cached_prompt_block( + static_cache=static_cache, + name="minimal_environment", + cache_scope="runtime_tail", + digest_inputs={ + "directory": session_directory, + "runtime_day": datetime.now().strftime("%Y-%m-%d"), + "platform": platform.system().lower(), + }, + builder=lambda: cls._build_minimal_environment(session_directory), + ) + if environment_block is not None: + blocks.append(environment_block) + + cls._append_turn_tail_blocks( + blocks=blocks, + turn_context=turn_context, + static_cache=static_cache, + session_id=session_id, + ) + return blocks @classmethod - async def build_system_prompts( + async def build_system_prompt_blocks( cls, *, session_id: str, @@ -1108,45 +1128,51 @@ async def build_system_prompts( model_id: str, execution_mode_prompt: Optional[str] = None, prompt_tool_names: Iterable[str] = (), - tool_revision: Optional[int] = None, memory_bootstrap_data: Optional[Dict[str, Any]] = None, static_cache: Optional[SystemPromptCache] = None, - sandbox_prompt_factory: Optional[AsyncPromptFactory] = None, - channel_context_prompt_factory: Optional[AsyncPromptFactory] = None, - tool_catalog_prompt_factory: Optional[StringPromptFactory] = None, - device_asset_prompt_factory: Optional[AsyncPromptFactory] = None, - device_revision: Optional[int] = None, + turn_context: Optional[TurnPromptContext] = None, use_text_tool_call_mode: bool = False, - ) -> List[str]: - """Build the runtime system prompt blocks for a session turn. + ) -> List[SystemPromptBlock]: + """Build the ordered system prompt blocks for a session turn. Stable identity and execution guidance come first, followed by session/workspace context, with runtime-only metadata kept at the - prompt tail. Cache mechanics are intentionally kept out of the block - construction below so this method reads as an ordered list of prompt - layers. + prompt tail. Runtime I/O is collected before this method so assembly is + deterministic and every downstream consumer sees the same blocks. """ + turn_context = turn_context or TurnPromptContext() vcs = "git" if session_directory else None - if await cls._is_builtin_system_subagent_session( - session_id=session_id, - agent_name=agent_name, - ): - prompts = await cls._build_subagent_minimal_prompts( + minimal_prompt = turn_context.minimal_prompt + if minimal_prompt is None: + minimal_prompt = await cls._is_builtin_system_subagent_session( + session_id=session_id, + agent_name=agent_name, + ) + if minimal_prompt: + blocks = cls._build_subagent_minimal_blocks( + session_id=session_id, session_directory=session_directory, agent_prompt=agent_prompt, + turn_context=turn_context, + static_cache=static_cache, ) cls._print_system_prompts_for_debug( session_id=session_id, agent_name=agent_name, provider_id=provider_id, model_id=model_id, - prompts=prompts, + blocks=blocks, ) - return prompts + return blocks normalized_tool_names = tuple(sorted(prompt_tool_names)) runtime_day = datetime.now().strftime("%Y-%m-%d") - custom_signature = SystemPrompt.custom_signature(directory=session_directory) + config_instructions = list(turn_context.config_instructions) + custom_signature = SystemPrompt.custom_signature( + directory=session_directory, + worktree=turn_context.worktree, + config_instructions=config_instructions, + ) memory_guidance = cls._build_memory_guidance_prompt( normalized_tool_names, memory_bootstrap_data, @@ -1158,7 +1184,11 @@ async def build_system_prompts( async def build_custom_context() -> Optional[str]: return cls._join_prompt_parts( - await SystemPrompt.custom(directory=session_directory), + await SystemPrompt.custom( + directory=session_directory, + worktree=turn_context.worktree, + config_instructions=config_instructions, + ), ) blocks: List[Optional[SystemPromptBlock]] = [ @@ -1215,23 +1245,32 @@ async def build_custom_context() -> Optional[str]: cache_scope="catalog", digest_inputs={ "agent_name": agent_name, - "tool_revision": tool_revision, + "tool_revision": turn_context.tool_revision, + "content": turn_context.tool_catalog or "", }, - builder=lambda: cls._build_optional_prompt(tool_catalog_prompt_factory) or "", + builder=lambda: cls._normalize_prompt_text( + turn_context.tool_catalog, + ), ), ] - if device_asset_prompt_factory: - blocks.append(await cls._build_cached_async_prompt_block( - static_cache=static_cache, - name="device_asset_hint", - cache_scope="runtime", - digest_inputs={ - "session_id": session_id, - "device_revision": device_revision, - }, - builder=device_asset_prompt_factory, - )) + if turn_context.device_asset_hint: + blocks.append( + cls._build_cached_prompt_block( + static_cache=static_cache, + name="device_asset_hint", + cache_scope="runtime", + digest_inputs={ + "session_id": session_id, + "device_revision": turn_context.device_revision, + "tool_revision": turn_context.tool_revision, + "content": turn_context.device_asset_hint, + }, + builder=lambda: cls._normalize_prompt_text( + turn_context.device_asset_hint, + ), + ) + ) blocks.append( cls._build_cached_prompt_block( @@ -1253,33 +1292,52 @@ async def build_custom_context() -> Optional[str]: static_cache=static_cache, name="context_files", cache_scope="workspace", - digest_inputs={"directory": session_directory, "signature": custom_signature}, - builder=build_custom_context, + digest_inputs={ + "directory": session_directory, + "worktree": turn_context.worktree, + "signature": custom_signature, + }, + loader=build_custom_context, ) blocks.append(custom_block) - if sandbox_prompt_factory: - blocks.append(await cls._build_cached_async_prompt_block( - static_cache=static_cache, - name="sandbox_context", - cache_scope="runtime", - digest_inputs={"session_id": session_id, "agent_name": agent_name}, - builder=sandbox_prompt_factory, - )) + if turn_context.sandbox_context: + blocks.append( + cls._build_cached_prompt_block( + static_cache=static_cache, + name="sandbox_context", + cache_scope="runtime_tail", + digest_inputs={ + "session_id": session_id, + "agent_name": agent_name, + "content": turn_context.sandbox_context, + }, + builder=lambda: cls._normalize_prompt_text( + turn_context.sandbox_context, + ), + ) + ) - if channel_context_prompt_factory: - blocks.append(await cls._build_cached_async_prompt_block( - static_cache=static_cache, - name="channel_context", - cache_scope="runtime", - digest_inputs={"session_id": session_id}, - builder=channel_context_prompt_factory, - )) + if turn_context.channel_context: + blocks.append( + cls._build_cached_prompt_block( + static_cache=static_cache, + name="channel_context", + cache_scope="runtime_tail", + digest_inputs={ + "session_id": session_id, + "content": turn_context.channel_context, + }, + builder=lambda: cls._normalize_prompt_text( + turn_context.channel_context, + ), + ) + ) blocks.append(cls._build_cached_prompt_block( static_cache=static_cache, name="runtime_metadata", - cache_scope="runtime", + cache_scope="runtime_tail", digest_inputs={ "session_id": session_id, "directory": session_directory, @@ -1298,28 +1356,55 @@ async def build_custom_context() -> Optional[str]: ), )) - cache_key = cls._system_prompt_cache_key( + resolved_blocks = [block for block in blocks if block is not None] + cls._append_turn_tail_blocks( + blocks=resolved_blocks, + turn_context=turn_context, + static_cache=static_cache, + session_id=session_id, + ) + cls._print_system_prompts_for_debug( session_id=session_id, agent_name=agent_name, provider_id=provider_id, model_id=model_id, - block_keys=[block.cache_key for block in blocks if block is not None], + blocks=resolved_blocks, ) - cached_prompts = cls._read_system_prompt_cache(static_cache, cache_key) - if cached_prompts is not None: - return cached_prompts + return resolved_blocks - prompts = cls._prompt_blocks_to_list(blocks) - cls._print_system_prompts_for_debug( + @classmethod + async def build_system_prompts( + cls, + *, + session_id: str, + session_directory: Optional[str], + agent_name: str, + agent_prompt: Optional[str], + provider_id: str, + model_id: str, + execution_mode_prompt: Optional[str] = None, + prompt_tool_names: Iterable[str] = (), + memory_bootstrap_data: Optional[Dict[str, Any]] = None, + static_cache: Optional[SystemPromptCache] = None, + turn_context: Optional[TurnPromptContext] = None, + use_text_tool_call_mode: bool = False, + ) -> List[str]: + """Compatibility API returning only the assembled prompt text.""" + blocks = await cls.build_system_prompt_blocks( session_id=session_id, + session_directory=session_directory, agent_name=agent_name, + agent_prompt=agent_prompt, provider_id=provider_id, model_id=model_id, - prompts=prompts, + execution_mode_prompt=execution_mode_prompt, + prompt_tool_names=prompt_tool_names, + memory_bootstrap_data=memory_bootstrap_data, + static_cache=static_cache, + turn_context=turn_context, + use_text_tool_call_mode=use_text_tool_call_mode, ) - - cls._write_system_prompt_cache(static_cache, cache_key, prompts) - return list(prompts) + return cls._prompt_blocks_to_list(blocks) @classmethod def _build_context_section(cls, context: ContextInfo) -> str: diff --git a/flocks/session/runtime/step_engine.py b/flocks/session/runtime/step_engine.py index c2b359421..b90040c1f 100644 --- a/flocks/session/runtime/step_engine.py +++ b/flocks/session/runtime/step_engine.py @@ -2,6 +2,7 @@ import asyncio import copy +import hashlib import json import os import re @@ -36,7 +37,7 @@ from flocks.utils.id import Identifier from flocks.session.session import Session, SessionInfo from flocks.session.message import Message, MessageInfo, MessageRole, TextPart -from flocks.session.prompt import SessionPrompt +from flocks.session.prompt import SessionPrompt, SystemPromptBlock, TurnPromptContext from flocks.session.core.status import SessionStatus, SessionStatusRetry, SessionStatusBusy from flocks.session.core.defaults import ( DEFAULT_MAX_TOOL_STEPS, @@ -275,6 +276,8 @@ def __init__( ] = {} self._turn: Optional[Any] = None self._model_policy: ModelRoutingPolicy = DEFAULT_MODEL_ROUTING_POLICY + self._step_agent: Optional[AgentInfo] = None + self._frozen_tool_request: Optional[ModelRequest[ChatMessage]] = None @classmethod def from_turn( @@ -1293,9 +1296,10 @@ async def _process_step( ) # Resolve agent agent_name = last_user.agent or self.agent_name - agent = getattr(self, "_step_agent", None) + agent = self._step_agent if agent is None: agent = await Agent.get(agent_name) or await Agent.get("rex") + assert agent is not None, "runtime agent invariant violated" if self._turn is not None: self._step_agent = agent @@ -1381,7 +1385,7 @@ async def _process_step( # Build prompts and tools tools_started_at = time.perf_counter() - frozen_tool_request = getattr(self, "_frozen_tool_request", None) + frozen_tool_request = self._frozen_tool_request if isinstance(frozen_tool_request, ModelRequest): tools = frozen_tool_request.provider_tools() else: @@ -1389,24 +1393,19 @@ async def _process_step( self._log_perf("runner.process_step.tools_ready", tools_started_at, tool_count=len(tools)) prompt_tool_names = self._get_prompt_tool_names_from_schema(tools) - async def sandbox_prompt_factory() -> Optional[str]: - return await self._build_sandbox_prompt(agent) - - async def channel_context_prompt_factory() -> Optional[str]: - return await self._build_channel_context_prompt() - - async def device_asset_prompt_factory() -> Optional[str]: - return await self._build_device_asset_hint() - - try: - from flocks.tool.device.store import device_revision as get_device_revision - - current_device_revision = get_device_revision() - except Exception: - current_device_revision = None - prompts_started_at = time.perf_counter() - system_prompts = await SessionPrompt.build_system_prompts( + minimal_prompt = await SessionPrompt._is_builtin_system_subagent_session( + session_id=self.session.id, + agent_name=agent.name, + ) + turn_prompt_context = await self._build_turn_prompt_context( + agent=agent, + messages=messages, + last_user=last_user, + tools=tools, + minimal_prompt=minimal_prompt, + ) + system_prompt_blocks = await SessionPrompt.build_system_prompt_blocks( session_id=self.session.id, session_directory=self.session.directory, agent_name=agent.name, @@ -1419,57 +1418,19 @@ async def device_asset_prompt_factory() -> Optional[str]: plan_file=self._turn_plan_file, ), prompt_tool_names=prompt_tool_names, - tool_revision=ToolRegistry.revision(), memory_bootstrap_data=self._memory_bootstrap_data, static_cache=self._static_cache, - sandbox_prompt_factory=sandbox_prompt_factory, - channel_context_prompt_factory=channel_context_prompt_factory, - tool_catalog_prompt_factory=lambda: self._build_tool_catalog_prompt(agent), - device_asset_prompt_factory=device_asset_prompt_factory, - device_revision=current_device_revision, + turn_context=turn_prompt_context, use_text_tool_call_mode=self._should_use_text_tool_call_mode(), ) - self._log_perf("runner.process_step.system_prompts_ready", prompts_started_at, prompt_count=len(system_prompts)) + self._log_perf( + "runner.process_step.system_prompts_ready", + prompts_started_at, + prompt_count=len(system_prompt_blocks), + ) await self._run_session_start_hook(agent) - if self._turn_additional_context: - system_prompts.append(self._turn_additional_context) - - if self._should_use_text_tool_call_mode() and tools: - text_tool_catalog = self._build_text_tool_call_catalog_prompt(tools) - if text_tool_catalog: - system_prompts.append(text_tool_catalog) - - # If the last assistant message only contains tool results and no text, - # force a direct answer to avoid repeated tool calls. - last_assistant_msg = None - for msg in reversed(messages): - if msg.role == MessageRole.ASSISTANT: - last_assistant_msg = msg - break - if last_assistant_msg: - parts = await Message.parts(last_assistant_msg.id, self.session.id) - has_text = any(getattr(p, "type", None) == "text" and getattr(p, "text", "").strip() for p in parts) - has_tool_result = any( - getattr(p, "type", None) == "tool" and - getattr(getattr(p, "state", None), "status", None) in ("completed", "error", "running") - for p in parts - ) - if has_tool_result and not has_text: - from flocks.session.prompt_strings import PROMPT_TOOL_RESULTS_AVAILABLE - system_prompts.append(PROMPT_TOOL_RESULTS_AVAILABLE) - - if has_tool_result and self._should_warn_about_tool_loop(last_user_id=last_user.id): - state = self._get_tool_loop_guard_state(last_user_id=last_user.id) - log.warn("runner.repeated_tool_calls_detected", { - "tool_name": state.get("last_signature", "").split(":", 1)[0], - "exact_count": state.get("exact_count", 0), - "step": self._step, - }) - from flocks.session.prompt_strings import PROMPT_REPEATED_TOOL_CALLS - system_prompts.append(PROMPT_REPEATED_TOOL_CALLS) - # Convert messages to chat format with error handling try: queued_user_message_ids = self._get_queued_user_message_ids(messages) @@ -1479,7 +1440,10 @@ async def device_asset_prompt_factory() -> Optional[str]: self._queued_user_message_ids = queued_user_message_ids chat_messages_started_at = time.perf_counter() try: - chat_messages = await self._to_chat_messages(messages, system_prompts) + chat_messages = await self._to_chat_messages( + messages, + system_prompt_blocks, + ) finally: if previous_queued_user_ids is None: if hasattr(self, "_queued_user_message_ids"): @@ -1985,6 +1949,162 @@ async def _record_usage_if_available( "error": str(exc), }) + async def _build_turn_prompt_context( + self, + *, + agent: AgentInfo, + messages: List[MessageInfo], + last_user: MessageInfo, + tools: List[Dict[str, Any]], + minimal_prompt: bool = False, + ) -> TurnPromptContext: + """Collect cached runtime values before deterministic prompt assembly.""" + if minimal_prompt: + return await self._add_turn_prompt_tail( + TurnPromptContext(minimal_prompt=True), + messages=messages, + last_user=last_user, + tools=tools, + ) + + from flocks.config import Config + from flocks.project.instance import Instance + + try: + from flocks.tool.device.store import device_revision + + current_device_revision = device_revision() + except Exception: + current_device_revision = None + + current_tool_revision = ToolRegistry.revision() + try: + config = await Config.get() + config_data = config.model_dump(by_alias=True, exclude_none=True) + config_instructions = tuple(config.instructions or ()) + except Exception as exc: + log.debug("runner.prompt_context.config_error", {"error": str(exc)}) + config_data = None + config_instructions = () + + worktree = Instance.get_worktree() + config_fingerprint = hashlib.sha256( + json.dumps( + config_data, + ensure_ascii=False, + sort_keys=True, + default=str, + ).encode("utf-8"), + ).hexdigest() + source_key = ( + self.session.id, + agent.name, + self.session.directory, + worktree, + current_tool_revision, + current_device_revision, + config_fingerprint, + ) + cached = self._static_cache.get("runtime_prompt_context") + source_context = None + if isinstance(cached, dict) and cached.get("key") == source_key: + candidate = cached.get("context") + if isinstance(candidate, TurnPromptContext): + source_context = candidate + + if source_context is None: + sandbox_context, channel_context, device_asset_hint = await asyncio.gather( + self._build_sandbox_prompt(agent, config_data=config_data), + self._build_channel_context_prompt(), + self._build_device_asset_hint(), + ) + source_context = TurnPromptContext( + tool_catalog=self._build_tool_catalog_prompt(agent), + device_asset_hint=device_asset_hint, + sandbox_context=sandbox_context, + channel_context=channel_context, + worktree=worktree, + config_instructions=config_instructions, + tool_revision=current_tool_revision, + device_revision=current_device_revision, + minimal_prompt=False, + ) + self._static_cache["runtime_prompt_context"] = { + "key": source_key, + "context": source_context, + } + + return await self._add_turn_prompt_tail( + source_context, + messages=messages, + last_user=last_user, + tools=tools, + ) + + async def _add_turn_prompt_tail( + self, + source_context: TurnPromptContext, + *, + messages: List[MessageInfo], + last_user: MessageInfo, + tools: List[Dict[str, Any]], + ) -> TurnPromptContext: + """Add uncached per-step context and reminders to a source snapshot.""" + text_tool_catalog = None + if self._should_use_text_tool_call_mode() and tools: + text_tool_catalog = self._build_text_tool_call_catalog_prompt(tools) + + tool_results_reminder = None + repeated_tool_calls_reminder = None + last_assistant_msg = next( + ( + message + for message in reversed(messages) + if message.role == MessageRole.ASSISTANT + ), + None, + ) + if last_assistant_msg is not None: + parts = await Message.parts(last_assistant_msg.id, self.session.id) + has_text = any( + getattr(part, "type", None) == "text" + and getattr(part, "text", "").strip() + for part in parts + ) + has_tool_result = any( + getattr(part, "type", None) == "tool" + and getattr(getattr(part, "state", None), "status", None) + in ("completed", "error", "running") + for part in parts + ) + if has_tool_result and not has_text: + from flocks.session.prompt_strings import ( + PROMPT_TOOL_RESULTS_AVAILABLE, + ) + + tool_results_reminder = PROMPT_TOOL_RESULTS_AVAILABLE + + if has_tool_result and self._should_warn_about_tool_loop( + last_user_id=last_user.id, + ): + state = self._get_tool_loop_guard_state(last_user_id=last_user.id) + log.warn("runner.repeated_tool_calls_detected", { + "tool_name": state.get("last_signature", "").split(":", 1)[0], + "exact_count": state.get("exact_count", 0), + "step": self._step, + }) + from flocks.session.prompt_strings import PROMPT_REPEATED_TOOL_CALLS + + repeated_tool_calls_reminder = PROMPT_REPEATED_TOOL_CALLS + + return replace( + source_context, + additional_context=self._turn_additional_context, + text_tool_catalog=text_tool_catalog, + tool_results_reminder=tool_results_reminder, + repeated_tool_calls_reminder=repeated_tool_calls_reminder, + ) + async def _build_device_asset_hint(self) -> Optional[str]: """Return concise device-aware tool guidance plus enabled device summary.""" try: @@ -2027,15 +2147,22 @@ async def _build_device_asset_hint(self) -> Optional[str]: "如果同类设备有多个候选，不要猜测，先询问用户选择。" ) - async def _build_sandbox_prompt(self, agent: AgentInfo) -> Optional[str]: + async def _build_sandbox_prompt( + self, + agent: AgentInfo, + *, + config_data: Optional[Dict[str, Any]] = None, + ) -> Optional[str]: """Build sandbox context prompt when sandboxing is active.""" try: - from flocks.config import Config from flocks.session.core.session_state import get_main_session_id from flocks.sandbox.system_prompt import build_sandbox_system_prompt - cfg = await Config.get() - config_data = cfg.model_dump(by_alias=True, exclude_none=True) + if config_data is None: + from flocks.config import Config + + config = await Config.get() + config_data = config.model_dump(by_alias=True, exclude_none=True) session_key = self.session.id main_session_key = get_main_session_id() or self.session.id return await build_sandbox_system_prompt( @@ -2550,14 +2677,26 @@ def _build_tool_output_text(self, part: Any, tool_name: str, ctx_window_tokens: def _build_system_message_content( self, - system_prompts: List[str], + system_prompts: List[SystemPromptBlock] | List[str], ) -> str | list[dict[str, Any]]: """Format system prompts for the active provider. Anthropic supports structured system blocks, which lets us place a conservative cache breakpoint before the dynamic runtime tail. """ - prompt_parts = [prompt for prompt in system_prompts if prompt and prompt.strip()] + typed_blocks = [ + block + for block in system_prompts + if isinstance(block, SystemPromptBlock) and block.content.strip() + ] + if typed_blocks: + prompt_parts = [block.content for block in typed_blocks] + else: + prompt_parts = [ + prompt + for prompt in system_prompts + if isinstance(prompt, str) and prompt.strip() + ] if not prompt_parts: return "" @@ -2565,7 +2704,19 @@ def _build_system_message_content( if "anthropic" not in provider_lower: return "\n\n".join(prompt_parts) - cache_break_index = max(0, len(prompt_parts) - 3) + if typed_blocks: + first_runtime_tail = next( + ( + index + for index, block in enumerate(typed_blocks) + if block.cache_scope == "runtime_tail" + ), + len(typed_blocks), + ) + cache_break_index = max(0, first_runtime_tail - 1) + else: + # Compatibility for callers still passing plain strings. + cache_break_index = max(0, len(prompt_parts) - 3) blocks: list[dict[str, Any]] = [] for index, prompt in enumerate(prompt_parts): block: dict[str, Any] = { @@ -2580,7 +2731,7 @@ def _build_system_message_content( async def _to_chat_messages( self, messages: List[MessageInfo], - system_prompts: List[str], + system_prompts: List[SystemPromptBlock] | List[str], ) -> List[ChatMessage]: """ Convert messages to chat format with tool calls. diff --git a/tests/session/test_prompt_tokens.py b/tests/session/test_prompt_tokens.py index cf130cc66..a14cfc317 100644 --- a/tests/session/test_prompt_tokens.py +++ b/tests/session/test_prompt_tokens.py @@ -25,6 +25,7 @@ PromptTemplate, SessionPrompt, SystemPrompt, + TurnPromptContext, ) from flocks.session import prompt_strings @@ -268,7 +269,9 @@ async def test_builtin_system_subagent_child_uses_minimal_prompt(self): agent_prompt="You are Rex Junior.", provider_id="anthropic", model_id="claude-sonnet", - tool_catalog_prompt_factory=lambda: "SHOULD_NOT_APPEAR", + turn_context=TurnPromptContext( + tool_catalog="SHOULD_NOT_APPEAR", + ), ) assert len(prompts) == 2 @@ -307,6 +310,38 @@ async def test_builtin_system_subagent_root_uses_full_prompt(self): assert len(prompts) > 2 assert any(PROMPT_DEFAULT.strip() in prompt for prompt in prompts) + @pytest.mark.asyncio + async def test_full_prompt_loads_worktree_and_config_instructions( + self, + tmp_path: Path, + ) -> None: + nested = tmp_path / "src" / "package" + nested.mkdir(parents=True) + (tmp_path / "AGENTS.md").write_text("project rules", encoding="utf-8") + (nested / "extra-rules.md").write_text("extra rules", encoding="utf-8") + + with patch.object( + SessionPrompt, + "_is_builtin_system_subagent_session", + AsyncMock(return_value=False), + ): + prompts = await SessionPrompt.build_system_prompts( + session_id="ses-instructions", + session_directory=str(nested), + agent_name="rex", + agent_prompt="agent prompt", + provider_id="openai", + model_id="gpt-5", + turn_context=TurnPromptContext( + worktree=str(tmp_path), + config_instructions=("extra-rules.md",), + ), + ) + + combined = "\n\n".join(prompts) + assert "project rules" in combined + assert "extra rules" in combined + # --------------------------------------------------------------------------- # SystemPrompt.provider() — returns List[str] diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index 25091d6c2..819975c73 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -33,7 +33,7 @@ ToolCall, ) from flocks.session.runtime.session_turn import LoopCallbacks, LoopContext -from flocks.session.prompt import SessionPrompt +from flocks.session.prompt import SessionPrompt, SystemPromptBlock, TurnPromptContext from flocks.session.core.defaults import DEFAULT_MAX_TOOL_STEPS from flocks.session.session import Session, SessionInfo from flocks.tool.registry import ToolCategory, ToolInfo @@ -641,14 +641,21 @@ async def test_build_system_prompts_reuses_loop_static_cache(self): env_mock = MagicMock(return_value=["env prompt"]) runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_mock = AsyncMock(return_value="device prompt") + turn_context = TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog", + device_asset_hint="device prompt", + tool_revision=1, + device_revision=7, + ) - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -657,13 +664,8 @@ async def test_build_system_prompts_reuses_loop_static_cache(self): provider_id=runner1.provider_id, model_id=runner1.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, - device_revision=7, + turn_context=turn_context, ) prompts2 = await SessionPrompt.build_system_prompts( session_id=session.id, @@ -673,22 +675,14 @@ async def test_build_system_prompts_reuses_loop_static_cache(self): provider_id=runner2.provider_id, model_id=runner2.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, - device_revision=7, + turn_context=turn_context, ) assert prompts1 == prompts2 env_mock.assert_called_once() runtime_mock.assert_called_once() custom_mock.assert_awaited_once() - sandbox_mock.assert_awaited_once() - channel_mock.assert_awaited_once() - device_mock.assert_awaited_once() @pytest.mark.asyncio async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(self): @@ -704,15 +698,13 @@ async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(sel "inject": True, }, } - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_mock = AsyncMock(return_value="device prompt") - - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch.object(SessionPrompt, "_build_tool_guidance_prompt", return_value="tool protocol"), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", return_value=["env prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", return_value=["runtime prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.custom", AsyncMock(return_value=["custom prompt"])): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch.object(SessionPrompt, "_build_tool_guidance_prompt", return_value="tool protocol"), + patch("flocks.session.prompt.SystemPrompt.environment_stable", return_value=["env prompt"]), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", return_value=["runtime prompt"]), + patch("flocks.session.prompt.SystemPrompt.custom", AsyncMock(return_value=["custom prompt"])), + ): prompts = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -730,11 +722,17 @@ async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(sel "write", ), memory_bootstrap_data=memory_bootstrap_data, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, - device_revision=3, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, + turn_context=TurnPromptContext( + tool_catalog="tool catalog", + device_asset_hint="device prompt", + sandbox_context="sandbox prompt", + channel_context="channel prompt", + additional_context="additional prompt", + text_tool_catalog="text tool catalog", + tool_results_reminder="tool results reminder", + repeated_tool_calls_reminder="tool loop reminder", + device_revision=3, + ), ) assert prompts == [ @@ -750,6 +748,10 @@ async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(sel "sandbox prompt", "channel prompt", "runtime prompt", + "additional prompt", + "text tool catalog", + "tool results reminder", + "tool loop reminder", ] @pytest.mark.asyncio @@ -763,16 +765,12 @@ async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): env_mock = MagicMock(return_value=["env prompt"]) runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_mock = AsyncMock(return_value="device prompt") - - catalog_prompts = iter(["tool catalog v1", "tool catalog v2"]) - - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -781,13 +779,15 @@ async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: next(catalog_prompts), - device_asset_prompt_factory=device_mock, - device_revision=1, + turn_context=TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog v1", + device_asset_hint="device prompt", + tool_revision=1, + device_revision=1, + ), ) agent.prompt = "agent prompt v2" prompts2 = await SessionPrompt.build_system_prompts( @@ -798,13 +798,15 @@ async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=2, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: next(catalog_prompts), - device_asset_prompt_factory=device_mock, - device_revision=1, + turn_context=TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog v2", + device_asset_hint="device prompt", + tool_revision=2, + device_revision=1, + ), ) assert prompts1 != prompts2 @@ -815,8 +817,6 @@ async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): env_mock.assert_called_once() runtime_mock.assert_called_once() custom_mock.assert_awaited_once() - sandbox_mock.assert_awaited_once() - channel_mock.assert_awaited_once() @pytest.mark.asyncio async def test_build_system_prompts_reuses_static_device_hint_cache(self): @@ -829,14 +829,20 @@ async def test_build_system_prompts_reuses_static_device_hint_cache(self): env_mock = MagicMock(return_value=["env prompt"]) runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_mock = AsyncMock(return_value="device prompt") + turn_context = TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog", + device_asset_hint="device prompt", + tool_revision=1, + ) - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -845,12 +851,8 @@ async def test_build_system_prompts_reuses_static_device_hint_cache(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, + turn_context=turn_context, ) prompts2 = await SessionPrompt.build_system_prompts( session_id=session.id, @@ -860,12 +862,8 @@ async def test_build_system_prompts_reuses_static_device_hint_cache(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, + turn_context=turn_context, ) assert prompts1 == prompts2 @@ -873,9 +871,6 @@ async def test_build_system_prompts_reuses_static_device_hint_cache(self): env_mock.assert_called_once() runtime_mock.assert_called_once() custom_mock.assert_awaited_once() - sandbox_mock.assert_awaited_once() - channel_mock.assert_awaited_once() - device_mock.assert_awaited_once() @pytest.mark.asyncio async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): @@ -888,14 +883,12 @@ async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): env_mock = MagicMock(return_value=["env prompt"]) runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_prompts = iter(["device prompt v1", "device prompt v2"]) - - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -904,13 +897,15 @@ async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=AsyncMock(side_effect=lambda: next(device_prompts)), - device_revision=1, + turn_context=TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog", + device_asset_hint="device prompt v1", + tool_revision=1, + device_revision=1, + ), ) prompts2 = await SessionPrompt.build_system_prompts( session_id=session.id, @@ -920,13 +915,15 @@ async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=AsyncMock(side_effect=lambda: next(device_prompts)), - device_revision=2, + turn_context=TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog", + device_asset_hint="device prompt v2", + tool_revision=1, + device_revision=2, + ), ) assert prompts1 != prompts2 @@ -935,8 +932,6 @@ async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): env_mock.assert_called_once() runtime_mock.assert_called_once() custom_mock.assert_awaited_once() - sandbox_mock.assert_awaited_once() - channel_mock.assert_awaited_once() @pytest.mark.asyncio async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): @@ -950,10 +945,12 @@ async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -962,8 +959,8 @@ async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, + turn_context=TurnPromptContext(tool_revision=1), ) agent.prompt = "agent prompt v2" prompts2 = await SessionPrompt.build_system_prompts( @@ -974,8 +971,8 @@ async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, + turn_context=TurnPromptContext(tool_revision=1), ) assert prompts1 != prompts2 @@ -1108,10 +1105,12 @@ async def test_filesystem_memory_guidance_depends_on_tool_names(self): runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts_with_memory = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -1127,9 +1126,9 @@ async def test_filesystem_memory_guidance_depends_on_tool_names(self): "read", "write", ), - tool_revision=1, memory_bootstrap_data=runner._memory_bootstrap_data, static_cache=shared_cache, + turn_context=TurnPromptContext(tool_revision=1), ) prompts_without_memory = await SessionPrompt.build_system_prompts( session_id=session.id, @@ -1139,9 +1138,9 @@ async def test_filesystem_memory_guidance_depends_on_tool_names(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, memory_bootstrap_data=runner._memory_bootstrap_data, static_cache=shared_cache, + turn_context=TurnPromptContext(tool_revision=1), ) assert prompts_with_memory != prompts_without_memory @@ -1475,15 +1474,31 @@ async def test_to_chat_messages_uses_structured_anthropic_system_blocks(monkeypa monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hello")) - chat_messages = await runner._to_chat_messages( - [message], - ["provider prompt", "agent prompt", "context prompt", "runtime prompt"], - ) + prompt_blocks = [ + SystemPromptBlock( + name=name, + content=content, + cache_scope=cache_scope, + digest_inputs={}, + cache_key=name, + ) + for name, content, cache_scope in ( + ("provider", "provider prompt", "global"), + ("agent", "agent prompt", "agent"), + ("context", "context prompt", "workspace"), + ("sandbox", "sandbox prompt", "runtime_tail"), + ("runtime", "runtime prompt", "runtime_tail"), + ("reminder", "reminder prompt", "runtime_tail"), + ) + ] + + chat_messages = await runner._to_chat_messages([message], prompt_blocks) assert chat_messages[0].role == "system" assert isinstance(chat_messages[0].content, list) - assert chat_messages[0].content[1]["cache_control"] == {"type": "ephemeral"} - assert chat_messages[0].content[-1]["text"] == "runtime prompt" + assert chat_messages[0].content[2]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in chat_messages[0].content[3] + assert chat_messages[0].content[-1]["text"] == "reminder prompt" @pytest.mark.asyncio @@ -2245,7 +2260,7 @@ async def fake_create(*args, **kwargs): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2303,7 +2318,7 @@ async def fake_to_chat_messages(_messages, _system_prompts): # noqa: ANN001 monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_to_chat_messages", fake_to_chat_messages) monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="queued")) @@ -2351,7 +2366,7 @@ async def fake_call_llm(*_args, **_kwargs): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2427,7 +2442,7 @@ async def fake_call_llm(*_args, **_kwargs): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2673,7 +2688,7 @@ async def publish_event(event_name, payload): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", staticmethod(lambda _provider_id: EmptyProvider())) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(StepEngine, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.SessionRetry, "sleep", AsyncMock(return_value=None)) @@ -2720,7 +2735,7 @@ async def test_process_step_uses_loaded_tool_schema_names_for_prompt_guidance(mo provider = MagicMock() provider.is_configured.return_value = True assistant_msg = SimpleNamespace(id="msg_assistant_prompt_guidance") - build_system_prompts = AsyncMock(return_value=[]) + build_system_prompt_blocks = AsyncMock(return_value=[]) tool_schema = [ {"type": "function", "function": {"name": "memory_search", "description": "", "parameters": {}}}, {"type": "function", "function": {"name": "bash", "description": "", "parameters": {}}}, @@ -2729,7 +2744,7 @@ async def test_process_step_uses_loaded_tool_schema_names_for_prompt_guidance(mo monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", build_system_prompts) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", build_system_prompt_blocks) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=tool_schema)) monkeypatch.setattr( runner, @@ -2749,8 +2764,8 @@ async def test_process_step_uses_loaded_tool_schema_names_for_prompt_guidance(mo result = await runner._process_step([last_user], last_user) assert result.content == "done" - build_system_prompts.assert_awaited_once() - assert build_system_prompts.await_args.kwargs["prompt_tool_names"] == ("bash", "memory_search") + build_system_prompt_blocks.assert_awaited_once() + assert build_system_prompt_blocks.await_args.kwargs["prompt_tool_names"] == ("bash", "memory_search") @pytest.mark.asyncio @@ -2778,7 +2793,7 @@ async def test_process_step_records_usage_after_success(monkeypatch): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2802,57 +2817,6 @@ async def test_process_step_records_usage_after_success(monkeypatch): update_mock.assert_any_await(runner.session.id, assistant_msg.id, finish="stop") -@pytest.mark.asyncio -async def test_process_step_passes_device_hint_factory_into_build_system_prompts(monkeypatch): - runner = _make_runner("ses_runner_device_hint_order") - runner.callbacks = LoopCallbacks(on_error=AsyncMock()) - - last_user = UserMessageInfo( - id="msg_user_device_hint_order", - sessionID=runner.session.id, - role="user", - time={"created": 1_000}, - agent="rex", - model={"providerID": "anthropic", "modelID": "claude-sonnet"}, - ) - - agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) - provider = MagicMock() - provider.is_configured.return_value = True - assistant_msg = SimpleNamespace(id="msg_assistant_device_hint_order") - build_system_prompts = AsyncMock(return_value=["provider", "tool catalog awareness", "device hint"]) - - monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) - monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) - monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", build_system_prompts) - monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) - device_hint_mock = AsyncMock(return_value="device hint") - monkeypatch.setattr(runner, "_build_device_asset_hint", device_hint_mock) - monkeypatch.setattr("flocks.tool.device.store.device_revision", lambda: 9) - monkeypatch.setattr( - runner, - "_to_chat_messages", - AsyncMock(return_value=[SimpleNamespace(role="user", content="hi")]), - ) - monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) - monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) - monkeypatch.setattr(runner_mod.Message, "create", AsyncMock(return_value=assistant_msg)) - monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) - monkeypatch.setattr( - runner, - "_call_llm", - AsyncMock(return_value=StepResult(action="stop", content="done")), - ) - - result = await runner._process_step([last_user], last_user) - - assert result.content == "done" - build_system_prompts.assert_awaited_once() - kwargs = build_system_prompts.await_args.kwargs - assert kwargs["device_revision"] == 9 - assert kwargs["device_asset_prompt_factory"] is not None - assert await kwargs["device_asset_prompt_factory"]() == "device hint" @pytest.mark.asyncio @@ -2883,7 +2847,7 @@ async def test_process_step_empty_retry_records_usage_per_attempt(monkeypatch): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2945,7 +2909,7 @@ async def test_process_step_retries_empty_transport_exception(monkeypatch): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2997,7 +2961,7 @@ async def _call_llm(*_args, **_kwargs): monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) monkeypatch.setattr( runner_mod.SessionPrompt, - "build_system_prompts", + "build_system_prompt_blocks", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -3045,7 +3009,7 @@ async def test_process_step_uses_default_max_steps_when_agent_steps_missing(monk monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=sentinel_tools)) monkeypatch.setattr( runner, @@ -3094,7 +3058,7 @@ async def test_process_step_respects_explicit_agent_steps_over_default(monkeypat monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=sentinel_tools)) monkeypatch.setattr( runner, @@ -3148,7 +3112,7 @@ async def fake_call_llm(self, provider, messages, tools, agent, assistant_msg): ))) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.Message, "create", create_mock) @@ -3274,3 +3238,141 @@ async def test_to_chat_messages_expands_workflow_node_ref_marker(monkeypatch): assert "node_id: query_fofa" in chat_messages[0].content assert "node_type: python" in chat_messages[0].content assert "只修改这个节点的代码并保留其他节点不变" in chat_messages[0].content + +@pytest.mark.asyncio +async def test_turn_prompt_context_reuses_runtime_snapshot(monkeypatch): + runner = _make_runner("ses_runtime_prompt_snapshot") + agent = _make_agent(name="rex") + last_user = UserMessageInfo( + id="msg_runtime_prompt_snapshot", + sessionID=runner.session.id, + role="user", + time={"created": 1_000}, + agent="rex", + model={"providerID": "anthropic", "modelID": "claude-sonnet"}, + ) + config = SimpleNamespace( + instructions=["rules.md"], + model_dump=lambda **_kwargs: {"sandbox": {"enabled": True}}, + ) + sandbox = AsyncMock(return_value="sandbox prompt") + channel = AsyncMock(return_value="channel prompt") + device = AsyncMock(return_value="device prompt") + catalog = MagicMock(return_value="tool catalog") + + monkeypatch.setattr("flocks.config.Config.get", AsyncMock(return_value=config)) + monkeypatch.setattr("flocks.project.instance.Instance.get_worktree", lambda: "/tmp") + monkeypatch.setattr("flocks.tool.device.store.device_revision", lambda: 3) + monkeypatch.setattr(runner_mod.ToolRegistry, "revision", lambda: 7) + monkeypatch.setattr(runner, "_build_sandbox_prompt", sandbox) + monkeypatch.setattr(runner, "_build_channel_context_prompt", channel) + monkeypatch.setattr(runner, "_build_device_asset_hint", device) + monkeypatch.setattr(runner, "_build_tool_catalog_prompt", catalog) + + first = await runner._build_turn_prompt_context( + agent=agent, + messages=[last_user], + last_user=last_user, + tools=[], + ) + second = await runner._build_turn_prompt_context( + agent=agent, + messages=[last_user], + last_user=last_user, + tools=[], + ) + + assert first == second + assert first.tool_revision == 7 + assert first.device_revision == 3 + assert first.config_instructions == ("rules.md",) + sandbox.assert_awaited_once() + channel.assert_awaited_once() + device.assert_awaited_once() + catalog.assert_called_once() + +@pytest.mark.asyncio +async def test_minimal_turn_prompt_context_skips_runtime_sources(monkeypatch): + runner = _make_runner("ses_minimal_runtime_prompt") + runner._turn_additional_context = "delegated context" + agent = _make_agent(name="rex-junior") + last_user = UserMessageInfo( + id="msg_minimal_runtime_prompt", + sessionID=runner.session.id, + role="user", + time={"created": 1_000}, + agent="rex-junior", + model={"providerID": "anthropic", "modelID": "claude-sonnet"}, + ) + sandbox = AsyncMock(return_value="sandbox prompt") + channel = AsyncMock(return_value="channel prompt") + device = AsyncMock(return_value="device prompt") + monkeypatch.setattr(runner, "_build_sandbox_prompt", sandbox) + monkeypatch.setattr(runner, "_build_channel_context_prompt", channel) + monkeypatch.setattr(runner, "_build_device_asset_hint", device) + + context = await runner._build_turn_prompt_context( + agent=agent, + messages=[last_user], + last_user=last_user, + tools=[], + minimal_prompt=True, + ) + + assert context.minimal_prompt is True + assert context.additional_context == "delegated context" + sandbox.assert_not_awaited() + channel.assert_not_awaited() + device.assert_not_awaited() + +@pytest.mark.asyncio +async def test_process_step_passes_device_hint_in_turn_prompt_context(monkeypatch): + runner = _make_runner("ses_runner_device_hint_order") + runner.callbacks = LoopCallbacks(on_error=AsyncMock()) + + last_user = UserMessageInfo( + id="msg_user_device_hint_order", + sessionID=runner.session.id, + role="user", + time={"created": 1_000}, + agent="rex", + model={"providerID": "anthropic", "modelID": "claude-sonnet"}, + ) + + agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) + provider = MagicMock() + provider.is_configured.return_value = True + assistant_msg = SimpleNamespace(id="msg_assistant_device_hint_order") + build_system_prompts = AsyncMock(return_value=["provider", "tool catalog awareness", "device hint"]) + + monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) + monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) + monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", build_system_prompts) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + device_hint_mock = AsyncMock(return_value="device hint") + monkeypatch.setattr(runner, "_build_device_asset_hint", device_hint_mock) + monkeypatch.setattr("flocks.tool.device.store.device_revision", lambda: 9) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hi")]), + ) + monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) + monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.Message, "create", AsyncMock(return_value=assistant_msg)) + monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) + monkeypatch.setattr( + runner, + "_call_llm", + AsyncMock(return_value=StepResult(action="stop", content="done")), + ) + + result = await runner._process_step([last_user], last_user) + + assert result.content == "done" + build_system_prompts.assert_awaited_once() + kwargs = build_system_prompts.await_args.kwargs + turn_context = kwargs["turn_context"] + assert turn_context.device_revision == 9 + assert turn_context.device_asset_hint == "device hint" diff --git a/tests/session/test_session_runner_tool_only_message.py b/tests/session/test_session_runner_tool_only_message.py index f8624de42..a9d898973 100644 --- a/tests/session/test_session_runner_tool_only_message.py +++ b/tests/session/test_session_runner_tool_only_message.py @@ -81,7 +81,7 @@ async def fake_get_prompt_tool_names(self, agent): # noqa: ANN001 del self, agent return () - async def fake_build_system_prompts(*args, **kwargs): # noqa: ANN002, ANN003 + async def fake_build_system_prompt_blocks(*args, **kwargs): # noqa: ANN002, ANN003 del args, kwargs return [] @@ -100,7 +100,11 @@ async def fake_call_llm(self, provider, messages, tools, agent, assistant_msg): monkeypatch.setattr(Provider, "apply_config", fake_apply_config) monkeypatch.setattr(Agent, "get", fake_agent_get) monkeypatch.setattr(StepEngine, "_get_prompt_tool_names", fake_get_prompt_tool_names) - monkeypatch.setattr(SessionPrompt, "build_system_prompts", fake_build_system_prompts) + monkeypatch.setattr( + SessionPrompt, + "build_system_prompt_blocks", + fake_build_system_prompt_blocks, + ) monkeypatch.setattr(StepEngine, "_build_callable_tool_schema", fake_build_callable_tool_schema) monkeypatch.setattr(StepEngine, "_to_chat_messages", fake_to_chat_messages) monkeypatch.setattr(StepEngine, "_call_llm", fake_call_llm) diff --git a/webui/src/components/common/SessionChat.test.ts b/webui/src/components/common/SessionChat.test.ts index 698138064..29978f308 100644 --- a/webui/src/components/common/SessionChat.test.ts +++ b/webui/src/components/common/SessionChat.test.ts @@ -3818,6 +3818,29 @@ describe('ChatToolPart bash rendering', () => { expect(screen.getByText('$').closest('pre')).toHaveClass('max-h-64'); expect(screen.getByText('tests passed').closest('pre')).toHaveClass('max-h-64'); }); + + it.each([ + 'Tool execution was interrupted', + 'Command failed with exit code 1', + ])('renders the bash error once: %s', (error) => { + render( + React.createElement(ChatToolPart, { + part: { + id: 'bash-error-part', + type: 'tool', + tool: 'bash', + callID: 'call-bash-error', + state: { + status: 'error', + input: { command: 'exit 1' }, + error, + }, + } as any, + }), + ); + + expect(screen.getAllByText(error)).toHaveLength(1); + }); }); describe('ChatToolPart question result rendering', () => { diff --git a/webui/src/components/common/SessionChat.tsx b/webui/src/components/common/SessionChat.tsx index 0be01168a..cf3d3cdcd 100644 --- a/webui/src/components/common/SessionChat.tsx +++ b/webui/src/components/common/SessionChat.tsx @@ -6275,7 +6275,7 @@ export function ChatToolPart({ part, pendingQuestion, onAnswer, onReject, proces )} - {status === 'error' && state.error && ( + {!isBashTool && status === 'error' && state.error && (
{state.error}
From 936de90c7daa34a6cf4d1a3ca86b7fe26fbafcff Mon Sep 17 00:00:00 2001 From: xiami762 Date: Sat, 8 Aug 2026 02:14:24 +0800 Subject: [PATCH 10/15] refactor(runtime): simplify session execution contracts --- flocks/session/prompt.py | 8 +- flocks/session/runtime/agent_loop.py | 41 ++-- flocks/session/runtime/continuation_policy.py | 2 +- flocks/session/runtime/contracts.py | 53 +---- flocks/session/runtime/session_turn.py | 42 +--- flocks/session/runtime/step_engine.py | 82 ++------ flocks/session/session_loop.py | 22 +- tests/agent/test_unified_session_loop.py | 42 +--- tests/session/runtime/test_agent_loop.py | 191 ++--------------- tests/session/runtime/test_contracts.py | 63 +----- tests/session/runtime/test_session_loop.py | 194 +----------------- tests/session/runtime/test_step_engine.py | 131 ------------ tests/session/test_actions.py | 71 ------- tests/session/test_auto_model_failover.py | 29 +-- tests/session/test_lifecycle_hooks.py | 2 +- tests/session/test_runner_step.py | 140 ------------- tests/session/test_session_context.py | 85 -------- .../test_session_loop_working_directory.py | 8 - tests/session_runtime_testkit.py | 6 - 19 files changed, 88 insertions(+), 1124 deletions(-) delete mode 100644 tests/session/runtime/test_step_engine.py delete mode 100644 tests/session/test_actions.py diff --git a/flocks/session/prompt.py b/flocks/session/prompt.py index 2d9f44592..d28643dd9 100644 --- a/flocks/session/prompt.py +++ b/flocks/session/prompt.py @@ -140,13 +140,11 @@ class ContextInfo(BaseModel): @dataclass(frozen=True) class SystemPromptBlock: - """Internal system prompt layer with cache metadata.""" + """Assembled system prompt layer.""" name: str content: str cache_scope: str - digest_inputs: Dict[str, Any] - cache_key: str @dataclass(frozen=True) @@ -843,8 +841,6 @@ def _build_cached_prompt_block( name=name, content=content, cache_scope=cache_scope, - digest_inputs=digest_inputs, - cache_key=cache_key, ) @classmethod @@ -870,8 +866,6 @@ async def _build_cached_async_prompt_block( name=name, content=content, cache_scope=cache_scope, - digest_inputs=digest_inputs, - cache_key=cache_key, ) @classmethod diff --git a/flocks/session/runtime/agent_loop.py b/flocks/session/runtime/agent_loop.py index da57a86d2..c8083a1a5 100644 --- a/flocks/session/runtime/agent_loop.py +++ b/flocks/session/runtime/agent_loop.py @@ -26,7 +26,7 @@ async def run( engine: StepEngine, ) -> AgentRunOutcome[MessageInfo]: """Run the current logical input to a session-level boundary.""" - state = turn.state + last_user = None last_message = None while not turn.aborted: @@ -36,22 +36,15 @@ async def run( if preparation.status == TurnPreparationStatus.COMPLETE: return AgentRunOutcome( status=AgentRunStatus.COMPLETED, - state=state, + last_user=last_user, last_message=preparation.last_message or last_message, ) - if preparation.status == TurnPreparationStatus.FATAL: - return AgentRunOutcome( - status=AgentRunStatus.FATAL_FAILURE, - state=state, - last_message=preparation.last_message or last_message, - error=preparation.error, - ) snapshot = preparation.snapshot if snapshot is None: return AgentRunOutcome( status=AgentRunStatus.FATAL_FAILURE, - state=state, + last_user=last_user, last_message=last_message, error=( "LoopContext returned READY without a model-turn " @@ -59,9 +52,7 @@ async def run( ), ) - state.active_model = snapshot.active_model - state.model_turn_index = snapshot.model_turn_index - state.messages = list(snapshot.messages) + last_user = snapshot.last_user try: step_result = await engine.run(snapshot) @@ -75,30 +66,26 @@ async def run( ) return AgentRunOutcome( status=AgentRunStatus.ABORTED, - state=state, + last_user=last_user, last_message=last_message, error="Aborted", ) - state.active_model = ( - step_result.effective_model or snapshot.active_model - ) boundary = await turn.commit_step(step_result) - state.messages = list(boundary.messages) last_message = boundary.last_message or last_message if turn.aborted: return AgentRunOutcome( status=AgentRunStatus.ABORTED, - state=state, + last_user=last_user, last_message=last_message, error=step_result.error, ) - if boundary.queued_inputs.messages: + if boundary.input_available: return AgentRunOutcome( status=AgentRunStatus.INPUT_AVAILABLE, - state=state, + last_user=last_user, last_message=last_message, error=step_result.error, step_result=step_result, @@ -116,7 +103,7 @@ async def run( ) return AgentRunOutcome( status=status, - state=state, + last_user=last_user, last_message=last_message, error=failure.message, step_result=step_result, @@ -127,7 +114,7 @@ async def run( if step_result.action == StepAction.COMPACT: return AgentRunOutcome( status=AgentRunStatus.CONTEXT_OVERFLOW, - state=state, + last_user=last_user, last_message=last_message, error=step_result.error, step_result=step_result, @@ -135,7 +122,7 @@ async def run( if step_result.action != StepAction.STOP: return AgentRunOutcome( status=AgentRunStatus.FATAL_FAILURE, - state=state, + last_user=last_user, last_message=last_message, error=f"Unknown step action: {step_result.action}", step_result=step_result, @@ -143,7 +130,7 @@ async def run( if step_result.error: return AgentRunOutcome( status=AgentRunStatus.FATAL_FAILURE, - state=state, + last_user=last_user, last_message=last_message, error=step_result.error, step_result=step_result, @@ -151,14 +138,14 @@ async def run( return AgentRunOutcome( status=AgentRunStatus.COMPLETED, - state=state, + last_user=last_user, last_message=last_message, step_result=step_result, ) return AgentRunOutcome( status=AgentRunStatus.ABORTED, - state=state, + last_user=last_user, last_message=last_message, error="Aborted", ) diff --git a/flocks/session/runtime/continuation_policy.py b/flocks/session/runtime/continuation_policy.py index 5fbc1a86d..f2caa7dd2 100644 --- a/flocks/session/runtime/continuation_policy.py +++ b/flocks/session/runtime/continuation_policy.py @@ -121,7 +121,7 @@ async def resolve( outcome: AgentRunOutcome[MessageInfo], ) -> ContinuationDecision[MessageInfo]: """Resolve goal and TurnFinish into a new logical turn.""" - last_user = outcome.state.metadata.get("last_user") + last_user = outcome.last_user last_message = outcome.last_message if last_user is None or last_message is None: await self.publish_turn_stopped( diff --git a/flocks/session/runtime/contracts.py b/flocks/session/runtime/contracts.py index fed280fda..51589a9f9 100644 --- a/flocks/session/runtime/contracts.py +++ b/flocks/session/runtime/contracts.py @@ -109,12 +109,9 @@ def provider_options(self) -> dict[str, Any]: class AttemptEffects: """Observable effects accumulated during one provider attempt.""" - request_sent: bool = False received_chunk: bool = False observable_output_started: bool = False tool_execution_started: bool = False - tool_execution_completed: bool = False - externally_visible: bool = False @property def replay_safe(self) -> bool: @@ -170,49 +167,16 @@ class StepResult: error: Optional[str] = None usage: Optional[dict[str, int]] = None failure: Optional[StepFailure] = None - effective_model: Optional[RuntimeModel] = None - - -@dataclass -class AgentRunState(Generic[MessageT]): - """Agent-loop-owned mutable state for one resumable agent run.""" - - session_id: str - agent_name: str - active_model: RuntimeModel - messages: list[MessageT] = field(default_factory=list) - model_turn_index: int = 0 - trace_step_offset: int = 0 - current_user_id: Optional[str] = None - metadata: dict[str, Any] = field(default_factory=dict) - - @property - def trace_step(self) -> int: - """Return the session-cumulative model-turn index.""" - return self.trace_step_offset + self.model_turn_index @dataclass(frozen=True) class ModelTurnSnapshot(Generic[MessageT]): """Immutable input presented to a step engine for one model turn.""" - session_id: str - agent_name: str active_model: RuntimeModel - model_turn_index: int trace_step: int messages: tuple[MessageT, ...] last_user: MessageT - metadata: Mapping[str, Any] = field(default_factory=dict) - - def __post_init__(self) -> None: - """Defensively freeze caller-owned collections.""" - object.__setattr__(self, "messages", tuple(self.messages)) - object.__setattr__( - self, - "metadata", - MappingProxyType(dict(self.metadata)), - ) class TurnPreparationStatus(str, Enum): @@ -221,7 +185,6 @@ class TurnPreparationStatus(str, Enum): READY = "ready" CONTINUE = "continue" COMPLETE = "complete" - FATAL = "fatal" @dataclass(frozen=True) @@ -231,25 +194,14 @@ class ModelTurnPreparation(Generic[MessageT]): status: TurnPreparationStatus snapshot: Optional[ModelTurnSnapshot[MessageT]] = None last_message: Optional[MessageT] = None - error: Optional[str] = None - - -@dataclass(frozen=True) -class QueuedInputBatch(Generic[MessageT]): - """New input made visible to the loop at a model-turn boundary.""" - - messages: tuple[MessageT, ...] = () @dataclass(frozen=True) class ModelTurnBoundary(Generic[MessageT]): """Committed session view after one model turn finishes.""" - messages: tuple[MessageT, ...] last_message: Optional[MessageT] = None - queued_inputs: QueuedInputBatch[MessageT] = field( - default_factory=QueuedInputBatch, - ) + input_available: bool = False @dataclass(frozen=True) @@ -281,7 +233,8 @@ class AgentRunOutcome(Generic[MessageT]): """Structured terminal result for a resumable agent-loop invocation.""" status: AgentRunStatus - state: AgentRunState[MessageT] + last_user: Optional[MessageT] = None last_message: Optional[MessageT] = None error: Optional[str] = None step_result: Optional[StepResult] = None + unhandled_error: bool = False diff --git a/flocks/session/runtime/session_turn.py b/flocks/session/runtime/session_turn.py index a358175a7..774f15b3b 100644 --- a/flocks/session/runtime/session_turn.py +++ b/flocks/session/runtime/session_turn.py @@ -14,11 +14,9 @@ from datetime import datetime from flocks.session.runtime.contracts import ( - AgentRunState, ModelTurnBoundary, ModelTurnPreparation, ModelTurnSnapshot, - QueuedInputBatch, RuntimeModel, StepAction, StepResult, @@ -129,21 +127,6 @@ class LoopContext: session_start_pending: bool = False model_policy: Optional[Any] = field(default=None, repr=False) continuation_policy: Optional[Any] = field(default=None, repr=False) - state: AgentRunState[MessageInfo] = field(init=False, repr=False) - - def __post_init__(self) -> None: - self.reset() - - def reset(self) -> None: - """Start a fresh AgentLoop state for the current logical input.""" - self.state = AgentRunState[MessageInfo]( - session_id=self.session.id, - agent_name=self.agent_name, - active_model=RuntimeModel(self.provider_id, self.model_id), - model_turn_index=self.step, - trace_step_offset=self.trace_step_offset, - current_user_id=self.turn_user_id, - ) @property def trace_step(self) -> int: @@ -204,10 +187,8 @@ async def prepare_step( self, ) -> ModelTurnPreparation[MessageInfo]: """Prepare one immutable model-turn snapshot from session state.""" - state = self.state SessionStatus.set(self.session.id, SessionStatusBusy()) self.step += 1 - state.model_turn_index = self.step turn_state = set_turn_state( self.session.id, step=self.step, @@ -315,8 +296,7 @@ async def prepare_step( last_message=last_assistant, ) - state.current_user_id = last_user.id - state.metadata["last_user"] = last_user + self.prepared_user_id = last_user.id await self._prepare_memory() self._schedule_title_generation(last_user, messages) @@ -338,15 +318,10 @@ async def prepare_step( return context_preparation active_model = RuntimeModel(self.provider_id, self.model_id) - state.active_model = active_model - state.messages = list(messages) return ModelTurnPreparation( status=TurnPreparationStatus.READY, snapshot=ModelTurnSnapshot( - session_id=self.session.id, - agent_name=self.agent_name, active_model=active_model, - model_turn_index=self.step, trace_step=self.trace_step, messages=tuple(messages), last_user=last_user, @@ -369,7 +344,15 @@ async def commit_step( else: post_messages = await Message.list(self.session.id) - last_user = self.state.metadata.get("last_user") + last_user = next( + ( + message + for message in reversed(post_messages) + if message.role == MessageRole.USER + and message.id == self.prepared_user_id + ), + None, + ) last_message = next( ( message @@ -448,11 +431,8 @@ async def commit_step( ) return ModelTurnBoundary( - messages=tuple(post_messages), last_message=last_message, - queued_inputs=QueuedInputBatch( - messages=(queued_user,) if queued_user is not None else (), - ), + input_available=queued_user is not None, ) async def has_late_input(self, processed_user_id: Optional[str]) -> bool: diff --git a/flocks/session/runtime/step_engine.py b/flocks/session/runtime/step_engine.py index b90040c1f..718caafb6 100644 --- a/flocks/session/runtime/step_engine.py +++ b/flocks/session/runtime/step_engine.py @@ -2,7 +2,6 @@ import asyncio import copy -import hashlib import json import os import re @@ -225,9 +224,6 @@ def _find_retryable_transport_exception(exception: Exception) -> Optional[Except return None -LlmAttemptState = AttemptEffects - - class StepCancelled(Exception): """Signal that the user cancelled the active session step.""" @@ -243,7 +239,6 @@ def __init__( agent_name: Optional[str] = None, callbacks: Optional[LoopCallbacks] = None, abort_event: Optional[asyncio.Event] = None, - session_store: Optional[Any] = None, memory_bootstrap_data: Optional[Dict[str, Any]] = None, static_cache: Optional[Dict[str, Any]] = None, defer_step_errors: bool = False, @@ -261,7 +256,6 @@ def __init__( self._external_abort = abort_event # External abort event (e.g. from SessionLoop) self._step = 0 self._recent_tool_calls: List[tuple[str, str]] = [] # Track recent (tool_name, args_json) for doom loop - self.session_store = session_store self._memory_bootstrap_data: Optional[Dict[str, Any]] = memory_bootstrap_data self._static_cache = static_cache if static_cache is not None else {} self._defer_step_errors = defer_step_errors @@ -269,7 +263,7 @@ def __init__( self._turn_additional_context = turn_additional_context self._session_start_pending = session_start_pending self._session_start_fired = False - self._attempt_state = LlmAttemptState() + self._attempt_state = AttemptEffects() self._hooked_model_requests: Dict[ str, Tuple[ModelRequest[ChatMessage], bool], @@ -293,7 +287,6 @@ def from_turn( agent_name=turn.agent_name, abort_event=turn.abort_event, callbacks=turn.callbacks, - session_store=turn.session_store, memory_bootstrap_data=turn.memory_bootstrap_data, static_cache=turn.step_static_cache, defer_step_errors=turn.auto_failover, @@ -325,7 +318,6 @@ async def run( result = await self._run_candidate( replace(snapshot, active_model=active_model), ) - result.effective_model = active_model failure = result.failure if not turn.auto_failover or failure is None: return result @@ -1255,13 +1247,10 @@ def _deferred_failure_result( decision: FailoverDecision, attempts: int, ) -> StepResult: - state = LlmAttemptState( - request_sent=self._attempt_state.request_sent, + state = AttemptEffects( received_chunk=self._attempt_state.received_chunk, observable_output_started=self._attempt_state.observable_output_started, tool_execution_started=self._attempt_state.tool_execution_started, - tool_execution_completed=self._attempt_state.tool_execution_completed, - externally_visible=self._attempt_state.externally_visible, ) return StepResult( action="stop", @@ -1283,7 +1272,7 @@ async def _process_step( last_user: MessageInfo, ) -> StepResult: """Process a single step in the loop with retry logic.""" - self._attempt_state = LlmAttemptState() + self._attempt_state = AttemptEffects() turn_execution_mode = runtime_execution_mode( getattr(last_user, "executionMode", None) ) @@ -1522,8 +1511,6 @@ async def _process_step( provider_id=self.provider_id, parent_id=last_user.id, ) - self._attempt_state.externally_visible = True - # Publish assistant message SSE event so frontends can show the message card if self.callbacks.event_publish_callback: import time as _time @@ -1988,51 +1975,22 @@ async def _build_turn_prompt_context( config_instructions = () worktree = Instance.get_worktree() - config_fingerprint = hashlib.sha256( - json.dumps( - config_data, - ensure_ascii=False, - sort_keys=True, - default=str, - ).encode("utf-8"), - ).hexdigest() - source_key = ( - self.session.id, - agent.name, - self.session.directory, - worktree, - current_tool_revision, - current_device_revision, - config_fingerprint, - ) - cached = self._static_cache.get("runtime_prompt_context") - source_context = None - if isinstance(cached, dict) and cached.get("key") == source_key: - candidate = cached.get("context") - if isinstance(candidate, TurnPromptContext): - source_context = candidate - - if source_context is None: - sandbox_context, channel_context, device_asset_hint = await asyncio.gather( - self._build_sandbox_prompt(agent, config_data=config_data), - self._build_channel_context_prompt(), - self._build_device_asset_hint(), - ) - source_context = TurnPromptContext( - tool_catalog=self._build_tool_catalog_prompt(agent), - device_asset_hint=device_asset_hint, - sandbox_context=sandbox_context, - channel_context=channel_context, - worktree=worktree, - config_instructions=config_instructions, - tool_revision=current_tool_revision, - device_revision=current_device_revision, - minimal_prompt=False, - ) - self._static_cache["runtime_prompt_context"] = { - "key": source_key, - "context": source_context, - } + sandbox_context, channel_context, device_asset_hint = await asyncio.gather( + self._build_sandbox_prompt(agent, config_data=config_data), + self._build_channel_context_prompt(), + self._build_device_asset_hint(), + ) + source_context = TurnPromptContext( + tool_catalog=self._build_tool_catalog_prompt(agent), + device_asset_hint=device_asset_hint, + sandbox_context=sandbox_context, + channel_context=channel_context, + worktree=worktree, + config_instructions=config_instructions, + tool_revision=current_tool_revision, + device_revision=current_device_revision, + minimal_prompt=False, + ) return await self._add_turn_prompt_tail( source_context, @@ -3400,7 +3358,6 @@ async def _on_tool_execution_end( tool_name: str, result: ToolResult, ) -> None: - self._attempt_state.tool_execution_completed = True if self.callbacks.on_tool_end: await self.callbacks.on_tool_end(tool_name, result) @@ -3564,7 +3521,6 @@ async def _on_tool_execution_end( "local_endpoint": stream_timeouts.is_local, }) try: - self._attempt_state.request_sent = True async for chunk in _iter_with_chunk_timeout( provider.chat_stream( model_id=self.model_id, diff --git a/flocks/session/session_loop.py b/flocks/session/session_loop.py index 5ef9cf31e..cf8e32438 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -21,7 +21,6 @@ ) from flocks.session.runtime.contracts import ( AgentRunOutcome, - AgentRunState, AgentRunStatus, RuntimeModel, ) @@ -195,9 +194,6 @@ async def run( turn.prepared_user_id or processed_user_id ) outcome = await cls._run_logical_input(turn) - processed_user_id = ( - outcome.state.current_user_id or processed_user_id - ) if await cls._should_continue( turn, continuation_policy, @@ -208,7 +204,6 @@ async def run( outcome = await cls._handle_execution_error( turn, exc, - processed_user_id, ) if await cls._settle_or_continue( @@ -230,7 +225,6 @@ async def _run_logical_input( turn: LoopContext, ) -> AgentRunOutcome[Any]: """Execute one prepared logical input through AgentLoop.""" - turn.reset() return await AgentLoop().run( turn, StepEngine.from_turn(turn), @@ -474,9 +468,7 @@ def _to_loop_result( } else None ) - unhandled_runtime_error = bool( - outcome.state.metadata.get("unhandled_runtime_error"), - ) + unhandled_runtime_error = outcome.unhandled_error return LoopResult( action=( "error" @@ -582,7 +574,6 @@ async def _recover_orphan_tools(session_id: str) -> None: async def _handle_execution_error( turn: LoopContext, error: Exception, - processed_user_id: Optional[str], ) -> AgentRunOutcome[Any]: session_id = turn.session.id log.error( @@ -610,19 +601,10 @@ async def _handle_execution_error( "session.error_event_failed", {"error": str(publish_error)}, ) - state = AgentRunState[Any]( - session_id=session_id, - agent_name=turn.agent_name, - active_model=RuntimeModel(turn.provider_id, turn.model_id), - model_turn_index=turn.step, - trace_step_offset=turn.trace_step_offset, - current_user_id=processed_user_id, - metadata={"unhandled_runtime_error": True}, - ) return AgentRunOutcome( status=AgentRunStatus.FATAL_FAILURE, - state=state, error=str(error), + unhandled_error=True, ) @classmethod diff --git a/tests/agent/test_unified_session_loop.py b/tests/agent/test_unified_session_loop.py index 8aa4bb81c..277913b0b 100644 --- a/tests/agent/test_unified_session_loop.py +++ b/tests/agent/test_unified_session_loop.py @@ -1,11 +1,7 @@ """ Tests for Phase 1: Unified UI entry via SessionLoop. -Verifies that: -1. LoopCallbacks carries model, tool, and event callbacks directly -2. StepEngine receives the same explicit callback object -3. The runtime has no reverse dependency on CLI callback globals -4. _resolve_model implements 5-level priority correctly +Verifies that _resolve_model implements its model-selection priority. """ import asyncio @@ -14,42 +10,6 @@ from unittest.mock import AsyncMock, MagicMock, patch from dataclasses import dataclass -from flocks.session.session_loop import LoopCallbacks - - -class TestLoopCallbacksFields: - """LoopCallbacks should carry all runtime callbacks directly.""" - - def test_event_publish_callback_field_exists(self): - cb = LoopCallbacks() - assert hasattr(cb, 'event_publish_callback') - assert cb.event_publish_callback is None - - def test_event_publish_callback_can_be_set(self): - publish = AsyncMock() - cb = LoopCallbacks(event_publish_callback=publish) - assert cb.event_publish_callback is publish - - def test_runtime_callbacks_are_flat(self): - on_text_delta = AsyncMock() - on_tool_start = AsyncMock() - callbacks = LoopCallbacks( - on_text_delta=on_text_delta, - on_tool_start=on_tool_start, - ) - assert callbacks.on_text_delta is on_text_delta - assert callbacks.on_tool_start is on_tool_start - - -class TestCallbackIdentity: - """The runtime should use the callbacks explicitly injected by callers.""" - - def test_explicit_callbacks_are_complete(self): - publish = AsyncMock() - cb = LoopCallbacks(event_publish_callback=publish) - assert cb.event_publish_callback is publish - - class TestResolveModel: """Test the _resolve_model 5-level priority.""" diff --git a/tests/session/runtime/test_agent_loop.py b/tests/session/runtime/test_agent_loop.py index 84acc85d5..5d614592e 100644 --- a/tests/session/runtime/test_agent_loop.py +++ b/tests/session/runtime/test_agent_loop.py @@ -3,21 +3,17 @@ from __future__ import annotations from collections import deque -from collections.abc import Callable from dataclasses import dataclass -from types import SimpleNamespace import pytest from flocks.session.runtime.agent_loop import AgentLoop from flocks.session.runtime.contracts import ( - AgentRunState, AgentRunStatus, AttemptEffects, ModelTurnBoundary, ModelTurnPreparation, ModelTurnSnapshot, - QueuedInputBatch, RuntimeModel, StepFailure, StepResult, @@ -31,12 +27,6 @@ class Message: content: str -PreparationFactory = Callable[ - [AgentRunState[Message]], - ModelTurnPreparation[Message], -] - - class FakeStepEngine: def __init__(self, results: list[StepResult]): self._results = deque(results) @@ -52,60 +42,33 @@ class FakeTurn: def __init__( self, - state: AgentRunState[Message], - preparations: list[ModelTurnPreparation[Message] | PreparationFactory], + preparations: list[ModelTurnPreparation[Message]], boundaries: list[ModelTurnBoundary[Message]], - *, - abort_after_commit: bool = False, ) -> None: - self.state = state self._preparations = deque(preparations) self._boundaries = deque(boundaries) - self.prepared_messages: list[tuple[Message, ...]] = [] self.aborted = False - self._abort_after_commit = abort_after_commit - self.session = SimpleNamespace(id=state.session_id) self.step = 0 async def prepare_step(self) -> ModelTurnPreparation[Message]: - self.prepared_messages.append(tuple(self.state.messages)) - preparation = self._preparations.popleft() - if callable(preparation): - return preparation(self.state) - return preparation + return self._preparations.popleft() async def commit_step( self, _step_result: StepResult, ) -> ModelTurnBoundary[Message]: - boundary = self._boundaries.popleft() - if self._abort_after_commit: - self.aborted = True - return boundary - - -def _state(messages: list[Message] | None = None) -> AgentRunState[Message]: - return AgentRunState( - session_id="session-1", - agent_name="rex", - active_model=RuntimeModel("provider-a", "model-a"), - messages=list(messages or [Message("user-1", "hello")]), - ) + return self._boundaries.popleft() def _ready( - state: AgentRunState[Message], + messages: tuple[Message, ...], *, turn: int = 0, ) -> ModelTurnPreparation[Message]: - messages = tuple(state.messages) return ModelTurnPreparation( status=TurnPreparationStatus.READY, snapshot=ModelTurnSnapshot( - session_id=state.session_id, - agent_name=state.agent_name, - active_model=state.active_model, - model_turn_index=turn, + active_model=RuntimeModel("provider-a", "model-a"), trace_step=turn, messages=messages, last_user=messages[-1], @@ -114,53 +77,12 @@ def _ready( async def _run( - state: AgentRunState[Message], engine: FakeStepEngine, preparations, boundaries, - *, - abort_after_commit: bool = False, ): - turn = FakeTurn( - state, - preparations, - boundaries, - abort_after_commit=abort_after_commit, - ) - outcome = await AgentLoop().run(turn, engine) - return outcome, turn - - -@pytest.mark.asyncio -async def test_loop_honors_deferred_preparation_then_completes() -> None: - user = Message("user-1", "hello") - assistant = Message("assistant-1", "done") - engine = FakeStepEngine([StepResult(action="stop")]) - outcome, turn = await _run( - _state([user]), - engine, - [ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE), _ready], - [ModelTurnBoundary(messages=(user, assistant), last_message=assistant)], - ) - - assert outcome.status == AgentRunStatus.COMPLETED - assert outcome.last_message == assistant - assert len(engine.snapshots) == 1 - assert len(turn.prepared_messages) == 2 - - -@pytest.mark.asyncio -async def test_loop_records_model_actually_used_by_step_engine() -> None: - user = Message("user-1", "hello") - assistant = Message("assistant-1", "done") - fallback = RuntimeModel("provider-b", "model-b") - outcome, _ = await _run( - _state([user]), - FakeStepEngine([StepResult(action="stop", effective_model=fallback)]), - [_ready], - [ModelTurnBoundary(messages=(user, assistant), last_message=assistant)], - ) - assert outcome.state.active_model == fallback + turn = FakeTurn(preparations, boundaries) + return await AgentLoop().run(turn, engine) @pytest.mark.asyncio @@ -171,48 +93,24 @@ async def test_loop_runs_another_step_after_tool_continue() -> None: engine = FakeStepEngine( [StepResult(action="continue"), StepResult(action="stop")], ) - outcome, _ = await _run( - _state([user]), + outcome = await _run( engine, - [_ready, lambda current: _ready(current, turn=1)], + [_ready((user,)), _ready((user, tool_result), turn=1)], [ - ModelTurnBoundary(messages=(user, tool_result), last_message=tool_result), - ModelTurnBoundary( - messages=(user, tool_result, assistant), - last_message=assistant, - ), + ModelTurnBoundary(last_message=tool_result), + ModelTurnBoundary(last_message=assistant), ], ) assert outcome.status == AgentRunStatus.COMPLETED + assert outcome.last_message == assistant assert len(engine.snapshots) == 2 assert engine.snapshots[1].messages == (user, tool_result) -@pytest.mark.asyncio -async def test_queued_input_yields_to_session_loop() -> None: - user = Message("user-1", "hello") - assistant = Message("assistant-1", "first answer") - queued_user = Message("user-2", "follow up") - outcome, _ = await _run( - _state([user]), - FakeStepEngine([StepResult(action="stop")]), - [_ready], - [ - ModelTurnBoundary( - messages=(user, assistant), - last_message=assistant, - queued_inputs=QueuedInputBatch(messages=(queued_user,)), - ), - ], - ) - assert outcome.status == AgentRunStatus.INPUT_AVAILABLE - - @pytest.mark.asyncio async def test_queued_input_precedes_final_step_failure() -> None: user = Message("user-1", "hello") failed = Message("assistant-1", "provider failed") - queued_user = Message("user-2", "try this instead") failure = StepFailure( message="provider failed", error_data={}, @@ -221,17 +119,15 @@ async def test_queued_input_precedes_final_step_failure() -> None: allow_fallback=False, attempt_state=AttemptEffects(observable_output_started=True), ) - outcome, _ = await _run( - _state([user]), + outcome = await _run( FakeStepEngine( [StepResult(action="stop", error=failure.message, failure=failure)], ), - [_ready], + [_ready((user,))], [ ModelTurnBoundary( - messages=(user, failed, queued_user), last_message=failed, - queued_inputs=QueuedInputBatch(messages=(queued_user,)), + input_available=True, ), ], ) @@ -263,62 +159,11 @@ async def test_failure_is_retryable_only_before_observable_effects( allow_fallback=True, attempt_state=effects, ) - outcome, _ = await _run( - _state([user]), + outcome = await _run( FakeStepEngine( [StepResult(action="stop", error=failure.message, failure=failure)], ), - [_ready], - [ModelTurnBoundary(messages=(user,))], + [_ready((user,))], + [ModelTurnBoundary()], ) assert outcome.status == expected_status - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("result", "expected_status"), - [ - ( - StepResult(action="compact", error="context overflow"), - AgentRunStatus.CONTEXT_OVERFLOW, - ), - (StepResult(action="unexpected"), AgentRunStatus.FATAL_FAILURE), - ], -) -async def test_loop_returns_structured_non_success_outcomes( - result: StepResult, - expected_status: AgentRunStatus, -) -> None: - user = Message("user-1", "hello") - outcome, _ = await _run( - _state([user]), - FakeStepEngine([result]), - [_ready], - [ModelTurnBoundary(messages=(user,))], - ) - assert outcome.status == expected_status - - -@pytest.mark.asyncio -async def test_loop_aborts_after_current_step_boundary() -> None: - user = Message("user-1", "hello") - outcome, _ = await _run( - _state([user]), - FakeStepEngine([StepResult(action="stop")]), - [_ready], - [ModelTurnBoundary(messages=(user,))], - abort_after_commit=True, - ) - assert outcome.status == AgentRunStatus.ABORTED - - -@pytest.mark.asyncio -async def test_ready_preparation_requires_snapshot() -> None: - outcome, _ = await _run( - _state(), - FakeStepEngine([]), - [ModelTurnPreparation(status=TurnPreparationStatus.READY)], - [], - ) - assert outcome.status == AgentRunStatus.FATAL_FAILURE - assert outcome.error == "LoopContext returned READY without a model-turn snapshot" diff --git a/tests/session/runtime/test_contracts.py b/tests/session/runtime/test_contracts.py index 0808ca940..eaf30b4fa 100644 --- a/tests/session/runtime/test_contracts.py +++ b/tests/session/runtime/test_contracts.py @@ -1,53 +1,10 @@ -"""Tests for session-neutral agent runtime contracts.""" - -import dataclasses - -import pytest +"""Tests for replay-safe runtime request contracts.""" from flocks.session.runtime.contracts import ( - AttemptEffects, ModelRequest, - ModelTurnSnapshot, - RuntimeModel, - StepAction, - StepResult, ) -def test_attempt_effects_allow_replay_only_before_observable_effects() -> None: - effects = AttemptEffects(received_chunk=True) - - assert effects.replay_safe is True - - effects.observable_output_started = True - assert effects.replay_safe is False - - effects.observable_output_started = False - effects.tool_execution_started = True - assert effects.replay_safe is False - - -def test_model_turn_snapshot_defensively_freezes_collections() -> None: - messages = ["user"] - metadata = {"tool_revision": 1} - snapshot = ModelTurnSnapshot( - session_id="session-1", - agent_name="rex", - active_model=RuntimeModel("provider", "model"), - model_turn_index=2, - trace_step=5, - messages=tuple(messages), - last_user="user", - metadata=metadata, - ) - - messages.append("new input") - metadata["tool_revision"] = 2 - - assert snapshot.messages == ("user",) - assert snapshot.metadata == {"tool_revision": 1} - - def test_model_request_freezes_and_isolates_provider_payloads() -> None: message = {"role": "user", "content": ["hello"]} tool = {"type": "function", "function": {"name": "read"}} @@ -67,21 +24,3 @@ def test_model_request_freezes_and_isolates_provider_payloads() -> None: assert request.provider_tools()[0]["function"]["name"] == "read" assert request.provider_options()["reasoning"]["effort"] == "high" - - -def test_model_request_identity_is_frozen() -> None: - request = ModelRequest( - provider_id="provider", - model_id="model", - messages=(), - tools=(), - options={}, - ) - - with pytest.raises(dataclasses.FrozenInstanceError): - request.model_id = "other" - - -def test_step_actions_are_explicit_but_unknown_adapter_values_remain_reportable() -> None: - assert StepResult(action=StepAction.CONTINUE).action == "continue" - assert StepResult(action="unexpected").action == "unexpected" diff --git a/tests/session/runtime/test_session_loop.py b/tests/session/runtime/test_session_loop.py index 540dba233..ee2f0e773 100644 --- a/tests/session/runtime/test_session_loop.py +++ b/tests/session/runtime/test_session_loop.py @@ -12,10 +12,8 @@ from flocks.session.runtime.agent_loop import AgentLoop from flocks.session.runtime.contracts import ( AgentRunOutcome, - AgentRunState, AgentRunStatus, ContinuationDecision, - RuntimeModel, StepResult, ) from flocks.session.session import Session, SessionInfo @@ -43,27 +41,14 @@ def _message(message_id: str) -> SimpleNamespace: def _outcome( - turn, - user_id: str, + user, label: str, - *, - status: AgentRunStatus = AgentRunStatus.COMPLETED, ) -> AgentRunOutcome: - state = AgentRunState( - session_id=turn.session.id, - agent_name=turn.agent_name, - active_model=RuntimeModel(turn.provider_id, turn.model_id), - current_user_id=user_id, - ) return AgentRunOutcome( - status=status, - state=state, + status=AgentRunStatus.COMPLETED, + last_user=user, last_message=SimpleNamespace(label=label), - step_result=( - StepResult(action="stop") - if status == AgentRunStatus.COMPLETED - else None - ), + step_result=StepResult(action="stop"), ) @@ -118,8 +103,7 @@ async def prepare_turn(turn): async def run_turn(turn, _engine): lease_ids.append(id(active[session.id])) return _outcome( - turn, - first_user.id if run.await_count == 1 else second_user.id, + first_user if run.await_count == 1 else second_user, "first" if run.await_count == 1 else "second", ) @@ -184,171 +168,3 @@ async def prepare(turn): assert result.error == "turn failed" assert run.await_count == 1 assert active == {} - - -@pytest.mark.asyncio -async def test_input_available_skips_terminal_continuation_resolution( - monkeypatch, - loop_io, -) -> None: - session, _ = loop_io - user = _message("msg_001") - continuation = SimpleNamespace( - prepare_logical_turn=AsyncMock(), - resolve=AsyncMock(return_value=ContinuationDecision()), - ) - monkeypatch.setattr(SessionLoop, "_continuation_policy", continuation) - run = AsyncMock() - - async def run_turn(turn, _engine): - return _outcome( - turn, - user.id, - "queued" if run.await_count == 1 else "done", - status=( - AgentRunStatus.INPUT_AVAILABLE - if run.await_count == 1 - else AgentRunStatus.COMPLETED - ), - ) - - run.side_effect = run_turn - monkeypatch.setattr(AgentLoop, "run", run) - monkeypatch.setattr( - "flocks.session.session_loop.Message.list", - AsyncMock(side_effect=[[], [user]]), - ) - - result = await SessionLoop.run( - session.id, - provider_id="provider", - model_id="model", - ) - - assert result.last_message.label == "done" - assert run.await_count == 2 - continuation.resolve.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_continuation_runs_under_same_lease( - monkeypatch, - loop_io, -) -> None: - session, active = loop_io - user = _message("msg_001") - continuation = SimpleNamespace( - prepare_logical_turn=AsyncMock(), - resolve=AsyncMock( - side_effect=[ - ContinuationDecision(messages=(user,), reason="goal"), - ContinuationDecision(), - ], - ), - ) - monkeypatch.setattr(SessionLoop, "_continuation_policy", continuation) - run = AsyncMock() - lease_ids: list[int] = [] - - async def run_turn(turn, _engine): - lease_ids.append(id(active[session.id])) - return _outcome( - turn, - user.id, - "first" if run.await_count == 1 else "second", - ) - - run.side_effect = run_turn - monkeypatch.setattr(AgentLoop, "run", run) - monkeypatch.setattr( - "flocks.session.session_loop.Message.list", - AsyncMock(side_effect=[[], [user]]), - ) - - result = await SessionLoop.run( - session.id, - provider_id="provider", - model_id="model", - ) - - assert result.last_message.label == "second" - assert len(set(lease_ids)) == 1 - assert continuation.resolve.await_count == 2 - - -@pytest.mark.asyncio -async def test_idle_is_visible_before_lease_release( - monkeypatch, - loop_io, -) -> None: - session, _ = loop_io - user = _message("msg_001") - continuation = SimpleNamespace( - prepare_logical_turn=AsyncMock(), - resolve=AsyncMock(return_value=ContinuationDecision()), - ) - monkeypatch.setattr(SessionLoop, "_continuation_policy", continuation) - monkeypatch.setattr( - AgentLoop, - "run", - AsyncMock(side_effect=lambda turn, _engine: _outcome( - turn, - user.id, - "done", - )), - ) - release_statuses: list[str] = [] - release = SessionLoop._leases.release - - def record_release(lease): - release_statuses.append(SessionStatus.get(session.id).type) - release(lease) - - monkeypatch.setattr(SessionLoop._leases, "release", record_release) - - async def touch_outside_lock(_project_id, session_id): - assert not Session.lifecycle_lock(session_id).locked() - - monkeypatch.setattr(Session, "touch", AsyncMock(side_effect=touch_outside_lock)) - monkeypatch.setattr( - "flocks.session.session_loop.Message.list", - AsyncMock(side_effect=[[], [user]]), - ) - - await SessionLoop.run( - session.id, - provider_id="provider", - model_id="model", - ) - - assert release_statuses == ["idle"] - - -@pytest.mark.asyncio -async def test_failed_busy_transition_releases_lease( - monkeypatch, - loop_io, -) -> None: - session, active = loop_io - monkeypatch.setattr( - "flocks.session.session_loop.Message.list", - AsyncMock(return_value=[]), - ) - monkeypatch.setattr( - SessionLoop, - "_mark_busy", - AsyncMock(side_effect=RuntimeError("busy failed")), - ) - run = AsyncMock() - monkeypatch.setattr(AgentLoop, "run", run) - - with pytest.raises(RuntimeError, match="busy failed"): - await SessionLoop.run( - session.id, - provider_id="provider", - model_id="model", - ) - - assert active == {} - assert SessionStatus.get(session.id).type == "idle" - run.assert_not_awaited() diff --git a/tests/session/runtime/test_step_engine.py b/tests/session/runtime/test_step_engine.py deleted file mode 100644 index 684701058..000000000 --- a/tests/session/runtime/test_step_engine.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Tests for the concrete session StepEngine.""" - -import asyncio -from types import SimpleNamespace -from unittest.mock import AsyncMock, patch - -import pytest - -from flocks.session.runtime.step_engine import ( - LlmAttemptState, - StepResult as LegacyStepResult, -) -from flocks.session.runtime.contracts import ( - AttemptEffects, - ModelTurnSnapshot, - RuntimeModel, - StepResult, -) -from flocks.session.runtime.session_turn import LoopCallbacks, LoopContext -from flocks.session.runtime.step_engine import StepCancelled, StepEngine -from flocks.session.session import SessionInfo - - -def _turn(*, aborted: bool = False) -> LoopContext: - abort_event = asyncio.Event() - if aborted: - abort_event.set() - return LoopContext( - session=SessionInfo.model_construct( - id="session-1", - projectID="project", - directory="/tmp/project", - agent="rex", - status="active", - ), - provider_id="provider", - model_id="model", - agent_name="rex", - callbacks=LoopCallbacks(event_publish_callback=None), - abort_event=abort_event, - model_candidates=[RuntimeModel("provider", "model")], - session_start_pending=True, - ) - - -def _snapshot(last_user) -> ModelTurnSnapshot: - return ModelTurnSnapshot( - session_id="session-1", - agent_name="rex", - active_model=RuntimeModel("provider", "model"), - model_turn_index=1, - trace_step=8, - messages=(last_user,), - last_user=last_user, - ) - - -@pytest.mark.asyncio -async def test_step_engine_executes_one_immutable_snapshot() -> None: - last_user = SimpleNamespace(id="user-1") - expected = StepResult(action="stop", content="done") - turn = _turn() - engine = StepEngine.from_turn(turn) - - async def execute(messages, user): - engine._session_start_fired = True - return expected - - process_step = AsyncMock(side_effect=execute) - - with patch.object(StepEngine, "_process_step", process_step): - result = await engine.run(_snapshot(last_user)) - - assert result is expected - assert engine._step == 8 - process_step.assert_awaited_once_with([last_user], last_user) - assert turn.session_start_pending is False - assert turn._current_step_task is None - - -@pytest.mark.asyncio -async def test_step_engine_refreshes_memory_loaded_after_construction() -> None: - last_user = SimpleNamespace(id="user-1") - turn = _turn() - engine = StepEngine.from_turn(turn) - loaded_memory = { - "instructions": "remember this", - "main_memory": {"content": "project context", "inject": True}, - } - turn.memory_bootstrap_data = loaded_memory - - async def execute(_messages, _user): - assert engine._memory_bootstrap_data is loaded_memory - return StepResult(action="stop", content="done") - - with patch.object(StepEngine, "_process_step", side_effect=execute): - await engine.run(_snapshot(last_user)) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("aborted", "expected_error"), - [ - (False, asyncio.CancelledError), - (True, StepCancelled), - ], -) -async def test_step_engine_only_translates_user_abort( - aborted, - expected_error, -) -> None: - last_user = SimpleNamespace(id="user-1") - turn = _turn(aborted=aborted) - engine = StepEngine.from_turn(turn) - - with ( - patch.object( - StepEngine, - "_process_step", - AsyncMock(side_effect=asyncio.CancelledError), - ), - pytest.raises(expected_error), - ): - await engine.run(_snapshot(last_user)) - - assert turn._current_step_task is None - - -def test_legacy_runner_contract_exports_remain_compatible() -> None: - assert LegacyStepResult is StepResult - assert LlmAttemptState is AttemptEffects diff --git a/tests/session/test_actions.py b/tests/session/test_actions.py deleted file mode 100644 index b6106edcb..000000000 --- a/tests/session/test_actions.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Tests for non-agent session actions.""" - -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import pytest - -from flocks.command.command import Command -from flocks.session import actions - - -@pytest.mark.asyncio -async def test_render_session_command_resolves_template(monkeypatch) -> None: - monkeypatch.setattr( - Command, - "get", - lambda _name: SimpleNamespace(template="Do $ARGUMENTS"), - ) - - result = await actions.render_session_command( - "session-1", - "test", - "the work", - ) - - assert result["template"] == "Do the work" - - -@pytest.mark.asyncio -async def test_run_session_shell_preserves_legacy_response(monkeypatch) -> None: - monkeypatch.setattr( - actions.Session, - "get_by_id", - AsyncMock( - return_value=SimpleNamespace(directory="/tmp/project"), - ), - ) - messages = [ - SimpleNamespace(id="user-1"), - SimpleNamespace(id="assistant-1"), - ] - monkeypatch.setattr( - actions.Message, - "create", - AsyncMock(side_effect=messages), - ) - process = SimpleNamespace( - communicate=AsyncMock(return_value=(b"done", b"")), - returncode=0, - ) - create_process = AsyncMock(return_value=process) - monkeypatch.setattr( - actions.asyncio, - "create_subprocess_shell", - create_process, - ) - - result = await actions.run_session_shell( - "session-1", - "rex", - "echo done", - ) - - create_process.assert_awaited_once_with( - "echo done", - stdout=actions.asyncio.subprocess.PIPE, - stderr=actions.asyncio.subprocess.PIPE, - cwd="/tmp/project", - ) - assert result["info"]["id"] == "assistant-1" - assert result["parts"][0]["state"]["output"] == "done" diff --git a/tests/session/test_auto_model_failover.py b/tests/session/test_auto_model_failover.py index 388b10c18..4890bcedc 100644 --- a/tests/session/test_auto_model_failover.py +++ b/tests/session/test_auto_model_failover.py @@ -8,14 +8,17 @@ from flocks.session.runtime.continuation_policy import DEFAULT_CONTINUATION_POLICY from flocks.session.runtime.agent_loop import AgentLoop -from flocks.session.runtime.contracts import ModelTurnSnapshot, RuntimeModel +from flocks.session.runtime.contracts import ( + AttemptEffects, + ModelTurnSnapshot, + RuntimeModel, +) from flocks.session.message import Message, MessageRole from flocks.session.runtime.model_policy import ( DEFAULT_MODEL_ROUTING_POLICY, AutoFailoverCooldown, ) from flocks.session.runtime.step_engine import ( - LlmAttemptState, StepEngine, StepFailure, StepResult, @@ -94,7 +97,7 @@ def _failure( reason: str = "server_error", safe: bool = True, ) -> StepResult: - state = LlmAttemptState(observable_output_started=not safe) + state = AttemptEffects(observable_output_started=not safe) message = "provider failed" return StepResult( action="stop", @@ -120,10 +123,7 @@ async def _process_step_with_failover( turn.callbacks = callbacks return await StepEngine.from_turn(turn).run( ModelTurnSnapshot( - session_id=turn.session.id, - agent_name=turn.agent_name, active_model=RuntimeModel(turn.provider_id, turn.model_id), - model_turn_index=turn.step, trace_step=turn.trace_step, messages=tuple(messages), last_user=last_user, @@ -218,7 +218,7 @@ async def test_auto_runner_uses_standard_retry_policy( monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.get", lambda _provider_id: provider) monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.apply_config", AsyncMock()) monkeypatch.setattr( - "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompts", + "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompt_blocks", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -284,7 +284,7 @@ async def test_last_auto_candidate_uses_standard_retry_policy( monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.get", lambda _provider_id: provider) monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.apply_config", AsyncMock()) monkeypatch.setattr( - "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompts", + "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompt_blocks", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -457,7 +457,7 @@ async def call_llm(*_args, **_kwargs): monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.get", lambda _provider_id: provider) monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.apply_config", AsyncMock()) monkeypatch.setattr( - "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompts", + "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompt_blocks", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -599,7 +599,7 @@ async def stream(): monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.get", lambda _provider_id: provider) monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.apply_config", AsyncMock()) monkeypatch.setattr( - "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompts", + "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompt_blocks", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -711,7 +711,7 @@ async def preflight_failure(_runner, _messages, _last_user): assistant_message_id=None, reason="provider_unavailable", allow_fallback=True, - attempt_state=LlmAttemptState(), + attempt_state=AttemptEffects(), attempts=0, ), ) @@ -1569,20 +1569,13 @@ async def test_unsupported_session_loop_ignores_auto_authorization( async def run_turn(_loop, ctx, _engine): from flocks.session.runtime.contracts import ( AgentRunOutcome, - AgentRunState, AgentRunStatus, ) nonlocal captured_ctx captured_ctx = ctx - state = AgentRunState( - session_id=ctx.session.id, - agent_name=ctx.agent_name, - active_model=RuntimeModel(ctx.provider_id, ctx.model_id), - ) return AgentRunOutcome( status=AgentRunStatus.ABORTED, - state=state, ) build_candidates = AsyncMock() diff --git a/tests/session/test_lifecycle_hooks.py b/tests/session/test_lifecycle_hooks.py index b3412b0a1..768a4562c 100644 --- a/tests/session/test_lifecycle_hooks.py +++ b/tests/session/test_lifecycle_hooks.py @@ -295,7 +295,7 @@ async def test_real_user_arriving_during_goal_evaluation_wins() -> None: ctx.callbacks = LoopCallbacks(event_publish_callback=AsyncMock()) create_message = AsyncMock() outcome = SimpleNamespace( - state=SimpleNamespace(metadata={"last_user": user}), + last_user=user, last_message=assistant, ) diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index 819975c73..032bbc7d4 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -1479,8 +1479,6 @@ async def test_to_chat_messages_uses_structured_anthropic_system_blocks(monkeypa name=name, content=content, cache_scope=cache_scope, - digest_inputs={}, - cache_key=name, ) for name, content, cache_scope in ( ("provider", "provider prompt", "global"), @@ -3238,141 +3236,3 @@ async def test_to_chat_messages_expands_workflow_node_ref_marker(monkeypatch): assert "node_id: query_fofa" in chat_messages[0].content assert "node_type: python" in chat_messages[0].content assert "只修改这个节点的代码并保留其他节点不变" in chat_messages[0].content - -@pytest.mark.asyncio -async def test_turn_prompt_context_reuses_runtime_snapshot(monkeypatch): - runner = _make_runner("ses_runtime_prompt_snapshot") - agent = _make_agent(name="rex") - last_user = UserMessageInfo( - id="msg_runtime_prompt_snapshot", - sessionID=runner.session.id, - role="user", - time={"created": 1_000}, - agent="rex", - model={"providerID": "anthropic", "modelID": "claude-sonnet"}, - ) - config = SimpleNamespace( - instructions=["rules.md"], - model_dump=lambda **_kwargs: {"sandbox": {"enabled": True}}, - ) - sandbox = AsyncMock(return_value="sandbox prompt") - channel = AsyncMock(return_value="channel prompt") - device = AsyncMock(return_value="device prompt") - catalog = MagicMock(return_value="tool catalog") - - monkeypatch.setattr("flocks.config.Config.get", AsyncMock(return_value=config)) - monkeypatch.setattr("flocks.project.instance.Instance.get_worktree", lambda: "/tmp") - monkeypatch.setattr("flocks.tool.device.store.device_revision", lambda: 3) - monkeypatch.setattr(runner_mod.ToolRegistry, "revision", lambda: 7) - monkeypatch.setattr(runner, "_build_sandbox_prompt", sandbox) - monkeypatch.setattr(runner, "_build_channel_context_prompt", channel) - monkeypatch.setattr(runner, "_build_device_asset_hint", device) - monkeypatch.setattr(runner, "_build_tool_catalog_prompt", catalog) - - first = await runner._build_turn_prompt_context( - agent=agent, - messages=[last_user], - last_user=last_user, - tools=[], - ) - second = await runner._build_turn_prompt_context( - agent=agent, - messages=[last_user], - last_user=last_user, - tools=[], - ) - - assert first == second - assert first.tool_revision == 7 - assert first.device_revision == 3 - assert first.config_instructions == ("rules.md",) - sandbox.assert_awaited_once() - channel.assert_awaited_once() - device.assert_awaited_once() - catalog.assert_called_once() - -@pytest.mark.asyncio -async def test_minimal_turn_prompt_context_skips_runtime_sources(monkeypatch): - runner = _make_runner("ses_minimal_runtime_prompt") - runner._turn_additional_context = "delegated context" - agent = _make_agent(name="rex-junior") - last_user = UserMessageInfo( - id="msg_minimal_runtime_prompt", - sessionID=runner.session.id, - role="user", - time={"created": 1_000}, - agent="rex-junior", - model={"providerID": "anthropic", "modelID": "claude-sonnet"}, - ) - sandbox = AsyncMock(return_value="sandbox prompt") - channel = AsyncMock(return_value="channel prompt") - device = AsyncMock(return_value="device prompt") - monkeypatch.setattr(runner, "_build_sandbox_prompt", sandbox) - monkeypatch.setattr(runner, "_build_channel_context_prompt", channel) - monkeypatch.setattr(runner, "_build_device_asset_hint", device) - - context = await runner._build_turn_prompt_context( - agent=agent, - messages=[last_user], - last_user=last_user, - tools=[], - minimal_prompt=True, - ) - - assert context.minimal_prompt is True - assert context.additional_context == "delegated context" - sandbox.assert_not_awaited() - channel.assert_not_awaited() - device.assert_not_awaited() - -@pytest.mark.asyncio -async def test_process_step_passes_device_hint_in_turn_prompt_context(monkeypatch): - runner = _make_runner("ses_runner_device_hint_order") - runner.callbacks = LoopCallbacks(on_error=AsyncMock()) - - last_user = UserMessageInfo( - id="msg_user_device_hint_order", - sessionID=runner.session.id, - role="user", - time={"created": 1_000}, - agent="rex", - model={"providerID": "anthropic", "modelID": "claude-sonnet"}, - ) - - agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) - provider = MagicMock() - provider.is_configured.return_value = True - assistant_msg = SimpleNamespace(id="msg_assistant_device_hint_order") - build_system_prompts = AsyncMock(return_value=["provider", "tool catalog awareness", "device hint"]) - - monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) - monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) - monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", build_system_prompts) - monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) - device_hint_mock = AsyncMock(return_value="device hint") - monkeypatch.setattr(runner, "_build_device_asset_hint", device_hint_mock) - monkeypatch.setattr("flocks.tool.device.store.device_revision", lambda: 9) - monkeypatch.setattr( - runner, - "_to_chat_messages", - AsyncMock(return_value=[SimpleNamespace(role="user", content="hi")]), - ) - monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) - monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) - monkeypatch.setattr(runner_mod.Message, "create", AsyncMock(return_value=assistant_msg)) - monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) - monkeypatch.setattr( - runner, - "_call_llm", - AsyncMock(return_value=StepResult(action="stop", content="done")), - ) - - result = await runner._process_step([last_user], last_user) - - assert result.content == "done" - build_system_prompts.assert_awaited_once() - kwargs = build_system_prompts.await_args.kwargs - turn_context = kwargs["turn_context"] - assert turn_context.device_revision == 9 - assert turn_context.device_asset_hint == "device hint" diff --git a/tests/session/test_session_context.py b/tests/session/test_session_context.py index 6216267d4..31ea064d4 100644 --- a/tests/session/test_session_context.py +++ b/tests/session/test_session_context.py @@ -5,16 +5,12 @@ 1. SessionContext protocol is properly defined 2. DefaultSessionContext implements all methods 3. DefaultSessionContext delegates to underlying session modules -4. LoopContext carries session_store -5. StepEngine accepts session_store """ import pytest from unittest.mock import AsyncMock, MagicMock, patch from flocks.session.core.context import SessionContext, DefaultSessionContext -from flocks.session.session_loop import LoopContext -from flocks.session.runtime.step_engine import StepEngine class TestSessionContextProtocol: @@ -135,84 +131,3 @@ async def test_touch_delegates_to_session(self): with patch("flocks.session.session.Session.touch", new_callable=AsyncMock) as mock_touch: await ctx.touch() mock_touch.assert_called_once_with("proj-1", "ses-123") - - -class TestLoopContextSessionStore: - """LoopContext should carry session_store.""" - - def test_loop_context_has_session_store_field(self): - import asyncio - session = MagicMock() - session.id = "test" - session.directory = "/test" - session.project_id = "proj" - - ctx = LoopContext( - session=session, - provider_id="anthropic", - model_id="claude-sonnet-4", - agent_name="rex", - ) - assert ctx.session_store is None - - def test_loop_context_with_session_store(self): - session = MagicMock() - session.id = "test" - session.directory = "/test" - session.project_id = "proj" - - session_store = DefaultSessionContext(session) - ctx = LoopContext( - session=session, - provider_id="anthropic", - model_id="claude-sonnet-4", - agent_name="rex", - session_store=session_store, - ) - assert ctx.session_store is session_store - assert ctx.session_store.session_id == "test" - - def test_loop_context_tracks_observed_prompt_tokens(self): - # B3 — LoopContext must expose ``last_observed_prompt_tokens`` so - # the overflow decision can prefer the provider's reported usage - # over our synthetic estimate. - session = MagicMock() - session.id = "test" - session.directory = "/test" - session.project_id = "proj" - - ctx = LoopContext( - session=session, - provider_id="anthropic", - model_id="claude-sonnet-4", - agent_name="rex", - ) - assert ctx.last_observed_prompt_tokens == 0 - ctx.last_observed_prompt_tokens = 123_456 - assert ctx.last_observed_prompt_tokens == 123_456 - - -class TestStepEngineSessionStore: - """StepEngine should accept session_store.""" - - def test_step_engine_accepts_session_store(self): - session = MagicMock() - session.id = "test" - session.directory = "/test" - session.project_id = "proj" - - session_store = DefaultSessionContext(session) - runner = StepEngine( - session=session, - session_store=session_store, - ) - assert runner.session_store is session_store - - def test_step_engine_session_store_defaults_to_none(self): - session = MagicMock() - session.id = "test" - session.directory = "/test" - session.project_id = "proj" - - runner = StepEngine(session=session) - assert runner.session_store is None diff --git a/tests/session/test_session_loop_working_directory.py b/tests/session/test_session_loop_working_directory.py index dbda31ed0..9794edfe8 100644 --- a/tests/session/test_session_loop_working_directory.py +++ b/tests/session/test_session_loop_working_directory.py @@ -6,9 +6,7 @@ from flocks.session.runtime.agent_loop import AgentLoop from flocks.session.runtime.contracts import ( AgentRunOutcome, - AgentRunState, AgentRunStatus, - RuntimeModel, ) from flocks.session.message import Message from flocks.session.session import Session, SessionInfo @@ -25,14 +23,8 @@ async def test_run_uses_runtime_working_directory(monkeypatch: pytest.MonkeyPatc ) async def run_agent_turn(context, _engine): - state = AgentRunState( - session_id=context.session.id, - agent_name=context.agent_name, - active_model=RuntimeModel(context.provider_id, context.model_id), - ) return AgentRunOutcome( status=AgentRunStatus.ABORTED, - state=state, ) run_turn = AsyncMock(side_effect=run_agent_turn) diff --git a/tests/session_runtime_testkit.py b/tests/session_runtime_testkit.py index 15cf0f4e4..0587b3d17 100644 --- a/tests/session_runtime_testkit.py +++ b/tests/session_runtime_testkit.py @@ -15,21 +15,15 @@ async def run_logical_turns( """Run logical turns without acquiring a persisted session lease.""" turn.callbacks = callbacks policy = turn.continuation_policy or SessionLoop._continuation_policy - processed_user_id = None while True: try: await policy.prepare_logical_turn(turn) - processed_user_id = turn.prepared_user_id or processed_user_id outcome = await SessionLoop._run_logical_input(turn) - processed_user_id = ( - outcome.state.current_user_id or processed_user_id - ) if await SessionLoop._should_continue(turn, policy, outcome): continue except Exception as exc: outcome = await SessionLoop._handle_execution_error( turn, exc, - processed_user_id, ) return SessionLoop._to_loop_result(turn, outcome) From 920f1c798c212804cf6c72d60e833e4493e660a0 Mon Sep 17 00:00:00 2001 From: xiami762 Date: Tue, 11 Aug 2026 11:38:05 +0800 Subject: [PATCH 11/15] chore: split runtime-adjacent changes --- flocks/cli/commands/task.py | 50 ++-- flocks/server/app.py | 12 +- flocks/server/routes/stats.py | 4 +- flocks/server/routes/task_entities.py | 92 ++++---- flocks/task/__init__.py | 4 +- flocks/task/background.py | 42 ++-- .../{schedule_task_manager.py => manager.py} | 12 +- flocks/task/plugin_sync.py | 4 +- flocks/task/scheduler.py | 4 +- flocks/tool/task/schedule_task_center.py | 64 +++--- .../test_task_queue_integration.py | 86 +++---- .../test_task_scheduler_context_route.py | 8 +- tests/server/test_lifespan.py | 4 +- tests/server/test_server.py | 44 ++-- .../storage/test_sqlite_connection_config.py | 6 +- tests/task/test_task.py | 216 +++++++++--------- tests/tool/test_task_center_compat.py | 38 +-- tests/tool/test_task_list_routing.py | 2 +- .../src/components/common/SessionChat.test.ts | 22 -- webui/src/components/common/SessionChat.tsx | 2 +- 20 files changed, 344 insertions(+), 372 deletions(-) rename flocks/task/{schedule_task_manager.py => manager.py} (99%) diff --git a/flocks/cli/commands/task.py b/flocks/cli/commands/task.py index cf7024b9c..25ec387c1 100644 --- a/flocks/cli/commands/task.py +++ b/flocks/cli/commands/task.py @@ -45,11 +45,11 @@ def task_dashboard(): async def _dashboard(): await Storage.init() - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.store import TaskStore await TaskStore.init() - counts = await ScheduleTaskManager.dashboard() + counts = await TaskManager.dashboard() panel_lines = [ f"🟢 Running: {counts.get('running', 0)}", @@ -63,7 +63,7 @@ async def _dashboard(): console.print(Panel("\n".join(panel_lines), title="📋 Task Center", border_style="cyan")) - unviewed = await ScheduleTaskManager.get_unviewed_results() + unviewed = await TaskManager.get_unviewed_results() if unviewed: console.print() console.print("[bold]Unviewed completed tasks:[/bold]") @@ -88,7 +88,7 @@ def task_list( async def _list_tasks(status_val, type_val, limit, fmt): await Storage.init() - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.models import SchedulerStatus, TaskStatus from flocks.task.store import TaskStore await TaskStore.init() @@ -99,7 +99,7 @@ async def _list_tasks(status_val, type_val, limit, fmt): scheduler_status = SchedulerStatus.ACTIVE elif status_val in ("paused", "disabled"): scheduler_status = SchedulerStatus.DISABLED - tasks, total = await ScheduleTaskManager.list_schedulers(status=scheduler_status, limit=limit) + tasks, total = await TaskManager.list_schedulers(status=scheduler_status, limit=limit) else: task_status = None if status_val: @@ -108,7 +108,7 @@ async def _list_tasks(status_val, type_val, limit, fmt): task_status = TaskStatus(mapped_status) except ValueError as exc: raise typer.BadParameter(f"Invalid execution status: {status_val}") from exc - tasks, total = await ScheduleTaskManager.list_executions( + tasks, total = await TaskManager.list_executions( status=task_status, limit=limit, ) @@ -160,13 +160,13 @@ def task_show(task_id: str = typer.Argument(..., help="Task ID")): async def _show_task(task_id: str): await Storage.init() - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.store import TaskStore await TaskStore.init() - task = await ScheduleTaskManager.get_execution(task_id) + task = await TaskManager.get_execution(task_id) if task is None: - task = await ScheduleTaskManager.get_scheduler(task_id) + task = await TaskManager.get_scheduler(task_id) if not task: console.print(f"[red]Task {task_id} not found[/red]") raise typer.Exit(1) @@ -243,7 +243,7 @@ def task_create( async def _create_task(title, description, task_type, priority, mode, agent, workflow, skills, cron, prompt): await Storage.init() - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.models import ( ExecutionMode, SchedulerMode, @@ -270,7 +270,7 @@ async def _create_task(title, description, task_type, priority, mode, agent, wor source = TaskSource(user_prompt=prompt) if prompt else None - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title=title, description=description, mode=scheduler_mode, @@ -284,7 +284,7 @@ async def _create_task(title, description, task_type, priority, mode, agent, wor ) console.print(f"[green]✅ Created scheduler:[/green] {scheduler.id} {scheduler.title}") if trigger.run_immediately: - executions, _ = await ScheduleTaskManager.list_scheduler_executions(scheduler.id, limit=1) + executions, _ = await TaskManager.list_scheduler_executions(scheduler.id, limit=1) if executions: console.print( f"[green]↳ execution:[/green] {executions[0].id} ({executions[0].status.value})" @@ -306,20 +306,20 @@ def task_queue( async def _queue(pause: bool, resume: bool): await Storage.init() - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.store import TaskStore await TaskStore.init() if pause: - ScheduleTaskManager.pause_queue() + TaskManager.pause_queue() console.print("[yellow]Queue paused[/yellow]") return if resume: - ScheduleTaskManager.resume_queue() + TaskManager.resume_queue() console.print("[green]Queue resumed[/green]") return - qs = await ScheduleTaskManager.queue_status() + qs = await TaskManager.queue_status() console.print(Panel( f"Paused: {qs['paused']}\n" f"Max concurrent: {qs['max_concurrent']}\n" @@ -342,11 +342,11 @@ def task_scheduled(): async def _scheduled(): await Storage.init() - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.store import TaskStore await TaskStore.init() - tasks, _ = await ScheduleTaskManager.list_schedulers(scheduled_only=True, limit=100) + tasks, _ = await TaskManager.list_schedulers(scheduled_only=True, limit=100) if not tasks: console.print("[dim]No scheduled tasks[/dim]") return @@ -384,11 +384,11 @@ def task_cancel(task_id: str = typer.Argument(..., help="Task ID")): async def _cancel(task_id: str): await Storage.init() - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.store import TaskStore await TaskStore.init() - task = await ScheduleTaskManager.cancel_execution(task_id) + task = await TaskManager.cancel_execution(task_id) if not task: console.print(f"[red]Task {task_id} not found[/red]") raise typer.Exit(1) @@ -403,11 +403,11 @@ def task_retry(task_id: str = typer.Argument(..., help="Task ID")): async def _retry(task_id: str): await Storage.init() - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.store import TaskStore await TaskStore.init() - task = await ScheduleTaskManager.retry_execution(task_id) + task = await TaskManager.retry_execution(task_id) if not task: console.print(f"[red]Task {task_id} not found or not failed[/red]") raise typer.Exit(1) @@ -422,13 +422,13 @@ def task_rerun(task_id: str = typer.Argument(..., help="Task ID")): async def _rerun(task_id: str): await Storage.init() - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.store import TaskStore await TaskStore.init() - task = await ScheduleTaskManager.rerun_execution(task_id) + task = await TaskManager.rerun_execution(task_id) if task is None: - task = await ScheduleTaskManager.rerun_scheduler(task_id) + task = await TaskManager.rerun_scheduler(task_id) if not task: console.print(f"[red]Task {task_id} not found[/red]") raise typer.Exit(1) diff --git a/flocks/server/app.py b/flocks/server/app.py index 9d47761ec..eb304cc7b 100644 --- a/flocks/server/app.py +++ b/flocks/server/app.py @@ -350,16 +350,16 @@ async def _sync_workflows_phase() -> None: # Start Task Center (scheduler + queue executor) try: - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager await _run_startup_phase( log, "schedule_task_manager.start", - ScheduleTaskManager.start, + TaskManager.start, ) log.info("schedule_task_manager.started") except Exception as e: - from flocks.task.schedule_task_manager import ScheduleTaskManager - ScheduleTaskManager.mark_start_failed(e) + from flocks.task.manager import TaskManager + TaskManager.mark_start_failed(e) log.warning("schedule_task_manager.start.failed", {"error": str(e)}) # Seed built-in scheduled tasks from .flocks/plugins/tasks/*.json (idempotent) @@ -538,9 +538,9 @@ async def _delayed_trigger_runtime_start() -> None: # Stop Task Center try: - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.store import TaskStore - await ScheduleTaskManager.stop() + await TaskManager.stop() await TaskStore.close() log.info("schedule_task_manager.stopped") except Exception as e: diff --git a/flocks/server/routes/stats.py b/flocks/server/routes/stats.py index 205ce1be9..af2e74766 100644 --- a/flocks/server/routes/stats.py +++ b/flocks/server/routes/stats.py @@ -19,7 +19,7 @@ from flocks.server.routes.provider import list_providers from flocks.server.routes.workflow import _list_workflows_from_fs, _migrate_storage_to_filesystem from flocks.skill.skill import Skill -from flocks.task.schedule_task_manager import ScheduleTaskManager +from flocks.task.manager import TaskManager from flocks.tool.registry import ToolRegistry from flocks.utils.log import Log @@ -70,7 +70,7 @@ def _should_count_agent(agent: Any) -> bool: async def _task_dashboard() -> dict[str, Any]: - return await ScheduleTaskManager.dashboard() + return await TaskManager.dashboard() async def _safe_dashboard(failures: list[str]) -> dict[str, Any]: diff --git a/flocks/server/routes/task_entities.py b/flocks/server/routes/task_entities.py index 18bbf35d6..ffbe42118 100644 --- a/flocks/server/routes/task_entities.py +++ b/flocks/server/routes/task_entities.py @@ -145,10 +145,10 @@ def _parse_task_type(task_type: str) -> str: @router.get("/task-system/notice") async def get_task_system_notice(): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager started_at = time.perf_counter() - notice = await ScheduleTaskManager.get_task_page_notice() + notice = await TaskManager.get_task_page_notice() log_route_timing(log, "task.notice.complete", started_at=started_at, extra={ "has_notice": bool(notice), }) @@ -157,10 +157,10 @@ async def get_task_system_notice(): @router.get("/task-system/dashboard") async def task_dashboard(): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager started_at = time.perf_counter() - payload = await ScheduleTaskManager.dashboard() + payload = await TaskManager.dashboard() log_route_timing(log, "task.dashboard.complete", started_at=started_at, extra={ "running": payload.get("running"), "queued": payload.get("queued"), @@ -171,10 +171,10 @@ async def task_dashboard(): @router.get("/task-system/queue/status") async def task_queue_status(): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager started_at = time.perf_counter() - payload = await ScheduleTaskManager.queue_status() + payload = await TaskManager.queue_status() log_route_timing(log, "task.queue_status.complete", started_at=started_at, extra={ "queued": payload.get("queued"), "running": payload.get("running"), @@ -185,17 +185,17 @@ async def task_queue_status(): @router.post("/task-system/queue/pause") async def pause_task_queue(): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - ScheduleTaskManager.pause_queue() + TaskManager.pause_queue() return {"paused": True} @router.post("/task-system/queue/resume") async def resume_task_queue(): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - ScheduleTaskManager.resume_queue() + TaskManager.resume_queue() return {"paused": False} @@ -209,9 +209,9 @@ async def list_schedulers( offset: int = Query(0, ge=0), limit: int = Query(20, ge=1, le=100), ): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - items, total = await ScheduleTaskManager.list_schedulers( + items, total = await TaskManager.list_schedulers( status=_parse_scheduler_status_filter(status_filter), priority=_parse_priority(priority), scheduled_only=scheduled_only, @@ -230,7 +230,7 @@ async def list_schedulers( @router.post("/task-schedulers", status_code=status.HTTP_201_CREATED) async def create_scheduler(req: SchedulerCreateRequest): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.models import ( SchedulerMode, TaskSource, @@ -268,7 +268,7 @@ async def create_scheduler(req: SchedulerCreateRequest): detail=str(exc), ) from exc - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title=req.title, description=req.description, mode=mode, @@ -289,9 +289,9 @@ async def create_scheduler(req: SchedulerCreateRequest): @router.get("/task-schedulers/{scheduler_id}") async def get_scheduler(scheduler_id: str): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - scheduler = await ScheduleTaskManager.get_scheduler(scheduler_id) + scheduler = await TaskManager.get_scheduler(scheduler_id) if not scheduler: raise HTTPException(404, "Task scheduler not found") return scheduler.model_dump(mode="json", by_alias=True) @@ -299,7 +299,7 @@ async def get_scheduler(scheduler_id: str): @router.put("/task-schedulers/{scheduler_id}") async def update_scheduler(scheduler_id: str, req: SchedulerUpdateRequest): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager fields = {k: v for k, v in req.model_dump(exclude_none=True).items()} if "priority" in fields: @@ -313,7 +313,7 @@ async def update_scheduler(scheduler_id: str, req: SchedulerUpdateRequest): run_at = fields.pop("run_at", None) user_prompt = fields.pop("user_prompt", None) try: - scheduler = await ScheduleTaskManager.update_scheduler_with_trigger( + scheduler = await TaskManager.update_scheduler_with_trigger( scheduler_id, fields=fields, cron=cron, @@ -335,18 +335,18 @@ async def update_scheduler(scheduler_id: str, req: SchedulerUpdateRequest): @router.delete("/task-schedulers/{scheduler_id}") async def delete_scheduler(scheduler_id: str): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - if not await ScheduleTaskManager.delete_scheduler(scheduler_id): + if not await TaskManager.delete_scheduler(scheduler_id): raise HTTPException(404, "Task scheduler not found") return {"ok": True} @router.post("/task-schedulers/{scheduler_id}/enable") async def enable_scheduler(scheduler_id: str): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - scheduler = await ScheduleTaskManager.enable_scheduler(scheduler_id) + scheduler = await TaskManager.enable_scheduler(scheduler_id) if not scheduler: raise HTTPException(404, "Task scheduler not found") return scheduler.model_dump(mode="json", by_alias=True) @@ -354,9 +354,9 @@ async def enable_scheduler(scheduler_id: str): @router.post("/task-schedulers/{scheduler_id}/disable") async def disable_scheduler(scheduler_id: str): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - scheduler = await ScheduleTaskManager.disable_scheduler(scheduler_id) + scheduler = await TaskManager.disable_scheduler(scheduler_id) if not scheduler: raise HTTPException(404, "Task scheduler not found") return scheduler.model_dump(mode="json", by_alias=True) @@ -368,9 +368,9 @@ async def list_scheduler_executions( offset: int = Query(0, ge=0), limit: int = Query(20, ge=1, le=100), ): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - items, total = await ScheduleTaskManager.list_scheduler_executions( + items, total = await TaskManager.list_scheduler_executions( scheduler_id, offset=offset, limit=limit ) return PaginatedResponse( @@ -383,9 +383,9 @@ async def list_scheduler_executions( @router.post("/task-schedulers/{scheduler_id}/run") async def run_scheduler(scheduler_id: str): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - execution = await ScheduleTaskManager.rerun_scheduler(scheduler_id) + execution = await TaskManager.rerun_scheduler(scheduler_id) if not execution: raise HTTPException(404, "Task scheduler not found") return execution.model_dump(mode="json", by_alias=True) @@ -402,9 +402,9 @@ async def list_executions( offset: int = Query(0, ge=0), limit: int = Query(20, ge=1, le=100), ): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - items, total = await ScheduleTaskManager.list_executions( + items, total = await TaskManager.list_executions( scheduler_id=scheduler_id, status=_parse_execution_status_filter(status_filter), priority=_parse_priority(priority), @@ -424,23 +424,23 @@ async def list_executions( @router.post("/task-executions/batch/cancel") async def batch_cancel(req: BatchRequest): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - return {"cancelled": await ScheduleTaskManager.batch_cancel(req.execution_ids)} + return {"cancelled": await TaskManager.batch_cancel(req.execution_ids)} @router.post("/task-executions/batch/delete") async def batch_delete(req: BatchRequest): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - return {"deleted": await ScheduleTaskManager.batch_delete(req.execution_ids)} + return {"deleted": await TaskManager.batch_delete(req.execution_ids)} @router.get("/task-executions/{execution_id}") async def get_execution(execution_id: str): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - execution = await ScheduleTaskManager.get_execution(execution_id) + execution = await TaskManager.get_execution(execution_id) if not execution: raise HTTPException(404, "Task execution not found") return execution.model_dump(mode="json", by_alias=True) @@ -448,9 +448,9 @@ async def get_execution(execution_id: str): @router.post("/task-executions/{execution_id}/viewed") async def mark_execution_viewed(execution_id: str): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - execution = await ScheduleTaskManager.mark_viewed(execution_id) + execution = await TaskManager.mark_viewed(execution_id) if not execution: raise HTTPException(404, "Task execution not found") return execution.model_dump(mode="json", by_alias=True) @@ -458,18 +458,18 @@ async def mark_execution_viewed(execution_id: str): @router.post("/task-executions/{execution_id}/cancel") async def cancel_execution(execution_id: str): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - execution = await ScheduleTaskManager.cancel_execution(execution_id) + execution = await TaskManager.cancel_execution(execution_id) if not execution: raise HTTPException(404, "Task execution not found") return execution.model_dump(mode="json", by_alias=True) @router.post("/task-executions/{execution_id}/retry") async def retry_execution(execution_id: str): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - execution = await ScheduleTaskManager.retry_execution(execution_id) + execution = await TaskManager.retry_execution(execution_id) if not execution: raise HTTPException(404, "Task execution not found") return execution.model_dump(mode="json", by_alias=True) @@ -477,9 +477,9 @@ async def retry_execution(execution_id: str): @router.post("/task-executions/{execution_id}/rerun") async def rerun_execution(execution_id: str): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - execution = await ScheduleTaskManager.rerun_execution(execution_id) + execution = await TaskManager.rerun_execution(execution_id) if not execution: raise HTTPException(404, "Task execution not found") return execution.model_dump(mode="json", by_alias=True) @@ -487,9 +487,9 @@ async def rerun_execution(execution_id: str): @router.delete("/task-executions/{execution_id}") async def delete_execution(execution_id: str): - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager - if not await ScheduleTaskManager.delete_execution(execution_id): + if not await TaskManager.delete_execution(execution_id): raise HTTPException(404, "Task execution not found") return {"ok": True} diff --git a/flocks/task/__init__.py b/flocks/task/__init__.py index 28b314739..180792a9d 100644 --- a/flocks/task/__init__.py +++ b/flocks/task/__init__.py @@ -20,7 +20,7 @@ SchedulerStatus, build_schedule, ) -from .schedule_task_manager import ScheduleTaskManager +from .manager import TaskManager from .store import TaskStore __all__ = [ @@ -30,7 +30,7 @@ "RetryConfig", "TaskExecution", "TaskExecutionQueueRef", - "ScheduleTaskManager", + "TaskManager", "TaskPriority", "TaskScheduler", "TaskTrigger", diff --git a/flocks/task/background.py b/flocks/task/background.py index 72781eec4..7e598f01f 100644 --- a/flocks/task/background.py +++ b/flocks/task/background.py @@ -237,27 +237,21 @@ async def _inject_parent_completion(self, task: BackgroundTask) -> None: "" ) try: - async def _persist_completion() -> None: - await Message.create( - session_id=task.parent_session_id, - role=MessageRole.USER, - content=content, - agent=task.parent_agent or "rex", - model=task.parent_model, - synthetic=True, - part_metadata={ - "kind": "background_task_result", - "task_id": task.id, - "session_id": task.session_id, - "status": state, - }, - ) - await self._update_parent_tool_part(task) - - await Session.run_active_write( - task.parent_session_id, - _persist_completion, + await Message.create( + session_id=task.parent_session_id, + role=MessageRole.USER, + content=content, + agent=task.parent_agent or "rex", + model=task.parent_model, + synthetic=True, + part_metadata={ + "kind": "background_task_result", + "task_id": task.id, + "session_id": task.session_id, + "status": state, + }, ) + await self._update_parent_tool_part(task) task.completion_injected = True self._schedule_parent_resume(task) except Exception as exc: @@ -271,6 +265,14 @@ def _schedule_parent_resume(self, task: BackgroundTask) -> None: """Kick the parent session so Rex consumes injected background results.""" if not task.parent_session_id: return + if task.status not in ("completed", "error"): + return + if SessionLoop.is_running(task.parent_session_id): + log.info("background.parent_resume.already_running", { + "task_id": task.id, + "parent_session_id": task.parent_session_id, + }) + return async def _run_parent() -> None: try: diff --git a/flocks/task/schedule_task_manager.py b/flocks/task/manager.py similarity index 99% rename from flocks/task/schedule_task_manager.py rename to flocks/task/manager.py index c769f4e6f..de038a25a 100644 --- a/flocks/task/schedule_task_manager.py +++ b/flocks/task/manager.py @@ -1,4 +1,4 @@ -"""Schedule task manager for the scheduler/execution domain.""" +"""Task Manager for scheduler/execution domain.""" import asyncio import json @@ -32,7 +32,7 @@ from .scheduler import TaskScheduler as SchedulerLoop from .store import TaskStore -log = Log.create(service="task.schedule_manager") +log = Log.create(service="task.manager") _TASK_EXPIRY_HOURS: int = 24 _CLEANUP_INTERVAL_S: int = 3600 @@ -49,8 +49,8 @@ class _TaskEventProps(_BaseModel): title: str -class ScheduleTaskManager: - _instance: Optional["ScheduleTaskManager"] = None +class TaskManager: + _instance: Optional["TaskManager"] = None _startup_error: Optional[str] = None def __init__( @@ -83,7 +83,7 @@ async def start( max_concurrent: int = 4, poll_interval: int = 5, scheduler_interval: int = 30, - ) -> "ScheduleTaskManager": + ) -> "TaskManager": if cls._instance and cls._instance._running: return cls._instance await TaskStore.init() @@ -128,7 +128,7 @@ async def stop(cls) -> None: log.info("manager.stopped") @classmethod - def get(cls) -> Optional["ScheduleTaskManager"]: + def get(cls) -> Optional["TaskManager"]: return cls._instance @classmethod diff --git a/flocks/task/plugin_sync.py b/flocks/task/plugin_sync.py index 1483ad538..1c27a582a 100644 --- a/flocks/task/plugin_sync.py +++ b/flocks/task/plugin_sync.py @@ -12,7 +12,7 @@ async def upsert_task_specs(specs: Sequence[TaskSpec]) -> int: - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.models import ( ExecutionMode, SchedulerMode, @@ -89,7 +89,7 @@ async def upsert_task_specs(specs: Sequence[TaskSpec]) -> int: log.warn("task.plugin.missing_cron", {"dedup_key": spec.dedup_key}) continue - await ScheduleTaskManager.create_scheduler( + await TaskManager.create_scheduler( title=spec.title, description=spec.description, mode=SchedulerMode.CRON, diff --git a/flocks/task/scheduler.py b/flocks/task/scheduler.py index 75acf6619..b8035dead 100644 --- a/flocks/task/scheduler.py +++ b/flocks/task/scheduler.py @@ -60,7 +60,7 @@ async def _loop(self) -> None: await asyncio.sleep(self._check_interval) async def _tick(self) -> None: - from .schedule_task_manager import ScheduleTaskManager + from .manager import TaskManager now = datetime.now(timezone.utc) schedulers = await TaskStore.list_due_schedulers() @@ -77,7 +77,7 @@ async def _tick(self) -> None: if scheduler.mode == SchedulerMode.ONCE else ExecutionTriggerType.SCHEDULED ) - await ScheduleTaskManager.create_execution_from_scheduler( + await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=trigger_type, enqueue=True, diff --git a/flocks/tool/task/schedule_task_center.py b/flocks/tool/task/schedule_task_center.py index 147624892..04c2c7d12 100644 --- a/flocks/tool/task/schedule_task_center.py +++ b/flocks/tool/task/schedule_task_center.py @@ -128,7 +128,7 @@ async def schedule_task_create( enabled: Optional[bool] = None, action: Optional[str] = None, ) -> ToolResult: - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.models import ( SchedulerMode, TaskPriority, @@ -178,7 +178,7 @@ async def schedule_task_create( user_prompt=user_prompt, ) - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title=title, description=description, mode=mode, @@ -187,7 +187,7 @@ async def schedule_task_create( trigger=trigger, ) if enabled is False: - scheduler = await ScheduleTaskManager.disable_scheduler(scheduler.id) or scheduler + scheduler = await TaskManager.disable_scheduler(scheduler.id) or scheduler display_tz = resolve_task_timezone_name(scheduler) output_lines = [ @@ -198,7 +198,7 @@ async def schedule_task_create( f"Priority: {scheduler.priority.value}", ] if scheduler.trigger.run_immediately: - executions, _ = await ScheduleTaskManager.list_scheduler_executions( + executions, _ = await TaskManager.list_scheduler_executions( scheduler.id, limit=1, ) @@ -238,7 +238,7 @@ async def schedule_task_list( type: Optional[str] = None, limit: int = 10, ) -> ToolResult: - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.models import SchedulerStatus, TaskStatus if type is not None and type not in _VALID_TYPES: @@ -286,7 +286,7 @@ async def schedule_task_list( scheduler_status = SchedulerStatus.ACTIVE elif status in ("disabled", "paused"): scheduler_status = SchedulerStatus.DISABLED - tasks, total = await ScheduleTaskManager.list_schedulers( + tasks, total = await TaskManager.list_schedulers( status=scheduler_status, scheduled_only=type != "scheduler", limit=limit, @@ -304,7 +304,7 @@ async def schedule_task_list( f"Valid values: {', '.join(s.value for s in TaskStatus)}." ), ) - tasks, total = await ScheduleTaskManager.list_executions( + tasks, total = await TaskManager.list_executions( status=task_status, limit=limit, ) @@ -322,22 +322,22 @@ async def schedule_task_status( task_id: str, resource_type: Optional[str] = None, ) -> ToolResult: - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager if resource_type == "scheduler": - task = await ScheduleTaskManager.get_scheduler(task_id) + task = await TaskManager.get_scheduler(task_id) elif resource_type == "execution": - task = await ScheduleTaskManager.get_execution(task_id) + task = await TaskManager.get_execution(task_id) else: - task = await ScheduleTaskManager.get_execution(task_id) + task = await TaskManager.get_execution(task_id) if task is None: - task = await ScheduleTaskManager.get_scheduler(task_id) + task = await TaskManager.get_scheduler(task_id) if ( resource_type != "scheduler" and task and getattr(getattr(task, "delivery_status", None), "value", None) == "unread" ): - await ScheduleTaskManager.mark_notified(task_id) + await TaskManager.mark_notified(task_id) if task is None: return ToolResult(success=False, error=f"Task {task_id} not found") @@ -363,7 +363,7 @@ async def schedule_task_update( user_prompt: Optional[str] = None, enabled: Optional[bool] = None, ) -> ToolResult: - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager from flocks.task.models import TaskPriority normalized_action = (action or "update").lower() @@ -373,13 +373,13 @@ async def schedule_task_update( normalized_action = "enable" if normalized_action == "cancel": - task = await ScheduleTaskManager.cancel_execution(task_id) + task = await TaskManager.cancel_execution(task_id) elif normalized_action == "retry": - task = await ScheduleTaskManager.retry_execution(task_id) + task = await TaskManager.retry_execution(task_id) elif normalized_action == "disable": - task = await ScheduleTaskManager.disable_scheduler(task_id) + task = await TaskManager.disable_scheduler(task_id) elif normalized_action == "enable": - task = await ScheduleTaskManager.enable_scheduler(task_id) + task = await TaskManager.enable_scheduler(task_id) elif normalized_action == "update": fields = {} if priority: @@ -389,7 +389,7 @@ async def schedule_task_update( if description is not None: fields["description"] = description try: - task = await ScheduleTaskManager.update_scheduler_with_trigger( + task = await TaskManager.update_scheduler_with_trigger( task_id, fields=fields, cron=cron, @@ -402,10 +402,10 @@ async def schedule_task_update( except ValueError as exc: return ToolResult(success=False, error=str(exc)) if enabled is False: - task = await ScheduleTaskManager.disable_scheduler(task_id) or task + task = await TaskManager.disable_scheduler(task_id) or task normalized_action = "disable" elif enabled is True: - task = await ScheduleTaskManager.enable_scheduler(task_id) or task + task = await TaskManager.enable_scheduler(task_id) or task normalized_action = "enable" else: return ToolResult(success=False, error=f"Unknown action: {action}") @@ -425,18 +425,18 @@ async def schedule_task_delete( task_id: str, resource_type: Optional[str] = None, ) -> ToolResult: - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager if resource_type == "execution": - ok = await ScheduleTaskManager.delete_execution(task_id) + ok = await TaskManager.delete_execution(task_id) elif resource_type == "scheduler": - ok = await ScheduleTaskManager.delete_scheduler(task_id) + ok = await TaskManager.delete_scheduler(task_id) else: - execution = await ScheduleTaskManager.get_execution(task_id) + execution = await TaskManager.get_execution(task_id) if execution is not None: - ok = await ScheduleTaskManager.delete_execution(task_id) + ok = await TaskManager.delete_execution(task_id) else: - ok = await ScheduleTaskManager.delete_scheduler(task_id) + ok = await TaskManager.delete_scheduler(task_id) if not ok: return ToolResult(success=False, error=f"Task {task_id} not found") return ToolResult(success=True, output=f"Task {task_id} deleted.") @@ -447,16 +447,16 @@ async def schedule_task_rerun( task_id: str, resource_type: Optional[str] = None, ) -> ToolResult: - from flocks.task.schedule_task_manager import ScheduleTaskManager + from flocks.task.manager import TaskManager if resource_type == "execution": - task = await ScheduleTaskManager.rerun_execution(task_id) + task = await TaskManager.rerun_execution(task_id) elif resource_type == "scheduler": - task = await ScheduleTaskManager.rerun_scheduler(task_id) + task = await TaskManager.rerun_scheduler(task_id) else: - task = await ScheduleTaskManager.rerun_execution(task_id) + task = await TaskManager.rerun_execution(task_id) if task is None: - task = await ScheduleTaskManager.rerun_scheduler(task_id) + task = await TaskManager.rerun_scheduler(task_id) if not task: return ToolResult(success=False, error=f"Task {task_id} not found") diff --git a/tests/integration/test_task_queue_integration.py b/tests/integration/test_task_queue_integration.py index 32c20d961..bde3ab7b0 100644 --- a/tests/integration/test_task_queue_integration.py +++ b/tests/integration/test_task_queue_integration.py @@ -8,12 +8,12 @@ import pytest -import flocks.task.schedule_task_manager as schedule_task_manager_module +import flocks.task.manager as task_manager_module from flocks.task.models import ExecutionMode from flocks.config.config import Config from flocks.storage.storage import Storage from flocks.task.executor import TaskExecutor -from flocks.task.schedule_task_manager import ScheduleTaskManager +from flocks.task.manager import TaskManager from flocks.task.models import ( ExecutionTriggerType, SchedulerMode, @@ -34,8 +34,8 @@ async def isolated_task_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): Config._cached_config = None Storage._db_path = None Storage._initialized = False - ScheduleTaskManager._instance = None - ScheduleTaskManager._startup_error = None + TaskManager._instance = None + TaskManager._startup_error = None TaskStore._initialized = False TaskStore._conn = None @@ -44,14 +44,14 @@ async def isolated_task_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): yield - await ScheduleTaskManager.stop() + await TaskManager.stop() await TaskStore.close() Config._global_config = None Config._cached_config = None Storage._db_path = None Storage._initialized = False - ScheduleTaskManager._instance = None - ScheduleTaskManager._startup_error = None + TaskManager._instance = None + TaskManager._startup_error = None TaskStore._initialized = False TaskStore._conn = None @@ -59,7 +59,7 @@ async def isolated_task_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): async def _wait_for_execution(execution_id: str, *, status: TaskStatus, timeout: float = 2.0): deadline = asyncio.get_running_loop().time() + timeout while asyncio.get_running_loop().time() < deadline: - execution = await ScheduleTaskManager.get_execution(execution_id) + execution = await TaskManager.get_execution(execution_id) if execution is not None and execution.status == status: return execution await asyncio.sleep(0.02) @@ -77,15 +77,15 @@ async def fake_dispatch(execution, scheduler): return await TaskStore.update_execution(execution) monkeypatch.setattr(TaskExecutor, "dispatch", fake_dispatch) - await ScheduleTaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) + await TaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="立即执行链路", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace"), ) - executions, total = await ScheduleTaskManager.list_scheduler_executions(scheduler.id) + executions, total = await TaskManager.list_scheduler_executions(scheduler.id) assert total == 1 completed = await _wait_for_execution(executions[0].id, status=TaskStatus.COMPLETED) @@ -108,18 +108,18 @@ async def fake_dispatch(execution, scheduler): return await TaskStore.update_execution(execution) monkeypatch.setattr(TaskExecutor, "dispatch", fake_dispatch) - await ScheduleTaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) + await TaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="重复执行", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace"), ) - first = (await ScheduleTaskManager.list_scheduler_executions(scheduler.id))[0][0] + first = (await TaskManager.list_scheduler_executions(scheduler.id))[0][0] first = await _wait_for_execution(first.id, status=TaskStatus.COMPLETED) - rerun = await ScheduleTaskManager.rerun_execution(first.id) + rerun = await TaskManager.rerun_execution(first.id) assert rerun is not None rerun = await _wait_for_execution(rerun.id, status=TaskStatus.COMPLETED) @@ -151,24 +151,24 @@ async def fake_dispatch(execution, scheduler): return await TaskStore.update_execution(execution) monkeypatch.setattr(TaskExecutor, "dispatch", fake_dispatch) - monkeypatch.setattr(schedule_task_manager_module, "_DISPATCH_GUARD_TIMEOUT_S", 0.05) - await ScheduleTaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) + monkeypatch.setattr(task_manager_module, "_DISPATCH_GUARD_TIMEOUT_S", 0.05) + await TaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) - first_scheduler = await ScheduleTaskManager.create_scheduler( + first_scheduler = await TaskManager.create_scheduler( title="超时任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace-1"), ) - second_scheduler = await ScheduleTaskManager.create_scheduler( + second_scheduler = await TaskManager.create_scheduler( title="后续任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace-2"), ) - first = (await ScheduleTaskManager.list_scheduler_executions(first_scheduler.id))[0][0] - second = (await ScheduleTaskManager.list_scheduler_executions(second_scheduler.id))[0][0] + first = (await TaskManager.list_scheduler_executions(first_scheduler.id))[0][0] + second = (await TaskManager.list_scheduler_executions(second_scheduler.id))[0][0] failed = await _wait_for_execution(first.id, status=TaskStatus.FAILED, timeout=1.0) completed = await _wait_for_execution(second.id, status=TaskStatus.COMPLETED, timeout=1.0) @@ -179,21 +179,21 @@ async def fake_dispatch(execution, scheduler): @pytest.mark.asyncio async def test_delete_scheduler_releases_claimed_queue_slot(tmp_path: Path): - await ScheduleTaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) + await TaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="删除释放槽位", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace-1"), ) - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) - manager = ScheduleTaskManager.get() + manager = TaskManager.get() assert manager is not None claimed = await manager.queue.dequeue() @@ -202,20 +202,20 @@ async def test_delete_scheduler_releases_claimed_queue_slot(tmp_path: Path): assert claimed.id == execution.id assert execution.id in manager.queue._running_ids - deleted = await ScheduleTaskManager.delete_scheduler(scheduler.id) + deleted = await TaskManager.delete_scheduler(scheduler.id) assert deleted is True assert execution.id not in manager.queue._running_ids - assert await ScheduleTaskManager.get_execution(execution.id) is None + assert await TaskManager.get_execution(execution.id) is None assert await TaskStore.get_queue_ref(execution.id) is None - next_scheduler = await ScheduleTaskManager.create_scheduler( + next_scheduler = await TaskManager.create_scheduler( title="后续可领取", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace-2"), ) - next_execution = await ScheduleTaskManager.create_execution_from_scheduler( + next_execution = await TaskManager.create_execution_from_scheduler( next_scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, @@ -236,21 +236,21 @@ async def fake_dispatch(_execution, _scheduler): raise asyncio.CancelledError() monkeypatch.setattr(TaskExecutor, "dispatch", fake_dispatch) - await ScheduleTaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) + await TaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="取消后回队", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) - manager = ScheduleTaskManager.get() + manager = TaskManager.get() assert manager is not None claimed = await manager.queue.dequeue() @@ -260,7 +260,7 @@ async def fake_dispatch(_execution, _scheduler): with pytest.raises(asyncio.CancelledError): await manager._run_execution(claimed) - refreshed = await ScheduleTaskManager.get_execution(execution.id) + refreshed = await TaskManager.get_execution(execution.id) queue_ref = await TaskStore.get_queue_ref(execution.id) assert refreshed is not None @@ -309,10 +309,10 @@ async def test_workflow_timeout_signals_cancel_and_stops_before_next_node( "read_workflow_from_fs", lambda workflow_id: {"workflowJson": workflow_json} if workflow_id == "wf_slow_cancel" else None, ) - monkeypatch.setattr(schedule_task_manager_module, "_DISPATCH_GUARD_TIMEOUT_S", 0.01) - await ScheduleTaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) + monkeypatch.setattr(task_manager_module, "_DISPATCH_GUARD_TIMEOUT_S", 0.01) + await TaskManager.start(max_concurrent=1, poll_interval=0.01, scheduler_interval=999) - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="可取消 workflow", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), @@ -321,7 +321,7 @@ async def test_workflow_timeout_signals_cancel_and_stops_before_next_node( workspace_directory=str(tmp_path / "workspace"), ) - execution = (await ScheduleTaskManager.list_scheduler_executions(scheduler.id))[0][0] + execution = (await TaskManager.list_scheduler_executions(scheduler.id))[0][0] failed = await _wait_for_execution(execution.id, status=TaskStatus.FAILED, timeout=1.0) await asyncio.sleep(0.15) @@ -466,8 +466,8 @@ async def test_standalone_legacy_migration_script_migrates_existing_tables(tmp_p stdout, stderr = await proc.communicate() assert proc.returncode == 0, (stdout or b"").decode() + (stderr or b"").decode() - scheduler = await ScheduleTaskManager.get_scheduler("task_legacy_1") - execution = await ScheduleTaskManager.get_execution("texec_legacy_1") + scheduler = await TaskManager.get_scheduler("task_legacy_1") + execution = await TaskManager.get_execution("texec_legacy_1") assert scheduler is not None assert scheduler.mode == SchedulerMode.ONCE @@ -475,7 +475,7 @@ async def test_standalone_legacy_migration_script_migrates_existing_tables(tmp_p assert execution.scheduler_id == "task_legacy_1" assert execution.status == TaskStatus.COMPLETED assert execution.session_id == "ses_123" - assert ScheduleTaskManager._legacy_tables_exist() is False + assert TaskManager._legacy_tables_exist() is False assert state_path.exists() is False @@ -581,8 +581,8 @@ async def test_standalone_legacy_migration_preserves_paused_scheduled_task_histo stdout, stderr = await proc.communicate() assert proc.returncode == 0, (stdout or b"").decode() + (stderr or b"").decode() - scheduler = await ScheduleTaskManager.get_scheduler("task_paused_sched") - execution = await ScheduleTaskManager.get_execution("legacy_exec_task_paused_sched") + scheduler = await TaskManager.get_scheduler("task_paused_sched") + execution = await TaskManager.get_execution("legacy_exec_task_paused_sched") assert scheduler is not None assert scheduler.status in (SchedulerStatus.ACTIVE, SchedulerStatus.DISABLED) diff --git a/tests/server/routes/test_task_scheduler_context_route.py b/tests/server/routes/test_task_scheduler_context_route.py index e31c85214..8df880d77 100644 --- a/tests/server/routes/test_task_scheduler_context_route.py +++ b/tests/server/routes/test_task_scheduler_context_route.py @@ -2,13 +2,13 @@ import pytest -from flocks.task.schedule_task_manager import ScheduleTaskManager +from flocks.task.manager import TaskManager from flocks.task.models import ExecutionMode, ExecutionTriggerType, SchedulerMode, TaskTrigger @pytest.mark.asyncio async def test_update_scheduler_accepts_context_for_workflow_inputs(client): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="工作流定时任务", mode=SchedulerMode.CRON, trigger=TaskTrigger(cron="0 9 * * *", timezone="Asia/Shanghai"), @@ -28,11 +28,11 @@ async def test_update_scheduler_accepts_context_for_workflow_inputs(client): assert response.status_code == 200 assert response.json()["context"] == {"keyword": "after", "limit": 5} - updated = await ScheduleTaskManager.get_scheduler(scheduler.id) + updated = await TaskManager.get_scheduler(scheduler.id) assert updated is not None assert updated.context == {"keyword": "after", "limit": 5} - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( updated, trigger_type=ExecutionTriggerType.SCHEDULED, enqueue=False, diff --git a/tests/server/test_lifespan.py b/tests/server/test_lifespan.py index 9e34b9fc4..ab2b4fa60 100644 --- a/tests/server/test_lifespan.py +++ b/tests/server/test_lifespan.py @@ -93,8 +93,8 @@ async def fake_async_noop(*_args, **_kwargs) -> None: ) monkeypatch.setitem( sys.modules, - "flocks.task.schedule_task_manager", - types.SimpleNamespace(ScheduleTaskManager=types.SimpleNamespace(start=fake_async_noop, stop=fake_async_noop)), + "flocks.task.manager", + types.SimpleNamespace(TaskManager=types.SimpleNamespace(start=fake_async_noop, stop=fake_async_noop)), ) monkeypatch.setitem( sys.modules, diff --git a/tests/server/test_server.py b/tests/server/test_server.py index 355849931..eb79f07c5 100644 --- a/tests/server/test_server.py +++ b/tests/server/test_server.py @@ -9,7 +9,7 @@ from fastapi import status from flocks.server.app import app -from flocks.task.schedule_task_manager import ScheduleTaskManager +from flocks.task.manager import TaskManager from flocks.task.store import TaskStore from flocks.task.models import ( DeliveryStatus, @@ -133,12 +133,12 @@ async def test_list_executions_invalid_priority_returns_422(client): @pytest.mark.asyncio async def test_task_schedulers_scheduled_only_excludes_immediate_queue_templates(client): - await ScheduleTaskManager.create_scheduler( + await TaskManager.create_scheduler( title="立即任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), ) - scheduled = await ScheduleTaskManager.create_scheduler( + scheduled = await TaskManager.create_scheduler( title="单次计划", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), @@ -155,12 +155,12 @@ async def test_task_schedulers_scheduled_only_excludes_immediate_queue_templates @pytest.mark.asyncio async def test_task_scheduler_list_accepts_legacy_paused_status_query(client): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="兼容旧 paused 调度查询", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), ) - await ScheduleTaskManager.disable_scheduler(scheduler.id) + await TaskManager.disable_scheduler(scheduler.id) response = await client.get("/api/task-schedulers", params={"status": "paused"}) @@ -171,13 +171,13 @@ async def test_task_scheduler_list_accepts_legacy_paused_status_query(client): @pytest.mark.asyncio async def test_task_schedulers_list_excludes_archived_builtin_after_delete(client): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="内置计划任务", mode=SchedulerMode.CRON, trigger=TaskTrigger(cron="*/5 * * * *", timezone="Asia/Shanghai"), dedup_key="builtin:test-scheduled-task", ) - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.SCHEDULED, enqueue=True, @@ -192,8 +192,8 @@ async def test_task_schedulers_list_excludes_archived_builtin_after_delete(clien ids = {item["id"] for item in data["items"]} assert scheduler.id not in ids - archived = await ScheduleTaskManager.get_scheduler(scheduler.id) - cancelled_execution = await ScheduleTaskManager.get_execution(execution.id) + archived = await TaskManager.get_scheduler(scheduler.id) + cancelled_execution = await TaskManager.get_execution(execution.id) assert archived is not None assert archived.status == SchedulerStatus.ARCHIVED assert cancelled_execution is not None @@ -202,13 +202,13 @@ async def test_task_schedulers_list_excludes_archived_builtin_after_delete(clien @pytest.mark.asyncio async def test_delete_scheduler_cleans_queue_state_for_non_builtin(client): - await ScheduleTaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await ScheduleTaskManager.create_scheduler( + await TaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) + scheduler = await TaskManager.create_scheduler( title="普通计划任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), ) - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, @@ -234,12 +234,12 @@ async def test_delete_scheduler_cleans_queue_state_for_non_builtin(client): @pytest.mark.asyncio async def test_task_execution_list_accepts_legacy_paused_status_query(client): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="兼容旧 paused 执行查询", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), ) - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -257,17 +257,17 @@ async def test_task_execution_list_accepts_legacy_paused_status_query(client): @pytest.mark.asyncio async def test_batch_cancel_endpoint_cancels_selected_executions(client): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="批量取消接口", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), ) - cancellable = await ScheduleTaskManager.create_execution_from_scheduler( + cancellable = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) - completed = await ScheduleTaskManager.create_execution_from_scheduler( + completed = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -284,19 +284,19 @@ async def test_batch_cancel_endpoint_cancels_selected_executions(client): assert response.status_code == status.HTTP_200_OK assert response.json()["cancelled"] == 1 - cancelled_execution = await ScheduleTaskManager.get_execution(cancellable.id) + cancelled_execution = await TaskManager.get_execution(cancellable.id) assert cancelled_execution is not None assert cancelled_execution.status == TaskStatus.CANCELLED @pytest.mark.asyncio async def test_execution_pause_and_resume_endpoints_are_removed(client): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="旧暂停接口", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), ) - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, @@ -311,12 +311,12 @@ async def test_execution_pause_and_resume_endpoints_are_removed(client): @pytest.mark.asyncio async def test_mark_execution_viewed_endpoint_updates_delivery_status(client): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="标记已读", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), ) - execution = (await ScheduleTaskManager.list_scheduler_executions(scheduler.id, limit=1))[0][0] + execution = (await TaskManager.list_scheduler_executions(scheduler.id, limit=1))[0][0] execution.status = TaskStatus.COMPLETED execution.delivery_status = DeliveryStatus.UNREAD await TaskStore.update_execution(execution) diff --git a/tests/storage/test_sqlite_connection_config.py b/tests/storage/test_sqlite_connection_config.py index eb1aa1725..19276d496 100644 --- a/tests/storage/test_sqlite_connection_config.py +++ b/tests/storage/test_sqlite_connection_config.py @@ -5,7 +5,7 @@ from flocks.config.config import Config from flocks.storage.storage import Storage -from flocks.task.schedule_task_manager import ScheduleTaskManager +from flocks.task.manager import TaskManager @pytest.fixture(autouse=True) @@ -73,10 +73,10 @@ def test_storage_connect_sync_applies_runtime_sqlite_pragmas() -> None: @pytest.mark.asyncio -async def test_schedule_task_manager_sync_connection_uses_storage_sqlite_contract() -> None: +async def test_task_manager_sync_connection_uses_storage_sqlite_contract() -> None: await Storage.init() - with ScheduleTaskManager._with_db_connection() as db: + with TaskManager._with_db_connection() as db: row = db.execute( "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'storage'" ).fetchone() diff --git a/tests/task/test_task.py b/tests/task/test_task.py index b1cb53520..f2b31aad7 100644 --- a/tests/task/test_task.py +++ b/tests/task/test_task.py @@ -10,7 +10,7 @@ from flocks.cli.commands import task as task_cli_commands import flocks.task.background as background_module -import flocks.task.schedule_task_manager as schedule_task_manager_module +import flocks.task.manager as task_manager_module import flocks.task.plugin_sync as plugin_sync_module from flocks.server.routes import question as question_routes from flocks.config.config import Config @@ -19,7 +19,7 @@ from flocks.task.background import BackgroundManager, BackgroundTask, LaunchInput from flocks.task.executor import TaskExecutor from flocks.task.formatting import format_task_datetime -from flocks.task.schedule_task_manager import ScheduleTaskManager +from flocks.task.manager import TaskManager from flocks.task.models import ( DeliveryStatus, ExecutionMode, @@ -49,8 +49,8 @@ async def isolated_task_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): Config._cached_config = None Storage._db_path = None Storage._initialized = False - ScheduleTaskManager._instance = None - ScheduleTaskManager._startup_error = None + TaskManager._instance = None + TaskManager._startup_error = None TaskStore._initialized = False TaskStore._conn = None @@ -59,21 +59,21 @@ async def isolated_task_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): yield - await ScheduleTaskManager.stop() + await TaskManager.stop() await TaskStore.close() Config._global_config = None Config._cached_config = None Storage._db_path = None Storage._initialized = False - ScheduleTaskManager._instance = None - ScheduleTaskManager._startup_error = None + TaskManager._instance = None + TaskManager._startup_error = None TaskStore._initialized = False TaskStore._conn = None @pytest.mark.asyncio async def test_immediate_scheduler_creates_single_queued_execution(tmp_path: Path): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="立即执行", description="创建后立刻入队", mode=SchedulerMode.ONCE, @@ -82,7 +82,7 @@ async def test_immediate_scheduler_creates_single_queued_execution(tmp_path: Pat workspace_directory=str(tmp_path / "workspace"), ) - executions, total = await ScheduleTaskManager.list_scheduler_executions(scheduler.id) + executions, total = await TaskManager.list_scheduler_executions(scheduler.id) assert total == 1 execution = executions[0] @@ -99,7 +99,7 @@ async def test_list_executions_tolerates_legacy_non_utf8_description( tmp_path: Path, ): legacy_description = "扫描 Windows 主机 192.168.254.1" - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="legacy encoding", description="placeholder", mode=SchedulerMode.ONCE, @@ -107,7 +107,7 @@ async def test_list_executions_tolerates_legacy_non_utf8_description( trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace"), ) - executions, _ = await ScheduleTaskManager.list_scheduler_executions(scheduler.id) + executions, _ = await TaskManager.list_scheduler_executions(scheduler.id) execution_id = executions[0].id db = await TaskStore.raw_db() @@ -121,7 +121,7 @@ async def test_list_executions_tolerates_legacy_non_utf8_description( ) await db.commit() - executions, total = await ScheduleTaskManager.list_executions(limit=20) + executions, total = await TaskManager.list_executions(limit=20) assert total == 1 assert executions[0].description == legacy_description @@ -129,7 +129,7 @@ async def test_list_executions_tolerates_legacy_non_utf8_description( @pytest.mark.asyncio async def test_once_scheduler_tick_creates_execution_and_disables_scheduler(tmp_path: Path): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="单次定时", mode=SchedulerMode.ONCE, trigger=TaskTrigger( @@ -142,8 +142,8 @@ async def test_once_scheduler_tick_creates_execution_and_disables_scheduler(tmp_ loop = SchedulerLoop() await loop._tick() - updated = await ScheduleTaskManager.get_scheduler(scheduler.id) - executions, total = await ScheduleTaskManager.list_scheduler_executions(scheduler.id) + updated = await TaskManager.get_scheduler(scheduler.id) + executions, total = await TaskManager.list_scheduler_executions(scheduler.id) assert updated is not None assert updated.status == SchedulerStatus.DISABLED @@ -154,7 +154,7 @@ async def test_once_scheduler_tick_creates_execution_and_disables_scheduler(tmp_ @pytest.mark.asyncio async def test_cron_scheduler_does_not_spawn_second_active_execution(tmp_path: Path): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="循环任务", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -172,7 +172,7 @@ async def test_cron_scheduler_does_not_spawn_second_active_execution(tmp_path: P await loop._tick() await loop._tick() - executions, total = await ScheduleTaskManager.list_scheduler_executions(scheduler.id, limit=10) + executions, total = await TaskManager.list_scheduler_executions(scheduler.id, limit=10) assert total == 1 assert executions[0].status == TaskStatus.QUEUED @@ -182,7 +182,7 @@ async def test_cron_scheduler_does_not_spawn_second_active_execution(tmp_path: P @pytest.mark.asyncio async def test_create_scheduler_rejects_six_field_cron(tmp_path: Path): with pytest.raises(ValueError, match="only 5-field cron is supported"): - await ScheduleTaskManager.create_scheduler( + await TaskManager.create_scheduler( title="非法 cron", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -196,7 +196,7 @@ async def test_create_scheduler_rejects_six_field_cron(tmp_path: Path): @pytest.mark.asyncio async def test_create_scheduler_rejects_out_of_range_five_field_cron(tmp_path: Path): with pytest.raises(ValueError, match="not a valid 5-field cron"): - await ScheduleTaskManager.create_scheduler( + await TaskManager.create_scheduler( title="越界 cron", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -214,7 +214,7 @@ async def test_create_scheduler_does_not_mutate_trigger_argument(tmp_path: Path) timezone="Asia/Shanghai", ) - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="不修改调用方 trigger", mode=SchedulerMode.CRON, trigger=trigger, @@ -276,7 +276,7 @@ async def test_plugin_sync_skips_new_scheduler_with_invalid_six_field_cron( async def test_plugin_sync_does_not_overwrite_existing_scheduler_with_invalid_six_field_cron( tmp_path: Path, ): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="原始内置任务", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -299,7 +299,7 @@ async def test_plugin_sync_does_not_overwrite_existing_scheduler_with_invalid_si ] ) - unchanged = await ScheduleTaskManager.get_scheduler(scheduler.id) + unchanged = await TaskManager.get_scheduler(scheduler.id) assert created == 0 assert unchanged is not None @@ -312,7 +312,7 @@ async def test_plugin_sync_does_not_overwrite_existing_scheduler_with_invalid_si async def test_update_scheduler_with_trigger_run_once_preserves_existing_cron_when_omitted( tmp_path: Path, ): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="切成单次任务时保留旧 cron", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -322,7 +322,7 @@ async def test_update_scheduler_with_trigger_run_once_preserves_existing_cron_wh workspace_directory=str(tmp_path / "workspace"), ) - updated = await ScheduleTaskManager.update_scheduler_with_trigger( + updated = await TaskManager.update_scheduler_with_trigger( scheduler.id, fields={}, run_once=True, @@ -338,8 +338,8 @@ async def test_update_scheduler_with_trigger_run_once_preserves_existing_cron_wh @pytest.mark.asyncio async def test_recover_stale_running_execution_unblocks_scheduler(tmp_path: Path): - manager = ScheduleTaskManager(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await ScheduleTaskManager.create_scheduler( + manager = TaskManager(max_concurrent=1, poll_interval=999, scheduler_interval=999) + scheduler = await TaskManager.create_scheduler( title="恢复阻塞循环任务", mode=SchedulerMode.CRON, trigger=TaskTrigger(cron="*/5 * * * *", timezone="Asia/Shanghai"), @@ -348,13 +348,13 @@ async def test_recover_stale_running_execution_unblocks_scheduler(tmp_path: Path scheduler.trigger.next_run = datetime.now(timezone.utc) - timedelta(seconds=1) await TaskStore.update_scheduler(scheduler) - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.SCHEDULED, enqueue=False, ) started_at = datetime.now(timezone.utc) - timedelta( - seconds=schedule_task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 5 + seconds=task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 5 ) execution.status = TaskStatus.RUNNING execution.started_at = started_at @@ -366,7 +366,7 @@ async def test_recover_stale_running_execution_unblocks_scheduler(tmp_path: Path recovered = await manager._recover_stale_active_executions() assert recovered == 1 - failed = await ScheduleTaskManager.get_execution(execution.id) + failed = await TaskManager.get_execution(execution.id) assert failed is not None assert failed.status == TaskStatus.FAILED assert "recovery threshold" in (failed.error or "") @@ -374,7 +374,7 @@ async def test_recover_stale_running_execution_unblocks_scheduler(tmp_path: Path loop = SchedulerLoop() await loop._tick() - executions, total = await ScheduleTaskManager.list_scheduler_executions(scheduler.id, limit=10) + executions, total = await TaskManager.list_scheduler_executions(scheduler.id, limit=10) assert total == 2 assert any(item.id == execution.id and item.status == TaskStatus.FAILED for item in executions) @@ -383,29 +383,29 @@ async def test_recover_stale_running_execution_unblocks_scheduler(tmp_path: Path @pytest.mark.asyncio async def test_recover_orphaned_queued_execution_restores_queue_ref(tmp_path: Path): - await ScheduleTaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) + await TaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) # Pause the execution loop so it cannot race the test by claiming the # execution before we manually simulate the orphan state (queued row # with no queue ref). - ScheduleTaskManager.pause_queue() - scheduler = await ScheduleTaskManager.create_scheduler( + TaskManager.pause_queue() + scheduler = await TaskManager.create_scheduler( title="恢复孤儿排队任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) await TaskStore.finish_queue_ref(execution.id) - manager = ScheduleTaskManager.get() + manager = TaskManager.get() assert manager is not None recovered = await manager._recover_orphaned_queued_executions() - refreshed = await ScheduleTaskManager.get_execution(execution.id) + refreshed = await TaskManager.get_execution(execution.id) assert recovered == 1 assert refreshed is not None @@ -420,54 +420,54 @@ async def test_recover_orphaned_queued_execution_restores_queue_ref(tmp_path: Pa @pytest.mark.asyncio async def test_queue_status_reports_stale_running_execution(tmp_path: Path): - await ScheduleTaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await ScheduleTaskManager.create_scheduler( + await TaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) + scheduler = await TaskManager.create_scheduler( title="阻塞诊断", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, ) started_at = datetime.now(timezone.utc) - timedelta( - seconds=schedule_task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 15 + seconds=task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 15 ) execution.status = TaskStatus.RUNNING execution.started_at = started_at execution.queued_at = started_at await TaskStore.update_execution(execution) - status = await ScheduleTaskManager.queue_status() + status = await TaskManager.queue_status() assert status["stale_running"] == 1 assert isinstance(status["oldest_running_seconds"], int) - assert status["oldest_running_seconds"] >= schedule_task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + assert status["oldest_running_seconds"] >= task_manager_module._RUNNING_RECOVERY_TIMEOUT_S @pytest.mark.asyncio async def test_queue_status_uses_largest_elapsed_running_time(tmp_path: Path): - await ScheduleTaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await ScheduleTaskManager.create_scheduler( + await TaskManager.start(max_concurrent=1, poll_interval=999, scheduler_interval=999) + scheduler = await TaskManager.create_scheduler( title="阻塞诊断-多任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - older_execution = await ScheduleTaskManager.create_execution_from_scheduler( + older_execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, ) - newer_execution = await ScheduleTaskManager.create_execution_from_scheduler( + newer_execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, ) older_started_at = datetime.now(timezone.utc) - timedelta( - seconds=schedule_task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 15 + seconds=task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 15 ) newer_started_at = datetime.now(timezone.utc) - timedelta(seconds=10) older_execution.status = TaskStatus.RUNNING @@ -479,11 +479,11 @@ async def test_queue_status_uses_largest_elapsed_running_time(tmp_path: Path): await TaskStore.update_execution(older_execution) await TaskStore.update_execution(newer_execution) - status = await ScheduleTaskManager.queue_status() + status = await TaskManager.queue_status() assert status["stale_running"] == 1 assert isinstance(status["oldest_running_seconds"], int) - assert status["oldest_running_seconds"] >= schedule_task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 15 + assert status["oldest_running_seconds"] >= task_manager_module._RUNNING_RECOVERY_TIMEOUT_S + 15 @pytest.mark.asyncio @@ -624,12 +624,6 @@ async def test_background_task_completion_injects_parent_context( monkeypatch.setattr(background_module.Message, "update_part", update_part) monkeypatch.setattr(background_module.SessionLoop, "run", parent_loop_run) - async def run_active_write(_session_id, operation, **_kwargs): - return await operation() - - active_write = AsyncMock(side_effect=run_active_write) - monkeypatch.setattr(background_module.Session, "run_active_write", active_write) - manager = BackgroundManager() task = BackgroundTask( id="bg_parent_inject", @@ -648,8 +642,6 @@ async def run_active_write(_session_id, operation, **_kwargs): await manager._inject_parent_completion(task) - active_write.assert_awaited_once() - assert active_write.await_args.args[0] == "ses-parent" create_message.assert_awaited_once() kwargs = create_message.await_args.kwargs assert kwargs["session_id"] == "ses-parent" @@ -670,11 +662,12 @@ async def run_active_write(_session_id, operation, **_kwargs): @pytest.mark.asyncio -async def test_background_task_completion_always_attempts_parent_resume( +async def test_background_task_completion_does_not_resume_running_parent( monkeypatch: pytest.MonkeyPatch, ): parent_loop_run = AsyncMock(return_value=SimpleNamespace(action="stop")) monkeypatch.setattr(background_module.SessionLoop, "run", parent_loop_run) + monkeypatch.setattr(background_module.SessionLoop, "is_running", lambda _session_id: True) manager = BackgroundManager() task = BackgroundTask( @@ -692,8 +685,7 @@ async def test_background_task_completion_always_attempts_parent_resume( manager._schedule_parent_resume(task) await asyncio.sleep(0) - parent_loop_run.assert_awaited_once() - assert parent_loop_run.await_args.kwargs["session_id"] == "ses-parent" + parent_loop_run.assert_not_awaited() @pytest.mark.asyncio @@ -747,15 +739,15 @@ def fake_run_workflow(*, workflow, inputs, **_kwargs): @pytest.mark.asyncio async def test_retry_queue_requeues_failed_execution(tmp_path: Path): - manager = ScheduleTaskManager(max_concurrent=1, poll_interval=999, scheduler_interval=999) + manager = TaskManager(max_concurrent=1, poll_interval=999, scheduler_interval=999) - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="失败重试", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -769,7 +761,7 @@ async def test_retry_queue_requeues_failed_execution(tmp_path: Path): await manager._process_retry_queue() - reloaded = await ScheduleTaskManager.get_execution(execution.id) + reloaded = await TaskManager.get_execution(execution.id) assert reloaded is not None assert reloaded.status == TaskStatus.QUEUED assert reloaded.retry.retry_after is None @@ -781,18 +773,18 @@ async def test_retry_queue_requeues_failed_execution(tmp_path: Path): @pytest.mark.asyncio async def test_queue_dequeue_respects_claimed_slots_before_running_status(tmp_path: Path): queue = TaskQueue(max_concurrent=1) - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="并发控制", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - first = await ScheduleTaskManager.create_execution_from_scheduler( + first = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) - second = await ScheduleTaskManager.create_execution_from_scheduler( + second = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, @@ -814,14 +806,14 @@ async def test_queue_dequeue_respects_claimed_slots_before_running_status(tmp_pa @pytest.mark.asyncio async def test_immediate_scheduler_dedup_does_not_create_duplicate_execution(tmp_path: Path): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="去重立即执行", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace"), dedup_key="dup-immediate", ) - duplicate = await ScheduleTaskManager.create_scheduler( + duplicate = await TaskManager.create_scheduler( title="去重立即执行", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), @@ -829,7 +821,7 @@ async def test_immediate_scheduler_dedup_does_not_create_duplicate_execution(tmp dedup_key="dup-immediate", ) - executions, total = await ScheduleTaskManager.list_scheduler_executions(scheduler.id, limit=10) + executions, total = await TaskManager.list_scheduler_executions(scheduler.id, limit=10) assert duplicate.id == scheduler.id assert total == 1 @@ -838,18 +830,18 @@ async def test_immediate_scheduler_dedup_does_not_create_duplicate_execution(tmp @pytest.mark.asyncio async def test_batch_cancel_counts_only_actual_cancellations(tmp_path: Path): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="批量取消", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - cancellable = await ScheduleTaskManager.create_execution_from_scheduler( + cancellable = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) - completed = await ScheduleTaskManager.create_execution_from_scheduler( + completed = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -858,7 +850,7 @@ async def test_batch_cancel_counts_only_actual_cancellations(tmp_path: Path): completed.completed_at = datetime.now(timezone.utc) await TaskStore.update_execution(completed) - cancelled = await ScheduleTaskManager.batch_cancel([cancellable.id, completed.id]) + cancelled = await TaskManager.batch_cancel([cancellable.id, completed.id]) assert cancelled == 1 @@ -874,28 +866,28 @@ async def fake_cancel_runtime(_cls, execution): cancelled_runtime_ids.append(execution.id) monkeypatch.setattr( - ScheduleTaskManager, + TaskManager, "_cancel_execution_runtime", classmethod(fake_cancel_runtime), ) - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="删除前清理普通计划", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - pending = await ScheduleTaskManager.create_execution_from_scheduler( + pending = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, ) - queued = await ScheduleTaskManager.create_execution_from_scheduler( + queued = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, ) - running = await ScheduleTaskManager.create_execution_from_scheduler( + running = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -906,7 +898,7 @@ async def fake_cancel_runtime(_cls, execution): running.session_id = "ses_delete_running" await TaskStore.update_execution(running) await TaskStore.enqueue_execution_ref(running.id) - completed = await ScheduleTaskManager.create_execution_from_scheduler( + completed = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -915,7 +907,7 @@ async def fake_cancel_runtime(_cls, execution): completed.completed_at = datetime.now(timezone.utc) await TaskStore.update_execution(completed) - deleted = await ScheduleTaskManager.delete_scheduler(scheduler.id) + deleted = await TaskManager.delete_scheduler(scheduler.id) assert deleted is True assert set(cancelled_runtime_ids) == { @@ -924,7 +916,7 @@ async def fake_cancel_runtime(_cls, execution): running.id, } for execution_id in (pending.id, queued.id, running.id, completed.id): - assert await ScheduleTaskManager.get_execution(execution_id) is None + assert await TaskManager.get_execution(execution_id) is None assert await TaskStore.get_queue_ref(queued.id) is None assert await TaskStore.get_queue_ref(running.id) is None @@ -946,13 +938,13 @@ async def fake_cancel_runtime(_cls, execution): @pytest.mark.asyncio async def test_store_init_normalizes_legacy_paused_execution(tmp_path: Path): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="兼容旧 paused 状态", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=True, @@ -974,7 +966,7 @@ async def test_store_init_normalizes_legacy_paused_execution(tmp_path: Path): await TaskStore.close() await TaskStore.init() - normalized = await ScheduleTaskManager.get_execution(execution.id) + normalized = await TaskManager.get_execution(execution.id) assert normalized is not None assert normalized.status == TaskStatus.CANCELLED @@ -985,12 +977,12 @@ async def test_store_init_normalizes_legacy_paused_execution(tmp_path: Path): @pytest.mark.asyncio async def test_cli_list_tasks_accepts_legacy_paused_status(monkeypatch: pytest.MonkeyPatch): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="CLI 兼容 paused 查询", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), ) - execution = await ScheduleTaskManager.create_execution_from_scheduler( + execution = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -998,7 +990,7 @@ async def test_cli_list_tasks_accepts_legacy_paused_status(monkeypatch: pytest.M execution.status = TaskStatus.CANCELLED execution.completed_at = datetime.now(timezone.utc) await TaskStore.update_execution(execution) - await ScheduleTaskManager.disable_scheduler(scheduler.id) + await TaskManager.disable_scheduler(scheduler.id) printed: list[object] = [] monkeypatch.setattr(task_cli_commands.console, "print", lambda *args, **kwargs: printed.append(args)) @@ -1057,24 +1049,24 @@ async def fake_cancel_runtime(_cls, execution): cancelled_runtime_ids.append(execution.id) monkeypatch.setattr( - ScheduleTaskManager, + TaskManager, "_cancel_execution_runtime", classmethod(fake_cancel_runtime), ) - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="删除前清理内置计划", mode=SchedulerMode.CRON, trigger=TaskTrigger(cron="*/5 * * * *", timezone="Asia/Shanghai"), workspace_directory=str(tmp_path / "workspace"), dedup_key="builtin:test-delete-cleanup", ) - queued = await ScheduleTaskManager.create_execution_from_scheduler( + queued = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.SCHEDULED, enqueue=True, ) - running = await ScheduleTaskManager.create_execution_from_scheduler( + running = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.SCHEDULED, enqueue=False, @@ -1085,7 +1077,7 @@ async def fake_cancel_runtime(_cls, execution): running.session_id = "ses_builtin_running" await TaskStore.update_execution(running) await TaskStore.enqueue_execution_ref(running.id) - completed = await ScheduleTaskManager.create_execution_from_scheduler( + completed = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.SCHEDULED, enqueue=False, @@ -1094,11 +1086,11 @@ async def fake_cancel_runtime(_cls, execution): completed.completed_at = datetime.now(timezone.utc) await TaskStore.update_execution(completed) - deleted = await ScheduleTaskManager.delete_scheduler(scheduler.id) - archived = await ScheduleTaskManager.get_scheduler(scheduler.id) - queued_execution = await ScheduleTaskManager.get_execution(queued.id) - running_execution = await ScheduleTaskManager.get_execution(running.id) - completed_execution = await ScheduleTaskManager.get_execution(completed.id) + deleted = await TaskManager.delete_scheduler(scheduler.id) + archived = await TaskManager.get_scheduler(scheduler.id) + queued_execution = await TaskManager.get_execution(queued.id) + running_execution = await TaskManager.get_execution(running.id) + completed_execution = await TaskManager.get_execution(completed.id) assert deleted is True assert archived is not None @@ -1116,39 +1108,39 @@ async def fake_cancel_runtime(_cls, execution): @pytest.mark.asyncio async def test_dashboard_counts_exclude_immediate_once_schedulers(tmp_path: Path): - await ScheduleTaskManager.create_scheduler( + await TaskManager.create_scheduler( title="队列任务模板", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=True), workspace_directory=str(tmp_path / "workspace-1"), ) - await ScheduleTaskManager.create_scheduler( + await TaskManager.create_scheduler( title="单次计划任务", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False, run_at=datetime.now(timezone.utc) + timedelta(hours=1)), workspace_directory=str(tmp_path / "workspace-2"), ) - await ScheduleTaskManager.create_scheduler( + await TaskManager.create_scheduler( title="循环计划任务", mode=SchedulerMode.CRON, trigger=TaskTrigger(cron="*/5 * * * *", timezone="Asia/Shanghai"), workspace_directory=str(tmp_path / "workspace-3"), ) - counts = await ScheduleTaskManager.dashboard() + counts = await TaskManager.dashboard() assert counts["scheduled_active"] == 2 @pytest.mark.asyncio async def test_get_unviewed_results_includes_unread_and_notified_only(tmp_path: Path): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="未读结果", mode=SchedulerMode.ONCE, trigger=TaskTrigger(run_immediately=False), workspace_directory=str(tmp_path / "workspace"), ) - unread = await ScheduleTaskManager.create_execution_from_scheduler( + unread = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -1158,7 +1150,7 @@ async def test_get_unviewed_results_includes_unread_and_notified_only(tmp_path: unread.delivery_status = DeliveryStatus.UNREAD await TaskStore.update_execution(unread) - notified = await ScheduleTaskManager.create_execution_from_scheduler( + notified = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -1168,7 +1160,7 @@ async def test_get_unviewed_results_includes_unread_and_notified_only(tmp_path: notified.delivery_status = DeliveryStatus.NOTIFIED await TaskStore.update_execution(notified) - viewed = await ScheduleTaskManager.create_execution_from_scheduler( + viewed = await TaskManager.create_execution_from_scheduler( scheduler, trigger_type=ExecutionTriggerType.RUN_ONCE, enqueue=False, @@ -1178,7 +1170,7 @@ async def test_get_unviewed_results_includes_unread_and_notified_only(tmp_path: viewed.delivery_status = DeliveryStatus.VIEWED await TaskStore.update_execution(viewed) - results = await ScheduleTaskManager.get_unviewed_results() + results = await TaskManager.get_unviewed_results() result_ids = {item.id for item in results} assert unread.id in result_ids @@ -1191,11 +1183,11 @@ async def test_task_page_notice_drops_legacy_tables_after_third_display(): db = await TaskStore.raw_db() await db.execute("CREATE TABLE tasks (id TEXT PRIMARY KEY)") await db.commit() - ScheduleTaskManager._write_migration_state({"failed": True, "notice_count": 0}) + TaskManager._write_migration_state({"failed": True, "notice_count": 0}) - first = await ScheduleTaskManager.get_task_page_notice() - second = await ScheduleTaskManager.get_task_page_notice() - third = await ScheduleTaskManager.get_task_page_notice() + first = await TaskManager.get_task_page_notice() + second = await TaskManager.get_task_page_notice() + third = await TaskManager.get_task_page_notice() assert first == { "message": "系统更新了任务表的存储，旧表自动迁移失败，请手动重建任务 scheduler", @@ -1203,4 +1195,4 @@ async def test_task_page_notice_drops_legacy_tables_after_third_display(): } assert second is not None and second["displayCount"] == 2 assert third is not None and third["displayCount"] == 3 - assert ScheduleTaskManager._legacy_tables_exist() is False + assert TaskManager._legacy_tables_exist() is False diff --git a/tests/tool/test_task_center_compat.py b/tests/tool/test_task_center_compat.py index c43f7609e..ad4ad61e5 100644 --- a/tests/tool/test_task_center_compat.py +++ b/tests/tool/test_task_center_compat.py @@ -7,7 +7,7 @@ import flocks.tool.task.schedule_task_center # noqa: F401 from flocks.config.config import Config from flocks.storage.storage import Storage -from flocks.task.schedule_task_manager import ScheduleTaskManager +from flocks.task.manager import TaskManager from flocks.task.models import SchedulerMode, TaskPriority, TaskScheduler, TaskTrigger from flocks.task.store import TaskStore from flocks.tool.registry import ToolContext, ToolRegistry @@ -27,8 +27,8 @@ async def isolated_task_env(tmp_path: pytest.TempPathFactory, monkeypatch: pytes Config._cached_config = None Storage._db_path = None Storage._initialized = False - ScheduleTaskManager._instance = None - ScheduleTaskManager._startup_error = None + TaskManager._instance = None + TaskManager._startup_error = None TaskStore._initialized = False TaskStore._conn = None @@ -37,14 +37,14 @@ async def isolated_task_env(tmp_path: pytest.TempPathFactory, monkeypatch: pytes yield - await ScheduleTaskManager.stop() + await TaskManager.stop() await TaskStore.close() Config._global_config = None Config._cached_config = None Storage._db_path = None Storage._initialized = False - ScheduleTaskManager._instance = None - ScheduleTaskManager._startup_error = None + TaskManager._instance = None + TaskManager._startup_error = None TaskStore._initialized = False TaskStore._conn = None @@ -107,7 +107,7 @@ async def test_task_create_accepts_legacy_schedule_type_alias(self): assert result.success is True - schedulers, total = await ScheduleTaskManager.list_schedulers(limit=10) + schedulers, total = await TaskManager.list_schedulers(limit=10) assert total == 1 scheduler = schedulers[0] assert scheduler.mode == SchedulerMode.CRON @@ -130,7 +130,7 @@ async def test_task_create_infers_scheduled_type_from_cron(self): assert result.success is True - schedulers, total = await ScheduleTaskManager.list_schedulers(limit=10) + schedulers, total = await TaskManager.list_schedulers(limit=10) assert total == 1 scheduler = schedulers[0] assert scheduler.mode == SchedulerMode.CRON @@ -152,7 +152,7 @@ async def test_task_create_accepts_legacy_schedule_and_enabled_fields(self): assert result.success is True - schedulers, total = await ScheduleTaskManager.list_schedulers(limit=10) + schedulers, total = await TaskManager.list_schedulers(limit=10) assert total == 1 scheduler = schedulers[0] assert scheduler.mode == SchedulerMode.CRON @@ -161,7 +161,7 @@ async def test_task_create_accepts_legacy_schedule_and_enabled_fields(self): @pytest.mark.asyncio async def test_task_update_accepts_schedule_fields(self): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="原始任务", description="原始描述", mode=SchedulerMode.ONCE, @@ -187,7 +187,7 @@ async def test_task_update_accepts_schedule_fields(self): assert result.success is True - updated = await ScheduleTaskManager.get_scheduler(scheduler.id) + updated = await TaskManager.get_scheduler(scheduler.id) assert updated is not None assert updated.mode == SchedulerMode.CRON assert updated.description == "更新后的描述" @@ -218,7 +218,7 @@ async def test_task_create_rejects_run_once_without_time_instead_of_immediate(se assert result.error is not None assert "run_at" in result.error or "cron" in result.error - _, total = await ScheduleTaskManager.list_schedulers(limit=10) + _, total = await TaskManager.list_schedulers(limit=10) assert total == 0 @pytest.mark.asyncio @@ -238,7 +238,7 @@ async def test_task_create_schedule_json_accepts_string_boolean_run_once(self): assert result.success is True - schedulers, total = await ScheduleTaskManager.list_schedulers(limit=10) + schedulers, total = await TaskManager.list_schedulers(limit=10) assert total == 1 scheduler = schedulers[0] assert scheduler.mode == SchedulerMode.CRON @@ -263,12 +263,12 @@ async def test_task_create_rejects_six_field_cron_with_hint(self): assert "only 5-field cron is supported" in result.error assert "`0 6 * * *`" in result.error - _, total = await ScheduleTaskManager.list_schedulers(limit=10) + _, total = await TaskManager.list_schedulers(limit=10) assert total == 0 @pytest.mark.asyncio async def test_task_update_rejects_six_field_cron_with_hint(self): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="待更新的定时任务", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -292,7 +292,7 @@ async def test_task_update_rejects_six_field_cron_with_hint(self): assert "only 5-field cron is supported" in result.error assert "`0 6 * * *`" in result.error - unchanged = await ScheduleTaskManager.get_scheduler(scheduler.id) + unchanged = await TaskManager.get_scheduler(scheduler.id) assert unchanged is not None assert unchanged.trigger.cron == "*/5 * * * *" @@ -327,7 +327,7 @@ async def test_task_status_formats_scheduler_times_in_schedule_timezone(self): @pytest.mark.asyncio async def test_task_update_can_disable_and_enable_scheduled_task(self): - scheduler = await ScheduleTaskManager.create_scheduler( + scheduler = await TaskManager.create_scheduler( title="可停止的定时任务", mode=SchedulerMode.CRON, trigger=TaskTrigger( @@ -345,7 +345,7 @@ async def test_task_update_can_disable_and_enable_scheduled_task(self): ) assert disable_result.success is True - disabled = await ScheduleTaskManager.get_scheduler(scheduler.id) + disabled = await TaskManager.get_scheduler(scheduler.id) assert disabled is not None assert disabled.status.value == "disabled" @@ -359,6 +359,6 @@ async def test_task_update_can_disable_and_enable_scheduled_task(self): ) assert enable_result.success is True - enabled = await ScheduleTaskManager.get_scheduler(scheduler.id) + enabled = await TaskManager.get_scheduler(scheduler.id) assert enabled is not None assert enabled.status.value == "active" diff --git a/tests/tool/test_task_list_routing.py b/tests/tool/test_task_list_routing.py index 7c8442c7e..1ffc21853 100644 --- a/tests/tool/test_task_list_routing.py +++ b/tests/tool/test_task_list_routing.py @@ -15,7 +15,7 @@ from flocks.tool.registry import ToolContext, ToolRegistry, ToolResult from flocks.tool.task.schedule_task_center import schedule_task, schedule_task_list -_TM_PATH = "flocks.task.schedule_task_manager.ScheduleTaskManager" +_TM_PATH = "flocks.task.manager.TaskManager" def _ctx() -> ToolContext: diff --git a/webui/src/components/common/SessionChat.test.ts b/webui/src/components/common/SessionChat.test.ts index 29978f308..3bbe50dca 100644 --- a/webui/src/components/common/SessionChat.test.ts +++ b/webui/src/components/common/SessionChat.test.ts @@ -3819,28 +3819,6 @@ describe('ChatToolPart bash rendering', () => { expect(screen.getByText('tests passed').closest('pre')).toHaveClass('max-h-64'); }); - it.each([ - 'Tool execution was interrupted', - 'Command failed with exit code 1', - ])('renders the bash error once: %s', (error) => { - render( - React.createElement(ChatToolPart, { - part: { - id: 'bash-error-part', - type: 'tool', - tool: 'bash', - callID: 'call-bash-error', - state: { - status: 'error', - input: { command: 'exit 1' }, - error, - }, - } as any, - }), - ); - - expect(screen.getAllByText(error)).toHaveLength(1); - }); }); describe('ChatToolPart question result rendering', () => { diff --git a/webui/src/components/common/SessionChat.tsx b/webui/src/components/common/SessionChat.tsx index cf3d3cdcd..0be01168a 100644 --- a/webui/src/components/common/SessionChat.tsx +++ b/webui/src/components/common/SessionChat.tsx @@ -6275,7 +6275,7 @@ export function ChatToolPart({ part, pendingQuestion, onAnswer, onReject, proces )} - {!isBashTool && status === 'error' && state.error && ( + {status === 'error' && state.error && (
{state.error}
From b94e1fbf9e4a4890ed49fbf67ecd68e956b17aea Mon Sep 17 00:00:00 2001 From: xiami762 Date: Tue, 11 Aug 2026 11:40:39 +0800 Subject: [PATCH 12/15] chore: isolate runtime refactor scope --- flocks/server/app.py | 10 +++++----- flocks/server/routes/task_entities.py | 2 +- tests/server/test_server.py | 1 - webui/src/components/common/SessionChat.test.ts | 1 - 4 files changed, 6 insertions(+), 8 deletions(-) diff --git a/flocks/server/app.py b/flocks/server/app.py index eb304cc7b..5a1635023 100644 --- a/flocks/server/app.py +++ b/flocks/server/app.py @@ -353,14 +353,14 @@ async def _sync_workflows_phase() -> None: from flocks.task.manager import TaskManager await _run_startup_phase( log, - "schedule_task_manager.start", + "task_manager.start", TaskManager.start, ) - log.info("schedule_task_manager.started") + log.info("task_manager.started") except Exception as e: from flocks.task.manager import TaskManager TaskManager.mark_start_failed(e) - log.warning("schedule_task_manager.start.failed", {"error": str(e)}) + log.warning("task_manager.start.failed", {"error": str(e)}) # Seed built-in scheduled tasks from .flocks/plugins/tasks/*.json (idempotent) try: @@ -542,9 +542,9 @@ async def _delayed_trigger_runtime_start() -> None: from flocks.task.store import TaskStore await TaskManager.stop() await TaskStore.close() - log.info("schedule_task_manager.stopped") + log.info("task_manager.stopped") except Exception as e: - log.warning("schedule_task_manager.stop.failed", {"error": str(e)}) + log.warning("task_manager.stop.failed", {"error": str(e)}) # Stop Skill file watcher try: diff --git a/flocks/server/routes/task_entities.py b/flocks/server/routes/task_entities.py index ffbe42118..d828e39a8 100644 --- a/flocks/server/routes/task_entities.py +++ b/flocks/server/routes/task_entities.py @@ -342,6 +342,7 @@ async def delete_scheduler(scheduler_id: str): return {"ok": True} + @router.post("/task-schedulers/{scheduler_id}/enable") async def enable_scheduler(scheduler_id: str): from flocks.task.manager import TaskManager @@ -492,4 +493,3 @@ async def delete_execution(execution_id: str): if not await TaskManager.delete_execution(execution_id): raise HTTPException(404, "Task execution not found") return {"ok": True} - diff --git a/tests/server/test_server.py b/tests/server/test_server.py index eb79f07c5..d326a7b69 100644 --- a/tests/server/test_server.py +++ b/tests/server/test_server.py @@ -9,7 +9,6 @@ from fastapi import status from flocks.server.app import app -from flocks.task.manager import TaskManager from flocks.task.store import TaskStore from flocks.task.models import ( DeliveryStatus, diff --git a/webui/src/components/common/SessionChat.test.ts b/webui/src/components/common/SessionChat.test.ts index 3bbe50dca..698138064 100644 --- a/webui/src/components/common/SessionChat.test.ts +++ b/webui/src/components/common/SessionChat.test.ts @@ -3818,7 +3818,6 @@ describe('ChatToolPart bash rendering', () => { expect(screen.getByText('$').closest('pre')).toHaveClass('max-h-64'); expect(screen.getByText('tests passed').closest('pre')).toHaveClass('max-h-64'); }); - }); describe('ChatToolPart question result rendering', () => { From e10fb040a3d830566e31d15d49fda8f70235b8fa Mon Sep 17 00:00:00 2001 From: xiami762 Date: Tue, 11 Aug 2026 11:42:08 +0800 Subject: [PATCH 13/15] chore: remove split diff noise --- flocks/server/routes/task_entities.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/flocks/server/routes/task_entities.py b/flocks/server/routes/task_entities.py index d828e39a8..69170c4dd 100644 --- a/flocks/server/routes/task_entities.py +++ b/flocks/server/routes/task_entities.py @@ -342,7 +342,6 @@ async def delete_scheduler(scheduler_id: str): return {"ok": True} - @router.post("/task-schedulers/{scheduler_id}/enable") async def enable_scheduler(scheduler_id: str): from flocks.task.manager import TaskManager @@ -493,3 +492,5 @@ async def delete_execution(execution_id: str): if not await TaskManager.delete_execution(execution_id): raise HTTPException(404, "Task execution not found") return {"ok": True} + + From 682fcf164a082aa30d945764089bbbba10ebe3e4 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Thu, 13 Aug 2026 11:27:49 +0800 Subject: [PATCH 14/15] fix(session): preserve hook continuation semantics --- flocks/session/runtime/continuation_policy.py | 64 ++--- flocks/session/runtime/session_turn.py | 1 - flocks/session/runtime/step_engine.py | 100 +++++++- tests/session/test_lifecycle_hooks.py | 241 +++++++----------- tests/session/test_runner_llm_hooks.py | 134 ++++++++++ 5 files changed, 337 insertions(+), 203 deletions(-) diff --git a/flocks/session/runtime/continuation_policy.py b/flocks/session/runtime/continuation_policy.py index f2caa7dd2..ab14e4224 100644 --- a/flocks/session/runtime/continuation_policy.py +++ b/flocks/session/runtime/continuation_policy.py @@ -80,7 +80,6 @@ async def prepare_logical_turn(self, context: Any) -> None: ) if is_real_user_turn: context.turn_additional_context = None - context.stop_hook_active = False await self.run_user_prompt_submit(context, last_user) context.prepared_user_id = last_user.id @@ -120,7 +119,7 @@ async def resolve( context: Any, outcome: AgentRunOutcome[MessageInfo], ) -> ContinuationDecision[MessageInfo]: - """Resolve goal and TurnFinish into a new logical turn.""" + """Resolve queued input and goal continuation, then observe turn completion.""" last_user = outcome.last_user last_message = outcome.last_message if last_user is None or last_message is None: @@ -232,13 +231,19 @@ async def resolve( return queued_decision if not context.should_abort() and getattr(last_message, "finish", None) == "stop": - hook_decision = await self.run_turn_finish( + await self.run_turn_after( context, last_user, last_message, ) - if hook_decision.should_continue: - return hook_decision + + queued_decision = await self._materialize_continuation( + context, + last_user, + last_message, + ) + if queued_decision.should_continue: + return queued_decision stop_reason = getattr(last_message, "finish", None) or "stop" await self.publish_turn_stopped( @@ -247,20 +252,20 @@ async def resolve( ) return ContinuationDecision() - async def run_turn_finish( + async def run_turn_after( self, context: Any, last_user: MessageInfo, last_message: MessageInfo, - ) -> ContinuationDecision[MessageInfo]: - """Run TurnFinish and materialize a blocked-stop continuation.""" + ) -> None: + """Publish terminal turn facts without changing continuation control flow.""" try: hook_user = last_user if context.turn_user_id: hook_user = await Message.get(context.session.id, context.turn_user_id) or last_user user_text = await Message.get_text_content(hook_user) assistant_text = await Message.get_text_content(last_message) - hook_context = await HookPipeline.run_turn_finish( + await HookPipeline.run_turn_after( { "sessionID": context.session.id, "workspace": context.session.directory, @@ -278,51 +283,21 @@ async def run_turn_finish( "id": last_message.id, "content": assistant_text, }, - "finishReason": "stop", - "stopHookActive": context.stop_hook_active, + "terminalOutcome": { + "status": "success", + "finish_reason": "stop", + }, } ) except Exception as exc: log.debug( - "session.hook.turn_finish_error", + "session.hook.turn_after_error", { "session_id": context.session.id, "message_id": getattr(last_message, "id", None), "error": str(exc), }, ) - return ContinuationDecision() - - decision = str(hook_context.output.get("decision") or "").strip().lower() - reason = str(hook_context.output.get("reason") or "").strip() - if decision != "block" or not reason or context.should_abort(): - return ContinuationDecision() - - allow_synthetic = await self._synthetic_continuation_allowed( - context, - last_message, - ) - continuation = await self._materialize_continuation( - context, - last_user, - last_message, - candidate_reason="turn_finish_hook", - content=reason, - agent=getattr(hook_user, "agent", None) or context.agent_name, - model={ - "providerID": context.provider_id, - "modelID": context.model_id, - }, - part_metadata={ - "turnFinishContinuation": True, - "stopHookActive": True, - "sourceAssistantMessageID": last_message.id, - }, - allow_synthetic=allow_synthetic, - ) - if continuation.reason == "turn_finish_hook": - context.stop_hook_active = True - return continuation @staticmethod async def _synthetic_continuation_allowed( @@ -448,7 +423,6 @@ async def _publish_continuation( message_id_key = { "queued_message": "queuedUserMessageID", "goal": "goalMessageID", - "turn_finish_hook": "turnFinishMessageID", }.get(reason, "continuationMessageID") await SessionEventSink.emit( context.callbacks, diff --git a/flocks/session/runtime/session_turn.py b/flocks/session/runtime/session_turn.py index 774f15b3b..3e782222d 100644 --- a/flocks/session/runtime/session_turn.py +++ b/flocks/session/runtime/session_turn.py @@ -121,7 +121,6 @@ class LoopContext: model_candidate_policy: Literal["fixed", "automatic", "configured"] = "automatic" turn_user_id: Optional[str] = None turn_additional_context: Optional[str] = None - stop_hook_active: bool = False prepared_user_id: Optional[str] = None prepared_messages: Optional[List[MessageInfo]] = field(default=None, repr=False) session_start_pending: bool = False diff --git a/flocks/session/runtime/step_engine.py b/flocks/session/runtime/step_engine.py index 00af784c5..7c5b30a35 100644 --- a/flocks/session/runtime/step_engine.py +++ b/flocks/session/runtime/step_engine.py @@ -49,6 +49,11 @@ SessionRetry, ) from flocks.session.lifecycle.compaction import SessionCompaction, CompactionPolicy +from flocks.session.llm_hook_utils import ( + StreamingTextReplacementBuffer, + restore_value_with_replacements, + stream_text_replacements_from_hook_output, +) from flocks.session.streaming.stream_processor import StreamProcessor from flocks.session.streaming.stream_events import ( StartEvent, @@ -98,6 +103,7 @@ TOOL_RESULT_CHAR_BUDGET_RATIO = 0.70 TOOL_RESULT_TURN_BUDGET_RATIO = 0.35 TOOL_RESULT_MIN_CHAR_BUDGET = 8_000 +STREAM_TEXT_REPLACEMENTS_METADATA_KEY = "llmHookStreamTextReplacements" def _annotate_with_provider_version(tool_info: Any, description: Optional[str]) -> str: @@ -3241,6 +3247,9 @@ async def _apply_before_model_hook( metadata.get("providerToolsEnabled"), ), ) + metadata[STREAM_TEXT_REPLACEMENTS_METADATA_KEY] = ( + stream_text_replacements_from_hook_output(hook_output) + ) return ModelRequest( provider_id=request.provider_id, model_id=request.model_id, @@ -3344,6 +3353,25 @@ def _build_llm_response_payload( self._active_model_request = request messages = request.provider_messages() tools = request.provider_tools() + replacements = [ + (replacement[0], replacement[1]) + for replacement in request.metadata.get( + STREAM_TEXT_REPLACEMENTS_METADATA_KEY, + (), + ) + if ( + isinstance(replacement, (list, tuple)) + and len(replacement) == 2 + and isinstance(replacement[0], str) + and isinstance(replacement[1], str) + ) + ] + stream_text_rewriter = ( + StreamingTextReplacementBuffer(replacements) if replacements else None + ) + stream_reasoning_rewriter = ( + StreamingTextReplacementBuffer(replacements) if replacements else None + ) # Create stream processor main_session_key = self.session.id @@ -3402,6 +3430,28 @@ async def _on_tool_execution_end( plan_permission_path=turn_plan_file.permission_path, ) + async def _flush_reasoning_rewriter() -> None: + if stream_reasoning_rewriter is None or not hasattr( + self, + "_current_reasoning_id", + ): + return + trailing_reasoning = stream_reasoning_rewriter.flush() + if not trailing_reasoning: + return + reasoning_metadata = getattr( + self, + "_current_reasoning_metadata", + {}, + ) or {} + await processor.process_event( + ReasoningDeltaEvent( + id=self._current_reasoning_id, + text=trailing_reasoning, + metadata=reasoning_metadata, + ) + ) + provider_options = request.provider_options() provider_tools = ( tools @@ -3578,24 +3628,30 @@ async def _on_tool_execution_end( # reasoning text. event_type = getattr(chunk, 'event_type', None) chunk_metadata = getattr(chunk, 'metadata', None) or {} + display_chunk_metadata = ( + restore_value_with_replacements(chunk_metadata, replacements) + if replacements + else chunk_metadata + ) reasoning_event_types = {"reasoning", "reasoning-start", "reasoning-end"} - if hasattr(self, '_current_reasoning_id') and chunk_metadata: + if hasattr(self, '_current_reasoning_id') and display_chunk_metadata: current_metadata = getattr(self, '_current_reasoning_metadata', {}) or {} - current_metadata.update(chunk_metadata) + current_metadata.update(display_chunk_metadata) self._current_reasoning_metadata = current_metadata if event_type == "reasoning-start" and not hasattr(self, '_current_reasoning_id'): self._attempt_state.observable_output_started = True reasoning_id_counter += 1 self._current_reasoning_id = f"reasoning-{reasoning_id_counter}" - self._current_reasoning_metadata = dict(chunk_metadata) + self._current_reasoning_metadata = dict(display_chunk_metadata) await processor.process_event(ReasoningStartEvent( id=self._current_reasoning_id, - metadata=chunk_metadata, + metadata=display_chunk_metadata, )) if event_type == "reasoning-end" and hasattr(self, '_current_reasoning_id'): + await _flush_reasoning_rewriter() reasoning_end_metadata = getattr(self, '_current_reasoning_metadata', {}) or {} await processor.process_event(ReasoningEndEvent( id=self._current_reasoning_id, @@ -3637,22 +3693,28 @@ async def _on_tool_execution_end( if not hasattr(self, '_current_reasoning_id'): reasoning_id_counter += 1 self._current_reasoning_id = f"reasoning-{reasoning_id_counter}" - self._current_reasoning_metadata = dict(chunk_metadata) + self._current_reasoning_metadata = dict(display_chunk_metadata) await processor.process_event(ReasoningStartEvent( id=self._current_reasoning_id, - metadata=chunk_metadata, + metadata=display_chunk_metadata, )) if chunk_reasoning: - await processor.process_event(ReasoningDeltaEvent( - id=self._current_reasoning_id, - text=chunk_reasoning, - metadata=chunk_metadata, - )) + if stream_reasoning_rewriter is not None: + reasoning_text = stream_reasoning_rewriter.feed( + reasoning_text, + ) + if reasoning_text: + await processor.process_event(ReasoningDeltaEvent( + id=self._current_reasoning_id, + text=reasoning_text, + metadata=display_chunk_metadata, + )) # 2) End reasoning block when this chunk also carries non-reasoning # content (or once the stream moves away from reasoning). if (chunk_text or chunk_tool_calls) and hasattr(self, '_current_reasoning_id'): + await _flush_reasoning_rewriter() reasoning_end_metadata = getattr(self, '_current_reasoning_metadata', {}) or {} await processor.process_event(ReasoningEndEvent( id=self._current_reasoning_id, @@ -3663,9 +3725,14 @@ async def _on_tool_execution_end( delattr(self, '_current_reasoning_metadata') # 3) Process text delta. - if chunk_text: + raw_chunk_text = chunk_text + if chunk_text and stream_text_rewriter is not None: + chunk_text = stream_text_rewriter.feed(chunk_text) + + if raw_chunk_text: self._attempt_state.observable_output_started = True chunk_counts["text"] += 1 + if chunk_text: if not text_started: await processor.process_event(TextStartEvent()) text_started = True @@ -3728,12 +3795,21 @@ async def _on_tool_execution_end( await tool_accumulator.flush_remaining(stream_finish_reason) + if stream_text_rewriter is not None: + trailing_text = stream_text_rewriter.flush() + if trailing_text: + if not text_started: + await processor.process_event(TextStartEvent()) + text_started = True + await processor.process_event(TextDeltaEvent(text=trailing_text)) + # End text block if started if text_started: await processor.process_event(TextEndEvent()) # End any remaining reasoning block if hasattr(self, '_current_reasoning_id'): + await _flush_reasoning_rewriter() reasoning_end_metadata = getattr(self, '_current_reasoning_metadata', {}) or {} await processor.process_event(ReasoningEndEvent( id=self._current_reasoning_id, diff --git a/tests/session/test_lifecycle_hooks.py b/tests/session/test_lifecycle_hooks.py index 768a4562c..ffbfdcd81 100644 --- a/tests/session/test_lifecycle_hooks.py +++ b/tests/session/test_lifecycle_hooks.py @@ -8,7 +8,6 @@ import pytest -from flocks.session.runtime.contracts import ContinuationDecision from flocks.hooks.pipeline import HookContext, HookStage from flocks.session.runtime.continuation_policy import DEFAULT_CONTINUATION_POLICY from flocks.session.goal import GoalDecision @@ -130,26 +129,31 @@ async def test_session_start_runs_only_when_pending() -> None: @pytest.mark.asyncio -async def test_turn_finish_block_creates_synthetic_continuation() -> None: - ctx = _loop_context("ses_turn_finish") +async def test_goal_waiting_cannot_be_overridden_by_turn_after_output() -> None: + ctx = _loop_context("ses_turn_after_goal_waiting") ctx.turn_user_id = "msg_user" user = SimpleNamespace( id="msg_user", agent="rex", + role="user", model={"providerID": "test-provider", "modelID": "test-model"}, ) assistant = SimpleNamespace( id="msg_assistant", agent="rex", + role="assistant", finish="stop", ) continuation = SimpleNamespace(id="msg_continuation") callbacks = LoopCallbacks(event_publish_callback=AsyncMock()) ctx.callbacks = callbacks + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock(return_value=[user, assistant]), + ) create_message = AsyncMock(return_value=continuation) run_hook = AsyncMock( return_value=HookContext( - stage=HookStage.TURN_FINISH, + stage=HookStage.TURN_AFTER, input={}, output={ "decision": "block", @@ -160,52 +164,58 @@ async def test_turn_finish_block_creates_synthetic_continuation() -> None: with ( patch( - "flocks.session.runtime.session_turn.Message.get", + "flocks.session.runtime.continuation_policy.Message.get", AsyncMock(return_value=user), ), patch( - "flocks.session.runtime.session_turn.Message.get_text_content", - AsyncMock(side_effect=["implement hooks", "implementation complete"]), + "flocks.session.runtime.continuation_policy.Message.get_text_content", + AsyncMock( + side_effect=lambda message: ( + "implement hooks" + if message.id == user.id + else "Please provide the missing input." + ) + ), ), patch( - "flocks.session.runtime.session_turn.Message.create", + "flocks.session.runtime.continuation_policy.Message.create", create_message, ), patch( - "flocks.hooks.pipeline.HookPipeline.run_turn_finish", + "flocks.hooks.pipeline.HookPipeline.run_turn_after", run_hook, ), patch( - "flocks.agent.registry.Agent.get", - AsyncMock(return_value=SimpleNamespace(steps=10)), + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", + AsyncMock( + return_value=GoalDecision( + status="active", + verdict="waiting", + should_continue=False, + reason="Waiting for user input.", + ) + ), ), ): - decision = await DEFAULT_CONTINUATION_POLICY.run_turn_finish( + decision = await DEFAULT_CONTINUATION_POLICY.resolve( ctx, - user, - assistant, - ) - continued = decision.should_continue - - assert continued is True - assert ctx.stop_hook_active is True - assert create_message.await_args.kwargs["content"] == ("Run the test suite before finishing.") - assert create_message.await_args.kwargs["synthetic"] is True - assert create_message.await_args.kwargs["part_metadata"] == { - "turnFinishContinuation": True, - "stopHookActive": True, - "sourceAssistantMessageID": assistant.id, - } + SimpleNamespace(last_user=user, last_message=assistant), + ) + + assert decision.should_continue is False + create_message.assert_not_awaited() hook_payload = run_hook.await_args.args[0] - assert hook_payload["finishReason"] == "stop" - assert hook_payload["stopHookActive"] is False + assert hook_payload["terminalOutcome"] == { + "status": "success", + "finish_reason": "stop", + } callbacks.event_publish_callback.assert_awaited_once() - assert callbacks.event_publish_callback.await_args.args[0] == "turn.continued" + assert callbacks.event_publish_callback.await_args.args[0] == "turn.stopped" @pytest.mark.asyncio -async def test_queued_prompt_arriving_during_turn_finish_wins() -> None: - ctx = _loop_context("ses_turn_finish_queue_race") +async def test_queued_prompt_arriving_during_turn_after_wins() -> None: + ctx = _loop_context("ses_turn_after_queue_race") ctx.turn_user_id = "msg_001" user = SimpleNamespace(id="msg_001", agent="rex", role="user") assistant = SimpleNamespace( @@ -215,48 +225,42 @@ async def test_queued_prompt_arriving_during_turn_finish_wins() -> None: finish="stop", ) queued_user = SimpleNamespace(id="msg_003", agent="rex", role="user") + messages = [user, assistant] ctx.session_store = SimpleNamespace( - get_messages=AsyncMock( - return_value=[user, assistant, queued_user], - ) + get_messages=AsyncMock(side_effect=lambda: messages), ) callbacks = LoopCallbacks(event_publish_callback=AsyncMock()) ctx.callbacks = callbacks - create_message = AsyncMock() + run_turn_after = AsyncMock(side_effect=lambda *_args: messages.append(queued_user)) with ( patch( - "flocks.session.runtime.session_turn.Message.get", - AsyncMock(return_value=user), - ), - patch( - "flocks.session.runtime.session_turn.Message.get_text_content", - AsyncMock(side_effect=["prompt", "response"]), - ), - patch( - "flocks.session.runtime.session_turn.Message.create", - create_message, + "flocks.session.runtime.continuation_policy.Message.get_text_content", + AsyncMock(return_value="response"), ), patch( - "flocks.hooks.pipeline.HookPipeline.run_turn_finish", + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", AsyncMock( - return_value=HookContext( - stage=HookStage.TURN_FINISH, - input={}, - output={"decision": "block", "reason": "continue"}, + return_value=GoalDecision( + status=None, + verdict="inactive", ) ), ), + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_after", + run_turn_after, + create=True, + ), ): - decision = await DEFAULT_CONTINUATION_POLICY.run_turn_finish( + decision = await DEFAULT_CONTINUATION_POLICY.resolve( ctx, - user, - assistant, + SimpleNamespace(last_user=user, last_message=assistant), ) - continued = decision.should_continue - assert continued is True - create_message.assert_not_awaited() + assert decision.should_continue is True + run_turn_after.assert_awaited_once_with(ctx, user, assistant) callbacks.event_publish_callback.assert_awaited_once() event_name, payload = callbacks.event_publish_callback.await_args.args assert event_name == "turn.continued" @@ -331,59 +335,6 @@ async def test_real_user_arriving_during_goal_evaluation_wins() -> None: create_message.assert_not_awaited() -@pytest.mark.asyncio -async def test_turn_finish_block_is_ignored_at_agent_step_limit() -> None: - ctx = _loop_context("ses_turn_finish_limit") - ctx.turn_user_id = "msg_user" - ctx.trace_step_offset = 2 - ctx.step = 1 - user = SimpleNamespace(id="msg_user", agent="rex") - assistant = SimpleNamespace( - id="msg_assistant", - agent="rex", - finish="stop", - ) - create_message = AsyncMock() - - with ( - patch( - "flocks.session.runtime.session_turn.Message.get", - AsyncMock(return_value=user), - ), - patch( - "flocks.session.runtime.session_turn.Message.get_text_content", - AsyncMock(side_effect=["prompt", "response"]), - ), - patch( - "flocks.session.runtime.session_turn.Message.create", - create_message, - ), - patch( - "flocks.hooks.pipeline.HookPipeline.run_turn_finish", - AsyncMock( - return_value=HookContext( - stage=HookStage.TURN_FINISH, - input={}, - output={"decision": "block", "reason": "continue"}, - ) - ), - ), - patch( - "flocks.agent.registry.Agent.get", - AsyncMock(return_value=SimpleNamespace(steps=3)), - ), - ): - decision = await DEFAULT_CONTINUATION_POLICY.run_turn_finish( - ctx, - user, - assistant, - ) - continued = decision.should_continue - - assert continued is False - create_message.assert_not_awaited() - - def _message( message_id: str, role: str, @@ -402,8 +353,8 @@ def _message( @pytest.mark.asyncio -async def test_turn_finish_runs_only_after_persisted_stop() -> None: - ctx = _loop_context("ses_turn_finish_integration") +async def test_turn_after_runs_only_after_persisted_stop() -> None: + ctx = _loop_context("ses_turn_after_integration") user = _message("msg_001", "user") assistant = _message("msg_002", "assistant", finish="stop") ctx.session_store = SimpleNamespace( @@ -414,7 +365,7 @@ async def test_turn_finish_runs_only_after_persisted_stop() -> None: ] ) ) - run_turn_finish = AsyncMock(return_value=ContinuationDecision()) + run_turn_after = AsyncMock() with ( patch( @@ -444,8 +395,8 @@ async def test_turn_finish_runs_only_after_persisted_stop() -> None: ), patch.object( DEFAULT_CONTINUATION_POLICY, - "run_turn_finish", - run_turn_finish, + "run_turn_after", + run_turn_after, ), patch( "flocks.session.runtime.step_engine.StepEngine._process_step", @@ -463,7 +414,7 @@ async def test_turn_finish_runs_only_after_persisted_stop() -> None: result = await run_logical_turns(ctx, LoopCallbacks()) assert result.action == "stop" - run_turn_finish.assert_awaited_once_with( + run_turn_after.assert_awaited_once_with( ctx, user, assistant, @@ -478,11 +429,11 @@ async def test_turn_finish_runs_only_after_persisted_stop() -> None: (StepResult(action="continue"), "tool-calls"), ], ) -async def test_turn_finish_skips_errors_and_tool_calls( +async def test_turn_after_skips_errors_and_tool_calls( step_result: StepResult, assistant_finish: str, ) -> None: - ctx = _loop_context(f"ses_turn_finish_{assistant_finish}") + ctx = _loop_context(f"ses_turn_after_{assistant_finish}") user = _message("msg_001", "user") assistant = _message("msg_002", "assistant", finish=assistant_finish) ctx.session_store = SimpleNamespace( @@ -493,7 +444,7 @@ async def test_turn_finish_skips_errors_and_tool_calls( ] ) ) - run_turn_finish = AsyncMock(return_value=ContinuationDecision()) + run_turn_after = AsyncMock() async def process_step(*_args, **_kwargs): if step_result.action == "continue": @@ -515,8 +466,8 @@ async def process_step(*_args, **_kwargs): ), patch.object( DEFAULT_CONTINUATION_POLICY, - "run_turn_finish", - run_turn_finish, + "run_turn_after", + run_turn_after, ), patch( "flocks.session.runtime.step_engine.StepEngine._process_step", @@ -534,12 +485,12 @@ async def process_step(*_args, **_kwargs): result = await run_logical_turns(ctx, LoopCallbacks()) assert result.action == "stop" - run_turn_finish.assert_not_awaited() + run_turn_after.assert_not_awaited() @pytest.mark.asyncio -async def test_queued_user_message_takes_priority_over_turn_finish() -> None: - ctx = _loop_context("ses_turn_finish_queue") +async def test_queued_user_message_takes_priority_over_turn_after() -> None: + ctx = _loop_context("ses_turn_after_queue") user = _message("msg_001", "user") assistant = _message("msg_002", "assistant", finish="stop") queued_user = _message("msg_003", "user") @@ -551,7 +502,7 @@ async def test_queued_user_message_takes_priority_over_turn_finish() -> None: ] ) ) - run_turn_finish = AsyncMock(return_value=ContinuationDecision()) + run_turn_after = AsyncMock() async def process_step(*_args, **_kwargs): ctx.signal_abort() @@ -572,8 +523,8 @@ async def process_step(*_args, **_kwargs): ), patch.object( DEFAULT_CONTINUATION_POLICY, - "run_turn_finish", - run_turn_finish, + "run_turn_after", + run_turn_after, ), patch( "flocks.session.runtime.step_engine.StepEngine._process_step", @@ -590,12 +541,12 @@ async def process_step(*_args, **_kwargs): ): await run_logical_turns(ctx, LoopCallbacks()) - run_turn_finish.assert_not_awaited() + run_turn_after.assert_not_awaited() @pytest.mark.asyncio -async def test_goal_continuation_takes_priority_over_turn_finish() -> None: - ctx = _loop_context("ses_turn_finish_goal") +async def test_goal_continuation_takes_priority_over_turn_after() -> None: + ctx = _loop_context("ses_turn_after_goal") user = _message("msg_001", "user") assistant = _message("msg_002", "assistant", finish="stop") goal_user = _message("msg_003", "user") @@ -607,7 +558,7 @@ async def test_goal_continuation_takes_priority_over_turn_finish() -> None: ] ) ) - run_turn_finish = AsyncMock(return_value=ContinuationDecision()) + run_turn_after = AsyncMock() async def process_step(*_args, **_kwargs): ctx.signal_abort() @@ -647,8 +598,8 @@ async def process_step(*_args, **_kwargs): ), patch.object( DEFAULT_CONTINUATION_POLICY, - "run_turn_finish", - run_turn_finish, + "run_turn_after", + run_turn_after, ), patch( "flocks.session.runtime.step_engine.StepEngine._process_step", @@ -665,15 +616,15 @@ async def process_step(*_args, **_kwargs): ): await run_logical_turns(ctx, LoopCallbacks()) - run_turn_finish.assert_not_awaited() + run_turn_after.assert_not_awaited() @pytest.mark.asyncio -async def test_abort_does_not_trigger_turn_finish() -> None: - ctx = _loop_context("ses_turn_finish_abort") +async def test_abort_does_not_trigger_turn_after() -> None: + ctx = _loop_context("ses_turn_after_abort") user = _message("msg_001", "user") ctx.session_store = SimpleNamespace(get_messages=AsyncMock(return_value=[user])) - run_turn_finish = AsyncMock(return_value=ContinuationDecision()) + run_turn_after = AsyncMock() async def cancel_for_user_abort(*_args, **_kwargs): ctx.signal_abort() @@ -694,8 +645,8 @@ async def cancel_for_user_abort(*_args, **_kwargs): ), patch.object( DEFAULT_CONTINUATION_POLICY, - "run_turn_finish", - run_turn_finish, + "run_turn_after", + run_turn_after, ), patch( "flocks.session.runtime.step_engine.StepEngine._process_step", @@ -712,12 +663,12 @@ async def cancel_for_user_abort(*_args, **_kwargs): ): await run_logical_turns(ctx, LoopCallbacks()) - run_turn_finish.assert_not_awaited() + run_turn_after.assert_not_awaited() @pytest.mark.asyncio -async def test_late_abort_after_step_completion_skips_turn_finish() -> None: - ctx = _loop_context("ses_turn_finish_late_abort") +async def test_late_abort_after_step_completion_skips_turn_after() -> None: + ctx = _loop_context("ses_turn_after_late_abort") user = _message("msg_001", "user") assistant = _message("msg_002", "assistant", finish="stop") ctx.session_store = SimpleNamespace( @@ -728,7 +679,7 @@ async def test_late_abort_after_step_completion_skips_turn_finish() -> None: ] ) ) - run_turn_finish = AsyncMock(return_value=ContinuationDecision()) + run_turn_after = AsyncMock() async def abort_after_step(_step: int) -> None: ctx.abort_event.set() @@ -761,8 +712,8 @@ async def abort_after_step(_step: int) -> None: ), patch.object( DEFAULT_CONTINUATION_POLICY, - "run_turn_finish", - run_turn_finish, + "run_turn_after", + run_turn_after, ), patch( "flocks.session.runtime.step_engine.StepEngine._process_step", @@ -782,5 +733,5 @@ async def abort_after_step(_step: int) -> None: LoopCallbacks(on_step_end=abort_after_step), ) - run_turn_finish.assert_not_awaited() + run_turn_after.assert_not_awaited() assert result.metadata["aborted"] is True diff --git a/tests/session/test_runner_llm_hooks.py b/tests/session/test_runner_llm_hooks.py index 9386cb853..5119d111d 100644 --- a/tests/session/test_runner_llm_hooks.py +++ b/tests/session/test_runner_llm_hooks.py @@ -39,6 +39,7 @@ class _FakeProcessor: def __init__(self, **_: object): self._text_parts: list[str] = [] self._reasoning_parts: list[str] = [] + self.reasoning_metadata: list[dict[str, object]] = [] self.finish_reason = "stop" self.tool_calls = {} self._langfuse_generation = None @@ -49,6 +50,7 @@ async def process_event(self, event) -> None: self._text_parts.append(event.text) elif event_name == "ReasoningDeltaEvent": self._reasoning_parts.append(event.text) + self.reasoning_metadata.append(event.metadata) elif event_name == "FinishEvent": self.finish_reason = event.finish_reason @@ -422,6 +424,138 @@ async def _gen(): assert "[[V_EMAIL_1]]" in str(generation_inputs) +@pytest.mark.asyncio +async def test_call_llm_restores_stream_replacements_across_chunks_and_retries( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner = _make_runner("ses_runner_stream_replacements") + assistant_msg = SimpleNamespace(id="msg_assistant_stream_replacements") + agent = SimpleNamespace(name="rex") + processors: list[_FakeProcessor] = [] + + async def _before(payload): + return SimpleNamespace( + output={ + "request": { + **payload["request"], + "messages": [ + {"role": "user", "content": "email [[V_EMAIL_1]]"} + ], + "providerOptions": {}, + }, + "redaction": { + "streamTextReplacements": [ + { + "placeholder": "[[V_EMAIL_1]]", + "value": "alice@example.com", + } + ], + }, + } + ) + + class _RecordingProcessor(_FakeProcessor): + def __init__(self, **kwargs: object): + super().__init__(**kwargs) + processors.append(self) + + monkeypatch.setattr(runner_mod, "StreamProcessor", _RecordingProcessor) + monkeypatch.setattr( + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock( + side_effect=lambda stage, _metadata=None: ( + stage == runner_mod.HookStage.LLM_BEFORE + ) + ), + ) + run_before = AsyncMock(side_effect=_before) + monkeypatch.setattr(runner_mod.HookPipeline, "run_llm_before", run_before) + monkeypatch.setattr(runner_mod, "langfuse_is_active", lambda: False) + monkeypatch.setattr( + "flocks.provider.options.build_provider_options", + lambda provider_id, model_id: {}, + ) + monkeypatch.setattr( + "flocks.session.streaming.tool_accumulator.ToolCallAccumulator", + _FakeToolAccumulator, + ) + monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) + + class _Provider: + def chat_stream(self, **kwargs): + assert kwargs["messages"][0].content == "email [[V_EMAIL_1]]" + + async def _gen(): + yield SimpleNamespace( + delta="", + reasoning="Contact [[V_EM", + metadata={"reasoningContent": "Contact [[V_EMAIL_1]]"}, + event_type="reasoning", + tool_calls=None, + finish_reason=None, + usage=None, + ) + yield SimpleNamespace( + delta="", + reasoning="AIL_1]]", + metadata={}, + event_type="reasoning", + tool_calls=None, + finish_reason=None, + usage=None, + ) + yield SimpleNamespace( + delta="Reply to [[V_EM", + reasoning=None, + metadata={}, + event_type=None, + tool_calls=None, + finish_reason=None, + usage=None, + ) + yield SimpleNamespace( + delta="AIL_1]]", + reasoning=None, + metadata={}, + event_type=None, + tool_calls=None, + finish_reason="stop", + usage=None, + ) + + return _gen() + + results = [] + for _ in range(2): + results.append( + await runner._call_llm( + provider=_Provider(), + messages=[ + ChatMessage(role="user", content="email alice@example.com") + ], + tools=[], + agent=agent, + assistant_msg=assistant_msg, + ) + ) + + assert [result.content for result in results] == [ + "Reply to alice@example.com", + "Reply to alice@example.com", + ] + assert [processor.get_reasoning_content() for processor in processors] == [ + "Contact alice@example.com", + "Contact alice@example.com", + ] + assert all( + "alice@example.com" in str(processor.reasoning_metadata) + and "[[V_EMAIL_1]]" not in str(processor.reasoning_metadata) + for processor in processors + ) + run_before.assert_awaited_once() + + @pytest.mark.asyncio async def test_call_llm_emits_after_hook_on_error(monkeypatch: pytest.MonkeyPatch): runner = _make_runner("ses_runner_llm_hooks_error") From 933f5bd12679abaf66f6d715d01583119ab3918e Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Thu, 13 Aug 2026 16:55:03 +0800 Subject: [PATCH 15/15] test(session): remove obsolete runner tests --- tests/permission/test_interactive.py | 38 ---------------------- tests/server/routes/test_session_routes.py | 36 -------------------- 2 files changed, 74 deletions(-) diff --git a/tests/permission/test_interactive.py b/tests/permission/test_interactive.py index 4521a34e8..f7eacbdaf 100644 --- a/tests/permission/test_interactive.py +++ b/tests/permission/test_interactive.py @@ -12,41 +12,3 @@ def test_auto_approve_enabled_reads_env(monkeypatch: pytest.MonkeyPatch) -> None assert auto_approve_enabled() is False monkeypatch.setenv("FLOCKS_AUTO_APPROVE", "true") assert auto_approve_enabled() is True - - -@pytest.mark.asyncio -async def test_runner_handle_permission_auto_allows_without_permission_next( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from flocks.session.runner import SessionRunner - - async def _unexpected_ask(*args, **kwargs): - raise AssertionError("PermissionNext.ask should not run for legacy tool permissions") - - monkeypatch.setattr( - "flocks.permission.next.PermissionNext.ask", - _unexpected_ask, - ) - - runner = SessionRunner.__new__(SessionRunner) - runner.session = type("Session", (), {"id": "ses_test"})() - runner._step = 1 - runner.callbacks = type( - "Callbacks", - (), - {"on_permission_request": None, "event_publish_callback": None}, - )() - - request = type( - "Request", - (), - { - "permission": "write", - "patterns": ["notes.md"], - "metadata": {}, - "message_id": "msg_1", - "always": ["*"], - }, - )() - - await runner._handle_permission(request) diff --git a/tests/server/routes/test_session_routes.py b/tests/server/routes/test_session_routes.py index d4798f1da..146eac9d7 100644 --- a/tests/server/routes/test_session_routes.py +++ b/tests/server/routes/test_session_routes.py @@ -21,7 +21,6 @@ from httpx import AsyncClient from flocks.auth.context import API_TOKEN_SERVICE_USER_ID, AuthUser from flocks.hooks.execution import ( - ExecutionStopped, current_execution_context, execution_context_scope, ) @@ -83,41 +82,6 @@ async def test_missing_session_directory_uses_cwd_and_publishes_notice( "fallbackDirectory": str(tmp_path), }, ) - - -@pytest.mark.asyncio -async def test_shell_route_maps_extension_stop_to_forbidden( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A Pro policy stop must not surface as an unhandled server error.""" - - monkeypatch.setattr(session_routes, "require_user", lambda _request: object()) - monkeypatch.setattr( - session_routes, - "_get_session_by_id_unfiltered", - AsyncMock(return_value=object()), - ) - monkeypatch.setattr( - session_routes, - "_require_session_write_access", - lambda _session, _user: None, - ) - monkeypatch.setattr( - "flocks.session.runner.SessionRunner.shell", - AsyncMock(side_effect=ExecutionStopped("hard_deny_system_delete")), - ) - - with pytest.raises(HTTPException) as error: - await session_routes.run_shell_command( - "ses_1", - session_routes.ShellRequest(agent="build", command="rm -rf /etc"), - SimpleNamespace(), - ) - - assert error.value.status_code == status.HTTP_403_FORBIDDEN - assert error.value.detail == "execution stopped by extension" - - @pytest.mark.asyncio async def test_background_session_task_preserves_execution_context() -> None: """Async session work retains opaque ingress context after scheduling."""