diff --git a/flocks/channel/inbound/dispatcher.py b/flocks/channel/inbound/dispatcher.py index b332110ce..f8eac3247 100644 --- a/flocks/channel/inbound/dispatcher.py +++ b/flocks/channel/inbound/dispatcher.py @@ -237,14 +237,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 @@ -514,8 +514,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 @@ -1223,13 +1223,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 2100b28dc..96465a333 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 c0dbf61f5..2ec5a3e8c 100644 --- a/flocks/server/routes/session.py +++ b/flocks/server/routes/session.py @@ -1740,30 +1740,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 @@ -1778,7 +1772,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, }) @@ -1835,7 +1828,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) @@ -1847,12 +1840,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}) @@ -3471,7 +3462,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 @@ -4945,7 +4935,7 @@ class ShellRequest(BaseModel): async def run_shell_command(sessionID: str, request: ShellRequest, http_request: Request): """Run shell command""" from flocks.hooks.execution import ExecutionStopped - 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) @@ -4956,17 +4946,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..b2657ad54 --- /dev/null +++ b/flocks/session/actions.py @@ -0,0 +1,180 @@ +"""Session actions that are independent of the agent execution loop.""" + +import asyncio +import os +from collections.abc import Mapping +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") + + cwd = session.directory or os.getcwd() + + async def _effect( + execution_command: str = command, + execution_cwd: str = cwd, + ) -> dict[str, Any]: + 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( + execution_command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=execution_cwd, + ) + 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": execution_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": execution_command}, + "output": output, + }, + }, + ], + } + + from flocks.session.tool_execution import ( + build_session_tool_execution_payload, + run_tool_execution_lifecycle, + ) + + payload = await build_session_tool_execution_payload( + session_id=session_id, + message_id=Identifier.create("message"), + agent=agent, + tool_name="shell", + tool_input={"command": command, "workdir": cwd}, + validated_input={"command": command, "workdir": cwd}, + tool_schema={ + "type": "object", + "properties": { + "command": {"type": "string"}, + "workdir": {"type": "string"}, + }, + "required": ["command"], + }, + tool_context_extra={ + "tool_source": "session_actions", + "tool_category": "command", + "workspace_dir": cwd, + "session_execution_profile": { + "entry": "session.shell", + "workspace_dir": cwd, + }, + }, + ) + + async def _patched_effect(patch: Mapping[str, Any]) -> dict[str, Any]: + patched_command = patch.get("command", command) + patched_cwd = patch.get("workdir", cwd) + if not isinstance(patched_command, str) or not isinstance(patched_cwd, str): + raise ValueError("Shell hook patch must contain string command and workdir") + return await _effect(patched_command, patched_cwd) + + return await run_tool_execution_lifecycle( + payload, + _effect, + patched_effect=_patched_effect, + ) 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/reminders.py b/flocks/session/features/reminders.py index b848fbfa5..cf2f7f9e8 100644 --- a/flocks/session/features/reminders.py +++ b/flocks/session/features/reminders.py @@ -1,9 +1,9 @@ """ Session Reminders data models. -Note: The SessionReminders business logic has been removed as it was dead code. -The _check_reminders() method in session_loop.py is defined but never wired -into _run_loop(). These data classes are kept for potential future use. +The runtime reminder injection path has been removed because it was never +wired into session execution. These exports remain as legacy import +compatibility until that public surface is retired separately. """ from dataclasses import dataclass @@ -35,10 +35,8 @@ class ReminderContext: current_focus: Optional[str] = None -# Minimal stub kept so session_loop._check_reminders() can still import these -# without error, even though _check_reminders() itself is never called. class SessionReminders: - """Reminder manager stub — business logic removed (was dead code).""" + """Compatibility stub for the removed reminder runtime.""" _last_reminder: dict = {} _last_step: dict = {} diff --git a/flocks/session/message.py b/flocks/session/message.py index a8c4b3e20..9eee957b7 100644 --- a/flocks/session/message.py +++ b/flocks/session/message.py @@ -1589,14 +1589,45 @@ async def create( if message.id not in cls._parts_cache[session_id]: cls._parts_cache[session_id][message.id] = [] cls._parts_cache[session_id][message.id].append(part) - - # Persist to storage - await cls._persist_indexed_state( - session_id, - message.id, - include_messages=True, - include_parts=True, + + serialized_cache_existed = ( + session_id in cls._parts_serialized_cache + ) + previous_serialized_parts = dict( + cls._parts_serialized_cache.get(session_id, {}) + ) + persisted_mids_existed = session_id in cls._parts_persisted_mids + previous_persisted_mids = set( + cls._parts_persisted_mids.get(session_id, set()) ) + + # Persistence is the commit point. Restore every in-memory index if + # the atomic storage mutation fails so callers cannot observe a + # message that does not exist in canonical storage. + try: + await cls._persist_indexed_state( + session_id, + message.id, + include_messages=True, + include_parts=True, + ) + except BaseException: + cls._messages_cache[session_id].pop() + cls._msg_id_index[session_id].pop(message.id, None) + cls._parts_cache[session_id].pop(message.id, None) + if serialized_cache_existed: + cls._parts_serialized_cache[session_id] = ( + previous_serialized_parts + ) + else: + cls._parts_serialized_cache.pop(session_id, None) + if persisted_mids_existed: + cls._parts_persisted_mids[session_id] = ( + previous_persisted_mids + ) + else: + cls._parts_persisted_mids.pop(session_id, None) + raise log.info("message.created", { "id": message.id, 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/session/runtime/agent_loop.py b/flocks/session/runtime/agent_loop.py new file mode 100644 index 000000000..21fdfc3d8 --- /dev/null +++ b/flocks/session/runtime/agent_loop.py @@ -0,0 +1,188 @@ +"""The control loop for one logical user input.""" + +from __future__ import annotations + +import asyncio + +from flocks.session.message import MessageInfo +from flocks.session.runtime.contracts import ( + AgentRunOutcome, + AgentRunStatus, + StepAction, + TurnPreparationStatus, +) +from flocks.session.runtime.session_turn import LoopContext +from flocks.session.runtime.step_engine import StepCancelled, StepEngine +from flocks.utils.log import Log + + +log = Log.create(service="session.agent_loop") + + +class AgentLoop: + """Decide whether one logical user input needs another model step.""" + + async def run( + self, + turn: LoopContext, + engine: StepEngine, + ) -> AgentRunOutcome[MessageInfo]: + """Run the current logical input to a session-level boundary.""" + last_user = None + last_message = None + + while not turn.aborted: + preparation = await turn.prepare_step() + if preparation.status == TurnPreparationStatus.CONTINUE: + continue + if preparation.status == TurnPreparationStatus.COMPLETE: + return AgentRunOutcome( + status=AgentRunStatus.COMPLETED, + last_user=last_user, + last_message=preparation.last_message or last_message, + ) + + snapshot = preparation.snapshot + if snapshot is None: + return AgentRunOutcome( + status=AgentRunStatus.FATAL_FAILURE, + last_user=last_user, + last_message=last_message, + error=( + "LoopContext returned READY without a model-turn " + "snapshot" + ), + ) + + last_user = snapshot.last_user + + try: + step_task = asyncio.create_task(engine.run(snapshot)) + turn._current_step_task = step_task + try: + step_result = await step_task + except asyncio.CancelledError as exc: + cleanup_task = asyncio.create_task( + engine.finalize_cancelled_attempt(), + name=( + "step-cancel-cleanup:" + f"{turn.session.id}:{turn.step}" + ), + ) + while True: + try: + await asyncio.shield(cleanup_task) + break + except asyncio.CancelledError: + if cleanup_task.cancelled(): + break + continue + except Exception as cleanup_error: + log.error( + "session.step.cancel_cleanup_failed", + { + "session_id": turn.session.id, + "step": turn.step, + "error": str(cleanup_error), + }, + ) + break + if turn.aborted: + raise StepCancelled from exc + raise + finally: + if turn._current_step_task is step_task: + turn._current_step_task = None + except StepCancelled: + log.info( + "session.step.cancelled", + { + "session_id": turn.session.id, + "step": turn.step, + }, + ) + return AgentRunOutcome( + status=AgentRunStatus.ABORTED, + last_user=last_user, + last_message=last_message, + error="Aborted", + ) + + boundary = await turn.commit_step(step_result) + last_message = boundary.last_message or last_message + + if turn.aborted: + return AgentRunOutcome( + status=AgentRunStatus.ABORTED, + last_user=last_user, + last_message=last_message, + error=step_result.error, + ) + + if boundary.input_available: + return AgentRunOutcome( + status=AgentRunStatus.INPUT_AVAILABLE, + last_user=last_user, + last_message=last_message, + error=step_result.error, + step_result=step_result, + ) + + 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, + last_user=last_user, + last_message=last_message, + error=failure.message, + step_result=step_result, + ) + + if step_result.action == StepAction.CONTINUE: + continue + if step_result.action == StepAction.COMPACT: + return AgentRunOutcome( + status=AgentRunStatus.CONTEXT_OVERFLOW, + last_user=last_user, + last_message=last_message, + error=step_result.error, + step_result=step_result, + ) + if step_result.action != StepAction.STOP: + return AgentRunOutcome( + status=AgentRunStatus.FATAL_FAILURE, + last_user=last_user, + last_message=last_message, + error=f"Unknown step action: {step_result.action}", + step_result=step_result, + ) + if step_result.error: + return AgentRunOutcome( + status=AgentRunStatus.FATAL_FAILURE, + last_user=last_user, + last_message=last_message, + error=step_result.error, + step_result=step_result, + ) + + return AgentRunOutcome( + status=AgentRunStatus.COMPLETED, + last_user=last_user, + last_message=last_message, + step_result=step_result, + ) + + return AgentRunOutcome( + status=AgentRunStatus.ABORTED, + 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 new file mode 100644 index 000000000..e3c196bb5 --- /dev/null +++ b/flocks/session/runtime/continuation_policy.py @@ -0,0 +1,472 @@ +"""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 ContinuationMaterializationError(RuntimeError): + """Raised when a selected continuation cannot be read or persisted.""" + + +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 + + last_assistant = next( + ( + message + for message in reversed(messages) + if message.role == MessageRole.ASSISTANT + ), + None, + ) + if last_assistant is not None: + from flocks.session.runtime.session_turn import ( + is_terminal_assistant_reply, + ) + + assistant_parts = await Message.parts( + last_assistant.id, + context.session.id, + ) + if is_terminal_assistant_reply( + last_user, + last_assistant, + assistant_parts, + ): + context.prepared_user_id = last_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 + 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, + "sessionCategory": context.session.category, + "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 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: + 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": + await self.run_turn_after( + context, + last_user, + last_message, + ) + + 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( + context, + stop_reason=stop_reason, + ) + return ContinuationDecision() + + async def run_turn_after( + self, + context: Any, + last_user: MessageInfo, + last_message: MessageInfo, + ) -> 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) + await HookPipeline.run_turn_after( + { + "sessionID": context.session.id, + "sessionCategory": context.session.category, + "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, + }, + "terminalOutcome": { + "status": "success", + "finish_reason": "stop", + }, + } + ) + except Exception as exc: + log.debug( + "session.hook.turn_after_error", + { + "session_id": context.session.id, + "message_id": getattr(last_message, "id", None), + "error": str(exc), + }, + ) + + @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)}, + ) + raise ContinuationMaterializationError( + "Failed to materialize continuation for " + f"session {context.session.id}: {exc}" + ) from exc + + 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", + }.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/session/runtime/contracts.py b/flocks/session/runtime/contracts.py new file mode 100644 index 000000000..056ca8974 --- /dev/null +++ b/flocks/session/runtime/contracts.py @@ -0,0 +1,265 @@ +"""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 +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 + +import copy +from dataclasses import dataclass, field +from enum import Enum +from types import MappingProxyType +from typing import Any, Generic, Mapping, Optional, TypeVar, cast + + +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 _freeze_provider_message(value: ProviderMessageT) -> ProviderMessageT: + """Freeze container payloads while retaining owned model messages. + + ``ModelRequest`` owns provider model objects such as ``ChatMessage`` for + the lifetime of one logical attempt. Provider adapters treat those + objects as read-only, so copying every model object (and its nested + content) on request construction would only duplicate the full context. + Plain container payloads keep the original isolation contract. + """ + if isinstance(value, (Mapping, list, tuple)): + return cast(ProviderMessageT, _freeze(value)) + return value + + +def _provider_message_view(value: ProviderMessageT) -> ProviderMessageT: + """Return a provider view without cloning owned model messages.""" + if isinstance(value, (Mapping, list, tuple)): + return cast(ProviderMessageT, _thaw(value)) + return value + + +@dataclass(frozen=True) +class RuntimeModel: + """Concrete provider/model selection for one model turn.""" + + provider_id: str + 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(_freeze_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 a fresh list containing read-only request messages.""" + return [_provider_message_view(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 cast(dict[str, Any], _thaw(self.options)) + + +@dataclass +class ActiveModelAttempt(Generic[ProviderMessageT]): + """Single request and hook state retained across bounded model retries.""" + + message_id: str + request: ModelRequest[ProviderMessageT] + hook_metadata: dict[str, Any] + llm_after_enabled: bool + outputs: list[dict[str, Any]] = field(default_factory=list) + hooks_initialized: bool = False + + +@dataclass +class AttemptEffects: + """Observable effects accumulated during one provider attempt.""" + + received_chunk: bool = False + observable_output_started: bool = False + tool_execution_started: 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) + + +@dataclass(frozen=True) +class FailoverDecision: + """Classification used by the session runtime'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] + + +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 runtime 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: 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 + + +@dataclass(frozen=True) +class ModelTurnSnapshot(Generic[MessageT]): + """Immutable input presented to a step engine for one model turn.""" + + active_model: RuntimeModel + trace_step: int + messages: tuple[MessageT, ...] + last_user: MessageT + + +class TurnPreparationStatus(str, Enum): + """Session preparation result before the next model turn.""" + + READY = "ready" + CONTINUE = "continue" + COMPLETE = "complete" + + +@dataclass(frozen=True) +class ModelTurnPreparation(Generic[MessageT]): + """Result of session-owned preparation at a model-turn boundary.""" + + status: TurnPreparationStatus + snapshot: Optional[ModelTurnSnapshot[MessageT]] = None + last_message: Optional[MessageT] = None + + +@dataclass(frozen=True) +class ModelTurnBoundary(Generic[MessageT]): + """Committed session view after one model turn finishes.""" + + last_message: Optional[MessageT] = None + input_available: bool = False + + +@dataclass(frozen=True) +class ContinuationDecision(Generic[MessageT]): + """Session-owned continuation policy result consumed by the outer 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 runtime.""" + + COMPLETED = "completed" + INPUT_AVAILABLE = "input_available" + 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 + 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/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..7b7451e03 --- /dev/null +++ b/flocks/session/runtime/session_turn.py @@ -0,0 +1,1027 @@ +"""State and persistence boundary for one logical session turn. + +Implements model-turn preparation with support for: +- Message processing +- Tool execution +- Compaction +""" + +import asyncio +import time +from typing import Optional, List, Dict, Any, Callable, Awaitable, Literal +from dataclasses import dataclass, field + +from flocks.session.runtime.contracts import ( + ModelTurnBoundary, + ModelTurnPreparation, + ModelTurnSnapshot, + RuntimeModel, + StepAction, + StepResult, + TurnPreparationStatus, +) +from flocks.utils.log import Log +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 + + +def is_terminal_assistant_reply( + last_user: MessageInfo, + last_assistant: Optional[MessageInfo], + last_assistant_parts: Optional[List[Any]] = None, +) -> bool: + """Return whether the latest user already has a completed reply.""" + if last_assistant is None: + return False + if any( + getattr(part, "type", None) == "tool" + for part in (last_assistant_parts or []) + ): + return False + if getattr(last_assistant, "finish", None) not in { + "tool-calls", + "unknown", + "summary", + None, + }: + return getattr(last_assistant, "parentID", None) == last_user.id + return False + + +@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 + # 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 LoopContext: + """Own state and persistence boundaries for one session loop run. + + Supports: + - Logical user-turn and message iteration + - Compaction triggers + """ + + 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 + 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) + + @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.""" + SessionStatus.set(self.session.id, SessionStatusBusy()) + self.step += 1 + 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 + 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: + 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": + pending_compactions.append(part) + log.debug( + "loop.message_scan_complete", + { + "session_id": self.session.id, + "step": self.step, + "compaction_count": len(pending_compactions), + "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, + ) + + self.prepared_user_id = last_user.id + await self._prepare_memory() + self._schedule_title_generation(last_user, messages) + + if pending_compactions: + compaction_preparation = await self._prepare_pending_compaction( + messages, + last_user, + pending_compactions.pop(), + ) + if compaction_preparation is not None: + return compaction_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) + return ModelTurnPreparation( + status=TurnPreparationStatus.READY, + snapshot=ModelTurnSnapshot( + active_model=active_model, + 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 = 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 + for message in reversed(post_messages) + if message.role == MessageRole.ASSISTANT + and last_user is not None + and 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( + last_message=last_message, + input_available=queued_user is not None, + ) + + 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_compaction( + self, + messages: List[MessageInfo], + last_user: MessageInfo, + compaction_part: Any, + ) -> Optional[ModelTurnPreparation[MessageInfo]]: + """Finish persisted compaction work before the model turn.""" + log.info( + "loop.compaction_pending", + { + "session_id": self.session.id, + "step": self.step, + "auto": getattr(compaction_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(compaction_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) + reasoning_tokens = tokens.get("reasoning", 0) + observed_prompt_tokens = input_tokens + cache_read + reported_total = observed_prompt_tokens + output_tokens + reasoning_tokens + if reported_total > 0: + self.last_observed_prompt_tokens = reported_total + + # Provider usage predates the latest assistant response and its tool + # results. Invalidate the estimate cached while tools were running and + # add only the messages created after that observation. + SessionPrompt.invalidate_message_cache(last_finished.id) + last_finished_index = next( + ( + index + for index, message in enumerate(messages) + if message.id == last_finished.id + ), + len(messages) - 1, + ) + if observed_prompt_tokens > 0: + tool_result_tokens = await SessionPrompt.estimate_tool_result_tokens( + self.session.id, + last_finished.id, + ) + later_tokens = await SessionPrompt.estimate_full_context_tokens( + self.session.id, + messages[last_finished_index + 1:], + policy=policy, + ) + estimated_component_tokens = tool_result_tokens + later_tokens + effective_tokens = reported_total + estimated_component_tokens + decision_source = "observed+estimated_delta" + else: + effective_tokens = await SessionPrompt.estimate_full_context_tokens( + self.session.id, + messages, + policy=policy, + ) + estimated_component_tokens = effective_tokens + decision_source = "estimated" + + tokens = { + "input": effective_tokens, + "output": 0, + "cache": {"read": 0, "write": 0}, + } + log.info( + "loop.tokens_decision", + { + "session_id": self.session.id, + "source": decision_source, + "effective_tokens": effective_tokens, + "observed_tokens": reported_total, + "estimated_component_tokens": estimated_component_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 + """ + return is_terminal_assistant_reply( + last_user, + last_assistant, + last_assistant_parts, + ) diff --git a/flocks/session/runner.py b/flocks/session/runtime/step_engine.py similarity index 82% rename from flocks/session/runner.py rename to flocks/session/runtime/step_engine.py index b07a35ed3..e9b21c12f 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 @@ -17,13 +8,31 @@ import sys import time from collections.abc import Mapping +from dataclasses import replace from datetime import datetime -from typing import Optional, Dict, Any, List, Callable, Awaitable, Tuple -from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple import httpcore import httpx +from flocks.session.runtime.contracts import ( + ActiveModelAttempt, + AttemptEffects, + FailoverDecision, + ModelRequest, + ModelTurnSnapshot, + RuntimeModel, + StepFailure, + StepResult, + ToolCall, +) +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 @@ -43,9 +52,7 @@ from flocks.session.lifecycle.compaction import SessionCompaction, CompactionPolicy from flocks.session.llm_hook_utils import ( StreamingTextReplacementBuffer, - apply_hook_request_output, restore_value_with_replacements, - serialize_chat_message, stream_text_replacements_from_hook_output, ) from flocks.session.streaming.stream_processor import StreamProcessor @@ -92,11 +99,12 @@ 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 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: @@ -122,6 +130,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 @@ -221,101 +231,21 @@ 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] +class StepCancelled(Exception): + """Signal that the user cancelled the active session step.""" -@dataclass -class LlmAttemptState: - """Observable side effects accumulated across retries for one model.""" +class StepEngine: + """Own one complete model/tool step, including retries and failover.""" - 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 - - -@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 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'] = {} - 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 memory_bootstrap_data: Optional[Dict[str, Any]] = None, static_cache: Optional[Dict[str, Any]] = None, defer_step_errors: bool = False, @@ -328,12 +258,10 @@ def __init__( 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._abort = asyncio.Event() - self._external_abort = abort_event # External abort event (e.g. from SessionLoop) + self.callbacks = callbacks or LoopCallbacks() + self._abort = abort_event or asyncio.Event() 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._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 @@ -341,7 +269,261 @@ 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._active_model_attempt: Optional[ + ActiveModelAttempt[ChatMessage] + ] = None + self._llm_retry_scope_active = False + self._turn: Optional[Any] = None + self._model_policy: ModelRoutingPolicy = DEFAULT_MODEL_ROUTING_POLICY + self._step_agent: Optional[AgentInfo] = None + self._frozen_tools: Optional[List[Dict[str, Any]]] = None + + @classmethod + def from_turn( + cls, + turn: Any, + model_policy: Optional[ModelRoutingPolicy] = None, + ) -> "StepEngine": + """Create the production engine for one stateful ``LoopContext``.""" + 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, + 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_tools = None + while True: + if turn.aborted: + raise StepCancelled + active_model = RuntimeModel(turn.provider_id, turn.model_id) + result = await self._run_candidate( + replace(snapshot, active_model=active_model), + ) + if turn.aborted: + raise StepCancelled + failure = result.failure + if not turn.auto_failover or failure is None: + return result + + 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 + + if turn.aborted: + raise StepCancelled + if not await self._remove_failed_attempt(failure): + await turn.finalize_failure(failure, snapshot.last_user) + return result + + if turn.aborted: + raise StepCancelled + 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 + + 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 + self._memory_bootstrap_data = turn.memory_bootstrap_data + + started_at = asyncio.get_running_loop().time() + try: + result = await self._process_step( + list(snapshot.messages), + snapshot.last_user, + ) + 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: + self._clear_model_request_state() + 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 _clear_model_request_state(self) -> None: + """Release request snapshots owned by the completed candidate.""" + self._active_model_attempt = None + self._llm_retry_scope_active = False + + 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() + if turn.aborted: + raise StepCancelled + 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: @@ -385,8 +567,6 @@ 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({ "sessionID": self.session.id, "workspace": self.session.directory, @@ -678,7 +858,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 = Provider.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: @@ -857,9 +1040,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 = Provider.get(self.provider_id) if provider is not None: for model in getattr(provider, "_config_models", []) or []: if model.id == self.model_id: @@ -977,292 +1158,14 @@ 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() - - async def _effect( - execution_command: str = command, - execution_cwd: str = cwd, - ) -> Dict[str, Any]: - 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( - execution_command, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - cwd=execution_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": execution_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": execution_command}, - "output": output, - }, - }], - } - - from flocks.session.tool_execution import ( - build_session_tool_execution_payload, - run_tool_execution_lifecycle, - ) - - payload = await build_session_tool_execution_payload( - session_id=session_id, - message_id=Identifier.create("message"), - agent=agent, - tool_name="shell", - tool_input={"command": command, "workdir": cwd}, - validated_input={"command": command, "workdir": cwd}, - tool_schema={ - "type": "object", - "properties": { - "command": {"type": "string"}, - "workdir": {"type": "string"}, - }, - "required": ["command"], - }, - tool_context_extra={ - "tool_source": "session_runner", - "tool_category": "command", - "workspace_dir": cwd, - "session_execution_profile": { - "entry": "session.shell", - "workspace_dir": cwd, - }, - }, - ) - - async def _patched_effect(patch: Mapping[str, Any]) -> Dict[str, Any]: - patched_command = patch.get("command", command) - patched_cwd = patch.get("workdir", cwd) - if not isinstance(patched_command, str) or not isinstance(patched_cwd, str): - raise ValueError("Shell hook patch must contain string command and workdir") - return await _effect(patched_command, patched_cwd) - - return await run_tool_execution_lifecycle( - payload, - _effect, - patched_effect=_patched_effect, - ) - 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).""" - if self._abort.is_set(): - return True - if self._external_abort is not None and self._external_abort.is_set(): - return True - return False + """Check if abort was signaled.""" + return self._abort.is_set() @staticmethod def classify_failover_error(error: Dict[str, Any]) -> FailoverDecision: @@ -1361,7 +1264,7 @@ def _deferred_failure_result( decision: FailoverDecision, attempts: int, ) -> StepResult: - state = LlmAttemptState( + 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, @@ -1379,14 +1282,14 @@ def _deferred_failure_result( attempts=attempts, ), ) - + async def _process_step( self, messages: List[MessageInfo], 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) ) @@ -1397,27 +1300,14 @@ 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 = 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 # Track session agent (Flocks compatibility) try: @@ -1425,11 +1315,11 @@ 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 = Provider.get(self.provider_id) if not provider: @@ -1443,7 +1333,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( @@ -1459,8 +1349,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) @@ -1471,7 +1359,7 @@ async def _process_step( "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: @@ -1483,7 +1371,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( @@ -1499,13 +1387,14 @@ 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 tools_started_at = time.perf_counter() - tools = await self._build_callable_tool_schema(agent, messages) + if self._frozen_tools is not None: + tools = copy.deepcopy(self._frozen_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) @@ -1549,7 +1438,11 @@ async def device_asset_prompt_factory() -> Optional[str]: device_revision=current_device_revision, 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_prompts), + ) await self._run_session_start_hook(agent) @@ -1579,7 +1472,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", { @@ -1599,7 +1492,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_prompts, + ) finally: if previous_queued_user_ids is None: if hasattr(self, "_queued_user_message_ids"): @@ -1619,7 +1515,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: @@ -1641,7 +1537,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 @@ -1649,17 +1545,25 @@ 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_tools is None: + self._frozen_tools = request.provider_tools() + + # Create the persisted assistant attempt after the request is frozen. assistant_msg = await Message.create( session_id=self.session.id, role=MessageRole.ASSISTANT, @@ -1669,7 +1573,12 @@ async def device_asset_prompt_factory() -> Optional[str]: provider_id=self.provider_id, parent_id=last_user.id, ) - + self._active_model_attempt = ActiveModelAttempt( + message_id=assistant_msg.id, + request=request, + hook_metadata={}, + llm_after_enabled=False, + ) # Publish assistant message SSE event so frontends can show the message card if self.callbacks.event_publish_callback: import time as _time @@ -1687,7 +1596,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 @@ -1699,6 +1608,7 @@ async def device_asset_prompt_factory() -> Optional[str]: MAX_EMPTY_RETRIES = 3 error_attempt = 0 empty_attempt = 0 + self._llm_retry_scope_active = True while not self.is_aborted: try: @@ -1715,6 +1625,7 @@ async def device_asset_prompt_factory() -> Optional[str]: ) if getattr(self, "_llm_call_aborted", False): + await self._emit_pending_llm_after(assistant_msg.id) await Message.update( self.session.id, assistant_msg.id, @@ -1794,16 +1705,15 @@ async def device_asset_prompt_factory() -> Optional[str]: "attempts": empty_attempt, }, } + await self._emit_pending_llm_after(assistant_msg.id) if self._defer_step_errors: return self._deferred_failure_result( 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: - await self.callbacks.on_error(empty_error_msg) await Message.update( self.session.id, assistant_msg.id, @@ -1850,6 +1760,7 @@ async def device_asset_prompt_factory() -> Optional[str]: finish = "tool-calls" if result.tool_calls else "stop" await Message.update(self.session.id, assistant_msg.id, finish=finish) await self._record_usage_if_available(result.usage, message_id=assistant_msg.id) + await self._emit_pending_llm_after(assistant_msg.id) # Note: Compaction check is now done in the main loop (run()) before processing step # This matches Flocks's logic: check lastFinished.tokens at loop start @@ -1902,7 +1813,7 @@ async def device_asset_prompt_factory() -> Optional[str]: "reason": retry_message, "max_retries": retry_limit, }) - + # Set retry status SessionStatus.set( self.session.id, @@ -1912,10 +1823,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: @@ -1943,6 +1854,7 @@ async def device_asset_prompt_factory() -> Optional[str]: final_error_message = CONNECTION_ERROR_DISPLAY_MESSAGE error_dict["data"]["displayMessage"] = CONNECTION_ERROR_DISPLAY_MESSAGE + await self._emit_pending_llm_after(assistant_msg.id) if self._defer_step_errors: return self._deferred_failure_result( message=final_error_message, @@ -1952,9 +1864,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, @@ -1970,10 +1879,11 @@ 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 + await self._emit_pending_llm_after(assistant_msg.id) return StepResult(action="stop", error="Aborted") @staticmethod @@ -2099,7 +2009,6 @@ 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: @@ -2334,7 +2243,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, @@ -2396,7 +2305,7 @@ 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 = ToolRegistry.get(tool_name) @@ -2404,11 +2313,11 @@ def _agent_declares_tool(self, agent: AgentInfo, tool_name: str) -> bool: 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 = { @@ -2433,7 +2342,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 @@ -2472,11 +2381,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) @@ -2485,7 +2394,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 [ @@ -2500,13 +2409,16 @@ 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, _, _ = Provider.resolve_model_info(self.provider_id, self.model_id) + ctx, _, _ = Provider.resolve_model_info( + self.provider_id, + self.model_id, + ) if ctx and ctx > 0: return ctx except Exception: @@ -2516,12 +2428,11 @@ def _get_context_window_tokens(self) -> int: def _message_conversion_cache_key( self, msg: MessageInfo, - parts: List[Any], *, + has_file_part: bool, is_latest_user_turn: bool, ) -> Tuple[Any, ...]: role = msg.role if isinstance(msg.role, str) else getattr(msg.role, "value", None) - has_file_part = any(getattr(part, "type", None) == "file" for part in parts) latest_user_marker = msg.id if (role == "user" and has_file_part and is_latest_user_turn) else ( "stale-file-user" if role == "user" and has_file_part else None ) @@ -2620,10 +2531,6 @@ def _wrap_queued_user_blocks( } return wrapped_blocks - @staticmethod - def _clone_cached_chat_messages(payloads: List[Dict[str, Any]]) -> List[ChatMessage]: - return [ChatMessage.model_validate(copy.deepcopy(payload)) for payload in payloads] - def _build_tool_output_text(self, part: Any, tool_name: str, ctx_window_tokens: int) -> Tuple[str, bool, bool]: state = getattr(part, "state", None) metadata = dict(getattr(state, "metadata", None) or {}) if state is not None else {} @@ -2696,7 +2603,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 @@ -2708,7 +2615,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 = Provider.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) @@ -2728,18 +2638,8 @@ async def _to_chat_messages( if _role == "user": last_user_msg_id = _msg.id - preloaded_parts: List[List[Any]] = [] message_signatures: List[Tuple[Any, ...]] = [] - for msg in messages: - parts = await Message.parts(msg.id, self.session.id) - preloaded_parts.append(parts) - message_signatures.append( - self._message_conversion_cache_key( - msg, - parts, - is_latest_user_turn=(msg.id == last_user_msg_id), - ) - ) + message_has_file_parts: List[bool] = [] system_content = self._build_system_message_content(system_prompts) if system_prompts else None system_cache_key = json.dumps(system_content, ensure_ascii=False, sort_keys=True, default=str) @@ -2748,15 +2648,36 @@ async def _to_chat_messages( resume_message_index = 0 if cached_context and cached_context.get("system_cache_key") == system_cache_key: cached_signatures = list(cached_context.get("message_signatures") or []) - if len(cached_signatures) <= len(message_signatures): + cached_has_file_parts = list( + cached_context.get("message_has_file_parts") or [], + ) + cached_chat_messages = cached_context.get("chat_messages") or () + cached_tool_result_refs = cached_context.get("tool_result_refs") or () + cache_shape_valid = ( + len(cached_signatures) == len(cached_has_file_parts) + and len(cached_signatures) <= len(messages) + and all( + isinstance(message, ChatMessage) + for message in cached_chat_messages + ) + ) + if cache_shape_valid: prefix_matches = True for idx, cached_signature in enumerate(cached_signatures): - if cached_signature != message_signatures[idx]: + current_signature = self._message_conversion_cache_key( + messages[idx], + has_file_part=cached_has_file_parts[idx], + is_latest_user_turn=(messages[idx].id == last_user_msg_id), + ) + if cached_signature != current_signature: prefix_matches = False break if prefix_matches: - chat_messages = self._clone_cached_chat_messages(cached_context.get("chat_messages") or []) + chat_messages = list(cached_chat_messages) + tool_result_refs = list(cached_tool_result_refs) resume_message_index = len(cached_signatures) + message_signatures = cached_signatures + message_has_file_parts = cached_has_file_parts turn_index = sum( 1 for msg in messages[:resume_message_index] @@ -2775,7 +2696,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: @@ -2785,8 +2706,20 @@ async def _to_chat_messages( turn_index += 1 is_latest_user_turn = msg.id == last_user_msg_id # Get message parts - parts = preloaded_parts[idx] - + parts = await Message.parts(msg.id, self.session.id) + has_file_part = any( + getattr(part, "type", None) == "file" + for part in parts + ) + message_has_file_parts.append(has_file_part) + message_signatures.append( + self._message_conversion_cache_key( + msg, + has_file_part=has_file_part, + is_latest_user_turn=is_latest_user_turn, + ) + ) + if not parts: # Fallback: use text content only content = await Message.get_text_content(msg) @@ -2799,7 +2732,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 @@ -2850,12 +2783,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" @@ -2875,7 +2802,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 @@ -2885,7 +2812,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( @@ -2899,7 +2826,7 @@ async def _to_chat_messages( else: # Non-AbortedError - skip continue - + assistant_content_parts = [] assistant_reasoning_parts = [] assistant_reasoning_content_parts = [] @@ -2910,11 +2837,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): @@ -2965,13 +2892,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, @@ -2985,7 +2912,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({ @@ -3013,7 +2940,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, @@ -3021,7 +2948,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 @@ -3039,7 +2966,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 @@ -3064,7 +2991,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 @@ -3099,7 +3026,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", { @@ -3117,10 +3044,9 @@ async def _to_chat_messages( context_cache["latest"] = { "system_cache_key": system_cache_key, "message_signatures": list(message_signatures), - "chat_messages": [ - message.model_dump(exclude_none=True) - for message in chat_messages - ], + "message_has_file_parts": list(message_has_file_parts), + "chat_messages": tuple(chat_messages), + "tool_result_refs": tuple(tool_result_refs), } self._log_perf( "runner.to_chat_messages.complete", @@ -3128,9 +3054,136 @@ 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"), + ), + ) + 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, + 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, @@ -3141,10 +3194,30 @@ 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. """ + active_attempt = self._active_model_attempt + if active_attempt is None: + request = self._build_model_request( + messages=messages, + tools=tools, + agent=agent, + ) + active_attempt = ActiveModelAttempt( + message_id=assistant_msg.id, + request=request, + hook_metadata={}, + llm_after_enabled=False, + ) + self._active_model_attempt = active_attempt + elif active_attempt.message_id != assistant_msg.id: + raise RuntimeError( + "A different model attempt is already active", + ) + else: + request = active_attempt.request def _build_llm_response_payload( *, content: str, @@ -3165,6 +3238,77 @@ 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, + }, + } + if not active_attempt.hooks_initialized: + 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.error("runner.hook.stage_probe.error", {"error": str(exc)}) + raise RuntimeError( + "LLM hook stage probe failed; request was not sent", + ) from exc + if llm_before_enabled: + try: + request = await self._apply_before_model_hook( + request, + llm_hook_metadata, + ) + except Exception as exc: + log.error("runner.hook.llm_before.error", {"error": str(exc)}) + raise RuntimeError( + "LLM before-hook failed; request was not sent", + ) from exc + active_attempt.request = request + active_attempt.hook_metadata = dict(llm_hook_metadata) + active_attempt.llm_after_enabled = llm_after_enabled + active_attempt.hooks_initialized = True + else: + request = active_attempt.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 try: @@ -3184,6 +3328,13 @@ 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: + 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) @@ -3191,12 +3342,12 @@ async def _on_tool_execution_start( session_id=self.session.id, assistant_message=assistant_msg, agent=agent, - abort_event=self._external_abort or self._abort, + abort_event=self._abort, permission_callback=self._handle_permission, 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, @@ -3214,11 +3365,35 @@ 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) + + 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 + 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'): @@ -3233,131 +3408,11 @@ async def _on_tool_execution_start( reasoning_id_counter = 0 stream_finish_reason: Optional[str] = None + # -- Observability: create trace & generation scopes (safe no-op when + # Langfuse is unconfigured). All observability calls are wrapped in + # try/except so they never break the core session flow. 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: - log.error("runner.call_llm.no_messages", { - "total_messages": len(messages), - "session_id": self.session.id, - }) - 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 - provider_tools = None if self._should_use_text_tool_call_mode() else (tools if tools else None) - if provider_tools is None and tools: - log.info("runner.text_tool_call_mode.enabled", { - "session_id": self.session.id, - "provider_id": self.provider_id, - "model_id": self.model_id, - "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 - replacements: list[tuple[str, str]] = [] - stream_text_rewriter: Optional[StreamingTextReplacementBuffer] = None - stream_reasoning_rewriter: Optional[StreamingTextReplacementBuffer] = None - self._llm_call_aborted = False - - 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, - )) - - 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.error("runner.hook.stage_probe.error", {"error": str(exc)}) - raise RuntimeError("LLM hook stage probe failed; request was not sent") from exc - - if llm_before_enabled: - llm_before_hook_input = { - **llm_hook_metadata, - "request": { - "messageCount": len(messages), - "messages": [serialize_chat_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() - llm_before_ctx = await HookPipeline.run_llm_before(llm_before_hook_input) - hook_output = getattr(llm_before_ctx, "output", None) or {} - replacements = stream_text_replacements_from_hook_output(hook_output) - if replacements: - stream_text_rewriter = StreamingTextReplacementBuffer(replacements) - stream_reasoning_rewriter = StreamingTextReplacementBuffer(replacements) - updated_request = hook_output.get("request") - if isinstance(updated_request, dict): - messages, provider_options = apply_hook_request_output( - messages, - provider_options, - hook_output, - ) - updated_tools = updated_request.get("tools") - if isinstance(updated_tools, list): - tools = copy.deepcopy(updated_tools) - provider_tools = None if self._should_use_text_tool_call_mode() else (tools if tools else None) - self._log_perf( - "runner.hook.llm_before.complete", - hook_started_at, - message_count=len(messages), - tool_count=len(tools), - ) - except Exception as exc: - log.error("runner.hook.llm_before.error", {"error": str(exc)}) - raise RuntimeError("LLM before-hook failed; request was not sent") from exc - - # -- Observability: create trace & generation scopes after llm_before, - # so previews use the same redacted messages that will be sent to the provider. - # All observability calls are wrapped in try/except so they never break - # the core session flow. if langfuse_is_active(): try: trace_tags = [ @@ -3375,7 +3430,7 @@ async def _flush_reasoning_rewriter() -> None: 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, @@ -3415,6 +3470,41 @@ async def _flush_reasoning_rewriter() -> None: 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: + log.error("runner.call_llm.no_messages", { + "total_messages": len(messages), + "session_id": self.session.id, + }) + 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", { + "session_id": self.session.id, + "provider_id": self.provider_id, + "model_id": self.model_id, + "tool_count": len(tools), + }) + + self._llm_call_aborted = False + llm_call_started_at = time.perf_counter() first_chunk_logged = False aborted_during_stream = False @@ -3547,7 +3637,9 @@ async def _flush_reasoning_rewriter() -> None: if chunk_reasoning: if stream_reasoning_rewriter is not None: - reasoning_text = stream_reasoning_rewriter.feed(reasoning_text) + reasoning_text = stream_reasoning_rewriter.feed( + reasoning_text, + ) if reasoning_text: await processor.process_event(ReasoningDeltaEvent( id=self._current_reasoning_id, @@ -3594,9 +3686,74 @@ async def _flush_reasoning_rewriter() -> None: for tc in chunk_tool_calls: await tool_accumulator.feed_chunk(tc) except asyncio.CancelledError: - # Foreground delegate tasks own child sessions. Let their - # cancellation/finalization finish before unwinding this step. - await processor.drain_parallel_tool_calls() + # Cancellation is still a terminal provider attempt. Complete the + # durable assistant state and the matching after-hook before the + # step task unwinds; repeated cancellation must not interrupt this + # cleanup and leave an unfinished assistant in history. + async def _finalize_cancelled_attempt() -> None: + await processor.drain_parallel_tool_calls() + partial_response = _build_llm_response_payload( + content=processor.get_text_content(), + reasoning=processor.get_reasoning_content(), + tool_calls=[], + ) + await self._record_llm_after_attempt( + message_id=assistant_msg.id, + output={ + "durationMs": int( + (time.perf_counter() - llm_call_started_at) * 1000 + ), + "error": { + "type": "CancelledError", + "message": "Model request cancelled", + }, + "response": partial_response, + "usage": stream_usage, + "chunkCounts": dict(chunk_counts), + }, + ) + await Message.update( + self.session.id, + assistant_msg.id, + error=self._build_message_aborted_error(), + finish="error", + ) + await self._record_usage_if_available( + stream_usage, + message_id=assistant_msg.id, + ) + await self._emit_pending_llm_after(assistant_msg.id) + self._end_observability( + generation_ctx, + trace_ctx, + output=partial_response, + usage=stream_usage, + metadata={"status": "cancelled"}, + trace_output=partial_response, + level="WARNING", + ) + + cleanup_task = asyncio.create_task( + _finalize_cancelled_attempt(), + name=f"llm-cancel-cleanup:{assistant_msg.id}", + ) + while True: + try: + await asyncio.shield(cleanup_task) + break + except asyncio.CancelledError: + if cleanup_task.cancelled(): + break + continue + except Exception as cleanup_error: + log.error( + "runner.llm.cancel_cleanup_failed", + { + "message_id": assistant_msg.id, + "error": str(cleanup_error), + }, + ) + break raise except Exception as exc: # A foreground delegate may already be running when the provider @@ -3608,25 +3765,21 @@ async def _flush_reasoning_rewriter() -> None: reasoning=processor.get_reasoning_content(), tool_calls=[], ) - if llm_after_enabled: - try: - await HookPipeline.run_llm_after( - llm_hook_metadata, - { - "durationMs": int((time.perf_counter() - llm_call_started_at) * 1000), - "error": { - "type": type(exc).__name__, - "message": str(exc), - }, - "response": partial_response, - "usage": stream_usage, - "chunkCounts": dict(chunk_counts), - }, - ) - except Exception as hook_exc: - log.debug("runner.hook.llm_after.error", {"error": str(hook_exc)}) + await self._record_llm_after_attempt( + message_id=assistant_msg.id, + output={ + "durationMs": int((time.perf_counter() - llm_call_started_at) * 1000), + "error": { + "type": type(exc).__name__, + "message": str(exc), + }, + "response": partial_response, + "usage": stream_usage, + "chunkCounts": dict(chunk_counts), + }, + ) raise - + log.debug("runner.stream.summary", { "total_chunks": chunk_counts["total"], "reasoning_chunks": chunk_counts["reasoning"], @@ -3646,11 +3799,11 @@ async def _flush_reasoning_rewriter() -> None: 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() @@ -3662,7 +3815,7 @@ async def _flush_reasoning_rewriter() -> None: 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() @@ -3672,11 +3825,11 @@ async def _flush_reasoning_rewriter() -> None: # 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: @@ -3693,7 +3846,7 @@ async def _flush_reasoning_rewriter() -> None: }) 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), @@ -3701,7 +3854,7 @@ async def _flush_reasoning_rewriter() -> None: "tool_calls": len(processor.tool_calls), "usage": stream_usage, }) - + # Update assistant message with content if content: await Message.update( @@ -3710,7 +3863,7 @@ async def _flush_reasoning_rewriter() -> None: 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 = [ @@ -3728,36 +3881,25 @@ async def _flush_reasoning_rewriter() -> None: reasoning=reasoning, tool_calls=tool_calls_for_result, ) - if llm_after_enabled: - try: - hook_started_at = time.perf_counter() - await HookPipeline.run_llm_after( - llm_hook_metadata, - { - "durationMs": int((time.perf_counter() - llm_call_started_at) * 1000), - "finishReason": processor.get_finish_reason(), - "contentLength": len(content), - "reasoningLength": len(reasoning), - "toolCallCount": len(tool_calls_for_result), - "toolCalls": [ - {"id": tool_call.id, "name": tool_call.name} - for tool_call in tool_calls_for_result[:30] - ], - "response": response_payload, - "usage": stream_usage, - "chunkCounts": dict(chunk_counts), - "action": result_action, - }, - ) - self._log_perf( - "runner.hook.llm_after.complete", - hook_started_at, - action=result_action, - tool_call_count=len(tool_calls_for_result), - ) - except Exception as exc: - log.debug("runner.hook.llm_after.error", {"error": str(exc)}) - + await self._record_llm_after_attempt( + message_id=assistant_msg.id, + output={ + "durationMs": int((time.perf_counter() - llm_call_started_at) * 1000), + "finishReason": processor.get_finish_reason(), + "contentLength": len(content), + "reasoningLength": len(reasoning), + "toolCallCount": len(tool_calls_for_result), + "toolCalls": [ + {"id": tool_call.id, "name": tool_call.name} + for tool_call in tool_calls_for_result[:30] + ], + "response": response_payload, + "usage": stream_usage, + "chunkCounts": dict(chunk_counts), + "action": result_action, + }, + ) + if tool_calls_for_result: response_payload = self._build_langfuse_response_payload( action="continue", @@ -3783,7 +3925,7 @@ async def _flush_reasoning_rewriter() -> None: tool_calls=tool_calls_for_result, usage=stream_usage, ) - + response_payload = self._build_langfuse_response_payload( action="stop", content=content, @@ -3803,7 +3945,88 @@ async def _flush_reasoning_rewriter() -> None: trace_output=response_payload, ) return StepResult(action=result_action, content=content, usage=stream_usage) - + + async def _record_llm_after_attempt( + self, + *, + message_id: str, + output: Dict[str, Any], + ) -> None: + """Record one provider attempt and close direct calls immediately.""" + active_attempt = self._active_model_attempt + if active_attempt is None or active_attempt.message_id != message_id: + raise RuntimeError( + "LLM output does not match the active model attempt", + ) + if active_attempt.llm_after_enabled: + active_attempt.outputs.append(output) + if not self._llm_retry_scope_active: + await self._emit_pending_llm_after(message_id) + + async def finalize_cancelled_attempt(self) -> None: + """Close an active attempt cancelled outside the provider stream.""" + active_attempt = self._active_model_attempt + if active_attempt is None: + return + message_id = active_attempt.message_id + await self._record_llm_after_attempt( + message_id=message_id, + output={ + "error": { + "type": "CancelledError", + "message": "Model step cancelled", + }, + "response": { + "role": "assistant", + "content": "", + "reasoning": "", + "toolCalls": [], + }, + }, + ) + try: + await Message.update( + self.session.id, + message_id, + error=self._build_message_aborted_error(), + finish="error", + ) + finally: + await self._emit_pending_llm_after(message_id) + + async def _emit_pending_llm_after(self, message_id: str) -> None: + """Emit one terminal after-hook for a logical model request.""" + self._llm_retry_scope_active = False + active_attempt = self._active_model_attempt + if active_attempt is None or active_attempt.message_id != message_id: + return + self._active_model_attempt = None + if not active_attempt.llm_after_enabled: + return + metadata = active_attempt.hook_metadata + attempts = active_attempt.outputs + if not attempts: + return + output = dict(attempts[-1]) + output["attemptCount"] = len(attempts) + output["failedAttempts"] = [ + dict(attempt) + for attempt in attempts[:-1] + if attempt.get("error") is not None + ] + try: + hook_started_at = time.perf_counter() + await HookPipeline.run_llm_after(metadata, output) + self._log_perf( + "runner.hook.llm_after.complete", + hook_started_at, + action=output.get("action"), + tool_call_count=output.get("toolCallCount", 0), + attempt_count=len(attempts), + ) + except Exception as exc: + log.debug("runner.hook.llm_after.error", {"error": str(exc)}) + @staticmethod def _end_observability( generation_ctx: Any, @@ -3883,41 +4106,3 @@ async def _handle_permission(self, request) -> None: ) if reply in {"deny", "reject", "never"}: raise PermissionError(f"Permission denied: {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/session.py b/flocks/session/session.py index d4352c36b..ea7a575b1 100644 --- a/flocks/session/session.py +++ b/flocks/session/session.py @@ -1073,13 +1073,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_loop.py b/flocks/session/session_loop.py index 53af1e89f..2fc1535b6 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -1,576 +1,98 @@ -""" -Session Loop Module +"""Public entry point and lifecycle owner for session execution.""" -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. -""" +from __future__ import annotations 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.utils.log import Log -from flocks.utils.id import Identifier +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, + AgentRunStatus, + RuntimeModel, +) +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 ( + LoopContext, + LoopCallbacks, + LoopResult, +) +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, SessionStatusIdle -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, - 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") - -MAX_OVERFLOW_COMPACTION_ATTEMPTS = 3 -POST_COMPACTION_COOLDOWN_STEPS = 2 -RATE_LIMIT_COOLDOWN_SECONDS = 60.0 -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.""" - - 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 + reasoning`` reported by the provider on - # the most recent finished assistant turn. Overflow decisions compare it - # with a current message estimate so tool output produced after that model - # call cannot be missed. - 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 - - @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) +class _SessionLease: + """One process-local ownership record.""" + session_id: str + turn: LoopContext -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] = {} - @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) +class _SessionLeaseRegistry: + """Keep lease bookkeeping out of the SessionLoop control flow.""" - @classmethod - async def validate_runtime_model( - 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" + def __init__(self, active_turns: MutableMapping[str, LoopContext]): + self._active_turns = active_turns - @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) - - 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 - - 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 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) - ) - candidates.append(selected) - return candidates + def get(self, session_id: str) -> Optional[LoopContext]: + return self._active_turns.get(session_id) - @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] - - @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 - - 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" - - @classmethod - def _active_cooldown_model( - cls, + def acquire( + self, 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) + turn: LoopContext, + ) -> Optional[_SessionLease]: + if session_id in self._active_turns: return None - return cooldown.model + self._active_turns[session_id] = turn + return _SessionLease(session_id=session_id, turn=turn) - @classmethod - def _cooldown_candidate_index( - 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 - - @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 - - @classmethod - def get_context(cls, session_id: str) -> Optional[LoopContext]: - """Get loop context for session""" - return cls._active_loops.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 - - @classmethod - def abort_children(cls, parent_session_id: str) -> int: - """Abort all child loops whose session.parent_id matches, recursively.""" - aborted = 0 - child_ids = [ - sid for sid, ctx in list(cls._active_loops.items()) - if getattr(ctx.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() - aborted += 1 - aborted += cls.abort_children(sid) - 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), - }) - - @classmethod - async def _publish_session_notice( - cls, - callbacks: "LoopCallbacks", - session_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)}) + 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) - @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 - ) + def owns(self, lease: _SessionLease) -> bool: + return self._active_turns.get(lease.session_id) is lease.turn - @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 +class SessionLoop: + """Decide whether a persistent session should continue or settle.""" + + _active_turns: ClassVar[dict[str, LoopContext]] = {} + _leases: ClassVar[_SessionLeaseRegistry] = _SessionLeaseRegistry( + _active_turns, + ) + _model_policy: ClassVar[ModelRoutingPolicy] = DEFAULT_MODEL_ROUTING_POLICY + _continuation_policy: ClassVar[ContinuationPolicy] = ( + DEFAULT_CONTINUATION_POLICY + ) + _release_tasks: ClassVar[dict[str, asyncio.Task[bool]]] = {} - 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, @@ -582,225 +104,193 @@ 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 + """Run one session until queued and synthetic continuations settle.""" + await cls._await_pending_release(session_id) + 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", ) - - # Get session + session = await Session.get_by_id(session_id) - if not session: - log.warning("loop.session_not_found", {"session_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("loop.session_not_active", { - "session_id": session_id, - "status": session.status, - }) + 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}) - - # Resolve model when not explicitly provided + session = session.model_copy( + update={"directory": working_directory}, + ) + if not provider_id or not model_id: resolved_provider, resolved_model = await cls._resolve_model( - session, provider_id, model_id + 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") + getattr(session, "category", "user"), ) ) + session.provider = provider_id + session.model = model_id - # 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( + trace_offset = await cls._load_trace_offset(session_id) + runtime_callbacks = callbacks or LoopCallbacks() + turn = LoopContext( session=session, provider_id=provider_id, model_id=model_id, agent_name=agent_name or session.agent or "rex", - session_ctx=session_ctx, + callbacks=runtime_callbacks, + session_store=DefaultSessionContext(session), trace_step_offset=trace_offset, auto_failover=auto_failover, auto_failover_allowed=auto_failover, - model_candidates=model_candidates, - candidate_index=candidate_index, + model_candidates=[primary_model], + candidate_index=0, session_start_pending=trace_offset == 0, + model_policy=cls._model_policy, + continuation_policy=cls._continuation_policy, ) - - # 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 + 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 - 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, - ) + await cls._mark_busy(session_id, runtime_callbacks) + await cls._recover_orphan_tools(session_id) + return await cls._run_owned_loop(lease, runtime_callbacks) 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 + if cls._leases.owns(lease): + await cls._release_session(lease, runtime_callbacks) + + @classmethod + async def _run_owned_loop( + cls, + lease: _SessionLease, + callbacks: LoopCallbacks, + ) -> LoopResult: + """Drive the session while this process owns its lease.""" + turn = lease.turn + processed_user_id: Optional[str] = None + while True: + continuation_policy = ( + turn.continuation_policy or cls._continuation_policy + ) try: - from flocks.bus.bus import Bus - from flocks.bus.events import SessionIdle - await Bus.publish(SessionIdle, {"sessionID": session_id}) + await continuation_policy.prepare_logical_turn(turn) + processed_user_id = ( + turn.prepared_user_id or processed_user_id + ) + outcome = await AgentLoop().run( + turn, + StepEngine.from_turn(turn), + ) + if await cls._should_continue( + turn, + continuation_policy, + outcome, + ): + continue except Exception as exc: - log.warn("loop.idle.event_error", {"error": str(exc)}) - + outcome = await cls._handle_execution_error( + turn, + exc, + ) + + if await cls._settle_or_continue( + lease, + callbacks, + processed_user_id, + allow_late_input=not outcome.unhandled_error, + ): + continue + + return cls._to_loop_result(turn, outcome) + + @staticmethod + async def _should_continue( + turn: LoopContext, + 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 is_running(cls, session_id: str) -> bool: + """Return whether this process owns the session.""" + return session_id in cls._active_turns + + @classmethod + def abort(cls, session_id: str) -> bool: + """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 active descendants of one parent session.""" + aborted = 0 + child_ids = [ + session_id + for session_id, turn in list(cls._active_turns.items()) + if getattr(turn.session, "parent_id", None) == parent_session_id + ] + 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(session_id) + return aborted + + @classmethod + 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 validate_auto_configuration(cls) -> tuple[bool, str]: + """Validate that Auto mode has an available primary model.""" + return await cls._model_policy.validate_auto_configuration() + @staticmethod async def _resolve_model( session: Any, @@ -809,105 +299,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": @@ -918,1658 +425,312 @@ 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: LoopContext, + 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 = outcome.unhandled_error + 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 {} + ), + }, ) - ctx.auto_failover = True - next_index = ( - 0 - if configured - else cls._cooldown_candidate_index( - ctx.session.id, - ctx.model_candidates, + + @staticmethod + def _authorize_auto_failover( + turn: LoopContext, + requested: bool, + ) -> None: + if requested and is_model_auto_session_category( + getattr(turn.session, "category", "user"), + ): + turn.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" ) - ) - cls._select_candidate(ctx, next_index) - return next_index + except Exception as exc: + log.debug("session.trace_offset.error", {"error": str(exc)}) + return 0 @classmethod - async def _prepare_auto_turn( + async def _acquire_lease( cls, - ctx: LoopContext, - last_user: MessageInfo, - ) -> bool: - """Synchronize routing when the loop advances to a real WebUI turn. + session_id: str, + turn: LoopContext, + ) -> _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", + ) + 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 - Returns: - True when ``last_user`` starts a new non-synthetic user turn. - """ - if last_user.id == ctx.turn_user_id: - return False + @staticmethod + async def _mark_busy( + session_id: str, + callbacks: LoopCallbacks, + ) -> None: + SessionStatus.set(session_id, SessionStatusBusy()) + await SessionEventSink.session_status(callbacks, session_id, "busy") - parts = await Message.parts(last_user.id, ctx.session.id) - if any(bool(getattr(part, "synthetic", False)) for part in parts): - return False + @staticmethod + async def _recover_orphan_tools(session_id: str) -> None: + try: + from flocks.session.orphan_tools import abort_orphan_running_parts - 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 + await abort_orphan_running_parts(session_id) + except Exception as exc: + log.warn( + "session.orphan_cleanup_failed", + {"session_id": session_id, "error": str(exc)}, + ) - primary = ctx.model_candidates[0] - config = await Config.get() - await cls._reset_auto_turn_candidates( - ctx, - primary, - last_user.id, - config=config, + @staticmethod + async def _handle_execution_error( + turn: LoopContext, + error: Exception, + ) -> AgentRunOutcome[Any]: + session_id = turn.session.id + log.error( + "session.execution_error", + {"session_id": session_id, "error": str(error)}, + ) + 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)}, ) - return True + try: + from flocks.bus.bus import Bus + from flocks.bus.events import SessionError - 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") + await Bus.publish( + SessionError, + {"sessionID": session_id, "error": str(error)}, ) - 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, + except Exception as publish_error: + log.warn( + "session.error_event_failed", + {"error": str(publish_error)}, + ) + return AgentRunOutcome( + status=AgentRunStatus.FATAL_FAILURE, + error=str(error), + unhandled_error=True, ) - 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_before_hook( + async def _release_session( cls, - ctx: LoopContext, - last_user: MessageInfo, + lease: _SessionLease, + callbacks: LoopCallbacks, ) -> None: - """Run UserPromptBefore 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_before({ - "sessionID": ctx.session.id, - "sessionCategory": ctx.session.category, - "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_before.error", { - "session_id": ctx.session.id, - "message_id": last_user.id, - "error": str(exc), - }) + release_task = cls._critical_release_task(lease, callbacks) + await cls._await_release_completion(release_task) - @classmethod - async def _run_turn_after_hook( - cls, - ctx: LoopContext, - callbacks: LoopCallbacks, - last_user: MessageInfo, - last_message: MessageInfo, + @staticmethod + async def _await_release_completion( + release_task: asyncio.Task[bool], ) -> bool: - """Run turn.after with terminal facts; never continue from hook output.""" - try: - from flocks.hooks.pipeline import HookPipeline + """Finish lease cleanup despite repeated cancellation.""" + interrupted: Optional[asyncio.CancelledError] = None + while True: + try: + released = await asyncio.shield(release_task) + break + except asyncio.CancelledError as exc: + if release_task.cancelled(): + raise + interrupted = exc - hook_user = last_user - if ctx.turn_user_id: - hook_user = ( - await Message.get(ctx.session.id, ctx.turn_user_id) - or last_user - ) - user_text = await Message.get_text_content(hook_user) - assistant_text = await Message.get_text_content(last_message) - await HookPipeline.run_turn_after({ - "sessionID": ctx.session.id, - "sessionCategory": ctx.session.category, - "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, - }, - "terminalOutcome": { - "status": "success", - "finish_reason": "stop", - }, - }) - except Exception as exc: - log.debug("loop.hook.turn_after.error", { - "session_id": ctx.session.id, - "message_id": getattr(last_message, "id", None), - "error": str(exc), - }) - return False - return False + if interrupted is not None: + raise interrupted + return released @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 + async def _await_pending_release(cls, session_id: str) -> None: + """Serialize a new busy event after the previous idle publication.""" + release_task = cls._release_tasks.get(session_id) + if release_task is None or release_task is asyncio.current_task(): return - await Message.update( - ctx.session.id, - failure.assistant_message_id, - error=failure.error_data, - finish="error", - ) + try: + await asyncio.shield(release_task) + except asyncio.CancelledError: + raise + except Exception as exc: + log.warn( + "session.previous_release_failed", + {"session_id": session_id, "error": str(exc)}, + ) @classmethod - async def _process_step_with_failover( + def _critical_release_task( cls, - ctx: LoopContext, + lease: _SessionLease, callbacks: LoopCallbacks, - messages: List[MessageInfo], - last_user: MessageInfo, - ) -> Any: - """Run one logical step, moving across candidates without replaying output.""" - from flocks.session.runner import RunnerCallbacks, SessionRunner + ) -> asyncio.Task[bool]: + existing = cls._release_tasks.get(lease.session_id) + if existing is not None and not existing.done(): + return existing + + async def release() -> bool: + async with Session.lifecycle_lock(lease.session_id): + if not cls._leases.owns(lease): + return False + cls._finalize_release_state_locked(lease) + await cls._publish_released(lease.turn, callbacks) + return True - 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, - ) - runner._step = ctx.trace_step - - step_result = await runner._process_step(messages, last_user) - if runner._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 + task = asyncio.create_task(release()) + cls._release_tasks[lease.session_id] = task - # 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, - }) - - previous = ctx.model_candidates[ctx.candidate_index] - next_candidate = ctx.model_candidates[next_index] - - 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 - - cls._select_candidate(ctx, next_index) - event_payload = { - "sessionID": ctx.session.id, - "from": { - "providerID": previous.provider_id, - "modelID": previous.model_id, - }, - "to": { - "providerID": next_candidate.provider_id, - "modelID": next_candidate.model_id, - }, - "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"], - }) - await cls._publish_runtime_event( - callbacks, - "session.model.fallback", - event_payload, - ) + def discard(completed: asyncio.Task[bool]) -> None: + if cls._release_tasks.get(lease.session_id) is completed: + cls._release_tasks.pop(lease.session_id, None) + + task.add_done_callback(discard) + return task @classmethod - async def _run_loop( + def _release_publication_task( cls, - ctx: LoopContext, + lease: _SessionLease, 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( - 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, - }) - - # 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 - - last_assistant_parts = ( - await Message.parts(last_assistant.id, ctx.session.id) - if last_assistant - else [] - ) - - # Check exit conditions (matching TUI lines 295-302) - 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 - ), - }) - last_message = last_assistant - break - - if await cls._prepare_auto_turn(ctx, last_user): - ctx.turn_additional_context = None - ctx.stop_hook_active = False - await cls._run_user_prompt_before_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) - reasoning_tokens = tokens_dict.get("reasoning", 0) - observed_prompt_tokens = input_tokens + cache_read - reported_total = observed_prompt_tokens + output_tokens + reasoning_tokens - - # Provider usage describes the prompt before the latest - # assistant response and its tool results. Always compare - # it with a lightweight estimate of the current messages - # so newly produced tool output cannot be missed. - if reported_total > 0: - ctx.last_observed_prompt_tokens = reported_total - # The assistant is marked ``tool-calls`` before its tools - # finish, so a concurrent UI estimate may have cached this - # message without the completed tool output. - SessionPrompt.invalidate_message_cache(last_finished.id) - last_finished_index = next( - ( - index - for index, message in enumerate(messages) - if message.id == last_finished.id - ), - len(messages) - 1, - ) + ) -> asyncio.Task[bool]: + """Track idle publication after settlement released the lease.""" - async def _estimate_effective_tokens() -> tuple[int, int, str]: - if observed_prompt_tokens > 0: - tool_result_tokens = ( - await SessionPrompt.estimate_tool_result_tokens( - ctx.session.id, - last_finished.id, - ) - ) - later_tokens = ( - await SessionPrompt.estimate_full_context_tokens( - ctx.session.id, - messages[last_finished_index + 1:], - policy=compaction_policy, - ) - ) - delta_tokens = tool_result_tokens + later_tokens - return ( - reported_total + delta_tokens, - delta_tokens, - "observed+estimated_delta", - ) - - estimated_tokens = ( - await SessionPrompt.estimate_full_context_tokens( - ctx.session.id, - messages, - policy=compaction_policy, - ) - ) - return max(reported_total, estimated_tokens), estimated_tokens, "estimated" - - ( - effective_tokens, - estimated_component_tokens, - decision_source, - ) = await _estimate_effective_tokens() - tokens_dict = { - "input": effective_tokens, - "output": 0, - "cache": {"read": 0, "write": 0}, - } - log.info("loop.tokens_decision", { - "session_id": ctx.session.id, - "source": decision_source, - "effective_tokens": effective_tokens, - "observed_tokens": reported_total, - "estimated_component_tokens": estimated_component_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: - message_tokens_before_cleanup = ( - await SessionPrompt.estimate_full_context_tokens( - ctx.session.id, - messages, - policy=compaction_policy, - ) - ) - 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, - }) - message_tokens_after_cleanup = ( - await SessionPrompt.estimate_full_context_tokens( - ctx.session.id, - messages, - policy=compaction_policy, - ) - ) - baseline_offset_tokens = max( - 0, - effective_tokens - message_tokens_before_cleanup, - ) - effective_tokens = ( - message_tokens_after_cleanup + baseline_offset_tokens - ) - tokens_dict["input"] = effective_tokens - current_input_tokens = effective_tokens - log.info("loop.pre_compact_cleanup_rechecked", { - "session_id": ctx.session.id, - "effective_tokens": effective_tokens, - "message_tokens": message_tokens_after_cleanup, - "baseline_offset_tokens": baseline_offset_tokens, - "overflow_threshold": ( - compaction_policy.overflow_threshold - ), - }) - if effective_tokens <= compaction_policy.overflow_threshold: - 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: - message_tokens_before_cleanup = ( - await SessionPrompt.estimate_full_context_tokens( - ctx.session.id, - messages, - policy=compaction_policy, - ) - ) - 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 - message_tokens_after_cleanup = ( - await SessionPrompt.estimate_full_context_tokens( - ctx.session.id, - messages, - policy=compaction_policy, - ) - ) - baseline_offset_tokens = max( - 0, - effective_tokens - message_tokens_before_cleanup, - ) - re_est = ( - message_tokens_after_cleanup - + baseline_offset_tokens - ) - 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, - ) - - 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 completed assistant belongs - # to the preceding user turn, so _should_exit() - # won't trigger for the continuation message. - 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, - ) - ) - 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) - last_message = next( - ( - msg - for msg in reversed(post_messages) - if msg.role == MessageRole.ASSISTANT - and getattr(msg, "parentID", None) == last_user.id - ), - 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, - }) - continue + async def publish() -> bool: + await cls._publish_released(lease.turn, callbacks) + return True - 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 - ) - 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, - }, - ) - 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, - }) - continue + task = asyncio.create_task(publish()) + cls._release_tasks[lease.session_id] = task - 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_after_hook( - ctx, - callbacks, - last_user, - last_message, - ) - ): - continue + def discard(completed: asyncio.Task[bool]) -> None: + if cls._release_tasks.get(lease.session_id) is completed: + cls._release_tasks.pop(lease.session_id, None) - stop_reason = step_result.error or (getattr(last_message, "finish", None) if last_message else None) or "stop" - turn_state = set_turn_state( - ctx.session.id, - step=ctx.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)) + task.add_done_callback(discard) + return task - 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, - ) - queued_user = await cls._detect_queued_user_message( - ctx.session.id, - post_messages, - last_user.id, - last_assistant_after_step, - ) - 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, - ) - 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", { - "session_id": ctx.session.id, - "action": step_result.action, - }) - break - - # Return result - return LoopResult( - action="error" if ctx.auto_failover and loop_error else "stop", - last_message=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 ctx.should_abort() 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( + async def _settle_or_continue( cls, - last_user: MessageInfo, - last_assistant: Optional[MessageInfo], - last_assistant_parts: Optional[List[Any]] = None, + lease: _SessionLease, + callbacks: LoopCallbacks, + processed_user_id: Optional[str], + *, + allow_late_input: bool = True, ) -> bool: - """ - Check if loop should exit - - Ported from original exit logic: - - Exit if assistant has responded with finish != tool-calls - - Exit if assistant is a response to the latest user message - """ - if not last_assistant: - return False + """Atomically keep ownership for late input or settle idle.""" + async with Session.lifecycle_lock(lease.session_id): + if allow_late_input and await lease.turn.has_late_input( + processed_user_id, + ): + log.info( + "session.continuing_for_late_input", + { + "session_id": lease.session_id, + "processed_user_id": processed_user_id, + }, + ) + return True + cls._finalize_release_state_locked(lease) + release_task = cls._release_publication_task(lease, callbacks) - 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 getattr(last_assistant, "parentID", None) == last_user.id: - # Assistant responded to this user turn - return True - + await cls._await_release_completion(release_task) return False - + @classmethod - async def _check_reminders( - cls, - ctx: LoopContext, - messages: List[MessageInfo], + 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 + async def _publish_released( + turn: LoopContext, 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 - + session_id = turn.session.id + await SessionEventSink.session_status(callbacks, session_id, "idle") 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, + await Session.touch(turn.session.project_id, session_id) + except Exception as exc: + log.warn( + "session.touch_failed", + {"session_id": session_id, "error": str(exc)}, ) - - # 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, - }) - + + 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("session.idle_event_failed", {"error": str(exc)}) -# Export __all__ = [ "SessionLoop", "LoopContext", 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..e1b23c4da 100644 --- a/flocks/task/background.py +++ b/flocks/task/background.py @@ -215,6 +215,7 @@ async def _inject_parent_completion(self, task: BackgroundTask) -> None: """Inject completed background task output into the parent context.""" if task.completion_injected or not task.parent_session_id: return + parent_session_id = task.parent_session_id state = task.status if state == "completed": @@ -237,21 +238,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=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( + 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 +272,7 @@ 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 + parent_session_id = task.parent_session_id async def _run_parent() -> None: try: @@ -281,7 +281,7 @@ async def _run_parent() -> None: model = task.parent_model or {} result = await SessionLoop.run( - session_id=task.parent_session_id, + session_id=parent_session_id, provider_id=model.get("providerID"), model_id=model.get("modelID"), agent_name=task.parent_agent, @@ -423,7 +423,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 +434,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/test_unified_session_loop.py b/tests/agent/test_unified_session_loop.py index 44aaa8fa0..16e514b6d 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. 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) -4. _resolve_model implements 5-level priority correctly +Verifies that _resolve_model implements its model-selection priority. """ import asyncio @@ -14,77 +10,27 @@ 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.""" - - def test_event_publish_callback_field_exists(self): - cb = RunnerCallbacks() - 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) - 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 - - -class TestCallbackPrecedence: - """SessionRunner should not override explicit callbacks with CLI fallback.""" +class TestResolveModel: + """Test the _resolve_model 5-level priority.""" - def test_explicit_callbacks_not_overridden(self): - """When event_publish_callback is set, CLI fallback should NOT be used.""" - 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 + @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() -class TestResolveModel: - """Test the _resolve_model 5-level priority.""" + monkeypatch.setattr( + Session, + "run_active_write", + classmethod(run_active_write), + ) @pytest.mark.asyncio async def test_priority_1_request_model(self): @@ -295,6 +241,9 @@ async def test_webui_auto_uses_default_primary_and_returns_actual_model( from types import SimpleNamespace from flocks.server.routes import session as session_routes + from flocks.session.runtime.model_policy import ( + DEFAULT_MODEL_ROUTING_POLICY, + ) from flocks.session.session_loop import LoopResult, SessionLoop request = session_routes.PromptRequest( @@ -350,7 +299,7 @@ async def test_webui_auto_uses_default_primary_and_returns_actual_model( AsyncMock(return_value=SimpleNamespace()), ) monkeypatch.setattr( - SessionLoop, + DEFAULT_MODEL_ROUTING_POLICY, "validate_runtime_model", AsyncMock(return_value=(True, "available")), ) @@ -407,7 +356,9 @@ async def test_unsupported_session_ignores_auto_flag(self, monkeypatch): from types import SimpleNamespace from flocks.server.routes import session as session_routes - from flocks.session.session_loop import SessionLoop + from flocks.session.runtime.model_policy import ( + DEFAULT_MODEL_ROUTING_POLICY, + ) request = session_routes.PromptRequest( parts=[{"type": "text", "text": "task input"}], @@ -452,7 +403,11 @@ async def test_unsupported_session_ignores_auto_flag(self, monkeypatch): "flocks.config.config.Config.get", AsyncMock(return_value=SimpleNamespace()), ) - monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) + monkeypatch.setattr( + DEFAULT_MODEL_ROUTING_POLICY, + "validate_runtime_model", + validate, + ) monkeypatch.setattr( "flocks.provider.provider.Provider._ensure_initialized", lambda: None, diff --git a/tests/channel/test_channel.py b/tests/channel/test_channel.py index 2f88a1659..37b5c1456 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 # ===================================================================== @@ -1046,7 +1070,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 @@ -1094,9 +1122,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 @@ -1143,24 +1176,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( @@ -1224,18 +1258,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( @@ -1263,18 +1297,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( @@ -2095,7 +2129,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 @@ -2136,9 +2174,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") @@ -2176,9 +2219,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") @@ -2216,3 +2264,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/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/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..8786c6980 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.""" @@ -2280,7 +2244,9 @@ async def test_prepare_auto_replay_defers_unavailable_primary_to_failover( category: str, ): from flocks.server.routes import session as session_routes - from flocks.session.session_loop import SessionLoop + from flocks.session.runtime.model_policy import ( + DEFAULT_MODEL_ROUTING_POLICY, + ) user_message = SimpleNamespace(agent="rex") monkeypatch.setattr( @@ -2306,7 +2272,11 @@ async def test_prepare_auto_replay_defers_unavailable_primary_to_failover( AsyncMock(return_value=SimpleNamespace()), ) validate = AsyncMock(return_value=(False, "provider_not_configured")) - monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) + monkeypatch.setattr( + DEFAULT_MODEL_ROUTING_POLICY, + "validate_runtime_model", + validate, + ) monkeypatch.setattr("flocks.provider.provider.Provider._ensure_initialized", lambda: None) monkeypatch.setattr("flocks.provider.provider.Provider.apply_config", AsyncMock()) monkeypatch.setattr("flocks.provider.provider.Provider.get", lambda _provider_id: None) @@ -2330,7 +2300,9 @@ async def test_prepare_replay_ignores_auto_on_unsupported_session( monkeypatch: pytest.MonkeyPatch, ): from flocks.server.routes import session as session_routes - from flocks.session.session_loop import SessionLoop + from flocks.session.runtime.model_policy import ( + DEFAULT_MODEL_ROUTING_POLICY, + ) user_message = SimpleNamespace(agent="rex") monkeypatch.setattr( @@ -2356,7 +2328,11 @@ async def test_prepare_replay_ignores_auto_on_unsupported_session( AsyncMock(return_value=SimpleNamespace()), ) validate = AsyncMock() - monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) + monkeypatch.setattr( + DEFAULT_MODEL_ROUTING_POLICY, + "validate_runtime_model", + validate, + ) monkeypatch.setattr( "flocks.provider.provider.Provider._ensure_initialized", lambda: None, diff --git a/tests/session/runtime/test_agent_loop.py b/tests/session/runtime/test_agent_loop.py new file mode 100644 index 000000000..5ef134584 --- /dev/null +++ b/tests/session/runtime/test_agent_loop.py @@ -0,0 +1,202 @@ +"""Tests for the logical-input AgentLoop.""" + +from __future__ import annotations + +import asyncio +from collections import deque +from dataclasses import dataclass + +import pytest + +from flocks.session.runtime.agent_loop import AgentLoop +from flocks.session.runtime.contracts import ( + AgentRunStatus, + AttemptEffects, + ModelTurnBoundary, + ModelTurnPreparation, + ModelTurnSnapshot, + RuntimeModel, + StepFailure, + StepResult, + TurnPreparationStatus, +) + + +@dataclass(frozen=True) +class Message: + id: str + content: str + + +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, + preparations: list[ModelTurnPreparation[Message]], + boundaries: list[ModelTurnBoundary[Message]], + ) -> None: + self._preparations = deque(preparations) + self._boundaries = deque(boundaries) + self.aborted = False + self.step = 0 + self._current_step_task = None + self.session = type("Session", (), {"id": "ses_agent_loop"})() + + async def prepare_step(self) -> ModelTurnPreparation[Message]: + return self._preparations.popleft() + + async def commit_step( + self, + _step_result: StepResult, + ) -> ModelTurnBoundary[Message]: + return self._boundaries.popleft() + + +def _ready( + messages: tuple[Message, ...], + *, + turn: int = 0, +) -> ModelTurnPreparation[Message]: + return ModelTurnPreparation( + status=TurnPreparationStatus.READY, + snapshot=ModelTurnSnapshot( + active_model=RuntimeModel("provider-a", "model-a"), + trace_step=turn, + messages=messages, + last_user=messages[-1], + ), + ) + + +async def _run( + engine: FakeStepEngine, + preparations, + boundaries, +): + turn = FakeTurn(preparations, boundaries) + return await AgentLoop().run(turn, engine) + + +@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( + engine, + [_ready((user,)), _ready((user, tool_result), turn=1)], + [ + 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_precedes_final_step_failure() -> None: + user = Message("user-1", "hello") + failed = Message("assistant-1", "provider failed") + 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( + FakeStepEngine( + [StepResult(action="stop", error=failure.message, failure=failure)], + ), + [_ready((user,))], + [ + ModelTurnBoundary( + last_message=failed, + input_available=True, + ), + ], + ) + assert outcome.status == AgentRunStatus.INPUT_AVAILABLE + assert outcome.step_result.failure is failure + + +@pytest.mark.asyncio +async def test_cancelled_step_is_finalized_before_abort_returns() -> None: + user = Message("user-1", "hello") + + class CancellableEngine: + def __init__(self) -> None: + self.started = asyncio.Event() + self.finalized = False + + async def run(self, _snapshot) -> StepResult: + self.started.set() + await asyncio.Event().wait() + raise AssertionError("unreachable") + + async def finalize_cancelled_attempt(self) -> None: + self.finalized = True + + engine = CancellableEngine() + turn = FakeTurn([_ready((user,))], []) + run = asyncio.create_task(AgentLoop().run(turn, engine)) + await asyncio.wait_for(engine.started.wait(), timeout=1) + turn.aborted = True + turn._current_step_task.cancel() + + outcome = await asyncio.wait_for(run, timeout=1) + + assert outcome.status == AgentRunStatus.ABORTED + assert engine.finalized is True + + +@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( + FakeStepEngine( + [StepResult(action="stop", error=failure.message, failure=failure)], + ), + [_ready((user,))], + [ModelTurnBoundary()], + ) + assert outcome.status == expected_status diff --git a/tests/session/runtime/test_contracts.py b/tests/session/runtime/test_contracts.py new file mode 100644 index 000000000..ac86e307e --- /dev/null +++ b/tests/session/runtime/test_contracts.py @@ -0,0 +1,54 @@ +"""Tests for replay-safe runtime request contracts.""" + +from flocks.provider.provider import ChatMessage +from flocks.session.runtime.contracts import ( + ModelRequest, +) + + +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" + message["content"].append("mutated source") + first_messages = request.provider_messages() + first_messages[0]["content"].append("mutated provider view") + first_tools = request.provider_tools() + first_tools[0]["function"]["name"] = "mutated" + + assert request.provider_messages()[0]["content"] == ["hello"] + assert request.provider_tools()[0]["function"]["name"] == "read" + assert request.provider_options()["reasoning"]["effort"] == "high" + + +def test_model_request_reuses_owned_chat_messages_for_provider_calls() -> None: + """Provider calls get a fresh list without copying the full history.""" + message = ChatMessage( + role="user", + content=[{"type": "text", "text": "large history entry"}], + ) + request = ModelRequest( + provider_id="provider", + model_id="model", + messages=(message,), + tools=(), + options={}, + ) + + first_messages = request.provider_messages() + second_messages = request.provider_messages() + + assert first_messages is not second_messages + assert request.messages[0] is message + assert first_messages[0] is message + assert second_messages[0] is message diff --git a/tests/session/runtime/test_session_loop.py b/tests/session/runtime/test_session_loop.py new file mode 100644 index 000000000..4ee38336d --- /dev/null +++ b/tests/session/runtime/test_session_loop.py @@ -0,0 +1,361 @@ +"""SessionLoop lifecycle and logical-turn ownership tests.""" + +from __future__ import annotations + +import asyncio +from contextlib import asynccontextmanager +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, + AgentRunStatus, + ContinuationDecision, + 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( + user, + label: str, +) -> AgentRunOutcome: + return AgentRunOutcome( + status=AgentRunStatus.COMPLETED, + last_user=user, + last_message=SimpleNamespace(label=label), + step_result=StepResult(action="stop"), + ) + + +@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( + first_user if run.await_count == 1 else second_user, + "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_continuation_error_is_not_masked_by_late_input( + 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( + side_effect=[ + RuntimeError("continuation store unavailable"), + ContinuationDecision(), + ], + ), + ) + monkeypatch.setattr(SessionLoop, "_continuation_policy", continuation) + run = AsyncMock( + side_effect=[ + _outcome(first_user, "first"), + _outcome(second_user, "second"), + ], + ) + 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.action == "error" + assert result.error == "continuation store unavailable" + assert run.await_count == 1 + assert active == {} + + +@pytest.mark.asyncio +async def test_repeated_cancellation_cannot_leak_session_lease( + monkeypatch, + loop_io, +) -> None: + session, active = loop_io + run_started = asyncio.Event() + release_started = asyncio.Event() + allow_release = asyncio.Event() + lock_count = 0 + + @asynccontextmanager + async def lifecycle_lock(_session_id: str): + nonlocal lock_count + lock_count += 1 + if lock_count > 1: + release_started.set() + await allow_release.wait() + yield + + async def run_until_cancelled(*_args): + run_started.set() + await asyncio.Event().wait() + + continuation = SimpleNamespace( + prepare_logical_turn=AsyncMock(), + resolve=AsyncMock(), + ) + monkeypatch.setattr(SessionLoop, "_continuation_policy", continuation) + monkeypatch.setattr(Session, "lifecycle_lock", lifecycle_lock) + monkeypatch.setattr( + "flocks.session.session_loop.Message.list", + AsyncMock(return_value=[]), + ) + monkeypatch.setattr(AgentLoop, "run", run_until_cancelled) + + run_task = asyncio.create_task( + SessionLoop.run( + session.id, + provider_id="provider", + model_id="model", + ) + ) + await asyncio.wait_for(run_started.wait(), timeout=1) + run_task.cancel() + await asyncio.wait_for(release_started.wait(), timeout=1) + run_task.cancel() + allow_release.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(run_task, timeout=1) + + assert active == {} + assert SessionStatus.get(session.id).type == "idle" + Session.touch.assert_awaited_once_with(session.project_id, session.id) + from flocks.bus.bus import Bus + + Bus.publish.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_next_run_waits_for_previous_idle_publication( + monkeypatch, + loop_io, +) -> None: + session, active = loop_io + first_user = _message("msg_001") + idle_started = asyncio.Event() + allow_idle = asyncio.Event() + events: list[str] = [] + + async def publish(event_name, payload): + status = payload.get("status", {}).get("type") + if event_name == "session.status" and status: + events.append(status) + if status == "idle" and events.count("idle") == 1: + idle_started.set() + await allow_idle.wait() + + async def prepare(turn): + turn.prepared_user_id = first_user.id + + continuation = SimpleNamespace( + prepare_logical_turn=AsyncMock(side_effect=prepare), + resolve=AsyncMock(return_value=ContinuationDecision()), + ) + monkeypatch.setattr(SessionLoop, "_continuation_policy", continuation) + monkeypatch.setattr( + "flocks.session.session_loop.Message.list", + AsyncMock(return_value=[first_user]), + ) + monkeypatch.setattr( + AgentLoop, + "run", + AsyncMock(return_value=_outcome(first_user, "done")), + ) + callbacks = SimpleNamespace( + event_publish_callback=publish, + on_error=None, + ) + + first_run = asyncio.create_task( + SessionLoop.run( + session.id, + provider_id="provider", + model_id="model", + callbacks=callbacks, + ) + ) + await asyncio.wait_for(idle_started.wait(), timeout=1) + second_run = asyncio.create_task( + SessionLoop.run( + session.id, + provider_id="provider", + model_id="model", + callbacks=callbacks, + ) + ) + await asyncio.sleep(0) + + assert events == ["busy", "idle"] + assert not second_run.done() + + allow_idle.set() + await asyncio.wait_for(first_run, timeout=1) + await asyncio.wait_for(second_run, timeout=1) + + assert events == ["busy", "idle", "busy", "idle"] + assert active == {} diff --git a/tests/session/test_auto_model_failover.py b/tests/session/test_auto_model_failover.py index b1d655b1f..174e4c109 100644 --- a/tests/session/test_auto_model_failover.py +++ b/tests/session/test_auto_model_failover.py @@ -1,27 +1,41 @@ """Focused tests for WebUI Auto runtime model failover.""" +import asyncio import time from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock 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 ( + AgentRunStatus, + AttemptEffects, + ModelTurnBoundary, + ModelTurnPreparation, + ModelTurnSnapshot, + RuntimeModel, + TurnPreparationStatus, +) from flocks.session.message import Message, MessageRole -from flocks.session.runner import ( - LlmAttemptState, - SessionRunner, +from flocks.session.runtime.model_policy import ( + DEFAULT_MODEL_ROUTING_POLICY, + AutoFailoverCooldown, +) +from flocks.session.runtime.step_engine import ( + 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,13 +80,29 @@ 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=DEFAULT_MODEL_ROUTING_POLICY.validate_runtime_model, + ) + + def _failure( *, assistant_id: str, 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", @@ -89,11 +119,28 @@ 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( + active_model=RuntimeModel(turn.provider_id, turn.model_id), + 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 +163,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 +193,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 +209,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 +237,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 +259,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 +275,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 +303,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 +315,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 +340,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 +349,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 +365,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 +411,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 +448,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 +477,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 +549,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 +579,7 @@ async def stream(): return stream() provider = FailingStreamProvider() - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="primary", model_id="primary-model", @@ -529,19 +590,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 +618,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 +638,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 +653,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], @@ -613,6 +674,63 @@ async def publish(event, payload): assert any(event == "session.model.fallback" for event, _ in events) +@pytest.mark.asyncio +async def test_abort_during_failed_attempt_cleanup_stops_failover( + monkeypatch, +) -> None: + ctx = _ctx() + last_user = SimpleNamespace(id="msg_user", agent="rex") + snapshot = ModelTurnSnapshot( + active_model=RuntimeModel(ctx.provider_id, ctx.model_id), + trace_step=ctx.trace_step, + messages=(last_user,), + last_user=last_user, + ) + ctx.prepare_step = AsyncMock( + return_value=ModelTurnPreparation( + status=TurnPreparationStatus.READY, + snapshot=snapshot, + ) + ) + ctx.commit_step = AsyncMock(return_value=ModelTurnBoundary()) + attempts: list[str] = [] + events: list[str] = [] + delete_started = asyncio.Event() + allow_delete = asyncio.Event() + + async def process_step(runner, _messages, _last_user): + attempts.append(runner.provider_id) + if runner.provider_id == "primary": + return _failure(assistant_id="msg_failed") + return StepResult(action="stop", content="unexpected fallback") + + async def delete_failed_attempt(*_args): + delete_started.set() + await allow_delete.wait() + return True + + async def publish(event, _payload): + events.append(event) + + monkeypatch.setattr(StepEngine, "_process_step", process_step) + monkeypatch.setattr(Message, "delete", delete_failed_attempt) + ctx.callbacks = LoopCallbacks(event_publish_callback=publish) + run_task = asyncio.create_task( + AgentLoop().run(ctx, StepEngine.from_turn(ctx)), + ) + + await asyncio.wait_for(delete_started.wait(), timeout=1) + ctx.signal_abort() + allow_delete.set() + outcome = await asyncio.wait_for(run_task, timeout=1) + + assert outcome.status == AgentRunStatus.ABORTED + assert attempts == ["primary"] + assert ctx.candidate_index == 0 + assert "session.model.fallback" not in events + ctx.commit_step.assert_not_awaited() + + @pytest.mark.asyncio async def test_queued_user_is_detected_before_replacement_assistant(): current_user = SimpleNamespace(id="msg_001", role=MessageRole.USER) @@ -622,7 +740,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, @@ -655,17 +773,17 @@ 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, ), ) 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 +814,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 +822,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 +856,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 +889,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 +906,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 +929,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 +943,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 +979,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 +988,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 +1011,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 +1044,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 +1055,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 +1077,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 +1105,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. @@ -1029,14 +1154,18 @@ async def validate(provider_id, _model_id, **_kwargs): available = provider_id != "missing" return available, "available" if available else "provider_not_configured" - monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) + monkeypatch.setattr( + DEFAULT_MODEL_ROUTING_POLICY, + "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 +1179,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 @@ -1085,14 +1216,14 @@ async def test_candidate_builder_keeps_active_cooldown_model_in_its_tier( lambda: model_manager, ) monkeypatch.setattr( - SessionLoop, + DEFAULT_MODEL_ROUTING_POLICY, "validate_runtime_model", AsyncMock(return_value=(True, "available")), ) 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 +1238,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, + DEFAULT_MODEL_ROUTING_POLICY, "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 +1274,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 +1284,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", @@ -1174,10 +1305,14 @@ async def validate(provider_id, _model_id, **_kwargs): available = provider_id != "missing" return available, "available" if available else "provider_not_configured" - monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) + monkeypatch.setattr( + DEFAULT_MODEL_ROUTING_POLICY, + "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,21 +1331,23 @@ 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(), ) monkeypatch.setattr( - SessionLoop, + DEFAULT_MODEL_ROUTING_POLICY, "validate_runtime_model", AsyncMock(return_value=(False, "provider_not_configured")), ) 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,24 +1359,24 @@ 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 -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( @@ -1248,7 +1385,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 +1412,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 +1436,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 +1456,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 +1489,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 +1511,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 +1538,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 +1551,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 +1579,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 +1595,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 +1619,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 +1636,29 @@ 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, + AgentRunStatus, + ) + nonlocal captured_ctx captured_ctx = ctx - return LoopResult(action="stop") + return AgentRunOutcome( + status=AgentRunStatus.ABORTED, + ) 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 +1680,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 +1700,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 +1717,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..9aa0b2ed2 100644 --- a/tests/session/test_execution_mode.py +++ b/tests/session/test_execution_mode.py @@ -399,13 +399,14 @@ 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) agent = SimpleNamespace( + name="rex", tools=[ "read", "bash", @@ -434,7 +435,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 b2db3efe9..68e9fd627 100644 --- a/tests/session/test_lifecycle_hooks.py +++ b/tests/session/test_lifecycle_hooks.py @@ -1,16 +1,24 @@ -"""Focused tests for session lifecycle hook integration.""" +"""Focused tests for the Python lifecycle-hook seams.""" from __future__ import annotations +import asyncio from types import SimpleNamespace -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from flocks.hooks.pipeline import HookContext, HookStage -from flocks.session.runner import SessionRunner +from flocks.session.runtime.continuation_policy import DEFAULT_CONTINUATION_POLICY +from flocks.session.goal import GoalDecision +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: @@ -25,12 +33,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 @@ -41,7 +50,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=[ [], @@ -50,26 +59,61 @@ async def test_real_user_turn_is_detected_once_and_synthetic_is_ignored() -> Non ), ), patch.object( - SessionLoop, - "_run_user_prompt_before_hook", + DEFAULT_CONTINUATION_POLICY, + "run_user_prompt_submit", AsyncMock(), - ) as prompt_hook, + ) 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_before_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 - prompt_hook.assert_awaited_once_with(ctx, first_user) + submit_hook.assert_awaited_once_with(ctx, first_user) @pytest.mark.asyncio -async def test_user_prompt_before_adds_ephemeral_turn_context() -> None: +async def test_settled_history_skips_user_prompt_submit() -> None: + ctx = _loop_context("ses_settled_history_hook") + user = _message("msg_user", "user") + assistant = _message("msg_assistant", "assistant", finish="stop") + assistant.parentID = user.id + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock(return_value=[user, assistant]), + ) + prepare_turn = AsyncMock(return_value=True) + submit_hook = AsyncMock() + + with ( + patch.object( + DEFAULT_CONTINUATION_POLICY._model_policy, + "prepare_turn", + prepare_turn, + ), + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_user_prompt_submit", + submit_hook, + ), + patch( + "flocks.session.runtime.continuation_policy.Message.parts", + AsyncMock(return_value=[]), + ), + ): + await DEFAULT_CONTINUATION_POLICY.prepare_logical_turn(ctx) + + assert ctx.prepared_user_id == user.id + prepare_turn.assert_not_awaited() + submit_hook.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_user_prompt_submit_adds_ephemeral_turn_context() -> None: ctx = _loop_context() user = SimpleNamespace(id="msg_user", agent="rex") run_hook = AsyncMock( return_value=HookContext( - stage=HookStage.USER_PROMPT_BEFORE, + stage=HookStage.USER_PROMPT_SUBMIT, input={}, output={"additionalContext": " current sprint context "}, ) @@ -77,25 +121,30 @@ async def test_user_prompt_before_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( - "flocks.hooks.pipeline.HookPipeline.run_user_prompt_before", + "flocks.hooks.pipeline.HookPipeline.run_user_prompt_submit", run_hook, ), ): - await SessionLoop._run_user_prompt_before_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] assert payload["messageID"] == user.id assert payload["prompt"] == "implement hooks" + assert payload["sessionCategory"] == "user" + assert payload["model"] == { + "providerID": "test-provider", + "modelID": "test-model", + } @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", @@ -104,7 +153,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")) @@ -112,39 +161,687 @@ async def test_session_start_runs_only_when_pending() -> None: run_hook.assert_awaited_once() assert runner._session_start_fired is True + assert run_hook.await_args.args[0]["sessionID"] == "ses_session_start" @pytest.mark.asyncio -async def test_turn_after_observes_terminal_outcome_without_continuation() -> None: - ctx = _loop_context("ses_turn_after") +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") - assistant = SimpleNamespace(id="msg_assistant", agent="rex", finish="stop") + 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()) - run_hook = AsyncMock(return_value=HookContext(stage=HookStage.TURN_AFTER, input={}, output={})) + 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_AFTER, + input={}, + output={ + "decision": "block", + "reason": "Run the test suite before finishing.", + }, + ) + ) with ( patch( - "flocks.session.session_loop.Message.get", + "flocks.session.runtime.continuation_policy.Message.get", AsyncMock(return_value=user), ), patch( - "flocks.session.session_loop.Message.get_text_content", - AsyncMock(side_effect=["prompt", "response"]), + "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.continuation_policy.Message.create", + create_message, ), patch( "flocks.hooks.pipeline.HookPipeline.run_turn_after", run_hook, ), + patch( + "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.resolve( + ctx, + 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["sessionCategory"] == "user" + 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.stopped" + + +@pytest.mark.asyncio +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( + id="msg_002", + agent="rex", + role="assistant", + finish="stop", + ) + queued_user = SimpleNamespace(id="msg_003", agent="rex", role="user") + messages = [user, assistant] + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock(side_effect=lambda: messages), + ) + callbacks = LoopCallbacks(event_publish_callback=AsyncMock()) + ctx.callbacks = callbacks + run_turn_after = AsyncMock(side_effect=lambda *_args: messages.append(queued_user)) + + with ( + patch( + "flocks.session.runtime.continuation_policy.Message.get_text_content", + AsyncMock(return_value="response"), + ), + patch( + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", + AsyncMock( + return_value=GoalDecision( + status=None, + verdict="inactive", + ) + ), + ), + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_after", + run_turn_after, + create=True, + ), ): - continued = await SessionLoop._run_turn_after_hook( + decision = await DEFAULT_CONTINUATION_POLICY.resolve( ctx, - callbacks, - user, - assistant, + SimpleNamespace(last_user=user, last_message=assistant), ) - assert continued is False - payload = run_hook.await_args.args[0] - assert payload["terminalOutcome"]["status"] == "success" - callbacks.event_publish_callback.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" + 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( + 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_goal_continuation_persistence_failure_is_terminal() -> None: + ctx = _loop_context("ses_goal_persistence_failure") + user = SimpleNamespace( + id="msg_user", + agent="rex", + role="user", + model={"providerID": "test-provider", "modelID": "test-model"}, + provider="test-provider", + ) + assistant = SimpleNamespace( + id="msg_assistant", + agent="rex", + role="assistant", + finish="stop", + parentID=user.id, + ) + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock(return_value=[user, assistant]), + ) + ctx.callbacks = LoopCallbacks(event_publish_callback=AsyncMock()) + + with ( + patch( + "flocks.session.runtime.continuation_policy.Message.get_text_content", + AsyncMock(return_value="work remains"), + ), + patch( + "flocks.session.runtime.continuation_policy.Message.create", + AsyncMock(side_effect=OSError("continuation store unavailable")), + ), + patch( + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", + AsyncMock( + return_value=GoalDecision( + status="active", + verdict="continue", + should_continue=True, + continuation_prompt="continue the goal", + ) + ), + ), + patch( + "flocks.agent.registry.Agent.get", + AsyncMock(return_value=SimpleNamespace(steps=10)), + ), + ): + with pytest.raises( + RuntimeError, + match="continuation store unavailable", + ): + await DEFAULT_CONTINUATION_POLICY.resolve( + ctx, + SimpleNamespace(last_user=user, last_message=assistant), + ) + + event_names = [ + call.args[0] + for call in ctx.callbacks.event_publish_callback.await_args_list + ] + assert "turn.stopped" not in event_names + + +def _message( + message_id: str, + role: str, + *, + finish: str | None = None, + parent_id: str | None = None, +) -> SimpleNamespace: + return SimpleNamespace( + id=message_id, + role=role, + finish=finish, + parentID=parent_id, + tokens=None, + summary=False, + agent="rex", + model={"providerID": "test-provider", "modelID": "test-model"}, + ) + + +@pytest.mark.asyncio +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", + parent_id=user.id, + ) + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user], + [user, assistant], + [user, assistant], + [user, assistant], + [user, assistant], + ] + ) + ) + run_turn_after = AsyncMock() + + with ( + patch( + "flocks.session.runtime.session_turn.Message.parts", + AsyncMock(return_value=[]), + ), + patch( + "flocks.session.runtime.session_turn.Message.get_text_content", + AsyncMock(return_value="final response"), + ), + patch( + "flocks.session.runtime.session_turn.Provider.resolve_model_info", + return_value=(0, 0, None), + ), + patch( + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", + AsyncMock( + return_value=GoalDecision( + status="inactive", + verdict="inactive", + ) + ), + ), + patch( + "flocks.session.runtime.continuation_policy.ContinuationPolicy.run_user_prompt_submit", + AsyncMock(), + ), + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_after", + run_turn_after, + ), + patch( + "flocks.session.runtime.step_engine.StepEngine._process_step", + AsyncMock(return_value=StepResult(action="stop")), + ), + patch( + "flocks.session.lifecycle.title.SessionTitle.ensure_title", + MagicMock(return_value=None), + ), + patch( + "flocks.session.runtime.session_turn.fire_and_forget", + MagicMock(), + ), + ): + result = await run_logical_turns(ctx, LoopCallbacks()) + + assert result.action == "stop" + run_turn_after.assert_awaited_once_with( + ctx, + user, + assistant, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("step_result", "assistant_finish"), + [ + (StepResult(action="stop", error="provider failed"), "error"), + (StepResult(action="continue"), "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_after_{assistant_finish}") + user = _message("msg_001", "user") + assistant = _message("msg_002", "assistant", finish=assistant_finish) + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user], + [user, assistant], + ] + ) + ) + run_turn_after = AsyncMock() + + async def process_step(*_args, **_kwargs): + if step_result.action == "continue": + ctx.signal_abort() + return step_result + + with ( + patch( + "flocks.session.runtime.session_turn.Message.parts", + AsyncMock(return_value=[]), + ), + patch( + "flocks.session.runtime.session_turn.Provider.resolve_model_info", + return_value=(0, 0, None), + ), + patch( + "flocks.session.runtime.continuation_policy.ContinuationPolicy.run_user_prompt_submit", + AsyncMock(), + ), + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_after", + run_turn_after, + ), + patch( + "flocks.session.runtime.step_engine.StepEngine._process_step", + AsyncMock(side_effect=process_step), + ), + patch( + "flocks.session.lifecycle.title.SessionTitle.ensure_title", + MagicMock(return_value=None), + ), + patch( + "flocks.session.runtime.session_turn.fire_and_forget", + MagicMock(), + ), + ): + result = await run_logical_turns(ctx, LoopCallbacks()) + + assert result.action == "stop" + run_turn_after.assert_not_awaited() + + +@pytest.mark.asyncio +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") + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user], + [user, assistant, queued_user], + ] + ) + ) + run_turn_after = AsyncMock() + + async def process_step(*_args, **_kwargs): + ctx.signal_abort() + return StepResult(action="stop") + + with ( + patch( + "flocks.session.runtime.session_turn.Message.parts", + AsyncMock(return_value=[]), + ), + patch( + "flocks.session.runtime.session_turn.Provider.resolve_model_info", + return_value=(0, 0, None), + ), + patch( + "flocks.session.runtime.continuation_policy.ContinuationPolicy.run_user_prompt_submit", + AsyncMock(), + ), + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_after", + run_turn_after, + ), + patch( + "flocks.session.runtime.step_engine.StepEngine._process_step", + AsyncMock(side_effect=process_step), + ), + patch( + "flocks.session.lifecycle.title.SessionTitle.ensure_title", + MagicMock(return_value=None), + ), + patch( + "flocks.session.runtime.session_turn.fire_and_forget", + MagicMock(), + ), + ): + await run_logical_turns(ctx, LoopCallbacks()) + + run_turn_after.assert_not_awaited() + + +@pytest.mark.asyncio +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") + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user], + [user, assistant], + ] + ) + ) + run_turn_after = AsyncMock() + + async def process_step(*_args, **_kwargs): + ctx.signal_abort() + return StepResult(action="stop") + + with ( + patch( + "flocks.session.runtime.session_turn.Message.parts", + AsyncMock(return_value=[]), + ), + patch( + "flocks.session.runtime.session_turn.Message.get_text_content", + AsyncMock(return_value="not done"), + ), + patch( + "flocks.session.runtime.session_turn.Message.create", + AsyncMock(return_value=goal_user), + ), + patch( + "flocks.session.runtime.session_turn.Provider.resolve_model_info", + return_value=(0, 0, None), + ), + patch( + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", + AsyncMock( + return_value=GoalDecision( + status="active", + verdict="continue", + should_continue=True, + continuation_prompt="continue the goal", + ) + ), + ), + patch( + "flocks.session.runtime.continuation_policy.ContinuationPolicy.run_user_prompt_submit", + AsyncMock(), + ), + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_after", + run_turn_after, + ), + patch( + "flocks.session.runtime.step_engine.StepEngine._process_step", + AsyncMock(side_effect=process_step), + ), + patch( + "flocks.session.lifecycle.title.SessionTitle.ensure_title", + MagicMock(return_value=None), + ), + patch( + "flocks.session.runtime.session_turn.fire_and_forget", + MagicMock(), + ), + ): + await run_logical_turns(ctx, LoopCallbacks()) + + run_turn_after.assert_not_awaited() + + +@pytest.mark.asyncio +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_after = AsyncMock() + + async def cancel_for_user_abort(*_args, **_kwargs): + ctx.signal_abort() + raise asyncio.CancelledError + + with ( + patch( + "flocks.session.runtime.session_turn.Message.parts", + AsyncMock(return_value=[]), + ), + patch( + "flocks.session.runtime.session_turn.Provider.resolve_model_info", + return_value=(0, 0, None), + ), + patch( + "flocks.session.runtime.continuation_policy.ContinuationPolicy.run_user_prompt_submit", + AsyncMock(), + ), + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_after", + run_turn_after, + ), + patch( + "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.runtime.session_turn.fire_and_forget", + MagicMock(), + ), + ): + await run_logical_turns(ctx, LoopCallbacks()) + + run_turn_after.assert_not_awaited() + + +@pytest.mark.asyncio +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( + get_messages=AsyncMock( + side_effect=[ + [user], + [user, assistant], + ] + ) + ) + run_turn_after = AsyncMock() + + async def abort_after_step(_step: int) -> None: + ctx.abort_event.set() + + with ( + patch( + "flocks.session.runtime.session_turn.Message.parts", + AsyncMock(return_value=[]), + ), + patch( + "flocks.session.runtime.session_turn.Message.get_text_content", + AsyncMock(return_value="final response"), + ), + patch( + "flocks.session.runtime.session_turn.Provider.resolve_model_info", + return_value=(0, 0, None), + ), + patch( + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", + AsyncMock( + return_value=GoalDecision( + status="inactive", + verdict="inactive", + ) + ), + ), + patch( + "flocks.session.runtime.continuation_policy.ContinuationPolicy.run_user_prompt_submit", + AsyncMock(), + ), + patch.object( + DEFAULT_CONTINUATION_POLICY, + "run_turn_after", + run_turn_after, + ), + patch( + "flocks.session.runtime.step_engine.StepEngine._process_step", + AsyncMock(return_value=StepResult(action="stop")), + ), + patch( + "flocks.session.lifecycle.title.SessionTitle.ensure_title", + MagicMock(return_value=None), + ), + patch( + "flocks.session.runtime.session_turn.fire_and_forget", + MagicMock(), + ), + ): + result = await run_logical_turns( + ctx, + LoopCallbacks(on_step_end=abort_after_step), + ) + + run_turn_after.assert_not_awaited() + assert result.metadata["aborted"] is True diff --git a/tests/session/test_message_parts_persistence.py b/tests/session/test_message_parts_persistence.py index a827828e2..e44bdea88 100644 --- a/tests/session/test_message_parts_persistence.py +++ b/tests/session/test_message_parts_persistence.py @@ -209,6 +209,40 @@ async def test_delete_restores_caches_when_message_persistence_fails( assert [message["id"] for message in stored_messages] == ["msg_a"] +@pytest.mark.asyncio +async def test_create_restores_caches_when_persistence_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session_id = "ses_parts_create_message_failure" + await Message.create( + session_id, + MessageRole.USER, + "existing", + id="msg_existing", + part_id="part_existing", + ) + monkeypatch.setattr( + Storage, + "mutate_many", + AsyncMock(side_effect=RuntimeError("message storage unavailable")), + ) + + with pytest.raises(RuntimeError, match="message storage unavailable"): + await Message.create( + session_id, + MessageRole.USER, + "ghost", + id="msg_ghost", + part_id="part_ghost", + ) + + assert await Message.get(session_id, "msg_ghost") is None + assert await Message.parts(session_id, "msg_ghost") == [] + assert [message.id for message in await Message.list(session_id)] == [ + "msg_existing" + ] + + @pytest.mark.asyncio async def test_delete_removes_message_and_parts_atomically() -> None: session_id = "ses_parts_delete_parts_failure" 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 29fa621ef..db9dfc61f 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,12 @@ 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.runtime.contracts import ActiveModelAttempt, ModelRequest 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 +28,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", @@ -39,6 +40,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 +51,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 @@ -184,7 +187,12 @@ async def _after(payload, result): AsyncMock(side_effect=_after), ) monkeypatch.setattr( - runner_mod.SessionRunner, + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock(return_value=True), + ) + monkeypatch.setattr( + runner_mod.StepEngine, "_end_observability", staticmethod(lambda *args, **kwargs: None), ) @@ -249,6 +257,81 @@ async def _gen(): assert order == ["before", "provider", "after"] +@pytest.mark.asyncio +async def test_call_llm_cancellation_closes_persisted_attempt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner = _make_runner("ses_runner_llm_cancelled") + assistant_msg = SimpleNamespace(id="msg_assistant_cancelled") + agent = SimpleNamespace(name="rex") + stream_waiting = asyncio.Event() + + monkeypatch.setattr(runner_mod, "StreamProcessor", _FakeProcessor) + monkeypatch.setattr( + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock(return_value=True), + ) + monkeypatch.setattr( + runner_mod.HookPipeline, + "run_llm_before", + AsyncMock(return_value=SimpleNamespace(input={}, output={})), + ) + run_after = AsyncMock() + monkeypatch.setattr(runner_mod.HookPipeline, "run_llm_after", run_after) + 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, + ) + update = AsyncMock(return_value=None) + monkeypatch.setattr(runner_mod.Message, "update", update) + + class _Provider: + def chat_stream(self, **_kwargs): + async def _gen(): + yield SimpleNamespace( + delta="partial", + reasoning=None, + tool_calls=None, + event_type=None, + finish_reason=None, + usage=None, + ) + stream_waiting.set() + await asyncio.Event().wait() + + return _gen() + + call = asyncio.create_task( + runner._call_llm( + provider=_Provider(), + messages=[ChatMessage(role="user", content="hello")], + tools=[], + agent=agent, + assistant_msg=assistant_msg, + ) + ) + await asyncio.wait_for(stream_waiting.wait(), timeout=1) + call.cancel() + + with pytest.raises(asyncio.CancelledError): + await call + + update.assert_awaited_once() + assert update.await_args.kwargs["finish"] == "error" + assert update.await_args.kwargs["error"]["name"] == "MessageAbortedError" + run_after.assert_awaited_once() + after_output = run_after.await_args.args[1] + assert after_output["error"]["type"] == "CancelledError" + assert after_output["response"]["content"] == "partial" + assert runner._active_model_attempt is None + + @pytest.mark.asyncio async def test_call_llm_blocks_provider_when_llm_before_hook_fails(monkeypatch: pytest.MonkeyPatch): runner = _make_runner("ses_runner_llm_before_fail_closed") @@ -417,6 +500,202 @@ 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_llm_after_aggregates_same_model_retry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner = _make_runner("ses_runner_llm_hook_retry_pair") + runner._llm_retry_scope_active = True + run_after = AsyncMock() + monkeypatch.setattr(runner_mod.HookPipeline, "run_llm_after", run_after) + metadata = {"messageID": "msg_assistant_hook_retry_pair"} + request = ModelRequest( + provider_id="anthropic", + model_id="claude-sonnet", + messages=(ChatMessage(role="user", content="retry request"),), + tools=(), + options={}, + ) + runner._active_model_attempt = ActiveModelAttempt( + message_id=metadata["messageID"], + request=request, + hook_metadata=metadata, + llm_after_enabled=True, + ) + await runner._record_llm_after_attempt( + message_id=metadata["messageID"], + output={"error": {"message": "first attempt failed"}}, + ) + await runner._record_llm_after_attempt( + message_id=metadata["messageID"], + output={"action": "stop"}, + ) + run_after.assert_not_awaited() + await runner._emit_pending_llm_after(metadata["messageID"]) + + run_after.assert_awaited_once() + after_payload = run_after.await_args.args[1] + assert after_payload["attemptCount"] == 2 + assert after_payload["failedAttempts"] == [ + {"error": {"message": "first attempt failed"}} + ] + + +@pytest.mark.asyncio +async def test_terminal_llm_after_releases_frozen_request() -> None: + runner = _make_runner("ses_runner_llm_request_cleanup") + message_id = "msg_assistant_request_cleanup" + request = ModelRequest( + provider_id="anthropic", + model_id="claude-sonnet", + messages=(ChatMessage(role="user", content="large request"),), + tools=(), + options={}, + ) + runner._active_model_attempt = ActiveModelAttempt( + message_id=message_id, + request=request, + hook_metadata={"messageID": message_id}, + llm_after_enabled=False, + ) + + await runner._emit_pending_llm_after(message_id) + + assert runner._active_model_attempt is None + + @pytest.mark.asyncio async def test_call_llm_emits_after_hook_on_error(monkeypatch: pytest.MonkeyPatch): runner = _make_runner("ses_runner_llm_hooks_error") @@ -452,7 +731,12 @@ async def _after(payload, result): AsyncMock(side_effect=_after), ) monkeypatch.setattr( - runner_mod.SessionRunner, + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock(return_value=True), + ) + monkeypatch.setattr( + runner_mod.StepEngine, "_end_observability", staticmethod(lambda *args, **kwargs: None), ) @@ -539,6 +823,10 @@ async def _execute_delegate(tool_name, ctx, **_kwargs): "flocks.session.streaming.stream_processor.ToolRegistry.execute", _execute_delegate, ) + monkeypatch.setattr( + "flocks.session.streaming.tool_accumulator.ToolRegistry.get_schema", + lambda _tool_name: None, + ) monkeypatch.setattr( StreamProcessor, "_resolve_sandbox_meta", 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_shell_hook.py b/tests/session/test_runner_shell_hook.py index 62e256cfd..59aec70c0 100644 --- a/tests/session/test_runner_shell_hook.py +++ b/tests/session/test_runner_shell_hook.py @@ -7,7 +7,7 @@ import pytest from flocks.hooks.pipeline import HookBase, HookPipeline -from flocks.session.runner import SessionRunner +from flocks.session.actions import run_session_shell from flocks.session.tool_execution import build_session_tool_execution_payload @@ -38,11 +38,11 @@ async def tool_before(self, ctx): ) create_process = AsyncMock(return_value=process) monkeypatch.setattr( - "flocks.session.runner.Session.get_by_id", + "flocks.session.actions.Session.get_by_id", AsyncMock(return_value=SimpleNamespace(directory=str(tmp_path))), ) monkeypatch.setattr( - "flocks.session.runner.Message.create", + "flocks.session.actions.Message.create", AsyncMock( side_effect=[ SimpleNamespace(id="msg_user"), @@ -51,12 +51,12 @@ async def tool_before(self, ctx): ), ) monkeypatch.setattr( - "flocks.session.runner.asyncio.create_subprocess_shell", + "flocks.session.actions.asyncio.create_subprocess_shell", create_process, ) HookPipeline.register("capture.action", CaptureAction()) - result = await SessionRunner.shell( + result = await run_session_shell( session_id="ses_1", agent="build", command="echo ok", diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index efee61f53..0a9466220 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, LoopContext from flocks.session.prompt import SessionPrompt, get_prompt_flocks_config_guard 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 = { @@ -757,7 +757,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" @@ -823,7 +823,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" @@ -882,7 +882,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" @@ -943,7 +943,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" @@ -989,7 +989,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", @@ -1033,7 +1033,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" @@ -1058,7 +1058,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", @@ -1094,7 +1094,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={ @@ -1158,7 +1158,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", @@ -1170,7 +1170,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", @@ -1205,13 +1205,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", @@ -1249,13 +1249,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", @@ -1289,7 +1289,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) @@ -1307,7 +1307,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"] @@ -1324,7 +1324,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"] @@ -1340,7 +1340,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"] @@ -1349,7 +1349,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", @@ -1358,7 +1358,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", @@ -1367,7 +1367,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", @@ -1376,7 +1376,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", @@ -1387,7 +1387,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", @@ -1396,7 +1396,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", @@ -1406,7 +1406,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", @@ -1433,7 +1433,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", @@ -1466,7 +1466,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", @@ -1489,7 +1489,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", @@ -1519,7 +1519,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], []) @@ -1553,6 +1553,30 @@ async def test_to_chat_messages_invalidates_shared_cache_when_message_parts_chan assert second_messages[1].content == "Error: Tool execution was interrupted" +@pytest.mark.asyncio +async def test_to_chat_messages_only_projects_appended_history(monkeypatch): + runner = StepEngine( + session=_make_session("ses_incremental_chat_projection"), + static_cache={}, + ) + first = SimpleNamespace(id="msg_first", role="user", content="first") + second = SimpleNamespace(id="msg_second", role="assistant", content="second") + parts = AsyncMock(return_value=[]) + text_content = AsyncMock( + side_effect=lambda message: message.content, + ) + monkeypatch.setattr(runner_mod.Message, "parts", parts) + monkeypatch.setattr(runner_mod.Message, "get_text_content", text_content) + + first_projection = await runner._to_chat_messages([first], []) + parts.reset_mock() + second_projection = await runner._to_chat_messages([first, second], []) + + assert [call.args[0] for call in parts.await_args_list] == [second.id] + assert first_projection[0] is second_projection[0] + assert [message.content for message in second_projection] == ["first", "second"] + + @pytest.mark.asyncio async def test_to_chat_messages_excludes_ignored_assistant_text(): session = await Session.create( @@ -1567,7 +1591,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], []) @@ -1585,7 +1609,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, @@ -1634,7 +1658,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" @@ -1701,7 +1725,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" @@ -1763,7 +1787,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" @@ -1829,7 +1853,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" @@ -1893,7 +1917,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" @@ -1958,7 +1982,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" @@ -2038,7 +2062,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" @@ -2113,7 +2137,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} @@ -2194,7 +2218,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", @@ -2220,7 +2244,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(), ) @@ -2274,7 +2298,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( @@ -2325,7 +2349,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", @@ -2372,7 +2396,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" @@ -2398,7 +2422,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", @@ -2461,7 +2485,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") @@ -2491,7 +2515,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 @@ -2530,7 +2554,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, @@ -2555,12 +2579,25 @@ 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, ) 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 = [ @@ -2609,7 +2646,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, ) @@ -2624,7 +2661,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 @@ -2662,12 +2699,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) @@ -2693,7 +2730,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", @@ -2744,7 +2781,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", @@ -2793,7 +2830,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", @@ -2848,7 +2885,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", @@ -2908,7 +2945,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", @@ -2957,7 +2994,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", @@ -3011,7 +3048,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( @@ -3049,7 +3086,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) @@ -3060,7 +3097,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( @@ -3098,7 +3135,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) @@ -3141,7 +3178,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", @@ -3153,8 +3190,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_session_abort_inject.py b/tests/session/test_session_abort_inject.py index 2f90b3da0..950ecd332 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,13 +14,19 @@ import pytest +from flocks.session.runtime.continuation_policy import DEFAULT_CONTINUATION_POLICY +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.prompt import SessionPrompt -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: @@ -56,65 +62,30 @@ 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.""" - external_event = asyncio.Event() + def test_runner_uses_supplied_abort_event(self): + """StepEngine should share the supplied turn abort event.""" + abort_event = asyncio.Event() session_info = _make_session_info() - runner = SessionRunner( + runner = StepEngine( session=session_info, - abort_event=external_event, + abort_event=abort_event, ) - # Initially not aborted - assert runner.is_aborted is False - - # Set external event → runner should report aborted - external_event.set() - assert runner.is_aborted is True - - def test_runner_internal_abort_still_works(self): - """SessionRunner's own abort() method should still work.""" - session_info = _make_session_info() - runner = SessionRunner(session=session_info) - - assert runner.is_aborted is False - runner.abort() - assert runner.is_aborted is True - - def test_runner_either_abort_triggers(self): - """Either internal or external abort should trigger is_aborted.""" - external_event = asyncio.Event() - session_info = _make_session_info() - - runner = SessionRunner( - session=session_info, - abort_event=external_event, - ) - - # Neither set → not aborted - assert runner.is_aborted is False - - # Only external set - external_event.set() - assert runner.is_aborted is True - - # Clear external, set internal - external_event.clear() - runner._abort.clear() + assert runner._abort is abort_event assert runner.is_aborted is False runner.abort() + assert abort_event.is_set() assert runner.is_aborted is True - def test_runner_without_external_event(self): - """Runner created without abort_event should still work normally.""" + def test_runner_creates_abort_event_when_omitted(self): + """Standalone engines should create their own abort event.""" 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 runner.abort() assert runner.is_aborted is True @@ -131,12 +102,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(), @@ -187,7 +158,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 @@ -196,10 +167,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") @@ -209,31 +180,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) - - def test_get_context(self): - """get_context should return the LoopContext for a running session.""" - session_info = _make_session_info("ctx_get_test") - ctx = LoopContext( - session=session_info, - provider_id="test", - model_id="test", - agent_name="test", - ) - SessionLoop._active_loops["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("running_test", None) # --------------------------------------------------------------------------- # _should_exit logic with injected messages @@ -243,61 +195,75 @@ class TestShouldExitWithInject: """Test that _should_exit correctly handles injected user messages.""" @staticmethod - def _make_msg(msg_id: str, role: str, finish: str = None): + def _make_msg( + msg_id: str, + role: str, + finish: str = None, + *, + parent_id: str | None = None, + ): """Create a minimal message-like object for testing.""" msg = type("Msg", (), {})() msg.id = msg_id msg.role = role msg.finish = finish - msg.parentID = None + msg.parentID = parent_id return msg - def test_exit_when_assistant_replies_to_user_and_finished(self): - """Should exit if last assistant is a finished reply to last user.""" + def test_exit_when_assistant_replies_to_user_with_non_monotonic_ids(self): + """Should use parent linkage instead of ordering generated IDs.""" last_user = self._make_msg("msg_002", "user") - last_assistant = self._make_msg("msg_001", "assistant", finish="stop") - last_assistant.parentID = last_user.id + last_assistant = self._make_msg( + "msg_001", + "assistant", + finish="stop", + parent_id=last_user.id, + ) - assert SessionLoop._should_exit(last_user, last_assistant) is True + assert RuntimeLoopContext._should_exit(last_user, last_assistant) is True - def test_no_exit_when_assistant_belongs_to_previous_user(self): - """Should NOT exit when a new user message follows the assistant. + def test_no_exit_when_user_injected_after_assistant(self): + """Should NOT exit when a new user message appears after the assistant. - This is the core inject scenario: the last assistant belongs to the - preceding user turn, so the loop should continue regardless of IDs. + This is the core inject scenario: the injected user message has a + higher ID than the last assistant message, so the loop should continue. """ - last_user = self._make_msg("msg_001", "user") - last_assistant = self._make_msg("msg_002", "assistant", finish="stop") - last_assistant.parentID = "msg_previous_user" + last_user = self._make_msg("msg_003", "user") # injected message + last_assistant = self._make_msg( + "msg_002", + "assistant", + finish="stop", + parent_id="msg_001", + ) - assert SessionLoop._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 SessionLoop._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 SessionLoop._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 SessionLoop._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 SessionLoop._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.""" @@ -305,7 +271,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 RuntimeLoopContext._should_exit( last_user, last_assistant, last_assistant_parts, @@ -324,7 +290,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, @@ -338,7 +304,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, @@ -350,14 +316,22 @@ async def test_detects_newer_user_when_step_failed_before_assistant_created(self class TestTurnLifecycle: @staticmethod - def _make_msg(msg_id: str, role: str, finish: str = None, *, tokens=None, summary: bool = False): + def _make_msg( + msg_id: str, + role: str, + finish: str = None, + *, + tokens=None, + summary: bool = False, + parent_id: str | None = None, + ): msg = type("Msg", (), {})() msg.id = msg_id msg.role = role msg.finish = finish - msg.parentID = None msg.tokens = tokens msg.summary = summary + msg.parentID = parent_id return msg @pytest.mark.asyncio @@ -374,11 +348,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] @@ -401,11 +375,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,17 +402,30 @@ async def test_run_loop_continues_for_active_goal_after_stop(self): agent_name="rex", ) user = self._make_msg("msg_001", "user") - assistant = self._make_msg("msg_002", "assistant", finish="stop") - assistant.parentID = user.id + assistant = self._make_msg( + "msg_002", + "assistant", + finish="stop", + parent_id=user.id, + ) goal_user = self._make_msg("msg_003", "user") - assistant_after_goal = self._make_msg("msg_004", "assistant", finish="stop") - assistant_after_goal.parentID = goal_user.id - ctx.session_ctx = SimpleNamespace( + assistant_after_goal = self._make_msg( + "msg_004", + "assistant", + finish="stop", + parent_id=goal_user.id, + ) + 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], + [user, assistant, goal_user, assistant_after_goal], ]) ) event_callback = AsyncMock() @@ -455,28 +442,27 @@ 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" - assert result.last_message is assistant_after_goal create_message.assert_awaited_once() assert create_message.await_args.kwargs["content"] == "continue toward goal" assert create_message.await_args.kwargs["synthetic"] is True @@ -507,31 +493,38 @@ async def test_run_loop_waits_for_user_input_after_goal_clarification(self): agent_name="rex", ) user = self._make_msg("msg_001", "user") - assistant = self._make_msg("msg_002", "assistant", finish="stop") - assistant.parentID = user.id - ctx.session_ctx = SimpleNamespace( + assistant = self._make_msg( + "msg_002", + "assistant", + finish="stop", + parent_id=user.id, + ) + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[ [user], [user, assistant], + [user, assistant], + [user, assistant], + [user, assistant], ]) ) 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="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", @@ -539,10 +532,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() @@ -564,12 +557,19 @@ async def test_run_loop_passes_pending_question_to_goal_judge(self): agent_name="rex", ) user = self._make_msg("msg_001", "user") - assistant = self._make_msg("msg_002", "assistant", finish="stop") - assistant.parentID = user.id - ctx.session_ctx = SimpleNamespace( + assistant = self._make_msg( + "msg_002", + "assistant", + finish="stop", + parent_id=user.id, + ) + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[ [user], [user, assistant], + [user, assistant], + [user, assistant], + [user, assistant], ]) ) event_callback = AsyncMock() @@ -582,28 +582,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() @@ -627,26 +627,38 @@ async def test_run_loop_publishes_goal_terminal_status(self): ) messages = [ self._make_msg("msg_001", "user"), - self._make_msg("msg_002", "assistant", finish="stop"), + self._make_msg( + "msg_002", + "assistant", + finish="stop", + parent_id="msg_001", + ), ] - messages[1].parentID = messages[0].id - ctx.session_ctx = SimpleNamespace( - get_messages=AsyncMock(side_effect=[[messages[0]], messages]) + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [messages[0]], + messages, + messages, + messages, + 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", @@ -654,10 +666,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] @@ -702,26 +714,35 @@ 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( - get_messages=AsyncMock(side_effect=[overflow_messages, normal_messages, normal_messages]) + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + overflow_messages, + normal_messages, + normal_messages, + normal_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", - AsyncMock(side_effect=[0, 50_000, 0, 0]), + "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] @@ -750,85 +771,47 @@ async def test_post_observation_tool_delta_combines_with_observed_prompt(self): model_id="test-model", agent_name="rex", ) - messages = [ - self._make_msg("stale_usage_user", "user"), - self._make_msg( - "stale_usage_assistant", - "assistant", - finish="tool-calls", - tokens={"input": 95_000, "output": 0, "cache": {"read": 0, "write": 0}}, - ), - ] - messages[0].content = "h" * 260_000 - ctx.session_ctx = SimpleNamespace(get_messages=AsyncMock(return_value=messages)) - tool_parts = [ - ToolPart( - sessionID=session.id, - messageID="stale_usage_assistant", - callID=f"call_delta_{index}", - tool="bash", - state=ToolStateCompleted( - input={"command": f"produce output {index}"}, - output="x" * 80_000, - title="bash", - metadata={}, - time={"start": index, "end": index + 1}, - ), - ) - for index in range(2) - ] - run_compaction = AsyncMock(return_value="stop") - parts_by_message = {"stale_usage_assistant": []} - truncation_calls = 0 - - async def truncate_one_tool_result(*args, **kwargs): # noqa: ARG001 - nonlocal truncation_calls - truncation_calls += 1 - if truncation_calls == 1: - tool_parts[0].state.time["compacted"] = 1 - return 1 - return 0 + assistant = self._make_msg( + "stale_usage_assistant", + "assistant", + finish="tool-calls", + tokens={ + "input": 95_000, + "output": 2_000, + "reasoning": 3_000, + "cache": {"read": 0, "write": 0}, + }, + ) + later_user = self._make_msg("stale_usage_later_user", "user") + messages = [assistant, later_user] + cleanup = AsyncMock(return_value=object()) with patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(128_000, 8_192, None), ), patch( - "flocks.session.session_loop.Message.parts", - AsyncMock( - side_effect=lambda message_id, _session_id: ( - list(parts_by_message.get(message_id, [])) - ), - ), - ), patch( - "flocks.session.session_loop.SessionCompaction.truncate_oversized_tool_outputs", - AsyncMock(side_effect=truncate_one_tool_result), - ), patch( - "flocks.session.session_loop.SessionCompaction.prune", - AsyncMock(), + "flocks.session.runtime.session_turn.SessionPrompt.invalidate_message_cache", + MagicMock(), ), patch( - "flocks.session.session_loop.run_compaction", - run_compaction, + "flocks.session.runtime.session_turn.SessionPrompt.estimate_tool_result_tokens", + AsyncMock(return_value=20_000), ), patch( - "flocks.session.runner.SessionRunner._process_step", - AsyncMock(return_value=StepResult(action="stop")), + "flocks.session.runtime.session_turn.SessionPrompt.estimate_full_context_tokens", + AsyncMock(return_value=5_000), + ), patch.object( + ctx, + "_prepare_tool_result_cleanup", + cleanup, ): - estimated_tokens = await SessionPrompt.estimate_full_context_tokens( - session.id, - messages, - ) - parts_by_message["stale_usage_assistant"] = tool_parts - result = await SessionLoop._run_loop(ctx, LoopCallbacks()) - current_message_tokens = await SessionPrompt.estimate_full_context_tokens( - session.id, + result = await ctx._prepare_context_window( messages, + later_user, + assistant, ) - assert estimated_tokens < int(128_000 * 0.85) - assert current_message_tokens < int(128_000 * 0.85) - assert result.action == "stop" - run_compaction.assert_awaited_once() - assert truncation_calls == 2 - assert ctx.last_observed_prompt_tokens == 95_000 + assert result is cleanup.return_value + assert cleanup.await_args.args[2] == 125_000 + assert ctx.last_observed_prompt_tokens == 100_000 @pytest.mark.asyncio async def test_run_loop_skips_exit_condition_when_assistant_has_tool_parts(self): @@ -846,11 +829,17 @@ async def test_run_loop_skips_exit_condition_when_assistant_has_tool_parts(self) ) messages = [ self._make_msg("msg_001", "user"), - self._make_msg("msg_002", "assistant", finish="stop"), + self._make_msg( + "msg_002", + "assistant", + finish="stop", + parent_id="msg_001", + ), ] - messages[1].parentID = messages[0].id - ctx.session_ctx = SimpleNamespace( - get_messages=AsyncMock(side_effect=[messages, messages]) + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[messages, messages, messages, messages, messages] + ) ) event_callback = AsyncMock() callbacks = LoopCallbacks(event_publish_callback=event_callback) @@ -858,25 +847,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] @@ -901,10 +890,14 @@ async def test_run_loop_breaks_on_exit_condition_without_tool_parts(self): ) messages = [ self._make_msg("msg_001", "user"), - self._make_msg("msg_002", "assistant", finish="stop"), + self._make_msg( + "msg_002", + "assistant", + finish="stop", + parent_id="msg_001", + ), ] - messages[1].parentID = messages[0].id - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(return_value=messages) ) event_callback = AsyncMock() @@ -913,16 +906,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] @@ -932,74 +925,9 @@ async def test_run_loop_breaks_on_exit_condition_without_tool_parts(self): assert event_names == ["turn.started"] @pytest.mark.asyncio - async def test_run_loop_processes_new_user_when_ids_are_not_monotonic( - self, - ) -> None: - session = SimpleNamespace( - id="loop_non_monotonic_id_session", - agent="rex", - directory="/tmp", - memory_enabled=False, - ) - ctx = LoopContext( - session=session, - provider_id="test-provider", - model_id="test-model", - agent_name="rex", - ) - previous_user = self._make_msg("msg_previous_user", "user") - previous_assistant = self._make_msg( - "msg_ffedcf5c6001TWU0fGZXuDeY00", "assistant", finish="stop" - ) - previous_assistant.parentID = previous_user.id - current_user = self._make_msg( - "msg_ffed09fd1001S0plX81SJ55NUz", - "user", - ) - current_assistant = self._make_msg( - "msg_current_assistant", "assistant", finish="stop" - ) - current_assistant.parentID = current_user.id - messages_before_step = [previous_user, previous_assistant, current_user] - messages_after_step = [*messages_before_step, current_assistant] - ctx.session_ctx = SimpleNamespace( - get_messages=AsyncMock( - side_effect=[messages_before_step, messages_after_step] - ) - ) - process_step = AsyncMock(return_value=StepResult(action="stop")) - - with patch( - "flocks.session.session_loop.Message.parts", - AsyncMock(return_value=[]), - ), patch( - "flocks.session.session_loop.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", - MagicMock(), - ), patch( - "flocks.session.runner.SessionRunner._process_step", - process_step, - ): - result = await SessionLoop._run_loop(ctx, LoopCallbacks()) - - assert result.last_message is current_assistant - process_step.assert_awaited_once() - - @pytest.mark.asyncio - async def test_run_loop_does_not_return_previous_reply_when_current_step_fails( - self, - ) -> None: - session = SimpleNamespace( - id="loop_failed_current_turn_session", - agent="rex", - directory="/tmp", - memory_enabled=False, - ) + async def test_commit_step_does_not_reuse_previous_assistant_reply(self): + """A failed current step must not surface an older assistant reply.""" + session = SimpleNamespace(id="stale_reply_session") ctx = LoopContext( session=session, provider_id="test-provider", @@ -1011,96 +939,23 @@ async def test_run_loop_does_not_return_previous_reply_when_current_step_fails( "msg_previous_assistant", "assistant", finish="stop", + parent_id=previous_user.id, ) - previous_assistant.parentID = previous_user.id current_user = self._make_msg("msg_current_user", "user") - messages = [previous_user, previous_assistant, current_user] - ctx.session_ctx = SimpleNamespace( - get_messages=AsyncMock(side_effect=[messages, messages]) - ) - on_error = AsyncMock() - process_step = AsyncMock( - return_value=StepResult(action="stop", error="provider failed") - ) - - with patch( - "flocks.session.session_loop.Message.parts", - AsyncMock(return_value=[]), - ), patch( - "flocks.session.session_loop.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", - MagicMock(), - ), patch( - "flocks.session.runner.SessionRunner._process_step", - process_step, - ): - result = await SessionLoop._run_loop( - ctx, - LoopCallbacks(on_error=on_error), + ctx.prepared_user_id = current_user.id + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock( + return_value=[ + previous_user, + previous_assistant, + current_user, + ] ) - - assert result.last_message is None - on_error.assert_awaited_once_with("provider failed") - process_step.assert_awaited_once() - - -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"}, - )) + boundary = await ctx.commit_step(StepResult(action="stop")) - 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.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) - - 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, - } + assert boundary.last_message is None # --------------------------------------------------------------------------- diff --git a/tests/session/test_session_context.py b/tests/session/test_session_context.py index 77388f586..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_ctx -5. SessionRunner accepts session_ctx """ 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.runner import SessionRunner 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 TestLoopContextSessionCtx: - """LoopContext should carry session_ctx.""" - - def test_loop_context_has_session_ctx_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_ctx is None - - def test_loop_context_with_session_ctx(self): - session = MagicMock() - session.id = "test" - session.directory = "/test" - session.project_id = "proj" - - session_ctx = DefaultSessionContext(session) - ctx = LoopContext( - session=session, - provider_id="anthropic", - model_id="claude-sonnet-4", - agent_name="rex", - session_ctx=session_ctx, - ) - assert ctx.session_ctx is session_ctx - assert ctx.session_ctx.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 TestRunnerSessionCtx: - """SessionRunner should accept session_ctx.""" - - def test_runner_accepts_session_ctx(self): - session = MagicMock() - session.id = "test" - session.directory = "/test" - session.project_id = "proj" - - session_ctx = DefaultSessionContext(session) - runner = SessionRunner( - session=session, - session_ctx=session_ctx, - ) - assert runner.session_ctx is session_ctx - - def test_runner_session_ctx_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 diff --git a/tests/session/test_session_loop_working_directory.py b/tests/session/test_session_loop_working_directory.py index 4c2cc5f3a..9794edfe8 100644 --- a/tests/session/test_session_loop_working_directory.py +++ b/tests/session/test_session_loop_working_directory.py @@ -3,9 +3,14 @@ import pytest from flocks.bus.bus import Bus +from flocks.session.runtime.agent_loop import AgentLoop +from flocks.session.runtime.contracts import ( + AgentRunOutcome, + AgentRunStatus, +) 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 +21,18 @@ 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): + return AgentRunOutcome( + status=AgentRunStatus.ABORTED, + ) + + 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 +47,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..e9e3b349a 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) @@ -113,4 +113,3 @@ async def fake_call_llm(self, provider, messages, tools, agent, assistant_msg): # Critical assertion: tools must remain available (not cleared to []). assert captured["tools"] == sentinel_tools - diff --git a/tests/session_runtime_testkit.py b/tests/session_runtime_testkit.py new file mode 100644 index 000000000..026484393 --- /dev/null +++ b/tests/session_runtime_testkit.py @@ -0,0 +1,35 @@ +"""Test-only helpers for exercising SessionLoop logical-turn control.""" + +from unittest.mock import AsyncMock, patch + +from flocks.session.session_loop import ( + LoopCallbacks, + LoopContext, + LoopResult, + SessionLoop, +) + + +async def run_logical_turns( + turn: LoopContext, + callbacks: LoopCallbacks, +) -> LoopResult: + """Run the production owned-loop path with an in-memory test lease.""" + turn.callbacks = callbacks + lease = SessionLoop._leases.acquire(turn.session.id, turn) + if lease is None: + raise RuntimeError(f"Session {turn.session.id} already has a test lease") + + with ( + patch.object( + turn, + "has_late_input", + AsyncMock(return_value=False), + ), + patch.object(SessionLoop, "_publish_released", AsyncMock()), + ): + try: + return await SessionLoop._run_owned_loop(lease, callbacks) + finally: + if SessionLoop._leases.owns(lease): + SessionLoop._finalize_release_state_locked(lease) 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 diff --git a/tui/flocks/command/index.test.ts b/tui/flocks/command/index.test.ts new file mode 100644 index 000000000..958ece053 --- /dev/null +++ b/tui/flocks/command/index.test.ts @@ -0,0 +1,28 @@ +import { describe, expect, test } from "bun:test" +import { mkdtemp, rm } from "node:fs/promises" +import { tmpdir } from "node:os" +import path from "node:path" +import { Agent } from "../agent/agent" +import { Instance } from "../project/instance" +import { Command } from "./index" + +describe("built-in commands", () => { + test("review uses an available primary agent", async () => { + const directory = await mkdtemp(path.join(tmpdir(), "flocks-review-command-")) + try { + await Instance.provide({ + directory, + fn: async () => { + const command = await Command.get(Command.Default.REVIEW) + const agentName = command.agent ?? (await Agent.defaultAgent()) + + expect(command.agent).toBeUndefined() + expect(await Agent.get(agentName)).toBeDefined() + }, + }) + } finally { + await Instance.disposeAll() + await rm(directory, { recursive: true, force: true }) + } + }) +})