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

aury/agents/react/tools.py

@@ -12,6 +12,7 @@ from ..core.types import (
     ToolContext,
     ToolResult,
     ToolInvocation,
+    ToolInvocationState,
 )
 from ..core.signals import SuspendSignal
 from ..llm import LLMMessage
@@ -39,12 +40,104 @@ def get_tool(agent: "ReactAgent", tool_name: str) -> "BaseTool | None":
     return None
 
 
-async def execute_tool(agent: "ReactAgent", invocation: ToolInvocation) -> ToolResult:
+async def _complete_tool_invocation(
+    agent: "ReactAgent",
+    invocation: ToolInvocation,
+    result: ToolResult,
+) -> None:
+    """Complete a tool invocation.
+    
+    This function:
+    1. Updates TOOL_USE block status (success/failed/aborted)
+    2. Publishes TOOL_END event
+    3. Adds result to LLM message history
+    
+    Args:
+        agent: ReactAgent instance
+        invocation: Tool invocation
+        result: Tool execution result
+    """
+    # 1. Update invocation state if not already marked
+    if not invocation.result:
+        invocation.mark_result(
+            result.output,
+            is_error=result.is_error,
+            truncated_result=result.truncated_output,
+        )
+    
+    # 2. Map invocation state to string status for frontend
+    status_map = {
+        ToolInvocationState.SUCCESS: "success",
+        ToolInvocationState.FAILED: "failed",
+        ToolInvocationState.ABORTED: "aborted",
+        ToolInvocationState.RESULT: "success",  # Backward compatibility
+    }
+    status = status_map.get(invocation.state, "failed")
+    
+    # 3. Update TOOL_USE block with final status
+    # Tools can PATCH their own block during execution for UI updates.
+    # Here we only update the final status.
+    if invocation.block_id:
+        patch_data: dict[str, Any] = {"status": status}
+        if invocation.is_error:
+            patch_data["error"] = result.output
+        
+        await agent.ctx.emit(BlockEvent(
+            block_id=invocation.block_id,
+            kind=BlockKind.TOOL_USE,
+            op=BlockOp.PATCH,
+            data=patch_data,
+        ))
+    
+    # 5. Publish TOOL_END event (for internal subscribers: monitoring, billing, etc.)
+    await agent.bus.publish(
+        Events.TOOL_END,
+        {
+            "call_id": invocation.tool_call_id,
+            "tool": invocation.tool_name,
+            "result": result.output[:500],
+            "is_error": invocation.is_error,
+            "status": status,
+            "duration_ms": invocation.duration_ms,
+        },
+    )
+    
+    # 6. Add to message history (for LLM, use truncated for context window)
+    agent._message_history.append(
+        LLMMessage(
+            role="tool",
+            content=invocation.result,
+            tool_call_id=invocation.tool_call_id,
+            name=invocation.tool_name,
+            truncated_content=invocation.truncated_result,
+        )
+    )
+    
+    logger.debug(
+        f"Tool completed: {invocation.tool_name} ({status})",
+        extra={
+            "invocation_id": agent._current_invocation.id if agent._current_invocation else "",
+            "call_id": invocation.tool_call_id,
+            "status": status,
+        },
+    )
+
+
+async def execute_tool(
+    agent: "ReactAgent",
+    invocation: ToolInvocation,
+    checkpoint: Any | None = None,
+) -> ToolResult:
     """Execute a single tool call.
     
+    Supports two execution modes:
+    1. Standard mode: Tool runs to completion
+    2. Continuation mode: Tool can be resumed from checkpoint
+    
     Args:
         agent: ReactAgent instance
         invocation: Tool invocation to execute
+        checkpoint: ToolCheckpoint for resuming (continuation mode only)
         
     Returns:
         ToolResult from tool execution
@@ -52,7 +145,11 @@ async def execute_tool(agent: "ReactAgent", invocation: ToolInvocation) -> ToolR
     # Check abort before execution
     if await agent._check_abort():
         error_msg = f"Tool {invocation.tool_name} aborted before execution"
-        invocation.mark_result(error_msg, is_error=True)
+        invocation.mark_result(
+            error_msg,
+            is_error=True,
+            status=ToolInvocationState.ABORTED,  # Explicit ABORTED status
+        )
         logger.info(
             f"Tool aborted before execution: {invocation.tool_name}",
             extra={
@@ -63,13 +160,17 @@ async def execute_tool(agent: "ReactAgent", invocation: ToolInvocation) -> ToolR
         return ToolResult.error(error_msg)
     
     invocation.mark_call_complete()
+    
+    is_resuming = checkpoint is not None
 
     logger.info(
-        f"Executing tool: {invocation.tool_name}",
+        f"Executing tool: {invocation.tool_name}" + (" (resuming from checkpoint)" if is_resuming else ""),
         extra={
             "invocation_id": agent._current_invocation.id if agent._current_invocation else "",
             "call_id": invocation.tool_call_id,
             "arguments": invocation.args,
+            "is_resuming": is_resuming,
+            "checkpoint_id": checkpoint.checkpoint_id if checkpoint else None,
         },
     )
 
@@ -85,8 +186,13 @@ async def execute_tool(agent: "ReactAgent", invocation: ToolInvocation) -> ToolR
             )
             return ToolResult.error(error_msg)
 
+        # Set tool context on ctx for middleware/tool access
+        agent.ctx.tool_call_id = invocation.tool_call_id
+        agent.ctx.tool_block_id = invocation.block_id
+        
         # === Middleware: on_tool_call ===
-        if agent.middleware:
+        # Skip middleware on resume - already processed on first call
+        if agent.middleware and not is_resuming:
             logger.debug(
                 f"Calling middleware: on_tool_call ({invocation.tool_name})",
                 extra={"invocation_id": agent._current_invocation.id, "call_id": invocation.tool_call_id},
@@ -111,29 +217,65 @@ async def execute_tool(agent: "ReactAgent", invocation: ToolInvocation) -> ToolR
                 invocation.args = hook_result.modified_data
 
         # Create ToolContext
-        # Get block_id for this tool call
-        tool_block_id = agent._tool_call_blocks.get(invocation.tool_call_id, "")
         tool_ctx = ToolContext(
             session_id=agent.session.id,
             invocation_id=agent._current_invocation.id if agent._current_invocation else "",
-            block_id=tool_block_id,
+            block_id=invocation.block_id,
             call_id=invocation.tool_call_id,
             agent=agent.config.name,
             abort_signal=agent._abort,
             update_metadata=agent._noop_update_metadata,
             middleware=agent.middleware,
         )
+        
+        # Emit running status before execution
+        if invocation.block_id:
+            await agent.ctx.emit(BlockEvent(
+                block_id=invocation.block_id,
+                kind=BlockKind.TOOL_USE,
+                op=BlockOp.PATCH,
+                data={"status": "running"},
+            ))
 
-        # Execute tool (with optional timeout from tool.config)
+        # Determine execution mode
+        # Use continuation mode if: tool supports it AND (resuming OR tool always uses continuation)
+        use_continuation = tool.supports_continuation
         timeout = tool.config.timeout
-        if timeout is not None:
-            result = await asyncio.wait_for(
-                tool.execute(invocation.args, tool_ctx),
-                timeout=timeout,
+        
+        # Set ctx on tool for self.emit() support
+        tool._set_ctx(tool_ctx)
+        
+        try:
+            if use_continuation:
+                # Continuation mode: use execute_resumable
+                if timeout is not None:
+                    result = await asyncio.wait_for(
+                        tool.execute_resumable(invocation.args, tool_ctx, checkpoint),
+                        timeout=timeout,
+                    )
+                else:
+                    result = await tool.execute_resumable(invocation.args, tool_ctx, checkpoint)
+            else:
+                # Standard mode: use execute
+                if timeout is not None:
+                    result = await asyncio.wait_for(
+                        tool.execute(invocation.args, tool_ctx),
+                        timeout=timeout,
+                    )
+                else:
+                    result = await tool.execute(invocation.args, tool_ctx)
+        finally:
+            # Clear ctx after execution
+            tool._set_ctx(None)
+        
+        # Check abort after execution
+        if await agent._check_abort():
+            invocation.mark_result(
+                "Tool execution interrupted by user",
+                is_error=True,
+                status=ToolInvocationState.ABORTED,
             )
-        else:
-            # No timeout - tool runs until completion
-            result = await tool.execute(invocation.args, tool_ctx)
+            return ToolResult.error("Tool execution interrupted by user")
 
         # === Middleware: on_tool_end ===
         if agent.middleware:
@@ -148,7 +290,12 @@ async def execute_tool(agent: "ReactAgent", invocation: ToolInvocation) -> ToolR
                     extra={"invocation_id": agent._current_invocation.id},
                 )
 
-        invocation.mark_result(result.output, is_error=result.is_error)
+        # Mark result - state will be auto-inferred (SUCCESS if not error, FAILED if error)
+        invocation.mark_result(
+            result.output,
+            is_error=result.is_error,
+            truncated_result=result.truncated_output,
+        )
         logger.info(
             f"Tool executed: {invocation.tool_name}",
             extra={
@@ -163,7 +310,11 @@ async def execute_tool(agent: "ReactAgent", invocation: ToolInvocation) -> ToolR
     except asyncio.TimeoutError:
         timeout = tool.config.timeout if tool else None
         error_msg = f"Tool {invocation.tool_name} timed out after {timeout}s"
-        invocation.mark_result(error_msg, is_error=True)
+        invocation.mark_result(
+            error_msg,
+            is_error=True,
+            status=ToolInvocationState.FAILED,  # Timeout is treated as FAILED
+        )
         logger.error(
             f"Tool timeout: {invocation.tool_name}",
             extra={
@@ -179,25 +330,38 @@ async def execute_tool(agent: "ReactAgent", invocation: ToolInvocation) -> ToolR
     
     except Exception as e:
         import traceback
+        import sys
         error_type = type(e).__name__
         error_msg = f"Tool execution error ({error_type}): {e}"
         stack_trace = traceback.format_exc()
-        invocation.mark_result(error_msg, is_error=True)
-        logger.error(
-            f"Tool execution failed: {invocation.tool_name} - {error_type}: {e}\n{stack_trace}",
-            extra={
-                "invocation_id": agent._current_invocation.id if agent._current_invocation else "",
-                "error_type": error_type,
-                "error": str(e),
-            },
+        invocation.mark_result(
+            error_msg,
+            is_error=True,
+            status=ToolInvocationState.FAILED,
         )
+        # Log error - wrap in try to ensure we always return
+        try:
+            logger.error(
+                f"Tool execution failed: {invocation.tool_name} - {error_type}: {e}\n{stack_trace}",
+                extra={
+                    "invocation_id": agent._current_invocation.id if agent._current_invocation else "",
+                    "error_type": error_type,
+                    "error": str(e),
+                },
+            )
+        except Exception as log_err:
+            # If standard logging fails, print directly to stderr
+            print(f"[TOOL_ERROR] Logging failed: {log_err}", file=sys.stderr)
+            print(f"[TOOL_ERROR] Original error: {invocation.tool_name} - {error_type}: {e}", file=sys.stderr)
+            print(stack_trace, file=sys.stderr)
         return ToolResult.error(error_msg)
 
 
 async def process_tool_results(agent: "ReactAgent") -> None:
-    """Execute tool calls and add results to history.
+    """Execute tool calls and emit results in real-time.
 
-    Executes tools in parallel or sequentially based on config.
+    For parallel execution, uses asyncio.as_completed to emit each tool's
+    result immediately upon completion (not waiting for all tools).
     
     This function directly modifies agent's internal state:
     - agent._message_history
@@ -223,25 +387,72 @@ async def process_tool_results(agent: "ReactAgent") -> None:
             "Tool execution aborted before starting",
             extra={"invocation_id": agent._current_invocation.id},
         )
-        # Return empty results - agent loop will handle abort
         return
     
     # Execute tools based on configuration
     if agent.config.parallel_tool_execution:
-        # Parallel execution using asyncio.gather with create_task
-        # create_task ensures each task gets its own ContextVar copy
-        tasks = [asyncio.create_task(execute_tool(agent, inv)) for inv in agent._tool_invocations]
-        results = await asyncio.gather(*tasks, return_exceptions=True)
+        # Parallel execution with real-time result emission
+        # Use as_completed to process and emit results as they finish
+        
+        # Wrapper to return both invocation and result
+        async def execute_with_invocation(inv: ToolInvocation) -> tuple[ToolInvocation, ToolResult]:
+            try:
+                result = await execute_tool(agent, inv)
+                return inv, result
+            except SuspendSignal:
+                # HITL signal - record placeholder before propagating
+                inv.mark_result("[等待用户输入]", is_error=False)
+                agent._message_history.append(
+                    LLMMessage(
+                        role="tool",
+                        content=inv.result,
+                        tool_call_id=inv.tool_call_id,
+                        name=inv.tool_name,
+                    )
+                )
+                await agent._save_tool_messages()
+                raise  # Re-raise to abort all tools
         
-        # Check abort after parallel execution - cancel remaining if aborted
+        tasks = [
+            asyncio.create_task(execute_with_invocation(inv))
+            for inv in agent._tool_invocations
+        ]
+        
+        # Process each tool as it completes
+        for coro in asyncio.as_completed(tasks):
+            try:
+                invocation, result = await coro
+                
+                # Complete invocation immediately (don't wait for other tools)
+                await _complete_tool_invocation(agent, invocation, result)
+                
+            except SuspendSignal:
+                # HITL signal - placeholder already recorded in execute_with_invocation
+                logger.info(
+                    "Tool execution suspended (HITL) during parallel execution",
+                    extra={"invocation_id": agent._current_invocation.id},
+                )
+                raise  # Propagate to abort all other tools
+            
+            except Exception as e:
+                # Unexpected error in the wrapper itself (not from execute_tool)
+                # This should be very rare as execute_tool catches all exceptions
+                logger.error(
+                    f"Unexpected error in parallel tool wrapper: {e}",
+                    extra={"invocation_id": agent._current_invocation.id},
+                    exc_info=True,
+                )
+                # Continue with other tools
+        
+        # Check abort after all parallel tools complete
         if await agent._check_abort():
             logger.info(
                 "Tool execution aborted after parallel execution",
                 extra={"invocation_id": agent._current_invocation.id},
             )
+    
     else:
         # Sequential execution with abort check between tools
-        results = []
         for inv in agent._tool_invocations:
             # Check abort before each tool
             if await agent._check_abort():
@@ -250,123 +461,42 @@ async def process_tool_results(agent: "ReactAgent") -> None:
                     extra={"invocation_id": agent._current_invocation.id},
                 )
                 # Mark remaining as aborted
-                error_result = ToolResult.error(f"Aborted before execution")
-                results.append(error_result)
-                inv.mark_result(error_result.output, is_error=True)
+                error_result = ToolResult.error("Aborted before execution")
+                inv.mark_result(
+                    error_result.output,
+                    is_error=True,
+                    status=ToolInvocationState.ABORTED,
+                )
+                await _complete_tool_invocation(agent, inv, error_result)
                 continue
                 
             try:
                 result = await execute_tool(agent, inv)
-                results.append(result)
-            except Exception as e:
-                results.append(e)
-
-    # Check for SuspendSignal - record tool_results first, then propagate
-    suspend_signal = None
-    for result in results:
-        if isinstance(result, SuspendSignal):
-            suspend_signal = result
-            break
-    
-    if suspend_signal:
-        # Must record tool_results to history BEFORE suspending
-        # Otherwise Claude API will reject next request (tool_use without tool_result)
-        for invocation, result in zip(agent._tool_invocations, results):
-            if isinstance(result, SuspendSignal):
-                # HITL tool - record placeholder result
-                invocation.mark_result(f"[等待用户输入: {suspend_signal.request_type}]", is_error=False)
-            else:
-                # Normal tool result - pass truncated_output
-                output = result.output if hasattr(result, 'output') else str(result)
-                truncated = getattr(result, 'truncated_output', None)
-                invocation.mark_result(
-                    output,
-                    is_error=getattr(result, 'is_error', False),
-                    truncated_result=truncated,
+                await _complete_tool_invocation(agent, inv, result)
+                
+            except SuspendSignal:
+                # HITL signal - record placeholder and propagate
+                inv.mark_result("[等待用户输入]", is_error=False)
+                agent._message_history.append(
+                    LLMMessage(
+                        role="tool",
+                        content=inv.result,
+                        tool_call_id=inv.tool_call_id,
+                        name=inv.tool_name,
+                    )
                 )
+                await agent._save_tool_messages()
+                raise
             
-            # Add to in-memory history (use truncated for context window)
-            agent._message_history.append(
-                LLMMessage(
-                    role="tool",
-                    content=invocation.result,
-                    tool_call_id=invocation.tool_call_id,
-                    name=invocation.tool_name,  # Required for Gemini
-                    truncated_content=invocation.truncated_result,
+            except Exception as e:
+                logger.error(
+                    f"Unexpected error in sequential tool execution: {e}",
+                    extra={"invocation_id": agent._current_invocation.id},
                 )
-            )
-        
-        # Persist via _save_tool_messages
-        await agent._save_tool_messages()
-        
-        logger.info(
-            "Tool execution suspended (HITL), tool_results recorded",
-            extra={"invocation_id": agent._current_invocation.id},
-        )
-        raise suspend_signal
-    
-    # Process results
-    tool_results = []
-
-    for invocation, result in zip(agent._tool_invocations, results):
-        # Handle exceptions from gather
-        if isinstance(result, Exception):
-            error_msg = f"Tool execution error: {str(result)}"
-            invocation.mark_result(error_msg, is_error=True)
-            result = ToolResult.error(error_msg)
-        else:
-            # Mark result with truncated_output
-            invocation.mark_result(
-                result.output,
-                is_error=result.is_error,
-                truncated_result=result.truncated_output,
-            )
-
-        # Get parent block_id from tool_call mapping
-        parent_block_id = agent._tool_call_blocks.get(invocation.tool_call_id)
-        
-        await agent.ctx.emit(BlockEvent(
-            kind=BlockKind.TOOL_RESULT,
-            op=BlockOp.APPLY,
-            parent_id=parent_block_id,
-            data={
-                "call_id": invocation.tool_call_id,
-                "content": result.output,
-                "is_error": invocation.is_error,
-            },
-        ))
-
-        await agent.bus.publish(
-            Events.TOOL_END,
-            {
-                "call_id": invocation.tool_call_id,
-                "tool": invocation.tool_name,
-                "result": result.output[:500],  # Truncate for event
-                "is_error": invocation.is_error,
-                "duration_ms": invocation.duration_ms,
-            },
-        )
-
-        tool_results.append(
-            {
-                "type": "tool_result",
-                "tool_use_id": invocation.tool_call_id,
-                "tool_name": invocation.tool_name,
-                "content": result.output,
-                "truncated_content": result.truncated_output,
-                "is_error": invocation.is_error,
-            }
-        )
-
-    # Add tool results as tool messages (OpenAI format)
-    for tr in tool_results:
-        print(f"[DEBUG _process_tool_results] Adding tool_result to history: {tr}")
-        agent._message_history.append(
-            LLMMessage(
-                role="tool",
-                content=tr["content"],
-                tool_call_id=tr["tool_use_id"],
-                name=tr["tool_name"],  # Required for Gemini
-                truncated_content=tr.get("truncated_content"),
-            )
-        )
+                error_result = ToolResult.error(f"Unexpected error: {str(e)}")
+                inv.mark_result(
+                    error_result.output,
+                    is_error=True,
+                    status=ToolInvocationState.FAILED,
+                )
+                await _complete_tool_invocation(agent, inv, error_result)

aury/agents/tool/builtin/ask_user.py

@@ -131,10 +131,6 @@ The conversation will pause until the user responds."""
         context: str | None,
     ) -> None:
         """Emit ASK block."""
