fix: allow litellm security patch

Author: scale-ballen

examples/tutorials/run_agent_test.sh

@@ -260,6 +260,18 @@ run_test() {
 
 
     # Run the tests with retry mechanism
+    local -a pytest_cmd=("uv" "run" "pytest")
+    if [ "$BUILD_CLI" = true ]; then
+        local wheel_file
+        wheel_file=$(ls /home/runner/work/*/*/dist/agentex_sdk-*.whl 2>/dev/null | head -n1)
+        if [[ -z "$wheel_file" ]]; then
+            wheel_file=$(ls "${SCRIPT_DIR}/../../dist/agentex_sdk-*.whl" 2>/dev/null | head -n1)
+        fi
+        if [[ -n "$wheel_file" ]]; then
+            pytest_cmd=("uv" "run" "--with" "$wheel_file" "pytest")
+        fi
+    fi
+
     local max_retries=5
     local retry_count=0
     local exit_code=1
@@ -270,7 +282,7 @@ run_test() {
         fi
 
         # Stream pytest output directly in real-time
-        uv run pytest tests/test_agent.py -v -s
+        "${pytest_cmd[@]}" tests/test_agent.py -v -s
         exit_code=$?
 
         if [ $exit_code -eq 0 ]; then

pyproject.toml

@@ -9,7 +9,7 @@ authors = [
 ]
 
 dependencies = [
-    "httpx>=0.27.2,<0.28",
+    "httpx>=0.28.1,<0.29",
     "pydantic>=2.0.0, <3",
   "typing-extensions>=4.14, <5",
   "anyio>=3.5.0, <5",
@@ -18,7 +18,9 @@ dependencies = [
     "typer>=0.16,<0.17",
     "questionary>=2.0.1,<3",
     "rich>=13.9.2,<14",
-    "fastapi>=0.115.0,<0.116",
+    "fastapi>=0.115.0",
+    "starlette>=0.49.1",
+    "tornado>=6.5.5",
     "uvicorn>=0.31.1",
     "watchfiles>=0.24.0,<1.0",
     "python-on-whales>=0.73.0,<0.74",
@@ -28,7 +30,7 @@ dependencies = [
     "temporalio>=1.26.0,<2",
     "aiohttp>=3.10.10,<4",
     "redis>=5.2.0,<6",
-    "litellm>=1.83.0,<2",
+    "litellm>=1.83.7,<2",
     "kubernetes>=25.0.0,<36.0.0",
     "jinja2>=3.1.3,<4",
     "mcp[cli]>=1.4.1",

src/agentex/lib/core/temporal/plugins/openai_agents/hooks/hooks.py

@@ -131,7 +131,7 @@ async def on_tool_start(self, context: RunContextWrapper, agent: Agent, tool: To
                 logger.warning(f"Failed to parse tool arguments: {tool_context.tool_arguments}")
                 tool_arguments = {}
 
-        await workflow.execute_activity_method(
+        await workflow.execute_activity(
             stream_lifecycle_content,
             args=[
                 self.task_id,
@@ -167,7 +167,7 @@ async def on_tool_end(
             else f"call_{id(tool)}"
         )
 
-        await workflow.execute_activity_method(
+        await workflow.execute_activity(
             stream_lifecycle_content,
             args=[
                 self.task_id,
@@ -195,7 +195,7 @@ async def on_handoff(
             from_agent: The agent transferring control
             to_agent: The agent receiving control
         """
-        await workflow.execute_activity_method(
+        await workflow.execute_activity(
             stream_lifecycle_content,
             args=[
                 self.task_id,

src/agentex/lib/core/tracing/processors/agentex_tracing_processor.py

@@ -26,6 +26,7 @@ def on_span_start(self, span: Span) -> None:
             input=span.input,
             output=span.output,
             parent_id=span.parent_id,
+            task_id=span.task_id,
         )
 
     @override
@@ -94,6 +95,7 @@ async def on_span_start(self, span: Span) -> None:
             input=span.input,
             output=span.output,
             data=span.data,
+            task_id=span.task_id,
         )
 
     @override

src/agentex/lib/sdk/fastacp/base/base_acp_server.py

@@ -11,8 +11,8 @@
 import uvicorn
 from fastapi import FastAPI, Request
 from pydantic import TypeAdapter, ValidationError
+from starlette.types import Send, Scope, ASGIApp, Receive
 from fastapi.responses import StreamingResponse
-from starlette.middleware.base import BaseHTTPMiddleware
 
 from agentex.lib.types.acp import (
     RPC_SYNC_METHODS,
@@ -44,17 +44,19 @@
 task_message_update_adapter = TypeAdapter(TaskMessageUpdate)
 
 
-class RequestIDMiddleware(BaseHTTPMiddleware):
-    """Middleware to extract or generate request IDs and add them to logs and response headers"""
+class RequestIDMiddleware:
+    """Pure ASGI middleware to set request IDs without buffering streaming responses."""
 
-    async def dispatch(self, request: Request, call_next):  # type: ignore[override]
-        # Extract request ID from header or generate a new one if there isn't one
-        request_id = request.headers.get("x-request-id") or uuid.uuid4().hex
-        # Store request ID in request state for access in handlers
-        ctx_var_request_id.set(request_id)
-        # Process request
-        response = await call_next(request)
-        return response
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] == "http":
+            headers = dict(scope.get("headers", []))
+            raw_request_id = headers.get(b"x-request-id", b"")
+            request_id = raw_request_id.decode() if raw_request_id else uuid.uuid4().hex
+            ctx_var_request_id.set(request_id)
+        await self.app(scope, receive, send)
 
 
 class BaseACPServer(FastAPI):

src/agentex/resources/agents.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import json
-from typing import Union, Optional, Generator, AsyncGenerator
+from typing import Any, Union, Optional, Generator, AsyncGenerator
 from typing_extensions import Literal
 
 import httpx
@@ -359,7 +359,7 @@ def create_task(
     ) -> CreateTaskResponse:
       if agent_id is not None and agent_name is not None:
         raise ValueError("Either agent_id or agent_name must be provided, but not both")
-      
+
       if agent_id is not None:
         raw_agent_rpc_response = self.rpc(
           agent_id=agent_id,
@@ -386,7 +386,7 @@ def create_task(
         )
       else:
         raise ValueError("Either agent_id or agent_name must be provided")
-      
+
       return CreateTaskResponse.model_validate(raw_agent_rpc_response, from_attributes=True)
 
     def cancel_task(
@@ -453,38 +453,71 @@ def send_message(
     ) -> SendMessageResponse:
       if agent_id is not None and agent_name is not None:
         raise ValueError("Either agent_id or agent_name must be provided, but not both")
-      
+
       if "stream" in params and params["stream"] == True:
         raise ValueError("If stream is set to True, use send_message_stream() instead")
+
+      if agent_id is not None:
+        raw_agent_rpc_response = self.with_streaming_response.rpc(
+          agent_id=agent_id,
+          method="message/send",
+          params=params,
+          id=id,
+          jsonrpc=jsonrpc,
+          extra_headers=extra_headers,
+          extra_query=extra_query,
+          extra_body=extra_body,
+          timeout=timeout,
+        )
+      elif agent_name is not None:
+        raw_agent_rpc_response = self.with_streaming_response.rpc_by_name(
+          agent_name=agent_name,
+          method="message/send",
+          params=params,
+          id=id,
+          jsonrpc=jsonrpc,
+          extra_headers=extra_headers,
+          extra_query=extra_query,
+          extra_body=extra_body,
+          timeout=timeout,
+        )
       else:
-        if agent_id is not None:
-          raw_agent_rpc_response = self.rpc(
-            agent_id=agent_id,
-            method="message/send",
-            params=params,
-            id=id,
-            jsonrpc=jsonrpc,
-            extra_headers=extra_headers,
-            extra_query=extra_query,
-            extra_body=extra_body,
-            timeout=timeout,
-          )
-        elif agent_name is not None:
-          raw_agent_rpc_response = self.rpc_by_name(
-            agent_name=agent_name,
-            method="message/send",
-            params=params,
-            id=id,
-            jsonrpc=jsonrpc,
-            extra_headers=extra_headers,
-            extra_query=extra_query,
-            extra_body=extra_body,
-            timeout=timeout,
-          )
-        else:
-          raise ValueError("Either agent_id or agent_name must be provided")
-        
-        return SendMessageResponse.model_validate(raw_agent_rpc_response, from_attributes=True)
+        raise ValueError("Either agent_id or agent_name must be provided")
+
+      task_messages: list[Any] = []
+      response_meta: dict[str, Any] = {}
+
+      with raw_agent_rpc_response as response:
+        for _line in response.iter_lines():
+          if not _line:
+            continue
+          line = _line.strip()
+          if line.startswith("data:"):
+            line = line[len("data:"):].strip()
+          if not line:
+            continue
+          try:
+            chunk = json.loads(line)
+            if not response_meta:
+              response_meta = {"id": chunk.get("id"), "jsonrpc": chunk.get("jsonrpc")}
+            try:
+              return SendMessageResponse.model_validate(chunk)
+            except ValidationError:
+              pass
+            chunk_stream = SendMessageStreamResponse.model_validate(chunk, from_attributes=True)
+            result = chunk_stream.result
+            if result is not None and getattr(result, "type", None) == "full":
+              parent = getattr(result, "parent_task_message", None)
+              if parent is not None:
+                task_messages.append(parent)
+          except (json.JSONDecodeError, ValidationError):
+            continue
+
+      return SendMessageResponse(
+        id=response_meta.get("id"),
+        jsonrpc=response_meta.get("jsonrpc"),
+        result=task_messages,
+      )
 
     def send_message_stream(
       self,
@@ -552,8 +585,8 @@ def send_message_stream(
               from_attributes=True
             )
             yield chunk_rpc_response
-          except json.JSONDecodeError:
-            # Skip invalid JSON lines
+          except (json.JSONDecodeError, ValidationError):
+            # Skip invalid JSON lines or lines that cannot be validated
             continue
 
     def send_event(
@@ -1021,38 +1054,71 @@ async def send_message(
     ) -> SendMessageResponse:
       if agent_id is not None and agent_name is not None:
         raise ValueError("Either agent_id or agent_name must be provided, but not both")
-      
+
       if "stream" in params and params["stream"] == True:
         raise ValueError("If stream is set to True, use send_message_stream() instead")
+
+      if agent_id is not None:
+        raw_agent_rpc_response = self.with_streaming_response.rpc(
+          agent_id=agent_id,
+          method="message/send",
+          params=params,
+          id=id,
+          jsonrpc=jsonrpc,
+          extra_headers=extra_headers,
+          extra_query=extra_query,
+          extra_body=extra_body,
+          timeout=timeout,
+        )
+      elif agent_name is not None:
+        raw_agent_rpc_response = self.with_streaming_response.rpc_by_name(
+          agent_name=agent_name,
+          method="message/send",
+          params=params,
+          id=id,
+          jsonrpc=jsonrpc,
+          extra_headers=extra_headers,
+          extra_query=extra_query,
+          extra_body=extra_body,
+          timeout=timeout,
+        )
       else:
-        if agent_id is not None:
-          raw_agent_rpc_response = await self.rpc(
-            agent_id=agent_id,
-            method="message/send",
-            params=params,
-            id=id,
-            jsonrpc=jsonrpc,
-            extra_headers=extra_headers,
-            extra_query=extra_query,
-            extra_body=extra_body,
-            timeout=timeout,
-          )
-        elif agent_name is not None:
-          raw_agent_rpc_response = await self.rpc_by_name(
-            agent_name=agent_name,
-            method="message/send",
-            params=params,
-            id=id,
-            jsonrpc=jsonrpc,
-            extra_headers=extra_headers,
-            extra_query=extra_query,
-            extra_body=extra_body,
-            timeout=timeout,
-          )
-        else:
-          raise ValueError("Either agent_id or agent_name must be provided")
-        
-        return SendMessageResponse.model_validate(raw_agent_rpc_response, from_attributes=True)
+        raise ValueError("Either agent_id or agent_name must be provided")
+
+      task_messages: list[Any] = []
+      response_meta: dict[str, Any] = {}
+
+      async with raw_agent_rpc_response as response:
+        async for _line in response.iter_lines():
+          if not _line:
+            continue
+          line = _line.strip()
+          if line.startswith("data:"):
+            line = line[len("data:"):].strip()
+          if not line:
+            continue
+          try:
+            chunk = json.loads(line)
+            if not response_meta:
+              response_meta = {"id": chunk.get("id"), "jsonrpc": chunk.get("jsonrpc")}
+            try:
+              return SendMessageResponse.model_validate(chunk)
+            except ValidationError:
+              pass
+            chunk_stream = SendMessageStreamResponse.model_validate(chunk, from_attributes=True)
+            result = chunk_stream.result
+            if result is not None and getattr(result, "type", None) == "full":
+              parent = getattr(result, "parent_task_message", None)
+              if parent is not None:
+                task_messages.append(parent)
+          except (json.JSONDecodeError, ValidationError):
+            continue
+
+      return SendMessageResponse(
+        id=response_meta.get("id"),
+        jsonrpc=response_meta.get("jsonrpc"),
+        result=task_messages,
+      )
 
     async def send_message_stream(
       self,