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

aury/agents/middleware/base.py

@@ -1,4 +1,11 @@
-"""Middleware protocol and base implementation."""
+"""Middleware protocol and base implementation.
+
+Middleware can access InvocationContext via get_current_ctx_or_none() for:
+- session_id, invocation_id, agent_id, agent_name
+- backends, metadata, etc.
+
+Middleware should use self._xxx for internal state between hooks.
+"""
 from __future__ import annotations
 
 from typing import Any, Protocol, runtime_checkable, TYPE_CHECKING
@@ -14,6 +21,7 @@ class Middleware(Protocol):
     """Middleware protocol for request/response processing.
     
     Includes both LLM request/response hooks and agent lifecycle hooks.
+    Use get_current_ctx_or_none() to access InvocationContext.
     """
     
     @property
@@ -26,13 +34,11 @@ class Middleware(Protocol):
     async def on_request(
         self,
         request: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Process request before LLM call.
         
         Args:
             request: The request to process
-            context: Execution context
             
         Returns:
             Modified request, or None to skip further processing
@@ -42,13 +48,11 @@ class Middleware(Protocol):
     async def on_response(
         self,
         response: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Process response after LLM call.
         
         Args:
             response: The response to process
-            context: Execution context
             
         Returns:
             Modified response, or None to skip further processing
@@ -58,13 +62,11 @@ class Middleware(Protocol):
     async def on_error(
         self,
         error: Exception,
-        context: dict[str, Any],
     ) -> Exception | None:
         """Handle errors.
         
         Args:
             error: The exception that occurred
-            context: Execution context
             
         Returns:
             Modified exception, or None to suppress
@@ -74,13 +76,11 @@ class Middleware(Protocol):
     async def on_model_stream(
         self,
         chunk: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Process streaming chunk (triggered by trigger_mode).
         
         Args:
             chunk: The streaming chunk
-            context: Execution context
             
         Returns:
             Modified chunk, or None to skip further processing
@@ -90,13 +90,11 @@ class Middleware(Protocol):
     async def on_thinking_stream(
         self,
         chunk: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Process thinking stream chunk.
         
         Args:
             chunk: The thinking chunk with {"delta": str}
-            context: Execution context
             
         Returns:
             Modified chunk, or None to skip
@@ -107,16 +105,14 @@ class Middleware(Protocol):
     
     async def on_agent_start(
         self,
-        agent_id: str,
         input_data: Any,
-        context: dict[str, Any],
     ) -> HookResult:
         """Called when agent starts processing.
         
+        Use get_current_ctx_or_none() to access agent_id, session_id, etc.
+        
         Args:
-            agent_id: The agent identifier
             input_data: Input to the agent
-            context: Execution context
             
         Returns:
             HookResult controlling execution flow
@@ -125,16 +121,14 @@ class Middleware(Protocol):
     
     async def on_agent_end(
         self,
-        agent_id: str,
         result: Any,
-        context: dict[str, Any],
     ) -> HookResult:
         """Called when agent completes processing.
         
+        Use get_current_ctx_or_none() to access agent_id, session_id, etc.
+        
         Args:
-            agent_id: The agent identifier
             result: Agent's result
-            context: Execution context
             
         Returns:
             HookResult (only CONTINUE/STOP meaningful here)
@@ -145,14 +139,12 @@ class Middleware(Protocol):
         self,
         tool: "BaseTool",
         params: dict[str, Any],
-        context: dict[str, Any],
     ) -> HookResult:
         """Called before tool execution.
         
         Args:
             tool: The tool to be called
             params: Tool parameters
-            context: Execution context
             
         Returns:
             HookResult - SKIP to skip tool, RETRY to modify params
@@ -164,7 +156,7 @@ class Middleware(Protocol):
         call_id: str,
         tool_name: str,
         delta: dict[str, Any],
-        context: dict[str, Any],
+        accumulated_args: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Called during streaming tool argument generation.
         
@@ -175,7 +167,7 @@ class Middleware(Protocol):
             call_id: Tool call identifier
             tool_name: Name of the tool being called
             delta: Incremental parameter update (e.g. {"content": "more text"})
-            context: Execution context with 'accumulated_args' containing current state
+            accumulated_args: Current accumulated arguments state
             
         Returns:
             Modified delta, or None to skip emitting this delta
@@ -186,14 +178,12 @@ class Middleware(Protocol):
         self,
         tool: "BaseTool",
         result: "ToolResult",
-        context: dict[str, Any],
     ) -> HookResult:
         """Called after tool execution.
         
         Args:
             tool: The tool that was called
             result: Tool execution result
-            context: Execution context
             
         Returns:
             HookResult - RETRY to re-execute tool
@@ -205,7 +195,6 @@ class Middleware(Protocol):
         parent_agent_id: str,
         child_agent_id: str,
         mode: str,  # "embedded" or "delegated"
-        context: dict[str, Any],
     ) -> HookResult:
         """Called when delegating to a sub-agent.
         
@@ -213,7 +202,6 @@ class Middleware(Protocol):
             parent_agent_id: Parent agent identifier
             child_agent_id: Child agent identifier
             mode: Delegation mode
-            context: Execution context
             
         Returns:
             HookResult - SKIP to skip delegation
@@ -225,7 +213,6 @@ class Middleware(Protocol):
         parent_agent_id: str,
         child_agent_id: str,
         result: Any,
-        context: dict[str, Any],
     ) -> HookResult:
         """Called when sub-agent completes.
         
@@ -233,7 +220,6 @@ class Middleware(Protocol):
             parent_agent_id: Parent agent identifier
             child_agent_id: Child agent identifier
             result: Sub-agent's result
-            context: Execution context
             
         Returns:
             HookResult (for post-processing)
@@ -243,7 +229,6 @@ class Middleware(Protocol):
     async def on_message_save(
         self,
         message: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Called before saving a message to history.
         
@@ -252,7 +237,6 @@ class Middleware(Protocol):
         
         Args:
             message: Message dict with 'role', 'content', etc.
-            context: Execution context
             
         Returns:
             Modified message, or None to skip saving
@@ -265,6 +249,9 @@ class BaseMiddleware:
     
     Subclass and override specific hooks as needed.
     All hooks have sensible pass-through defaults.
+    
+    Use get_current_ctx_or_none() to access InvocationContext.
+    Use self._xxx for internal state between hooks.
     """
     
     _config: MiddlewareConfig = MiddlewareConfig()
