aury-agent 0.0.12__py3-none-any.whl → 0.0.14__py3-none-any.whl

aury/agents/hitl/exceptions.py

@@ -35,16 +35,16 @@ class HITLRequest:
     
     Stored in invocation for persistence.
     """
-    request_id: str
-    request_type: str  # ask_user, permission, form, workflow_human
+    hitl_id: str
+    hitl_type: str  # ask_user, confirm, permission, external_auth, workflow_human
     
-    # Display
-    message: str | None = None
-    options: list[str] | None = None
+    # Type-specific data
+    data: dict[str, Any] = field(default_factory=dict)  # {message, options, ...}
     
     # Context
     tool_name: str | None = None  # If triggered by tool
     node_id: str | None = None    # If triggered by workflow node
+    block_id: str | None = None   # Associated UI block
     
     # Metadata
     metadata: dict[str, Any] = field(default_factory=dict)
@@ -52,12 +52,12 @@ class HITLRequest:
     def to_dict(self) -> dict[str, Any]:
         """Convert to dictionary for serialization."""
         return {
-            "request_id": self.request_id,
-            "request_type": self.request_type,
-            "message": self.message,
-            "options": self.options,
+            "hitl_id": self.hitl_id,
+            "hitl_type": self.hitl_type,
+            "data": self.data,
             "tool_name": self.tool_name,
             "node_id": self.node_id,
+            "block_id": self.block_id,
             "metadata": self.metadata,
         }
     
@@ -65,16 +65,216 @@ class HITLRequest:
     def from_dict(cls, data: dict[str, Any]) -> "HITLRequest":
         """Create from dictionary."""
         return cls(
-            request_id=data["request_id"],
-            request_type=data.get("request_type", "ask_user"),
-            message=data.get("message"),
-            options=data.get("options"),
+            hitl_id=data["hitl_id"],
+            hitl_type=data.get("hitl_type", "ask_user"),
+            data=data.get("data", {}),
             tool_name=data.get("tool_name"),
             node_id=data.get("node_id"),
+            block_id=data.get("block_id"),
             metadata=data.get("metadata", {}),
         )
 
 
+@dataclass
+class ToolCheckpoint:
+    """Tool execution checkpoint for continuation mode.
+    
+    When a tool raises HITLSuspend with resume_mode="continuation",
+    the framework creates a ToolCheckpoint to save the tool's execution
+    state. When the user responds, the tool is resumed from this checkpoint.
+    
+    Use cases:
+    - OAuth authorization flow (wait for callback)
+    - Payment confirmation (wait for payment gateway)
+    - Multi-step wizards with user confirmation
+    - External system integration with async callbacks
+    
+    Storage:
+    - Stored via CheckpointBackend (Redis/DB)
+    - Keyed by checkpoint_id and callback_id
+    - Has TTL for automatic expiration
+    
+    Example:
+        # In tool execution:
+        raise HITLSuspend(
+            request_id="hitl_123",
+            request_type="external_auth",
+            resume_mode="continuation",
+            tool_state={"step": 2, "partial_data": {...}},
+            metadata={"auth_url": "https://...", "callback_id": "cb_456"},
+        )
+        
+        # Framework creates ToolCheckpoint and saves it
+        # When callback arrives, framework loads checkpoint and resumes tool
+    """
+    
+    # Identity
+    checkpoint_id: str
+    callback_id: str | None = None  # For external callback matching
+    
+    # Association
+    session_id: str | None = None
+    invocation_id: str | None = None
+    block_id: str | None = None  # Frontend HITL block
+    
+    # Tool execution context
+    tool_name: str = ""
+    tool_call_id: str = ""
+    params: dict[str, Any] = field(default_factory=dict)  # Original params
+    tool_state: dict[str, Any] = field(default_factory=dict)  # Internal state
+    
+    # HITL info
+    hitl_id: str = ""
+    hitl_type: str = ""  # ask_user, confirm, external_auth, etc.
+    
+    # Status
+    status: str = "pending"  # pending | completed | expired | failed | cancelled
+    expires_at: int | None = None  # Unix timestamp
+    
+    # User response (filled after callback/response)
+    user_response: Any | None = None
+    error: str | None = None
+    
+    # Timestamps
+    created_at: int = 0
+    updated_at: int = 0
+    
+    def __post_init__(self):
+        import time
+        now = int(time.time())
+        if not self.created_at:
+            self.created_at = now
+        if not self.updated_at:
+            self.updated_at = now
+    
+    @property
+    def is_expired(self) -> bool:
+        """Check if checkpoint has expired."""
+        if self.expires_at is None:
+            return False
+        import time
+        return time.time() > self.expires_at
+    
+    @property
+    def is_pending(self) -> bool:
+        """Check if checkpoint is waiting for response."""
+        return self.status == "pending" and not self.is_expired
+    
+    def mark_completed(self, response: Any) -> None:
+        """Mark checkpoint as completed with user response."""
+        import time
+        self.status = "completed"
+        self.user_response = response
+        self.updated_at = int(time.time())
+    
+    def mark_failed(self, error: str) -> None:
+        """Mark checkpoint as failed."""
+        import time
+        self.status = "failed"
+        self.error = error
+        self.updated_at = int(time.time())
+    
+    def mark_cancelled(self, reason: str = "user_cancelled") -> None:
+        """Mark checkpoint as cancelled."""
+        import time
+        self.status = "cancelled"
+        self.error = reason
+        self.updated_at = int(time.time())
+    
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "checkpoint_id": self.checkpoint_id,
+            "callback_id": self.callback_id,
+            "session_id": self.session_id,
+            "invocation_id": self.invocation_id,
+            "block_id": self.block_id,
+            "tool_name": self.tool_name,
+            "tool_call_id": self.tool_call_id,
+            "params": self.params,
+            "tool_state": self.tool_state,
+            "hitl_id": self.hitl_id,
+            "hitl_type": self.hitl_type,
+            "status": self.status,
+            "expires_at": self.expires_at,
+            "user_response": self.user_response,
+            "error": self.error,
+            "created_at": self.created_at,
+            "updated_at": self.updated_at,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ToolCheckpoint":
+        """Create from dictionary."""
+        return cls(
+            checkpoint_id=data["checkpoint_id"],
+            callback_id=data.get("callback_id"),
+            session_id=data.get("session_id"),
+            invocation_id=data.get("invocation_id"),
+            block_id=data.get("block_id"),
+            tool_name=data.get("tool_name", ""),
+            tool_call_id=data.get("tool_call_id", ""),
+            params=data.get("params", {}),
+            tool_state=data.get("tool_state", {}),
+            hitl_id=data.get("hitl_id", ""),
+            hitl_type=data.get("hitl_type", ""),
+            status=data.get("status", "pending"),
+            expires_at=data.get("expires_at"),
+            user_response=data.get("user_response"),
+            error=data.get("error"),
+            created_at=data.get("created_at", 0),
+            updated_at=data.get("updated_at", 0),
+        )
+    
+    @classmethod
+    def from_hitl_suspend(
+        cls,
+        suspend: HITLSuspend,
+        *,
+        tool_call_id: str,
+        params: dict[str, Any],
+        session_id: str | None = None,
+        invocation_id: str | None = None,
+        block_id: str | None = None,
+        expires_in: int | None = 600,  # Default 10 minutes
+    ) -> "ToolCheckpoint":
+        """Create checkpoint from HITLSuspend signal.
+        
+        Args:
+            suspend: The HITLSuspend signal
+            tool_call_id: Tool call ID
+            params: Original tool parameters
+            session_id: Session ID
+            invocation_id: Invocation ID
+            block_id: Frontend block ID
+            expires_in: Expiration in seconds (None = no expiration)
+        """
+        from ..core.types.session import generate_id
+        import time
+        
+        checkpoint_id = suspend.checkpoint_id or generate_id("ckpt")
+        callback_id = suspend.metadata.get("callback_id")
+        
+        expires_at = None
+        if expires_in is not None:
+            expires_at = int(time.time()) + expires_in
+        
+        return cls(
+            checkpoint_id=checkpoint_id,
+            callback_id=callback_id,
+            session_id=session_id,
+            invocation_id=invocation_id,
+            block_id=block_id,
+            tool_name=suspend.tool_name or "",
+            tool_call_id=tool_call_id,
+            params=params,
+            tool_state=suspend.tool_state or {},
+            hitl_id=suspend.hitl_id,
+            hitl_type=suspend.hitl_type,
+            expires_at=expires_at,
+        )
+
+
 __all__ = [
     # Signals
     "SuspendSignal",
@@ -84,4 +284,5 @@ __all__ = [
     "HITLCancelledError",
     # Types
     "HITLRequest",
+    "ToolCheckpoint",
 ]

aury/agents/react/agent.py

@@ -377,6 +377,42 @@ class ReactAgent(BaseAgent):
                     ))
                     break
 
+                # Re-fetch context from providers (providers decide whether to update)
+                logger.debug(
+                    "Re-fetching agent context for step",
+                    extra={
+                        "invocation_id": self._current_invocation.id,
+                        "step": self._current_step,
+                    },
+                )
+                self._agent_context = await ctx_helpers.fetch_agent_context(
+                    self._ctx,
+                    input,
+                    self._context_providers,
+                    self._tools,
+                    self._delegate_tool_class,
+                    self._middleware_chain,
+                )
+                
+                # Update system message with new context (in case providers updated system_content)
+                if self._message_history and self._message_history[0].role == "system":
+                    # Rebuild system message using helper
+                    final_system_prompt = ctx_helpers.build_system_message(
+                        self._agent_context,
+                        self.config.system_prompt,
+                        input,
+                    )
+                    
+                    # Log if context was injected
+                    if self._agent_context.system_content:
+                        logger.info(
+                            f"Updated system message with context (length: {len(self._agent_context.system_content)})",
+                            extra={"invocation_id": self._current_invocation.id, "step": self._current_step},
+                        )
+                    
+                    # Update the system message
+                    self._message_history[0] = LLMMessage(role="system", content=final_system_prompt)
+
                 # Take snapshot before step
                 snapshot_id = None
                 if self.snapshot:
@@ -528,6 +564,17 @@ class ReactAgent(BaseAgent):
             if self._current_invocation:
                 self._current_invocation.state = InvocationState.SUSPENDED
                 
+                # Save agent_state for resume (only if persist_hitl_state is enabled)
+                if self.config.persist_hitl_state:
+                    self._current_invocation.agent_state = {
+                        "step": self._current_step,
+                        "message_history": [
+                            {"role": m.role, "content": m.content} for m in self._message_history
+                        ],
+                        "text_buffer": self._text_buffer,
+                    }
+                    self._current_invocation.step_count = self._current_step
+                
                 # Save invocation state
                 if self.ctx.backends and self.ctx.backends.invocation:
                     await self.ctx.backends.invocation.update(

aury/agents/react/context.py

@@ -221,31 +221,11 @@ async def build_messages(
     messages = []
     
     # System message: config.system_prompt + agent_context.system_content
-    final_system_prompt = system_prompt or default_system_prompt(agent_context.tools)
-    
-    # Format system_prompt with dynamic variables
-    now = datetime.now()
-    
-    # Build template variables: datetime + custom vars from input
-    template_vars = {
-        "current_date": now.strftime("%Y-%m-%d"),
-        "current_time": now.strftime("%H:%M:%S"),
-        "current_datetime": now.strftime("%Y-%m-%d %H:%M:%S"),
-    }
-    
-    # Add custom variables from PromptInput (user_name, tenant, etc.)
-    if hasattr(input, 'vars') and input.vars:
-        template_vars.update(input.vars)
-    
-    try:
-        final_system_prompt = final_system_prompt.format(**template_vars)
-    except KeyError as e:
-        # Log missing variable but continue
-        logger.debug(f"System prompt template variable not found: {e}")
-        pass
-    
-    if agent_context.system_content:
-        final_system_prompt = final_system_prompt + "\n\n" + agent_context.system_content
+    final_system_prompt = build_system_message(
+        agent_context,
+        system_prompt,
+        input,
+    )
     messages.append(LLMMessage(role="system", content=final_system_prompt))
     
     # Historical messages from AgentContext (provided by MessageContextProvider)
@@ -283,6 +263,52 @@ async def build_messages(
     return messages
 
 
+def build_system_message(
+    agent_context: AgentContext,
+    base_system_prompt: str | None,
+    input: "PromptInput | None" = None,
+) -> str:
+    """Build system message with agent context.
+    
+    Args:
+        agent_context: Agent context with system_content, tools, etc.
+        base_system_prompt: Base system prompt (or None for default)
+        input: Prompt input for custom template variables
+        
+    Returns:
+        Final system prompt string
+    """
+    from datetime import datetime
+    
+    # Get base prompt
+    final_system_prompt = base_system_prompt or default_system_prompt(agent_context.tools)
+    
+    # Build template variables: datetime + custom vars from input
+    now = datetime.now()
+    template_vars = {
+        "current_date": now.strftime("%Y-%m-%d"),
+        "current_time": now.strftime("%H:%M:%S"),
+        "current_datetime": now.strftime("%Y-%m-%d %H:%M:%S"),
+    }
+    
+    # Add custom variables from PromptInput
+    if input and hasattr(input, 'vars') and input.vars:
+        template_vars.update(input.vars)
+    
+    # Format with template variables
+    try:
+        final_system_prompt = final_system_prompt.format(**template_vars)
+    except KeyError as e:
+        logger.debug(f"System prompt template variable not found: {e}")
+        pass
+    
+    # Append system_content if available
+    if agent_context.system_content:
+        final_system_prompt = final_system_prompt + "\n\n" + agent_context.system_content
+    
+    return final_system_prompt
+
+
 def default_system_prompt(tools: list["BaseTool"]) -> str:
     """Generate default system prompt with tool descriptions.
     

