agentex-sdk 0.6.7__py3-none-any.whl → 0.7.1__py3-none-any.whl

agentex/_streaming.py

@@ -54,11 +54,12 @@ class Stream(Generic[_T]):
         process_data = self._client._process_response_data
         iterator = self._iter_events()
 
-        for sse in iterator:
-            yield process_data(data=sse.json(), cast_to=cast_to, response=response)
-
-        # As we might not fully consume the response stream, we need to close it explicitly
-        response.close()
+        try:
+            for sse in iterator:
+                yield process_data(data=sse.json(), cast_to=cast_to, response=response)
+        finally:
+            # Ensure the response is closed even if the consumer doesn't read all data
+            response.close()
 
     def __enter__(self) -> Self:
         return self
@@ -117,11 +118,12 @@ class AsyncStream(Generic[_T]):
         process_data = self._client._process_response_data
         iterator = self._iter_events()
 
-        async for sse in iterator:
-            yield process_data(data=sse.json(), cast_to=cast_to, response=response)
-
-        # As we might not fully consume the response stream, we need to close it explicitly
-        await response.aclose()
+        try:
+            async for sse in iterator:
+                yield process_data(data=sse.json(), cast_to=cast_to, response=response)
+        finally:
+            # Ensure the response is closed even if the consumer doesn't read all data
+            await response.aclose()
 
     async def __aenter__(self) -> Self:
         return self

agentex/_types.py

@@ -243,6 +243,9 @@ _T_co = TypeVar("_T_co", covariant=True)
 if TYPE_CHECKING:
     # This works because str.__contains__ does not accept object (either in typeshed or at runtime)
     # https://github.com/hauntsaninja/useful_types/blob/5e9710f3875107d068e7679fd7fec9cfab0eff3b/useful_types/__init__.py#L285
+    #
+    # Note: index() and count() methods are intentionally omitted to allow pyright to properly
+    # infer TypedDict types when dict literals are used in lists assigned to SequenceNotStr.
     class SequenceNotStr(Protocol[_T_co]):
         @overload
         def __getitem__(self, index: SupportsIndex, /) -> _T_co: ...
@@ -251,8 +254,6 @@ if TYPE_CHECKING:
         def __contains__(self, value: object, /) -> bool: ...
         def __len__(self) -> int: ...
         def __iter__(self) -> Iterator[_T_co]: ...
-        def index(self, value: Any, start: int = 0, stop: int = ..., /) -> int: ...
-        def count(self, value: Any, /) -> int: ...
         def __reversed__(self) -> Iterator[_T_co]: ...
 else:
     # just point this to a normal `Sequence` at runtime to avoid having to special case

agentex/_version.py

@@ -1,4 +1,4 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
 __title__ = "agentex"
-__version__ = "0.6.7"  # x-release-please-version
+__version__ = "0.7.1"  # x-release-please-version

agentex/lib/core/temporal/plugins/claude_agents/__init__.py

