union-app-chat-stream 1.1.6

package/app/agent/runtime/activity.py

@@ -0,0 +1,393 @@
+"""收集根执行和子执行消息，并生成标准 AG-UI Activity 快照。"""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import Any
+
+from ag_ui.core import (
+    ActivitySnapshotEvent,
+    EventType,
+    ReasoningMessageContentEvent,
+    ReasoningMessageStartEvent,
+    TextMessageContentEvent,
+    TextMessageStartEvent,
+    ToolCallArgsEvent,
+    ToolCallResultEvent,
+    ToolCallStartEvent,
+)
+
+from app.agent.runtime.delegation import Delegation, PlanReferenceError
+
+
+class MessageReducer:
+    """Reduce standard AG-UI message events by message/tool-call identity."""
+
+    def __init__(self, on_new_message=None) -> None:
+        self.messages: list[dict[str, Any]] = []
+        self._by_id: dict[str, dict[str, Any]] = {}
+        self._tool_calls: dict[str, dict[str, Any]] = {}
+        self._on_new_message = on_new_message
+
+    def add_messages(self, messages: list[dict[str, Any]]) -> None:
+        for raw in messages:
+            message = dict(raw)
+            message_id = str(message.get("id") or "")
+            if not message_id:
+                continue
+            existing = self._by_id.get(message_id)
+            if existing is None:
+                self._append(message)
+            else:
+                existing.clear()
+                existing.update(message)
+            self._index_tool_calls(self._by_id[message_id])
+
+    def apply(self, event: Any) -> None:
+        if isinstance(event, TextMessageStartEvent):
+            message = {"id": event.message_id, "role": event.role, "content": ""}
+            if event.name:
+                message["name"] = event.name
+            self._ensure(message)
+        elif isinstance(event, TextMessageContentEvent):
+            self._content(event.message_id, event.delta, "assistant")
+        elif isinstance(event, ReasoningMessageStartEvent):
+            self._ensure({"id": event.message_id, "role": "reasoning", "content": ""})
+        elif isinstance(event, ReasoningMessageContentEvent):
+            self._content(event.message_id, event.delta, "reasoning")
+        elif isinstance(event, ToolCallStartEvent):
+            message_id = event.parent_message_id or f"tool-parent:{event.tool_call_id}"
+            message = self._ensure(
+                {"id": message_id, "role": "assistant", "content": "", "toolCalls": []}
+            )
+            calls = message.setdefault("toolCalls", [])
+            call = {
+                "id": event.tool_call_id,
+                "type": "function",
+                "function": {"name": event.tool_call_name, "arguments": ""},
+            }
+            calls.append(call)
+            self._tool_calls[event.tool_call_id] = call
+        elif isinstance(event, ToolCallArgsEvent):
+            call = self._tool_calls.get(event.tool_call_id)
+            if call is not None:
+                call["function"]["arguments"] += event.delta
+        elif isinstance(event, ToolCallResultEvent):
+            message = {
+                "id": event.message_id,
+                "role": "tool",
+                "toolCallId": event.tool_call_id,
+                "content": event.content,
+            }
+            self._ensure(message)
+
+    def _content(self, message_id: str, delta: str, role: str) -> None:
+        message = self._ensure({"id": message_id, "role": role, "content": ""})
+        message["content"] = str(message.get("content") or "") + delta
+
+    def _ensure(self, message: dict[str, Any]) -> dict[str, Any]:
+        existing = self._by_id.get(str(message["id"]))
+        if existing is not None:
+            return existing
+        self._append(message)
+        return message
+
+    def _append(self, message: dict[str, Any]) -> None:
+        message_id = str(message["id"])
+        self.messages.append(message)
+        self._by_id[message_id] = message
+        if self._on_new_message is not None:
+            self._on_new_message(message)
+        self._index_tool_calls(message)
+
+    def _index_tool_calls(self, message: dict[str, Any]) -> None:
+        for call in message.get("toolCalls") or []:
+            call_id = call.get("id")
+            if call_id:
+                self._tool_calls[str(call_id)] = call
+
+
+@dataclass
+class ChildExecution:
+    delegation: Delegation
+    status: str = "running"
+    error_code: str | None = None
+    reducer: MessageReducer | None = None
+    transformer: Any | None = None
+
+
+@dataclass
+class RunCollector:
+    root_run_id: str
+    conversation_id: str
+    root_agent_name: str = "UnionCoordinatorAgent"
+    root_status: str = "running"
+    root_error_code: str | None = None
+    children: dict[str, ChildExecution] = field(default_factory=dict)
+    ordered_messages: list[tuple[str, dict[str, Any]]] = field(default_factory=list)
+    canonical_root_messages: list[dict[str, Any]] | None = None
+    plans: dict[str, dict[str, Any]] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        self.root = MessageReducer(
+            lambda message: self.ordered_messages.append((self.root_run_id, message))
+        )
+
+    def start_child(self, delegation: Delegation) -> ChildExecution:
+        child = ChildExecution(delegation=delegation)
+        child.reducer = MessageReducer(
+            lambda message: self.ordered_messages.append((delegation.run_id, message))
+        )
+        self.children[delegation.run_id] = child
+        return child
+
+    def finish_child(
+        self, run_id: str, status: str, error_code: str | None = None
+    ) -> ChildExecution:
+        child = self.children[run_id]
+        child.status = status
+        child.error_code = error_code
+        return child
+
+    def validate_plan_step(self, index: int) -> None:
+        plan = self.plans.get(self.root_run_id)
+        if plan is None or index < 1 or index > len(plan["items"]):
+            raise PlanReferenceError("plan step does not exist")
+        if any(
+            child.status == "running" and child.delegation.plan_index == index
+            for child in self.children.values()
+        ):
+            raise PlanReferenceError("plan step is already running")
+
+    def set_root_messages(self, messages: list[dict[str, Any]]) -> None:
+        """Combine trusted ingress users with the Agent's official message dump."""
+        ingress_users = [
+            deepcopy(message)
+            for message in self.root.messages
+            if message.get("role") == "user"
+        ]
+        seen_ids = {
+            str(message["id"])
+            for message in ingress_users
+            if message.get("id") is not None
+        }
+        self.canonical_root_messages = ingress_users
+        for message in messages:
+            message_id = message.get("id")
+            if message_id is not None and str(message_id) in seen_ids:
+                continue
+            self.canonical_root_messages.append(deepcopy(message))
+            if message_id is not None:
+                seen_ids.add(str(message_id))
+
+    def snapshot(self, run_id: str) -> ActivitySnapshotEvent:
+        child = self.children[run_id]
+        metadata = child.delegation
+        return ActivitySnapshotEvent(
+            message_id=f"subagent:{run_id}",
+            activity_type="subagent_execution",
+            replace=True,
+            content={
+                "runId": run_id,
+                "parentRunId": metadata.parent_run_id,
+                "agentName": metadata.agent_name,
+                "delegationToolCallId": metadata.delegation_tool_call_id,
+                "task": metadata.task,
+                "status": child.status,
+                "messages": child.reducer.messages if child.reducer else [],
+                "errorCode": child.error_code,
+            },
+        )
+
+    def update_plan(
+        self,
+        *,
+        run_id: str,
+        parent_run_id: str | None,
+        agent_name: str,
+        skill_ids: list[str],
+        items: list[dict[str, str]],
+    ) -> ActivitySnapshotEvent:
+        """Project one complete plan as a replaceable AG-UI Activity."""
+        if run_id == self.root_run_id and run_id in self.plans:
+            existing_items = self.plans[run_id]["items"]
+            for child in self.children.values():
+                index = child.delegation.plan_index
+                if index is None:
+                    continue
+                if (
+                    index > len(items)
+                    or index > len(existing_items)
+                    or items[index - 1]["content"]
+                    != existing_items[index - 1]["content"]
+                ):
+                    raise PlanReferenceError("bound plan step changed")
+        current = next(
+            (index for index, item in enumerate(items, 1) if item["status"] == "in_progress"),
+            None,
+        )
+        if current is None:
+            current = next(
+                (index for index, item in enumerate(items, 1) if item["status"] == "pending"),
+                len(items) or None,
+            )
+        status = (
+            "running"
+            if any(item["status"] in {"pending", "in_progress"} for item in items)
+            else "completed"
+        )
+        content = {
+            "runId": run_id,
+            "parentRunId": parent_run_id,
+            "agentName": agent_name,
+            "skillIds": skill_ids,
+            "currentIndex": current,
+            "total": len(items),
+            "status": status,
+            "items": [dict(item) for item in items],
+        }
+        self.plans[run_id] = content
+        event = self.plan_snapshot(run_id)
+        assert event is not None
+        return event
+
+    def plan_snapshot(self, run_id: str) -> ActivitySnapshotEvent | None:
+        content = self.plans.get(run_id)
+        if content is None:
+            return None
+        projected = deepcopy(content)
+        if run_id == self.root_run_id:
+            for child in self.children.values():
+                index = child.delegation.plan_index
+                if index is None or index > len(projected["items"]):
+                    continue
+                projected["items"][index - 1]["status"] = (
+                    "in_progress" if child.status == "running" else child.status
+                )
+        current = next(
+            (
+                index
+                for index, item in enumerate(projected["items"], 1)
+                if item["status"] == "in_progress"
+            ),
+            None,
+        )
+        if current is None:
+            current = next(
+                (
+                    index
+                    for index, item in enumerate(projected["items"], 1)
+                    if item["status"] == "pending"
+                ),
+                len(projected["items"]) or None,
+            )
+        statuses = {item["status"] for item in projected["items"]}
+        projected["currentIndex"] = current
+        projected["status"] = (
+            "running"
+            if statuses & {"pending", "in_progress"}
+            else "failed"
+            if "failed" in statuses
+            else "cancelled"
+            if "cancelled" in statuses
+            else "completed"
+        )
+        return ActivitySnapshotEvent(
+            message_id=f"task-plan:{run_id}",
+            activity_type="task_plan",
+            replace=True,
+            content=projected,
+        )
+
+    def finish_plan(
+        self, run_id: str, status: str
+    ) -> ActivitySnapshotEvent | None:
+        content = self.plans.get(run_id)
+        if content is None or content["status"] == status:
+            return None
+        content = {**content, "status": status}
+        self.plans[run_id] = content
+        return ActivitySnapshotEvent(
+            message_id=f"task-plan:{run_id}",
+            activity_type="task_plan",
+            replace=True,
+            content=content,
+        )
+
+    def apply_root_event(self, event: Any) -> None:
+        self.root.apply(event)
+        if getattr(event, "type", None) == EventType.RUN_ERROR:
+            self.root_status = "failed"
+            self.root_error_code = getattr(event, "code", None) or "run_error"
+        elif getattr(event, "type", None) == EventType.RUN_FINISHED:
+            self.root_status = "completed"
+
+    def completion_payload(self) -> dict[str, Any]:
+        executions = [
+            {
+                "runId": self.root_run_id,
+                "parentRunId": None,
+                "agentName": self.root_agent_name,
+                "delegationToolCallId": None,
+                "task": None,
+                "status": self.root_status,
+                "errorCode": self.root_error_code,
+            }
+        ]
+        executions.extend(
+            {
+                "runId": child.delegation.run_id,
+                "parentRunId": child.delegation.parent_run_id,
+                "agentName": child.delegation.agent_name,
+                "delegationToolCallId": child.delegation.delegation_tool_call_id,
+                "task": child.delegation.task,
+                "status": child.status,
+                "errorCode": child.error_code,
+            }
+            for child in self.children.values()
+        )
+        root_messages = (
+            self.canonical_root_messages
+            if self.canonical_root_messages is not None
+            else self.root.messages
+        )
+        messages = [
+            {"runId": self.root_run_id, "message": persisted}
+            for message in root_messages
+            if (
+                persisted := self._persisted_message(
+                    self.root_run_id,
+                    message,
+                )
+            )
+            is not None
+        ]
+        messages.extend(
+            {"runId": run_id, "message": persisted}
+            for run_id, message in self.ordered_messages
+            if run_id != self.root_run_id
+            and (persisted := self._persisted_message(run_id, message)) is not None
+        )
+        return {
+            "conversationId": self.conversation_id,
+            "runId": self.root_run_id,
+            "status": self.root_status,
+            "errorCode": self.root_error_code,
+            "executions": executions,
+            "messages": messages,
+        }
+
+    def _persisted_message(
+        self, run_id: str, message: dict[str, Any]
+    ) -> dict[str, Any] | None:
+        value = deepcopy(message)
+        if (
+            run_id == self.root_run_id
+            and self.root_status != "completed"
+            and value.get("role") == "assistant"
+        ):
+            if not value.get("toolCalls"):
+                return None
+            value["content"] = ""
+        return value

