hdsp-jupyter-extension 2.0.26__py3-none-any.whl → 2.0.28__py3-none-any.whl

agent_server/langchain/middleware/subagent_middleware.py

@@ -2,29 +2,43 @@
 SubAgentMiddleware
 
 Middleware that enables subagent delegation via the `task` tool.
-Based on Deep Agents library pattern (benchmarked, not installed).
-
-Key features:
-- Provides task(agent_name, description) tool for subagent invocation
-- Context isolation: subagents run in clean context
-- Synchronous execution: subagent returns result directly to caller
-- Nested subagent support: python_developer can call athena_query
-- Subagent caching: compiled agents are cached to avoid recompilation overhead
+Benchmarked from Deep Agents library pattern.
+
+Architecture:
+- Uses ToolRuntime for state access (instead of global variables)
+- Returns Command for state updates (instead of plain strings)
+- Extracts generated code/SQL to state fields (avoids JSON escaping)
+- ContentInjectionMiddleware handles injection into target tools
+
+Key patterns from Deep Agents:
+- _EXCLUDED_STATE_KEYS: isolates messages/todos between parent and subagent
+- Command(update={...}): state update return from tools
+- ToolRuntime.state: access graph state from within tools
 """
 
+import contextvars
 import hashlib
 import json
 import logging
-from typing import TYPE_CHECKING, Any, Dict, List, Optional
+import re
+import time
+import uuid
+from typing import Any, Callable, Dict, List, Optional, Union
 
-from langchain_core.tools import tool
+from langchain_core.messages import HumanMessage, ToolMessage
+from langchain_core.tools import StructuredTool
 from pydantic import BaseModel, Field
 
-if TYPE_CHECKING:
-    pass
-
 logger = logging.getLogger(__name__)
 
+# Deep Agents pattern: state keys excluded from bidirectional sharing
+_EXCLUDED_STATE_KEYS = ("messages", "todos")
+
+# Context variable to track the current main agent's thread_id
+_current_thread_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
+    "current_thread_id", default=None
+)
+
 # Global registry for subagent factories (set by AgentFactory)
 _subagent_factory = None
 _current_llm_config = None
@@ -32,15 +46,15 @@ _current_llm_config = None
 _subagent_cache: Dict[str, Any] = {}
 
 
+# ---------------------------------------------------------------------------
+# Global factory management (backward-compatible)
+# ---------------------------------------------------------------------------
+
 def set_subagent_factory(factory_func, llm_config: Dict[str, Any]):
-    """
-    Set the subagent factory function.
-    Called by AgentFactory during initialization.
-    """
+    """Set the subagent factory function. Called by AgentFactory."""
     global _subagent_factory, _current_llm_config, _subagent_cache
     _subagent_factory = factory_func
     _current_llm_config = llm_config
-    # Clear cache when factory changes (new LLM config)
     _subagent_cache.clear()
     logger.info("SubAgentMiddleware factory initialized (cache cleared)")
 
@@ -59,12 +73,7 @@ def _get_config_hash(llm_config: Dict[str, Any]) -> str:
 def get_or_create_subagent(
     agent_name: str, factory_func, llm_config: Dict[str, Any]
 ) -> Any:
-    """
-    Get cached subagent or create new one.
-
-    Caching avoids expensive recompilation of LangGraph agents.
-    Cache key = "{agent_name}_{config_hash}" to handle different LLM configs.
-    """
+    """Get cached subagent or create new one."""
     global _subagent_cache
 
     config_hash = _get_config_hash(llm_config)
@@ -80,43 +89,330 @@ def get_or_create_subagent(
     logger.info(
         f"Cached subagent '{agent_name}' (total cached: {len(_subagent_cache)})"
     )
-
     return subagent
 
 
 def clear_subagent_cache():
-    """Clear the subagent cache. Useful for testing or config changes."""
+    """Clear the subagent cache."""
     global _subagent_cache
     count = len(_subagent_cache)
     _subagent_cache.clear()
     logger.info(f"Subagent cache cleared ({count} entries removed)")
 
 
+# ---------------------------------------------------------------------------
+# Thread ID tracking (for code history middleware)
+# ---------------------------------------------------------------------------
+
+def set_current_thread_id(thread_id: str) -> None:
+    """Set the current main agent's thread_id for code history tracking."""
+    _current_thread_id.set(thread_id)
+    logger.debug(f"Set current thread_id: {thread_id}")
+
+
+def get_current_thread_id() -> Optional[str]:
+    """Get the current main agent's thread_id."""
+    return _current_thread_id.get()
+
+
+# ---------------------------------------------------------------------------
+# Content extraction utilities
+# ---------------------------------------------------------------------------
+
+def extract_code_block(response: str, lang: str = "python") -> Optional[str]:
+    """Extract code/SQL from response using [CODE]/[SQL] markers or fenced blocks.
+
+    Tries in order:
+    1. [CODE] or [SQL] section marker
+    2. ```python or ```sql fenced code block
+    3. ``` generic fenced code block
+
+    Args:
+        response: Full response text from subagent
+        lang: Language type ("python" or "sql")
+
+    Returns:
+        Extracted code string, or None if not found
+    """
+    if not response:
+        return None
+
+    # Pattern 1: [CODE] or [SQL] section marker with fenced block
+    marker = "[CODE]" if lang == "python" else "[SQL]"
+    marker_pattern = re.escape(marker) + r'\s*\n\s*```(?:\w+)?\s*\n(.*?)```'
+    match = re.search(marker_pattern, response, re.DOTALL)
+    if match:
+        code = match.group(1).strip()
+        if code:
+            return code
+
+    # Pattern 2: [CODE] or [SQL] marker with content after it (no fenced block)
+    marker_pattern2 = re.escape(marker) + r'\s*\n(.*?)(?=\n\[|\Z)'
+    match = re.search(marker_pattern2, response, re.DOTALL)
+    if match:
+        code = match.group(1).strip()
+        # Remove fenced block markers if present
+        if code.startswith(f"```{lang}"):
+            code = code[len(f"```{lang}"):].strip()
+        if code.startswith("```"):
+            code = code[3:].strip()
+        if code.endswith("```"):
+            code = code[:-3].strip()
+        if code:
+            return code
+
+    # Pattern 3: Fenced code block with language
+    fenced_pattern = rf'```{lang}\s*\n(.*?)```'
+    match = re.search(fenced_pattern, response, re.DOTALL)
+    if match:
+        code = match.group(1).strip()
+        if code:
+            return code
+
+    # Pattern 4: Generic fenced code block (fallback)
+    generic_pattern = r'```\s*\n(.*?)```'
+    match = re.search(generic_pattern, response, re.DOTALL)
+    if match:
+        code = match.group(1).strip()
+        if code and len(code) > 10:  # Minimum meaningful code
+            return code
+
+    return None
+
+
+def extract_description(response: str) -> Optional[str]:
+    """Extract [DESCRIPTION] section from subagent response.
+
+    Args:
+        response: Full response text from subagent
+
+    Returns:
+        Description string, or None if not found
+    """
+    if not response:
+        return None
+
+    patterns = [
+        r'\[DESCRIPTION\]\s*\n(.*?)(?=\n\s*\[(?:CODE|SQL)\])',
+        r'\[DESCRIPTION\]\s*(.*?)(?=\s*\[(?:CODE|SQL)\])',
+        r'\[DESCRIPTION\]\s*\n(.*?)(?=\n\s*```)',
+        r'\[DESCRIPTION\]\s*\n(.+?)(?=\n\n|\Z)',
+    ]
+
+    for pattern in patterns:
+        match = re.search(pattern, response, re.DOTALL | re.IGNORECASE)
+        if match:
+            description = match.group(1).strip()
+            if description and len(description) > 5:
+                return description
+
+    return None
+
+
+def extract_generated_content(
+    response: str, agent_type: str
+) -> tuple:
+    """Extract code/SQL from subagent response based on agent type.
+
+    Args:
+        response: Full response text from subagent
+        agent_type: "python_developer", "athena_query", etc.
+
+    Returns:
+        Tuple of (content, content_type) or (None, None)
+    """
+    if agent_type == "python_developer":
+        code = extract_code_block(response, lang="python")
+        return (code, "python") if code else (None, None)
+    elif agent_type == "athena_query":
+        sql = extract_code_block(response, lang="sql")
+        return (sql, "sql") if sql else (None, None)
+    return (None, None)
+
+
+def _build_summary(
+    response: str,
+    agent_type: str,
+    content_extracted: bool = False,
+) -> str:
+    """Build ToolMessage content for the caller.
+
+    When content is extracted to state, returns a short summary.
+    Otherwise returns the full response.
+
+    Args:
+        response: Full response from subagent
+        agent_type: Type of the subagent
+        content_extracted: Whether content was extracted to state
+
+    Returns:
+        Summary or full response string
+    """
+    if not content_extracted:
+        return response
+
+    # Content extracted - return summary only (actual content is in state)
+    desc = extract_description(response)
+    desc_part = f"\n설명: {desc}" if desc else ""
+
+    if agent_type == "python_developer":
+        return (
+            f"python_developer 코드 생성 완료.{desc_part}\n"
+            "jupyter_cell_tool() 또는 write_file_tool()로 실행/저장하세요. "
+            "(코드는 자동 주입됩니다)"
+        )
+    elif agent_type == "athena_query":
+        return (
+            f"athena_query SQL 생성 완료.{desc_part}\n"
+            "markdown_tool()로 표시하세요. "
+            "(SQL은 자동 주입됩니다)"
+        )
+    return response
+
+
+# ---------------------------------------------------------------------------
+# Task tool creation
+# ---------------------------------------------------------------------------
+
+def _invoke_subagent(
+    caller_name: str,
+    agent_name: str,
+    description: str,
+    context: Optional[str],
+    runtime_state: Optional[Dict[str, Any]] = None,
+    thread_id: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Execute a subagent and return its result dict.
+
+    Shared logic between Main Agent task tool and nested task tool.
+
+    Args:
+        caller_name: Name of the calling agent
+        agent_name: Name of the subagent to invoke
+        description: Task description
+        context: Optional additional context
+        runtime_state: Optional state from ToolRuntime (for state sharing)
+        thread_id: Optional main agent thread_id (for code history lookup)
+
+    Returns:
+        Result dict from subagent.invoke()
+    """
+    from agent_server.langchain.middleware.subagent_events import (
+        clear_current_subagent,
+        emit_subagent_complete,
+        emit_subagent_start,
+        set_current_subagent,
+    )
+
+    emit_subagent_start(agent_name, description)
+
+    factory_func, llm_config = get_subagent_factory()
+    if factory_func is None:
+        raise RuntimeError("SubAgentMiddleware not initialized. Call set_subagent_factory first.")
+
+    try:
+        set_current_subagent(agent_name)
+
+        t0 = time.time()
+        subagent = get_or_create_subagent(agent_name, factory_func, llm_config)
+        logger.info(f"[TIMING] get_or_create_subagent took {time.time()-t0:.2f}s")
+
+        # Inject code history for python_developer
+        enhanced_context = context
+        if agent_name == "python_developer":
+            try:
+                from agent_server.langchain.middleware.code_history_middleware import (
+                    get_code_history_tracker,
+                    get_context_with_history,
+                )
+                main_thread_id = thread_id or get_current_thread_id()
+                tracker = get_code_history_tracker(main_thread_id)
+                if tracker.get_entry_count() > 0:
+                    enhanced_context = get_context_with_history(context, main_thread_id)
+                    logger.info(
+                        f"[TIMING] code history injected "
+                        f"(entries={tracker.get_entry_count()}, "
+                        f"context_len={len(enhanced_context) if enhanced_context else 0})"
+                    )
+            except Exception as e:
+                logger.warning(f"Failed to inject code history: {e}")
+
+        # Build message content
+        if enhanced_context:
+            message_content = f"## Task\n{description}\n\n## Context (provided by Main Agent)\n{enhanced_context}"
+        else:
+            message_content = description
+
+        logger.info(f"[{caller_name}] Subagent message length: {len(message_content)}")
+
+        # Build subagent input state (Deep Agents pattern)
+        if runtime_state is not None:
+            subagent_state = {
+                k: v for k, v in runtime_state.items()
+                if k not in _EXCLUDED_STATE_KEYS
+            }
+            subagent_state["messages"] = [HumanMessage(content=message_content)]
+        else:
+            subagent_state = {"messages": [{"role": "user", "content": message_content}]}
+
+        subagent_thread_id = f"subagent-{agent_name}-{uuid.uuid4().hex[:8]}"
+        subagent_config = {"configurable": {"thread_id": subagent_thread_id}}
+
+        t_invoke = time.time()
+        logger.info(f"[TIMING] About to invoke subagent '{agent_name}'...")
+        result = subagent.invoke(subagent_state, config=subagent_config)
+        logger.info(f"[TIMING] subagent.invoke() took {time.time()-t_invoke:.2f}s")
+
+        # Extract response text
+        messages = result.get("messages", [])
+        if messages:
+            final_message = messages[-1]
+            response_text = (
+                final_message.content
+                if hasattr(final_message, "content")
+                else str(final_message)
+            )
+        else:
+            response_text = "Subagent completed but returned no messages."
+
+        logger.info(
+            f"[{caller_name}] Subagent '{agent_name}' returned: {str(response_text)[:200]}..."
+        )
+
+        emit_subagent_complete(agent_name, str(response_text)[:100])
+        return result
+
+    except Exception as e:
+        error_msg = f"Subagent '{agent_name}' failed: {str(e)}"
+        logger.error(error_msg, exc_info=True)
+        emit_subagent_complete(agent_name, f"Error: {str(e)[:50]}")
+        raise
+    finally:
+        clear_current_subagent()
+
+
 def create_task_tool(
     caller_name: str,
     allowed_subagents: Optional[List[str]] = None,
-):
-    """
-    Create a task tool for calling subagents.
+) -> StructuredTool:
+    """Create a task tool for nested subagent calls (returns plain string).
 
-    The task tool executes subagents synchronously and returns their result.
-    Subagents like python_developer return generated code/analysis,
-    which the Main Agent can then execute if needed.
+    Used by subagents that can call other subagents (e.g., python_developer → athena_query).
+    Returns full response text as string (no Command/state extraction).
 
     Args:
-        caller_name: Name of the agent creating this tool (for logging/validation)
-        allowed_subagents: Optional list of subagent names this agent can call.
-                          If None, all subagents are allowed (for Main Agent).
+        caller_name: Name of the agent creating this tool
+        allowed_subagents: Optional list of allowed subagent names
 
     Returns:
-        A tool that can be used to delegate tasks to subagents.
+        StructuredTool for subagent delegation
     """
     from agent_server.langchain.subagents.base import (
         SUBAGENT_CONFIGS,
         get_subagent_config,
     )
 