-        emit = getattr(ctx, 'emit', None)
-        if emit is None:
-            return
-        
         block = BlockEvent(
             kind=BlockKind.ASK,
             op=BlockOp.APPLY,
@@ -149,7 +145,7 @@ The conversation will pause until the user responds."""
             invocation_id=ctx.invocation_id,
         )
         
-        await emit(block)
+        await self.emit(block)
 
 
 __all__ = ["AskUserTool"]

aury/agents/tool/builtin/delegate.py

@@ -648,14 +648,10 @@ Specify the agent key and task data."""
         event: Any,
     ) -> None:
         """发送 sub-agent 进度事件（transient，不持久化）."""
-        emit = getattr(ctx, 'emit', None)
-        if emit is None:
-            return
-        
         # 直接转发 event，但标记为 sub-agent 的
         if hasattr(event, 'to_dict'):
             # 已经是 BlockEvent，直接转发
-            await emit(event)
+            await self.emit(event)
     
     async def _execute_delegated(
         self,
@@ -705,10 +701,6 @@ Specify the agent key and task data."""
         branch: str | None = None,
     ) -> None:
         """Emit SUB_AGENT block."""
-        emit = getattr(ctx, 'emit', None)
-        if emit is None:
-            return
-        
         block = BlockEvent(
             block_id=block_id,
             kind=BlockKind.SUB_AGENT,
@@ -723,7 +715,7 @@ Specify the agent key and task data."""
             invocation_id=ctx.invocation_id,
         )
         
