feat: Implement waiting timeout handling in BaseAgent and AgentState

This commit is contained in:
Ahmed Allam
2025-10-20 21:13:27 -07:00
committed by Ahmed Allam
parent c78f7d37de
commit 49df6ef8e0
4 changed files with 43 additions and 7 deletions

View File

@@ -206,6 +206,26 @@ class BaseAgent(metaclass=AgentMeta):
async def _wait_for_input(self) -> None: async def _wait_for_input(self) -> None:
import asyncio import asyncio
if self.state.has_waiting_timeout():
self.state.resume_from_waiting()
self.state.add_message("assistant", "Waiting timeout reached. Resuming execution.")
from strix.cli.tracer import get_global_tracer
tracer = get_global_tracer()
if tracer:
tracer.update_agent_status(self.state.agent_id, "running")
try:
from strix.tools.agents_graph.agents_graph_actions import _agent_graph
if self.state.agent_id in _agent_graph["nodes"]:
_agent_graph["nodes"][self.state.agent_id]["status"] = "running"
except (ImportError, KeyError):
pass
return
await asyncio.sleep(0.5) await asyncio.sleep(0.5)
async def _enter_waiting_state( async def _enter_waiting_state(

View File

@@ -24,6 +24,7 @@ class AgentState(BaseModel):
stop_requested: bool = False stop_requested: bool = False
waiting_for_input: bool = False waiting_for_input: bool = False
llm_failed: bool = False llm_failed: bool = False
waiting_start_time: datetime | None = None
final_result: dict[str, Any] | None = None final_result: dict[str, Any] | None = None
messages: list[dict[str, Any]] = Field(default_factory=list) messages: list[dict[str, Any]] = Field(default_factory=list)
@@ -88,12 +89,13 @@ class AgentState(BaseModel):
def enter_waiting_state(self, llm_failed: bool = False) -> None: def enter_waiting_state(self, llm_failed: bool = False) -> None:
self.waiting_for_input = True self.waiting_for_input = True
self.stop_requested = False self.waiting_start_time = datetime.now(UTC)
self.llm_failed = llm_failed self.llm_failed = llm_failed
self.last_updated = datetime.now(UTC).isoformat() self.last_updated = datetime.now(UTC).isoformat()
def resume_from_waiting(self, new_task: str | None = None) -> None: def resume_from_waiting(self, new_task: str | None = None) -> None:
self.waiting_for_input = False self.waiting_for_input = False
self.waiting_start_time = None
self.stop_requested = False self.stop_requested = False
self.completed = False self.completed = False
self.llm_failed = False self.llm_failed = False
@@ -104,6 +106,21 @@ class AgentState(BaseModel):
def has_reached_max_iterations(self) -> bool: def has_reached_max_iterations(self) -> bool:
return self.iteration >= self.max_iterations return self.iteration >= self.max_iterations
def has_waiting_timeout(self) -> bool:
if not self.waiting_for_input or not self.waiting_start_time:
return False
if (
self.stop_requested
or self.llm_failed
or self.completed
or self.has_reached_max_iterations()
):
return False
elapsed = (datetime.now(UTC) - self.waiting_start_time).total_seconds()
return elapsed > 120
def has_empty_last_messages(self, count: int = 3) -> bool: def has_empty_last_messages(self, count: int = 3) -> bool:
if len(self.messages) < count: if len(self.messages) < count:
return False return False

View File

@@ -603,5 +603,6 @@ def wait_for_message(
"Message from another agent", "Message from another agent",
"Message from user", "Message from user",
"Direct communication", "Direct communication",
"Waiting timeout reached",
], ],
} }

View File

@@ -198,9 +198,12 @@ NOTE: If you are waiting for an agent that is NOT your subagent, you first tell
- Another agent sends it a message via send_message_to_agent - Another agent sends it a message via send_message_to_agent
- A user sends it a direct message through the CLI - A user sends it a direct message through the CLI
- Any other form of inter-agent or user communication occurs - Any other form of inter-agent or user communication occurs
- Waiting timeout is reached
The agent will automatically resume from where it left off once a message is received. The agent will automatically resume from where it left off once a message is received.
This is particularly useful for parent agents waiting for subagent results or for coordination points in multi-agent workflows.</details> This is particularly useful for parent agents waiting for subagent results or for coordination points in multi-agent workflows.
NOTE: If you finished your task, and you do NOT have any child agents running, you should NEVER use this tool, and just call finish tool instead.
</details>
<parameters> <parameters>
<parameter name="reason" type="string" required="false"> <parameter name="reason" type="string" required="false">
<description>Explanation for why the agent is waiting (for logging and monitoring purposes)</description> <description>Explanation for why the agent is waiting (for logging and monitoring purposes)</description>
@@ -215,11 +218,6 @@ NOTE: If you are waiting for an agent that is NOT your subagent, you first tell
<parameter=reason>Waiting for subdomain enumeration and port scanning subagents to complete their tasks and report findings</parameter> <parameter=reason>Waiting for subdomain enumeration and port scanning subagents to complete their tasks and report findings</parameter>
</function> </function>
# Wait for user input on next steps
<function=wait_for_message>
<parameter=reason>Waiting for user decision on whether to proceed with exploitation of discovered SQL injection vulnerability</parameter>
</function>
# Coordinate with other agents # Coordinate with other agents
<function=wait_for_message> <function=wait_for_message>
<parameter=reason>Waiting for vulnerability assessment agent to share discovered attack vectors before proceeding with exploitation phase</parameter> <parameter=reason>Waiting for vulnerability assessment agent to share discovered attack vectors before proceeding with exploitation phase</parameter>