package/app/agent/runtime/delegation.py

@@ -0,0 +1,80 @@
+"""定义 delegation 子执行的元数据、观察器和任务局部生命周期。"""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from contextlib import asynccontextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Any
+
+
+class PlanReferenceError(ValueError):
+    """Reject an invalid live binding between one root-plan step and a child."""
+
+
+@dataclass
+class Delegation:
+    run_id: str
+    parent_run_id: str
+    agent_name: str
+    delegation_tool_call_id: str
+    task: str
+    plan_index: int | None = None
+    terminal_status: str | None = None
+    terminal_error_code: str | None = None
+
+
+@dataclass(frozen=True)
+class DelegationObserver:
+    on_start: Callable[[Delegation], Awaitable[None]]
+    on_event: Callable[[Delegation, Any], Awaitable[None]]
+    on_finish: Callable[[Delegation, str, str | None], Awaitable[None]]
+
+
+current_delegation: ContextVar[Delegation | None] = ContextVar(
+    "current_delegation",
+    default=None,
+)
+
+
+@asynccontextmanager
+async def delegation_scope(ctx, delegation: Delegation):
+    """Notify one observer exactly once around a delegate_task call."""
+    observer = ctx.deps.delegation_observer
+    token = current_delegation.set(delegation)
+    started = False
+    try:
+        if observer is not None:
+            await observer.on_start(delegation)
+            started = True
+        yield
+    except asyncio.CancelledError:
+        if ctx.deps.cancelled.is_set():
+            delegation.terminal_status = "cancelled"
+            delegation.terminal_error_code = "cancelled"
+        if observer is not None and started:
+            await observer.on_finish(
+                delegation,
+                delegation.terminal_status or "failed",
+                delegation.terminal_error_code or "subagent_timeout",
+            )
+        raise
+    except Exception as exc:
+        if observer is not None and started:
+            await observer.on_finish(
+                delegation,
+                delegation.terminal_status or "failed",
+                delegation.terminal_error_code or type(exc).__name__,
+            )
+        raise
+    else:
+        if observer is not None and started:
+            await observer.on_finish(
+                delegation,
+                delegation.terminal_status or "completed",
+                delegation.terminal_error_code,
+            )
+    finally:
+        current_delegation.reset(token)

package/app/agent/runtime/deps.py

@@ -0,0 +1,34 @@
+"""定义服务端为一次 Agent 运行注入的可信依赖。"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from app.agent.runtime.delegation import DelegationObserver
+    from app.utils.api_client import ApiClient
+
+
+@dataclass
+class RunDeps:
+    """Trusted, server-created dependencies for one Agent run."""
+
+    user_id: str
+    permissions: tuple[str, ...]
+    state_client: Any
+    api_client: ApiClient
+    cancelled: Any
+    authentication_type: str = "CAS"
+    org_code: str = ""
+    scheduled_run_id: int | None = None
+    rag_service: Any | None = None
+    effective_at: str | None = None
+    effective_timezone: str | None = None
+    planned_run_ids: set[str] = field(default_factory=set)
+    delegated_agents_by_run: dict[str, set[str]] = field(default_factory=dict)
+    delegation_observer: DelegationObserver | None = None
+    plan_observer: Callable[[Any, list[dict[str, str]]], Awaitable[None]] | None = None
+    delegation_plan_validator: Callable[[int], None] | None = None