-    # Build description based on allowed subagents
+    # Build description
     if allowed_subagents:
         available = [
             f"- {name}: {SUBAGENT_CONFIGS[name].description}"
@@ -124,261 +420,251 @@ def create_task_tool(
             if name in SUBAGENT_CONFIGS
         ]
     else:
-        # Main Agent (planner) can call non-restricted subagents OR those explicitly allowing "planner"
         available = [
             f"- {config.name}: {config.description}"
             for config in SUBAGENT_CONFIGS.values()
             if not config.callable_by or caller_name in config.callable_by
         ]
-
     available_str = "\n".join(available)
 
-    # Create Pydantic schema for the task tool (required for Gemini compatibility)
-    class TaskInput(BaseModel):
-        """Input schema for task tool"""
-
-        agent_name: str = Field(
-            description=f"Name of the subagent to invoke. Available: {', '.join(allowed_subagents) if allowed_subagents else 'python_developer, researcher, athena_query'}"
-        )
-        description: str = Field(
-            description="Detailed task description for the subagent (Korean preferred)"
-        )
-        context: Optional[str] = Field(
-            default=None,
-            description="Additional context for the subagent: resource info (file sizes, memory), previous code, variable state, etc.",
-        )
-
-    @tool(args_schema=TaskInput)
     def task_tool(
-        agent_name: str, description: str, context: Optional[str] = None
+        agent_name: str,
+        description: str,
+        context: Optional[str] = None,
     ) -> str:
-        """
-        Delegate a task to a specialized subagent.
-
-        The subagent will execute the task and return its result (code, analysis, etc.).
-        Code execution tools (jupyter_cell_tool, write_file_tool) are handled by Main Agent.
-
-        Args:
-            agent_name: Name of the subagent to invoke
-            description: Detailed task description for the subagent
-            context: Additional context (resource info, previous code, etc.)
-
-        Returns:
-            Result from the subagent execution (string summary or generated code)
-        """
-        # Validate subagent exists
+        """Delegate a task to a specialized subagent (nested call)."""
         if agent_name not in SUBAGENT_CONFIGS:
-            return f"Error: Unknown agent '{agent_name}'. Available agents:\n{available_str}"
+            return f"Error: Unknown agent '{agent_name}'. Available:\n{available_str}"
 
-        # Validate caller is allowed to call this subagent
         config = get_subagent_config(agent_name)
         if allowed_subagents and agent_name not in allowed_subagents:
             return f"Error: '{caller_name}' cannot call '{agent_name}'. Allowed: {allowed_subagents}"
-
         if config.callable_by and caller_name not in config.callable_by:
             return f"Error: '{agent_name}' can only be called by: {config.callable_by}"
 
-        logger.info(
-            f"[{caller_name}] Invoking subagent '{agent_name}': {description[:100]}..."
-        )
+        try:
+            result = _invoke_subagent(caller_name, agent_name, description, context)
+            messages = result.get("messages", [])
+            if messages:
+                final = messages[-1]
+                return final.content if hasattr(final, "content") else str(final)
+            return "Subagent completed but returned no messages."
+        except Exception as e:
+            return f"Error: Subagent '{agent_name}' failed: {str(e)}"
 
-        # Import subagent event emitters
-        from agent_server.langchain.middleware.subagent_events import (
-            clear_current_subagent,
-            emit_subagent_complete,
-            emit_subagent_start,
-            set_current_subagent,
-        )
+    return StructuredTool.from_function(
+        name="task_tool",
+        func=task_tool,
+        description=f"Delegate a task to a specialized subagent.\n\nAvailable agents:\n{available_str}",
+    )
 
-        # Emit subagent start event for UI
-        emit_subagent_start(agent_name, description)
 
-        # Get the factory and config
-        factory_func, llm_config = get_subagent_factory()
-        if factory_func is None:
-            return "Error: SubAgentMiddleware not initialized. Call set_subagent_factory first."
+def _create_main_task_tool(
+    caller_name: str,
+    allowed_subagents: Optional[List[str]] = None,
+) -> StructuredTool:
+    """Create a task tool for Main Agent (returns Command with state update).
 
-        try:
-            import time
-
-            # Set current subagent context for tool call tracking
-            set_current_subagent(agent_name)
-
-            # Get or create the subagent (cached for performance)
-            # Avoids expensive LangGraph recompilation on each call
-            t0 = time.time()
-            subagent = get_or_create_subagent(agent_name, factory_func, llm_config)
-            t1 = time.time()
-            logger.info(f"[TIMING] get_or_create_subagent took {t1-t0:.2f}s")
-
-            # Execute subagent synchronously with clean context
-            # The subagent runs in isolation, receiving task description + optional context
-            import uuid
-
-            subagent_thread_id = f"subagent-{agent_name}-{uuid.uuid4().hex[:8]}"
-            subagent_config = {
-                "configurable": {
-                    "thread_id": subagent_thread_id,
-                }
-            }
+    Uses ToolRuntime for state access and returns Command to update
+    generated_content/generated_content_type/content_description in state.
 
-            # Inject code history for python_developer
-            enhanced_context = context
-            if agent_name == "python_developer":
-                try:
-                    t2 = time.time()
-                    from agent_server.langchain.middleware.code_history_middleware import (
-                        get_code_history_tracker,
-                        get_context_with_history,
-                    )
+    Args:
+        caller_name: Name of the calling agent (usually "planner")
+        allowed_subagents: Optional list of allowed subagent names
 
-                    tracker = get_code_history_tracker()
-                    if tracker.get_entry_count() > 0:
-                        enhanced_context = get_context_with_history(context)
-                        t3 = time.time()
-                        logger.info(
-                            f"[TIMING] code history injection took {t3-t2:.2f}s "
-                            f"(entries={tracker.get_entry_count()}, "
-                            f"context_len={len(enhanced_context) if enhanced_context else 0})"
-                        )
-                except Exception as e:
-                    logger.warning(f"Failed to inject code history: {e}")
+    Returns:
+        StructuredTool that returns Command
+    """
+    # Lazy import to avoid circular dependency
+    from langchain.tools import ToolRuntime
+    from langgraph.types import Command
+
+    from agent_server.langchain.subagents.base import (
+        SUBAGENT_CONFIGS,
+        get_subagent_config,
+    )
+
+    # Build description
+    if allowed_subagents:
+        available = [
+            f"- {name}: {SUBAGENT_CONFIGS[name].description}"
+            for name in allowed_subagents
+            if name in SUBAGENT_CONFIGS
+        ]
+    else:
+        available = [
+            f"- {config.name}: {config.description}"
+            for config in SUBAGENT_CONFIGS.values()
+            if not config.callable_by or caller_name in config.callable_by
+        ]
+    available_str = "\n".join(available)
 
-            # Build the message content with optional context
-            if enhanced_context:
-                message_content = f"""## Task
-{description}
+    def task(
+        description: str,
+        subagent_type: str,
+        runtime: ToolRuntime,
+    ) -> str:
+        """Delegate a task to a specialized subagent.
 
-## Context (provided by Main Agent)
-{enhanced_context}"""
-            else:
-                message_content = description
+        The subagent executes the task and returns its result.
+        For python_developer and athena_query, generated code/SQL is stored
+        in state and auto-injected into target tools by ContentInjectionMiddleware.
 
-            logger.info(
-                f"[{caller_name}] Subagent message length: {len(message_content)}"
-            )
+        Args:
+            description: Detailed task description for the subagent
+            subagent_type: Name of the subagent to invoke
+            runtime: ToolRuntime (auto-injected by LangGraph)
 
-            # Execute the subagent
-            t_invoke_start = time.time()
-            logger.info(f"[TIMING] About to invoke subagent '{agent_name}'...")
-            result = subagent.invoke(
-                {"messages": [{"role": "user", "content": message_content}]},
-                config=subagent_config,
-            )
-            t_invoke_end = time.time()
-            logger.info(
-                f"[TIMING] subagent.invoke() took {t_invoke_end-t_invoke_start:.2f}s"
+        Returns:
+            Command with state update, or error string
+        """
+        if subagent_type not in SUBAGENT_CONFIGS:
+            return f"Error: Unknown agent '{subagent_type}'. Available:\n{available_str}"
+
+        config = get_subagent_config(subagent_type)
+        if allowed_subagents and subagent_type not in allowed_subagents:
+            return f"Error: '{caller_name}' cannot call '{subagent_type}'."
+        if config.callable_by and caller_name not in config.callable_by:
+            return f"Error: '{subagent_type}' can only be called by: {config.callable_by}"
+
+        try:
+            # Extract thread_id from runtime config for code history lookup
+            main_thread_id = None
+            if hasattr(runtime, "config") and runtime.config:
+                configurable = runtime.config.get("configurable", {})
+                main_thread_id = configurable.get("thread_id")
+
+            # Invoke subagent with state sharing (Deep Agents pattern)
+            result = _invoke_subagent(
+                caller_name,
+                subagent_type,
+                description,
+                context=None,
+                runtime_state=dict(runtime.state) if runtime.state else None,
+                thread_id=main_thread_id,
             )
 
-            # Extract the final message from the result
+            # Extract response text
             messages = result.get("messages", [])
+            response_text = ""
             if messages:
-                final_message = messages[-1]
-                if hasattr(final_message, "content"):
-                    response = final_message.content
-                else:
-                    response = str(final_message)
-            else:
-                response = "Subagent completed but returned no messages."
-
-            logger.info(
-                f"[{caller_name}] Subagent '{agent_name}' returned: {str(response)[:200]}..."
-            )
-
-            # Extract description from python_developer response for auto-injection
-            if agent_name == "python_developer":
-                try:
-                    from agent_server.langchain.middleware.description_injector import (
-                        process_task_tool_response,
-                    )
+                final = messages[-1]
+                response_text = (
+                    final.content if hasattr(final, "content") else str(final)
+                )
+
+            # Extract generated content to state (Deep Agents Command pattern)
+            content, content_type = extract_generated_content(response_text, subagent_type)
+            desc = extract_description(response_text)
+
+            # Build state update (exclude messages/todos like Deep Agents)
+            state_update: Dict[str, Any] = {
+                k: v for k, v in result.items()
+                if k not in _EXCLUDED_STATE_KEYS
+            }
 
-                    process_task_tool_response(agent_name, str(response))
-                except Exception as e:
-                    logger.warning(f"Failed to extract description: {e}")
+            if content:
+                state_update["generated_content"] = content
+                state_update["generated_content_type"] = content_type
+                state_update["content_description"] = desc
+                logger.info(
+                    f"[State] Extracted {content_type} content "
+                    f"({len(content)} chars) to state"
+                )
+
+            # Build ToolMessage (summary for Main Agent, since content is in state)
+            summary = _build_summary(
+                response_text,
+                subagent_type,
+                content_extracted=bool(content),
+            )
 
-            # Emit subagent complete event for UI
-            emit_subagent_complete(agent_name, str(response)[:100])
+            if not runtime.tool_call_id:
+                raise ValueError("Tool call ID is required for subagent invocation")
 
-            return response
+            return Command(
+                update={
+                    **state_update,
+                    "messages": [
+                        ToolMessage(
+                            content=summary,
+                            tool_call_id=runtime.tool_call_id,
+                        )
+                    ],
+                }
+            )
 
         except Exception as e:
-            error_msg = f"Subagent '{agent_name}' failed: {str(e)}"
+            error_msg = f"Error: Subagent '{subagent_type}' failed: {str(e)}"
             logger.error(error_msg, exc_info=True)
-            # Emit complete event even on error
-            emit_subagent_complete(agent_name, f"Error: {str(e)[:50]}")
-            return f"Error: {error_msg}"
-        finally:
-            # Always clear subagent context
-            clear_current_subagent()
-
-    # Update tool docstring with available agents
-    task_tool.__doc__ = f"""Delegate a task to a specialized subagent.
-
-Available agents:
-{available_str}
-
-The subagent will analyze the task and return its result.
-- python_developer: Returns generated Python code and analysis
-- researcher: Returns search results and findings
-- athena_query: Returns SQL query string
-
-For code execution (running Python, writing files), use Main Agent's tools directly.
-
-IMPORTANT: For python_developer, ALWAYS provide context with:
-- Resource info (file sizes, memory) from check_resource_tool
-- Previous code context if building on existing work
-- Variable names and their current state
-
-Args:
-    agent_name: Name of the subagent to invoke
-    description: Detailed task description for the subagent
-    context: Additional context (resource info, previous code, etc.)
-
-Returns:
-    Result from the subagent execution
-"""
+            return error_msg
+
+    return StructuredTool.from_function(
+        name="task_tool",
+        func=task,
+        description=(
+            f"Delegate a task to a specialized subagent.\n\n"
+            f"Available agents:\n{available_str}\n\n"
+            "Generated code/SQL is automatically injected into execution tools.\n"
+            "After calling task_tool, use the appropriate tool to execute/display:\n"
+            "- python_developer → jupyter_cell_tool() or write_file_tool()\n"
+            "- athena_query → markdown_tool()\n"
+            "- researcher → direct response"
+        ),
+    )
 
-    return task_tool
 
+# ---------------------------------------------------------------------------
+# SubAgentMiddleware (extends AgentMiddleware)
+# ---------------------------------------------------------------------------
 
 class SubAgentMiddleware:
-    """
-    Middleware that adds subagent delegation capability.
+    """Middleware that adds subagent delegation via the task tool.
 
-    This middleware:
-    1. Adds the `task` tool to the agent's toolset
-    2. The task tool executes subagents synchronously
-    3. Subagent results are returned directly to the caller
+    Benchmarked from Deep Agents library pattern:
+    - Extends AgentMiddleware (tools auto-registered, wrap_model_call hook)
+    - Task tool uses ToolRuntime for state access
+    - Returns Command for state updates (generated_content in state)
+    - ContentInjectionMiddleware handles injection into target tools
 
     Usage:
         middleware = SubAgentMiddleware(
-            caller_name="main_agent",
-            allowed_subagents=["python_developer", "researcher"],
+            caller_name="planner",
+            allowed_subagents=None,
         )
 
         agent = create_agent(
             model=llm,
             tools=tools,
             middleware=[middleware, ...],
+            state_schema=HDSPAgentState,
         )
     """
 
+    TASK_SYSTEM_PROMPT = """## task_tool 사용법
+
+task_tool로 서브에이전트에게 작업을 위임하세요.
+서브에이전트가 생성한 코드/SQL은 자동으로 실행 도구에 주입됩니다.
+
+task_tool 호출 후 처리:
+- python_developer → jupyter_cell_tool() 호출 (코드 자동 주입)
+- python_developer (파일 저장) → write_file_tool() 호출 (코드 자동 주입)
+- athena_query → markdown_tool() 호출 (SQL 자동 주입)
+- researcher → 응답 내용 직접 활용"""
+
     def __init__(
         self,
         caller_name: str,
         allowed_subagents: Optional[List[str]] = None,
     ):
-        """
-        Initialize SubAgentMiddleware.
-
-        Args:
-            caller_name: Name of the agent using this middleware
-            allowed_subagents: List of subagents this agent can call.
-                             None means all non-restricted subagents.
-        """
         self.caller_name = caller_name
         self.allowed_subagents = allowed_subagents
-        self.task_tool = create_task_tool(caller_name, allowed_subagents)
+
+        # Create task tool (auto-registered via self.tools)
+        self.task_tool = _create_main_task_tool(caller_name, allowed_subagents)
+
+        # AgentMiddleware convention: tools attribute is auto-registered by create_agent
+        self.tools = [self.task_tool]
 
         logger.info(
             f"SubAgentMiddleware initialized for '{caller_name}' "
@@ -386,13 +672,5 @@ class SubAgentMiddleware:
         )
 
     def get_tools(self) -> List[Any]:
-        """Get the tools provided by this middleware."""
+        """Get the tools provided by this middleware (backward compat)."""
         return [self.task_tool]
-
-    def __call__(self, tools: List[Any]) -> List[Any]:
-        """
-        Add task tool to the agent's toolset.
-
-        This is called during agent creation to augment the tool list.
-        """
-        return tools + [self.task_tool]