-        await emit(block)
+        await self.emit(block)
     
     async def _emit_parallel_block(
         self,
@@ -733,10 +725,6 @@ Specify the agent key and task data."""
         block_id: str,
     ) -> None:
         """Emit PARALLEL block for parallel execution."""
-        emit = getattr(ctx, 'emit', None)
-        if emit is None:
-            return
-        
         block = BlockEvent(
             block_id=block_id,
             kind=BlockKind.SUB_AGENT,  # Use SUB_AGENT with parallel flag
@@ -750,7 +738,7 @@ Specify the agent key and task data."""
             invocation_id=ctx.invocation_id,
         )
         
-        await emit(block)
+        await self.emit(block)
 
 
 __all__ = ["DelegateTool"]

aury/agents/tool/builtin/plan.py

@@ -238,10 +238,6 @@ Use this to track your progress on complex tasks."""
     
     async def _emit_plan_block(self, ctx: ToolContext, plan: dict, action: str) -> None:
         """Emit PLAN block event."""
-        emit = getattr(ctx, 'emit', None)
-        if emit is None:
-            return
-        
         # Get or generate block_id
         block_id = plan.get("block_id")
         if not block_id:
@@ -262,7 +258,7 @@ Use this to track your progress on complex tasks."""
             invocation_id=ctx.invocation_id,
         )
         