@@ -0,0 +1,72 @@
+"""Claude Agents SDK integration with Temporal.
+
+This plugin provides integration between Claude Agents SDK and AgentEx's
+Temporal-based orchestration platform.
+
+Features:
+- Temporal activity wrapper for Claude SDK calls
+- Real-time streaming to Redis/UI
+- Session resume for conversation context
+- Tool call visibility (Read, Write, Bash, etc.)
+- Subagent support with nested tracing
+- Workspace isolation per task
+
+Architecture:
+- activities.py: Temporal activity definitions
+- message_handler.py: Message parsing and streaming logic
+- Reuses OpenAI's ContextInterceptor for context threading
+
+Usage:
+    from agentex.lib.core.temporal.plugins.claude_agents import (
+        run_claude_agent_activity,
+        create_workspace_directory,
+        ContextInterceptor,
+    )
+
+    # In worker
+    worker = AgentexWorker(
+        task_queue=queue_name,
+        interceptors=[ContextInterceptor()],
+    )
+
+    activities = get_all_activities()
+    activities.extend([run_claude_agent_activity, create_workspace_directory])
+
+    await worker.run(activities=activities, workflow=YourWorkflow)
+"""
+
+from agentex.lib.core.temporal.plugins.claude_agents.hooks import (
+    TemporalStreamingHooks,
+    create_streaming_hooks,
+)
+from agentex.lib.core.temporal.plugins.claude_agents.activities import (
+    run_claude_agent_activity,
+    create_workspace_directory,
+)
+from agentex.lib.core.temporal.plugins.claude_agents.message_handler import (
+    ClaudeMessageHandler,
+)
+
+# Reuse OpenAI's context threading - this is the key to streaming!
+from agentex.lib.core.temporal.plugins.openai_agents.interceptors.context_interceptor import (
+    ContextInterceptor,
+    streaming_task_id,
+    streaming_trace_id,
+    streaming_parent_span_id,
+)
+
+__all__ = [
+    # Activities
+    "run_claude_agent_activity",
+    "create_workspace_directory",
+    # Message handling
+    "ClaudeMessageHandler",
+    # Hooks
+    "create_streaming_hooks",
+    "TemporalStreamingHooks",
+    # Context threading (reused from OpenAI)
+    "ContextInterceptor",
+    "streaming_task_id",
+    "streaming_trace_id",
+    "streaming_parent_span_id",
+]

agentex/lib/core/temporal/plugins/claude_agents/activities.py

