delegate-tool 0.1.0__py3-none-any.whl

delegate_tool/main.py

@@ -0,0 +1,912 @@
+"""
+delegate Tool
+=============
+Delegates one or more tasks to a child agent node.
+
+Parameters:
+  tasks    - List of tasks/queries to delegate (required).
+"""
+
+from openchadpy.context import pipeline_ctx
+from openchad.chat.main import Chat
+import json
+import logging
+import os
+import re
+import sys
+import asyncio
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from openchadpy.context import agent_ctx, model_id_ctx, fields_ctx, additional_args_ctx
+from openchadpy.pipeline_base import object_to_text_tree, PipelineBase
+from openchadpy.tool_base import ToolBase, ToolRegistry
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Module-level helpers (mirrors logic from openchadpy/main.py and
+# Pipeline/openchad/chat/main.py so the tool is self-contained)
+# ---------------------------------------------------------------------------
+
+_RE_BACKTICK_FENCE_PB = re.compile(
+    r"```(?:[a-zA-Z0-9_+\-]*)?\r?\n(.*?)```",
+    re.DOTALL,
+)
+_RE_THINK_PB  = re.compile(r"<think>.*?</think>|<thinking>.*?</thinking>", re.DOTALL | re.IGNORECASE)
+_RE_TOOL_PB   = re.compile(r"<ToolCall\b[^>]*/>", re.DOTALL)
+_RE_CBLOCK_PB = re.compile(r"<CodeBlock\b[^>]*>.*?</CodeBlock>", re.DOTALL)
+
+def _extract_code(text: str) -> str:
+    """Strip MDX render tags then return the first fenced code block, or the
+    whole cleaned text when no fence is present.
+
+    NOTE: end() runs before finalize(), so the parser's pending buffer may
+    still hold the closing ``` when this is called.  We handle both the
+    complete-fence case (regex match) and the incomplete-fence case (opening
+    fence only) so a missing closing ``` never breaks compilation.
+    """
+    import textwrap
+    text = _RE_THINK_PB.sub("", text)
+    text = _RE_TOOL_PB.sub("", text)
+    text = _RE_CBLOCK_PB.sub(
+        lambda m: re.sub(r"<CodeBlock\b[^>]*>", "", m.group(0)).replace("</CodeBlock>", ""),
+        text,
+    )
+    text = text.strip("\r\n")
+    matches = _RE_BACKTICK_FENCE_PB.findall(text)
+    if matches:
+        raw_code = matches[0]
+    else:
+        # Fallback: closing ``` may still be in the parser's pending buffer.
+        # Strip just the opening fence and treat everything after as the body.
+        open_match = re.match(r"```(?:[a-zA-Z0-9_+\-]*)?\r?\n(.*)", text, re.DOTALL)
+        if open_match:
+            raw_code = open_match.group(1)
+            # Remove any trailing incomplete ``` that made it in
+            raw_code = re.sub(r"\n?```\s*$", "", raw_code, flags=re.DOTALL)
+        else:
+            raw_code = text
+    return textwrap.dedent(raw_code.strip("\r\n")).strip()
+
+
+def _build_agent_tree(flat: dict) -> dict:
+    """Normalise a flat agent dict and return a nested tree.
+
+    Handles two storage formats for ``children``:
+
+    * **List of IDs** ``["child_id", ...]``  – each child is a separate top-level
+      key in *flat* (legacy format).
+    * **Inline dict** ``{"child_id": {<node_data>}, ...}`` – child node data is
+      embedded directly inside the parent (current DB format).
+
+    Both formats are normalised into a flat ``normalized`` registry first so
+    that ``build_node`` can always do a simple lookup.
+    """
+    normalized: dict = {}
+
+    def _normalize_node(node_id: str, raw_node: Any) -> None:
+        """Register *node_id* and recursively register any inline children."""
+        if node_id in normalized or not isinstance(raw_node, dict):
+            return
+        node = dict(raw_node)
+
+        # Parse JSON-string fields.
+        for key in ("tools", "children"):
+            val = node.get(key)
+            if isinstance(val, str):
+                try:
+                    node[key] = json.loads(val)
+                except Exception:
+                    pass
+
+        # Normalise tools to a list.
+        if not isinstance(node.get("tools"), list):
+            node["tools"] = []
+
+        # Normalise children: convert inline dict → list of IDs and register
+        # each child node so build_node can look them up.
+        children_val = node.get("children")
+        if isinstance(children_val, dict):
+            # Inline format: {"child_id": {<node>}, ...}
+            node["children"] = list(children_val.keys())
+            for cid, cdata in children_val.items():
+                _normalize_node(cid, cdata)
+        elif isinstance(children_val, list):
+            # ID-list format: ["child_id", ...]
+            node["children"] = children_val
+        else:
+            node["children"] = []
+
+        normalized[node_id] = node
+
+    for node_id, raw_node in flat.items():
+        _normalize_node(node_id, raw_node)
+
+    all_child_ids = {
+        cid
+        for n in normalized.values()
+        for cid in n.get("children", [])
+    }
+    root_ids = [nid for nid in normalized if nid not in all_child_ids]
+
+    def build_node(nid: str) -> dict:
+        if nid not in normalized:
+            logger.warning("[_build_agent_tree] node '%s' not found – skipping", nid)
+            return {}
+        n = normalized[nid]
+        return {
+            **{k: v for k, v in n.items() if k != "children"},
+            "children": {
+                cid: build_node(cid)
+                for cid in n.get("children", [])
+            },
+        }
+
+    return {nid: build_node(nid) for nid in root_ids}
+
+
+
+def _find_in_tree(tree: dict, target_id: str) -> Optional[dict]:
+    """Recursively search a nested agent tree for *target_id*.
+    Returns the matching node dict, or None if not found.
+    """
+    for nid, node in tree.items():
+        if nid == target_id:
+            return node
+        children = node.get("children", {})
+        if isinstance(children, dict) and children:
+            found = _find_in_tree(children, target_id)
+            if found is not None:
+                return found
+    return None
+
+
+async def _find_agent_root(
+    workspace: str,
+    agent_id: str,
+) -> Optional[Tuple[str, dict]]:
+    """Find which root agent tab owns *agent_id* (which may be a child node).
+
+    Strategy:
+    1. Query the raw ``agents`` SQLite table for all root tab IDs.
+    2. For each root, load its flat agent list via ``Database`` and build the
+       nested tree with ``_build_agent_tree``.
+    3. Recursively search the tree for *agent_id*.
+
+    Returns ``(root_tab_id, agent_node)`` on success, or ``None`` if not found
+    in any root.
+    """
+    from openchadpy.sqlite import sqlite
+    from openchadpy.database import Database
+
+    # 1. All root tab IDs live in the raw `agents` table (inserted by the frontend).
+    result = await sqlite({
+        "db": workspace,
+        "command": "query",
+        "sql": "SELECT id FROM agents",
+    })
+    rows = result.get("data", [])
+    root_ids: List[str] = [row["id"] for row in rows if row.get("id")]
+
+    logger.debug(f"[delegate] scanning {len(root_ids)} root agent(s) for '{agent_id}'")
+
+    # 2+3. Load each root's tree and search for agent_id.
+    for root_id in root_ids:
+        try:
+            agent_db = Database(workspace=workspace, tab_id=root_id)
+            flat_agents = await agent_db.get("agents")
+            if not flat_agents:
+                continue
+            agent_tree = _build_agent_tree(flat_agents)
+            node = _find_in_tree(agent_tree, agent_id)
+            if node is not None:
+                logger.info(f"[delegate] found '{agent_id}' under root '{root_id}'")
+                return root_id, node
+        except Exception as exc:
+            logger.warning(f"[delegate] error scanning root '{root_id}': {exc}")
+            continue
+
+    return None
+
+
+def _parse_action_result(raw: Any) -> Optional[Dict[str, Any]]:
+    """Normalise a sandbox return value into a plain ActionResult-shaped dict.
+
+    Accepts:
+    - An ActionResult dataclass instance (has ``.result`` attribute)
+    - A plain dict with 'result', 'next_tasks', 'next_branch' keys
+    - None or anything else → returns None
+    """
+    if raw is None:
+        return None
+    if hasattr(raw, "result"):
+        return {
+            "result":        raw.result,
+            "next_tasks":    list(getattr(raw, "next_tasks", []) or []),
+            "next_branch":   getattr(raw, "next_branch", None),
+            "next_branches": dict(getattr(raw, "next_branches", {}) or {}),
+        }
+    if isinstance(raw, dict) and "result" in raw:
+        return {
+            "result":        raw.get("result", {}),
+            "next_tasks":    list(raw.get("next_tasks") or []),
+            "next_branch":   raw.get("next_branch"),
+            "next_branches": dict(raw.get("next_branches") or {}),
+        }
+    return None
+
+
+def _serialize_ans(ans: Any) -> Any:
+    """Ensure that the dict returned by ``pipeline.llm_tool`` / code_sandbox is
+    fully JSON-serialisable.
+
+    ``code_sandbox.execute()`` returns ``{"result": <ActionResult dataclass>, ...}``.
+    That raw dataclass is not JSON-serialisable, which causes the downstream
+    ``json.dumps(tool_result)`` call in the pipeline to raise a ``TypeError``.
+
+    We convert any dataclass stored under the ``"result"`` key to a plain dict
+    via ``dataclasses.asdict``, leaving every other value untouched.
+    """
+    if not isinstance(ans, dict):
+        return ans
+    result_val = ans.get("result")
+    if result_val is not None and hasattr(result_val, "__dataclass_fields__"):
+        from dataclasses import asdict as _dc_asdict
+        try:
+            ans = dict(ans)
+            ans["result"] = _dc_asdict(result_val)
+        except Exception:
+            # Last-resort: convert each field individually.
+            ans = dict(ans)
+            serialized: Dict[str, Any] = {}
+            for field_name in result_val.__dataclass_fields__:
+                try:
+                    serialized[field_name] = getattr(result_val, field_name)
+                except Exception:
+                    serialized[field_name] = "<unreadable>"
+            ans["result"] = serialized
+    return ans
+
+
+def _extract_response(ans: Any) -> Any:
+    """Flatten the raw ``exec_result`` / ``code_sandbox`` wrapper so the parent
+    agent sees only the *meaningful* payload — not the internal bookkeeping.
+
+    ``pipeline.llm_tool`` in programmatic mode returns the raw sandbox dict::
+
+        {"success": True, "result": <ActionResult>, "output": "", "error": None}
+
+    Returning this verbatim means the parent has to dig through
+    ``response.result.result`` (or deeper) to find actual data.  This helper
+    unwraps it so the parent sees:
+
+    * On success  → the value stored in ``ActionResult.result`` (the sandbox
+      code's actual return payload, e.g. ``{"path": "..."}``)
+    * On error    → ``{"error": "<message>"}``
+    * Non-sandbox → the value unchanged
+    """
+    if not isinstance(ans, dict):
+        return ans
+
+    # Not a sandbox exec_result — pass through as-is.
+    if "success" not in ans:
+        return ans
+
+    if not ans.get("success"):
+        return {"error": ans.get("error") or "Unknown error"}
+
+    result = ans.get("result")
+    if result is None:
+        output = ans.get("output", "").strip()
+        return {"output": output} if output else {}
+
+    # Unwrap ActionResult dataclass.
+    if hasattr(result, "__dataclass_fields__"):
+        from dataclasses import asdict as _dc_asdict
+        try:
+            result = _dc_asdict(result)
+        except Exception:
+            pass
+
+    # Unwrap serialized ActionResult dict: extract only the .result payload
+    # so callers don't need to know about next_branch / next_tasks / etc.
+    if isinstance(result, dict) and "result" in result and (
+        "next_branch" in result or "next_branches" in result
+    ):
+        return result.get("result") or {}
+
+    return result
+
+
+def _parse_bool(val: Any) -> bool:
+    if isinstance(val, bool):
+        return val
+    if isinstance(val, str):
+        return val.lower() == "true"
+    return False
+
+def _create_agent_query_tool(pipeline: "PipelineBase", agent_node: dict) -> ToolRegistry:
+    allow_multiple = _parse_bool(agent_node.get("allowMultiple"))
+    if allow_multiple:
+        schema = {
+            "type": "function",
+            "function": {
+                "name": "agent_query",
+                "description": "Send queries to one or more agents and retrieve their responses. Use this to delegate subtasks or gather information from specialized agents.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "queries": {
+                            "type": "array",
+                            "description": "List of queries to send to agents. Multiple entries allow querying several agents in a single call.",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "agent_id": {
+                                        "type": "string",
+                                        "description": "The unique identifier of the target agent."
+                                    },
+                                    "tasks": {
+                                        "type": "array",
+                                        "description": "List of tasks or questions to delegate to the target agent.",
+                                        "items": {
+                                            "type": "string",
+                                            "description": "A specific question or instruction for the agent."
+                                        }
+                                    }
+                                },
+                                "required": ["tasks"]
+                            },
+                            "minItems": 1
+                        }
+                    },
+                    "required": ["queries"]
+                }
+            }
+        }
+    else:
+        schema = {
+            "type": "function",
+            "function": {
+                "name": "agent_query",
+                "description": "Send a list of tasks to a specific agent and retrieve its response. Use this to delegate subtasks or gather information from specialized agents.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "agent_id": {
+                            "type": "string",
+                            "description": "The unique identifier of the target agent."
+                        },
+                        "tasks": {
+                            "type": "array",
+                            "description": "List of tasks or questions to delegate to the target agent.",
+                            "items": {
+                                "type": "string",
+                                "description": "A specific question or instruction for the agent."
+                            }
+                        }
+                    },
+                    "required": ["agent_id", "tasks"]
+                }
+            }
+        }
+
+    async def execute_tool(**kwargs) -> Dict[str, Any]:
+        queries = kwargs.get("queries")
+        if not queries:
+            query_val = kwargs.get("query") or kwargs.get("tasks")
+            queries = [{"agent_id": kwargs.get("agent_id"), "tasks": query_val}]
+        
+        async def run_agent_queries(q_item: dict) -> List[dict]:
+            sub_agent_id = q_item.get("agent_id")
+            sub_tasks = q_item.get("tasks") or q_item.get("query")
+            if not sub_agent_id or not sub_tasks:
+                return []
+            if isinstance(sub_tasks, str):
+                sub_tasks = [sub_tasks]
+            
+            agent_results = []
+            sub_agent_tree = agent_node.get("children", {}).get(sub_agent_id)
+            if sub_agent_tree:
+                old_agent = agent_ctx.get()
+                old_model = model_id_ctx.get()
+                old_fields = fields_ctx.get()
+                old_additional_args = additional_args_ctx.get()
+                
+                agent_ctx.set({sub_agent_id: sub_agent_tree})
+
+                sub_model = sub_agent_tree.get("model")
+                if sub_model:
+                    model_id_ctx.set(sub_model)
+                
+                fields_ctx.set(json.loads(sub_agent_tree.get("toolValues", "{}")))
+                additional_args_ctx.set(json.loads(sub_agent_tree.get("additionalArgs", "{}"))) 
+                
+                try:
+                    sub_agent_query = _create_agent_query_tool(pipeline, sub_agent_tree) if sub_agent_tree.get("children") else None
+                    
+                    for task in sub_tasks:
+                        if pipeline.cancel_event and pipeline.cancel_event.is_set():
+                            logger.info(
+                                "[%s] cancel_event set before agent_query sub-task '%s' – aborting",
+                                pipeline.__class__.__name__, task,
+                            )
+                            break
+                        
+                        tool_registry = {'agent_query': sub_agent_query} if sub_agent_query else None
+                        try:
+                            llm_task = asyncio.get_event_loop().create_task(
+                                pipeline.llm_tool(
+                                    history_id=sub_agent_id,
+                                    query=task,
+                                    tool_registry=tool_registry,
+                                ),
+                                name=f"agent_query_llm_tool_{id(pipeline)}",
+                            )
+                            cancel_sentinel = None
+                            if pipeline.cancel_event:
+                                cancel_sentinel = asyncio.get_event_loop().create_task(
+                                    pipeline.cancel_event.wait(),
+                                    name=f"cancel_sentinel_aq_{id(pipeline)}",
+                                )
+                            try:
+                                if cancel_sentinel is not None:
+                                    done, _ = await asyncio.wait(
+                                        {llm_task, cancel_sentinel},
+                                        return_when=asyncio.FIRST_COMPLETED,
+                                    )
+                                    if cancel_sentinel in done:
+                                        logger.info(
+                                            "[%s] agent_query: cancel_event fired – cancelling sub-task",
+                                            pipeline.__class__.__name__,
+                                        )
+                                        llm_task.cancel()
+                                        try:
+                                            await llm_task
+                                        except (asyncio.CancelledError, Exception):
+                                            pass
+                                        break
+                                    ans = llm_task.result()
+                                else:
+                                    ans = await llm_task
+                            finally:
+                                if cancel_sentinel is not None and not cancel_sentinel.done():
+                                    cancel_sentinel.cancel()
+                                    try:
+                                        await cancel_sentinel
+                                    except (asyncio.CancelledError, Exception):
+                                        pass
+                            
+                            entry = {
+                                "agent_id": sub_agent_id,
+                                "response": _extract_response(ans)
+                            }
+                            # Recursive fan-out check for programmatic child
+                            if isinstance(ans, dict) and ans.get("success") and ans.get("result") is not None:
+                                child_action = _parse_action_result(ans["result"])
+                                if child_action:
+                                    has_single = child_action.get("next_branch") and child_action.get("next_tasks")
+                                    has_multi  = bool(child_action.get("next_branches"))
+                                    if has_single or has_multi:
+                                        sub_results = await _fan_out_branch(pipeline, child_action)
+                                        entry["sub_branch_results"] = sub_results
+                            agent_results.append(entry)
+                        except Exception as e:
+                            logger.exception("Error executing sub-agent task: %s", task)
+                            agent_results.append({
+                                "agent_id": sub_agent_id,
+                                "response": f"Error: {str(e)}"
+                            })
+                finally:
+                    agent_ctx.set(old_agent)
+                    model_id_ctx.set(old_model)
+                    fields_ctx.set(old_fields)
+                    additional_args_ctx.set(old_additional_args)
+            else:
+                agent_results.append({
+                    "agent_id": sub_agent_id,
+                    "response": f"Error: Agent '{sub_agent_id}' not found."
+                })
+            return agent_results
+
+        results = await asyncio.gather(*(run_agent_queries(q) for q in queries))
+        flat_results = []
+        for r in results:
+            flat_results.extend(r)
+        return {"results": flat_results} if allow_multiple else (flat_results[0] if flat_results else {})
+
+    return ToolRegistry(call=execute_tool, schema=schema)
+
+
+async def _fan_out_single_branch(
+    pipeline: "PipelineBase",
+    child_id: str,
+    tasks: List[str],
+    cancel_event: Any,
+) -> List[Dict[str, Any]]:
+    """Run *tasks* against a single child branch (``child_id``).
+
+    Expects ``agent_ctx`` to already be set to the parent node so the child
+    node can be resolved from its ``children`` dict.  The context is swapped
+    to the child for the duration of the call and restored before returning.
+
+    Returns a list of per-task result dicts.
+    """
+    agent = agent_ctx.get()
+    if not agent or not isinstance(agent, dict):
+        logger.warning("[_fan_out_single_branch] no agent_ctx, skipping branch '%s'", child_id)
+        return []
+
+    current_id = next(iter(agent))
+    current_node = agent[current_id]
+    child_node = current_node.get("children", {}).get(child_id)
+    if not child_node:
+        logger.warning(
+            "[_fan_out_single_branch] child node '%s' not found in '%s'",
+            child_id, current_id,
+        )
+        return []
+
+    # Swap agent_ctx / model_id_ctx to the child node for the duration.
+    old_agent = agent_ctx.get()
+    old_model = model_id_ctx.get()
+    old_fields = fields_ctx.get()
+    old_additional_args = additional_args_ctx.get()
+    
+    agent_ctx.set({child_id: child_node})
+    child_model = child_node.get("model")
+    if child_model:
+        model_id_ctx.set(child_model)
+    
+    fields_ctx.set(json.loads(child_node.get("toolValues", "{}")))
+    additional_args_ctx.set(json.loads(child_node.get("additionalArgs", "{}")))
+
+    branch_results: List[Dict[str, Any]] = []
+    try:
+        is_child_prog = _parse_bool(child_node.get("enableProgrammaticToolCalling"))
+        has_children = bool(child_node.get("children"))
+        child_agent_query = None
+        if not is_child_prog and has_children:
+            child_agent_query = _create_agent_query_tool(pipeline, child_node)
+
+        import inspect
+        try:
+            sig = inspect.signature(pipeline.llm_tool)
+            has_tool_registry = "tool_registry" in sig.parameters or any(
+                p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
+            )
+        except Exception:
+            has_tool_registry = False
+
+        for task in tasks:
+            # Pre-task cancel check.
+            if cancel_event and cancel_event.is_set():
+                logger.info(
+                    "[_fan_out_single_branch] cancel_event set before task '%s' – aborting remaining tasks for '%s'",
+                    task, child_id,
+                )
+                break
+            try:
+                # ── Cancel-aware LLM call (mirrors tool sentinel pattern) ──────
+                _llm_sentinel: Optional[asyncio.Task] = None
+                if cancel_event:
+                    _llm_sentinel = asyncio.get_event_loop().create_task(
+                        cancel_event.wait(),
+                        name=f"cancel_sentinel_fan_out_{child_id}_{id(pipeline)}",
+                    )
+                try:
+                    llm_kwargs: Dict[str, Any] = {
+                        "query": task,
+                        "history_id": child_id,
+                    }
+                    if child_agent_query and has_tool_registry:
+                        llm_kwargs["tool_registry"] = {"agent_query": child_agent_query}
+
+                    _llm_task: asyncio.Task = asyncio.get_event_loop().create_task(
+                        pipeline.llm_tool(**llm_kwargs),
+                        name=f"llm_tool_fan_out_{child_id}_{id(pipeline)}",
+                    )
+                    try:
+                        if _llm_sentinel is not None:
+                            # Race: whichever finishes first wins.
+                            _done, _ = await asyncio.wait(
+                                {_llm_task, _llm_sentinel},
+                                return_when=asyncio.FIRST_COMPLETED,
+                            )
+                            if _llm_sentinel in _done:
+                                # Cancel was requested while the LLM call was in flight.
+                                logger.info(
+                                    "[_fan_out_single_branch] cancel_event fired during llm_tool for task '%s' in branch '%s' – cancelling",
+                                    task, child_id,
+                                )
+                                _llm_task.cancel()
+                                try:
+                                    await _llm_task
+                                except (asyncio.CancelledError, Exception):
+                                    pass
+                                break  # exit the task loop
+                            # LLM call finished normally – retrieve result.
+                            ans = _llm_task.result()
+                        else:
+                            ans = await _llm_task
+                    except asyncio.CancelledError:
+                        _llm_task.cancel()
+                        raise
+                    except Exception:
+                        raise
+                finally:
+                    # Always clean up the sentinel so it never leaks.
+                    if _llm_sentinel is not None and not _llm_sentinel.done():
+                        _llm_sentinel.cancel()
+                        try:
+                            await _llm_sentinel
+                        except (asyncio.CancelledError, Exception):
+                            pass
+                # ─────────────────────────────────────────────────────────────
+                entry: Dict[str, Any] = {"agent_id": child_id, "task": task, "response": _extract_response(ans)}
+
+                # Recursive fan-out: if the child returned an ActionResult with
+                # next_branch+next_tasks OR next_branches, recurse.
+                if isinstance(ans, dict) and ans.get("success") and ans.get("result") is not None:
+                    child_action = _parse_action_result(ans["result"])
+                    if child_action:
+                        has_single = child_action.get("next_branch") and child_action.get("next_tasks")
+                        has_multi  = bool(child_action.get("next_branches"))
+                        if has_single or has_multi:
+                            # Pre-recursion cancel check.
+                            if cancel_event and cancel_event.is_set():
+                                logger.info(
+                                    "[_fan_out_single_branch] cancel_event set before recursive fan-out from '%s' – skipping",
+                                    child_id,
+                                )
+                            else:
+                                logger.info(
+                                    "[_fan_out_single_branch] recursive fan-out from '%s': next_branch=%s next_branches=%s tasks=%s",
+                                    child_id,
+                                    child_action.get("next_branch"),
+                                    list((child_action.get("next_branches") or {}).keys()),
+                                    child_action.get("next_tasks"),
+                                )
+                                # agent_ctx is currently {child_id: child_node},
+                                # so the recursive call resolves grandchildren
+                                # from child_node["children"].
+                                sub_results = await _fan_out_branch(pipeline, child_action)
+                                entry["sub_branch_results"] = sub_results
+
+                branch_results.append(entry)
+            except Exception as exc:
+                logger.exception("[_fan_out_single_branch] task failed | task=%s child=%s", task, child_id)
+                branch_results.append({"agent_id": child_id, "task": task, "error": str(exc)})
+    finally:
+        agent_ctx.set(old_agent)
+        model_id_ctx.set(old_model)
+        fields_ctx.set(old_fields)
+        additional_args_ctx.set(old_additional_args)
+
+    logger.info(
+        "[_fan_out_single_branch] completed %d branch tasks for '%s'",
+        len(branch_results),
+        child_id,
+    )
+    return branch_results
+
+
+async def _fan_out_branch(pipeline: "PipelineBase", action_result: Dict[str, Any]) -> List[Dict[str, Any]]:
+    """Delegate tasks to one or more child nodes via ``pipeline.llm_tool``.
+
+    Supports two routing shapes in *action_result*:
+
+    * **Single-branch**: ``{"next_branch": str, "next_tasks": [str, ...]}``
+    – all tasks are sent to the one named child.
+
+    * **Multi-branch**: ``{"next_branches": {branch_id: [str, ...], ...}}``
+    – tasks are grouped per branch and each group is dispatched in order.
+
+    Both shapes may coexist; ``next_branches`` is processed first, then
+    ``next_branch`` / ``next_tasks`` if present.
+
+    Returns a flat list of branch result dicts.  If a child's result is itself
+    an ActionResult with ``next_branch`` / ``next_tasks`` or ``next_branches``,
+    this function recurses automatically.
+
+    Cancellation: checks ``pipeline.cancel_event`` before every branch and
+    every task.  On cancel, returns whatever results were already collected so
+    no work is orphaned and the caller can propagate the signal.
+    """
+    cancel_event = getattr(pipeline, "cancel_event", None)
+    all_results: List[Dict[str, Any]] = []
+
+    # ── Multi-branch path (concurrent) ───────────────────────────────────────
+    next_branches: Dict[str, List[str]] = action_result.get("next_branches") or {}
+    if next_branches:
+        if cancel_event and cancel_event.is_set():
+            logger.info(
+                "[_fan_out_branch] cancel_event set before multi-branch dispatch – aborting",
+            )
+            return all_results
+
+        async def _run_branch(br_id: str, br_tasks: List[str]) -> List[Dict[str, Any]]:
+            if cancel_event and cancel_event.is_set():
+                logger.info(
+                    "[_fan_out_branch] cancel_event set before branch '%s' – skipping",
+                    br_id,
+                )
+                return []
+            logger.info(
+                "[_fan_out_branch] multi-branch dispatch (concurrent): branch='%s', tasks=%s",
+                br_id, br_tasks,
+            )
+            return await _fan_out_single_branch(pipeline, br_id, br_tasks, cancel_event)
+
+        branch_results = await asyncio.gather(
+            *(_run_branch(br_id, br_tasks) for br_id, br_tasks in next_branches.items()),
+            return_exceptions=True,
+        )
+        for item in branch_results:
+            if isinstance(item, BaseException):
+                logger.error("[_fan_out_branch] branch raised an exception: %s", item)
+            else:
+                all_results.extend(item)
+
+    # ── Single-branch path ───────────────────────────────────────────────────
+    single_id: str | None = action_result.get("next_branch")
+    single_tasks: List[str] = list(action_result.get("next_tasks") or [])
+    if single_id and single_tasks:
+        if cancel_event and cancel_event.is_set():
+            logger.info(
+                "[_fan_out_branch] cancel_event set before single branch '%s' – aborting",
+                single_id,
+            )
+            return all_results
+        logger.info(
+            "[_fan_out_branch] single-branch dispatch: branch='%s', tasks=%s",
+            single_id, single_tasks,
+        )
+        results = await _fan_out_single_branch(pipeline, single_id, single_tasks, cancel_event)
+        all_results.extend(results)
+
+    logger.info("[_fan_out_branch] completed total %d results", len(all_results))
+    return all_results
+
+
+class Tool(ToolBase):
+    """Delegate tasks to a agent node."""
+
+    name = "delegate_tool"
+    description = (
+        "Delegate a list of tasks or questions to a specialized agent node. "
+        "The agent will execute these tasks and return their results."
+    )
+    input_schema = {
+        "type": "object",
+        "properties": {
+            "tasks": {
+                "type": "array",
+                "items": {
+                    "type": "string",
+                },
+                "description": "List of tasks or queries to run on the delegated agent.",
+            },
+        },
+        "required": ["tasks"],
+    }
+
+    allowed_callers = ["direct", "code_execution", "mcp_client"]
+
+    fields = [
+        {
+            'name': 'Target Agent',
+            'value': {
+                'type': 'string',
+                'placeholder': 'Agent ID...'
+            }
+        },
+    ]
+
+    async def run_code(self, code: str, task: str) -> Dict[str, Any]:
+        try:
+            if self.code_sandbox and self.workspace and self.tab_id:
+                return await self.code_sandbox.execute(code, task, workspace=self.workspace, tab_id=self.tab_id)
+            else:
+                return {"error": "Setup code_sandbox failed"}
+        except Exception as e:
+            logger.error(f"Failed to run code: {e}")
+            return {"error": str(e)}    
+    
+    async def execute(self, **kwargs) -> Dict[str, Any]:  # noqa: C901
+        from openchadpy.database import Database
+        
+        agent_id: str = self.get_field("Target Agent")
+        tasks: List[str] = kwargs.get("tasks", [])
+        pipeline: "PipelineBase | None" = pipeline_ctx.get()
+
+        if pipeline is None:
+            return {"error": "pipeline is required"}
+
+
+        logger.info(f"[delegate] agent_id={agent_id!r} tasks={tasks!r}")
+
+        if not agent_id:
+            return {"error": "agent_id is required and must not be empty."}
+        if not tasks:
+            return {"error": "tasks array is required and must not be empty."}
+
+        workspace = self.workspace or "global"
+
+        # ── Load agent tree ───────────────────────────────────────────────────
+        # agent_id may be a root tab ID or a child node ID.  We try the fast
+        # path first (treat agent_id as a root); if that yields nothing we scan
+        # all roots in the workspace to find the owning tree.
+        try:
+            agent_node: Optional[dict] = None
+
+            # Fast path: agent_id is itself a root tab ID.
+            agent_db = Database(workspace=workspace, tab_id=agent_id)
+            flat_agents = await agent_db.get("agents")
+            if flat_agents:
+                agent_tree = _build_agent_tree(flat_agents)
+                agent_node = _find_in_tree(agent_tree, agent_id)
+
+            # Slow path: agent_id is a child node — scan all root trees.
+            if agent_node is None:
+                logger.info(
+                    f"[delegate] '{agent_id}' not found as root; scanning all roots in workspace '{workspace}'"
+                )
+                found = await _find_agent_root(workspace, agent_id)
+                if found is None:
+                    return {"error": f"Agent '{agent_id}' not found in any agent tree in workspace '{workspace}'."}
+                _, agent_node = found
+
+        except Exception as e:
+            logger.error(f"[delegate] failed to load agent tree for '{agent_id}': {e}", exc_info=True)
+            return {"error": f"Failed to load agent tree: {e}"}
+
+        # ── Swap context vars for the duration of delegation ─────────────────
+        old_agent  = agent_ctx.get()
+        old_model  = model_id_ctx.get()
+        old_fields = fields_ctx.get()
+        old_additional_args = additional_args_ctx.get()
+
+        try:
+            agent_ctx.set({agent_id: agent_node})
+
+            sub_model = agent_node.get("model")
+            if sub_model:
+                model_id_ctx.set(sub_model)
+
+            fields_ctx.set(json.loads(agent_node.get("toolValues", "{}")))
+            additional_args_ctx.set(json.loads(agent_node.get("additionalArgs", "{}"))) 
+
+            results: List[Dict[str, Any]] = []
+            for task in tasks:
+                try:
+                    ans = await pipeline.llm_tool(query=task, history_id=agent_id)
+                    # Fan-out: if the answer is an ActionResult-like dict, recurse
+                    entry: Dict[str, Any] = {
+                        "agent_id": agent_id,
+                        "task": task,
+                        "response": _extract_response(ans),
+                    }
+                    if isinstance(ans, dict) and ans.get("success") and ans.get("result") is not None:
+                        child_action = _parse_action_result(ans["result"])
+                        if child_action:
+                            has_single = child_action.get("next_branch") and child_action.get("next_tasks")
+                            has_multi  = bool(child_action.get("next_branches"))
+                            if has_single or has_multi:
+                                sub_results = await _fan_out_branch(pipeline, child_action)
+                                entry["sub_branch_results"] = sub_results
+                    results.append(entry)
+                except Exception as e:
+                    logger.error(f"[delegate] task failed for agent '{agent_id}': {e}", exc_info=True)
+                    results.append({"agent_id": agent_id, "task": task, "error": str(e)})
+
+            return {"results": results}
+        finally:
+            agent_ctx.set(old_agent)
+            model_id_ctx.set(old_model)
+            fields_ctx.set(old_fields)
+            additional_args_ctx.set(old_additional_args)
+
+
+

delegate_tool-0.1.0.dist-info/METADATA

@@ -0,0 +1,7 @@
+Metadata-Version: 2.4
+Name: delegate_tool
+Version: 0.1.0
+Summary: Tool to delegating tasks to other agent.
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+Requires-Dist: openchadpy>=0.1.23

delegate_tool-0.1.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+delegate_tool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+delegate_tool/main.py,sha256=Ijz9IUvKNSORRPpoG5gGpHUrdnmYe1VArJSAgRkU8qI,39563
+delegate_tool-0.1.0.dist-info/METADATA,sha256=7JKmBU_QDLgsYElHa2Mf4RI2izvXW2alMSbJ90bpcAc,205
+delegate_tool-0.1.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+delegate_tool-0.1.0.dist-info/top_level.txt,sha256=VAR9kkpIF6eQDBmvPJaudWq9IUWQkcnr_-O7rSxmh00,14
+delegate_tool-0.1.0.dist-info/RECORD,,

delegate_tool-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (83.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

delegate_tool-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+delegate_tool