-        await emit(block)
+        await self.emit(block)
     
     def _format_plan(self, plan: dict) -> str:
         """Format plan as text."""

aury/agents/tool/builtin/thinking.py

@@ -92,11 +92,6 @@ This makes your reasoning visible and traceable."""
         invocation_id: int | None = None,
     ) -> None:
         """Emit THINKING block."""
-        emit = getattr(ctx, 'emit', None)
-        if emit is None:
-            logger.debug(f"ThinkingTool: emit not available, invocation_id={invocation_id}")
-            return
-        
         logger.debug(f"ThinkingTool: emitting block, category={category}, invocation_id={invocation_id}")
         
         block = BlockEvent(
@@ -110,7 +105,7 @@ This makes your reasoning visible and traceable."""
             invocation_id=ctx.invocation_id,
         )
         
-        await emit(block)
+        await self.emit(block)
         logger.debug(f"ThinkingTool: block emitted successfully, invocation_id={invocation_id}")
 
 

aury/agents/tool/builtin/yield_result.py

@@ -106,11 +106,6 @@ This will end your current session and resume the parent."""
         invocation_id: int | None = None,
     ) -> None:
         """Emit YIELD block."""
-        emit = getattr(ctx, 'emit', None)
-        if emit is None:
-            logger.debug(f"YieldResultTool: emit not available, invocation_id={invocation_id}")
-            return
-        
         logger.debug(f"YieldResultTool: emitting yield block, status={status}, invocation_id={invocation_id}")
         
         block = BlockEvent(
@@ -126,7 +121,7 @@ This will end your current session and resume the parent."""
             invocation_id=ctx.invocation_id,
         )
         
-        await emit(block)
+        await self.emit(block)
         logger.debug(f"YieldResultTool: yield block emitted successfully, invocation_id={invocation_id}")