aury/agents/react/factory.py

@@ -47,6 +47,8 @@ def create_react_agent(
     delegate_tool_class: "type[BaseTool] | None" = None,
     # Context metadata
     context_metadata: dict | None = None,
+    # HITL resume support
+    invocation_id: str | None = None,
 ) -> "ReactAgent":
     """Create ReactAgent with minimal boilerplate.
     

aury/agents/react/pause.py

@@ -91,8 +91,9 @@ async def resume_agent_internal(agent: "ReactAgent", invocation_id: str) -> None
 
     invocation = Invocation.from_dict(inv_data)
 
-    if invocation.state != InvocationState.PAUSED:
-        raise ValueError(f"Invocation is not paused: {invocation.state}")
+    # Support both PAUSED and SUSPENDED (HITL) states
+    if invocation.state not in (InvocationState.PAUSED, InvocationState.SUSPENDED):
+        raise ValueError(f"Invocation is not paused/suspended: {invocation.state}")
 
     # Restore state
     agent._current_invocation = invocation
@@ -154,6 +155,16 @@ async def resume_agent_internal(agent: "ReactAgent", invocation_id: str) -> None
         if not agent._paused:
             agent._current_invocation.state = InvocationState.COMPLETED
             agent._current_invocation.finished_at = __import__("datetime").datetime.now()
+            
+            # Clear agent_state after successful completion (save space)
+            agent._current_invocation.agent_state = None
+            
+            # Update invocation to database
+            if agent.ctx.backends and agent.ctx.backends.invocation:
+                await agent.ctx.backends.invocation.update(
+                    agent._current_invocation.id,
+                    agent._current_invocation.to_dict(),
+                )
 
     except Exception as e:
         agent._current_invocation.state = InvocationState.FAILED

aury/agents/react/step.py

@@ -289,8 +289,12 @@ async def execute_step(agent: "ReactAgent") -> str | None:
         extra={"invocation_id": agent._current_invocation.id, "model": agent.llm.model},
     )
     
+    # Track whether we aborted mid-stream
+    aborted = False
+    
     async for event in agent.llm.complete(**llm_kwargs):
         if await agent._check_abort():
+            aborted = True
             break
 
         if event.type == "content":
@@ -360,6 +364,8 @@ async def execute_step(agent: "ReactAgent") -> str | None:
 
         elif event.type == "thinking_completed":
             # Thinking completed - emit block completed status
+            # Note: thinking_completed from LLM means it finished naturally,
+            # so we always use "completed" here (not aborted)
             if agent._current_thinking_block_id and not thinking_completed_emitted:
                 await agent.ctx.emit(BlockEvent(
                     block_id=agent._current_thinking_block_id,
@@ -386,14 +392,19 @@ async def execute_step(agent: "ReactAgent") -> str | None:
                 block_id = generate_id("blk")
                 agent._tool_call_blocks[tc.id] = block_id
                 
+                # Get display_name from tool if available
+                tool = agent._get_tool(tc.name)
+                display_name = tool.display_name if tool else tc.name
+                
                 await agent.ctx.emit(BlockEvent(
                     block_id=block_id,
                     kind=BlockKind.TOOL_USE,
                     op=BlockOp.APPLY,
                     data={
                         "name": tc.name,
+                        "display_name": display_name,
                         "call_id": tc.id,
-                        "status": "streaming",  # Indicate arguments are streaming
+                        "status": "pending",  # Initial status, arguments pending
                     },
                 ))
 
@@ -498,9 +509,13 @@ async def execute_step(agent: "ReactAgent") -> str | None:
             # Tool call complete (arguments fully received)
             if event.tool_call:
                 tc = event.tool_call
+                # Strict mode: tool_call_start must have been received
+                block_id = agent._tool_call_blocks[tc.id]  # Will raise KeyError if not found
+                
                 invocation = ToolInvocation(
                     tool_call_id=tc.id,
                     tool_name=tc.name,
+                    block_id=block_id,
                     args_raw=tc.arguments,
                     state=ToolInvocationState.CALL,
                 )
@@ -512,18 +527,19 @@ async def execute_step(agent: "ReactAgent") -> str | None:
                     invocation.args = {}
 
                 agent._tool_invocations.append(invocation)
-
-                # Strict mode: tool_call_start must have been received
-                block_id = agent._tool_call_blocks[tc.id]  # Will raise KeyError if not found
+                
+                # Build patch data
+                patch_data: dict[str, Any] = {
+                    "call_id": tc.id,
+                    "arguments": invocation.args,
+                    "status": "ready",
+                }
+                
                 await agent.ctx.emit(BlockEvent(
                     block_id=block_id,
                     kind=BlockKind.TOOL_USE,
                     op=BlockOp.PATCH,
-                    data={
-                        "call_id": tc.id,
-                        "arguments": invocation.args,
-                        "status": "ready",
-                    },
+                    data=patch_data,
                 ))
 
                 await agent.bus.publish(
@@ -605,16 +621,27 @@ async def execute_step(agent: "ReactAgent") -> str | None:
                     extra={"invocation_id": agent._current_invocation.id},
                 )
     
-    # Emit text block completed status
+    # Emit thinking block final status if streaming and not yet completed
+    if agent._current_thinking_block_id and not thinking_completed_emitted:
+        status = "aborted" if aborted else "completed"
+        await agent.ctx.emit(BlockEvent(
+            block_id=agent._current_thinking_block_id,
+            kind=BlockKind.THINKING,
+            op=BlockOp.PATCH,
+            data={"status": status},
+        ))
+    
+    # Emit text block final status (completed or aborted)
     if agent._current_text_block_id:
+        status = "aborted" if aborted else "completed"
         await agent.ctx.emit(BlockEvent(
             block_id=agent._current_text_block_id,
             kind=BlockKind.TEXT,
             op=BlockOp.PATCH,
-            data={"status": "completed"},
+            data={"status": status},
         ))
     
-    # If thinking was buffered, emit it now
+    # If thinking was buffered, emit it now (non-streaming mode)
     if agent._thinking_buffer and not agent.config.stream_thinking:
         await agent.ctx.emit(BlockEvent(
             kind=BlockKind.THINKING,