@@ -278,7 +265,6 @@ class BaseMiddleware:
     async def on_request(
         self,
         request: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Default: pass through."""
         return request
@@ -286,7 +272,6 @@ class BaseMiddleware:
     async def on_response(
         self,
         response: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Default: pass through."""
         return response
@@ -294,7 +279,6 @@ class BaseMiddleware:
     async def on_error(
         self,
         error: Exception,
-        context: dict[str, Any],
     ) -> Exception | None:
         """Default: re-raise error."""
         return error
@@ -302,7 +286,6 @@ class BaseMiddleware:
     async def on_model_stream(
         self,
         chunk: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Default: pass through."""
         return chunk
@@ -310,7 +293,6 @@ class BaseMiddleware:
     async def on_thinking_stream(
         self,
         chunk: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Default: pass through."""
         return chunk
@@ -319,18 +301,14 @@ class BaseMiddleware:
     
     async def on_agent_start(
         self,
-        agent_id: str,
         input_data: Any,
-        context: dict[str, Any],
     ) -> HookResult:
         """Default: continue."""
         return HookResult.proceed()
     
     async def on_agent_end(
         self,
-        agent_id: str,
         result: Any,
-        context: dict[str, Any],
     ) -> HookResult:
         """Default: continue."""
         return HookResult.proceed()
@@ -339,7 +317,6 @@ class BaseMiddleware:
         self,
         tool: "BaseTool",
         params: dict[str, Any],
-        context: dict[str, Any],
     ) -> HookResult:
         """Default: continue."""
         return HookResult.proceed()
@@ -349,7 +326,7 @@ class BaseMiddleware:
         call_id: str,
         tool_name: str,
         delta: dict[str, Any],
-        context: dict[str, Any],
+        accumulated_args: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Default: pass through."""
         return delta
@@ -358,7 +335,6 @@ class BaseMiddleware:
         self,
         tool: "BaseTool",
         result: "ToolResult",
-        context: dict[str, Any],
     ) -> HookResult:
         """Default: continue."""
         return HookResult.proceed()
@@ -368,7 +344,6 @@ class BaseMiddleware:
         parent_agent_id: str,
         child_agent_id: str,
         mode: str,
-        context: dict[str, Any],
     ) -> HookResult:
         """Default: continue."""
         return HookResult.proceed()
@@ -378,7 +353,6 @@ class BaseMiddleware:
         parent_agent_id: str,
         child_agent_id: str,
         result: Any,
-        context: dict[str, Any],
     ) -> HookResult:
         """Default: continue."""
         return HookResult.proceed()
@@ -386,7 +360,6 @@ class BaseMiddleware:
     async def on_message_save(
         self,
         message: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Default: pass through."""
         return message

aury/agents/middleware/chain.py

@@ -112,14 +112,13 @@ class MiddlewareChain:
     async def process_request(
         self,
         request: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Process request through all middlewares."""
         current = request
         logger.debug(f"Processing request through {len(self._middlewares)} middlewares")
         
         for i, mw in enumerate(self._middlewares):
-            result = await mw.on_request(current, context)
+            result = await mw.on_request(current)
             if result is None:
                 logger.info(f"Middleware #{i} blocked request")
                 return None
@@ -131,14 +130,13 @@ class MiddlewareChain:
     async def process_response(
         self,
         response: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Process response through all middlewares (reverse order)."""
         current = response
         logger.debug(f"Processing response through {len(self._middlewares)} middlewares (reverse order)")
         
         for i, mw in enumerate(reversed(self._middlewares)):
-            result = await mw.on_response(current, context)
+            result = await mw.on_response(current)
             if result is None:
                 logger.info(f"Middleware #{i} blocked response")
                 return None
@@ -150,14 +148,13 @@ class MiddlewareChain:
     async def process_error(
         self,
         error: Exception,
-        context: dict[str, Any],
     ) -> Exception | None:
         """Process error through all middlewares."""
         current = error
         logger.debug(f"Processing error {type(error).__name__} through {len(self._middlewares)} middlewares")
         
         for i, mw in enumerate(self._middlewares):
-            result = await mw.on_error(current, context)
+            result = await mw.on_error(current)
             if result is None:
                 logger.info(f"Middleware #{i} suppressed error")
                 return None
@@ -169,13 +166,11 @@ class MiddlewareChain:
     async def process_stream_chunk(
         self,
         chunk: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Process streaming chunk through middlewares based on trigger mode."""
         text = chunk.get("text", chunk.get("delta", ""))
         self._token_buffer += text
         self._token_count += 1
-        logger.debug(f"Processing stream chunk, token_count={self._token_count}, triggered_middlewares=?")
         
         current = chunk
         triggered_count = 0
@@ -185,14 +180,16 @@ class MiddlewareChain:
             
             if should_trigger:
                 triggered_count += 1
-                logger.debug(f"Middleware #{i} triggered, mode={mw.config.trigger_mode}")
-                result = await mw.on_model_stream(current, context)
+                result = await mw.on_model_stream(current)
                 if result is None:
                     logger.info(f"Middleware #{i} blocked stream chunk")
                     return None
                 current = result
         
-        logger.debug(f"Stream chunk processing completed, {triggered_count} middlewares triggered")
+        # Log only every 50 tokens to reduce noise
+        if self._token_count % 50 == 0:
+            logger.debug(f"Stream progress: token_count={self._token_count}, middlewares={len(self._middlewares)}")
+        
         return current
     
     async def process_tool_call_delta(
@@ -200,7 +197,7 @@ class MiddlewareChain:
         call_id: str,
         tool_name: str,
         delta: dict[str, Any],
-        context: dict[str, Any],
+        accumulated_args: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Process tool call delta through all middlewares.
         
@@ -208,7 +205,7 @@ class MiddlewareChain:
             call_id: Tool call identifier
             tool_name: Name of the tool being called
             delta: Incremental parameter update
-            context: Execution context with 'accumulated_args'
+            accumulated_args: Current accumulated arguments state
             
         Returns:
             Modified delta, or None to skip emitting
@@ -217,7 +214,7 @@ class MiddlewareChain:
         logger.debug(f"Processing tool_call_delta for {tool_name} (call_id={call_id}) through {len(self._middlewares)} middlewares")
         
         for i, mw in enumerate(self._middlewares):
-            result = await mw.on_tool_call_delta(call_id, tool_name, current, context)
+            result = await mw.on_tool_call_delta(call_id, tool_name, current, accumulated_args)
             if result is None:
                 logger.info(f"Middleware #{i} blocked tool_call_delta")
                 return None
@@ -259,19 +256,17 @@ class MiddlewareChain:
     
     async def process_agent_start(
         self,
-        agent_id: str,
         input_data: Any,
-        context: dict[str, Any],
     ) -> HookResult:
         """Process agent start through all middlewares.
         
         Returns:
             First non-CONTINUE result, or CONTINUE if all pass
         """
-        logger.debug(f"Processing agent_start for agent_id={agent_id}, {len(self._middlewares)} middlewares")
+        logger.debug(f"Processing agent_start, {len(self._middlewares)} middlewares")
         for i, mw in enumerate(self._middlewares):
             if hasattr(mw, 'on_agent_start'):
-                result = await mw.on_agent_start(agent_id, input_data, context)
+                result = await mw.on_agent_start(input_data)
                 if result.action != HookAction.CONTINUE:
                     logger.info(f"Middleware #{i} returned {result.action} on agent_start")
                     return result
@@ -280,15 +275,13 @@ class MiddlewareChain:
     
     async def process_agent_end(
         self,
-        agent_id: str,
         result: Any,
-        context: dict[str, Any],
     ) -> HookResult:
         """Process agent end through all middlewares (reverse order)."""
-        logger.debug(f"Processing agent_end for agent_id={agent_id}, {len(self._middlewares)} middlewares (reverse order)")
+        logger.debug(f"Processing agent_end, {len(self._middlewares)} middlewares (reverse order)")
         for i, mw in enumerate(reversed(self._middlewares)):
             if hasattr(mw, 'on_agent_end'):
-                hook_result = await mw.on_agent_end(agent_id, result, context)
+                hook_result = await mw.on_agent_end(result)
                 if hook_result.action != HookAction.CONTINUE:
                     logger.info(f"Middleware #{i} returned {hook_result.action} on agent_end")
                     return hook_result
@@ -299,7 +292,6 @@ class MiddlewareChain:
         self,
         tool: "BaseTool",
         params: dict[str, Any],
-        context: dict[str, Any],
     ) -> HookResult:
         """Process tool call through all middlewares.
         
@@ -309,7 +301,7 @@ class MiddlewareChain:
         logger.debug(f"Processing tool_call for tool={tool.name}, {len(self._middlewares)} middlewares")
         for i, mw in enumerate(self._middlewares):
             if hasattr(mw, 'on_tool_call'):
-                result = await mw.on_tool_call(tool, params, context)
+                result = await mw.on_tool_call(tool, params)
                 if result.action != HookAction.CONTINUE:
                     logger.info(f"Middleware #{i} returned {result.action} on tool_call for tool={tool.name}")
                     return result
@@ -320,13 +312,12 @@ class MiddlewareChain:
         self,
         tool: "BaseTool",
         result: "ToolResult",
-        context: dict[str, Any],
     ) -> HookResult:
         """Process tool end through all middlewares (reverse order)."""
         logger.debug(f"Processing tool_end for tool={tool.name}, {len(self._middlewares)} middlewares (reverse order)")
         for i, mw in enumerate(reversed(self._middlewares)):
             if hasattr(mw, 'on_tool_end'):
-                hook_result = await mw.on_tool_end(tool, result, context)
+                hook_result = await mw.on_tool_end(tool, result)
                 if hook_result.action != HookAction.CONTINUE:
                     logger.info(f"Middleware #{i} returned {hook_result.action} on tool_end for tool={tool.name}")
                     return hook_result
@@ -338,14 +329,13 @@ class MiddlewareChain:
         parent_agent_id: str,
         child_agent_id: str,
         mode: str,
-        context: dict[str, Any],
     ) -> HookResult:
         """Process sub-agent start through all middlewares."""
         logger.debug(f"Processing subagent_start, parent={parent_agent_id}, child={child_agent_id}, mode={mode}, {len(self._middlewares)} middlewares")
         for i, mw in enumerate(self._middlewares):
             if hasattr(mw, 'on_subagent_start'):
                 result = await mw.on_subagent_start(
-                    parent_agent_id, child_agent_id, mode, context
+                    parent_agent_id, child_agent_id, mode
                 )
                 if result.action != HookAction.CONTINUE:
                     logger.info(f"Middleware #{i} returned {result.action} on subagent_start")
@@ -358,14 +348,13 @@ class MiddlewareChain:
         parent_agent_id: str,
         child_agent_id: str,
         result: Any,
-        context: dict[str, Any],
     ) -> HookResult:
         """Process sub-agent end through all middlewares (reverse order)."""
         logger.debug(f"Processing subagent_end, parent={parent_agent_id}, child={child_agent_id}, {len(self._middlewares)} middlewares (reverse order)")
         for i, mw in enumerate(reversed(self._middlewares)):
             if hasattr(mw, 'on_subagent_end'):
                 hook_result = await mw.on_subagent_end(
-                    parent_agent_id, child_agent_id, result, context
+                    parent_agent_id, child_agent_id, result
                 )
                 if hook_result.action != HookAction.CONTINUE:
                     logger.info(f"Middleware #{i} returned {hook_result.action} on subagent_end")