@@ -0,0 +1,154 @@
+"""Temporal activities for Claude Agents SDK integration."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+from temporalio import activity
+from claude_agent_sdk import AgentDefinition, ClaudeSDKClient, ClaudeAgentOptions
+
+from agentex.lib.utils.logging import make_logger
+from agentex.lib.core.temporal.plugins.claude_agents.hooks import create_streaming_hooks
+from agentex.lib.core.temporal.plugins.claude_agents.message_handler import ClaudeMessageHandler
+from agentex.lib.core.temporal.plugins.openai_agents.interceptors.context_interceptor import (
+    streaming_task_id,
+    streaming_trace_id,
+    streaming_parent_span_id,
+)
+
+logger = make_logger(__name__)
+
+
+@activity.defn
+async def create_workspace_directory(task_id: str, workspace_root: str | None = None) -> str:
+    """Create workspace directory for task - runs as Temporal activity
+
+    Args:
+        task_id: Task ID for workspace directory name
+        workspace_root: Root directory for workspaces (defaults to .claude-workspace/ in cwd)
+
+    Returns:
+        Absolute path to created workspace
+    """
+    if workspace_root is None:
+        # Default to .claude-workspace in current directory
+        # Follows Claude SDK's .claude/ convention
+        workspace_root = os.path.join(os.getcwd(), ".claude-workspace")
+
+    workspace_path = os.path.join(workspace_root, task_id)
+    os.makedirs(workspace_path, exist_ok=True)
+    logger.info(f"Created workspace: {workspace_path}")
+    return workspace_path
+
+
+@activity.defn(name="run_claude_agent_activity")
+async def run_claude_agent_activity(
+    prompt: str,
+    workspace_path: str,
+    allowed_tools: list[str],
+    permission_mode: str = "acceptEdits",
+    system_prompt: str | None = None,
+    resume_session_id: str | None = None,
+    agents: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Execute Claude SDK - wrapped in Temporal activity
+
+    This activity:
+    1. Gets task_id from ContextVar (set by ContextInterceptor)
+    2. Configures Claude with workspace isolation and session resume
+    3. Runs Claude SDK and processes messages via ClaudeMessageHandler
+    4. Streams messages to UI in real-time
+    5. Returns session_id, usage, and cost for next turn
+
+    Args:
+        prompt: User message to send to Claude
+        workspace_path: Directory for file operations (cwd)
+        allowed_tools: List of tools Claude can use (include "Task" for subagents)
+        permission_mode: Permission mode (default: acceptEdits)
+        system_prompt: Optional system prompt override
+        resume_session_id: Optional session ID to resume conversation context
+        agents: Optional dict of subagent definitions for Task tool
+
+    Returns:
+        dict with "messages", "session_id", "usage", and "cost_usd" keys
+    """
+
+    # Get streaming context from ContextVars (set by interceptor)
+    task_id = streaming_task_id.get()
+    trace_id = streaming_trace_id.get()
+    parent_span_id = streaming_parent_span_id.get()
+
+    logger.info(
+        f"[run_claude_agent_activity] Starting - "
+        f"task_id={task_id}, workspace={workspace_path}, tools={allowed_tools}, "
+        f"resume={'YES' if resume_session_id else 'NO (new session)'}, "
+        f"subagents={list(agents.keys()) if agents else 'NONE'}"
+    )
+
+    # Reconstruct AgentDefinition objects from serialized dicts
+    # Temporal serializes dataclasses to dicts, need to recreate them
+    agent_defs = None
+    if agents:
+        agent_defs = {}
+        for name, agent_data in agents.items():
+            if isinstance(agent_data, AgentDefinition):
+                agent_defs[name] = agent_data
+            else:
+                # Reconstruct from dict
+                agent_defs[name] = AgentDefinition(
+                    description=agent_data.get('description', ''),
+                    prompt=agent_data.get('prompt', ''),
+                    tools=agent_data.get('tools'),
+                    model=agent_data.get('model'),
+                )
+
+    # Create hooks for streaming tool calls and subagent execution
+    hooks = create_streaming_hooks(
+        task_id=task_id,
+        trace_id=trace_id,
+        parent_span_id=parent_span_id,
+    )
+
+    # Configure Claude with workspace isolation, session resume, subagents, and hooks
+    options = ClaudeAgentOptions(
+        cwd=workspace_path,
+        allowed_tools=allowed_tools,
+        permission_mode=permission_mode,  # type: ignore
+        system_prompt=system_prompt,
+        resume=resume_session_id,
+        agents=agent_defs,
+        hooks=hooks,  # Tool lifecycle hooks for streaming!
+    )
+
+    # Create message handler for streaming
+    handler = ClaudeMessageHandler(
+        task_id=task_id,
+        trace_id=trace_id,
+        parent_span_id=parent_span_id,
+    )
+
+    # Run Claude and process messages
+    try:
+        await handler.initialize()
+
+        async with ClaudeSDKClient(options=options) as client:
+            await client.query(prompt)
+
+            # Use receive_response() instead of receive_messages()
+            # receive_response() yields messages until ResultMessage, then stops
+            # receive_messages() is infinite and never completes!
+            async for message in client.receive_response():
+                await handler.handle_message(message)
+
+        logger.debug(f"Message loop completed, cleaning up...")
+        await handler.cleanup()
+
+        results = handler.get_results()
+        logger.debug(f"Returning results with keys: {results.keys()}")
+        return results
+
+    except Exception as e:
+        logger.error(f"[run_claude_agent_activity] Error: {e}", exc_info=True)
+        await handler.cleanup()
+        raise

agentex/lib/core/temporal/plugins/claude_agents/hooks/__init__.py

@@ -0,0 +1,11 @@
+"""Claude SDK hooks for streaming lifecycle events to AgentEx UI."""
+
+from agentex.lib.core.temporal.plugins.claude_agents.hooks.hooks import (
+    TemporalStreamingHooks,
+    create_streaming_hooks,
+)
+
+__all__ = [
+    "create_streaming_hooks",
+    "TemporalStreamingHooks",
+]

agentex/lib/core/temporal/plugins/claude_agents/hooks/hooks.py

