aury-agent 0.0.6__py3-none-any.whl → 0.0.7__py3-none-any.whl

aury/agents/middleware/message_container.py

@@ -49,10 +49,6 @@ class MessageContainerMiddleware(BaseMiddleware):
                        parent_id. Defaults to {"thinking", "text"}.
     """
     
-    # Key to store the token in middleware context
-    _TOKEN_KEY = "_message_container_token"
-    _BLOCK_ID_KEY = "_message_container_block_id"
-    
     # Default kinds that should be grouped under message container
     DEFAULT_KINDS = {"thinking", "text"}
     
@@ -64,13 +60,17 @@ class MessageContainerMiddleware(BaseMiddleware):
         """
         super().__init__()
         self.apply_to_kinds = apply_to_kinds or self.DEFAULT_KINDS
+        self._token: Any = None
+        self._block_id: str | None = None
     
     async def on_request(
         self,
         request: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Create message container block and set parent_id."""
+        from ..core.context import get_current_ctx_or_none
+        ctx = get_current_ctx_or_none()
+        
         # Generate container block ID
         message_block_id = generate_id("blk")
         
@@ -86,41 +86,33 @@ class MessageContainerMiddleware(BaseMiddleware):
             op=BlockOp.APPLY,
             data={
                 "type": "llm_response",
-                "step": context.get("step"),
             },
-            session_id=context.get("session_id"),
-            invocation_id=context.get("invocation_id"),
+            session_id=ctx.session_id if ctx else None,
+            invocation_id=ctx.invocation_id if ctx else None,
         ))
         
         # Set parent_id in ContextVar with apply_to_kinds filter
-        # Only blocks matching apply_to_kinds will inherit this parent_id
-        token = set_parent_id(message_block_id, apply_to_kinds=self.apply_to_kinds)
-        context[self._TOKEN_KEY] = token
-        context[self._BLOCK_ID_KEY] = message_block_id
+        self._token = set_parent_id(message_block_id, apply_to_kinds=self.apply_to_kinds)
+        self._block_id = message_block_id
         
         return request
     
     async def on_response(
         self,
         response: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Reset parent_id to previous value."""
-        token = context.get(self._TOKEN_KEY)
-        if token is not None:
-            reset_parent_id(token)
+        if self._token is not None:
+            reset_parent_id(self._token)
+            self._token = None
         return response
     
     async def on_error(
         self,
         error: Exception,
-        context: dict[str, Any],
     ) -> Exception | None:
         """Reset parent_id on error too."""
-        token = context.get(self._TOKEN_KEY)
-        if token is not None:
-            reset_parent_id(token)
+        if self._token is not None:
+            reset_parent_id(self._token)
+            self._token = None
         return error
-
-
-__all__ = ["MessageContainerMiddleware"]

aury/agents/middleware/raw_message.py

@@ -70,18 +70,18 @@ class RawMessageMiddleware(BaseMiddleware):
     async def on_message_save(
         self,
         message: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Store complete message to RawMessageStore.
         
         Args:
             message: Complete message dict with 'role', 'content', etc.
-            context: Execution context with 'invocation_id', etc.
             
         Returns:
             The message with added 'raw_msg_id' field
         """
-        invocation_id = context.get("invocation_id", "")
+        from ..core.context import get_current_ctx_or_none
+        ctx = get_current_ctx_or_none()
+        invocation_id = ctx.invocation_id if ctx else ""
         if not invocation_id:
             return message
         
@@ -108,7 +108,6 @@ class RawMessageMiddleware(BaseMiddleware):
         self,
         agent_id: str,
         result: Any,
-        context: dict[str, Any],
     ) -> HookResult:
         """Clean up raw messages when invocation completes.
         
@@ -117,7 +116,9 @@ class RawMessageMiddleware(BaseMiddleware):
         if self.persist_raw:
             return HookResult.proceed()
         
-        invocation_id = context.get("invocation_id", "")
+        from ..core.context import get_current_ctx_or_none
+        ctx = get_current_ctx_or_none()
+        invocation_id = ctx.invocation_id if ctx else ""
         if invocation_id:
             await self._cleanup_invocation(invocation_id)
         

aury/agents/middleware/truncation.py

@@ -48,13 +48,11 @@ class MessageTruncationMiddleware(BaseMiddleware):
     async def on_message_save(
         self,
         message: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Truncate message content before saving.
         
         Args:
             message: Message dict with 'role', 'content', etc.
-            context: Execution context
             
         Returns:
             Modified message with truncated content

aury/agents/react/agent.py

@@ -93,6 +93,7 @@ class ReactAgent(BaseAgent):
         enable_history: bool = True,
         history_limit: int = 50,
         delegate_tool_class: "type[BaseTool] | None" = None,
+        context_metadata: dict | None = None,
     ) -> "ReactAgent":
         """Create ReactAgent with minimal boilerplate. See factory.create_react_agent for details."""
         return create_react_agent(
@@ -110,6 +111,7 @@ class ReactAgent(BaseAgent):
             enable_history=enable_history,
             history_limit=history_limit,
             delegate_tool_class=delegate_tool_class,
+            context_metadata=context_metadata,
         )
     
     @classmethod
@@ -263,41 +265,38 @@ class ReactAgent(BaseAgent):
             "Starting ReactAgent run",
             extra={
                 "session_id": self.session.id,
-                "agent": self.name,
+                "agent_id": self.ctx.agent_id,
             }
         )
 
-        # Build middleware context
-        from ..core.context import emit as global_emit
-        mw_context = {
-            "session_id": self.session.id,
-            "agent_id": self.name,
-            "agent_type": self.agent_type,
-            "emit": global_emit,
-            "backends": self.ctx.backends,
-        }
-
         try:
-            # Create new invocation
+            # Create new invocation using ctx.invocation_id
             self._current_invocation = Invocation(
-                id=generate_id("inv"),
+                id=self.ctx.invocation_id,
                 session_id=self.session.id,
+                agent_id=self.ctx.agent_id,
                 state=InvocationState.RUNNING,
                 started_at=datetime.now(),
             )
-            mw_context["invocation_id"] = self._current_invocation.id
 
             logger.info("Created invocation", extra={"invocation_id": self._current_invocation.id})
 
+            # Persist invocation immediately (so we have record even if agent fails)
+            if self.ctx.backends and self.ctx.backends.invocation:
+                await self.ctx.backends.invocation.create(
+                    self._current_invocation.id,
+                    self.session.id,
+                    self._current_invocation.to_dict(),
+                    agent_id=self.ctx.agent_id,
+                )
+
             # === Middleware: on_agent_start ===
             if self.middleware:
                 logger.info(
                     "Calling middleware: on_agent_start",
                     extra={"invocation_id": self._current_invocation.id},
                 )
-                hook_result = await self.middleware.process_agent_start(
-                    self.name, input, mw_context
-                )
+                hook_result = await self.middleware.process_agent_start(input)
                 if hook_result.action == HookAction.STOP:
                     logger.warning("Agent stopped by middleware on_agent_start", extra={"invocation_id": self._current_invocation.id})
                     await self.ctx.emit(BlockEvent(
@@ -449,6 +448,13 @@ class ReactAgent(BaseAgent):
             
             # Complete invocation
             if is_aborted:
+                # Save current buffer content before marking as aborted
+                if self._text_buffer or self._thinking_buffer or self._tool_invocations:
+                    await persist_helpers.save_assistant_message(self)
+                # Save completed tool results
+                if self._tool_invocations:
+                    await persist_helpers.save_tool_messages(self)
+                # Mark invocation as aborted
                 self._current_invocation.state = InvocationState.ABORTED
                 logger.info(
                     "Invocation aborted by user",
@@ -490,9 +496,7 @@ class ReactAgent(BaseAgent):
             # === Middleware: on_agent_end ===
             if self.middleware:
                 await self.middleware.process_agent_end(
-                    self.name,
                     {"steps": self._current_step, "finish_reason": finish_reason},
-                    mw_context,
                 )
 
             await self.bus.publish(
@@ -551,7 +555,7 @@ class ReactAgent(BaseAgent):
 
             # === Middleware: on_error ===
             if self.middleware:
-                processed_error = await self.middleware.process_error(e, mw_context)
+                processed_error = await self.middleware.process_error(e)
                 if processed_error is None:
                     logger.warning(
                         "Error suppressed by middleware",

aury/agents/react/factory.py

@@ -44,6 +44,8 @@ def create_react_agent(
     history_limit: int = 50,
     # Tool customization
     delegate_tool_class: "type[BaseTool] | None" = None,
+    # Context metadata
+    context_metadata: dict | None = None,
 ) -> "ReactAgent":
     """Create ReactAgent with minimal boilerplate.
     
@@ -159,16 +161,26 @@ def create_react_agent(
     all_providers = default_providers + (context_providers or [])
     
     # Build context
+    # agent_id: use config.id if provided, fallback to config.code, then "react_agent"
+    # agent_name: use config.name (display name)
+    agent_id = (
+        config.id if config and config.id 
+        else (config.code if config and config.code else "react_agent")
+    )
+    agent_name = config.name if config else None
+    
     ctx = InvocationContext(
         session=session,
         invocation_id=generate_id("inv"),
-        agent_id=config.name if config else "react_agent",
+        agent_id=agent_id,
+        agent_name=agent_name,
         backends=backends,
         bus=bus,
         llm=llm,
         middleware=middleware_chain,
         memory=memory,
         snapshot=snapshot,
+        metadata=context_metadata or {},
     )
     
     agent = ReactAgent(ctx, config)

aury/agents/react/persistence.py

@@ -81,14 +81,7 @@ async def trigger_message_save(agent: "ReactAgent", message: dict) -> dict | Non
     if not agent.middleware:
         return message
     
-    namespace = getattr(agent, '_message_namespace', None)
-    mw_context = {
-        "session_id": agent.session.id,
-        "agent_id": agent.name,
-        "namespace": namespace,
-    }
-    
-    return await agent.middleware.process_message_save(message, mw_context)
+    return await agent.middleware.process_message_save(message)
 
 
 async def save_user_message(agent: "ReactAgent", input: "PromptInput") -> None:

aury/agents/react/step.py

@@ -190,6 +190,7 @@ async def execute_step(agent: "ReactAgent") -> str | None:
     agent._text_buffer = ""
     agent._thinking_buffer = ""  # Buffer for non-streaming thinking
     agent._tool_invocations = []
+    agent._last_usage = None  # Store usage for middleware
     
     # Reset block IDs for this step (each step gets fresh block IDs)
     agent._current_text_block_id = None
@@ -199,20 +200,9 @@ async def execute_step(agent: "ReactAgent") -> str | None:
     agent._call_id_to_tool = {}
     agent._tool_call_blocks = {}
     
-    # Track accumulated arguments for streaming tool calls (for middleware context)
+    # Track accumulated arguments for streaming tool calls
     tool_call_accumulated_args: dict[str, dict[str, Any]] = {}
 
-    # Build middleware context for this step
-    mw_context = {
-        "session_id": agent.session.id,
-        "invocation_id": agent._current_invocation.id if agent._current_invocation else "",
-        "step": agent._current_step,
-        "agent_id": agent.name,
-        "emit": global_emit,  # For middleware to emit BlockEvent/ActionEvent
-        "backends": agent.ctx.backends,
-        "tool_mode": effective_tool_mode.value,  # Add tool mode to context
-    }
-
     # Build LLM call kwargs
     # Note: temperature, max_tokens, timeout, retries are configured on LLMProvider
     llm_kwargs: dict[str, Any] = {
@@ -244,7 +234,7 @@ async def execute_step(agent: "ReactAgent") -> str | None:
             "Calling middleware: on_request",
             extra={"invocation_id": agent._current_invocation.id},
         )
-        llm_kwargs = await agent.middleware.process_request(llm_kwargs, mw_context)
+        llm_kwargs = await agent.middleware.process_request(llm_kwargs)
         if llm_kwargs is None:
             logger.warning(
                 "LLM request cancelled by middleware",
@@ -307,7 +297,7 @@ async def execute_step(agent: "ReactAgent") -> str | None:
                 stream_chunk = {"delta": event.delta, "type": "content"}
                 if agent.middleware:
                     stream_chunk = await agent.middleware.process_stream_chunk(
-                        stream_chunk, mw_context
+                        stream_chunk
                     )
                     if stream_chunk is None:
                         continue  # Skip this chunk
@@ -342,7 +332,7 @@ async def execute_step(agent: "ReactAgent") -> str | None:
                 stream_chunk = {"delta": event.delta, "type": "thinking"}
                 if agent.middleware:
                     stream_chunk = await agent.middleware.process_stream_chunk(
-                        stream_chunk, mw_context
+                        stream_chunk
                     )
                     if stream_chunk is None:
                         continue  # Skip this chunk
@@ -450,12 +440,9 @@ async def execute_step(agent: "ReactAgent") -> str | None:
                             # === Middleware: on_tool_call_delta ===
                             processed_delta = arguments_delta
                             if agent.middleware:
-                                delta_context = {
-                                    **mw_context,
-                                    "accumulated_args": tool_call_accumulated_args.get(call_id, {}),
-                                }
+                                accumulated_args = tool_call_accumulated_args.get(call_id, {})
                                 processed_delta = await agent.middleware.process_tool_call_delta(
-                                    call_id, tool_name, arguments_delta, delta_context
+                                    call_id, tool_name, arguments_delta, accumulated_args
                                 )
                                 if processed_delta is None:
                                     continue  # Skip this delta
@@ -539,17 +526,19 @@ async def execute_step(agent: "ReactAgent") -> str | None:
 
         elif event.type == "usage":
             if event.usage:
+                # Store usage for middleware
+                agent._last_usage = {
+                    "provider": agent.llm.provider,
+                    "model": agent.llm.model,
+                    "input_tokens": event.usage.input_tokens,
+                    "output_tokens": event.usage.output_tokens,
+                    "cache_read_tokens": event.usage.cache_read_tokens,
+                    "cache_write_tokens": event.usage.cache_write_tokens,
+                    "reasoning_tokens": event.usage.reasoning_tokens,
+                }
                 await agent.bus.publish(
                     Events.USAGE_RECORDED,
-                    {
-                        "provider": agent.llm.provider,
-                        "model": agent.llm.model,
-                        "input_tokens": event.usage.input_tokens,
-                        "output_tokens": event.usage.output_tokens,
-                        "cache_read_tokens": event.usage.cache_read_tokens,
-                        "cache_write_tokens": event.usage.cache_write_tokens,
-                        "reasoning_tokens": event.usage.reasoning_tokens,
-                    },
+                    agent._last_usage,
                 )
 
         elif event.type == "error":
@@ -624,6 +613,7 @@ async def execute_step(agent: "ReactAgent") -> str | None:
         "thinking": agent._thinking_buffer,
         "tool_calls": len(agent._tool_invocations),
         "finish_reason": finish_reason,
+        "usage": agent._last_usage,  # Include usage for middleware
     }
     if agent.middleware:
         logger.debug(
@@ -635,7 +625,7 @@ async def execute_step(agent: "ReactAgent") -> str | None:
             },
         )
         llm_response_data = await agent.middleware.process_response(
-            llm_response_data, mw_context
+            llm_response_data
         )
 
     await agent.bus.publish(

aury/agents/react/tools.py

@@ -49,6 +49,19 @@ async def execute_tool(agent: "ReactAgent", invocation: ToolInvocation) -> ToolR
     Returns:
         ToolResult from tool execution
     """
+    # 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)
+        logger.info(
+            f"Tool aborted before execution: {invocation.tool_name}",
+            extra={
+                "invocation_id": agent._current_invocation.id if agent._current_invocation else "",
+                "call_id": invocation.tool_call_id,
+            },
+        )
+        return ToolResult.error(error_msg)
+    
     invocation.mark_call_complete()
 
     logger.info(
@@ -60,14 +73,6 @@ async def execute_tool(agent: "ReactAgent", invocation: ToolInvocation) -> ToolR
         },
     )
 
-    # Build middleware context
-    mw_context = {
-        "session_id": agent.session.id,
-        "invocation_id": agent._current_invocation.id if agent._current_invocation else "",
-        "tool_call_id": invocation.tool_call_id,
-        "agent_id": agent.name,
-    }
-
     try:
         # Get tool from agent context
         tool = get_tool(agent, invocation.tool_name)
@@ -87,7 +92,7 @@ async def execute_tool(agent: "ReactAgent", invocation: ToolInvocation) -> ToolR
                 extra={"invocation_id": agent._current_invocation.id, "call_id": invocation.tool_call_id},
             )
             hook_result = await agent.middleware.process_tool_call(
-                tool, invocation.args, mw_context
+                tool, invocation.args
             )
             if hook_result.action == HookAction.SKIP:
                 logger.warning(
@@ -136,7 +141,7 @@ async def execute_tool(agent: "ReactAgent", invocation: ToolInvocation) -> ToolR
                 f"Calling middleware: on_tool_end ({invocation.tool_name})",
                 extra={"invocation_id": agent._current_invocation.id},
             )
-            hook_result = await agent.middleware.process_tool_end(tool, result, mw_context)
+            hook_result = await agent.middleware.process_tool_end(tool, result)
             if hook_result.action == HookAction.RETRY:
                 logger.info(
                     f"Tool {invocation.tool_name} retry requested by middleware",
@@ -212,16 +217,44 @@ async def process_tool_results(agent: "ReactAgent") -> None:
         },
     )
 
+    # Check abort before starting tool execution
+    if await agent._check_abort():
+        logger.info(
+            "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)
+        
+        # Check abort after parallel execution - cancel remaining if aborted
+        if await agent._check_abort():
+            logger.info(
+                "Tool execution aborted after parallel execution",
+                extra={"invocation_id": agent._current_invocation.id},
+            )
     else:
-        # Sequential execution
+        # Sequential execution with abort check between tools
         results = []
         for inv in agent._tool_invocations:
+            # Check abort before each tool
+            if await agent._check_abort():
+                logger.info(
+                    f"Tool execution aborted before {inv.tool_name}",
+                    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)
+                continue
+                
             try:
                 result = await execute_tool(agent, inv)
                 results.append(result)

aury/agents/tool/builtin/delegate.py

@@ -153,12 +153,7 @@ Specify the agent key and task data."""
         
         # Get dynamic agents from middleware (progressive disclosure)
         if self.middleware and ctx:
-            mw_context = {
-                "session_id": ctx.session_id,
-                "invocation_id": ctx.invocation_id,
-                "agent_id": ctx.agent,
-            }
-            dynamic_agents = await self.middleware.get_dynamic_subagents(mw_context)
+            dynamic_agents = await self.middleware.get_dynamic_subagents()
             if dynamic_agents:
                 # Store dynamic agents for later lookup
                 for config in dynamic_agents:

aury/agents/tool/decorator.py

@@ -53,7 +53,8 @@ def _type_to_schema(t: type) -> dict[str, Any]:
         bool: {"type": "boolean"},
         list: {"type": "array"},
         dict: {"type": "object"},
-        Any: {},
+        # Any type: JSON Schema 2020-12 requires valid type, use object as fallback
+        Any: {"type": "object"},
     }
     
     return type_map.get(t, {"type": "string"})

aury/agents/workflow/adapter.py

@@ -162,21 +162,10 @@ class WorkflowAgent(BaseAgent):
         inputs = input if isinstance(input, dict) else {"input": input}
         logger.info(f"WorkflowAgent executing, workflow={self.workflow.spec.name}, invocation_id={self.ctx.invocation_id}")
         
-        # Build middleware context
-        mw_context = {
-            "session_id": self.session.id,
-            "invocation_id": self.ctx.invocation_id,
-            "agent_id": self.name,
-            "agent_type": self.agent_type,
-            "workflow_name": self.workflow.spec.name,
-        }
-        
         # === Middleware: on_agent_start ===
         if self.middleware:
             logger.debug(f"WorkflowAgent: processing on_agent_start hooks, invocation_id={self.ctx.invocation_id}")
-            hook_result = await self.middleware.process_agent_start(
-                self.name, inputs, mw_context
-            )
+            hook_result = await self.middleware.process_agent_start(inputs)
             if hook_result.action == HookAction.STOP:
                 logger.warning(f"Workflow stopped by middleware on_agent_start, invocation_id={self.ctx.invocation_id}")
                 await self.ctx.emit(BlockEvent(
@@ -205,16 +194,14 @@ class WorkflowAgent(BaseAgent):
             # === Middleware: on_agent_end ===
             if self.middleware:
                 logger.debug(f"WorkflowAgent: processing on_agent_end hooks, invocation_id={self.ctx.invocation_id}")
-                await self.middleware.process_agent_end(
-                    self.name, result, mw_context
-                )
+                await self.middleware.process_agent_end(result)
         
         except Exception as e:
             logger.error(f"WorkflowAgent: execution error, error={type(e).__name__}, workflow={self.workflow.spec.name}, invocation_id={self.ctx.invocation_id}", exc_info=True)
             # === Middleware: on_error ===
             if self.middleware:
                 logger.debug(f"WorkflowAgent: processing on_error hooks, invocation_id={self.ctx.invocation_id}")
-                processed_error = await self.middleware.process_error(e, mw_context)
+                processed_error = await self.middleware.process_error(e)
                 if processed_error is None:
                     logger.warning(f"WorkflowAgent: error suppressed by middleware, invocation_id={self.ctx.invocation_id}")
                     return

aury/agents/workflow/executor.py

@@ -598,24 +598,12 @@ class WorkflowExecutor:
         # Get effective middleware for this node
         effective_middleware = self._get_effective_middleware(node)
         
-        # Build middleware context
-        mw_context = {
-            "session_id": self.ctx.session_id,
-            "invocation_id": self.ctx.invocation_id,
-            "parent_agent_id": self.workflow.spec.name,
-            "child_agent_id": node.agent,
-            "node_id": node.id,
-            "parent_block_id": parent_block_id,
-            "has_node_middleware": bool(node.middleware),
-        }
-        
         # === Middleware: on_subagent_start ===
         if effective_middleware:
             hook_result = await effective_middleware.process_subagent_start(
                 self.workflow.spec.name,
                 node.agent,
                 "embedded",  # Workflow nodes are embedded execution
-                mw_context,
             )
             if hook_result.action == HookAction.SKIP:
                 logger.info(f"SubAgent {node.agent} skipped by middleware")
@@ -628,8 +616,13 @@ class WorkflowExecutor:
         try:
             # Create child context for sub-agent with effective middleware
             # Note: parent_block_id is already set via ContextVar above
+            # Get agent name from factory if available
+            agent_class = self.agent_factory.get_class(node.agent)
+            agent_name = getattr(agent_class, 'name', node.agent) if agent_class else node.agent
+            
             child_ctx = self.ctx.create_child(
                 agent_id=node.agent,
+                agent_name=agent_name,
                 middleware=effective_middleware,
             )
             
@@ -662,7 +655,6 @@ class WorkflowExecutor:
                     self.workflow.spec.name,
                     node.agent,
                     result,
-                    mw_context,
                 )
             
             return result

{aury_agent-0.0.6.dist-info → aury_agent-0.0.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aury-agent
-Version: 0.0.6
+Version: 0.0.7
 Summary: Aury Agent Framework - React Agent and Workflow orchestration
 Author: Aury Team
 License: MIT