@@ -376,13 +365,11 @@ class MiddlewareChain:
     async def process_message_save(
         self,
         message: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Process message save through all middlewares.
         
         Args:
             message: Message to be saved
-            context: Execution context
             
         Returns:
             Modified message, or None to skip saving
@@ -392,7 +379,7 @@ class MiddlewareChain:
         
         for i, mw in enumerate(self._middlewares):
             if hasattr(mw, 'on_message_save'):
-                result = await mw.on_message_save(current, context)
+                result = await mw.on_message_save(current)
                 if result is None:
                     logger.info(f"Middleware #{i} blocked message save for role={message.get('role')}")
                     return None

aury/agents/middleware/message.py

@@ -56,37 +56,34 @@ class MessageBackendMiddleware(BaseMiddleware):
     async def on_message_save(
         self,
         message: dict[str, Any],
-        context: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Save message via backends.message.
         
         Args:
             message: Message dict with 'role', 'content', etc.
-            context: Execution context with 'session_id', 'agent_id', 'backends'
             
         Returns:
             The message (pass through to other middlewares)
         """
         from ..core.context import get_current_ctx_or_none
         
-        session_id = context.get("session_id", "")
-        if not session_id:
-            return message
-        
         # Get MessageBackend from context
         ctx = get_current_ctx_or_none()
         if ctx is None or ctx.backends is None or ctx.backends.message is None:
             # No backend available, pass through
             return message
         
+        session_id = ctx.session_id or ""
+        if not session_id:
+            return message
+        
         backend = ctx.backends.message
         
         # Extract message fields
         role = message.get("role", "")
         content = message.get("content", "")
-        invocation_id = context.get("invocation_id", "")
-        agent_id = context.get("agent_id")
-        namespace = context.get("namespace")
+        invocation_id = ctx.invocation_id or ""
+        agent_id = ctx.agent_id
         tool_call_id = message.get("tool_call_id")
         
         # Build message dict for backend
@@ -107,7 +104,6 @@ class MessageBackendMiddleware(BaseMiddleware):
             message=msg_dict,
             type="truncated",
             agent_id=agent_id,
-            namespace=namespace,
             invocation_id=invocation_id,
         )
         
@@ -118,7 +114,6 @@ class MessageBackendMiddleware(BaseMiddleware):
                 message=message,  # Full original message
                 type="raw",
                 agent_id=agent_id,
-                namespace=namespace,
                 invocation_id=invocation_id,
             )