diff --git a/flocks/agent/agents/hephaestus/prompt_builder.py b/flocks/agent/agents/hephaestus/prompt_builder.py index 9617affcb..3d220164b 100644 --- a/flocks/agent/agents/hephaestus/prompt_builder.py +++ b/flocks/agent/agents/hephaestus/prompt_builder.py @@ -29,7 +29,6 @@ def inject( available_agents=available_agents, available_tools=tools, available_skills=skills, - use_task_system=False, ) @@ -37,7 +36,6 @@ def build_hephaestus_prompt( available_agents: List["AvailableAgent"], available_tools: List["AvailableTool"], available_skills: List["AvailableSkill"], - use_task_system: bool = False, ) -> str: from flocks.agent.prompt_utils import ( build_agent_selection_table, @@ -62,7 +60,7 @@ def build_hephaestus_prompt( oracle_section = build_oracle_section(available_agents) hard_blocks = build_hard_blocks_section() anti_patterns = build_anti_patterns_section() - todo_discipline = _todo_discipline_section(use_task_system) + todo_discipline = _todo_discipline_section() template = """You are Hephaestus, an autonomous deep worker for software engineering. @@ -245,44 +243,7 @@ def build_hephaestus_prompt( return prompt -def _todo_discipline_section(use_task_system: bool) -> str: - if use_task_system: - return """## Task Discipline (NON-NEGOTIABLE) - -**Track ALL multi-step work with tasks. This is your execution backbone.** - -### When to Create Tasks (MANDATORY) - -| Trigger | Action | -|---------|--------| -| 2+ step task | `TaskCreate` FIRST, atomic breakdown | -| Uncertain scope | `TaskCreate` to clarify thinking | -| Complex single task | Break down into trackable steps | - -### Workflow (STRICT) - -1. **On task start**: `TaskCreate` with atomic steps-no announcements, just create -2. **Before each step**: `TaskUpdate(status="in_progress")` (ONE at a time) -3. **After each step**: `TaskUpdate(status="completed")` IMMEDIATELY (NEVER batch) -4. **Scope changes**: Update tasks BEFORE proceeding - -### Why This Matters - -- **Execution anchor**: Tasks prevent drift from original request -- **Recovery**: If interrupted, tasks enable seamless continuation -- **Accountability**: Each task = explicit commitment to deliver - -### Anti-Patterns (BLOCKING) - -| Violation | Why It Fails | -|-----------|--------------| -| Skipping tasks on multi-step work | Steps get forgotten, user has no visibility | -| Batch-completing multiple tasks | Defeats real-time tracking purpose | -| Proceeding without `in_progress` | No indication of current work | -| Finishing without completing tasks | Task appears incomplete | - -**NO TASKS ON MULTI-STEP WORK = INCOMPLETE WORK.**""" - +def _todo_discipline_section() -> str: return """## Todo Discipline (NON-NEGOTIABLE) **Track ALL multi-step work with todos. This is your execution backbone.** diff --git a/flocks/agent/agents/rex/prompt_builder.py b/flocks/agent/agents/rex/prompt_builder.py index e74413457..d4bb8d1ee 100644 --- a/flocks/agent/agents/rex/prompt_builder.py +++ b/flocks/agent/agents/rex/prompt_builder.py @@ -30,7 +30,6 @@ def inject( available_tools=tools, available_skills=skills, available_workflows=workflows or [], - use_task_system=False, ) @@ -39,7 +38,6 @@ def build_dynamic_rex_prompt( available_tools: List["AvailableTool"], available_skills: List["AvailableSkill"], available_workflows: Optional[List["AvailableWorkflow"]] = None, - use_task_system: bool = False, ) -> str: from flocks.agent.prompt_utils import ( build_agent_selection_table, @@ -58,12 +56,8 @@ def build_dynamic_rex_prompt( im_send_section = _build_im_send_pointer_section() anti_patterns = _build_rex_anti_patterns_section() command_guidance_section = _build_command_guidance_section() - task_management_section = _task_management_section(use_task_system) - todo_hook_note = ( - "YOUR TASK CREATION WOULD BE TRACKED BY HOOK([SYSTEM REMINDER - TASK CONTINUATION])" - if use_task_system - else "YOUR TODO CREATION WOULD BE TRACKED BY HOOK([SYSTEM REMINDER - TODO CONTINUATION])" - ) + task_management_section = _task_management_section() + todo_hook_note = "YOUR TODO CREATION WOULD BE TRACKED BY HOOK([SYSTEM REMINDER - TODO CONTINUATION])" template = """ You are "Rex" - Powerful AI orchestrator for security operations. @@ -143,7 +137,7 @@ def build_dynamic_rex_prompt( - Match existing codebase patterns when editing. - Fix bugs minimally; do not refactor during a bugfix unless required. - Keep search bounded: stop when you have enough context, when results repeat, or when direct evidence already answers the question. -- For independent parallel branches whose results are needed this turn, emit multiple foreground `delegate_task` / `task` tool calls in the same assistant turn. The runtime executes those sibling tool calls concurrently and returns all tool results before you continue. +- For independent parallel branches whose results are needed this turn, emit multiple foreground `delegate_task` tool calls in the same assistant turn. The runtime executes those sibling tool calls concurrently and returns all tool results before you continue. - Do not use `run_in_background=true`; background subagent execution is disabled. ## 5. Verify @@ -278,55 +272,42 @@ def _build_clarification_protocol() -> str: ```""" -def _task_management_section(use_task_system: bool) -> str: - title = "Task Management" if use_task_system else "Todo Management" - unit = "tasks" if use_task_system else "todos" - create_action = "`TaskCreate`" if use_task_system else '`todo(action="write")`' - progress_action = ( - '`TaskUpdate(status="in_progress")`' - if use_task_system - else "mark `in_progress`" - ) - complete_action = ( - '`TaskUpdate(status="completed")`' - if use_task_system - else "mark `completed`" - ) +def _task_management_section() -> str: clarification_protocol = _build_clarification_protocol() - return f""" -## {title} + return f""" +## Todo Management -Use {unit} as the primary coordination mechanism for non-trivial execution work. +Use todos as the primary coordination mechanism for non-trivial execution work. ### When They Are Mandatory | Trigger | Action | |---------|--------| -| Multi-step work (2+ steps) | Create {unit} first | -| Uncertain scope | Create {unit} to structure the work | -| User request with multiple items | Create {unit} first | -| Complex single task | Break it into {unit} | +| Multi-step work (2+ steps) | Create todos first | +| Uncertain scope | Create todos to structure the work | +| User request with multiple items | Create todos first | +| Complex single task | Break it into todos | ### Operating Rules -1. Start with {create_action} before implementation work begins. -2. ONLY add {unit} when the user wants execution, not when they only want analysis or planning. -3. Before each step, {progress_action}. Keep only one item in progress. -4. After each step, {complete_action} immediately. Never batch updates. -5. If scope changes, update the {unit} before continuing. +1. Start with `todo(action="write")` before implementation work begins. +2. ONLY add todos when the user wants execution, not when they only want analysis or planning. +3. Before each step, mark it `in_progress`. Keep only one item in progress. +4. After each step, mark it `completed` immediately. Never batch updates. +5. If scope changes, update the todos before continuing. ### Failure Modes | Violation | Why It Breaks the Workflow | |-----------|----------------------------| -| Skipping {unit} on non-trivial work | The user loses progress visibility and steps get dropped | -| Batch-completing multiple {unit} | Real-time tracking becomes meaningless | +| Skipping todos on non-trivial work | The user loses progress visibility and steps get dropped | +| Batch-completing multiple todos | Real-time tracking becomes meaningless | | Proceeding without an in-progress item | It is unclear what is being worked on | | Finishing without closing items | The work appears incomplete | {clarification_protocol} -""" +""" def _build_security_priority_section(available_agents: List["AvailableAgent"]) -> str: diff --git a/flocks/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/commands/import_.py b/flocks/cli/commands/import_.py index fe04b4845..2dd701c8a 100644 --- a/flocks/cli/commands/import_.py +++ b/flocks/cli/commands/import_.py @@ -98,6 +98,17 @@ def _normalize_part_data( metadata = normalized.get("metadata") metadata_dict = metadata if isinstance(metadata, dict) else {} + if part_type == "subtask": + normalized["type"] = "text" + normalized["text"] = "" + normalized["ignored"] = True + normalized["metadata"] = { + **metadata_dict, + "legacyPartType": "subtask", + } + part_type = "text" + metadata_dict = normalized["metadata"] + if "content" in normalized and "text" not in normalized: normalized["text"] = normalized.get("content", "") @@ -135,10 +146,6 @@ def _normalize_part_data( ) elif part_type == "agent": normalized.setdefault("name", metadata_dict.get("name") or normalized.get("content") or "agent") - elif part_type == "subtask": - normalized.setdefault("prompt", metadata_dict.get("prompt") or normalized.get("content", "")) - normalized.setdefault("description", metadata_dict.get("description") or "") - normalized.setdefault("agent", metadata_dict.get("agent") or "agent") elif part_type == "retry": normalized.setdefault("attempt", metadata_dict.get("attempt") or 1) normalized.setdefault("error", metadata_dict.get("error") or {}) diff --git a/flocks/cli/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/command/command.py b/flocks/command/command.py index b75ee92d8..32e06b177 100644 --- a/flocks/command/command.py +++ b/flocks/command/command.py @@ -24,7 +24,6 @@ class CommandDef: template: str agent: Optional[str] = None model: Optional[str] = None - subtask: Optional[bool] = None hidden: bool = False aliases: Tuple[str, ...] = field(default_factory=tuple) visible_surfaces: Tuple[CommandSurface, ...] = ("webui", "tui", "acp", "cli") diff --git a/flocks/config/config.py b/flocks/config/config.py index 0f63e1d13..a8fdc1588 100644 --- a/flocks/config/config.py +++ b/flocks/config/config.py @@ -158,7 +158,6 @@ class CommandConfig(BaseModel): description: Optional[str] = None agent: Optional[str] = None model: Optional[str] = None - subtask: Optional[bool] = None # ==================== Provider Configuration ==================== diff --git a/flocks/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/misc.py b/flocks/server/routes/misc.py index 2605ced40..230e0af2d 100644 --- a/flocks/server/routes/misc.py +++ b/flocks/server/routes/misc.py @@ -151,7 +151,6 @@ async def list_commands() -> List[Dict[str, Any]]: "template": cmd.template, "agent": cmd.agent, "model": cmd.model, - "subtask": cmd.subtask, "hidden": cmd.hidden, "aliases": list(cmd.aliases), "visible_surfaces": list(cmd.visible_surfaces), @@ -192,7 +191,6 @@ async def get_command(name: str) -> Dict[str, Any]: "template": cmd.template, "agent": cmd.agent, "model": cmd.model, - "subtask": cmd.subtask, "hidden": cmd.hidden, "aliases": list(cmd.aliases), "visible_surfaces": list(cmd.visible_surfaces), @@ -241,4 +239,3 @@ async def list_experimental_resources() -> Dict[str, Any]: # Return empty dict - resources are not implemented yet return {} - diff --git a/flocks/server/routes/session.py b/flocks/server/routes/session.py index c0dbf61f5..5841e4e4d 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}) @@ -2091,15 +2082,6 @@ class AgentPartInput(BaseModel): name: str = Field(..., description="Agent name") -class SubtaskPartInput(BaseModel): - """Subtask part input for API compatibility""" - type: Literal["subtask"] = "subtask" - id: Optional[str] = Field(None, description="Part ID") - agent: str = Field(..., description="Agent name") - prompt: str = Field(..., description="Subtask prompt") - description: Optional[str] = Field(None, description="Subtask description") - - class PromptRequest(BaseModel): """ Request to send a prompt/message @@ -3471,7 +3453,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 +4926,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 +4937,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/server/routes/skill.py b/flocks/server/routes/skill.py index 28d0e3d7f..4d16ee70c 100644 --- a/flocks/server/routes/skill.py +++ b/flocks/server/routes/skill.py @@ -164,7 +164,6 @@ class CommandResponse(BaseModel): template: str = Field(..., description="Command template") agent: Optional[str] = Field(None, description="Preferred agent") model: Optional[str] = Field(None, description="Preferred model") - subtask: Optional[bool] = Field(None, description="Run as subtask") hidden: bool = Field(False, description="Hidden from UI") aliases: List[str] = Field(default_factory=list, description="Alternate slash aliases") visible_surfaces: List[str] = Field(default_factory=list, description="Surfaces where the command is visible") @@ -182,7 +181,6 @@ def _command_to_response(cmd: CommandInfo) -> CommandResponse: template=cmd.template, agent=cmd.agent, model=cmd.model, - subtask=cmd.subtask, hidden=cmd.hidden, aliases=list(cmd.aliases), visible_surfaces=list(cmd.visible_surfaces), diff --git a/flocks/session/__init__.py b/flocks/session/__init__.py index 54840e2fb..086208d66 100644 --- a/flocks/session/__init__.py +++ b/flocks/session/__init__.py @@ -30,18 +30,10 @@ ReasoningPart, PatchPart, AgentPart, - SubtaskPart, ) 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, @@ -53,11 +45,6 @@ ReminderConfig, ReminderContext, ) -from flocks.session.features.subtask import ( - SessionSubtask, - SubtaskInfo, - SubtaskResult, -) from flocks.session.lifecycle.revert import ( SessionRevertManager, RevertInput, @@ -93,7 +80,6 @@ "ReasoningPart", "PatchPart", "AgentPart", - "SubtaskPart", # Prompt "SessionPrompt", "SystemPrompt", @@ -104,12 +90,6 @@ # Summary "SessionSummary", "FileDiff", - # Runner - "SessionRunner", - "RunnerCallbacks", - "ToolCall", - "StepResult", - "run_session", # Session Loop "SessionLoop", "LoopContext", @@ -119,10 +99,6 @@ "SessionReminders", "ReminderConfig", "ReminderContext", - # Subtask - "SessionSubtask", - "SubtaskInfo", - "SubtaskResult", # Revert "SessionRevertManager", "RevertInput", 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/context_usage.py b/flocks/session/context_usage.py index ca35e09d1..422249a77 100644 --- a/flocks/session/context_usage.py +++ b/flocks/session/context_usage.py @@ -15,7 +15,7 @@ from flocks.provider.provider import Provider from flocks.session.message import Message -from flocks.session.prompt import SessionPrompt +from flocks.session.prompt import SessionPrompt, TurnPromptContext from flocks.session.session import SessionInfo from flocks.utils.log import Log @@ -23,7 +23,7 @@ log = Log.create(service="context-usage") UsageSource = Literal["observed", "estimated"] -DELEGATION_TOOLS = {"delegate_task", "task"} +DELEGATION_TOOLS = {"delegate_task"} ZERO_VISIBLE_SEGMENTS = {"agentDelegation"} @@ -309,7 +309,16 @@ async def _estimate_system_prompt_tokens( if agent is None: agent = await Agent.get("rex") - prompts = await SessionPrompt.build_system_prompts( + from flocks.config import Config + from flocks.project.instance import Instance + + try: + config = await Config.get() + config_instructions = tuple(config.instructions or ()) + except Exception: + config_instructions = () + + prompt_blocks = await SessionPrompt.build_system_prompt_blocks( session_id=session_id, session_directory=getattr(session, "directory", None) if session is not None else None, agent_name=getattr(agent, "name", agent_name) if agent is not None else agent_name, @@ -317,9 +326,16 @@ async def _estimate_system_prompt_tokens( provider_id=provider_id, model_id=model_id, prompt_tool_names=prompt_tool_names, - tool_revision=ToolRegistry.revision(), + turn_context=TurnPromptContext( + worktree=Instance.get_worktree(), + config_instructions=config_instructions, + tool_revision=ToolRegistry.revision(), + ), + ) + return sum( + SessionPrompt.count_tokens(block.content) + for block in prompt_blocks ) - return sum(SessionPrompt.count_tokens(prompt) for prompt in prompts) except Exception as exc: log.debug("context_usage.system_prompt_estimate_failed", { "session_id": session_id, @@ -441,8 +457,8 @@ async def _estimate_message_breakdown(session_id: str, messages: List[Any]) -> t if part_type in {"reasoning", "thinking"}: tokens_by_key["reasoning"] += SessionPrompt.count_tokens(_field_value(part, "text", "") or "") continue - if part_type in {"agent", "subtask"}: - tokens_by_key["agentDelegation"] += _estimate_subtask_part_tokens(part) + if part_type == "agent": + tokens_by_key["agentDelegation"] += _estimate_agent_part_tokens(part) continue if part_type != "tool": continue @@ -499,7 +515,7 @@ def _context_key_for_tool(tool_name: str) -> str: return "tools" -def _estimate_subtask_part_tokens(part: Any) -> int: +def _estimate_agent_part_tokens(part: Any) -> int: total = 0 for field in ("prompt", "description", "name"): value = _field_value(part, field, "") diff --git a/flocks/session/execution_mode.py b/flocks/session/execution_mode.py index 13d12b154..00952ff5c 100644 --- a/flocks/session/execution_mode.py +++ b/flocks/session/execution_mode.py @@ -28,7 +28,7 @@ class SessionExecutionMode(str, Enum): "run_slash_command", } ) -PLAN_DELEGATION_TOOL_NAMES = frozenset({"delegate_task", "task"}) +PLAN_DELEGATION_TOOL_NAMES = frozenset({"delegate_task"}) PLAN_DELEGATABLE_AGENT_NAMES = frozenset({"explore", "librarian"}) PLAN_PATH_SCOPED_TOOL_NAMES = frozenset({"apply_patch", "edit", "write"}) diff --git a/flocks/session/features/activity_forwarder.py b/flocks/session/features/activity_forwarder.py index a4eeff11e..5898e229b 100644 --- a/flocks/session/features/activity_forwarder.py +++ b/flocks/session/features/activity_forwarder.py @@ -69,15 +69,12 @@ def build_callbacks(self, event_publish_callback=None): server layer. Avoids session → server reverse dependency. """ from flocks.session.session_loop import LoopCallbacks - from flocks.session.runner import RunnerCallbacks return LoopCallbacks( event_publish_callback=event_publish_callback, - runner_callbacks=RunnerCallbacks( - on_tool_start=self._on_tool_start, - on_tool_end=self._on_tool_end, - on_text_delta=self._on_text_delta, - ), + on_tool_start=self._on_tool_start, + on_tool_end=self._on_tool_end, + on_text_delta=self._on_text_delta, ) # ------------------------------------------------------------------ diff --git a/flocks/session/features/subtask.py b/flocks/session/features/subtask.py deleted file mode 100644 index f1231cea5..000000000 --- a/flocks/session/features/subtask.py +++ /dev/null @@ -1,62 +0,0 @@ -""" -Session Subtask data models. - -Note: The SessionSubtask business logic has been removed as it was dead code. -The live subtask execution path is session_loop.py::_execute_subtask(), which -handles the full lifecycle inline without using this module. - -These data classes are kept because they are exported from session/__init__.py -and may be referenced by external consumers. -""" - -from dataclasses import dataclass, field -from datetime import datetime -from typing import Any, Dict, Optional - - -@dataclass -class SubtaskInfo: - """Information about a subtask""" - id: str - parent_session_id: str - child_session_id: Optional[str] = None - task_description: str = "" - agent: Optional[str] = None - model: Optional[str] = None - status: str = "pending" # pending, running, completed, error - result: Optional[str] = None - error: Optional[str] = None - created_at: int = field(default_factory=lambda: int(datetime.now().timestamp() * 1000)) - completed_at: Optional[int] = None - - -@dataclass -class SubtaskResult: - """Result of subtask execution""" - subtask_id: str - success: bool - output: str - error: Optional[str] = None - metadata: Dict[str, Any] = field(default_factory=dict) - - -# Minimal stub so imports of SessionSubtask don't break existing code. -class SessionSubtask: - """Subtask manager stub — business logic removed (was dead code). - - The active execution path is SessionLoop._execute_subtask() in session_loop.py. - """ - - @classmethod - async def execute_subtask(cls, *args, **kwargs) -> SubtaskResult: - raise NotImplementedError( - "SessionSubtask.execute_subtask() is deprecated. " - "Subtask execution is handled by SessionLoop._execute_subtask()." - ) - - -__all__ = [ - "SessionSubtask", - "SubtaskInfo", - "SubtaskResult", -] diff --git a/flocks/session/message.py b/flocks/session/message.py index a8c4b3e20..f3ac17bef 100644 --- a/flocks/session/message.py +++ b/flocks/session/message.py @@ -259,21 +259,6 @@ class AgentPart(BaseModel): source: Optional[Dict[str, Any]] = Field(None, description="Source information") -class SubtaskPart(BaseModel): - """Subtask/subagent part - Flocks compatible""" - model_config = ConfigDict(populate_by_name=True, by_alias=True) - - id: str = Field(default_factory=lambda: Identifier.ascending("part")) - sessionID: str = Field(..., description="Session ID") - messageID: str = Field(..., description="Message ID") - type: Literal["subtask"] = "subtask" - prompt: str = Field(..., description="Task prompt") - description: str = Field(..., description="Task description") - agent: str = Field(..., description="Agent name") - model: Optional[Dict[str, str]] = Field(None, description="Model configuration") - command: Optional[str] = Field(None, description="Command to execute") - - class RetryPart(BaseModel): """Retry part - Flocks compatible""" model_config = ConfigDict(populate_by_name=True, by_alias=True) @@ -301,7 +286,6 @@ class CompactionPart(BaseModel): # Union type for all parts - matches Flocks MessageV2.Part PartType = Union[ TextPart, - SubtaskPart, ReasoningPart, FilePart, ToolPart, @@ -1122,6 +1106,18 @@ def _normalize_part_data( metadata = normalized.get("metadata") metadata_dict = metadata if isinstance(metadata, dict) else {} + if part_type == "subtask": + normalized["type"] = "text" + normalized["text"] = "" + normalized["ignored"] = True + normalized["metadata"] = { + **metadata_dict, + "legacyPartType": "subtask", + } + part_type = "text" + metadata = normalized["metadata"] + metadata_dict = metadata + if "content" in normalized and "text" not in normalized: normalized["text"] = normalized.get("content", "") @@ -1174,10 +1170,6 @@ def _normalize_part_data( normalized.setdefault("tokens", metadata_dict.get("tokens") or cls._default_token_usage()) elif part_type == "agent": normalized.setdefault("name", metadata_dict.get("name") or normalized.get("content") or "agent") - elif part_type == "subtask": - normalized.setdefault("prompt", metadata_dict.get("prompt") or normalized.get("content", "")) - normalized.setdefault("description", metadata_dict.get("description") or "") - normalized.setdefault("agent", metadata_dict.get("agent") or "agent") elif part_type == "retry": normalized.setdefault("attempt", metadata_dict.get("attempt") or 1) normalized.setdefault("error", metadata_dict.get("error") or {}) @@ -1255,7 +1247,6 @@ def deserialize_part( 'step-start': StepStartPart, 'step-finish': StepFinishPart, 'agent': AgentPart, - 'subtask': SubtaskPart, 'retry': RetryPart, 'compaction': CompactionPart, } diff --git a/flocks/session/prompt.py b/flocks/session/prompt.py index 75efef655..f0c827a3a 100644 --- a/flocks/session/prompt.py +++ b/flocks/session/prompt.py @@ -35,8 +35,7 @@ # Output token maximum OUTPUT_TOKEN_MAX = int(os.getenv("FLOCKS_OUTPUT_TOKEN_MAX", "32000")) SystemPromptCache = Dict[str, Any] -AsyncPromptFactory = Callable[[], Awaitable[Optional[str]]] -StringPromptFactory = Callable[[], Optional[str]] +AsyncPromptLoader = Callable[[], Awaitable[Optional[str]]] # Prompt template directory (same structure as Flocks) @@ -139,13 +138,30 @@ class ContextInfo(BaseModel): @dataclass(frozen=True) class SystemPromptBlock: - """Internal system prompt layer with cache metadata.""" + """Assembled system prompt layer.""" name: str content: str cache_scope: str - digest_inputs: Dict[str, Any] - cache_key: str + + +@dataclass(frozen=True) +class TurnPromptContext: + """Runtime prompt values collected once before deterministic assembly.""" + + tool_catalog: Optional[str] = None + device_asset_hint: Optional[str] = None + sandbox_context: Optional[str] = None + channel_context: Optional[str] = None + additional_context: Optional[str] = None + text_tool_catalog: Optional[str] = None + tool_results_reminder: Optional[str] = None + repeated_tool_calls_reminder: Optional[str] = None + worktree: Optional[str] = None + config_instructions: tuple[str, ...] = () + tool_revision: Optional[int] = None + device_revision: Optional[int] = None + minimal_prompt: Optional[bool] = None class SystemPrompt: @@ -804,49 +820,6 @@ def _layer_cache_key( """Build a layer cache key for one prompt block.""" return f"system_prompt_block:{name}:{cls._system_prompt_cache_digest(digest_inputs)}" - @classmethod - def _system_prompt_cache_key( - cls, - *, - session_id: str, - agent_name: str, - provider_id: str, - model_id: str, - block_keys: Iterable[str], - ) -> str: - """Build the cache key for the composed system prompt snapshot.""" - cache_digest = cls._system_prompt_cache_digest({ - "block_keys": tuple(block_keys), - }) - return f"system_prompts:{session_id}:{agent_name}:{provider_id}:{model_id}:{cache_digest}" - - @classmethod - def _read_system_prompt_cache( - cls, - static_cache: Optional[SystemPromptCache], - cache_key: Optional[str], - ) -> Optional[List[str]]: - """Return a defensive copy of cached prompt blocks when available.""" - if static_cache is None or cache_key is None: - return None - - cached = static_cache.get(cache_key) - if cached is None: - return None - return list(cached) - - @classmethod - def _write_system_prompt_cache( - cls, - static_cache: Optional[SystemPromptCache], - cache_key: Optional[str], - prompts: List[str], - ) -> None: - """Store a defensive copy of prompt blocks in the session cache.""" - if static_cache is None or cache_key is None: - return - static_cache[cache_key] = list(prompts) - @classmethod def _read_cached_prompt_block( cls, @@ -909,8 +882,6 @@ def _build_cached_prompt_block( name=name, content=content, cache_scope=cache_scope, - digest_inputs=digest_inputs, - cache_key=cache_key, ) @classmethod @@ -921,22 +892,21 @@ async def _build_cached_async_prompt_block( name: str, cache_scope: str, digest_inputs: Dict[str, Any], - builder: AsyncPromptFactory, + loader: AsyncPromptLoader, ) -> Optional[SystemPromptBlock]: """Build or reuse a cached async prompt block.""" cache_key = cls._layer_cache_key(name=name, digest_inputs=digest_inputs) content = cls._read_cached_prompt_block(static_cache, cache_key) if content is None: - content = cls._normalize_prompt_text(await builder()) - cls._write_cached_prompt_block(static_cache, cache_key, content) + content = cls._normalize_prompt_text(await loader()) + if content: + cls._write_cached_prompt_block(static_cache, cache_key, content) if not content: return None return SystemPromptBlock( name=name, content=content, cache_scope=cache_scope, - digest_inputs=digest_inputs, - cache_key=cache_key, ) @classmethod @@ -1044,26 +1014,6 @@ def _prompt_blocks_to_list( if block is not None and block.content.strip() ] - @classmethod - async def _build_optional_async_prompt( - cls, - prompt_factory: Optional[AsyncPromptFactory], - ) -> Optional[str]: - """Run an optional async prompt factory.""" - if not prompt_factory: - return None - return await prompt_factory() - - @classmethod - def _build_optional_prompt( - cls, - prompt_factory: Optional[StringPromptFactory], - ) -> Optional[str]: - """Run an optional synchronous prompt factory.""" - if not prompt_factory: - return None - return prompt_factory() - @classmethod def _print_system_prompts_for_debug( cls, @@ -1072,7 +1022,7 @@ def _print_system_prompts_for_debug( agent_name: str, provider_id: str, model_id: str, - prompts: List[str], + blocks: Iterable[SystemPromptBlock], ) -> None: """Print prompt blocks when FLOCKS_PRINT_SYSTEM_PROMPT is enabled.""" if os.getenv("FLOCKS_PRINT_SYSTEM_PROMPT", "").lower() not in ("1", "true", "yes"): @@ -1083,8 +1033,12 @@ def _print_system_prompts_for_debug( f"agent={agent_name} model={provider_id}/{model_id} ===" ) print(header, file=sys.stderr) - for idx, prompt in enumerate(prompts): - print(f"\n--- prompt[{idx}] ---\n{prompt}\n", file=sys.stderr) + for idx, block in enumerate(blocks): + print( + f"\n--- prompt[{idx}] {block.name} scope={block.cache_scope} " + f"---\n{block.content}\n", + file=sys.stderr, + ) print("=== end system_prompt ===\n", file=sys.stderr) @classmethod @@ -1141,22 +1095,91 @@ async def _is_builtin_system_subagent_session( return False @classmethod - async def _build_subagent_minimal_prompts( + def _append_turn_tail_blocks( cls, *, + blocks: List[SystemPromptBlock], + turn_context: TurnPromptContext, + static_cache: Optional[SystemPromptCache], + session_id: str, + ) -> None: + """Append per-turn values in the exact order sent to the model.""" + tail_values = [ + ("turn_additional_context", turn_context.additional_context), + ("text_tool_catalog", turn_context.text_tool_catalog), + ("tool_results_reminder", turn_context.tool_results_reminder), + ( + "repeated_tool_calls_reminder", + turn_context.repeated_tool_calls_reminder, + ), + ] + for name, content in tail_values: + block = cls._build_cached_prompt_block( + static_cache=static_cache, + name=name, + cache_scope="runtime_tail", + digest_inputs={"session_id": session_id, "content": content or ""}, + builder=lambda value=content: cls._normalize_prompt_text(value), + ) + if block is not None: + blocks.append(block) + + @classmethod + def _build_subagent_minimal_blocks( + cls, + *, + session_id: str, session_directory: Optional[str], agent_prompt: Optional[str], - ) -> List[str]: - """Build minimal system prompts for built-in system subagents.""" - prompts = [ - get_prompt_flocks_config_guard().strip(), - cls._normalize_prompt_text(agent_prompt), - cls._build_minimal_environment(session_directory), - ] - return [prompt for prompt in prompts if prompt] + turn_context: TurnPromptContext, + static_cache: Optional[SystemPromptCache], + ) -> List[SystemPromptBlock]: + """Build minimal prompt blocks for built-in system subagents.""" + blocks: List[SystemPromptBlock] = [] + guard_block = cls._build_cached_prompt_block( + static_cache=static_cache, + name="flocks_config_guard", + cache_scope="global", + digest_inputs={"prompt": get_prompt_flocks_config_guard()}, + builder=lambda: get_prompt_flocks_config_guard().strip(), + ) + if guard_block is not None: + blocks.append(guard_block) + + agent_block = cls._build_cached_prompt_block( + static_cache=static_cache, + name="agent_identity", + cache_scope="agent", + digest_inputs={"agent_prompt": agent_prompt or ""}, + builder=lambda: cls._normalize_prompt_text(agent_prompt), + ) + if agent_block is not None: + blocks.append(agent_block) + + environment_block = cls._build_cached_prompt_block( + static_cache=static_cache, + name="minimal_environment", + cache_scope="runtime_tail", + digest_inputs={ + "directory": session_directory, + "runtime_day": datetime.now().strftime("%Y-%m-%d"), + "platform": platform.system().lower(), + }, + builder=lambda: cls._build_minimal_environment(session_directory), + ) + if environment_block is not None: + blocks.append(environment_block) + + cls._append_turn_tail_blocks( + blocks=blocks, + turn_context=turn_context, + static_cache=static_cache, + session_id=session_id, + ) + return blocks @classmethod - async def build_system_prompts( + async def build_system_prompt_blocks( cls, *, session_id: str, @@ -1167,45 +1190,51 @@ async def build_system_prompts( model_id: str, execution_mode_prompt: Optional[str] = None, prompt_tool_names: Iterable[str] = (), - tool_revision: Optional[int] = None, memory_bootstrap_data: Optional[Dict[str, Any]] = None, static_cache: Optional[SystemPromptCache] = None, - sandbox_prompt_factory: Optional[AsyncPromptFactory] = None, - channel_context_prompt_factory: Optional[AsyncPromptFactory] = None, - tool_catalog_prompt_factory: Optional[StringPromptFactory] = None, - device_asset_prompt_factory: Optional[AsyncPromptFactory] = None, - device_revision: Optional[int] = None, + turn_context: Optional[TurnPromptContext] = None, use_text_tool_call_mode: bool = False, - ) -> List[str]: - """Build the runtime system prompt blocks for a session turn. + ) -> List[SystemPromptBlock]: + """Build the ordered system prompt blocks for a session turn. Stable identity and execution guidance come first, followed by session/workspace context, with runtime-only metadata kept at the - prompt tail. Cache mechanics are intentionally kept out of the block - construction below so this method reads as an ordered list of prompt - layers. + prompt tail. Runtime I/O is collected before this method so assembly is + deterministic and every downstream consumer sees the same blocks. """ + turn_context = turn_context or TurnPromptContext() vcs = "git" if session_directory else None - if await cls._is_builtin_system_subagent_session( - session_id=session_id, - agent_name=agent_name, - ): - prompts = await cls._build_subagent_minimal_prompts( + minimal_prompt = turn_context.minimal_prompt + if minimal_prompt is None: + minimal_prompt = await cls._is_builtin_system_subagent_session( + session_id=session_id, + agent_name=agent_name, + ) + if minimal_prompt: + blocks = cls._build_subagent_minimal_blocks( + session_id=session_id, session_directory=session_directory, agent_prompt=agent_prompt, + turn_context=turn_context, + static_cache=static_cache, ) cls._print_system_prompts_for_debug( session_id=session_id, agent_name=agent_name, provider_id=provider_id, model_id=model_id, - prompts=prompts, + blocks=blocks, ) - return prompts + return blocks normalized_tool_names = tuple(sorted(prompt_tool_names)) runtime_day = datetime.now().strftime("%Y-%m-%d") - custom_signature = SystemPrompt.custom_signature(directory=session_directory) + config_instructions = list(turn_context.config_instructions) + custom_signature = SystemPrompt.custom_signature( + directory=session_directory, + worktree=turn_context.worktree, + config_instructions=config_instructions, + ) memory_guidance = cls._build_memory_guidance_prompt( normalized_tool_names, memory_bootstrap_data, @@ -1217,7 +1246,11 @@ async def build_system_prompts( async def build_custom_context() -> Optional[str]: return cls._join_prompt_parts( - await SystemPrompt.custom(directory=session_directory), + await SystemPrompt.custom( + directory=session_directory, + worktree=turn_context.worktree, + config_instructions=config_instructions, + ), ) blocks: List[Optional[SystemPromptBlock]] = [ @@ -1281,23 +1314,32 @@ async def build_custom_context() -> Optional[str]: cache_scope="catalog", digest_inputs={ "agent_name": agent_name, - "tool_revision": tool_revision, + "tool_revision": turn_context.tool_revision, + "content": turn_context.tool_catalog or "", }, - builder=lambda: cls._build_optional_prompt(tool_catalog_prompt_factory) or "", + builder=lambda: cls._normalize_prompt_text( + turn_context.tool_catalog, + ), ), ] - if device_asset_prompt_factory: - blocks.append(await cls._build_cached_async_prompt_block( - static_cache=static_cache, - name="device_asset_hint", - cache_scope="runtime", - digest_inputs={ - "session_id": session_id, - "device_revision": device_revision, - }, - builder=device_asset_prompt_factory, - )) + if turn_context.device_asset_hint: + blocks.append( + cls._build_cached_prompt_block( + static_cache=static_cache, + name="device_asset_hint", + cache_scope="runtime", + digest_inputs={ + "session_id": session_id, + "device_revision": turn_context.device_revision, + "tool_revision": turn_context.tool_revision, + "content": turn_context.device_asset_hint, + }, + builder=lambda: cls._normalize_prompt_text( + turn_context.device_asset_hint, + ), + ) + ) blocks.append( cls._build_cached_prompt_block( @@ -1319,33 +1361,52 @@ async def build_custom_context() -> Optional[str]: static_cache=static_cache, name="context_files", cache_scope="workspace", - digest_inputs={"directory": session_directory, "signature": custom_signature}, - builder=build_custom_context, + digest_inputs={ + "directory": session_directory, + "worktree": turn_context.worktree, + "signature": custom_signature, + }, + loader=build_custom_context, ) blocks.append(custom_block) - if sandbox_prompt_factory: - blocks.append(await cls._build_cached_async_prompt_block( - static_cache=static_cache, - name="sandbox_context", - cache_scope="runtime", - digest_inputs={"session_id": session_id, "agent_name": agent_name}, - builder=sandbox_prompt_factory, - )) + if turn_context.sandbox_context: + blocks.append( + cls._build_cached_prompt_block( + static_cache=static_cache, + name="sandbox_context", + cache_scope="runtime_tail", + digest_inputs={ + "session_id": session_id, + "agent_name": agent_name, + "content": turn_context.sandbox_context, + }, + builder=lambda: cls._normalize_prompt_text( + turn_context.sandbox_context, + ), + ) + ) - if channel_context_prompt_factory: - blocks.append(await cls._build_cached_async_prompt_block( - static_cache=static_cache, - name="channel_context", - cache_scope="runtime", - digest_inputs={"session_id": session_id}, - builder=channel_context_prompt_factory, - )) + if turn_context.channel_context: + blocks.append( + cls._build_cached_prompt_block( + static_cache=static_cache, + name="channel_context", + cache_scope="runtime_tail", + digest_inputs={ + "session_id": session_id, + "content": turn_context.channel_context, + }, + builder=lambda: cls._normalize_prompt_text( + turn_context.channel_context, + ), + ) + ) blocks.append(cls._build_cached_prompt_block( static_cache=static_cache, name="runtime_metadata", - cache_scope="runtime", + cache_scope="runtime_tail", digest_inputs={ "session_id": session_id, "directory": session_directory, @@ -1364,28 +1425,55 @@ async def build_custom_context() -> Optional[str]: ), )) - cache_key = cls._system_prompt_cache_key( + resolved_blocks = [block for block in blocks if block is not None] + cls._append_turn_tail_blocks( + blocks=resolved_blocks, + turn_context=turn_context, + static_cache=static_cache, + session_id=session_id, + ) + cls._print_system_prompts_for_debug( session_id=session_id, agent_name=agent_name, provider_id=provider_id, model_id=model_id, - block_keys=[block.cache_key for block in blocks if block is not None], + blocks=resolved_blocks, ) - cached_prompts = cls._read_system_prompt_cache(static_cache, cache_key) - if cached_prompts is not None: - return cached_prompts + return resolved_blocks - prompts = cls._prompt_blocks_to_list(blocks) - cls._print_system_prompts_for_debug( + @classmethod + async def build_system_prompts( + cls, + *, + session_id: str, + session_directory: Optional[str], + agent_name: str, + agent_prompt: Optional[str], + provider_id: str, + model_id: str, + execution_mode_prompt: Optional[str] = None, + prompt_tool_names: Iterable[str] = (), + memory_bootstrap_data: Optional[Dict[str, Any]] = None, + static_cache: Optional[SystemPromptCache] = None, + turn_context: Optional[TurnPromptContext] = None, + use_text_tool_call_mode: bool = False, + ) -> List[str]: + """Compatibility API returning only the assembled prompt text.""" + blocks = await cls.build_system_prompt_blocks( session_id=session_id, + session_directory=session_directory, agent_name=agent_name, + agent_prompt=agent_prompt, provider_id=provider_id, model_id=model_id, - prompts=prompts, + execution_mode_prompt=execution_mode_prompt, + prompt_tool_names=prompt_tool_names, + memory_bootstrap_data=memory_bootstrap_data, + static_cache=static_cache, + turn_context=turn_context, + use_text_tool_call_mode=use_text_tool_call_mode, ) - - cls._write_system_prompt_cache(static_cache, cache_key, prompts) - return list(prompts) + return cls._prompt_blocks_to_list(blocks) @classmethod def _build_context_section(cls, context: ContextInfo) -> str: diff --git a/flocks/session/prompt/anthropic-20250930.txt b/flocks/session/prompt/anthropic-20250930.txt index a8ada5ede..ec4e65c94 100644 --- a/flocks/session/prompt/anthropic-20250930.txt +++ b/flocks/session/prompt/anthropic-20250930.txt @@ -122,10 +122,10 @@ I've found existing rules. Let me mark the first todo as in_progress and start d Users may configure 'hooks', shell commands that execute in response to events like tool calls, in settings. Treat feedback from hooks, including , as coming from the user. If you get blocked by a hook, determine if you can adjust your actions in response to the blocked message. If not, ask the user to check their hooks configuration. # Tool usage policy -- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - Tool results and user messages may include tags. tags contain useful information and reminders. They are automatically added by the system, and bear no direct relation to the specific tool results or user messages in which they appear. - When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. diff --git a/flocks/session/prompt/anthropic.txt b/flocks/session/prompt/anthropic.txt index 655551871..2455ef584 100644 --- a/flocks/session/prompt/anthropic.txt +++ b/flocks/session/prompt/anthropic.txt @@ -66,22 +66,22 @@ I've found existing SIGMA rules. Let me mark the first todo as in_progress and s # Tool usage policy -- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - Tool results and user messages may include tags. tags contain useful information and reminders. They are automatically added by the system, and bear no direct relation to the specific tool results or user messages in which they appear. - When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. - You can call multiple tools in a single response. If you intend to call multiple tools and there are no dependencies between them, make all independent tool calls in parallel. Maximize use of parallel tool calls where possible to increase efficiency. However, if some tool calls depend on previous calls to inform dependent values, do NOT call these tools in parallel and instead call them sequentially. For instance, if one operation must complete before another starts, run these operations sequentially instead. Never use placeholders or guess missing parameters in tool calls. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. -- VERY IMPORTANT: When exploring security logs, configurations, or investigating incidents that require context gathering, use the Task tool for complex searches instead of running multiple search commands directly. +- VERY IMPORTANT: When exploring security logs, configurations, or investigating incidents that require context gathering, use `delegate_task` for complex searches instead of running multiple search commands directly. - IMPORTANT: Always respond in the same language as the user. user: Where are authentication failures logged in our application? -assistant: [Uses the Task tool to find authentication logging locations instead of using Glob or Grep directly] +assistant: [Uses `delegate_task` to find authentication logging locations instead of using Glob or Grep directly] user: Find all places where user input is processed without validation -assistant: [Uses the Task tool to comprehensively search for input validation gaps] +assistant: [Uses `delegate_task` to comprehensively search for input validation gaps] IMPORTANT: Always use `todo(action="write")` to plan and track tasks throughout the conversation. diff --git a/flocks/session/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..c8083a1a5 --- /dev/null +++ b/flocks/session/runtime/agent_loop.py @@ -0,0 +1,151 @@ +"""The control loop for one logical user input.""" + +from __future__ import annotations + +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_result = await engine.run(snapshot) + except StepCancelled: + log.info( + "session.step.cancelled", + { + "session_id": turn.session.id, + "step": turn.step, + }, + ) + return AgentRunOutcome( + status=AgentRunStatus.ABORTED, + 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..5620a7abe --- /dev/null +++ b/flocks/session/runtime/continuation_policy.py @@ -0,0 +1,440 @@ +"""Session-level logical turn preparation and continuation policy.""" + +from __future__ import annotations + +import inspect +from typing import Any, Optional + +from flocks.hooks.pipeline import HookPipeline +from flocks.session.runtime.contracts import ( + AgentRunOutcome, + ContinuationDecision, +) +from flocks.session.core.turn_state import set_turn_state +from flocks.session.runtime.event_sink import SessionEventSink +from flocks.session.goal import GoalManager +from flocks.session.message import Message, MessageInfo, MessageRole +from flocks.session.runtime.model_policy import ( + DEFAULT_MODEL_ROUTING_POLICY, + ModelRoutingPolicy, +) +from flocks.utils.log import Log + + +log = Log.create(service="session.continuation_policy") + + +class ContinuationPolicy: + """Own boundaries between durable logical user turns.""" + + def __init__(self, model_policy: ModelRoutingPolicy) -> None: + self._model_policy = model_policy + + async def publish_turn_stopped( + self, + turn: Any, + *, + stop_reason: str, + ) -> None: + """Publish the terminal state of one logical turn.""" + await SessionEventSink.turn_stopped( + turn.callbacks, + turn.session.id, + step=turn.step, + stop_reason=stop_reason, + ) + + @staticmethod + async def detect_queued_user_message( + _session_id: str, + post_messages: list[MessageInfo], + current_user_id: str, + _last_message: Optional[MessageInfo], + ) -> Optional[MessageInfo]: + """Return the newest user message after the current logical input.""" + newest_user = next( + (message for message in reversed(post_messages) if message.role == MessageRole.USER), + None, + ) + if newest_user is None or newest_user.id <= current_user_id: + return None + return newest_user + + async def prepare_logical_turn(self, context: Any) -> None: + """Prepare model routing and UserPromptSubmit once per logical input.""" + if context.session_store: + messages = await context.session_store.get_messages() + else: + messages = await Message.list(context.session.id) + context.prepared_messages = list(messages) + last_user = next( + (message for message in reversed(messages) if message.role == MessageRole.USER), + None, + ) + if last_user is None or last_user.id == context.prepared_user_id: + return + + is_real_user_turn = await self._model_policy.prepare_turn( + context, + last_user, + ) + if is_real_user_turn: + context.turn_additional_context = None + 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)}, + ) + return ContinuationDecision() + + if selected.should_continue: + await self._publish_continuation( + context, + selected, + event_metadata=event_metadata, + ) + return selected + + @staticmethod + async def _publish_continuation( + context: Any, + decision: ContinuationDecision[MessageInfo], + *, + event_metadata: Optional[dict[str, Any]] = None, + ) -> None: + """Publish the one continuation selected by the lifecycle boundary.""" + reason = decision.reason + message = decision.messages[0] + queued = reason == "queued_message" + turn_state = set_turn_state( + context.session.id, + step=context.step, + status="continued", + continue_reason=reason, + queued_message_detected=queued, + ) + message_id_key = { + "queued_message": "queuedUserMessageID", + "goal": "goalMessageID", + }.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..51589a9f9 --- /dev/null +++ b/flocks/session/runtime/contracts.py @@ -0,0 +1,240 @@ +"""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 + + +MessageT = TypeVar("MessageT") +ProviderMessageT = TypeVar("ProviderMessageT") + + +def _freeze(value: Any) -> Any: + """Recursively freeze request mappings and sequences.""" + if isinstance(value, Mapping): + return MappingProxyType( + {key: _freeze(item) for key, item in value.items()}, + ) + if isinstance(value, (list, tuple)): + return tuple(_freeze(item) for item in value) + if isinstance(value, set): + return frozenset(_freeze(item) for item in value) + if isinstance(value, (str, bytes, int, float, bool, type(None))): + return value + if hasattr(value, "model_copy"): + return value.model_copy(deep=True) + return copy.copy(value) + + +def _thaw(value: Any) -> Any: + """Return a provider-owned mutable copy of a frozen request value.""" + if isinstance(value, Mapping): + return {key: _thaw(item) for key, item in value.items()} + if isinstance(value, (tuple, list)): + return [_thaw(item) for item in value] + if isinstance(value, frozenset): + return {_thaw(item) for item in value} + if isinstance(value, (str, bytes, int, float, bool, type(None))): + return value + if hasattr(value, "model_copy"): + return value.model_copy(deep=True) + return copy.copy(value) + + +def _clone_provider_message(value: ProviderMessageT) -> ProviderMessageT: + if isinstance(value, (Mapping, list, tuple)): + return _thaw(_freeze(value)) + if hasattr(value, "model_copy"): + return value.model_copy(deep=True) + return copy.copy(value) + + +@dataclass(frozen=True) +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(_clone_provider_message(message) for message in self.messages), + ) + object.__setattr__( + self, + "tools", + tuple(_freeze(tool) for tool in self.tools), + ) + object.__setattr__(self, "options", _freeze(self.options)) + object.__setattr__(self, "metadata", _freeze(self.metadata)) + + def provider_messages(self) -> list[ProviderMessageT]: + """Return an isolated mutable copy for one provider invocation.""" + return [_clone_provider_message(message) for message in self.messages] + + def provider_tools(self) -> list[dict[str, Any]]: + """Return an isolated mutable tool-schema payload.""" + return [_thaw(tool) for tool in self.tools] + + def provider_options(self) -> dict[str, Any]: + """Return isolated provider options for one invocation.""" + return _thaw(self.options) + + +@dataclass +class AttemptEffects: + """Observable effects accumulated during one provider attempt.""" + + 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..3e782222d --- /dev/null +++ b/flocks/session/runtime/session_turn.py @@ -0,0 +1,1048 @@ +"""State and persistence boundary for one logical session turn. + +Implements model-turn preparation with support for: +- Message processing +- Tool execution +- Compaction +- Reminders +""" + +import asyncio +import time +from typing import Optional, List, Dict, Any, Callable, Awaitable, Literal +from dataclasses import dataclass, field +from datetime import datetime + +from flocks.session.runtime.contracts import ( + 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 + + +@dataclass +class LoopCallbacks: + """Callbacks for loop events""" + + on_step_start: Optional[Callable[[int], Awaitable[None]]] = None + on_step_end: Optional[Callable[[int], Awaitable[None]]] = None + on_text_delta: Optional[Callable[[str], Awaitable[None]]] = None + on_reasoning_delta: Optional[Callable[[str], Awaitable[None]]] = None + on_tool_start: Optional[ + Callable[[str, Dict[str, Any]], Awaitable[None]] + ] = None + on_tool_end: Optional[Callable[[str, Any], Awaitable[None]]] = None + on_permission_request: Optional[ + Callable[[Any], Awaitable[bool]] + ] = None + on_compaction: Optional[Callable[[], Awaitable[None]]] = None + on_error: Optional[Callable[[str], Awaitable[None]]] = None + on_reminder: Optional[Callable[[str], Awaitable[None]]] = None + # SSE event publishing callback (for TUI/WebUI real-time updates) + event_publish_callback: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None + + +@dataclass +class LoopResult: + """Result of loop execution""" + + action: str # "stop", "continue", "compact", "error", "queued" + last_message: Optional[MessageInfo] = None + error: Optional[str] = None + provider_id: Optional[str] = None + model_id: Optional[str] = None + metadata: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class LoopContext: + """Own state and persistence boundaries for one session loop run. + + Supports: + - Logical user-turn and message iteration + - Compaction triggers + - Reminder injection + """ + + session: SessionInfo + provider_id: str + model_id: str + agent_name: str + callbacks: LoopCallbacks = field(default_factory=LoopCallbacks, repr=False) + step: int = 0 + abort_event: asyncio.Event = field(default_factory=asyncio.Event) + session_store: Optional[Any] = None + trace_step_offset: int = 0 + _current_step_task: Optional[asyncio.Task] = field(default=None, repr=False) + memory_bootstrap_data: Optional[Dict[str, Any]] = field(default=None, repr=False) + step_static_cache: Dict[str, Any] = field(default_factory=dict, repr=False) + overflow_compaction_attempts: int = 0 + tool_result_truncation_attempted: bool = False + last_compaction_step: Optional[int] = None + last_cleanup_step: Optional[int] = None + last_observed_prompt_tokens: int = 0 + auto_failover: bool = False + auto_failover_allowed: bool = False + model_candidates: List[RuntimeModel] = field(default_factory=list) + candidate_index: int = 0 + model_candidate_policy: Literal["fixed", "automatic", "configured"] = "automatic" + turn_user_id: Optional[str] = None + turn_additional_context: Optional[str] = None + 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 ( + not self.auto_failover + or last_user is None + or getattr(message, "parentID", None) == last_user.id + ) + ), + None, + ) + + queued_user = None + if last_user is not None: + policy = self.continuation_policy + if policy is None: + from flocks.session.runtime.continuation_policy import ( + DEFAULT_CONTINUATION_POLICY, + ) + + policy = DEFAULT_CONTINUATION_POLICY + queued_user = await policy.detect_queued_user_message( + self.session.id, + post_messages, + last_user.id, + last_message, + ) + + if queued_user is not None: + turn_state = set_turn_state( + self.session.id, + step=self.step, + status="continued", + continue_reason="queued_message", + queued_message_detected=True, + ) + await SessionEventSink.emit( + self.callbacks, + "turn.continued", + { + **turn_state.model_dump(by_alias=True), + "queuedUserMessageID": queued_user.id, + }, + ) + log.info( + "session.turn.queued_input", + { + "session_id": self.session.id, + "queued_user_id": queued_user.id, + "last_assistant_id": ( + last_message.id if last_message else None + ), + }, + ) + elif step_result.action == StepAction.CONTINUE: + turn_state = set_turn_state( + self.session.id, + step=self.step, + status="continued", + continue_reason="tool_calls", + queued_message_detected=False, + ) + await SessionEventSink.emit( + self.callbacks, + "turn.continued", + turn_state.model_dump(by_alias=True), + ) + elif step_result.error: + await SessionEventSink.turn_stopped( + self.callbacks, + self.session.id, + step=self.step, + stop_reason=step_result.error, + ) + + return ModelTurnBoundary( + 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) + reported_total = input_tokens + cache_read + output_tokens + if reported_total > 0: + self.last_observed_prompt_tokens = reported_total + log.info( + "loop.tokens_decision", + { + "session_id": self.session.id, + "source": "observed", + "effective_tokens": input_tokens + cache_read, + "overflow_threshold": policy.overflow_threshold, + }, + ) + else: + estimated_tokens = await SessionPrompt.estimate_full_context_tokens( + self.session.id, + messages, + policy=policy, + ) + tokens = { + "input": estimated_tokens, + "output": 0, + "cache": {"read": 0, "write": 0}, + } + log.info( + "loop.tokens_decision", + { + "session_id": self.session.id, + "source": "estimated", + "effective_tokens": estimated_tokens, + "message_count": len(messages), + "overflow_threshold": policy.overflow_threshold, + }, + ) + + try: + cache = tokens.get("cache") or {} + current_input_tokens = tokens.get("input", 0) + (cache.get("read", 0) if isinstance(cache, dict) else 0) + recent_compaction = self._has_recent_compaction_cooldown() + near_overflow = current_input_tokens >= policy.preemptive_threshold + if near_overflow and self.last_cleanup_step != self.step: + cleanup_result = await self._prepare_tool_result_cleanup( + model_context, + policy, + current_input_tokens, + recent_compaction, + ) + if cleanup_result is not None: + return cleanup_result + + is_overflow = await SessionCompaction.is_overflow( + tokens=tokens, + model_context=model_context, + policy=policy, + ) + if not is_overflow: + return None + + log.info( + "loop.context_overflow_detected", + { + "session_id": self.session.id, + "step": self.step, + "tokens": tokens, + "tier": policy.tier.value, + "overflow_compaction_attempts": (self.overflow_compaction_attempts), + }, + ) + if self.overflow_compaction_attempts >= MAX_OVERFLOW_COMPACTION_ATTEMPTS: + await self._report_compaction_exhausted( + tokens, + ) + return ModelTurnPreparation( + status=TurnPreparationStatus.COMPLETE, + ) + + if not self.tool_result_truncation_attempted: + self.tool_result_truncation_attempted = True + try: + truncation_count = await SessionCompaction.truncate_oversized_tool_outputs( + self.session.id, + context_window_tokens=model_context, + ) + if truncation_count > 0: + log.info( + "loop.oversized_tool_truncated", + { + "session_id": self.session.id, + "truncated": truncation_count, + }, + ) + estimated_tokens = await SessionPrompt.estimate_full_context_tokens( + self.session.id, + messages, + policy=policy, + ) + still_overflow = await SessionCompaction.is_overflow( + tokens={ + "input": estimated_tokens, + "output": 0, + "cache": {"read": 0, "write": 0}, + }, + model_context=model_context, + policy=policy, + ) + if not still_overflow: + log.info( + "loop.overflow_resolved_by_truncation", + {"session_id": self.session.id}, + ) + return ModelTurnPreparation( + status=TurnPreparationStatus.CONTINUE, + ) + except Exception as exc: + log.warn( + "loop.oversized_truncation_error", + {"session_id": self.session.id, "error": str(exc)}, + ) + + return await self._prepare_full_compaction( + messages, + last_user, + policy, + ) + except Exception as exc: + log.error( + "loop.compaction_overflow_check_error", + {"error": str(exc)}, + ) + return None + + @staticmethod + def _normalise_token_usage(message: MessageInfo) -> Dict[str, Any]: + """Normalise provider token usage into the legacy mapping shape.""" + raw_tokens = getattr(message, "tokens", None) + if not raw_tokens: + return {} + if isinstance(raw_tokens, dict): + return raw_tokens + if hasattr(raw_tokens, "model_dump"): + return raw_tokens.model_dump() + if hasattr(raw_tokens, "__dict__"): + return vars(raw_tokens) + return {} + + async def _prepare_tool_result_cleanup( + self, + model_context: int, + policy: CompactionPolicy, + current_input_tokens: int, + recent_compaction: bool, + ) -> Optional[ModelTurnPreparation[MessageInfo]]: + """Apply the cheap tool-result cleanup before full compaction.""" + try: + truncation_count = await SessionCompaction.truncate_oversized_tool_outputs( + self.session.id, + context_window_tokens=model_context, + ) + self.last_cleanup_step = self.step + if truncation_count <= 0: + return None + + set_context_state( + self.session.id, + tool_results_compacted=True, + last_compaction_step=self.last_compaction_step, + last_compaction_reason="pre_compact_cleanup", + ) + await SessionEventSink.emit( + self.callbacks, + "context.compacted", + { + "sessionID": self.session.id, + "step": self.step, + "reason": "pre_compact_cleanup", + "truncatedToolResults": truncation_count, + "cooldownActive": recent_compaction, + }, + ) + log.info( + "loop.pre_compact_cleanup_applied", + { + "session_id": self.session.id, + "step": self.step, + "truncated": truncation_count, + "preemptive_threshold": policy.preemptive_threshold, + "input_tokens": current_input_tokens, + "cooldown_active": recent_compaction, + }, + ) + turn_state = set_turn_state( + self.session.id, + step=self.step, + status="continued", + continue_reason="pre_compact_cleanup", + queued_message_detected=False, + ) + await SessionEventSink.emit( + self.callbacks, + "turn.continued", + turn_state.model_dump(by_alias=True), + ) + return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) + except Exception as exc: + log.warn( + "loop.pre_compact_cleanup_error", + {"session_id": self.session.id, "error": str(exc)}, + ) + return None + + async def _report_compaction_exhausted( + self, + tokens: Dict[str, Any], + ) -> None: + """Surface whether exhaustion came from context or provider health.""" + history = _get_compaction_history(self.session.id) + provider_error = history.summary_last_error + in_cooldown = history.summary_cooldown_until > 0 and history.summary_cooldown_until > time.monotonic() + cooldown_seconds = max( + 0, + round(history.summary_cooldown_until - time.monotonic()), + ) + if in_cooldown or provider_error: + notice = ( + "摘要模型暂时不可用，上下文压缩跳过了本轮压缩。" + + (f"冷却剩余约 {cooldown_seconds} 秒，" if in_cooldown else "") + + "建议稍后继续，或切换到其他模型重试。" + ) + error = ( + "Compaction skipped: summary provider unavailable " + f"({provider_error or 'cooldown active'})." + + (f" Cooldown expires in ~{cooldown_seconds}s." if in_cooldown else "") + + " Wait for the provider to recover or switch models." + ) + else: + notice = "当前任务上下文过重，已经多次 compact 仍接近上限。建议收敛工具输出、缩小搜索范围，或开启新会话。" + error = ( + "Context overflow: prompt too large for the model after " + f"{self.overflow_compaction_attempts} compaction attempts. " + "Try starting a new session or use a larger-context model." + ) + + await SessionEventSink.notice( + self.callbacks, + self.session.id, + level="warning", + message=notice, + details={ + "attempts": self.overflow_compaction_attempts, + "maxAttempts": MAX_OVERFLOW_COMPACTION_ATTEMPTS, + "tokens": tokens, + "providerError": provider_error or None, + "cooldownRemainingSeconds": (cooldown_seconds if in_cooldown else 0), + }, + ) + log.error( + "loop.overflow_compaction_exhausted", + { + "session_id": self.session.id, + "attempts": self.overflow_compaction_attempts, + "max": MAX_OVERFLOW_COMPACTION_ATTEMPTS, + "tokens": tokens, + "in_cooldown": in_cooldown, + "provider_error": provider_error or None, + }, + ) + if self.callbacks.on_error: + await self.callbacks.on_error(error) + + async def _prepare_full_compaction( + self, + messages: List[MessageInfo], + last_user: MessageInfo, + policy: CompactionPolicy, + ) -> ModelTurnPreparation[MessageInfo]: + """Run full compaction and request preparation to reload the session.""" + self.overflow_compaction_attempts += 1 + if self.overflow_compaction_attempts >= 2: + await SessionEventSink.notice( + self.callbacks, + self.session.id, + level="info", + message=("本轮上下文持续接近模型上限，系统将优先尝试压缩历史工具输出。"), + details={ + "attempt": self.overflow_compaction_attempts, + "threshold": policy.overflow_threshold, + "buffer": policy.overflow_buffer, + }, + ) + log.warn( + "loop.overflow_compaction_attempt", + { + "session_id": self.session.id, + "attempt": self.overflow_compaction_attempts, + "max": MAX_OVERFLOW_COMPACTION_ATTEMPTS, + }, + ) + if self.callbacks.on_compaction: + await self.callbacks.on_compaction() + await SessionCompaction.prune(self.session.id, policy=policy) + + publish = self.callbacks.event_publish_callback + progress_callback = None + if publish is not None: + + async def progress_callback(stage: str, data: dict) -> None: + await publish( + "session.compaction_progress", + { + "sessionID": self.session.id, + "stage": stage, + "data": data, + }, + ) + + result = await run_compaction( + self.session.id, + parent_message_id=last_user.id, + messages=messages, + provider_id=self.provider_id, + model_id=self.model_id, + auto=True, + event_publish_callback=publish, + status_after="busy", + policy=policy, + progress_callback=progress_callback, + ) + if result == "stop": + log.error( + "loop.compaction_failed", + {"session_id": self.session.id}, + ) + if self.callbacks.on_error: + await self.callbacks.on_error("Compaction failed") + return ModelTurnPreparation(status=TurnPreparationStatus.COMPLETE) + if result == "skipped": + log.info( + "loop.compaction_skipped", + {"session_id": self.session.id, "step": self.step}, + ) + else: + self.last_compaction_step = self.step + set_context_state( + self.session.id, + compaction_performed=True, + last_compaction_step=self.step, + last_compaction_reason="full_compaction", + ) + await SessionEventSink.emit( + self.callbacks, + "context.compacted", + { + "sessionID": self.session.id, + "step": self.step, + "reason": "full_compaction", + "attempt": self.overflow_compaction_attempts, + "cooldownUntilStep": (self.step + POST_COMPACTION_COOLDOWN_STEPS), + }, + ) + return ModelTurnPreparation(status=TurnPreparationStatus.CONTINUE) + + def _build_compaction_policy(self) -> CompactionPolicy: + """ + Construct a CompactionPolicy from the current model's info. + + Falls back to ``CompactionPolicy.default()`` when the model info + cannot be resolved (e.g. unknown provider or missing context_window). + """ + return build_compaction_policy(self.provider_id, self.model_id) + + @staticmethod + def _should_exit( + last_user: MessageInfo, + last_assistant: Optional[MessageInfo], + last_assistant_parts: Optional[List[Any]] = None, + ) -> bool: + """ + Check if loop should exit + + Ported from original exit logic: + - Exit if assistant has responded with finish != tool-calls + - Exit if assistant message is after user message + """ + if not last_assistant: + return False + + if any(getattr(part, "type", None) == "tool" for part in (last_assistant_parts or [])): + return False + + # Check finish reason + if last_assistant.finish: + if last_assistant.finish not in ("tool-calls", "unknown", "summary"): + # Assistant finished with stop/error/etc + if last_user.id < last_assistant.id: + # Assistant responded after user + return True + + return False + + async def _check_reminders( + self, + messages: List[MessageInfo], + ) -> None: + """ + Check and inject reminders (P1 feature) + + Reminders are system messages injected periodically to: + - Remind agent of task goals + - Prevent drift from original intent + - Nudge towards completion + """ + from flocks.session.features.reminders import SessionReminders, ReminderContext + + # Calculate elapsed time + if messages: + first_msg = messages[0] + if hasattr(first_msg, "time") and hasattr(first_msg.time, "created"): + first_time = first_msg.time.created + current_time = int(datetime.now().timestamp() * 1000) + elapsed_ms = current_time - first_time + else: + elapsed_ms = 0 + else: + elapsed_ms = 0 + + # Extract original task + original_task = await SessionReminders.extract_original_task(messages) + + # Create reminder context + reminder_ctx = ReminderContext( + session_id=self.session.id, + step_count=self.step, + message_count=len(messages), + elapsed_ms=elapsed_ms, + original_task=original_task, + ) + + # Check if reminder should be injected + if SessionReminders.should_remind(self.session.id, reminder_ctx): + # Create and inject reminder + reminder_msg = await SessionReminders.create_reminder( + self.session.id, + reminder_ctx, + ) + + if reminder_msg and self.callbacks.on_reminder: + await self.callbacks.on_reminder( + await Message.get_text_content(reminder_msg), + ) diff --git a/flocks/session/runner.py b/flocks/session/runtime/step_engine.py similarity index 84% rename from flocks/session/runner.py rename to flocks/session/runtime/step_engine.py index b07a35ed3..7c5b30a35 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 @@ -18,17 +9,34 @@ import time from collections.abc import Mapping from datetime import datetime -from typing import Optional, Dict, Any, List, Callable, Awaitable, Tuple -from dataclasses import dataclass, field +from typing import Optional, Dict, Any, List, Tuple +from dataclasses import replace import httpcore import httpx +from flocks.session.runtime.contracts import ( + 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 from flocks.session.message import Message, MessageInfo, MessageRole, TextPart -from flocks.session.prompt import SessionPrompt +from flocks.session.prompt import SessionPrompt, SystemPromptBlock, TurnPromptContext from flocks.session.core.status import SessionStatus, SessionStatusRetry, SessionStatusBusy from flocks.session.core.defaults import ( DEFAULT_MAX_TOOL_STEPS, @@ -43,9 +51,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 +98,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 +129,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 +230,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 +257,11 @@ 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.callbacks = callbacks or LoopCallbacks() self._abort = asyncio.Event() self._external_abort = abort_event # External abort event (e.g. from SessionLoop) self._step = 0 self._recent_tool_calls: List[tuple[str, str]] = [] # Track recent (tool_name, args_json) for doom loop - self.session_ctx = session_ctx # SessionContext interface for decoupled access self._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,247 @@ def __init__( self._turn_additional_context = turn_additional_context self._session_start_pending = session_start_pending self._session_start_fired = False - self._attempt_state = LlmAttemptState() + self._attempt_state = AttemptEffects() + self._hooked_model_requests: Dict[ + str, + Tuple[ModelRequest[ChatMessage], bool], + ] = {} + self._turn: Optional[Any] = None + self._model_policy: ModelRoutingPolicy = DEFAULT_MODEL_ROUTING_POLICY + self._step_agent: Optional[AgentInfo] = None + self._frozen_tool_request: Optional[ModelRequest[ChatMessage]] = None + + @classmethod + def from_turn( + 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_tool_request = None + while True: + active_model = RuntimeModel(turn.provider_id, turn.model_id) + result = await self._run_candidate( + replace(snapshot, active_model=active_model), + ) + 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 not await self._remove_failed_attempt(failure): + await turn.finalize_failure(failure, snapshot.last_user) + return result + + await self._switch_candidate(next_index, failure.reason) + + def _require_turn(self) -> Any: + if self._turn is None: + raise RuntimeError( + "StepEngine.run() requires StepEngine.from_turn()", + ) + return self._turn + + 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 + + task = asyncio.create_task( + self._process_step(list(snapshot.messages), snapshot.last_user), + ) + turn._current_step_task = task + started_at = asyncio.get_running_loop().time() + try: + result = await task + if self._session_start_fired: + turn.session_start_pending = False + return result + except asyncio.CancelledError as exc: + if turn.aborted: + raise StepCancelled from exc + raise + finally: + turn._current_step_task = None + log.debug( + "session.step.complete", + { + "session_id": turn.session.id, + "step": turn.step, + "duration_ms": int( + ( + asyncio.get_running_loop().time() + - started_at + ) + * 1000 + ), + }, + ) + + def _record_chain_exhaustion( + self, + failure: Any, + has_next: bool, + ) -> None: + turn = self._require_turn() + if not ( + turn.model_candidate_policy == "automatic" + and failure.allow_fallback + and failure.attempt_state.replay_safe + and not has_next + and turn.candidate_index > 0 + and failure.reason not in {"rate_limit", "billing"} + ): + return + + expires_at = time.monotonic() + CHAIN_EXHAUSTION_COOLDOWN_SECONDS + existing = self._model_policy.cooldowns.get(turn.session.id) + if existing and existing.expires_at > expires_at: + return + self._model_policy.cooldowns[turn.session.id] = AutoFailoverCooldown( + model=turn.model_candidates[turn.candidate_index], + primary=turn.model_candidates[0], + expires_at=expires_at, + reason="chain_exhausted", + ) + + async def _remove_failed_attempt(self, failure: Any) -> bool: + turn = self._require_turn() + message_id = failure.assistant_message_id + if not message_id: + return True + try: + deleted = await Message.delete(turn.session.id, message_id) + except Exception as exc: + deleted = False + log.error( + "session.model.fallback_cleanup_failed", + { + "session_id": turn.session.id, + "message_id": message_id, + "error": str(exc), + }, + ) + if not deleted: + return False + await SessionEventSink.emit( + turn.callbacks, + "message.removed", + { + "sessionID": turn.session.id, + "messageID": message_id, + }, + ) + return True + + async def _switch_candidate(self, next_index: int, reason: str) -> None: + turn = self._require_turn() + previous = turn.model_candidates[turn.candidate_index] + next_candidate = turn.model_candidates[next_index] + if turn.model_candidate_policy == "automatic": + if turn.candidate_index == 0 and reason in { + "rate_limit", + "billing", + }: + self._model_policy.cooldowns[turn.session.id] = ( + AutoFailoverCooldown( + model=next_candidate, + primary=turn.model_candidates[0], + expires_at=( + time.monotonic() + + RATE_LIMIT_COOLDOWN_SECONDS + ), + reason=reason, + ) + ) + else: + cooldown = self._model_policy.cooldowns.get(turn.session.id) + if cooldown and cooldown.expires_at > time.monotonic(): + cooldown.model = next_candidate + + self._model_policy.select_candidate(turn, next_index) + payload = { + "sessionID": turn.session.id, + "from": { + "providerID": previous.provider_id, + "modelID": previous.model_id, + }, + "to": { + "providerID": next_candidate.provider_id, + "modelID": next_candidate.model_id, + }, + "reason": reason, + "candidateIndex": next_index, + } + log.warn( + "session.model.fallback", + { + "from": payload["from"], + "to": payload["to"], + "reason": reason, + "candidateIndex": next_index, + }, + ) + await SessionEventSink.emit( + turn.callbacks, + "session.model.fallback", + payload, + ) @staticmethod def _canonical_tool_signature(tool_name: str, arguments: Dict[str, Any]) -> str: @@ -385,8 +553,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 +844,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 +1026,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,284 +1144,10 @@ def _append_file_content_block( placeholder = placeholder[:MAX_PLACEHOLDER_CHARS] + "…" text_fallbacks.append(placeholder) - @classmethod - async def loop(cls, session_id: str) -> Optional['MessageInfo']: - """ - Start or continue session processing loop. - - This is the main entry point for session execution, - matching Flocks' SessionPrompt.loop() behavior. - - Now delegates to SessionLoop for better separation of concerns. - - Args: - session_id: Session ID to process - - Returns: - Last assistant message with parts - """ - # Delegate to SessionLoop (new architecture) - from flocks.session.session_loop import SessionLoop - - result = await SessionLoop.run(session_id) - return result.last_message - - @classmethod - def cancel(cls, session_id: str) -> bool: - """ - Cancel a running session. - - Args: - session_id: Session ID to cancel - - Returns: - True if session was cancelled - """ - from flocks.session.core.status import SessionStatus - - runner = cls._active_sessions.get(session_id) - if runner: - runner.abort() - del cls._active_sessions[session_id] - log.info("runner.cancelled", {"session_id": session_id}) - - # Set status to idle (Flocks compatibility) - from flocks.session.core.status import SessionStatusIdle - SessionStatus.set(session_id, SessionStatusIdle()) - return True - - @classmethod - def cancel_children(cls, parent_session_id: str) -> int: - """Cancel all runners whose session.parent_id matches, recursively.""" - from flocks.session.core.status import SessionStatus, SessionStatusIdle - - cancelled = 0 - child_ids = [ - sid for sid, runner in list(cls._active_sessions.items()) - if getattr(runner.session, 'parent_id', None) == parent_session_id - ] - for sid in child_ids: - runner = cls._active_sessions.pop(sid, None) - if runner: - runner.abort() - SessionStatus.set(sid, SessionStatusIdle()) - cancelled += 1 - log.info("runner.child_cancelled", { - "session_id": sid, - "parent_session_id": parent_session_id, - }) - cancelled += cls.cancel_children(sid) - return cancelled - - @classmethod - async def command( - cls, - session_id: str, - command: str, - arguments: str = "", - message_id: Optional[str] = None, - agent: Optional[str] = None, - model: Optional[str] = None, - variant: Optional[str] = None, - ) -> Dict[str, Any]: - """ - Execute a slash command in a session. - - Args: - session_id: Session ID - command: Command name (e.g., "init", "help") - arguments: Command arguments - message_id: Optional message ID - agent: Optional agent name - model: Optional model string (provider/model) - variant: Optional model variant - - Returns: - Command execution result - """ - from flocks.command.command import Command - - # Get command definition - cmd = Command.get(command) - if not cmd: - raise ValueError(f"Command '{command}' not found") - - # Parse model if provided - provider_id, model_id = None, None - if model: - parts = model.split("/", 1) - if len(parts) == 2: - provider_id, model_id = parts - - # Execute command template - template = cmd.template - - # Replace placeholders - template = template.replace("$ARGUMENTS", arguments) - - # Create prompt request - parts = [{"type": "text", "text": template}] - - log.info("runner.command", { - "session_id": session_id, - "command": command, - "arguments": arguments[:50] if arguments else "", - }) - - return { - "command": command, - "arguments": arguments, - "template": template, - } - - @classmethod - async def shell( - cls, - session_id: str, - agent: str, - command: str, - model: Optional[Dict[str, str]] = None, - ) -> Dict[str, Any]: - """ - Execute a shell command in session context. - - Args: - session_id: Session ID - agent: Agent name - command: Shell command to execute - model: Optional model info - - Returns: - Shell execution result - """ - session = await Session.get_by_id(session_id) - if not session: - raise ValueError(f"Session {session_id} not found") - - cwd = session.directory or os.getcwd() - - 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).""" @@ -1361,7 +1254,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 +1272,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 +1290,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 +1305,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 +1323,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 +1339,6 @@ async def _process_step( error_dict=error_dict, visible_text=error_dict["data"]["displayMessage"], ) - if self.callbacks.on_error: - await self.callbacks.on_error(error_dict["data"]["displayMessage"]) return StepResult(action="stop", error=error_dict["data"]["displayMessage"]) # Apply config-based provider options (api_key/base_url) @@ -1471,7 +1349,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 +1361,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,34 +1377,31 @@ 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) + frozen_tool_request = self._frozen_tool_request + if isinstance(frozen_tool_request, ModelRequest): + tools = frozen_tool_request.provider_tools() + else: + tools = await self._build_callable_tool_schema(agent, messages) self._log_perf("runner.process_step.tools_ready", tools_started_at, tool_count=len(tools)) prompt_tool_names = self._get_prompt_tool_names_from_schema(tools) - async def sandbox_prompt_factory() -> Optional[str]: - return await self._build_sandbox_prompt(agent) - - async def channel_context_prompt_factory() -> Optional[str]: - return await self._build_channel_context_prompt() - - async def device_asset_prompt_factory() -> Optional[str]: - return await self._build_device_asset_hint() - - try: - from flocks.tool.device.store import device_revision as get_device_revision - - current_device_revision = get_device_revision() - except Exception: - current_device_revision = None - prompts_started_at = time.perf_counter() - system_prompts = await SessionPrompt.build_system_prompts( + minimal_prompt = await SessionPrompt._is_builtin_system_subagent_session( + session_id=self.session.id, + agent_name=agent.name, + ) + turn_prompt_context = await self._build_turn_prompt_context( + agent=agent, + messages=messages, + last_user=last_user, + tools=tools, + minimal_prompt=minimal_prompt, + ) + system_prompt_blocks = await SessionPrompt.build_system_prompt_blocks( session_id=self.session.id, session_directory=self.session.directory, agent_name=agent.name, @@ -1539,57 +1414,19 @@ async def device_asset_prompt_factory() -> Optional[str]: plan_file=self._turn_plan_file, ), prompt_tool_names=prompt_tool_names, - tool_revision=ToolRegistry.revision(), memory_bootstrap_data=self._memory_bootstrap_data, static_cache=self._static_cache, - sandbox_prompt_factory=sandbox_prompt_factory, - channel_context_prompt_factory=channel_context_prompt_factory, - tool_catalog_prompt_factory=lambda: self._build_tool_catalog_prompt(agent), - device_asset_prompt_factory=device_asset_prompt_factory, - device_revision=current_device_revision, + turn_context=turn_prompt_context, use_text_tool_call_mode=self._should_use_text_tool_call_mode(), ) - self._log_perf("runner.process_step.system_prompts_ready", prompts_started_at, prompt_count=len(system_prompts)) + self._log_perf( + "runner.process_step.system_prompts_ready", + prompts_started_at, + prompt_count=len(system_prompt_blocks), + ) await self._run_session_start_hook(agent) - if self._turn_additional_context: - system_prompts.append(self._turn_additional_context) - - if self._should_use_text_tool_call_mode() and tools: - text_tool_catalog = self._build_text_tool_call_catalog_prompt(tools) - if text_tool_catalog: - system_prompts.append(text_tool_catalog) - - # If the last assistant message only contains tool results and no text, - # force a direct answer to avoid repeated tool calls. - last_assistant_msg = None - for msg in reversed(messages): - if msg.role == MessageRole.ASSISTANT: - last_assistant_msg = msg - break - if last_assistant_msg: - parts = await Message.parts(last_assistant_msg.id, self.session.id) - has_text = any(getattr(p, "type", None) == "text" and getattr(p, "text", "").strip() for p in parts) - has_tool_result = any( - getattr(p, "type", None) == "tool" and - getattr(getattr(p, "state", None), "status", None) in ("completed", "error", "running") - for p in parts - ) - if has_tool_result and not has_text: - from flocks.session.prompt_strings import PROMPT_TOOL_RESULTS_AVAILABLE - system_prompts.append(PROMPT_TOOL_RESULTS_AVAILABLE) - - if has_tool_result and self._should_warn_about_tool_loop(last_user_id=last_user.id): - state = self._get_tool_loop_guard_state(last_user_id=last_user.id) - log.warn("runner.repeated_tool_calls_detected", { - "tool_name": state.get("last_signature", "").split(":", 1)[0], - "exact_count": state.get("exact_count", 0), - "step": self._step, - }) - from flocks.session.prompt_strings import PROMPT_REPEATED_TOOL_CALLS - system_prompts.append(PROMPT_REPEATED_TOOL_CALLS) - # Convert messages to chat format with error handling try: queued_user_message_ids = self._get_queued_user_message_ids(messages) @@ -1599,7 +1436,10 @@ async def device_asset_prompt_factory() -> Optional[str]: self._queued_user_message_ids = queued_user_message_ids chat_messages_started_at = time.perf_counter() try: - chat_messages = await self._to_chat_messages(messages, system_prompts) + chat_messages = await self._to_chat_messages( + messages, + system_prompt_blocks, + ) finally: if previous_queued_user_ids is None: if hasattr(self, "_queued_user_message_ids"): @@ -1619,7 +1459,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 +1481,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 +1489,26 @@ async def device_asset_prompt_factory() -> Optional[str]: role="assistant", content=PROMPT_MAX_STEPS, )) - + log.warn("runner.max_steps_reached", { "step": self._step, "max_steps": max_steps, "session_id": self.session.id, }) - + # Disable tools when max steps reached tools = [] - - # Create assistant message (will be reused across retries) + + request = self._build_model_request( + messages=chat_messages, + tools=tools, + agent=agent, + ) + if self._turn is not None and self._frozen_tool_request is None: + self._frozen_tool_request = request + self._active_model_request = request + + # Create the persisted assistant attempt after the request is frozen. assistant_msg = await Message.create( session_id=self.session.id, role=MessageRole.ASSISTANT, @@ -1669,7 +1518,6 @@ async def device_asset_prompt_factory() -> Optional[str]: provider_id=self.provider_id, parent_id=last_user.id, ) - # Publish assistant message SSE event so frontends can show the message card if self.callbacks.event_publish_callback: import time as _time @@ -1687,7 +1535,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 @@ -1799,11 +1647,9 @@ async def device_asset_prompt_factory() -> Optional[str]: message=empty_error_msg, error_data=empty_error_dict, assistant_message_id=assistant_msg.id, - decision=FailoverDecision(True, "empty_response", 3), + decision=FailoverDecision(True, "empty_response"), attempts=empty_attempt, ) - if self.callbacks.on_error: - await self.callbacks.on_error(empty_error_msg) await Message.update( self.session.id, assistant_msg.id, @@ -1902,7 +1748,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 +1758,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: @@ -1952,9 +1798,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,9 +1813,9 @@ async def device_asset_prompt_factory() -> Optional[str]: error_dict=error_dict, text_part=text_part, ) - + return StepResult(action="stop", error=final_error_message) - + # Aborted return StepResult(action="stop", error="Aborted") @@ -2099,7 +1942,134 @@ async def _record_usage_if_available( "model_id": self.model_id, "error": str(exc), }) - + + async def _build_turn_prompt_context( + self, + *, + agent: AgentInfo, + messages: List[MessageInfo], + last_user: MessageInfo, + tools: List[Dict[str, Any]], + minimal_prompt: bool = False, + ) -> TurnPromptContext: + """Collect cached runtime values before deterministic prompt assembly.""" + if minimal_prompt: + return await self._add_turn_prompt_tail( + TurnPromptContext(minimal_prompt=True), + messages=messages, + last_user=last_user, + tools=tools, + ) + + from flocks.config import Config + from flocks.project.instance import Instance + + try: + from flocks.tool.device.store import device_revision + + current_device_revision = device_revision() + except Exception: + current_device_revision = None + + current_tool_revision = ToolRegistry.revision() + try: + config = await Config.get() + config_data = config.model_dump(by_alias=True, exclude_none=True) + config_instructions = tuple(config.instructions or ()) + except Exception as exc: + log.debug("runner.prompt_context.config_error", {"error": str(exc)}) + config_data = None + config_instructions = () + + worktree = Instance.get_worktree() + sandbox_context, channel_context, device_asset_hint = await asyncio.gather( + self._build_sandbox_prompt(agent, config_data=config_data), + self._build_channel_context_prompt(), + self._build_device_asset_hint(), + ) + source_context = TurnPromptContext( + tool_catalog=self._build_tool_catalog_prompt(agent), + device_asset_hint=device_asset_hint, + sandbox_context=sandbox_context, + channel_context=channel_context, + worktree=worktree, + config_instructions=config_instructions, + tool_revision=current_tool_revision, + device_revision=current_device_revision, + minimal_prompt=False, + ) + + return await self._add_turn_prompt_tail( + source_context, + messages=messages, + last_user=last_user, + tools=tools, + ) + + async def _add_turn_prompt_tail( + self, + source_context: TurnPromptContext, + *, + messages: List[MessageInfo], + last_user: MessageInfo, + tools: List[Dict[str, Any]], + ) -> TurnPromptContext: + """Add uncached per-step context and reminders to a source snapshot.""" + text_tool_catalog = None + if self._should_use_text_tool_call_mode() and tools: + text_tool_catalog = self._build_text_tool_call_catalog_prompt(tools) + + tool_results_reminder = None + repeated_tool_calls_reminder = None + last_assistant_msg = next( + ( + message + for message in reversed(messages) + if message.role == MessageRole.ASSISTANT + ), + None, + ) + if last_assistant_msg is not None: + parts = await Message.parts(last_assistant_msg.id, self.session.id) + has_text = any( + getattr(part, "type", None) == "text" + and getattr(part, "text", "").strip() + for part in parts + ) + has_tool_result = any( + getattr(part, "type", None) == "tool" + and getattr(getattr(part, "state", None), "status", None) + in ("completed", "error", "running") + for part in parts + ) + if has_tool_result and not has_text: + from flocks.session.prompt_strings import ( + PROMPT_TOOL_RESULTS_AVAILABLE, + ) + + tool_results_reminder = PROMPT_TOOL_RESULTS_AVAILABLE + + if has_tool_result and self._should_warn_about_tool_loop( + last_user_id=last_user.id, + ): + state = self._get_tool_loop_guard_state(last_user_id=last_user.id) + log.warn("runner.repeated_tool_calls_detected", { + "tool_name": state.get("last_signature", "").split(":", 1)[0], + "exact_count": state.get("exact_count", 0), + "step": self._step, + }) + from flocks.session.prompt_strings import PROMPT_REPEATED_TOOL_CALLS + + repeated_tool_calls_reminder = PROMPT_REPEATED_TOOL_CALLS + + return replace( + source_context, + additional_context=self._turn_additional_context, + text_tool_catalog=text_tool_catalog, + tool_results_reminder=tool_results_reminder, + repeated_tool_calls_reminder=repeated_tool_calls_reminder, + ) + async def _build_device_asset_hint(self) -> Optional[str]: """Return concise device-aware tool guidance plus enabled device summary.""" try: @@ -2142,15 +2112,22 @@ async def _build_device_asset_hint(self) -> Optional[str]: "如果同类设备有多个候选，不要猜测，先询问用户选择。" ) - async def _build_sandbox_prompt(self, agent: AgentInfo) -> Optional[str]: + async def _build_sandbox_prompt( + self, + agent: AgentInfo, + *, + config_data: Optional[Dict[str, Any]] = None, + ) -> Optional[str]: """Build sandbox context prompt when sandboxing is active.""" try: - from flocks.config import Config from flocks.session.core.session_state import get_main_session_id from flocks.sandbox.system_prompt import build_sandbox_system_prompt - cfg = await Config.get() - config_data = cfg.model_dump(by_alias=True, exclude_none=True) + if config_data is None: + from flocks.config import Config + + config = await Config.get() + config_data = config.model_dump(by_alias=True, exclude_none=True) session_key = self.session.id main_session_key = get_main_session_id() or self.session.id return await build_sandbox_system_prompt( @@ -2334,7 +2311,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 +2373,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 +2381,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 +2410,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 +2449,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 +2462,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 +2477,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: @@ -2662,14 +2642,26 @@ def _build_tool_output_text(self, part: Any, tool_name: str, ctx_window_tokens: def _build_system_message_content( self, - system_prompts: List[str], + system_prompts: List[SystemPromptBlock] | List[str], ) -> str | list[dict[str, Any]]: """Format system prompts for the active provider. Anthropic supports structured system blocks, which lets us place a conservative cache breakpoint before the dynamic runtime tail. """ - prompt_parts = [prompt for prompt in system_prompts if prompt and prompt.strip()] + typed_blocks = [ + block + for block in system_prompts + if isinstance(block, SystemPromptBlock) and block.content.strip() + ] + if typed_blocks: + prompt_parts = [block.content for block in typed_blocks] + else: + prompt_parts = [ + prompt + for prompt in system_prompts + if isinstance(prompt, str) and prompt.strip() + ] if not prompt_parts: return "" @@ -2677,7 +2669,19 @@ def _build_system_message_content( if "anthropic" not in provider_lower: return "\n\n".join(prompt_parts) - cache_break_index = max(0, len(prompt_parts) - 3) + if typed_blocks: + first_runtime_tail = next( + ( + index + for index, block in enumerate(typed_blocks) + if block.cache_scope == "runtime_tail" + ), + len(typed_blocks), + ) + cache_break_index = max(0, first_runtime_tail - 1) + else: + # Compatibility for callers still passing plain strings. + cache_break_index = max(0, len(prompt_parts) - 3) blocks: list[dict[str, Any]] = [] for index, prompt in enumerate(prompt_parts): block: dict[str, Any] = { @@ -2692,11 +2696,11 @@ def _build_system_message_content( async def _to_chat_messages( self, messages: List[MessageInfo], - system_prompts: List[str], + system_prompts: List[SystemPromptBlock] | List[str], ) -> List[ChatMessage]: """ Convert messages to chat format with tool calls. - + Ported from original MessageV2.toModelMessage() logic: - Include text parts - Include tool calls and results @@ -2708,7 +2712,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) @@ -2775,7 +2782,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: @@ -2786,7 +2793,7 @@ async def _to_chat_messages( is_latest_user_turn = msg.id == last_user_msg_id # Get message parts parts = preloaded_parts[idx] - + if not parts: # Fallback: use text content only content = await Message.get_text_content(msg) @@ -2799,7 +2806,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 +2857,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 +2876,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 +2886,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 +2900,7 @@ async def _to_chat_messages( else: # Non-AbortedError - skip continue - + assistant_content_parts = [] assistant_reasoning_parts = [] assistant_reasoning_content_parts = [] @@ -2910,11 +2911,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 +2966,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 +2986,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 +3014,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 +3022,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 +3040,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 +3065,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 +3100,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", { @@ -3128,9 +3129,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 +3269,17 @@ async def _call_llm( ) -> StepResult: """ Call LLM and process response with event-driven streaming. - + Uses StreamProcessor to handle events and execute tools synchronously. Ported from Flocks' SessionProcessor.process() behavior. """ + request = getattr(self, "_active_model_request", None) + if not isinstance(request, ModelRequest): + request = self._build_model_request( + messages=messages, + tools=tools, + agent=agent, + ) def _build_llm_response_payload( *, content: str, @@ -3165,6 +3300,79 @@ def _build_llm_response_payload( ], } + llm_hook_metadata = { + "sessionID": self.session.id, + "messageID": assistant_msg.id, + "workspace": self.session.directory, + "agent": agent.name, + "step": self._step, + "model": { + "providerID": request.provider_id, + "modelID": request.model_id, + }, + } + cached_request = self._hooked_model_requests.get(assistant_msg.id) + if cached_request is None: + llm_before_enabled = False + llm_after_enabled = False + try: + llm_before_enabled = ( + await HookPipeline.has_stage_handlers( + HookStage.LLM_BEFORE, + llm_hook_metadata, + ) + ) + llm_after_enabled = ( + await HookPipeline.has_stage_handlers( + HookStage.LLM_AFTER, + llm_hook_metadata, + ) + ) + except Exception as exc: + log.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 + self._hooked_model_requests[assistant_msg.id] = ( + request, + llm_after_enabled, + ) + else: + request, llm_after_enabled = cached_request + self._active_model_request = request + messages = request.provider_messages() + tools = request.provider_tools() + 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 +3392,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) @@ -3196,7 +3411,7 @@ async def _on_tool_execution_start( text_delta_callback=self.callbacks.on_text_delta, reasoning_delta_callback=self.callbacks.on_reasoning_delta, tool_start_callback=_on_tool_execution_start, - tool_end_callback=self.callbacks.on_tool_end, + tool_end_callback=_on_tool_execution_end, event_publish_callback=self.callbacks.event_publish_callback, session_key=self.session.id, main_session_key=main_session_key, @@ -3214,11 +3429,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 +3472,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 +3494,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 +3534,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 +3701,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, @@ -3626,7 +3782,7 @@ async def _flush_reasoning_rewriter() -> None: except Exception as hook_exc: log.debug("runner.hook.llm_after.error", {"error": str(hook_exc)}) raise - + log.debug("runner.stream.summary", { "total_chunks": chunk_counts["total"], "reasoning_chunks": chunk_counts["reasoning"], @@ -3646,11 +3802,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 +3818,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 +3828,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 +3849,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 +3857,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 +3866,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 = [ @@ -3757,7 +3913,7 @@ async def _flush_reasoning_rewriter() -> None: ) except Exception as exc: log.debug("runner.hook.llm_after.error", {"error": str(exc)}) - + if tool_calls_for_result: response_payload = self._build_langfuse_response_payload( action="continue", @@ -3783,7 +3939,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 +3959,7 @@ async def _flush_reasoning_rewriter() -> None: trace_output=response_payload, ) return StepResult(action=result_action, content=content, usage=stream_usage) - + @staticmethod def _end_observability( generation_ctx: Any, @@ -3883,41 +4039,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 b61e0ba69..cf8e32438 100644 --- a/flocks/session/session_loop.py +++ b/flocks/session/session_loop.py @@ -1,576 +1,96 @@ -""" -Session Loop Module - -Core session execution loop logic extracted from runner.py. -Implements the main session processing loop with support for: -- Message processing -- Tool execution -- Compaction -- Subtask handling -- Reminders - -Ported from original SessionPrompt.loop() pattern. -""" - -import asyncio -import hashlib -import inspect -import time -from typing import Optional, List, Dict, Any, Callable, Awaitable, Literal -from dataclasses import dataclass, field -from datetime import datetime +"""Public entry point and lifecycle owner for session execution.""" -from flocks.utils.log import Log -from flocks.utils.id import Identifier +from __future__ import annotations + +from collections.abc import MutableMapping +from dataclasses import dataclass +from typing import Any, ClassVar, Optional + +from flocks.session.core.context import DefaultSessionContext +from flocks.session.core.status import ( + SessionStatus, + SessionStatusBusy, + SessionStatusIdle, +) +from flocks.session.core.turn_state import clear_turn_state +from flocks.session.message import Message +from flocks.session.runtime.agent_loop import AgentLoop +from flocks.session.runtime.continuation_policy import ( + DEFAULT_CONTINUATION_POLICY, + ContinuationPolicy, +) +from flocks.session.runtime.contracts import ( + AgentRunOutcome, + 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 - - @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] + def get(self, session_id: str) -> Optional[LoopContext]: + return self._active_turns.get(session_id) - @classmethod - async def validate_auto_configuration(cls) -> tuple[bool, str]: - """Validate that a newly selected Auto mode has a usable chain.""" - from flocks.config.config import Config - - 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: + turn: LoopContext, + ) -> Optional[_SessionLease]: + if session_id in self._active_turns: return None - if cooldown.expires_at <= time.monotonic() or cooldown.primary != primary: - cls._auto_failover_cooldowns.pop(session_id, None) - return None - return cooldown.model - - @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 + self._active_turns[session_id] = turn + return _SessionLease(session_id=session_id, turn=turn) - @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), - }) + 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 - 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), - ) + def owns(self, lease: _SessionLease) -> bool: + return self._active_turns.get(lease.session_id) is lease.turn - @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)}) - - @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 - ) - - @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 +class SessionLoop: + """Decide whether a persistent session should continue or settle.""" - newest_user = None - for msg in reversed(post_messages): - if msg.role == MessageRole.USER: - newest_user = msg - break + _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 + ) - 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 +102,223 @@ 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.""" + 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), - }) - + settled = False + processed_user_id: Optional[str] = None 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: + await cls._mark_busy(session_id, runtime_callbacks) + await cls._recover_orphan_tools(session_id) + + while True: + continuation_policy = ( + turn.continuation_policy or cls._continuation_policy + ) 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 continuation_policy.prepare_logical_turn(turn) + processed_user_id = ( + turn.prepared_user_id or processed_user_id + ) + outcome = await cls._run_logical_input(turn) + if await cls._should_continue( + turn, + continuation_policy, + outcome, + ): + continue + except Exception as exc: + outcome = await cls._handle_execution_error( + turn, + exc, + ) + + if await cls._settle_or_continue( + lease, + runtime_callbacks, + processed_user_id, + ): + continue + + settled = True + return cls._to_loop_result(turn, outcome) finally: - # Clean up - if session_id in cls._active_loops: - del cls._active_loops[session_id] - clear_turn_state(session_id) - - # Set status to idle - SessionStatus.set(session_id, SessionStatusIdle()) - await cls._publish_session_status(callbacks or LoopCallbacks(), session_id, "idle") - - # Touch session (update timestamp) - await Session.touch(session.project_id, session_id) - - # Publish idle event - try: - from flocks.bus.bus import Bus - from flocks.bus.events import SessionIdle - await Bus.publish(SessionIdle, {"sessionID": session_id}) - except Exception as exc: - log.warn("loop.idle.event_error", {"error": str(exc)}) - + if not settled and cls._leases.owns(lease): + await cls._release_session(lease, runtime_callbacks) + + @classmethod + async def _run_logical_input( + cls, + turn: LoopContext, + ) -> AgentRunOutcome[Any]: + """Execute one prepared logical input through AgentLoop.""" + return await AgentLoop().run( + turn, + StepEngine.from_turn(turn), + ) + + @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 get_context(cls, session_id: str) -> Optional[LoopContext]: + """Return the active turn used by public session controls.""" + return cls._active_turns.get(session_id) + + @classmethod + def abort(cls, session_id: str) -> bool: + """Abort 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_runtime_model( + cls, + provider_id: str, + model_id: str, + *, + config: Optional[Any] = None, + ) -> tuple[bool, str]: + """Validate one provider/model candidate.""" + return await cls._model_policy.validate_runtime_model( + provider_id, + model_id, + config=config, + ) + + @classmethod + async def validate_auto_configuration(cls) -> tuple[bool, str]: + """Validate that Auto mode has an available primary model.""" + 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 cls.validate_runtime_model( + default_llm["provider_id"], + default_llm["model_id"], + ) + if not available: + return False, f"primary_{reason}" + return True, "available" + @staticmethod async def _resolve_model( session: Any, @@ -809,105 +327,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 +453,223 @@ 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, - ) - ctx.model_candidate_policy = ( - "configured" if configured else "automatic" - ) - ctx.auto_failover = True - next_index = ( - 0 - if configured - else cls._cooldown_candidate_index( - ctx.session.id, - ctx.model_candidates, - ) + 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 ) - cls._select_candidate(ctx, next_index) - return next_index - - @classmethod - async def _prepare_auto_turn( - cls, - ctx: LoopContext, - last_user: MessageInfo, - ) -> bool: - """Synchronize routing when the loop advances to a real WebUI turn. - - Returns: - True when ``last_user`` starts a new non-synthetic user turn. - """ - if last_user.id == ctx.turn_user_id: - return False - - parts = await Message.parts(last_user.id, ctx.session.id) - if any(bool(getattr(part, "synthetic", False)) for part in parts): - return False - - if ctx.turn_user_id is None: - ctx.turn_user_id = last_user.id - if ctx.auto_failover and ctx.auto_failover_allowed: - from flocks.config.config import Config - - primary = ctx.model_candidates[0] - config = await Config.get() - await cls._reset_auto_turn_candidates( - ctx, - primary, - last_user.id, - config=config, + unhandled_runtime_error = outcome.unhandled_error + return LoopResult( + action=( + "error" + if ( + unhandled_runtime_error + or (turn.auto_failover and loop_error) ) - return True - - ctx.turn_user_id = last_user.id - persisted_session = await Session.get_by_id(ctx.session.id) - persisted_model_auto = bool( - persisted_session - and is_model_auto_session_category( - getattr(persisted_session, "category", "user") - ) - and getattr(persisted_session, "model_auto", False) - ) - persisted_auto = persisted_model_auto and ctx.auto_failover_allowed - - user_model = getattr(last_user, "model", None) - user_provider_id = None - user_model_id = None - if isinstance(user_model, dict): - user_provider_id = user_model.get("providerID") or user_model.get("provider_id") - user_model_id = user_model.get("modelID") or user_model.get("model_id") - - if not persisted_auto: - ctx.auto_failover = False - if not persisted_model_auto: - cls.clear_auto_failover_state(ctx.session.id) - ctx.auto_failover_allowed = False - provider_id = ( - getattr(persisted_session, "provider", None) - if Session.has_pinned_model(persisted_session) - else user_provider_id - ) or ctx.provider_id - model_id = ( - getattr(persisted_session, "model", None) - if Session.has_pinned_model(persisted_session) - else user_model_id - ) or ctx.model_id - ctx.model_candidates = [RuntimeModel(provider_id, model_id)] - ctx.model_candidate_policy = "fixed" - cls._select_candidate(ctx, 0) - log.info("session.model.auto_disabled_for_turn", { - "session_id": ctx.session.id, - "provider_id": provider_id, - "model_id": model_id, - }) - return True - - from flocks.config.config import Config - - config = await Config.get() - previous = RuntimeModel(ctx.provider_id, ctx.model_id) - default_llm = await Config.resolve_default_llm() - primary = RuntimeModel( - provider_id=(default_llm or {}).get("provider_id") or user_provider_id or ctx.provider_id, - model_id=(default_llm or {}).get("model_id") or user_model_id or ctx.model_id, - ) - next_index = await cls._reset_auto_turn_candidates( - ctx, - primary, - last_user.id, - config=config, + 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 {} + ), + }, ) - 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( - cls, - ctx: LoopContext, - last_user: MessageInfo, + @staticmethod + def _authorize_auto_failover( + turn: LoopContext, + requested: bool, ) -> None: - """Run UserPromptBefore once for a newly observed real user turn.""" + 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: - 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() + messages = await Message.list(session_id) + return sum( + 1 for message in messages if message.role == "assistant" + ) 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), - }) + log.debug("session.trace_offset.error", {"error": str(exc)}) + return 0 @classmethod - async def _run_turn_after_hook( + async def _acquire_lease( cls, - ctx: LoopContext, + 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 + + @staticmethod + async def _mark_busy( + session_id: str, callbacks: LoopCallbacks, - last_user: MessageInfo, - last_message: MessageInfo, - ) -> bool: - """Run turn.after with terminal facts; never continue from hook output.""" + ) -> None: + SessionStatus.set(session_id, SessionStatusBusy()) + await SessionEventSink.session_status(callbacks, session_id, "busy") + + @staticmethod + async def _recover_orphan_tools(session_id: str) -> None: try: - from flocks.hooks.pipeline import HookPipeline + from flocks.session.orphan_tools import abort_orphan_running_parts - 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", - }, - }) + await abort_orphan_running_parts(session_id) 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 - - @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", + log.warn( + "session.orphan_cleanup_failed", + {"session_id": session_id, "error": str(exc)}, ) - failure.assistant_message_id = assistant.id - return - await Message.update( - ctx.session.id, - failure.assistant_message_id, - error=failure.error_data, - finish="error", + + @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)}, + ) + try: + from flocks.bus.bus import Bus + from flocks.bus.events import SessionError - @classmethod - async def _process_step_with_failover( - cls, - ctx: LoopContext, - 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 - - 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, + await Bus.publish( + SessionError, + {"sessionID": session_id, "error": str(error)}, ) - 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 - - # 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, + 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, + ) @classmethod - async def _run_loop( + async def _release_session( 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, - ) - - 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 new user message flips the - # ID ordering so _should_exit() won't trigger. - continue - except Exception as e: - log.error("loop.compaction_overflow_check_error", {"error": str(e)}) - - # Process single step — wrap in a Task so abort() can cancel it immediately - # rather than waiting for the current tool call to finish. - step_task = asyncio.create_task( - cls._process_step_with_failover( - ctx, - callbacks, - messages, - last_user, - ) - ) - ctx._current_step_task = step_task - step_started_at = asyncio.get_event_loop().time() - try: - step_result = await step_task - except asyncio.CancelledError: - log.info("loop.step_cancelled", {"session_id": ctx.session.id, "step": ctx.step}) - break - finally: - ctx._current_step_task = None - log.debug("loop.step_complete", { - "session_id": ctx.session.id, - "step": ctx.step, - "duration_ms": int((asyncio.get_event_loop().time() - step_started_at) * 1000), - }) - - # Callback: step end - if callbacks.on_step_end: - await callbacks.on_step_end(ctx.step) - - # Handle result - if step_result.action == "stop": - loop_error = step_result.error - # Report error if step failed - if step_result.error and callbacks.on_error: - await callbacks.on_error(step_result.error) - - # Get last assistant message via SessionContext - if ctx.session_ctx: - post_messages = await ctx.session_ctx.get_messages() - else: - post_messages = await Message.list(ctx.session.id) - for msg in reversed(post_messages): - if ( - msg.role == MessageRole.ASSISTANT - and ( - not ctx.auto_failover - or getattr(msg, "parentID", None) == last_user.id - ) - ): - last_message = msg - break - - queued_user = await cls._detect_queued_user_message( - ctx.session.id, - post_messages, - last_user.id, - last_message, - ) - if queued_user is not None: - turn_state = set_turn_state( - ctx.session.id, - step=ctx.step, - status="continued", - continue_reason="queued_message", - queued_message_detected=True, - ) - await cls._publish_runtime_event(callbacks, "turn.continued", { - **turn_state.model_dump(by_alias=True), - "queuedUserMessageID": queued_user.id, - }) - log.info("loop.continuing_for_queued_message", { - "session_id": ctx.session.id, - "queued_user_id": queued_user.id, - "last_assistant_id": last_message.id if last_message else None, - }) - continue - - 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 - - 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 + ) -> None: + async with Session.lifecycle_lock(lease.session_id): + cls._finalize_release_state_locked(lease) + await cls._publish_released(lease.turn, callbacks) - 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)) - - 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], ) -> bool: - """ - Check if loop should exit - - Ported from original exit logic: - - Exit if assistant has responded with finish != tool-calls - - Exit if assistant message is after user message - """ - if not last_assistant: - return False + """Atomically keep ownership for late input or settle idle.""" + async with Session.lifecycle_lock(lease.session_id): + if await lease.turn.has_late_input(processed_user_id): + log.info( + "session.continuing_for_late_input", + { + "session_id": lease.session_id, + "processed_user_id": processed_user_id, + }, + ) + return True + cls._finalize_release_state_locked(lease) - if any( - getattr(part, "type", None) == "tool" - for part in (last_assistant_parts or []) - ): - return False - - # Check finish reason - if last_assistant.finish: - if last_assistant.finish not in ("tool-calls", "unknown", "summary"): - # Assistant finished with stop/error/etc - if last_user.id < last_assistant.id: - # Assistant responded after user - return True - + await cls._publish_released(lease.turn, callbacks) 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/streaming/stream_processor.py b/flocks/session/streaming/stream_processor.py index dc4e4b967..bab48f342 100644 --- a/flocks/session/streaming/stream_processor.py +++ b/flocks/session/streaming/stream_processor.py @@ -1526,7 +1526,7 @@ async def _handle_text_end(self, event: TextEndEvent) -> None: def _should_run_tool_call_parallel(self, event: ToolCallEvent) -> bool: """Return true for independent foreground subagent tool-calls.""" - if event.tool_name not in {"delegate_task", "task"}: + if event.tool_name != "delegate_task": return False tool_input = event.input if isinstance(event.input, dict) else {} if tool_input.get("run_in_background") is True: @@ -1778,7 +1778,7 @@ def _parse_dsml_text_tool_calls(self, text: str) -> list[dict]: re.DOTALL | re.IGNORECASE, ): body = match.group(1).strip() - if not body or not body[:1] in "{[": + if not body or body[:1] not in "{[": continue try: 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..7e598f01f 100644 --- a/flocks/task/background.py +++ b/flocks/task/background.py @@ -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/flocks/tool/agent/delegate_task.py b/flocks/tool/agent/delegate_task.py index eb8ea2195..c467f0db3 100644 --- a/flocks/tool/agent/delegate_task.py +++ b/flocks/tool/agent/delegate_task.py @@ -217,7 +217,7 @@ def _derive_task_description( - Background subagent execution is disabled. Do not set run_in_background=true. - Foreground execution is always used: the tool waits for completion and returns results inline. - For independent parallel work needed this turn, emit multiple sibling - foreground delegate_task/task tool calls in the same assistant response. + foreground delegate_task tool calls in the same assistant response. The runtime executes them concurrently and the webui renders each as its own DelegateTaskCard. @@ -231,6 +231,7 @@ def _derive_task_description( name="delegate_task", description=DESCRIPTION, category=ToolCategory.SYSTEM, + native=True, parameters=[ ToolParameter( name="load_skills", @@ -283,9 +284,8 @@ async def delegate_task_tool( load_skills: Optional[List[str]] = None, description: Optional[str] = None, # Internal-only: not exposed in the public schema. The registry rejects - # `run_in_background=True` at the schema layer for any caller, but legacy - # in-process call paths (e.g. `task.py` alias) may still pass it through. - # This guard is the second line of defense. + # `run_in_background=True` at the schema layer for any caller. This guard + # also protects direct in-process callers that bypass the registry. run_in_background: bool = False, subagent_type: Optional[str] = None, session_id: Optional[str] = None, @@ -297,7 +297,7 @@ async def delegate_task_tool( success=False, error=( "Background subagent execution is disabled. " - "Use foreground delegate_task/task calls; emit multiple sibling calls " + "Use foreground delegate_task calls; emit multiple sibling calls " "in the same assistant turn for parallel work." ), ) diff --git a/flocks/tool/agent/task.py b/flocks/tool/agent/task.py deleted file mode 100644 index 891231729..000000000 --- a/flocks/tool/agent/task.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Compatibility alias for delegate_task. - -The runtime keeps ``task`` as a registered tool name for workflow/backward -compatibility, but all scheduling behavior lives in ``delegate_task``. -Background subagent execution is disabled; run synchronously and emit -multiple sibling tool calls in one assistant turn for parallel work. -""" - -from __future__ import annotations - -from typing import Optional - -from flocks.tool.agent.delegate_task import delegate_task_tool -from flocks.tool.registry import ( - ParameterType, - ToolCategory, - ToolContext, - ToolParameter, - ToolRegistry, - ToolResult, -) - - -DESCRIPTION = """Compatibility alias for delegate_task. - -Use delegate_task directly for new prompts. Workflows may continue using task; -it accepts the same single-subagent shape and forwards it to delegate_task. -Background subagent execution is disabled; run synchronously and emit multiple -sibling tool calls in one assistant turn for parallel work. -""" - - -@ToolRegistry.register_function( - name="task", - description=DESCRIPTION, - category=ToolCategory.SYSTEM, - native=False, - parameters=[ - ToolParameter( - name="description", - type=ParameterType.STRING, - description="Optional short task description (3-5 words)", - required=False, - ), - ToolParameter( - name="prompt", - type=ParameterType.STRING, - description="Detailed prompt for the subagent.", - required=True, - ), - ToolParameter( - name="subagent_type", - type=ParameterType.STRING, - description="Delegatable agent name. Required for new tasks; omit when continuing with session_id.", - required=False, - ), - ToolParameter( - name="load_skills", - type=ParameterType.ARRAY, - description="Optional skill names to inject into the delegated agent", - required=False, - default=[], - ), - ToolParameter( - name="session_id", - type=ParameterType.STRING, - description="Existing subagent session to continue", - required=False, - ), - ToolParameter( - name="command", - type=ParameterType.STRING, - description="Deprecated command name retained for caller compatibility", - required=False, - ), - ToolParameter( - name="model", - type=ParameterType.STRING, - description="Optional model override (provider/model or model)", - required=False, - ), - ], -) -async def task_tool( - ctx: ToolContext, - description: Optional[str] = None, - prompt: Optional[str] = None, - subagent_type: Optional[str] = None, - load_skills: Optional[list] = None, - run_in_background: bool = False, - session_id: Optional[str] = None, - command: Optional[str] = None, - model: Optional[str] = None, -) -> ToolResult: - """Forward legacy task calls to delegate_task.""" - return await delegate_task_tool( - ctx=ctx, - prompt=prompt, - load_skills=load_skills, - description=description, - run_in_background=run_in_background, - subagent_type=subagent_type, - session_id=session_id, - command=command, - model=model, - ) diff --git a/flocks/tool/catalog.py b/flocks/tool/catalog.py index 40cffccb2..578c71422 100644 --- a/flocks/tool/catalog.py +++ b/flocks/tool/catalog.py @@ -37,7 +37,6 @@ class ToolCatalogMetadata(BaseModel): "webfetch": ["web", "http-fetch"], "websearch": ["web", "research"], "delegate_task": ["agent", "delegation"], - "task": ["agent", "delegation"], "schedule_task": ["scheduled-task", "scheduler-management"], "todo": ["task-management", "progress-tracking"], "run_workflow": ["workflow", "execution"], diff --git a/flocks/tool/registry.py b/flocks/tool/registry.py index 5b902c9cf..39b8cf6e0 100644 --- a/flocks/tool/registry.py +++ b/flocks/tool/registry.py @@ -1845,7 +1845,7 @@ def _register_builtin_tools(cls) -> None: # web/ — internet access ("flocks.tool.web", ["webfetch", "websearch"]), # agent/ — agent delegation/coordination - ("flocks.tool.agent", ["delegate_task", "task"]), + ("flocks.tool.agent", ["delegate_task"]), # task/ — task/workflow ("flocks.tool.task", [ "schedule_task_center", diff --git a/flocks/utils/id.py b/flocks/utils/id.py index 7257ac0db..a7d682582 100644 --- a/flocks/utils/id.py +++ b/flocks/utils/id.py @@ -25,7 +25,6 @@ "call", # cal "step", # stp "agent", # agt - "subtask", # stk "event", # evt "tqref", # tqr "chbind", # chb (channel session binding) @@ -54,7 +53,6 @@ class Identifier: "call": "cal", "step": "stp", "agent": "agt", - "subtask": "stk", "event": "evt", "tqref": "tqr", "task": "tsk", diff --git a/flocks/workflow/tool_context.py b/flocks/workflow/tool_context.py index 8526a410e..8982cbc4b 100644 --- a/flocks/workflow/tool_context.py +++ b/flocks/workflow/tool_context.py @@ -35,7 +35,7 @@ async def build_workflow_tool_context( Prefer the caller-provided session/message. When absent, create a temporary parent session and synthetic user message so workflow-internal tools such as - ``task`` / ``delegate_task`` can resolve a valid parent session. + ``delegate_task`` can resolve a valid parent session. """ effective_session_id = str(session_id or "").strip() diff --git a/tests/agent/test_unified_session_loop.py b/tests/agent/test_unified_session_loop.py index 44aaa8fa0..277913b0b 100644 --- a/tests/agent/test_unified_session_loop.py +++ b/tests/agent/test_unified_session_loop.py @@ -1,11 +1,7 @@ """ Tests for Phase 1: Unified UI entry via SessionLoop. -Verifies that: -1. 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): 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/integration/test_real_tool_calls.py b/tests/integration/test_real_tool_calls.py index 86adb3f5f..7b88006b2 100644 --- a/tests/integration/test_real_tool_calls.py +++ b/tests/integration/test_real_tool_calls.py @@ -191,16 +191,11 @@ async def on_tool_start(tool_name, args): async def on_tool_end(tool_name, result): tool_ends.append((tool_name, result)) - from flocks.session.runner import RunnerCallbacks - runner_callbacks = RunnerCallbacks( + callbacks = LoopCallbacks( on_tool_start=on_tool_start, on_tool_end=on_tool_end, ) - callbacks = LoopCallbacks( - runner_callbacks=runner_callbacks - ) - # Mock LLM with patch('flocks.provider.provider.Provider.chat') as mock_chat: first_response = MagicMock() diff --git a/tests/observability/test_langfuse_observability.py b/tests/observability/test_langfuse_observability.py index 8dcc02ef8..fa476e53a 100644 --- a/tests/observability/test_langfuse_observability.py +++ b/tests/observability/test_langfuse_observability.py @@ -125,7 +125,7 @@ def test_create_trace_forwards_tags(monkeypatch): monkeypatch.setattr(lf, "_get_client", lambda: client) obs = lf.create_trace( - name="SessionRunner.step", + name="StepEngine.step", session_id="s1", tags=["session:s1", "step:2", "session_step:s1:2"], input={"step": 2}, @@ -141,7 +141,7 @@ def test_create_trace_uses_start_observation_for_new_sdk(monkeypatch): monkeypatch.setattr(lf, "_get_client", lambda: client) obs = lf.create_trace( - name="SessionRunner.step", + name="StepEngine.step", session_id="s1", user_id="u1", tags=["session:s1", "step:2"], @@ -157,7 +157,7 @@ def test_create_trace_uses_start_observation_for_new_sdk(monkeypatch): assert client.start_observation_payload["metadata"]["session_id"] == "s1" assert client.start_observation_payload["metadata"]["user_id"] == "u1" assert client.start_observation_payload["metadata"]["tags"] == ["session:s1", "step:2"] - assert obs._otel_span.attributes["langfuse.trace.name"] == "SessionRunner.step" + assert obs._otel_span.attributes["langfuse.trace.name"] == "StepEngine.step" assert obs._otel_span.attributes["session.id"] == "s1" assert obs._otel_span.attributes["user.id"] == "u1" assert obs._otel_span.attributes["langfuse.trace.tags"] == ["session:s1", "step:2"] @@ -167,7 +167,7 @@ def test_generation_and_span_inherit_trace_dimensions_from_parent(monkeypatch): monkeypatch.setattr(lf, "_get_client", lambda: object()) parent = _TrackingObservation("trace", {"name": "trace"}) - parent._otel_span.attributes["langfuse.trace.name"] = "SessionRunner.step" + parent._otel_span.attributes["langfuse.trace.name"] = "StepEngine.step" parent._otel_span.attributes["session.id"] = "s1" parent._otel_span.attributes["user.id"] = "u1" parent._otel_span.attributes["langfuse.trace.tags"] = ["session:s1", "step:2"] @@ -175,11 +175,11 @@ def test_generation_and_span_inherit_trace_dimensions_from_parent(monkeypatch): gen = lf.create_generation(parent=parent, name="LLM.generate", model="gpt-5", input={"x": 1}) span = lf.create_span(parent=parent, name="Tool.execute.read", input={"path": "/tmp/a"}) - assert gen._otel_span.attributes["langfuse.trace.name"] == "SessionRunner.step" + assert gen._otel_span.attributes["langfuse.trace.name"] == "StepEngine.step" assert gen._otel_span.attributes["session.id"] == "s1" assert gen._otel_span.attributes["user.id"] == "u1" assert gen._otel_span.attributes["langfuse.trace.tags"] == ["session:s1", "step:2"] - assert span._otel_span.attributes["langfuse.trace.name"] == "SessionRunner.step" + assert span._otel_span.attributes["langfuse.trace.name"] == "StepEngine.step" assert span._otel_span.attributes["session.id"] == "s1" assert span._otel_span.attributes["user.id"] == "u1" assert span._otel_span.attributes["langfuse.trace.tags"] == ["session:s1", "step:2"] diff --git a/tests/permission/test_interactive.py b/tests/permission/test_interactive.py index 4521a34e8..f7eacbdaf 100644 --- a/tests/permission/test_interactive.py +++ b/tests/permission/test_interactive.py @@ -12,41 +12,3 @@ def test_auto_approve_enabled_reads_env(monkeypatch: pytest.MonkeyPatch) -> None assert auto_approve_enabled() is False monkeypatch.setenv("FLOCKS_AUTO_APPROVE", "true") assert auto_approve_enabled() is True - - -@pytest.mark.asyncio -async def test_runner_handle_permission_auto_allows_without_permission_next( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from flocks.session.runner import SessionRunner - - async def _unexpected_ask(*args, **kwargs): - raise AssertionError("PermissionNext.ask should not run for legacy tool permissions") - - monkeypatch.setattr( - "flocks.permission.next.PermissionNext.ask", - _unexpected_ask, - ) - - runner = SessionRunner.__new__(SessionRunner) - runner.session = type("Session", (), {"id": "ses_test"})() - runner._step = 1 - runner.callbacks = type( - "Callbacks", - (), - {"on_permission_request": None, "event_publish_callback": None}, - )() - - request = type( - "Request", - (), - { - "permission": "write", - "patterns": ["notes.md"], - "metadata": {}, - "message_id": "msg_1", - "always": ["*"], - }, - )() - - await runner._handle_permission(request) diff --git a/tests/server/routes/test_session_routes.py b/tests/server/routes/test_session_routes.py index d4798f1da..146eac9d7 100644 --- a/tests/server/routes/test_session_routes.py +++ b/tests/server/routes/test_session_routes.py @@ -21,7 +21,6 @@ from httpx import AsyncClient from flocks.auth.context import API_TOKEN_SERVICE_USER_ID, AuthUser from flocks.hooks.execution import ( - ExecutionStopped, current_execution_context, execution_context_scope, ) @@ -83,41 +82,6 @@ async def test_missing_session_directory_uses_cwd_and_publishes_notice( "fallbackDirectory": str(tmp_path), }, ) - - -@pytest.mark.asyncio -async def test_shell_route_maps_extension_stop_to_forbidden( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """A Pro policy stop must not surface as an unhandled server error.""" - - monkeypatch.setattr(session_routes, "require_user", lambda _request: object()) - monkeypatch.setattr( - session_routes, - "_get_session_by_id_unfiltered", - AsyncMock(return_value=object()), - ) - monkeypatch.setattr( - session_routes, - "_require_session_write_access", - lambda _session, _user: None, - ) - monkeypatch.setattr( - "flocks.session.runner.SessionRunner.shell", - AsyncMock(side_effect=ExecutionStopped("hard_deny_system_delete")), - ) - - with pytest.raises(HTTPException) as error: - await session_routes.run_shell_command( - "ses_1", - session_routes.ShellRequest(agent="build", command="rm -rf /etc"), - SimpleNamespace(), - ) - - assert error.value.status_code == status.HTTP_403_FORBIDDEN - assert error.value.detail == "execution stopped by extension" - - @pytest.mark.asyncio async def test_background_session_task_preserves_execution_context() -> None: """Async session work retains opaque ingress context after scheduling.""" diff --git a/tests/session/runtime/test_agent_loop.py b/tests/session/runtime/test_agent_loop.py new file mode 100644 index 000000000..5d614592e --- /dev/null +++ b/tests/session/runtime/test_agent_loop.py @@ -0,0 +1,169 @@ +"""Tests for the logical-input AgentLoop.""" + +from __future__ import annotations + +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 + + 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 +@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..eaf30b4fa --- /dev/null +++ b/tests/session/runtime/test_contracts.py @@ -0,0 +1,26 @@ +"""Tests for replay-safe runtime request contracts.""" + +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" + first_tools = request.provider_tools() + first_tools[0]["function"]["name"] = "mutated" + + assert request.provider_tools()[0]["function"]["name"] == "read" + assert request.provider_options()["reasoning"]["effort"] == "high" diff --git a/tests/session/runtime/test_session_loop.py b/tests/session/runtime/test_session_loop.py new file mode 100644 index 000000000..ee2f0e773 --- /dev/null +++ b/tests/session/runtime/test_session_loop.py @@ -0,0 +1,170 @@ +"""SessionLoop lifecycle and logical-turn ownership tests.""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from flocks.session.core.status import SessionStatus +from flocks.session.message import MessageRole +from flocks.session.runtime.agent_loop import AgentLoop +from flocks.session.runtime.contracts import ( + AgentRunOutcome, + 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 == {} diff --git a/tests/session/test_auto_model_failover.py b/tests/session/test_auto_model_failover.py index b1d655b1f..4890bcedc 100644 --- a/tests/session/test_auto_model_failover.py +++ b/tests/session/test_auto_model_failover.py @@ -6,22 +6,31 @@ 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 ( + AttemptEffects, + ModelTurnSnapshot, + RuntimeModel, +) 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 +75,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=SessionLoop.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 +114,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 +158,12 @@ def test_failover_classifier( message: str, reason: str, ): - decision = SessionRunner.classify_failover_error({ - "name": "APIError", - "data": {"message": message, "statusCode": status_code}, - }) + decision = StepEngine.classify_failover_error( + { + "name": "APIError", + "data": {"message": message, "statusCode": status_code}, + } + ) assert decision.eligible is True assert decision.reason == reason @@ -144,7 +188,7 @@ async def test_auto_runner_uses_standard_retry_policy( status_code: int, expected_calls: int, ): - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="primary", model_id="primary-model", @@ -160,19 +204,21 @@ async def test_auto_runner_uses_standard_retry_policy( call_llm = AsyncMock(side_effect=failure) monkeypatch.setattr( - "flocks.session.runner.Agent.get", - AsyncMock(return_value=SimpleNamespace( - name="rex", - steps=None, - mode="primary", - prompt="", - tools=[], - )), - ) - monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) - monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + "flocks.session.runtime.step_engine.Agent.get", + AsyncMock( + return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + ) + ), + ) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.apply_config", AsyncMock()) monkeypatch.setattr( - "flocks.session.runner.SessionPrompt.build_system_prompts", + "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompt_blocks", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -186,7 +232,7 @@ async def test_auto_runner_uses_standard_retry_policy( monkeypatch.setattr(Message, "create", AsyncMock(return_value=assistant)) monkeypatch.setattr(Message, "update", AsyncMock()) monkeypatch.setattr(runner, "_call_llm", call_llm) - monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", AsyncMock()) + monkeypatch.setattr("flocks.session.runtime.step_engine.SessionRetry.sleep", AsyncMock()) result = await runner._process_step([last_user], last_user) @@ -208,7 +254,7 @@ async def test_last_auto_candidate_uses_standard_retry_policy( expected_calls: int, ): """The last candidate uses the same retry policy as every other mode.""" - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="fallback", model_id="fallback-model", @@ -224,19 +270,21 @@ async def test_last_auto_candidate_uses_standard_retry_policy( call_llm = AsyncMock(side_effect=failure) monkeypatch.setattr( - "flocks.session.runner.Agent.get", - AsyncMock(return_value=SimpleNamespace( - name="rex", - steps=None, - mode="primary", - prompt="", - tools=[], - )), - ) - monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) - monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + "flocks.session.runtime.step_engine.Agent.get", + AsyncMock( + return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + ) + ), + ) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.apply_config", AsyncMock()) monkeypatch.setattr( - "flocks.session.runner.SessionPrompt.build_system_prompts", + "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompt_blocks", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -250,7 +298,7 @@ async def test_last_auto_candidate_uses_standard_retry_policy( monkeypatch.setattr(Message, "create", AsyncMock(return_value=assistant)) monkeypatch.setattr(Message, "update", AsyncMock()) monkeypatch.setattr(runner, "_call_llm", call_llm) - monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", AsyncMock()) + monkeypatch.setattr("flocks.session.runtime.step_engine.SessionRetry.sleep", AsyncMock()) result = await runner._process_step([last_user], last_user) @@ -262,9 +310,7 @@ async def test_last_auto_candidate_uses_standard_retry_policy( ("exception", "status_code", "reason"), [ ( - type("GoogleSdkError", (RuntimeError,), {"code": 429})( - "Resource exhausted" - ), + type("GoogleSdkError", (RuntimeError,), {"code": 429})("Resource exhausted"), 429, "rate_limit", ), @@ -289,7 +335,7 @@ def test_exception_status_is_normalized_from_sdk_shapes( status_code: int, reason: str, ): - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="primary", model_id="primary-model", @@ -298,16 +344,14 @@ def test_exception_status_is_normalized_from_sdk_shapes( error = runner._exception_to_error_dict(exception) assert error["data"]["statusCode"] == status_code - assert SessionRunner.classify_failover_error(error).reason == reason + assert StepEngine.classify_failover_error(error).reason == reason def test_exception_status_is_normalized_from_cause_chain(): - inner = type("GoogleSdkError", (RuntimeError,), {"code": 401})( - "Unauthenticated" - ) + inner = type("GoogleSdkError", (RuntimeError,), {"code": 401})("Unauthenticated") outer = RuntimeError("Provider wrapper failed") outer.__cause__ = inner - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="primary", model_id="primary-model", @@ -316,34 +360,40 @@ def test_exception_status_is_normalized_from_cause_chain(): error = runner._exception_to_error_dict(outer) assert error["data"]["statusCode"] == 401 - assert SessionRunner.classify_failover_error(error).reason == "auth" + assert StepEngine.classify_failover_error(error).reason == "auth" def test_local_validation_error_never_fails_over(): - decision = SessionRunner.classify_failover_error({ - "name": "ValidationError", - "data": {"message": "Local prompt schema validation failed"}, - }) + decision = StepEngine.classify_failover_error( + { + "name": "ValidationError", + "data": {"message": "Local prompt schema validation failed"}, + } + ) assert decision.eligible is False assert decision.reason == "local_error" def test_model_not_found_without_status_fails_over(): - decision = SessionRunner.classify_failover_error({ - "name": "ValueError", - "data": {"message": "Model acme-v2 not found for provider custom"}, - }) + decision = StepEngine.classify_failover_error( + { + "name": "ValueError", + "data": {"message": "Model acme-v2 not found for provider custom"}, + } + ) assert decision.eligible is True assert decision.reason == "model_not_found" def test_content_filter_error_fails_over_immediately(): - decision = SessionRunner.classify_failover_error({ - "name": "BadRequestError", - "data": {"message": "Response blocked by content_filter"}, - }) + decision = StepEngine.classify_failover_error( + { + "name": "BadRequestError", + "data": {"message": "Response blocked by content_filter"}, + } + ) assert decision.eligible is True assert decision.reason == "content_policy" @@ -356,22 +406,24 @@ def test_candidate_switch_keeps_tool_loop_guard_only(): "signature": "same-tool-call", "count": 2, } - ctx.runner_static_cache.update({ - "tool_loop_guard": tool_loop_guard, - "tool_schema_cache": {"primary": "schema"}, - "chat_context_cache": {"primary": "context"}, - "system_prompt": "primary prompt", - }) + ctx.step_static_cache.update( + { + "tool_loop_guard": tool_loop_guard, + "tool_schema_cache": {"primary": "schema"}, + "chat_context_cache": {"primary": "context"}, + "system_prompt": "primary prompt", + } + ) - SessionLoop._select_candidate(ctx, 1) + DEFAULT_MODEL_ROUTING_POLICY.select_candidate(ctx, 1) - assert ctx.runner_static_cache == {"tool_loop_guard": tool_loop_guard} - assert ctx.runner_static_cache["tool_loop_guard"] is tool_loop_guard + assert ctx.step_static_cache == {"tool_loop_guard": tool_loop_guard} + assert ctx.step_static_cache["tool_loop_guard"] is tool_loop_guard @pytest.mark.asyncio async def test_reasoning_only_empty_response_is_not_replayed(monkeypatch): - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="primary", model_id="primary-model", @@ -391,19 +443,21 @@ async def call_llm(*_args, **_kwargs): return StepResult(action="stop", content="") monkeypatch.setattr( - "flocks.session.runner.Agent.get", - AsyncMock(return_value=SimpleNamespace( - name="rex", - steps=None, - mode="primary", - prompt="", - tools=[], - )), - ) - monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) - monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + "flocks.session.runtime.step_engine.Agent.get", + AsyncMock( + return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + ) + ), + ) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.apply_config", AsyncMock()) monkeypatch.setattr( - "flocks.session.runner.SessionPrompt.build_system_prompts", + "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompt_blocks", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -418,7 +472,7 @@ async def call_llm(*_args, **_kwargs): monkeypatch.setattr(Message, "update", AsyncMock()) monkeypatch.setattr(runner, "_call_llm", call_llm) sleep = AsyncMock() - monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", sleep) + monkeypatch.setattr("flocks.session.runtime.step_engine.SessionRetry.sleep", sleep) result = await runner._process_step([last_user], last_user) @@ -490,11 +544,13 @@ def get_reasoning_content(self): reasoning=None, event_type=None, metadata={}, - tool_calls=[{ - "index": 0, - "id": "call_1", - "function": {"name": "example_tool", "arguments": "{}"}, - }], + tool_calls=[ + { + "index": 0, + "id": "call_1", + "function": {"name": "example_tool", "arguments": "{}"}, + } + ], finish_reason=None, usage=None, ) @@ -518,7 +574,7 @@ async def stream(): return stream() provider = FailingStreamProvider() - runner = SessionRunner( + runner = StepEngine( session=_session(), provider_id="primary", model_id="primary-model", @@ -529,19 +585,21 @@ async def stream(): assistant = SimpleNamespace(id="msg_assistant") monkeypatch.setattr( - "flocks.session.runner.Agent.get", - AsyncMock(return_value=SimpleNamespace( - name="rex", - steps=None, - mode="primary", - prompt="", - tools=[], - )), - ) - monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) - monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + "flocks.session.runtime.step_engine.Agent.get", + AsyncMock( + return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + ) + ), + ) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runtime.step_engine.Provider.apply_config", AsyncMock()) monkeypatch.setattr( - "flocks.session.runner.SessionPrompt.build_system_prompts", + "flocks.session.runtime.step_engine.SessionPrompt.build_system_prompt_blocks", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -555,18 +613,18 @@ async def stream(): monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr(Message, "create", AsyncMock(return_value=assistant)) monkeypatch.setattr(Message, "update", AsyncMock()) - monkeypatch.setattr("flocks.session.runner.StreamProcessor", FakeStreamProcessor) + monkeypatch.setattr("flocks.session.runtime.step_engine.StreamProcessor", FakeStreamProcessor) monkeypatch.setattr( - "flocks.session.runner.HookPipeline.has_stage_handlers", + "flocks.session.runtime.step_engine.HookPipeline.has_stage_handlers", AsyncMock(return_value=False), ) - monkeypatch.setattr("flocks.session.runner.langfuse_is_active", lambda: False) + monkeypatch.setattr("flocks.session.runtime.step_engine.langfuse_is_active", lambda: False) monkeypatch.setattr( "flocks.provider.options.build_provider_options", lambda _provider_id, _model_id: {}, ) sleep = AsyncMock() - monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", sleep) + monkeypatch.setattr("flocks.session.runtime.step_engine.SessionRetry.sleep", sleep) result = await runner._process_step([last_user], last_user) @@ -575,9 +633,7 @@ async def stream(): assert result.failure.allow_fallback is False assert result.failure.attempt_state.received_chunk is True assert result.failure.attempt_state.observable_output_started is True - assert result.failure.attempt_state.tool_execution_started is ( - chunk_kind == "tool" - ) + assert result.failure.attempt_state.tool_execution_started is (chunk_kind == "tool") sleep.assert_not_awaited() @@ -592,14 +648,14 @@ async def process_step(runner, _messages, _last_user): return _failure(assistant_id="msg_failed") return StepResult(action="stop", content="recovered") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) delete = AsyncMock(return_value=True) monkeypatch.setattr(Message, "delete", delete) async def publish(event, payload): events.append((event, payload)) - result = await SessionLoop._process_step_with_failover( + result = await _process_step_with_failover( ctx, LoopCallbacks(event_publish_callback=publish), [last_user], @@ -622,7 +678,7 @@ async def test_queued_user_is_detected_before_replacement_assistant(): role=MessageRole.ASSISTANT, ) - detected = await SessionLoop._detect_queued_user_message( + detected = await DEFAULT_CONTINUATION_POLICY.detect_queued_user_message( "ses_auto", [current_user, queued_user, replacement_assistant], current_user.id, @@ -655,17 +711,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 +752,7 @@ async def test_failed_blank_message_deletion_stops_switch(monkeypatch): ctx = _ctx() last_user = SimpleNamespace(id="msg_user", agent="rex") monkeypatch.setattr( - SessionRunner, + StepEngine, "_process_step", AsyncMock(return_value=_failure(assistant_id="msg_failed")), ) @@ -704,7 +760,7 @@ async def test_failed_blank_message_deletion_stops_switch(monkeypatch): update = AsyncMock() monkeypatch.setattr(Message, "update", update) - result = await SessionLoop._process_step_with_failover( + result = await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], @@ -738,11 +794,11 @@ async def process_step(runner, _messages, _last_user): return _failure(assistant_id=f"msg_{runner.provider_id}") return StepResult(action="stop", content="recovered") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) delete = AsyncMock(return_value=True) monkeypatch.setattr(Message, "delete", delete) - result = await SessionLoop._process_step_with_failover( + result = await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], @@ -771,13 +827,13 @@ async def test_chain_exhaustion_finalizes_only_last_candidate(monkeypatch): async def process_step(runner, _messages, _last_user): return _failure(assistant_id=f"msg_{runner.provider_id}") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) delete = AsyncMock(return_value=True) update = AsyncMock() monkeypatch.setattr(Message, "delete", delete) monkeypatch.setattr(Message, "update", update) - result = await SessionLoop._process_step_with_failover( + result = await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], @@ -788,7 +844,7 @@ async def process_step(runner, _messages, _last_user): assert delete.await_count == 2 update.assert_awaited_once() assert update.await_args.args[1] == "msg_fallback-2" - cooldown = SessionLoop._auto_failover_cooldowns[ctx.session.id] + cooldown = DEFAULT_MODEL_ROUTING_POLICY.cooldowns[ctx.session.id] assert cooldown.model == RuntimeModel("fallback-2", "model-2") assert cooldown.reason == "chain_exhausted" @@ -811,11 +867,13 @@ async def test_full_loop_reports_chain_exhaustion_once(monkeypatch): parentID=user.id, finish="error", ) - ctx.session_ctx = SimpleNamespace( - get_messages=AsyncMock(side_effect=[ - [user], - [user, final_assistant], - ]) + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user], + [user, final_assistant], + ] + ) ) attempts = [] @@ -823,14 +881,14 @@ async def process_step(runner, _messages, _last_user): attempts.append((runner.provider_id, runner.model_id)) return _failure(assistant_id=f"msg_{runner.provider_id}") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) update = AsyncMock() monkeypatch.setattr(Message, "update", update) on_error = AsyncMock() - result = await SessionLoop._run_loop( + result = await run_logical_turns( ctx, LoopCallbacks( on_error=on_error, @@ -859,7 +917,7 @@ async def test_observable_failure_is_finalized_without_replay(monkeypatch): ctx = _ctx() last_user = SimpleNamespace(id="msg_user", agent="rex") monkeypatch.setattr( - SessionRunner, + StepEngine, "_process_step", AsyncMock(return_value=_failure(assistant_id="msg_partial", safe=False)), ) @@ -868,7 +926,7 @@ async def test_observable_failure_is_finalized_without_replay(monkeypatch): monkeypatch.setattr(Message, "delete", delete) monkeypatch.setattr(Message, "update", update) - result = await SessionLoop._process_step_with_failover( + result = await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], @@ -891,23 +949,26 @@ async def process_step(runner, _messages, _last_user): return _failure(assistant_id="msg_rate", reason="rate_limit") return StepResult(action="stop", content="recovered") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) - await SessionLoop._process_step_with_failover( + await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], last_user, ) - cooldown = SessionLoop._auto_failover_cooldowns[ctx.session.id] + cooldown = DEFAULT_MODEL_ROUTING_POLICY.cooldowns[ctx.session.id] assert cooldown.model == RuntimeModel("fallback", "fallback-model") assert cooldown.reason == "rate_limit" - assert SessionLoop._cooldown_candidate_index( - ctx.session.id, - ctx.model_candidates, - ) == 1 + assert ( + DEFAULT_MODEL_ROUTING_POLICY.cooldown_candidate_index( + ctx.session.id, + ctx.model_candidates, + ) + == 1 + ) @pytest.mark.asyncio @@ -921,10 +982,10 @@ async def process_step(runner, _messages, _last_user): return _failure(assistant_id="msg_rate", reason="rate_limit") return StepResult(action="stop", content="recovered") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) - await SessionLoop._process_step_with_failover( + await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], @@ -932,15 +993,17 @@ async def process_step(runner, _messages, _last_user): ) assert (ctx.provider_id, ctx.model_id) == ("fallback", "fallback-model") - assert ctx.session.id not in SessionLoop._auto_failover_cooldowns + assert ctx.session.id not in DEFAULT_MODEL_ROUTING_POLICY.cooldowns @pytest.mark.asyncio async def test_403_quota_failure_sets_primary_cooldown(monkeypatch): - decision = SessionRunner.classify_failover_error({ - "name": "APIError", - "data": {"message": "Quota exceeded", "statusCode": 403}, - }) + decision = StepEngine.classify_failover_error( + { + "name": "APIError", + "data": {"message": "Quota exceeded", "statusCode": 403}, + } + ) ctx = _ctx() last_user = SimpleNamespace(id="msg_user", agent="rex") @@ -952,17 +1015,17 @@ async def process_step(runner, _messages, _last_user): ) return StepResult(action="stop", content="recovered") - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) - await SessionLoop._process_step_with_failover( + await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], last_user, ) - cooldown = SessionLoop._auto_failover_cooldowns[ctx.session.id] + cooldown = DEFAULT_MODEL_ROUTING_POLICY.cooldowns[ctx.session.id] assert cooldown.reason == "rate_limit" assert cooldown.model == RuntimeModel("fallback", "fallback-model") assert cooldown.expires_at > time.monotonic() + 50 @@ -980,18 +1043,18 @@ async def process_step(runner, _messages, _last_user): reason=reason, ) - monkeypatch.setattr(SessionRunner, "_process_step", process_step) + monkeypatch.setattr(StepEngine, "_process_step", process_step) monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) monkeypatch.setattr(Message, "update", AsyncMock()) - await SessionLoop._process_step_with_failover( + await _process_step_with_failover( ctx, LoopCallbacks(), [last_user], last_user, ) - cooldown = SessionLoop._auto_failover_cooldowns[ctx.session.id] + cooldown = DEFAULT_MODEL_ROUTING_POLICY.cooldowns[ctx.session.id] assert cooldown.reason == "rate_limit" assert cooldown.model == RuntimeModel("fallback", "fallback-model") # A 5s anti-replay window must not replace the primary's 60s cooldown. @@ -1032,11 +1095,11 @@ async def validate(provider_id, _model_id, **_kwargs): monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) primary = RuntimeModel("primary", "primary-model") - first = await SessionLoop._build_model_candidates( + first = await _build_model_candidates( primary, route_seed="ses_auto:msg_1", ) - repeated = await SessionLoop._build_model_candidates( + repeated = await _build_model_candidates( primary, route_seed="ses_auto:msg_1", ) @@ -1050,10 +1113,12 @@ async def validate(provider_id, _model_id, **_kwargs): assert all(candidate.provider_id != "missing" for candidate in first) selections = { - tuple(await SessionLoop._build_model_candidates( - primary, - route_seed=f"ses_auto:msg_{index}", - )) + tuple( + await _build_model_candidates( + primary, + route_seed=f"ses_auto:msg_{index}", + ) + ) for index in range(12) } assert len(selections) > 1 @@ -1092,7 +1157,7 @@ async def test_candidate_builder_keeps_active_cooldown_model_in_its_tier( primary = RuntimeModel("primary", "primary-model") cooldown_model = RuntimeModel("other", "other-b") - candidates = await SessionLoop._build_model_candidates( + candidates = await _build_model_candidates( primary, route_seed="ses_auto:new-turn", preferred=cooldown_model, @@ -1107,21 +1172,19 @@ async def test_candidate_builder_keeps_active_cooldown_model_in_its_tier( async def test_auto_configuration_only_requires_available_primary(monkeypatch): monkeypatch.setattr( "flocks.config.config.Config.resolve_default_llm", - AsyncMock(return_value={ - "provider_id": "primary", - "model_id": "primary-model", - }), + AsyncMock( + return_value={ + "provider_id": "primary", + "model_id": "primary-model", + } + ), ) monkeypatch.setattr( SessionLoop, "validate_runtime_model", AsyncMock(return_value=(True, "available")), ) - build_candidates = AsyncMock() - monkeypatch.setattr(SessionLoop, "_build_model_candidates", build_candidates) - assert await SessionLoop.validate_auto_configuration() == (True, "available") - build_candidates.assert_not_awaited() @pytest.mark.asyncio @@ -1145,7 +1208,7 @@ async def test_candidate_builder_allows_primary_only_chain(monkeypatch): primary = RuntimeModel("primary", "primary-model") - assert await SessionLoop._build_model_candidates( + assert await _build_model_candidates( primary, route_seed="ses_auto:msg_primary_only", ) == [primary] @@ -1155,11 +1218,13 @@ async def test_candidate_builder_allows_primary_only_chain(monkeypatch): async def test_candidate_builder_uses_configured_order_without_discovery( monkeypatch, ): - config = SimpleNamespace(fallback_providers=[ - SimpleNamespace(provider_id="other", model_id="model-b"), - SimpleNamespace(provider_id="primary", model_id="model-a"), - SimpleNamespace(provider_id="missing", model_id="missing-model"), - ]) + config = SimpleNamespace( + fallback_providers=[ + SimpleNamespace(provider_id="other", model_id="model-b"), + SimpleNamespace(provider_id="primary", model_id="model-a"), + SimpleNamespace(provider_id="missing", model_id="missing-model"), + ] + ) model_manager = MagicMock() monkeypatch.setattr( "flocks.provider.provider.Provider.apply_config", @@ -1177,7 +1242,7 @@ async def validate(provider_id, _model_id, **_kwargs): monkeypatch.setattr(SessionLoop, "validate_runtime_model", validate) primary = RuntimeModel("primary", "primary-model") - candidates = await SessionLoop._build_model_candidates( + candidates = await _build_model_candidates( primary, route_seed="unused-for-configured", preferred=RuntimeModel("other", "model-b"), @@ -1196,9 +1261,11 @@ async def validate(provider_id, _model_id, **_kwargs): async def test_configured_chain_with_no_available_fallbacks_keeps_primary_only( monkeypatch, ): - config = SimpleNamespace(fallback_providers=[ - SimpleNamespace(provider_id="missing", model_id="missing-model"), - ]) + config = SimpleNamespace( + fallback_providers=[ + SimpleNamespace(provider_id="missing", model_id="missing-model"), + ] + ) monkeypatch.setattr( "flocks.provider.provider.Provider.apply_config", AsyncMock(), @@ -1210,7 +1277,7 @@ async def test_configured_chain_with_no_available_fallbacks_keeps_primary_only( ) primary = RuntimeModel("primary", "primary-model") - assert await SessionLoop._build_model_candidates( + assert await _build_model_candidates( primary, route_seed="unused-for-configured", config=config, @@ -1222,24 +1289,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 +1315,7 @@ async def test_synthetic_subtask_continuation_keeps_fallback(monkeypatch): AsyncMock(return_value=[SimpleNamespace(synthetic=True)]), ) - await SessionLoop._prepare_auto_turn(ctx, synthetic_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, synthetic_user) assert ctx.auto_failover is True assert ctx.turn_user_id == "msg_real" @@ -1275,9 +1342,9 @@ async def test_first_real_turn_builds_stable_chain_from_user_id(monkeypatch): AsyncMock(return_value=config), ) build = AsyncMock(return_value=rebuilt) - monkeypatch.setattr(SessionLoop, "_build_model_candidates", build) + monkeypatch.setattr(DEFAULT_MODEL_ROUTING_POLICY, "build_candidates", build) - await SessionLoop._prepare_auto_turn(ctx, first_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, first_user) assert ctx.turn_user_id == "msg_first" assert ctx.model_candidates == rebuilt @@ -1299,14 +1366,16 @@ async def test_configured_first_real_turn_ignores_cooldown_and_starts_primary( id="msg_first", model={"providerID": "primary", "modelID": "primary-model"}, ) - config = SimpleNamespace(fallback_providers=[ - SimpleNamespace(provider_id="fallback", model_id="fallback-model"), - ]) + config = SimpleNamespace( + fallback_providers=[ + SimpleNamespace(provider_id="fallback", model_id="fallback-model"), + ] + ) rebuilt = [ RuntimeModel("primary", "primary-model"), RuntimeModel("fallback", "fallback-model"), ] - SessionLoop._auto_failover_cooldowns[ctx.session.id] = AutoFailoverCooldown( + DEFAULT_MODEL_ROUTING_POLICY.cooldowns[ctx.session.id] = AutoFailoverCooldown( model=rebuilt[1], primary=rebuilt[0], expires_at=float("inf"), @@ -1317,17 +1386,17 @@ async def test_configured_first_real_turn_ignores_cooldown_and_starts_primary( AsyncMock(return_value=config), ) monkeypatch.setattr( - SessionLoop, - "_build_model_candidates", + DEFAULT_MODEL_ROUTING_POLICY, + "build_candidates", AsyncMock(return_value=rebuilt), ) - await SessionLoop._prepare_auto_turn(ctx, first_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, first_user) assert ctx.model_candidate_policy == "configured" assert ctx.candidate_index == 0 assert (ctx.provider_id, ctx.model_id) == ("primary", "primary-model") - assert ctx.session.id not in SessionLoop._auto_failover_cooldowns + assert ctx.session.id not in DEFAULT_MODEL_ROUTING_POLICY.cooldowns @pytest.mark.asyncio @@ -1350,7 +1419,7 @@ async def test_queued_explicit_model_disables_auto(monkeypatch): AsyncMock(return_value=persisted), ) - await SessionLoop._prepare_auto_turn(ctx, queued_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, queued_user) assert ctx.auto_failover is False assert ctx.model_candidates == [RuntimeModel("explicit", "explicit-model")] @@ -1372,7 +1441,7 @@ async def test_non_webui_loop_cannot_activate_persisted_auto(monkeypatch): AsyncMock(return_value=_session(model_auto=True)), ) - await SessionLoop._prepare_auto_turn(ctx, queued_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, queued_user) assert ctx.auto_failover is False assert ctx.auto_failover_allowed is False @@ -1399,10 +1468,12 @@ async def test_queued_webui_turn_rebuilds_auto_chain(monkeypatch): ) monkeypatch.setattr( "flocks.config.config.Config.resolve_default_llm", - AsyncMock(return_value={ - "provider_id": "primary", - "model_id": "primary-model", - }), + AsyncMock( + return_value={ + "provider_id": "primary", + "model_id": "primary-model", + } + ), ) config = SimpleNamespace(fallback_providers=None) monkeypatch.setattr( @@ -1410,9 +1481,9 @@ async def test_queued_webui_turn_rebuilds_auto_chain(monkeypatch): AsyncMock(return_value=config), ) build = AsyncMock(return_value=rebuilt) - monkeypatch.setattr(SessionLoop, "_build_model_candidates", build) + monkeypatch.setattr(DEFAULT_MODEL_ROUTING_POLICY, "build_candidates", build) - await SessionLoop._prepare_auto_turn(ctx, queued_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, queued_user) assert ctx.auto_failover is True assert ctx.model_candidates == rebuilt @@ -1438,9 +1509,11 @@ async def test_queued_configured_turn_restarts_from_primary(monkeypatch): RuntimeModel("primary", "primary-model"), RuntimeModel("fallback", "fallback-model"), ] - config = SimpleNamespace(fallback_providers=[ - SimpleNamespace(provider_id="fallback", model_id="fallback-model"), - ]) + config = SimpleNamespace( + fallback_providers=[ + SimpleNamespace(provider_id="fallback", model_id="fallback-model"), + ] + ) monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr( "flocks.session.session.Session.get_by_id", @@ -1452,18 +1525,20 @@ async def test_queued_configured_turn_restarts_from_primary(monkeypatch): ) monkeypatch.setattr( "flocks.config.config.Config.resolve_default_llm", - AsyncMock(return_value={ - "provider_id": "primary", - "model_id": "primary-model", - }), + AsyncMock( + return_value={ + "provider_id": "primary", + "model_id": "primary-model", + } + ), ) monkeypatch.setattr( - SessionLoop, - "_build_model_candidates", + DEFAULT_MODEL_ROUTING_POLICY, + "build_candidates", AsyncMock(return_value=rebuilt), ) - await SessionLoop._prepare_auto_turn(ctx, queued_user) + await DEFAULT_MODEL_ROUTING_POLICY.prepare_turn(ctx, queued_user) assert ctx.model_candidate_policy == "configured" assert ctx.candidate_index == 0 @@ -1474,11 +1549,11 @@ async def test_queued_configured_turn_restarts_from_primary(monkeypatch): @pytest.mark.parametrize("category", ["user", "entity-config", "workflow"]) async def test_queued_webui_auto_authorizes_active_loop(category): ctx = _ctx(auto=False, category=category) - SessionLoop._active_loops[ctx.session.id] = ctx + SessionLoop._active_turns[ctx.session.id] = ctx try: result = await SessionLoop.run(ctx.session.id, auto_failover=True) finally: - SessionLoop._active_loops.pop(ctx.session.id, None) + SessionLoop._active_turns.pop(ctx.session.id, None) assert result.action == "queued" assert ctx.auto_failover_allowed is True @@ -1491,19 +1566,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 +1610,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 +1630,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 +1647,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_context_usage.py b/tests/session/test_context_usage.py index 2bc9db71a..303dcf88a 100644 --- a/tests/session/test_context_usage.py +++ b/tests/session/test_context_usage.py @@ -275,7 +275,7 @@ async def test_context_usage_splits_skill_and_delegation_tools(context_usage_moc ), SimpleNamespace( type="tool", - tool="task", + tool="delegate_task", state=SimpleNamespace(input={}, output="t" * 80, time={"start": 3}), ), SimpleNamespace( @@ -288,20 +288,16 @@ async def test_context_usage_splits_skill_and_delegation_tools(context_usage_moc metadata={"tool": "skill_load"}, state=SimpleNamespace(input={}, output="m" * 40, time={"start": 5}), ), - SimpleNamespace( - type="subtask", - prompt="p" * 40, - description="q" * 40, - ), ] } snapshot = await context_usage.build_context_usage_snapshot("sess-1") assert [(segment.key, segment.tokens) for segment in snapshot.segments] == [ + ("conversation", 20), ("tools", 30), ("skillLoad", 30), - ("agentDelegation", 50), + ("agentDelegation", 30), ] tools_segment = next(segment for segment in snapshot.segments if segment.key == "tools") assert tools_segment.tokens == 30 diff --git a/tests/session/test_execution_mode.py b/tests/session/test_execution_mode.py index 0ef33c1de..637d3cfcd 100644 --- a/tests/session/test_execution_mode.py +++ b/tests/session/test_execution_mode.py @@ -89,7 +89,6 @@ def test_plan_uses_read_only_permission_rules() -> None: assert is_tool_allowed(SessionExecutionMode.PLAN, "edit") assert is_tool_allowed(SessionExecutionMode.PLAN, "write") assert is_tool_allowed(SessionExecutionMode.PLAN, "unknown_plugin_tool") - assert is_tool_allowed(SessionExecutionMode.PLAN, "task") assert is_tool_allowed(SessionExecutionMode.PLAN, "delegate_task") assert not is_tool_allowed(SessionExecutionMode.PLAN, "run_slash_command") @@ -102,14 +101,13 @@ def test_plan_uses_read_only_permission_rules() -> None: assert execution_mode_prompt("build") == "" -@pytest.mark.parametrize("tool_name", ["task", "delegate_task"]) -def test_plan_delegation_only_allows_explore_and_librarian(tool_name) -> None: +def test_plan_delegation_only_allows_explore_and_librarian() -> None: ctx = ToolContext(session_id="session-1", message_id="message-1") for subagent_type in ("explore", "librarian"): assert tool_call_denial_reason( SessionExecutionMode.PLAN, - tool_name, + "delegate_task", {"subagent_type": subagent_type}, ctx, ) is None @@ -121,7 +119,7 @@ def test_plan_delegation_only_allows_explore_and_librarian(tool_name) -> None: ): reason = tool_call_denial_reason( SessionExecutionMode.PLAN, - tool_name, + "delegate_task", arguments, ctx, ) @@ -245,7 +243,7 @@ async def handler(_ctx, **_kwargs): tool = Tool( info=ToolInfo( - name="task", + name="delegate_task", description="Delegation test tool", category=ToolCategory.FILE, ), @@ -258,7 +256,7 @@ async def handler(_ctx, **_kwargs): ) explore = await ToolRegistry.execute( - "task", + "delegate_task", ctx=ToolContext( session_id="session-1", message_id="message-1", @@ -399,9 +397,9 @@ async def test_read_only_sandbox_allows_only_plan_artifact_write(tmp_path) -> No @pytest.mark.asyncio async def test_runner_filters_tools_with_message_mode(monkeypatch) -> None: - from flocks.session.runner import SessionRunner + from flocks.session.runtime.step_engine import StepEngine - runner = object.__new__(SessionRunner) + runner = object.__new__(StepEngine) runner.session = SimpleNamespace(id="session-1") runner._step = 1 runner.callbacks = SimpleNamespace(event_publish_callback=None) @@ -411,7 +409,6 @@ async def test_runner_filters_tools_with_message_mode(monkeypatch) -> None: "bash", "write", "edit", - "task", "delegate_task", "run_slash_command", ] @@ -423,7 +420,6 @@ async def test_runner_filters_tools_with_message_mode(monkeypatch) -> None: SimpleNamespace(name="bash"), SimpleNamespace(name="write"), SimpleNamespace(name="edit"), - SimpleNamespace(name="task"), SimpleNamespace(name="delegate_task"), SimpleNamespace(name="run_slash_command"), ], @@ -434,7 +430,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( @@ -465,7 +461,6 @@ async def list_tools(**_kwargs): "bash", "write", "edit", - "task", "delegate_task", "plan_exit", ] @@ -476,6 +471,5 @@ async def list_tools(**_kwargs): "edit", "plan_exit", "read", - "task", "write", ] diff --git a/tests/session/test_lifecycle_hooks.py b/tests/session/test_lifecycle_hooks.py index b2db3efe9..607a6c5a4 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,26 @@ 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_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 +86,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 +118,7 @@ async def test_session_start_runs_only_when_pending() -> None: run_hook = AsyncMock() with patch( - "flocks.session.runner.HookPipeline.run_session_start", + "flocks.session.runtime.step_engine.HookPipeline.run_session_start", run_hook, ): await runner._run_session_start_hook(SimpleNamespace(name="rex")) @@ -112,39 +126,614 @@ 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.", + ) + ), + ), ): - 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 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, + ), + ): + decision = await DEFAULT_CONTINUATION_POLICY.resolve( + ctx, + 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() + + +def _message( + message_id: str, + role: str, + *, + finish: str | None = None, +) -> SimpleNamespace: + return SimpleNamespace( + id=message_id, + role=role, + finish=finish, + 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") + ctx.session_store = SimpleNamespace( + get_messages=AsyncMock( + side_effect=[ + [user], + [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.py b/tests/session/test_message_parts.py index 89504ba62..c28a7d1c0 100644 --- a/tests/session/test_message_parts.py +++ b/tests/session/test_message_parts.py @@ -27,7 +27,6 @@ SnapshotPart, StepFinishPart, StepStartPart, - SubtaskPart, TextPart, TokenCache, TokenUsage, @@ -313,21 +312,9 @@ def test_creation(self): # --------------------------------------------------------------------------- -# SubtaskPart / AgentPart +# AgentPart # --------------------------------------------------------------------------- -class TestSubtaskPart: - def test_creation(self): - part = SubtaskPart( - sessionID=SID, - messageID=MID, - prompt="Summarize findings", - description="Summarize", - agent="rex", - ) - assert part.type == "subtask" - assert part.agent == "rex" - class TestAgentPart: def test_creation(self): @@ -416,6 +403,24 @@ def test_deserialize_reasoning_part(self): assert deserialized is not None assert deserialized.type == "reasoning" + def test_deserialize_legacy_subtask_as_ignored_text(self): + deserialized = Message.deserialize_part( + { + "id": "part_legacy_subtask", + "sessionID": SID, + "messageID": MID, + "type": "subtask", + "prompt": "old delegated command", + "description": "legacy", + "agent": "rex", + } + ) + + assert deserialized.type == "text" + assert deserialized.text == "" + assert deserialized.ignored is True + assert deserialized.metadata == {"legacyPartType": "subtask"} + def test_deserialize_unknown_type_falls_back_to_text(self): # Unknown type falls back to TextPart; missing required fields raise exception with pytest.raises(Exception): @@ -537,11 +542,11 @@ async def test_store_part_does_not_downgrade_terminal_tool_state(self, monkeypat sessionID=sid, messageID=msg.id, callID="call_terminal_guard", - tool="task", + tool="delegate_task", state=ToolStateCompleted( input={"prompt": "run"}, output="done", - title="task", + title="delegate_task", metadata={"sessionId": "ses_child_done"}, time={"start": 1000, "end": 2000}, ), @@ -551,10 +556,10 @@ async def test_store_part_does_not_downgrade_terminal_tool_state(self, monkeypat sessionID=sid, messageID=msg.id, callID="call_terminal_guard", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "run"}, - title="task", + title="delegate_task", metadata={"sessionId": "ses_child_done", "status": "running"}, time={"start": 1000}, ), diff --git a/tests/session/test_prompt_tokens.py b/tests/session/test_prompt_tokens.py index a1f44d984..9c3d92dba 100644 --- a/tests/session/test_prompt_tokens.py +++ b/tests/session/test_prompt_tokens.py @@ -26,6 +26,7 @@ PromptTemplate, SessionPrompt, SystemPrompt, + TurnPromptContext, ) from flocks.session import prompt_strings @@ -269,7 +270,9 @@ async def test_builtin_system_subagent_child_uses_minimal_prompt(self): agent_prompt="You are Rex Junior.", provider_id="anthropic", model_id="claude-sonnet", - tool_catalog_prompt_factory=lambda: "SHOULD_NOT_APPEAR", + turn_context=TurnPromptContext( + tool_catalog="SHOULD_NOT_APPEAR", + ), ) assert len(prompts) == 3 @@ -309,6 +312,38 @@ async def test_builtin_system_subagent_root_uses_full_prompt(self): assert len(prompts) > 2 assert any(PROMPT_DEFAULT.strip() in prompt for prompt in prompts) + @pytest.mark.asyncio + async def test_full_prompt_loads_worktree_and_config_instructions( + self, + tmp_path: Path, + ) -> None: + nested = tmp_path / "src" / "package" + nested.mkdir(parents=True) + (tmp_path / "AGENTS.md").write_text("project rules", encoding="utf-8") + (nested / "extra-rules.md").write_text("extra rules", encoding="utf-8") + + with patch.object( + SessionPrompt, + "_is_builtin_system_subagent_session", + AsyncMock(return_value=False), + ): + prompts = await SessionPrompt.build_system_prompts( + session_id="ses-instructions", + session_directory=str(nested), + agent_name="rex", + agent_prompt="agent prompt", + provider_id="openai", + model_id="gpt-5", + turn_context=TurnPromptContext( + worktree=str(tmp_path), + config_instructions=("extra-rules.md",), + ), + ) + + combined = "\n\n".join(prompts) + assert "project rules" in combined + assert "extra rules" in combined + @pytest.mark.asyncio async def test_evolution_subagent_child_uses_full_prompt(self): agent = AgentInfo( 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..5119d111d 100644 --- a/tests/session/test_runner_llm_hooks.py +++ b/tests/session/test_runner_llm_hooks.py @@ -1,4 +1,4 @@ -"""Tests for LLM lifecycle hooks in SessionRunner and HookPipeline.""" +"""Tests for LLM lifecycle hooks in StepEngine and HookPipeline.""" from __future__ import annotations @@ -8,11 +8,11 @@ import pytest -import flocks.session.runner as runner_mod +import flocks.session.runtime.step_engine as runner_mod from flocks.hooks.pipeline import HookBase, HookPipeline from flocks.provider.provider import ChatMessage from flocks.session.streaming.stream_processor import StreamProcessor -from flocks.session.runner import SessionRunner +from flocks.session.runtime.step_engine import StepEngine from flocks.session.session import SessionInfo from flocks.tool.registry import ToolResult @@ -27,8 +27,8 @@ def _make_session(session_id: str = "ses_runner_llm_hooks") -> SessionInfo: ) -def _make_runner(session_id: str = "ses_runner_llm_hooks") -> SessionRunner: - return SessionRunner( +def _make_runner(session_id: str = "ses_runner_llm_hooks") -> StepEngine: + return StepEngine( session=_make_session(session_id), provider_id="anthropic", model_id="claude-sonnet", @@ -39,6 +39,7 @@ class _FakeProcessor: def __init__(self, **_: object): self._text_parts: list[str] = [] self._reasoning_parts: list[str] = [] + self.reasoning_metadata: list[dict[str, object]] = [] self.finish_reason = "stop" self.tool_calls = {} self._langfuse_generation = None @@ -49,6 +50,7 @@ async def process_event(self, event) -> None: self._text_parts.append(event.text) elif event_name == "ReasoningDeltaEvent": self._reasoning_parts.append(event.text) + self.reasoning_metadata.append(event.metadata) elif event_name == "FinishEvent": self.finish_reason = event.finish_reason @@ -184,7 +186,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), ) @@ -417,6 +424,138 @@ async def _gen(): assert "[[V_EMAIL_1]]" in str(generation_inputs) +@pytest.mark.asyncio +async def test_call_llm_restores_stream_replacements_across_chunks_and_retries( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner = _make_runner("ses_runner_stream_replacements") + assistant_msg = SimpleNamespace(id="msg_assistant_stream_replacements") + agent = SimpleNamespace(name="rex") + processors: list[_FakeProcessor] = [] + + async def _before(payload): + return SimpleNamespace( + output={ + "request": { + **payload["request"], + "messages": [ + {"role": "user", "content": "email [[V_EMAIL_1]]"} + ], + "providerOptions": {}, + }, + "redaction": { + "streamTextReplacements": [ + { + "placeholder": "[[V_EMAIL_1]]", + "value": "alice@example.com", + } + ], + }, + } + ) + + class _RecordingProcessor(_FakeProcessor): + def __init__(self, **kwargs: object): + super().__init__(**kwargs) + processors.append(self) + + monkeypatch.setattr(runner_mod, "StreamProcessor", _RecordingProcessor) + monkeypatch.setattr( + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock( + side_effect=lambda stage, _metadata=None: ( + stage == runner_mod.HookStage.LLM_BEFORE + ) + ), + ) + run_before = AsyncMock(side_effect=_before) + monkeypatch.setattr(runner_mod.HookPipeline, "run_llm_before", run_before) + monkeypatch.setattr(runner_mod, "langfuse_is_active", lambda: False) + monkeypatch.setattr( + "flocks.provider.options.build_provider_options", + lambda provider_id, model_id: {}, + ) + monkeypatch.setattr( + "flocks.session.streaming.tool_accumulator.ToolCallAccumulator", + _FakeToolAccumulator, + ) + monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) + + class _Provider: + def chat_stream(self, **kwargs): + assert kwargs["messages"][0].content == "email [[V_EMAIL_1]]" + + async def _gen(): + yield SimpleNamespace( + delta="", + reasoning="Contact [[V_EM", + metadata={"reasoningContent": "Contact [[V_EMAIL_1]]"}, + event_type="reasoning", + tool_calls=None, + finish_reason=None, + usage=None, + ) + yield SimpleNamespace( + delta="", + reasoning="AIL_1]]", + metadata={}, + event_type="reasoning", + tool_calls=None, + finish_reason=None, + usage=None, + ) + yield SimpleNamespace( + delta="Reply to [[V_EM", + reasoning=None, + metadata={}, + event_type=None, + tool_calls=None, + finish_reason=None, + usage=None, + ) + yield SimpleNamespace( + delta="AIL_1]]", + reasoning=None, + metadata={}, + event_type=None, + tool_calls=None, + finish_reason="stop", + usage=None, + ) + + return _gen() + + results = [] + for _ in range(2): + results.append( + await runner._call_llm( + provider=_Provider(), + messages=[ + ChatMessage(role="user", content="email alice@example.com") + ], + tools=[], + agent=agent, + assistant_msg=assistant_msg, + ) + ) + + assert [result.content for result in results] == [ + "Reply to alice@example.com", + "Reply to alice@example.com", + ] + assert [processor.get_reasoning_content() for processor in processors] == [ + "Contact alice@example.com", + "Contact alice@example.com", + ] + assert all( + "alice@example.com" in str(processor.reasoning_metadata) + and "[[V_EMAIL_1]]" not in str(processor.reasoning_metadata) + for processor in processors + ) + run_before.assert_awaited_once() + + @pytest.mark.asyncio async def test_call_llm_emits_after_hook_on_error(monkeypatch: pytest.MonkeyPatch): runner = _make_runner("ses_runner_llm_hooks_error") @@ -452,7 +591,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), ) 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..da81ffe5e 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,13 +27,18 @@ ToolStateRunning, UserMessageInfo, ) -from flocks.session.runner import ( - RunnerCallbacks, - SessionRunner, +from flocks.session.runtime.step_engine import ( + StepEngine, StepResult, ToolCall, ) -from flocks.session.prompt import SessionPrompt, get_prompt_flocks_config_guard +from flocks.session.runtime.session_turn import LoopCallbacks, LoopContext +from flocks.session.prompt import ( + SessionPrompt, + SystemPromptBlock, + TurnPromptContext, + get_prompt_flocks_config_guard, +) from flocks.session.core.defaults import DEFAULT_MAX_TOOL_STEPS from flocks.session.session import Session, SessionInfo from flocks.tool.registry import ToolCategory, ToolInfo @@ -62,7 +67,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 +172,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 +192,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 +387,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 +452,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 +474,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 +495,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 +529,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 +554,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 +567,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 +590,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 +617,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,22 +638,29 @@ 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" env_mock = MagicMock(return_value=["env prompt"]) runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_mock = AsyncMock(return_value="device prompt") + turn_context = TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog", + device_asset_hint="device prompt", + tool_revision=1, + device_revision=7, + ) - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -657,13 +669,8 @@ async def test_build_system_prompts_reuses_loop_static_cache(self): provider_id=runner1.provider_id, model_id=runner1.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, - device_revision=7, + turn_context=turn_context, ) prompts2 = await SessionPrompt.build_system_prompts( session_id=session.id, @@ -673,27 +680,19 @@ async def test_build_system_prompts_reuses_loop_static_cache(self): provider_id=runner2.provider_id, model_id=runner2.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, - device_revision=7, + turn_context=turn_context, ) assert prompts1 == prompts2 env_mock.assert_called_once() runtime_mock.assert_called_once() custom_mock.assert_awaited_once() - sandbox_mock.assert_awaited_once() - channel_mock.assert_awaited_once() - device_mock.assert_awaited_once() @pytest.mark.asyncio async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(self): 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 = { @@ -704,15 +703,13 @@ async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(sel "inject": True, }, } - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_mock = AsyncMock(return_value="device prompt") - - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch.object(SessionPrompt, "_build_tool_guidance_prompt", return_value="tool protocol"), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", return_value=["env prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", return_value=["runtime prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.custom", AsyncMock(return_value=["custom prompt"])): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch.object(SessionPrompt, "_build_tool_guidance_prompt", return_value="tool protocol"), + patch("flocks.session.prompt.SystemPrompt.environment_stable", return_value=["env prompt"]), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", return_value=["runtime prompt"]), + patch("flocks.session.prompt.SystemPrompt.custom", AsyncMock(return_value=["custom prompt"])), + ): prompts = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -730,11 +727,17 @@ async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(sel "write", ), memory_bootstrap_data=memory_bootstrap_data, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, - device_revision=3, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, + turn_context=TurnPromptContext( + tool_catalog="tool catalog", + device_asset_hint="device prompt", + sandbox_context="sandbox prompt", + channel_context="channel prompt", + additional_context="additional prompt", + text_tool_catalog="text tool catalog", + tool_results_reminder="tool results reminder", + repeated_tool_calls_reminder="tool loop reminder", + device_revision=3, + ), ) assert prompts == [ @@ -751,29 +754,29 @@ async def test_build_system_prompts_orders_stable_prefix_before_runtime_tail(sel "sandbox prompt", "channel prompt", "runtime prompt", + "additional prompt", + "text tool catalog", + "tool results reminder", + "tool loop reminder", ] @pytest.mark.asyncio 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" env_mock = MagicMock(return_value=["env prompt"]) runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_mock = AsyncMock(return_value="device prompt") - - catalog_prompts = iter(["tool catalog v1", "tool catalog v2"]) - - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -782,13 +785,15 @@ async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: next(catalog_prompts), - device_asset_prompt_factory=device_mock, - device_revision=1, + turn_context=TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog v1", + device_asset_hint="device prompt", + tool_revision=1, + device_revision=1, + ), ) agent.prompt = "agent prompt v2" prompts2 = await SessionPrompt.build_system_prompts( @@ -799,13 +804,15 @@ async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=2, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: next(catalog_prompts), - device_asset_prompt_factory=device_mock, - device_revision=1, + turn_context=TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog v2", + device_asset_hint="device prompt", + tool_revision=2, + device_revision=1, + ), ) assert prompts1 != prompts2 @@ -816,28 +823,32 @@ async def test_build_system_prompts_rebuilds_when_tool_revision_changes(self): env_mock.assert_called_once() runtime_mock.assert_called_once() custom_mock.assert_awaited_once() - sandbox_mock.assert_awaited_once() - channel_mock.assert_awaited_once() @pytest.mark.asyncio async def test_build_system_prompts_reuses_static_device_hint_cache(self): 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" env_mock = MagicMock(return_value=["env prompt"]) runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_mock = AsyncMock(return_value="device prompt") + turn_context = TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog", + device_asset_hint="device prompt", + tool_revision=1, + ) - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -846,12 +857,8 @@ async def test_build_system_prompts_reuses_static_device_hint_cache(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, + turn_context=turn_context, ) prompts2 = await SessionPrompt.build_system_prompts( session_id=session.id, @@ -861,12 +868,8 @@ async def test_build_system_prompts_reuses_static_device_hint_cache(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=device_mock, + turn_context=turn_context, ) assert prompts1 == prompts2 @@ -874,29 +877,24 @@ async def test_build_system_prompts_reuses_static_device_hint_cache(self): env_mock.assert_called_once() runtime_mock.assert_called_once() custom_mock.assert_awaited_once() - sandbox_mock.assert_awaited_once() - channel_mock.assert_awaited_once() - device_mock.assert_awaited_once() @pytest.mark.asyncio async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): 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" env_mock = MagicMock(return_value=["env prompt"]) runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - sandbox_mock = AsyncMock(return_value="sandbox prompt") - channel_mock = AsyncMock(return_value="channel prompt") - device_prompts = iter(["device prompt v1", "device prompt v2"]) - - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -905,13 +903,15 @@ async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=AsyncMock(side_effect=lambda: next(device_prompts)), - device_revision=1, + turn_context=TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog", + device_asset_hint="device prompt v1", + tool_revision=1, + device_revision=1, + ), ) prompts2 = await SessionPrompt.build_system_prompts( session_id=session.id, @@ -921,13 +921,15 @@ async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, - sandbox_prompt_factory=sandbox_mock, - channel_context_prompt_factory=channel_mock, - tool_catalog_prompt_factory=lambda: "tool catalog", - device_asset_prompt_factory=AsyncMock(side_effect=lambda: next(device_prompts)), - device_revision=2, + turn_context=TurnPromptContext( + sandbox_context="sandbox prompt", + channel_context="channel prompt", + tool_catalog="tool catalog", + device_asset_hint="device prompt v2", + tool_revision=1, + device_revision=2, + ), ) assert prompts1 != prompts2 @@ -936,14 +938,12 @@ async def test_build_system_prompts_rebuilds_when_device_revision_changes(self): env_mock.assert_called_once() runtime_mock.assert_called_once() custom_mock.assert_awaited_once() - sandbox_mock.assert_awaited_once() - channel_mock.assert_awaited_once() @pytest.mark.asyncio async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): 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" @@ -951,10 +951,12 @@ async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts1 = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -963,8 +965,8 @@ async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, + turn_context=TurnPromptContext(tool_revision=1), ) agent.prompt = "agent prompt v2" prompts2 = await SessionPrompt.build_system_prompts( @@ -975,8 +977,8 @@ async def test_build_system_prompts_rebuilds_when_agent_prompt_changes(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, static_cache=shared_cache, + turn_context=TurnPromptContext(tool_revision=1), ) assert prompts1 != prompts2 @@ -989,7 +991,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 +1035,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 +1060,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 +1096,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={ @@ -1109,10 +1111,12 @@ async def test_filesystem_memory_guidance_depends_on_tool_names(self): runtime_mock = MagicMock(return_value=["runtime prompt"]) custom_mock = AsyncMock(return_value=["custom prompt"]) - with patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), \ - patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), \ - patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), \ - patch("flocks.session.prompt.SystemPrompt.custom", custom_mock): + with ( + patch("flocks.session.prompt.SystemPrompt.provider", return_value=["provider prompt"]), + patch("flocks.session.prompt.SystemPrompt.environment_stable", env_mock), + patch("flocks.session.prompt.SystemPrompt.runtime_metadata", runtime_mock), + patch("flocks.session.prompt.SystemPrompt.custom", custom_mock), + ): prompts_with_memory = await SessionPrompt.build_system_prompts( session_id=session.id, session_directory=session.directory, @@ -1128,9 +1132,9 @@ async def test_filesystem_memory_guidance_depends_on_tool_names(self): "read", "write", ), - tool_revision=1, memory_bootstrap_data=runner._memory_bootstrap_data, static_cache=shared_cache, + turn_context=TurnPromptContext(tool_revision=1), ) prompts_without_memory = await SessionPrompt.build_system_prompts( session_id=session.id, @@ -1140,9 +1144,9 @@ async def test_filesystem_memory_guidance_depends_on_tool_names(self): provider_id=runner.provider_id, model_id=runner.model_id, prompt_tool_names=("read",), - tool_revision=1, memory_bootstrap_data=runner._memory_bootstrap_data, static_cache=shared_cache, + turn_context=TurnPromptContext(tool_revision=1), ) assert prompts_with_memory != prompts_without_memory @@ -1158,7 +1162,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 +1174,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 +1209,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 +1253,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 +1293,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 +1311,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 +1328,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 +1344,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 +1353,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 +1362,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 +1371,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 +1380,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 +1391,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 +1400,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 +1410,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 +1437,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 +1470,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", @@ -1476,20 +1480,34 @@ async def test_to_chat_messages_uses_structured_anthropic_system_blocks(monkeypa monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hello")) - chat_messages = await runner._to_chat_messages( - [message], - ["provider prompt", "agent prompt", "context prompt", "runtime prompt"], - ) + prompt_blocks = [ + SystemPromptBlock( + name=name, + content=content, + cache_scope=cache_scope, + ) + for name, content, cache_scope in ( + ("provider", "provider prompt", "global"), + ("agent", "agent prompt", "agent"), + ("context", "context prompt", "workspace"), + ("sandbox", "sandbox prompt", "runtime_tail"), + ("runtime", "runtime prompt", "runtime_tail"), + ("reminder", "reminder prompt", "runtime_tail"), + ) + ] + + chat_messages = await runner._to_chat_messages([message], prompt_blocks) assert chat_messages[0].role == "system" assert isinstance(chat_messages[0].content, list) - assert chat_messages[0].content[1]["cache_control"] == {"type": "ephemeral"} - assert chat_messages[0].content[-1]["text"] == "runtime prompt" + assert chat_messages[0].content[2]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in chat_messages[0].content[3] + assert chat_messages[0].content[-1]["text"] == "reminder prompt" @pytest.mark.asyncio 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 +1537,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], []) @@ -1534,7 +1552,7 @@ async def test_to_chat_messages_invalidates_shared_cache_when_message_parts_chan sessionID=session.id, messageID=assistant_message.id, callID="call_cache_fix", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1547,7 +1565,7 @@ async def test_to_chat_messages_invalidates_shared_cache_when_message_parts_chan assert len(second_messages) == 2 assert second_messages[0].role == "assistant" assert second_messages[0].tool_calls is not None - assert second_messages[0].tool_calls[0]["function"]["name"] == "task" + assert second_messages[0].tool_calls[0]["function"]["name"] == "delegate_task" assert second_messages[1].role == "tool" assert second_messages[1].tool_call_id == "call_cache_fix" assert second_messages[1].content == "Error: Tool execution was interrupted" @@ -1567,7 +1585,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 +1603,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, @@ -1604,7 +1622,7 @@ async def test_to_chat_messages_preserves_assistant_reasoning_for_replay(): sessionID=session.id, messageID=assistant_message.id, callID="call_reasoning_replay", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1618,7 +1636,7 @@ async def test_to_chat_messages_preserves_assistant_reasoning_for_replay(): assert chat_messages[0].role == "assistant" assert chat_messages[0].reasoning == "Need to call the tool first." assert chat_messages[0].tool_calls is not None - assert chat_messages[0].tool_calls[0]["function"]["name"] == "task" + assert chat_messages[0].tool_calls[0]["function"]["name"] == "delegate_task" assert chat_messages[1].role == "tool" assert chat_messages[1].tool_call_id == "call_reasoning_replay" @@ -1634,7 +1652,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" @@ -1673,7 +1691,7 @@ async def test_to_chat_messages_restores_provider_reasoning_fields_from_metadata sessionID=session.id, messageID=assistant_message.id, callID="call_reasoning_metadata", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1687,7 +1705,7 @@ async def test_to_chat_messages_restores_provider_reasoning_fields_from_metadata assert chat_messages[0].reasoning == "Need to call the tool first." assert chat_messages[0].reasoning_content == "Need to call the tool first." assert chat_messages[0].reasoning_source == "native_reasoning_content" - assert chat_messages[0].tool_calls[0]["function"]["name"] == "task" + assert chat_messages[0].tool_calls[0]["function"]["name"] == "delegate_task" @pytest.mark.asyncio @@ -1701,7 +1719,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" @@ -1735,7 +1753,7 @@ async def test_to_chat_messages_restores_redacted_anthropic_thinking_blocks(monk sessionID=session.id, messageID=assistant_message.id, callID="call_redacted_reasoning", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1749,7 +1767,7 @@ async def test_to_chat_messages_restores_redacted_anthropic_thinking_blocks(monk assert chat_messages[0].custom_settings["anthropic_thinking_blocks"] == [ {"type": "redacted_thinking", "data": "opaque_blob"} ] - assert chat_messages[0].tool_calls[0]["function"]["name"] == "task" + assert chat_messages[0].tool_calls[0]["function"]["name"] == "delegate_task" @pytest.mark.asyncio @@ -1763,7 +1781,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" @@ -1797,7 +1815,7 @@ async def test_to_chat_messages_restores_signed_anthropic_thinking_blocks(monkey sessionID=session.id, messageID=assistant_message.id, callID="call_signed_reasoning", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1815,7 +1833,7 @@ async def test_to_chat_messages_restores_signed_anthropic_thinking_blocks(monkey "signature": "sig123", } ] - assert chat_messages[0].tool_calls[0]["function"]["name"] == "task" + assert chat_messages[0].tool_calls[0]["function"]["name"] == "delegate_task" @pytest.mark.asyncio @@ -1829,7 +1847,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" @@ -1862,7 +1880,7 @@ async def test_to_chat_messages_restores_unsigned_anthropic_thinking_blocks(monk sessionID=session.id, messageID=assistant_message.id, callID="call_unsigned_reasoning", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1879,7 +1897,7 @@ async def test_to_chat_messages_restores_unsigned_anthropic_thinking_blocks(monk "thinking": "Unsigned plan before tool use.", } ] - assert chat_messages[0].tool_calls[0]["function"]["name"] == "task" + assert chat_messages[0].tool_calls[0]["function"]["name"] == "delegate_task" @pytest.mark.asyncio @@ -1893,7 +1911,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" @@ -1927,7 +1945,7 @@ async def test_runner_history_round_trip_formats_anthropic_payload(monkeypatch): sessionID=session.id, messageID=assistant_message.id, callID="call_signed_reasoning", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -1958,7 +1976,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" @@ -2011,7 +2029,7 @@ async def test_to_chat_messages_prefers_provider_specific_interleaved_resolution sessionID=session.id, messageID=assistant_message.id, callID="call_provider_specific_interleaved", - tool="task", + tool="delegate_task", state=ToolStateRunning( input={"prompt": "continue"}, time={"start": 1}, @@ -2038,7 +2056,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 +2131,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 +2212,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 +2238,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(), ) @@ -2246,7 +2264,7 @@ async def fake_create(*args, **kwargs): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2274,7 +2292,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( @@ -2304,7 +2322,7 @@ async def fake_to_chat_messages(_messages, _system_prompts): # noqa: ANN001 monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_to_chat_messages", fake_to_chat_messages) monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="queued")) @@ -2325,7 +2343,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", @@ -2352,7 +2370,7 @@ async def fake_call_llm(*_args, **_kwargs): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2372,7 +2390,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 +2416,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", @@ -2428,7 +2446,7 @@ async def fake_call_llm(*_args, **_kwargs): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2461,7 +2479,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 +2509,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 +2548,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 +2573,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 +2640,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 +2655,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 @@ -2661,13 +2692,13 @@ async def publish_event(event_name, payload): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", staticmethod(lambda _provider_id: EmptyProvider())) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) - monkeypatch.setattr(SessionRunner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) + monkeypatch.setattr(StepEngine, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.SessionRetry, "sleep", AsyncMock(return_value=None)) 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 +2724,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", @@ -2708,7 +2739,7 @@ async def test_process_step_uses_loaded_tool_schema_names_for_prompt_guidance(mo provider = MagicMock() provider.is_configured.return_value = True assistant_msg = SimpleNamespace(id="msg_assistant_prompt_guidance") - build_system_prompts = AsyncMock(return_value=[]) + build_system_prompt_blocks = AsyncMock(return_value=[]) tool_schema = [ {"type": "function", "function": {"name": "memory_search", "description": "", "parameters": {}}}, {"type": "function", "function": {"name": "bash", "description": "", "parameters": {}}}, @@ -2717,7 +2748,7 @@ async def test_process_step_uses_loaded_tool_schema_names_for_prompt_guidance(mo monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", build_system_prompts) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", build_system_prompt_blocks) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=tool_schema)) monkeypatch.setattr( runner, @@ -2737,14 +2768,14 @@ async def test_process_step_uses_loaded_tool_schema_names_for_prompt_guidance(mo result = await runner._process_step([last_user], last_user) assert result.content == "done" - build_system_prompts.assert_awaited_once() - assert build_system_prompts.await_args.kwargs["prompt_tool_names"] == ("bash", "memory_search") + build_system_prompt_blocks.assert_awaited_once() + assert build_system_prompt_blocks.await_args.kwargs["prompt_tool_names"] == ("bash", "memory_search") @pytest.mark.asyncio 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", @@ -2766,7 +2797,7 @@ async def test_process_step_records_usage_after_success(monkeypatch): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2790,57 +2821,6 @@ async def test_process_step_records_usage_after_success(monkeypatch): update_mock.assert_any_await(runner.session.id, assistant_msg.id, finish="stop") -@pytest.mark.asyncio -async def test_process_step_passes_device_hint_factory_into_build_system_prompts(monkeypatch): - runner = _make_runner("ses_runner_device_hint_order") - runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) - - last_user = UserMessageInfo( - id="msg_user_device_hint_order", - sessionID=runner.session.id, - role="user", - time={"created": 1_000}, - agent="rex", - model={"providerID": "anthropic", "modelID": "claude-sonnet"}, - ) - - agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) - provider = MagicMock() - provider.is_configured.return_value = True - assistant_msg = SimpleNamespace(id="msg_assistant_device_hint_order") - build_system_prompts = AsyncMock(return_value=["provider", "tool catalog awareness", "device hint"]) - - monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) - monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) - monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", build_system_prompts) - monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) - device_hint_mock = AsyncMock(return_value="device hint") - monkeypatch.setattr(runner, "_build_device_asset_hint", device_hint_mock) - monkeypatch.setattr("flocks.tool.device.store.device_revision", lambda: 9) - monkeypatch.setattr( - runner, - "_to_chat_messages", - AsyncMock(return_value=[SimpleNamespace(role="user", content="hi")]), - ) - monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) - monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) - monkeypatch.setattr(runner_mod.Message, "create", AsyncMock(return_value=assistant_msg)) - monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) - monkeypatch.setattr( - runner, - "_call_llm", - AsyncMock(return_value=StepResult(action="stop", content="done")), - ) - - result = await runner._process_step([last_user], last_user) - - assert result.content == "done" - build_system_prompts.assert_awaited_once() - kwargs = build_system_prompts.await_args.kwargs - assert kwargs["device_revision"] == 9 - assert kwargs["device_asset_prompt_factory"] is not None - assert await kwargs["device_asset_prompt_factory"]() == "device hint" @pytest.mark.asyncio @@ -2848,7 +2828,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", @@ -2871,7 +2851,7 @@ async def test_process_step_empty_retry_records_usage_per_attempt(monkeypatch): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2908,7 +2888,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", @@ -2933,7 +2913,7 @@ async def test_process_step_retries_empty_transport_exception(monkeypatch): monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) monkeypatch.setattr( runner, @@ -2957,7 +2937,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", @@ -2985,7 +2965,7 @@ async def _call_llm(*_args, **_kwargs): monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) monkeypatch.setattr( runner_mod.SessionPrompt, - "build_system_prompts", + "build_system_prompt_blocks", AsyncMock(return_value=[]), ) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) @@ -3011,7 +2991,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( @@ -3033,7 +3013,7 @@ async def test_process_step_uses_default_max_steps_when_agent_steps_missing(monk monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=sentinel_tools)) monkeypatch.setattr( runner, @@ -3049,7 +3029,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 +3040,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( @@ -3082,7 +3062,7 @@ async def test_process_step_respects_explicit_agent_steps_over_default(monkeypat monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=sentinel_tools)) monkeypatch.setattr( runner, @@ -3098,7 +3078,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) @@ -3136,12 +3116,12 @@ async def fake_call_llm(self, provider, messages, tools, agent, assistant_msg): ))) monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) - monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) monkeypatch.setattr(runner_mod.Message, "create", create_mock) 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 +3133,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 fd72c0855..bc7693b67 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,14 +62,14 @@ def _make_completed_tool_part(message_id: str) -> ToolPart: # --------------------------------------------------------------------------- class TestAbortPropagation: - """Test that abort_event propagates from SessionLoop to SessionRunner.""" + """Test that abort_event propagates from SessionLoop to StepEngine.""" def test_runner_accepts_external_abort_event(self): - """SessionRunner should accept an optional external abort_event.""" + """StepEngine should accept an optional external abort_event.""" external_event = asyncio.Event() session_info = _make_session_info() - runner = SessionRunner( + runner = StepEngine( session=session_info, abort_event=external_event, ) @@ -76,9 +82,9 @@ def test_runner_accepts_external_abort_event(self): assert runner.is_aborted is True def test_runner_internal_abort_still_works(self): - """SessionRunner's own abort() method should still work.""" + """StepEngine's own abort() method should still work.""" session_info = _make_session_info() - runner = SessionRunner(session=session_info) + runner = StepEngine(session=session_info) assert runner.is_aborted is False runner.abort() @@ -89,7 +95,7 @@ def test_runner_either_abort_triggers(self): external_event = asyncio.Event() session_info = _make_session_info() - runner = SessionRunner( + runner = StepEngine( session=session_info, abort_event=external_event, ) @@ -112,7 +118,7 @@ def test_runner_either_abort_triggers(self): def test_runner_without_external_event(self): """Runner created without abort_event should still work normally.""" session_info = _make_session_info() - runner = SessionRunner(session=session_info) + runner = StepEngine(session=session_info) assert runner._external_abort is None assert runner.is_aborted is False @@ -131,12 +137,12 @@ async def test_session_loop_run_publishes_busy_and_idle_status_events(self): ), patch( "flocks.session.session_loop.Message.list", AsyncMock(return_value=[]), + ), patch( + "flocks.session.runtime.session_turn.Message.list", + AsyncMock(return_value=[]), ), patch( "flocks.session.orphan_tools.abort_orphan_running_parts", AsyncMock(return_value=0), - ), patch( - "flocks.session.session_loop.SessionLoop._run_loop", - AsyncMock(return_value=LoopResult(action="stop")), ), patch( "flocks.session.session_loop.Session.touch", AsyncMock(), @@ -187,7 +193,7 @@ def test_abort_running_session(self): ) # Register the context - SessionLoop._active_loops["test_loop_abort"] = ctx + SessionLoop._active_turns["test_loop_abort"] = ctx try: assert ctx.should_abort() is False @@ -196,10 +202,10 @@ def test_abort_running_session(self): assert ctx.should_abort() is True finally: # Clean up - SessionLoop._active_loops.pop("test_loop_abort", None) + SessionLoop._active_turns.pop("test_loop_abort", None) def test_is_running(self): - """is_running should reflect _active_loops state.""" + """is_running should reflect the runtime lease registry.""" assert SessionLoop.is_running("not_there") is False session_info = _make_session_info("running_test") @@ -209,12 +215,12 @@ def test_is_running(self): model_id="test", agent_name="test", ) - SessionLoop._active_loops["running_test"] = ctx + SessionLoop._active_turns["running_test"] = ctx try: assert SessionLoop.is_running("running_test") is True finally: - SessionLoop._active_loops.pop("running_test", None) + SessionLoop._active_turns.pop("running_test", None) def test_get_context(self): """get_context should return the LoopContext for a running session.""" @@ -225,14 +231,14 @@ def test_get_context(self): model_id="test", agent_name="test", ) - SessionLoop._active_loops["ctx_get_test"] = ctx + SessionLoop._active_turns["ctx_get_test"] = ctx try: retrieved = SessionLoop.get_context("ctx_get_test") assert retrieved is ctx assert SessionLoop.get_context("nonexistent") is None finally: - SessionLoop._active_loops.pop("ctx_get_test", None) + SessionLoop._active_turns.pop("ctx_get_test", None) # --------------------------------------------------------------------------- @@ -257,7 +263,7 @@ def test_exit_when_assistant_after_user_and_finished(self): last_assistant = self._make_msg("msg_002", "assistant", finish="stop") # assistant.id > user.id → user.id < assistant.id → True → should exit - assert SessionLoop._should_exit(last_user, last_assistant) is True + assert RuntimeLoopContext._should_exit(last_user, last_assistant) is True def test_no_exit_when_user_injected_after_assistant(self): """Should NOT exit when a new user message appears after the assistant. @@ -269,34 +275,34 @@ def test_no_exit_when_user_injected_after_assistant(self): last_assistant = self._make_msg("msg_002", "assistant", finish="stop") # user.id > assistant.id → user.id < assistant.id → False → don't exit - assert SessionLoop._should_exit(last_user, last_assistant) is False + assert 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.""" @@ -304,7 +310,7 @@ def test_no_exit_when_assistant_has_completed_tool_parts(self): last_assistant = self._make_msg("msg_002", "assistant", finish="stop") last_assistant_parts = [_make_completed_tool_part(last_assistant.id)] - assert SessionLoop._should_exit( + assert RuntimeLoopContext._should_exit( last_user, last_assistant, last_assistant_parts, @@ -323,7 +329,7 @@ def _make_msg(msg_id: str, role: str): async def test_does_not_treat_current_user_as_queued_when_no_assistant_exists(self): current_user = self._make_msg("msg_001", "user") - queued = await SessionLoop._detect_queued_user_message( + queued = await DEFAULT_CONTINUATION_POLICY.detect_queued_user_message( "session-1", [current_user], current_user.id, @@ -337,7 +343,7 @@ async def test_detects_newer_user_when_step_failed_before_assistant_created(self current_user = self._make_msg("msg_001", "user") newer_user = self._make_msg("msg_002", "user") - queued = await SessionLoop._detect_queued_user_message( + queued = await DEFAULT_CONTINUATION_POLICY.detect_queued_user_message( "session-1", [current_user, newer_user], current_user.id, @@ -372,11 +378,11 @@ async def test_run_loop_stops_turn_when_messages_are_empty(self): model_id="test-model", agent_name="rex", ) - ctx.session_ctx = SimpleNamespace(get_messages=AsyncMock(return_value=[])) + ctx.session_store = SimpleNamespace(get_messages=AsyncMock(return_value=[])) event_callback = AsyncMock() callbacks = LoopCallbacks(event_publish_callback=event_callback) - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" event_names = [call.args[0] for call in event_callback.await_args_list] @@ -399,11 +405,11 @@ async def test_run_loop_stops_turn_when_no_user_message_exists(self): agent_name="rex", ) assistant = self._make_msg("msg_001", "assistant", finish="stop") - ctx.session_ctx = SimpleNamespace(get_messages=AsyncMock(return_value=[assistant])) + ctx.session_store = SimpleNamespace(get_messages=AsyncMock(return_value=[assistant])) event_callback = AsyncMock() callbacks = LoopCallbacks(event_publish_callback=event_callback) - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" event_names = [call.args[0] for call in event_callback.await_args_list] @@ -429,12 +435,16 @@ async def test_run_loop_continues_for_active_goal_after_stop(self): assistant = self._make_msg("msg_002", "assistant", finish="stop") goal_user = self._make_msg("msg_003", "user") assistant_after_goal = self._make_msg("msg_004", "assistant", finish="stop") - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[ [user], [user, assistant], + [user, assistant], + [user, assistant], [user, assistant, goal_user], [user, assistant, goal_user, assistant_after_goal], + [user, assistant, goal_user, assistant_after_goal], + [user, assistant, goal_user, assistant_after_goal], ]) ) event_callback = AsyncMock() @@ -451,25 +461,25 @@ async def test_run_loop_continues_for_active_goal_after_stop(self): ] with patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", MagicMock(return_value="still working"), ), patch( - "flocks.session.session_loop.Message.create", + "flocks.session.runtime.session_turn.Message.create", AsyncMock(return_value=goal_user), ) as create_message, patch( - "flocks.session.session_loop.GoalManager.evaluate_after_turn", + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", AsyncMock(side_effect=goal_decisions), ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(side_effect=[StepResult(action="stop"), StepResult(action="stop")]), ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" create_message.assert_awaited_once() @@ -503,7 +513,7 @@ async def test_run_loop_waits_for_user_input_after_goal_clarification(self): ) user = self._make_msg("msg_001", "user") assistant = self._make_msg("msg_002", "assistant", finish="stop") - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[ [user], [user, assistant], @@ -513,19 +523,19 @@ async def test_run_loop_waits_for_user_input_after_goal_clarification(self): callbacks = LoopCallbacks(event_publish_callback=event_callback) with patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", MagicMock(return_value="Please clarify what tests to write."), ), patch( - "flocks.session.session_loop.Message.create", + "flocks.session.runtime.session_turn.Message.create", AsyncMock(), ) as create_message, patch( - "flocks.session.session_loop.GoalManager.evaluate_after_turn", + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", AsyncMock(return_value=GoalDecision( status="active", verdict="waiting", @@ -533,10 +543,10 @@ async def test_run_loop_waits_for_user_input_after_goal_clarification(self): reason="waiting for user clarification", )), ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(return_value=StepResult(action="stop")), ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" create_message.assert_not_awaited() @@ -559,7 +569,7 @@ async def test_run_loop_passes_pending_question_to_goal_judge(self): ) user = self._make_msg("msg_001", "user") assistant = self._make_msg("msg_002", "assistant", finish="stop") - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[ [user], [user, assistant], @@ -575,28 +585,28 @@ async def test_run_loop_passes_pending_question_to_goal_judge(self): )) with patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", MagicMock(return_value="Please provide the input."), ), patch( "flocks.server.routes.question.has_pending_questions", MagicMock(return_value=True), ), patch( - "flocks.session.session_loop.Message.create", + "flocks.session.runtime.session_turn.Message.create", AsyncMock(), ) as create_message, patch( - "flocks.session.session_loop.GoalManager.evaluate_after_turn", + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", evaluate_goal, ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(return_value=StepResult(action="stop")), ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" create_message.assert_not_awaited() @@ -622,23 +632,23 @@ async def test_run_loop_publishes_goal_terminal_status(self): self._make_msg("msg_001", "user"), self._make_msg("msg_002", "assistant", finish="stop"), ] - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[[messages[0]], messages]) ) event_callback = AsyncMock() callbacks = LoopCallbacks(event_publish_callback=event_callback) with patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.Message.get_text_content", + "flocks.session.runtime.session_turn.Message.get_text_content", MagicMock(return_value="Goal complete: done"), ), patch( - "flocks.session.session_loop.GoalManager.evaluate_after_turn", + "flocks.session.runtime.continuation_policy.GoalManager.evaluate_after_turn", AsyncMock(return_value=GoalDecision( status="completed", verdict="complete", @@ -646,10 +656,10 @@ async def test_run_loop_publishes_goal_terminal_status(self): objective="finish work", )), ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", AsyncMock(return_value=StepResult(action="stop")), ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" event_names = [call.args[0] for call in event_callback.await_args_list] @@ -694,26 +704,26 @@ async def test_pre_compact_cleanup_emits_turn_continued_before_next_iteration(se tokens={"input": 0, "output": 0, "cache": {"read": 0, "write": 0}}, ), ] - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[overflow_messages, normal_messages, normal_messages]) ) event_callback = AsyncMock() callbacks = LoopCallbacks(event_publish_callback=event_callback) with patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(20000, 1024, None), ), patch( - "flocks.session.session_loop.SessionCompaction.truncate_oversized_tool_outputs", + "flocks.session.runtime.session_turn.SessionCompaction.truncate_oversized_tool_outputs", AsyncMock(return_value=1), ), patch( - "flocks.session.session_loop.SessionPrompt.estimate_full_context_tokens", - 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] @@ -728,100 +738,6 @@ async def test_pre_compact_cleanup_emits_turn_continued_before_next_iteration(se assert cleanup_turn["continue_reason"] == "pre_compact_cleanup" assert cleanup_turn["status"] == "continued" - @pytest.mark.asyncio - async def test_post_observation_tool_delta_combines_with_observed_prompt(self): - session = SimpleNamespace( - id="turn_stale_usage_session", - agent="rex", - directory="/tmp", - memory_enabled=False, - ) - ctx = LoopContext( - session=session, - provider_id="test-provider", - 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 - - with patch( - "flocks.session.session_loop.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(), - ), patch( - "flocks.session.session_loop.run_compaction", - run_compaction, - ), patch( - "flocks.session.runner.SessionRunner._process_step", - AsyncMock(return_value=StepResult(action="stop")), - ): - 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, - messages, - ) - - 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 - @pytest.mark.asyncio async def test_run_loop_skips_exit_condition_when_assistant_has_tool_parts(self): session = SimpleNamespace( @@ -840,7 +756,7 @@ async def test_run_loop_skips_exit_condition_when_assistant_has_tool_parts(self) self._make_msg("msg_001", "user"), self._make_msg("msg_002", "assistant", finish="stop"), ] - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(side_effect=[messages, messages]) ) event_callback = AsyncMock() @@ -849,25 +765,25 @@ async def test_run_loop_skips_exit_condition_when_assistant_has_tool_parts(self) log_info = MagicMock() with patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[_make_completed_tool_part("msg_002")]), ), patch( - "flocks.session.session_loop.Provider.resolve_model_info", + "flocks.session.runtime.session_turn.Provider.resolve_model_info", return_value=(0, 0, None), ), patch( "flocks.session.lifecycle.title.SessionTitle.ensure_title", MagicMock(return_value=None), ), patch( - "flocks.session.session_loop.fire_and_forget", + "flocks.session.runtime.session_turn.fire_and_forget", MagicMock(), ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", process_step, ), patch( - "flocks.session.session_loop.log.info", + "flocks.session.runtime.session_turn.log.info", log_info, ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" assert result.last_message is messages[1] @@ -894,7 +810,7 @@ async def test_run_loop_breaks_on_exit_condition_without_tool_parts(self): self._make_msg("msg_001", "user"), self._make_msg("msg_002", "assistant", finish="stop"), ] - ctx.session_ctx = SimpleNamespace( + ctx.session_store = SimpleNamespace( get_messages=AsyncMock(return_value=messages) ) event_callback = AsyncMock() @@ -903,16 +819,16 @@ async def test_run_loop_breaks_on_exit_condition_without_tool_parts(self): log_info = MagicMock() with patch( - "flocks.session.session_loop.Message.parts", + "flocks.session.runtime.session_turn.Message.parts", AsyncMock(return_value=[]), ), patch( - "flocks.session.session_loop.log.info", + "flocks.session.runtime.session_turn.log.info", log_info, ), patch( - "flocks.session.runner.SessionRunner._process_step", + "flocks.session.runtime.step_engine.StepEngine._process_step", process_step, ): - result = await SessionLoop._run_loop(ctx, callbacks) + result = await run_logical_turns(ctx, callbacks) assert result.action == "stop" assert result.last_message is messages[1] @@ -922,60 +838,6 @@ async def test_run_loop_breaks_on_exit_condition_without_tool_parts(self): assert event_names == ["turn.started"] -class TestExecuteSubtask: - @pytest.mark.asyncio - async def test_execute_subtask_passes_tool_context_first(self): - session_info = _make_session_info("subtask_exec_test") - ctx = LoopContext( - session=session_info, - provider_id="test-provider", - model_id="test-model", - agent_name="rex", - ) - last_user = SimpleNamespace( - id="msg_parent", - agent="rex", - model={"providerID": "test-provider", "modelID": "test-model"}, - provider="test-provider", - ) - task_part = SimpleNamespace( - agent="helper", - prompt="do the thing", - description="test task", - command=None, - model=None, - ) - - task_tool = MagicMock() - task_tool.execute = AsyncMock(return_value=SimpleNamespace( - output="done", - title="task complete", - metadata={"sessionId": "child-session"}, - )) - - assistant_msg = SimpleNamespace(id="msg_assistant") - synthetic_msg = SimpleNamespace(id="msg_synthetic") - - with patch("flocks.agent.registry.Agent.get", AsyncMock(return_value=SimpleNamespace(name="helper"))), \ - patch("flocks.tool.registry.ToolRegistry.get", return_value=task_tool), \ - patch("flocks.session.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, - } - - # --------------------------------------------------------------------------- # LoopContext tests # --------------------------------------------------------------------------- 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..a9d898973 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 @@ -81,7 +81,7 @@ async def fake_get_prompt_tool_names(self, agent): # noqa: ANN001 del self, agent return () - async def fake_build_system_prompts(*args, **kwargs): # noqa: ANN002, ANN003 + async def fake_build_system_prompt_blocks(*args, **kwargs): # noqa: ANN002, ANN003 del args, kwargs return [] @@ -99,13 +99,17 @@ 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(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, "_get_prompt_tool_names", fake_get_prompt_tool_names) + monkeypatch.setattr( + SessionPrompt, + "build_system_prompt_blocks", + fake_build_system_prompt_blocks, + ) + monkeypatch.setattr(StepEngine, "_build_callable_tool_schema", fake_build_callable_tool_schema) + monkeypatch.setattr(StepEngine, "_to_chat_messages", fake_to_chat_messages) + monkeypatch.setattr(StepEngine, "_call_llm", fake_call_llm) - runner = SessionRunner(session=session, provider_id="test-provider", model_id="test-model", agent_name="rex") + runner = StepEngine(session=session, provider_id="test-provider", model_id="test-model", agent_name="rex") runner._step = 2 # ensure reminder wrapping branch doesn't break assumptions result = await runner._process_step(messages=messages, last_user=user_2) diff --git a/tests/session_runtime_testkit.py b/tests/session_runtime_testkit.py new file mode 100644 index 000000000..0587b3d17 --- /dev/null +++ b/tests/session_runtime_testkit.py @@ -0,0 +1,29 @@ +"""Test-only helpers for exercising SessionLoop logical-turn control.""" + +from flocks.session.session_loop import ( + LoopCallbacks, + LoopContext, + LoopResult, + SessionLoop, +) + + +async def run_logical_turns( + turn: LoopContext, + callbacks: LoopCallbacks, +) -> LoopResult: + """Run logical turns without acquiring a persisted session lease.""" + turn.callbacks = callbacks + policy = turn.continuation_policy or SessionLoop._continuation_policy + while True: + try: + await policy.prepare_logical_turn(turn) + outcome = await SessionLoop._run_logical_input(turn) + if await SessionLoop._should_continue(turn, policy, outcome): + continue + except Exception as exc: + outcome = await SessionLoop._handle_execution_error( + turn, + exc, + ) + return SessionLoop._to_loop_result(turn, outcome) diff --git a/tests/tool/test_builtin_management_tools.py b/tests/tool/test_builtin_management_tools.py index 9c649404a..ea830c20e 100644 --- a/tests/tool/test_builtin_management_tools.py +++ b/tests/tool/test_builtin_management_tools.py @@ -49,13 +49,13 @@ def test_lsp_remains_non_native_by_default() -> None: assert tool.info.native is False -def test_task_remains_non_native_when_declared() -> None: +def test_delegate_task_remains_native_when_declared() -> None: ToolRegistry.init() - tool = ToolRegistry.get("task") + tool = ToolRegistry.get("delegate_task") assert tool is not None - assert tool.info.native is False + assert tool.info.native is True def test_model_config_tools_remain_non_native_by_default() -> None: diff --git a/tests/tool/test_delegate_task_compat.py b/tests/tool/test_delegate_task_compat.py index 644718c10..43ceeea6b 100644 --- a/tests/tool/test_delegate_task_compat.py +++ b/tests/tool/test_delegate_task_compat.py @@ -20,9 +20,8 @@ def test_delegate_description_requires_material_delegation_benefit(self): assert "Do not delegate trivial edits" in DESCRIPTION assert "The task requires multiple steps or research" not in DESCRIPTION - @pytest.mark.parametrize("tool_name", ["delegate_task", "task"]) - def test_delegate_schema_exposes_only_subagent_routing(self, tool_name): - schema = ToolRegistry.get_schema(tool_name) + def test_delegate_schema_exposes_only_subagent_routing(self): + schema = ToolRegistry.get_schema("delegate_task") assert schema is not None assert "prompt" in schema.required assert "subagent_type" in schema.properties @@ -75,7 +74,6 @@ async def test_delegate_task_derives_description_and_ignores_blank_skills(self): permissions = create_session.await_args.kwargs["permission"] denied_permissions = {rule.permission for rule in permissions if rule.action == "deny"} assert "delegate_task" not in denied_permissions - assert "task" not in denied_permissions @pytest.mark.asyncio async def test_delegate_task_explicit_model_override_is_pinned(self): @@ -115,10 +113,9 @@ async def test_delegate_task_explicit_model_override_is_pinned(self): assert loop_run.await_args.kwargs["model_id"] == "claude-haiku-4-5" @pytest.mark.asyncio - @pytest.mark.parametrize("tool_name", ["delegate_task", "task"]) - async def test_delegate_tools_reject_removed_category_parameter(self, tool_name): + async def test_delegate_task_rejects_removed_category_parameter(self): result = await ToolRegistry.execute( - tool_name, + "delegate_task", ctx=_make_ctx(), category="quick", prompt="Summarize the diff", @@ -128,10 +125,9 @@ async def test_delegate_tools_reject_removed_category_parameter(self, tool_name) assert "unknown parameters: category" in (result.error or "") @pytest.mark.asyncio - @pytest.mark.parametrize("tool_name", ["delegate_task", "task"]) - async def test_delegate_tools_accept_deprecated_command_parameter(self, tool_name): + async def test_delegate_task_accepts_deprecated_command_parameter(self): result = await ToolRegistry.execute( - tool_name, + "delegate_task", ctx=_make_ctx(), command="legacy-tracking-command", prompt="Summarize the diff", diff --git a/tests/tool/test_task_model_pinning.py b/tests/tool/test_task_model_pinning.py deleted file mode 100644 index 4f75e10ab..000000000 --- a/tests/tool/test_task_model_pinning.py +++ /dev/null @@ -1,61 +0,0 @@ -from unittest.mock import AsyncMock, patch - -import pytest - -from flocks.tool.agent.task import task_tool -from flocks.tool.registry import ToolContext, ToolRegistry, ToolResult - - -def _make_ctx() -> ToolContext: - return ToolContext(session_id="test-session", message_id="test-message", agent="rex") - - -class TestTaskCompatibilityAlias: - def test_task_schema_does_not_expose_background_execution(self): - schema = ToolRegistry.get_schema("task") - assert schema is not None - assert "run_in_background" not in schema.properties - # Legacy batch shape is gone. - assert "tasks" not in schema.properties - - @pytest.mark.asyncio - async def test_task_tool_rejects_background_execution_when_called_directly(self): - result = await task_tool( - _make_ctx(), - description="delegate explore", - prompt="Inspect the repository", - subagent_type="explore", - run_in_background=True, - ) - - assert result.success is False - assert "Background subagent execution is disabled" in (result.error or "") - - @pytest.mark.asyncio - async def test_task_tool_forwards_single_call_to_delegate_task(self): - delegate_result = ToolResult( - success=True, - output="ok", - metadata={"sessionId": "ses-child"}, - ) - - with patch( - "flocks.tool.agent.task.delegate_task_tool", - AsyncMock(return_value=delegate_result), - ) as delegate: - result = await task_tool( - _make_ctx(), - description="delegate explore", - prompt="Inspect the repository", - subagent_type="explore", - model="openai/gpt-5", - ) - - assert result is delegate_result - delegate.assert_awaited_once() - kwargs = delegate.await_args.kwargs - assert kwargs["description"] == "delegate explore" - assert kwargs["prompt"] == "Inspect the repository" - assert kwargs["subagent_type"] == "explore" - assert kwargs["run_in_background"] is False - assert kwargs["model"] == "openai/gpt-5" diff --git a/tests/tool/test_tool_catalog.py b/tests/tool/test_tool_catalog.py index f9e440ad7..ded19df6e 100644 --- a/tests/tool/test_tool_catalog.py +++ b/tests/tool/test_tool_catalog.py @@ -78,12 +78,8 @@ def test_catalog_uses_real_builtin_tool_names_for_metadata_keys() -> None: assert name in TOOL_TAGS -def test_task_tool_tags_reflect_agent_delegation() -> None: - metadata = get_tool_catalog_metadata("task") - - assert "agent" in metadata.tags - assert "delegation" in metadata.tags - assert "planning" not in metadata.tags +def test_retired_task_tool_has_no_catalog_entry() -> None: + assert "task" not in TOOL_TAGS def test_schedule_task_and_todo_use_distinct_management_tags() -> None: diff --git a/tests/tool/test_tools.py b/tests/tool/test_tools.py index dfba07484..4ce0a37ee 100644 --- a/tests/tool/test_tools.py +++ b/tests/tool/test_tools.py @@ -156,7 +156,7 @@ def test_expected_tools_registered(self): # P1 tools "webfetch", "todo", "question", # P2 tools - "task", "lsp", "skill_load", + "delegate_task", "lsp", "skill_load", # P3 tools (2) "websearch", "apply_patch", ] @@ -801,13 +801,13 @@ async def test_webfetch_schema(self): # P2 Tools Tests # ============================================================================= -class TestTaskTool: - """Test the task tool""" +class TestDelegateTaskTool: + """Test the delegate_task tool""" @pytest.mark.asyncio - async def test_task_exists(self): - """Test that task tool is registered""" - tool = ToolRegistry.get("task") + async def test_delegate_task_exists(self): + """Test that delegate_task tool is registered""" + tool = ToolRegistry.get("delegate_task") assert tool is not None diff --git a/tests/utils/test_id_compatibility.py b/tests/utils/test_id_compatibility.py index 22b1691a2..c26d9c509 100644 --- a/tests/utils/test_id_compatibility.py +++ b/tests/utils/test_id_compatibility.py @@ -26,7 +26,6 @@ def test_prefix_mappings(self): "call": "cal", "step": "stp", "agent": "agt", - "subtask": "stk", "event": "evt", "tqref": "tqr", "task": "tsk", diff --git a/tui/flocks/cli/cmd/tui/routes/session/index.tsx b/tui/flocks/cli/cmd/tui/routes/session/index.tsx index bade13c9d..899a1e6b1 100644 --- a/tui/flocks/cli/cmd/tui/routes/session/index.tsx +++ b/tui/flocks/cli/cmd/tui/routes/session/index.tsx @@ -40,7 +40,7 @@ import type { GrepTool } from "@/tool/grep" import type { EditTool } from "@/tool/edit" import type { ApplyPatchTool } from "@/tool/apply_patch" import type { WebFetchTool } from "@/tool/webfetch" -import type { TaskTool } from "@/tool/task" +import type { DelegateTaskTool } from "@/tool/delegate-task" import type { QuestionTool } from "@/tool/question" import { useKeyboard, useRenderer, useTerminalDimensions, type JSX } from "@opentui/solid" import { useSDK } from "@tui/context/sdk" @@ -1861,7 +1861,7 @@ function SubagentActivity(props: { ) } -function Task(props: ToolProps) { +function Task(props: ToolProps) { const { theme } = useTheme() const keybind = useKeybind() const { navigate } = useRoute() @@ -1977,7 +1977,7 @@ function DelegateTask(props: ToolProps) { navigate({ type: "session", sessionID: sessionId()! }) : undefined} part={props.part} > @@ -1997,7 +1997,7 @@ function DelegateTask(props: ToolProps) { > - {delegateInput.description || "subtask"} + {delegateInput.description || "delegated task"} {isBackground() ? " (background)" : ""} {statusText()} @@ -2016,7 +2016,7 @@ function DelegateTask(props: ToolProps) { part={props.part} > {agentName()}{" "} - "{delegateInput.description || "subtask"}" + "{delegateInput.description || "delegated task"}" {isBackground() ? " (bg)" : ""} diff --git a/tui/flocks/command/index.ts b/tui/flocks/command/index.ts index 976f1cd51..56120ce30 100644 --- a/tui/flocks/command/index.ts +++ b/tui/flocks/command/index.ts @@ -30,7 +30,6 @@ export namespace Command { // workaround for zod not supporting async functions natively so we use getters // https://zod.dev/v4/changelog?id=zfunction template: z.promise(z.string()).or(z.string()), - subtask: z.boolean().optional(), hints: z.array(z.string()), }) .meta({ @@ -70,10 +69,10 @@ export namespace Command { [Default.REVIEW]: { name: Default.REVIEW, description: "review changes [commit|branch|pr], defaults to uncommitted", + agent: "oracle", get template() { return PROMPT_REVIEW.replace("${path}", Instance.worktree) }, - subtask: true, hints: hints(PROMPT_REVIEW), }, } @@ -87,7 +86,6 @@ export namespace Command { get template() { return command.template }, - subtask: command.subtask, hints: hints(command.template), } } diff --git a/tui/flocks/config/config.ts b/tui/flocks/config/config.ts index 5085330ed..10a8e4327 100644 --- a/tui/flocks/config/config.ts +++ b/tui/flocks/config/config.ts @@ -559,7 +559,6 @@ export namespace Config { description: z.string().optional(), agent: z.string().optional(), model: z.string().optional(), - subtask: z.boolean().optional(), }) export type Command = z.infer diff --git a/tui/flocks/session/message-v2.ts b/tui/flocks/session/message-v2.ts index f2f3331e5..ff596603b 100644 --- a/tui/flocks/session/message-v2.ts +++ b/tui/flocks/session/message-v2.ts @@ -163,21 +163,6 @@ export namespace MessageV2 { }) export type CompactionPart = z.infer - export const SubtaskPart = PartBase.extend({ - type: z.literal("subtask"), - prompt: z.string(), - description: z.string(), - agent: z.string(), - model: z - .object({ - providerID: z.string(), - modelID: z.string(), - }) - .optional(), - command: z.string().optional(), - }) - export type SubtaskPart = z.infer - export const RetryPart = PartBase.extend({ type: z.literal("retry"), attempt: z.number(), @@ -329,7 +314,6 @@ export namespace MessageV2 { export const Part = z .discriminatedUnion("type", [ TextPart, - SubtaskPart, ReasoningPart, FilePart, ToolPart, @@ -466,12 +450,6 @@ export namespace MessageV2 { text: "What did we do so far?", }) } - if (part.type === "subtask") { - userMessage.parts.push({ - type: "text", - text: "The following tool was executed by the user", - }) - } } } diff --git a/tui/flocks/session/prompt.ts b/tui/flocks/session/prompt.ts index befddd369..c0caae740 100644 --- a/tui/flocks/session/prompt.ts +++ b/tui/flocks/session/prompt.ts @@ -34,7 +34,6 @@ import { SessionSummary } from "./summary" import { NamedError } from "@flocks-ai/util/error" import { fn } from "@/util/fn" import { SessionProcessor } from "./processor" -import { TaskTool } from "@/tool/task" import { Tool } from "@/tool/tool" import { PermissionNext } from "@/permission/next" import { SessionStatus } from "./status" @@ -200,16 +199,6 @@ export namespace SessionPrompt { .meta({ ref: "AgentPartInput", }), - MessageV2.SubtaskPart.omit({ - messageID: true, - sessionID: true, - }) - .partial({ - id: true, - }) - .meta({ - ref: "SubtaskPartInput", - }), ]), ), }) @@ -344,7 +333,7 @@ export namespace SessionPrompt { let lastUser: MessageV2.User | undefined let lastAssistant: MessageV2.Assistant | undefined let lastFinished: MessageV2.Assistant | undefined - let tasks: (MessageV2.CompactionPart | MessageV2.SubtaskPart)[] = [] + const pendingCompactions: MessageV2.CompactionPart[] = [] for (let i = msgs.length - 1; i >= 0; i--) { const msg = msgs[i] if (!lastUser && msg.info.role === "user") lastUser = msg.info as MessageV2.User @@ -352,9 +341,9 @@ export namespace SessionPrompt { if (!lastFinished && msg.info.role === "assistant" && msg.info.finish) lastFinished = msg.info as MessageV2.Assistant if (lastUser && lastFinished) break - const task = msg.parts.filter((part) => part.type === "compaction" || part.type === "subtask") - if (task && !lastFinished) { - tasks.push(...task) + const compactions = msg.parts.filter((part) => part.type === "compaction") + if (compactions.length > 0 && !lastFinished) { + pendingCompactions.push(...compactions) } } @@ -378,183 +367,16 @@ export namespace SessionPrompt { }) const model = await Provider.getModel(lastUser.model.providerID, lastUser.model.modelID) - const task = tasks.pop() - - // pending subtask - // TODO: centralize "invoke tool" logic - if (task?.type === "subtask") { - const taskTool = await TaskTool.init() - const taskModel = task.model ? await Provider.getModel(task.model.providerID, task.model.modelID) : model - const assistantMessage = (await Session.updateMessage({ - id: Identifier.ascending("message"), - role: "assistant", - parentID: lastUser.id, - sessionID, - mode: task.agent, - agent: task.agent, - path: { - cwd: Instance.directory, - root: Instance.worktree, - }, - cost: 0, - tokens: { - input: 0, - output: 0, - reasoning: 0, - cache: { read: 0, write: 0 }, - }, - modelID: taskModel.id, - providerID: taskModel.providerID, - time: { - created: Date.now(), - }, - })) as MessageV2.Assistant - let part = (await Session.updatePart({ - id: Identifier.ascending("part"), - messageID: assistantMessage.id, - sessionID: assistantMessage.sessionID, - type: "tool", - callID: ulid(), - tool: TaskTool.id, - state: { - status: "running", - input: { - prompt: task.prompt, - description: task.description, - subagent_type: task.agent, - command: task.command, - }, - time: { - start: Date.now(), - }, - }, - })) as MessageV2.ToolPart - const taskArgs = { - prompt: task.prompt, - description: task.description, - subagent_type: task.agent, - command: task.command, - } - await Plugin.trigger( - "tool.execute.before", - { - tool: "task", - sessionID, - callID: part.id, - }, - { args: taskArgs }, - ) - let executionError: Error | undefined - const taskAgent = await Agent.get(task.agent) - const taskCtx: Tool.Context = { - agent: task.agent, - messageID: assistantMessage.id, - sessionID: sessionID, - abort, - callID: part.callID, - extra: { bypassAgentCheck: true }, - async metadata(input) { - await Session.updatePart({ - ...part, - type: "tool", - state: { - ...part.state, - ...input, - }, - } satisfies MessageV2.ToolPart) - }, - async ask(req) { - await PermissionNext.ask({ - ...req, - sessionID: sessionID, - ruleset: PermissionNext.merge(taskAgent.permission, session.permission ?? []), - }) - }, - } - const result = await taskTool.execute(taskArgs, taskCtx).catch((error) => { - executionError = error - log.error("subtask execution failed", { error, agent: task.agent, description: task.description }) - return undefined - }) - await Plugin.trigger( - "tool.execute.after", - { - tool: "task", - sessionID, - callID: part.id, - }, - result, - ) - assistantMessage.finish = "tool-calls" - assistantMessage.time.completed = Date.now() - await Session.updateMessage(assistantMessage) - if (result && part.state.status === "running") { - await Session.updatePart({ - ...part, - state: { - status: "completed", - input: part.state.input, - title: result.title, - metadata: result.metadata, - output: result.output, - attachments: result.attachments, - time: { - ...part.state.time, - end: Date.now(), - }, - }, - } satisfies MessageV2.ToolPart) - } - if (!result) { - await Session.updatePart({ - ...part, - state: { - status: "error", - error: executionError ? `Tool execution failed: ${executionError.message}` : "Tool execution failed", - time: { - start: part.state.status === "running" ? part.state.time.start : Date.now(), - end: Date.now(), - }, - metadata: part.metadata, - input: part.state.input, - }, - } satisfies MessageV2.ToolPart) - } - - // Add synthetic user message to prevent certain reasoning models from erroring - // If we create assistant messages w/ out user ones following mid loop thinking signatures - // will be missing and it can cause errors for models like gemini for example - const summaryUserMsg: MessageV2.User = { - id: Identifier.ascending("message"), - sessionID, - role: "user", - time: { - created: Date.now(), - }, - agent: lastUser.agent, - model: lastUser.model, - } - await Session.updateMessage(summaryUserMsg) - await Session.updatePart({ - id: Identifier.ascending("part"), - messageID: summaryUserMsg.id, - sessionID, - type: "text", - text: "Summarize the task tool output above and continue with your task.", - synthetic: true, - } satisfies MessageV2.TextPart) - - continue - } + const pendingCompaction = pendingCompactions.pop() // pending compaction - if (task?.type === "compaction") { + if (pendingCompaction) { const result = await SessionCompaction.process({ messages: msgs, parentID: lastUser.id, abort, sessionID, - auto: task.auto, + auto: pendingCompaction.auto, }) if (result === "stop") break continue @@ -1742,30 +1564,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the } const templateParts = await resolvePromptParts(template) - const isSubtask = (agent.mode === "subagent" && command.subtask !== false) || command.subtask === true - const parts = isSubtask - ? [ - { - type: "subtask" as const, - agent: agent.name, - description: command.description ?? "", - command: input.command, - model: { - providerID: taskModel.providerID, - modelID: taskModel.modelID, - }, - // TODO: how can we make task tool accept a more complex input? - prompt: templateParts.find((y) => y.type === "text")?.text ?? "", - }, - ] - : [...templateParts, ...(input.parts ?? [])] - - const userAgent = isSubtask ? (input.agent ?? (await Agent.defaultAgent())) : agentName - const userModel = isSubtask - ? input.model - ? Provider.parseModel(input.model) - : await lastModel(input.sessionID) - : taskModel + const parts = [...templateParts, ...(input.parts ?? [])] await Plugin.trigger( "command.execute.before", @@ -1780,8 +1579,8 @@ NOTE: At any point in time through this workflow you should feel free to ask the const result = (await prompt({ sessionID: input.sessionID, messageID: input.messageID, - model: userModel, - agent: userAgent, + model: taskModel, + agent: agentName, parts, variant: input.variant, })) as MessageV2.WithParts @@ -1817,15 +1616,9 @@ NOTE: At any point in time through this workflow you should feel free to ask the if (!isFirst) return // Gather all messages up to and including the first real user message for context - // This includes any shell/subtask executions that preceded the user's first prompt const contextMessages = input.history.slice(0, firstRealUserIdx + 1) const firstRealUser = contextMessages[firstRealUserIdx] - // For subtask-only messages (from command invocations), extract the prompt directly - // since toModelMessage converts subtask parts to generic "The following tool was executed by the user" - const subtaskParts = firstRealUser.parts.filter((p) => p.type === "subtask") as MessageV2.SubtaskPart[] - const hasOnlySubtaskParts = subtaskParts.length > 0 && firstRealUser.parts.every((p) => p.type === "subtask") - const agent = await Agent.get("title") if (!agent) return const result = await LLM.stream({ @@ -1848,9 +1641,7 @@ NOTE: At any point in time through this workflow you should feel free to ask the role: "user", content: "Generate a title for this conversation:\n", }, - ...(hasOnlySubtaskParts - ? [{ role: "user" as const, content: subtaskParts.map((p) => p.prompt).join("\n") }] - : MessageV2.toModelMessage(contextMessages)), + ...MessageV2.toModelMessage(contextMessages), ], }) const text = await result.text.catch((err) => log.error("failed to generate title", { error: err })) diff --git a/tui/flocks/session/prompt/anthropic-20250930.txt b/tui/flocks/session/prompt/anthropic-20250930.txt index 676c4d8dc..01ac8e5b5 100644 --- a/tui/flocks/session/prompt/anthropic-20250930.txt +++ b/tui/flocks/session/prompt/anthropic-20250930.txt @@ -129,10 +129,10 @@ The user will primarily request you perform software engineering tasks. This inc # Tool usage policy -- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. diff --git a/tui/flocks/session/prompt/anthropic.txt b/tui/flocks/session/prompt/anthropic.txt index 7a0e5fd5c..0709a5c30 100644 --- a/tui/flocks/session/prompt/anthropic.txt +++ b/tui/flocks/session/prompt/anthropic.txt @@ -73,20 +73,20 @@ The user will primarily request you perform SecOps tasks. This includes security # Tool usage policy -- You should proactively use the Task tool with specialized agents when the task at hand matches the agent's description. +- You should proactively use `delegate_task` with specialized agents when the task at hand matches the agent's description. - When WebFetch returns a message about a redirect to a different host, you should immediately make a new WebFetch request with the redirect URL provided in the response. - You can call multiple tools in a single response. If you intend to call multiple tools and there are no dependencies between them, make all independent tool calls in parallel. Maximize use of parallel tool calls where possible to increase efficiency. However, if some tool calls depend on previous calls to inform dependent values, do NOT call these tools in parallel and instead call them sequentially. For instance, if one operation must complete before another starts, run these operations sequentially instead. Never use placeholders or guess missing parameters in tool calls. -- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple Task tool calls. +- If the user specifies that they want you to run tools "in parallel", you MUST send a single message with multiple tool use content blocks. For example, if you need to launch multiple agents in parallel, send a single message with multiple `delegate_task` calls. - Use specialized tools instead of bash commands when possible, as this provides a better user experience. For file operations, use dedicated tools: Read for reading files instead of cat/head/tail, Edit for editing instead of sed/awk, and Write for creating files instead of cat with heredoc or echo redirection. Reserve bash tools exclusively for actual system commands and terminal operations that require shell execution. NEVER use bash echo or other command-line tools to communicate thoughts, explanations, or instructions to the user. Output all communication directly in your response text instead. -- VERY IMPORTANT: When exploring the codebase to gather context or to answer a question that is not a needle query for a specific file/class/function, it is CRITICAL that you use the Task tool instead of running search commands directly. +- VERY IMPORTANT: When exploring the codebase to gather context or to answer a question that is not a needle query for a specific file/class/function, it is CRITICAL that you use `delegate_task` instead of running search commands directly. user: Where are errors from the client handled? -assistant: [Uses the Task tool to find the files that handle client errors instead of using Glob or Grep directly] +assistant: [Uses `delegate_task` to find the files that handle client errors instead of using Glob or Grep directly] user: What is the codebase structure? -assistant: [Uses the Task tool] +assistant: [Uses `delegate_task`] IMPORTANT: Always use `todo(action="write")` to plan and track tasks throughout the conversation. diff --git a/tui/flocks/tool/task.ts b/tui/flocks/tool/delegate-task.ts similarity index 97% rename from tui/flocks/tool/task.ts rename to tui/flocks/tool/delegate-task.ts index f98316b39..e2b9c5d92 100644 --- a/tui/flocks/tool/task.ts +++ b/tui/flocks/tool/delegate-task.ts @@ -1,5 +1,5 @@ import { Tool } from "./tool" -import DESCRIPTION from "./task.txt" +import DESCRIPTION from "./delegate-task.txt" import z from "zod" import { Session } from "../session" import { Bus } from "../bus" @@ -20,7 +20,7 @@ const parameters = z.object({ command: z.string().describe("The command that triggered this task").optional(), }) -export const TaskTool = Tool.define("task", async (ctx) => { +export const DelegateTaskTool = Tool.define("delegate_task", async (ctx) => { const agents = await Agent.list().then((x) => x.filter((a) => a.mode !== "primary")) // Filter agents by permissions if agent provided @@ -41,7 +41,6 @@ export const TaskTool = Tool.define("task", async (ctx) => { async execute(params: z.infer, ctx) { const config = await Config.get() - // Skip permission check when user explicitly invoked via @ or command subtask if (!ctx.extra?.bypassAgentCheck) { await ctx.ask({ permission: "task", diff --git a/tui/flocks/tool/task.txt b/tui/flocks/tool/delegate-task.txt similarity index 79% rename from tui/flocks/tool/task.txt rename to tui/flocks/tool/delegate-task.txt index 7af2a6f60..21258b793 100644 --- a/tui/flocks/tool/task.txt +++ b/tui/flocks/tool/delegate-task.txt @@ -1,17 +1,17 @@ -Launch a new agent to handle complex, multistep tasks autonomously. +Delegate a complex, multistep task to another agent. Available agent types and the tools they have access to: {agents} -When using the Task tool, you must specify a subagent_type parameter to select which agent type to use. +When using the delegate_task tool, you must specify a subagent_type parameter to select which agent type to use. -When to use the Task tool: -- When you are instructed to execute custom slash commands. Use the Task tool with the slash command invocation as the entire prompt. The slash command can take arguments. For example: Task(description="Check the file", prompt="/check-file path/to/file.py") +When to use the delegate_task tool: +- When you are instructed to execute custom slash commands. Use delegate_task with the slash command invocation as the entire prompt. The slash command can take arguments. For example: delegate_task(description="Check the file", prompt="/check-file path/to/file.py") -When NOT to use the Task tool: -- If you want to read a specific file path, use the Read or Glob tool instead of the Task tool, to find the match more quickly +When NOT to use the delegate_task tool: +- If you want to read a specific file path, use the Read or Glob tool instead of delegate_task, to find the match more quickly - If you are searching for a specific class definition like "class Foo", use the Glob tool instead, to find the match more quickly -- If you are searching for code within a specific file or set of 2-3 files, use the Read tool instead of the Task tool, to find the match more quickly +- If you are searching for code within a specific file or set of 2-3 files, use the Read tool instead of delegate_task, to find the match more quickly - Other tasks that are not related to the agent descriptions above @@ -48,7 +48,7 @@ function isPrime(n) { Since a significant piece of code was written and the task was completed, now use the code-reviewer agent to review the code assistant: Now let me use the code-reviewer agent to review the code -assistant: Uses the Task tool to launch the code-reviewer agent +assistant: Uses delegate_task to launch the code-reviewer agent @@ -56,5 +56,5 @@ user: "Hello" Since the user is greeting, use the greeting-responder agent to respond with a friendly joke -assistant: "I'm going to use the Task tool to launch the with the greeting-responder agent" +assistant: "I'm going to use delegate_task to launch the greeting-responder agent" diff --git a/tui/flocks/tool/registry.ts b/tui/flocks/tool/registry.ts index 4f4eed7a0..0714a55c7 100644 --- a/tui/flocks/tool/registry.ts +++ b/tui/flocks/tool/registry.ts @@ -4,7 +4,7 @@ import { EditTool } from "./edit" import { GlobTool } from "./glob" import { GrepTool } from "./grep" import { ReadTool } from "./read" -import { TaskTool } from "./task" +import { DelegateTaskTool } from "./delegate-task" import { TodoTool } from "./todo" import { WebFetchTool } from "./webfetch" import { WriteTool } from "./write" @@ -99,7 +99,7 @@ export namespace ToolRegistry { GrepTool, EditTool, WriteTool, - TaskTool, + DelegateTaskTool, WebFetchTool, TodoTool, WebSearchTool, diff --git a/tui/sdk/gen/types.gen.ts b/tui/sdk/gen/types.gen.ts index 8ac5c7342..ed4ba1843 100644 --- a/tui/sdk/gen/types.gen.ts +++ b/tui/sdk/gen/types.gen.ts @@ -383,15 +383,6 @@ export type CompactionPart = { export type Part = | TextPart - | { - id: string - sessionID: string - messageID: string - type: "subtask" - prompt: string - description: string - agent: string - } | ReasoningPart | FilePart | ToolPart @@ -1217,7 +1208,6 @@ export type Config = { description?: string agent?: string model?: string - subtask?: boolean } } watcher?: { @@ -1432,21 +1422,12 @@ export type AgentPartInput = { } } -export type SubtaskPartInput = { - id?: string - type: "subtask" - prompt: string - description: string - agent: string -} - export type Command = { name: string description?: string agent?: string model?: string template: string - subtask?: boolean } export type Model = { @@ -2591,7 +2572,7 @@ export type SessionPromptData = { tools?: { [key: string]: boolean } - parts: Array + parts: Array } path: { /** @@ -2686,7 +2667,7 @@ export type SessionPromptAsyncData = { tools?: { [key: string]: boolean } - parts: Array + parts: Array } path: { /** diff --git a/tui/sdk/v2/gen/sdk.gen.ts b/tui/sdk/v2/gen/sdk.gen.ts index a84c70938..033e235b6 100644 --- a/tui/sdk/v2/gen/sdk.gen.ts +++ b/tui/sdk/v2/gen/sdk.gen.ts @@ -134,7 +134,6 @@ import type { SessionUnshareResponses, SessionUpdateErrors, SessionUpdateResponses, - SubtaskPartInput, TextPartInput, ToolIdsErrors, ToolIdsResponses, @@ -1364,7 +1363,7 @@ export class Session extends HeyApiClient { } system?: string variant?: string - parts?: Array + parts?: Array }, options?: Options, ) { @@ -1452,7 +1451,7 @@ export class Session extends HeyApiClient { } system?: string variant?: string - parts?: Array + parts?: Array }, options?: Options, ) { diff --git a/tui/sdk/v2/gen/types.gen.ts b/tui/sdk/v2/gen/types.gen.ts index 77f869dc6..0109b8e50 100644 --- a/tui/sdk/v2/gen/types.gen.ts +++ b/tui/sdk/v2/gen/types.gen.ts @@ -429,20 +429,6 @@ export type CompactionPart = { export type Part = | TextPart - | { - id: string - sessionID: string - messageID: string - type: "subtask" - prompt: string - description: string - agent: string - model?: { - providerID: string - modelID: string - } - command?: string - } | ReasoningPart | FilePart | ToolPart @@ -1617,7 +1603,6 @@ export type Config = { description?: string agent?: string model?: string - subtask?: boolean } } watcher?: { @@ -1953,19 +1938,6 @@ export type AgentPartInput = { } } -export type SubtaskPartInput = { - id?: string - type: "subtask" - prompt: string - description: string - agent: string - model?: { - providerID: string - modelID: string - } - command?: string -} - export type ProviderAuthMethod = { type: "oauth" | "api" label: string @@ -2071,7 +2043,6 @@ export type Command = { model?: string mcp?: boolean template: string - subtask?: boolean hints: Array } @@ -3226,7 +3197,7 @@ export type SessionPromptData = { } system?: string variant?: string - parts: Array + parts: Array } path: { /** @@ -3413,7 +3384,7 @@ export type SessionPromptAsyncData = { } system?: string variant?: string - parts: Array + parts: Array } path: { /** diff --git a/webui/src/api/skill.ts b/webui/src/api/skill.ts index 254520cec..00c2f9dad 100644 --- a/webui/src/api/skill.ts +++ b/webui/src/api/skill.ts @@ -39,7 +39,6 @@ export interface Command { template: string; agent?: string; model?: string; - subtask?: boolean; hidden: boolean; aliases: string[]; visible_surfaces: string[];