@@ -0,0 +1,212 @@
+"""Claude SDK hooks for streaming tool calls and subagent execution to AgentEx UI.
+
+This module provides hook callbacks that integrate with Claude SDK's hooks system
+to stream tool execution lifecycle events in real-time.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from claude_agent_sdk import HookMatcher
+
+from agentex.lib import adk
+from agentex.lib.utils.logging import make_logger
+from agentex.types.task_message_update import StreamTaskMessageFull
+from agentex.types.tool_request_content import ToolRequestContent
+from agentex.types.tool_response_content import ToolResponseContent
+
+logger = make_logger(__name__)
+
+
+class TemporalStreamingHooks:
+    """Hooks for streaming Claude SDK lifecycle events to AgentEx UI.
+
+    Implements Claude SDK hook callbacks:
+    - PreToolUse: Called before tool execution → stream tool request
+    - PostToolUse: Called after tool execution → stream tool result
+
+    Also handles subagent detection and nested tracing.
+    """
+
+    def __init__(
+        self,
+        task_id: str | None,
+        trace_id: str | None = None,
+        parent_span_id: str | None = None,
+    ):
+        """Initialize streaming hooks.
+
+        Args:
+            task_id: AgentEx task ID for routing streams
+            trace_id: Trace ID for nested spans
+            parent_span_id: Parent span ID for subagent spans
+        """
+        self.task_id = task_id
+        self.trace_id = trace_id
+        self.parent_span_id = parent_span_id
+
+        # Track active subagent spans
+        self.subagent_spans: dict[str, Any] = {}  # tool_call_id → (ctx, span)
+
+    async def pre_tool_use(
+        self,
+        input_data: dict[str, Any],
+        tool_use_id: str | None,
+        _context: Any,
+    ) -> dict[str, Any]:
+        """Hook called before tool execution.
+
+        Args:
+            input_data: Contains tool_name, tool_input from Claude SDK
+            tool_use_id: Unique ID for this tool call
+            context: Hook context from Claude SDK
+
+        Returns:
+            Empty dict (allow execution to proceed)
+        """
+        if not self.task_id or not tool_use_id:
+            return {}
+
+        tool_name = input_data.get("tool_name", "unknown")
+        tool_input = input_data.get("tool_input", {})
+
+        logger.info(f"🔧 Tool request: {tool_name}")
+
+        # Special handling for Task tool (subagents) - create nested span
+        if tool_name == "Task" and self.trace_id and self.parent_span_id:
+            subagent_type = tool_input.get("subagent_type", "unknown")
+            logger.info(f"🤖 Subagent started: {subagent_type}")
+
+            # Create nested trace span for subagent
+            subagent_ctx = adk.tracing.span(
+                trace_id=self.trace_id,
+                parent_id=self.parent_span_id,
+                name=f"Subagent: {subagent_type}",
+                input=tool_input,
+            )
+            subagent_span = await subagent_ctx.__aenter__()
+            self.subagent_spans[tool_use_id] = (subagent_ctx, subagent_span)
+
+        # Stream tool request to UI
+        try:
+            async with adk.streaming.streaming_task_message_context(
+                task_id=self.task_id,
+                initial_content=ToolRequestContent(
+                    author="agent",
+                    name=tool_name,
+                    arguments=tool_input,
+                    tool_call_id=tool_use_id,
+                )
+            ) as tool_ctx:
+                await tool_ctx.stream_update(
+                    StreamTaskMessageFull(
+                        parent_task_message=tool_ctx.task_message,
+                        content=ToolRequestContent(
+                            author="agent",
+                            name=tool_name,
+                            arguments=tool_input,
+                            tool_call_id=tool_use_id,
+                        ),
+                        type="full"
+                    )
+                )
+        except Exception as e:
+            logger.warning(f"Failed to stream tool request: {e}")
+
+        return {}  # Allow execution
+
+    async def post_tool_use(
+        self,
+        input_data: dict[str, Any],
+        tool_use_id: str | None,
+        _context: Any,
+    ) -> dict[str, Any]:
+        """Hook called after tool execution.
+
+        Args:
+            input_data: Contains tool_name, tool_output from Claude SDK
+            tool_use_id: Unique ID for this tool call
+            context: Hook context from Claude SDK
+
+        Returns:
+            Empty dict
+        """
+        if not self.task_id or not tool_use_id:
+            return {}
+
+        tool_name = input_data.get("tool_name", "unknown")
+        tool_output = input_data.get("tool_output", "")
+
+        logger.info(f"✅ Tool result: {tool_name}")
+
+        # If this was a subagent, close the nested span
+        if tool_use_id in self.subagent_spans:
+            subagent_ctx, subagent_span = self.subagent_spans[tool_use_id]
+            subagent_span.output = {"result": tool_output}
+            await subagent_ctx.__aexit__(None, None, None)
+            logger.info(f"🤖 Subagent completed: {tool_name}")
+            del self.subagent_spans[tool_use_id]
+
+        # Stream tool response to UI
+        try:
+            async with adk.streaming.streaming_task_message_context(
+                task_id=self.task_id,
+                initial_content=ToolResponseContent(
+                    author="agent",
+                    name=tool_name,
+                    content=tool_output,
+                    tool_call_id=tool_use_id,
+                )
+            ) as tool_ctx:
+                await tool_ctx.stream_update(
+                    StreamTaskMessageFull(
+                        parent_task_message=tool_ctx.task_message,
+                        content=ToolResponseContent(
+                            author="agent",
+                            name=tool_name,
+                            content=tool_output,
+                            tool_call_id=tool_use_id,
+                        ),
+                        type="full"
+                    )
+                )
+        except Exception as e:
+            logger.warning(f"Failed to stream tool response: {e}")
+
+        return {}
+
+
+def create_streaming_hooks(
+    task_id: str | None,
+    trace_id: str | None = None,
+    parent_span_id: str | None = None,
+) -> dict[str, list[HookMatcher]]:
+    """Create Claude SDK hooks configuration for streaming.
+
+    Returns hooks dict suitable for ClaudeAgentOptions(hooks=...).
+
+    Args:
+        task_id: AgentEx task ID for streaming
+        trace_id: Trace ID for nested spans
+        parent_span_id: Parent span ID for subagent spans
+
+    Returns:
+        Dict with PreToolUse and PostToolUse hook configurations
+    """
+    hooks_instance = TemporalStreamingHooks(task_id, trace_id, parent_span_id)
+
+    return {
+        "PreToolUse": [
+            HookMatcher(
+                matcher=None,  # Match all tools
+                hooks=[hooks_instance.pre_tool_use]
+            )
+        ],
+        "PostToolUse": [
+            HookMatcher(
+                matcher=None,  # Match all tools
+                hooks=[hooks_instance.post_tool_use]
+            )
+        ],
+    }

agentex/lib/core/temporal/plugins/claude_agents/message_handler.py

@@ -0,0 +1,178 @@
+"""Message handling and streaming for Claude Agents SDK.
+
+Simplified message handler that focuses on:
+- Streaming text content to UI
+- Extracting session_id for conversation continuity
+- Extracting usage and cost information
+
+Tool requests/responses are handled by Claude SDK hooks (see hooks/hooks.py).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from claude_agent_sdk import (
+    TextBlock,
+    ResultMessage,
+    SystemMessage,
+    AssistantMessage,
+)
+
+from agentex.lib import adk
+from agentex.lib.utils.logging import make_logger
+from agentex.types.text_content import TextContent
+from agentex.types.task_message_delta import TextDelta
+from agentex.types.task_message_update import StreamTaskMessageDelta
+
+logger = make_logger(__name__)
+
+
+class ClaudeMessageHandler:
+    """Handles Claude SDK messages and streams them to AgentEx UI.
+
+    Simplified handler focused on:
+    - Streaming text blocks to UI
+    - Extracting session_id from SystemMessage/ResultMessage
+    - Extracting usage and cost from ResultMessage
+    - Serializing responses for Temporal
+
+    Note: Tool lifecycle events (requests/responses) are handled by
+    TemporalStreamingHooks, not this class.
+    """
+
+    def __init__(
+        self,
+        task_id: str | None,
+        trace_id: str | None,
+        parent_span_id: str | None,
+    ):
+        self.task_id = task_id
+        self.trace_id = trace_id
+        self.parent_span_id = parent_span_id
+
+        # Message tracking
+        self.messages: list[Any] = []
+        self.serialized_messages: list[dict] = []
+
+        # Streaming context for text
+        self.streaming_ctx = None
+
+        # Result data
+        self.session_id: str | None = None
+        self.usage_info: dict | None = None
+        self.cost_info: float | None = None
+
+    async def initialize(self):
+        """Initialize streaming context if task_id is available."""
+        if self.task_id:
+            logger.debug(f"Creating streaming context for task: {self.task_id}")
+            self.streaming_ctx = await adk.streaming.streaming_task_message_context(
+                task_id=self.task_id,
+                initial_content=TextContent(
+                    author="agent",
+                    content="",
+                    format="markdown"
+                )
+            ).__aenter__()
+
+    async def handle_message(self, message: Any):
+        """Process a single message from Claude SDK."""
+        self.messages.append(message)
+        msg_num = len(self.messages)
+
+        # Debug logging (verbose - only for troubleshooting)
+        logger.debug(f"📨 [{msg_num}] Message type: {type(message).__name__}")
+        if isinstance(message, AssistantMessage):
+            block_types = [type(b).__name__ for b in message.content]
+            logger.debug(f"   [{msg_num}] Content blocks: {block_types}")
+
+        # Route to specific handlers
+        # Note: Tool requests/responses are handled by hooks, not here!
+        if isinstance(message, AssistantMessage):
+            await self._handle_assistant_message(message, msg_num)
+        elif isinstance(message, SystemMessage):
+            await self._handle_system_message(message)
+        elif isinstance(message, ResultMessage):
+            await self._handle_result_message(message)
+
+    async def _handle_assistant_message(self, message: AssistantMessage, _msg_num: int):
+        """Handle AssistantMessage - contains text blocks.
+
+        Note: Tool calls (ToolUseBlock/ToolResultBlock) are handled by hooks, not here.
+        We only process TextBlock for streaming text to UI.
+        """
+        # Stream text blocks to UI
+        for block in message.content:
+            if isinstance(block, TextBlock):
+                await self._handle_text_block(block)
+
+        # Collect text for final response
+        text_content = []
+        for block in message.content:
+            if isinstance(block, TextBlock):
+                text_content.append(block.text)
+
+        if text_content:
+            self.serialized_messages.append({
+                "role": "assistant",
+                "content": "\n".join(text_content)
+            })
+
+    async def _handle_text_block(self, block: TextBlock):
+        """Handle text content block."""
+        if not block.text or not self.streaming_ctx:
+            return
+
+        logger.debug(f"💬 Text block: {block.text[:50]}...")
+
+        delta = TextDelta(type="text", text_delta=block.text)
+
+        try:
+            await self.streaming_ctx.stream_update(
+                StreamTaskMessageDelta(
+                    parent_task_message=self.streaming_ctx.task_message,
+                    delta=delta,
+                    type="delta"
+                )
+            )
+        except Exception as e:
+            logger.warning(f"Failed to stream text delta: {e}")
+
+    async def _handle_system_message(self, message: SystemMessage):
+        """Handle system message - extract session_id."""
+        if message.subtype == "init":
+            self.session_id = message.data.get("session_id")
+            logger.debug(f"Session initialized: {self.session_id[:16] if self.session_id else 'unknown'}...")
+        else:
+            logger.debug(f"SystemMessage: {message.subtype}")
+
+    async def _handle_result_message(self, message: ResultMessage):
+        """Handle result message - extract usage and cost."""
+        self.usage_info = message.usage
+        self.cost_info = message.total_cost_usd
+
+        # Update session_id if available
+        if message.session_id:
+            self.session_id = message.session_id
+
+        logger.info(f"💰 Cost: ${self.cost_info:.4f}, Duration: {message.duration_ms}ms, Turns: {message.num_turns}")
+
+    async def cleanup(self):
+        """Clean up open streaming contexts."""
+        if self.streaming_ctx:
+            try:
+                await self.streaming_ctx.close()
+                logger.debug(f"Closed streaming context")
+            except Exception as e:
+                logger.warning(f"Failed to close streaming context: {e}")
+
+    def get_results(self) -> dict[str, Any]:
+        """Get final results for Temporal."""
+        return {
+            "messages": self.serialized_messages,
+            "task_id": self.task_id,
+            "session_id": self.session_id,
+            "usage": self.usage_info,
+            "cost_usd": self.cost_info,
+        }