agentengine-sdk-python 0.8.3__py3-none-any.whl → 0.8.5__py3-none-any.whl

ksadk/harness/engine/subagents.py

@@ -0,0 +1,343 @@
+"""Real child execution, durable approval bridging and cancellation lineage."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import asdict, replace
+
+from ksadk.harness.engine import budgets
+from ksadk.harness.engine.base import ExecutionEngineError
+from ksadk.harness.events import RuntimeEvent
+from ksadk.harness.state import RunStatus
+from ksadk.harness.subagent import (
+    SubAgentExecutionError,
+    SubAgentResult,
+    child_spec,
+    resequence_child_events,
+    validate_subagent_output,
+)
+from ksadk.runtime import ResumePayload, ResumeTarget, RunHandle, StartRequest
+
+
+class ChildApprovalPending(BaseException):
+    """Return a checkpointable child wait to the parent graph routing node."""
+
+    def __init__(self, detail):
+        self.detail = detail
+        super().__init__("child is awaiting an approval decision")
+
+
+async def open_child(*, engine, parent_run, sub, task, call_id):
+    from ksadk.harness.engine.langgraph import ManagedLangGraphEngine
+    from ksadk.harness.engine.thread_ids import encode_thread_id
+
+    child_id = f"{parent_run.handle.run_id}:sub:{sub.name}:{call_id or 'delegation'}"
+    active = engine._active_subagent_runs.setdefault(parent_run.handle.run_id, {})
+    if call_id in active:
+        child_engine, handle = active[call_id]
+        return child_engine, handle, False
+    child_agent = f"{parent_run.state.agent_id}:{sub.name}"
+    checkpoint_session = f"{parent_run.state.session_id}:sub:{sub.name}:{call_id or 'delegation'}"
+    spec = child_spec(parent_run.compiled.spec, sub)
+    if sub.max_total_tokens == 0 or (
+        sub.max_artifacts == 0 and sub.meta.get("requires_artifacts")
+    ):
+        raise SubAgentExecutionError("budget_exhausted", "sub-agent budget is 0")
+    policy = parent_run.execution_policy
+    child_context = (
+        policy.child_system_context
+        if policy is not None and policy.child_system_context is not None
+        else policy.system_context
+        if policy is not None
+        else ""
+    )
+    if child_context:
+        spec = spec.model_copy(
+            update={
+                "prompt": spec.prompt.model_copy(
+                    update={"instructions": sub.instructions + "\n\n" + child_context}
+                )
+            }
+        )
+    # Parent controller ceilings also constrain children without a local config.
+    config = dict(spec.execution_strategy.config)
+    for kind, key in (
+        ("tokens", "max_total_tokens"),
+        ("tools", "max_tool_calls"),
+        ("models", "max_model_calls"),
+        ("artifacts", "max_artifacts"),
+    ):
+        inherited = budgets.limit(parent_run, kind)
+        if inherited is not None:
+            local = config.get(key)
+            config[key] = inherited if local is None else min(inherited, int(local))
+    spec = spec.model_copy(
+        update={"execution_strategy": spec.execution_strategy.model_copy(update={"config": config})}
+    )
+    artifact_store = getattr(engine._mcp_disclosure, "_artifact_store", None)
+    if artifact_store is not None and sub.max_artifacts is not None:
+        from ksadk.harness.artifact_store import BudgetedArtifactStore
+
+        artifact_store = BudgetedArtifactStore(
+            artifact_store, run_id=child_id, max_artifacts=sub.max_artifacts
+        )
+    child_engine = ManagedLangGraphEngine(
+        reasoner=engine._reasoner,
+        checkpointer=engine._checkpointer,
+        tenant_id=engine._tenant_id,
+        tools={name: parent_run.tools[name] for name in sub.tools if name in parent_run.tools},
+        approval_required=set(parent_run.approval_required) & set(sub.tools),
+        capability_runtime=engine._capability_runtime,
+        context_engine=engine._context_engine,
+        memory_runtime=engine._memory_runtime,
+        skill_runtime=engine._skill_runtime if sub.inherit_skills else None,
+        mcp_runtime=engine._mcp_runtime if sub.inherit_mcp else None,
+        artifact_store=artifact_store,
+        mcp_offload_policy=engine._mcp_disclosure._offload_policy,
+        max_reasoning_turns=min(sub.max_turns, engine._max_reasoning_turns),
+    )
+    child_engine._inherited_policy_tool_names = parent_run.policy_tool_names & set(sub.tools)
+    compiled = await child_engine.compile(spec)
+    request = StartRequest(
+        input=task,
+        user_id=parent_run.state.user_id,
+        session_id=parent_run.state.session_id,
+        agent_id=child_agent,
+        runtime_type="managed-langgraph",
+        metadata={
+            # Preserve host context, but never inject the parent's conversation as
+            # the child's task or reuse the parent's checkpoint namespace.
+            **{
+                key: value
+                for key, value in parent_run.request.metadata.items()
+                if key
+                not in {"conversation_history", "invocation_id", "run_id", "checkpoint_session_id"}
+            },
+            "invocation_id": child_id,
+            "parent_run_id": parent_run.handle.run_id,
+            "checkpoint_session_id": checkpoint_session,
+            "authorization_agent_id": parent_run.request.metadata.get(
+                "authorization_agent_id", parent_run.state.agent_id
+            ),
+        },
+    )
+    candidate = RunHandle(
+        run_id=child_id,
+        session_id=parent_run.state.session_id,
+        runtime_type="managed-langgraph",
+        native_ref={
+            "thread_id": encode_thread_id(
+                tenant_id=engine._tenant_id,
+                user_id=parent_run.state.user_id,
+                agent_id=child_agent,
+                session_id=checkpoint_session,
+                run_id=child_id,
+            ),
+            "user_id": parent_run.state.user_id,
+            "agent_id": child_agent,
+            "parent_run_id": parent_run.handle.run_id,
+            "authorization_agent_id": request.metadata["authorization_agent_id"],
+        },
+    )
+    options = dict(
+        execution_policy=parent_run.execution_policy,
+        execution_policy_resolver=parent_run.execution_policy_resolver,
+        execution_policy_request=parent_run.execution_policy_request,
+    )
+    handle, recovered = None, False
+    if engine._checkpointer is not None:
+        try:
+            handle = await child_engine.attach(candidate, compiled, allow_completed=True, **options)
+            recovered = True
+        except ExecutionEngineError as exc:
+            if "无未决 Checkpoint" not in str(exc):
+                raise
+    if handle is None:
+        handle = await child_engine.start(request, compiled, **options)
+    child_run = child_engine._runs[handle.run_id]
+    child_run.request = child_run.request.model_copy(update={
+        "metadata": {**child_run.request.metadata, **request.metadata}
+    })
+    budgets.adopt_child(parent_run, child_run)
+    active[call_id] = (child_engine, handle)
+    return child_engine, handle, recovered
+
+
+async def execute_subagent(*, engine, parent_run, sub, task, call_id):
+    child_engine, handle, recovered = await open_child(
+        engine=engine, parent_run=parent_run, sub=sub, task=task, call_id=call_id
+    )
+    child_run = child_engine._runs[handle.run_id]
+    collected: list[RuntimeEvent] = []
+    final_text = next(
+        (m.content for m in reversed(child_run.state.messages) if m.role.value == "assistant"), ""
+    )
+    failure, category = "", ""
+
+    def forward(event):
+        nonlocal final_text, failure, category
+        collected.append(event)
+        resequence_child_events(parent_run, [event], observe=True)
+        if event.event_type == "text.completed" and event.phase == "final_answer":
+            final_text = str(event.payload.get("text") or "")
+        if event.event_type == "run.failed" and not failure:
+            failure = str(event.payload.get("error") or "child run failed")
+            category = str(event.payload.get("error_category") or "child_failed")
+        if (event.event_type == "tool.call.end"
+                and event.payload.get("error_category") == "budget_exhausted"):
+            failure, category = str(event.payload["error"]), "budget_exhausted"
+
+    try:
+        while not child_run.done:
+            if child_run.state.status is RunStatus.AWAITING_APPROVAL:
+                if engine._checkpointer is None:
+                    raise SubAgentExecutionError(
+                        "approval_unavailable", "child approval requires checkpointing"
+                    )
+                info = dict(child_run.pending_approval)
+                # The parent interrupt persists the exact child Run/Call. Its
+                # decision is consumed only by that child's pending native call.
+                detail = {
+                    "call_id": call_id,
+                    "name": sub.name,
+                    "risk": "high",
+                    "kind": "child_tool",
+                    "child_run_id": handle.run_id,
+                    "parent_run_id": parent_run.handle.run_id,
+                    "child_call_id": child_run.pending_approval_call_id,
+                    "child_handle": handle.model_dump(mode="json"),
+                    "child_approval": info,
+                    "child_spec": asdict(sub),
+                }
+                resolved = parent_run.child_approval_decision
+                if not resolved:
+                    raise ChildApprovalPending(detail)
+                if (
+                    resolved.get("child_run_id") != handle.run_id
+                    or resolved.get("child_call_id") != child_run.pending_approval_call_id
+                ):
+                    raise ExecutionEngineError(
+                        "child approval decision does not match pending execution"
+                    )
+                decision = resolved["decision"]
+                parent_run.child_approval_decision = None
+                await child_engine.resume(
+                    handle,
+                    ResumeTarget(kind="thread_id", id=handle.native_ref["thread_id"]),
+                    ResumePayload(
+                        kind="approval_decision",
+                        call_id=child_run.pending_approval_call_id,
+                        data=decision,
+                    ),
+                )
+            elif child_run.state.status is RunStatus.PAUSED:
+                await child_engine.resume(
+                    handle, ResumeTarget(kind="thread_id", id=handle.native_ref["thread_id"]), None
+                )
+            # Approval wait time is excluded. Each executing segment remains
+            # bounded and child limits remain cumulative across checkpoints.
+            remaining_seconds = sub.timeout_seconds - child_run.execution_elapsed_seconds
+            if remaining_seconds <= 0:
+                raise TimeoutError
+
+            async def consume_segment():
+                async for event in child_engine.stream(handle):
+                    forward(event)
+                    if category == "budget_exhausted" and not child_run.done:
+                        await child_engine.cancel(handle)
+
+            await asyncio.wait_for(consume_segment(), timeout=remaining_seconds)
+            if child_run.state.status is RunStatus.AWAITING_APPROVAL:
+                continue
+            if not child_run.done:
+                raise SubAgentExecutionError(
+                    "child_interrupted", "child stopped without a terminal state"
+                )
+    except (TimeoutError, asyncio.TimeoutError):
+        failure, category = f"sub-agent {sub.name!r} timed out", "timeout"
+        await cancel_and_drain(child_engine, handle, forward)
+    except asyncio.CancelledError:
+        # Cancellation must await the child's native terminal event before the
+        # parent publishes its terminal state; no detached child continues.
+        await cancel_and_drain(child_engine, handle, forward)
+        raise
+    finally:
+        if child_run.done:
+            active = engine._active_subagent_runs.get(parent_run.handle.run_id, {})
+            active.pop(call_id, None)
+            if not active:
+                engine._active_subagent_runs.pop(parent_run.handle.run_id, None)
+    if child_run.state.status is RunStatus.CANCELED and not failure:
+        failure, category = "child run canceled", "canceled"
+    if child_run.state.status is RunStatus.FAILED and not failure:
+        failure, category = "child run failed", "child_failed"
+    if (
+        failure
+        and sub.failure_policy == "retry"
+        and sub.max_retries > 0
+        and category in {"child_failed", "timeout"}
+    ):
+        result, events = await execute_subagent(
+            engine=engine,
+            parent_run=parent_run,
+            sub=replace(sub, max_retries=sub.max_retries - 1),
+            task=task,
+            call_id=f"{call_id}:retry-{sub.max_retries}",
+        )
+        return result, [*collected, *events]
+    if failure and sub.failure_policy != "return_error":
+        raise SubAgentExecutionError(category, failure)
+    output = None
+    if not failure:
+        try:
+            output = validate_subagent_output(final_text, sub.output_schema)
+        except ValueError as exc:
+            raise SubAgentExecutionError("output_validation_failed", str(exc)) from exc
+    return SubAgentResult(
+        status="failed" if failure else "completed",
+        output=output,
+        artifact_refs=tuple(child_run.artifact_refs),
+        usage={
+            "total_tokens": child_run.budget_usage.get("tokens", 0),
+            "input_tokens": child_run.budget_usage.get("input_tokens", 0),
+            "output_tokens": child_run.budget_usage.get("output_tokens", 0),
+        },
+        evidence={
+            "child_agent_id": child_run.state.agent_id,
+            "child_run_id": handle.run_id,
+            "parent_run_id": parent_run.handle.run_id,
+            "child_handle": handle.model_dump(mode="json"),
+            "recovered_from_checkpoint": recovered,
+        },
+        error_category=category or None,
+        error=failure or None,
+    ).model_dump(mode="json"), collected
+
+
+async def cancel_and_drain(child_engine, handle, forward):
+    await child_engine.cancel(handle)
+    run = child_engine._runs[handle.run_id]
+    if run.task is not None and not run.task.done():
+        await asyncio.gather(run.task, return_exceptions=True)
+    async for event in child_engine.stream(handle):
+        forward(event)
+
+
+async def restore_pending_child(engine, run):
+    info = run.pending_approval
+    if info.get("kind") != "child_tool":
+        return
+    sub = run.sub_agents.get(str(info.get("name") or ""))
+    if sub is None and info.get("child_spec"):
+        from ksadk.harness.spec import SubAgentBinding
+        from ksadk.harness.subagent import SubAgentSpec
+
+        declaration = dict(info["child_spec"])
+        metadata = declaration.pop("meta", {})
+        sub = replace(
+            SubAgentSpec.from_binding(SubAgentBinding.model_validate(declaration)),
+            meta=metadata,
+        )
+    if sub is None:
+        raise ExecutionEngineError("pending child definition is unavailable")
+    await open_child(engine=engine, parent_run=run, sub=sub, task="", call_id=str(info["call_id"]))

ksadk/harness/engine/thread_ids.py

@@ -0,0 +1,112 @@
+"""Checkpointer thread_id 租户复合编码（plan §6.2.2，Kernel 决策）。
+
+LangGraph Checkpointer 的 ``thread_id`` 是扁平字符串，无租户语义。本模块
+是唯一编码实现，禁止各处手拼。Approval 跨进程恢复依赖该编码。
+"""
+
+from __future__ import annotations
+
+import re
+
+from pydantic import BaseModel, ConfigDict, Field
+
+_PATTERN = re.compile(
+    r"^tenant:(?P<tenant>[^/]+)"
+    r"/user:(?P<user>[^/]+)"
+    r"/agent:(?P<agent>[^/]+)"
+    r"/session:(?P<session>[^/]+)"
+    r"/run:(?P<run>[^/]+)$"
+)
+
+
+class ThreadIdError(ValueError):
+    """thread_id 编码不合法或与声明的租户不匹配。"""
+
+
+class HarnessThreadId(BaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    tenant_id: str = Field(min_length=1, max_length=128)
+    user_id: str = Field(min_length=1, max_length=128)
+    agent_id: str = Field(min_length=1, max_length=160)
+    session_id: str = Field(min_length=1, max_length=160)
+    run_id: str = Field(min_length=1, max_length=160)
+
+    def encode(self) -> str:
+        return (
+            f"tenant:{self.tenant_id}/user:{self.user_id}/agent:{self.agent_id}"
+            f"/session:{self.session_id}/run:{self.run_id}"
+        )
+
+    def tenant_prefix(self) -> str:
+        """租户隔离过滤前缀（Checkpointer 读取侧使用）。"""
+        return f"tenant:{self.tenant_id}/"
+
+
+def encode_thread_id(
+    *,
+    tenant_id: str,
+    user_id: str,
+    agent_id: str,
+    session_id: str,
+    run_id: str,
+) -> str:
+    """构造复合 thread_id。任何字段含 ``/`` 视为注入攻击，拒绝。"""
+    for name, value in (
+        ("tenant_id", tenant_id),
+        ("user_id", user_id),
+        ("agent_id", agent_id),
+        ("session_id", session_id),
+        ("run_id", run_id),
+    ):
+        if not value or "/" in value:
+            raise ThreadIdError(f"{name} 不能为空且不能包含 '/': {value!r}")
+    return HarnessThreadId(
+        tenant_id=tenant_id,
+        user_id=user_id,
+        agent_id=agent_id,
+        session_id=session_id,
+        run_id=run_id,
+    ).encode()
+
+
+def decode_thread_id(thread_id: str) -> HarnessThreadId:
+    match = _PATTERN.match(thread_id)
+    if match is None:
+        raise ThreadIdError(
+            f"thread_id 不符合租户复合编码 tenant:*/user:*/agent:*/session:*/run:*: {thread_id!r}"
+        )
+    groups = match.groupdict()
+    return HarnessThreadId(
+        tenant_id=groups["tenant"],
+        user_id=groups["user"],
+        agent_id=groups["agent"],
+        session_id=groups["session"],
+        run_id=groups["run"],
+    )
+
+
+def validate_tenant(thread_id: str, *, tenant_id: str) -> None:
+    """写入/读取侧强制校验：thread_id 必须落在声明租户下。"""
+    decoded = decode_thread_id(thread_id)
+    if decoded.tenant_id != tenant_id:
+        raise ThreadIdError(f"thread_id 租户不匹配: {decoded.tenant_id!r} != {tenant_id!r}")
+
+
+def matches_tenant(thread_id: str, *, tenant_id: str) -> bool:
+    """读取侧前缀过滤（宽松版，编码不合法返回 False 而非抛错）。"""
+    try:
+        validate_tenant(thread_id, tenant_id=tenant_id)
+    except ThreadIdError:
+        return False
+    return True
+
+
+__all__ = [
+    "HarnessThreadId",
+    "ThreadIdError",
+    "decode_thread_id",
+    "encode_thread_id",
+    "matches_tenant",
+    "validate_tenant",
+]

ksadk/harness/engine/tool_dispatch.py

@@ -0,0 +1,144 @@
+"""Governed tool dispatch shared by root and child Managed executions."""
+
+from __future__ import annotations
+
+import inspect
+from typing import Any
+
+from ksadk.harness.engine.mcp_disclosure import McpDisclosureCursors
+from ksadk.harness.events import EventType, RuntimeEvent
+from ksadk.harness.run_control import ControlAction, RunControlReplan, RunControlStop
+
+
+async def invoke_tool(
+    self,
+    name: str,
+    arguments: dict[str, Any],
+    *,
+    run: Any = None,
+    mcp_cursors: McpDisclosureCursors | None = None,
+    call_id: str = "",
+) -> Any:
+    if run is not None:
+        from ksadk.harness.engine import budgets
+
+        # A model that reports over-budget usage cannot spend a separate tool
+        # allowance to perform a side effect after exhausting the tree budget.
+        budgets.check(run, "tokens", 0)
+        if run.controller is not None and (not call_id or call_id not in run.budget_tool_calls):
+            exempt_tools = set(
+                run.compiled.spec.execution_strategy.config.get("stagnation_exempt_tools", ())
+            )
+            decision = run.controller.before_tool(
+                name,
+                arguments,
+                stagnation_exempt=name in exempt_tools,
+            )
+            if decision.action == ControlAction.STOP:
+                raise RunControlStop(decision.reason)
+            if decision.action == ControlAction.REPLAN:
+                raise RunControlReplan(decision.reason)
+        if not call_id or call_id not in run.budget_tool_calls:
+            budgets.check(run, "tools")
+            budgets.spend(run, "tools")
+            if call_id:
+                run.budget_tool_calls.add(call_id)
+            run.tool_calls_started += 1
+    # 收口 6：子 Agent 即工具——内联运行到完成，子事件并入父流。
+    sub = (run.sub_agents if run is not None else self._sub_agents).get(name)
+    if sub is not None and run is not None:
+        from ksadk.harness.subagent import run_subagent
+
+        text, _ = await run_subagent(
+            engine=self,
+            parent_run=run,
+            sub=sub,
+            task=str((arguments or {}).get("task") or ""),
+            call_id=call_id,
+        )
+        return text
+    if self._skill_disclosure.is_tool(name):
+        return self._invoke_skill_tool(name, arguments, run=run)
+    if self._mcp_disclosure.is_tool(name):
+        if mcp_cursors is None:
+            mcp_cursors = McpDisclosureCursors()
+        return await self._invoke_mcp_tool(
+            name,
+            arguments,
+            run=run,
+            cursors=mcp_cursors,
+            call_id=call_id,
+        )
+    if self._delegation_runtime is not None and self._delegation_runtime.is_tool(name):
+        if run is None:
+            raise RuntimeError("dynamic delegation requires an active parent run")
+        result, child_events = await self._delegation_runtime.invoke(
+            engine=self, parent_run=run, arguments=arguments, call_id=call_id or name,
+        )
+        self._pending_subagent_events.setdefault(run.handle.run_id, {})[call_id or name] = (
+            child_events
+        )
+        return result
+    tool = (run.tools if run is not None else self._tools).get(name)
+    if tool is None:
+        raise RuntimeError(
+            f"engine tool {name!r} is not available; it may be filtered or unpublished"
+        )
+    if run is not None and callable(getattr(tool, "drain_artifacts", None)):
+        artifact_budget = run.compiled.spec.execution_strategy.config.get("max_artifacts")
+        declared_cost = getattr(tool, "artifact_budget_cost", None)
+        if artifact_budget is not None and declared_cost is None:
+            from ksadk.harness.subagent import SubAgentExecutionError
+
+            raise SubAgentExecutionError(
+                "budget_exhausted",
+                "artifact-producing tool lacks artifact_budget_cost; "
+                "finite budget cannot be proven before execution",
+            )
+        if artifact_budget is not None:
+            artifact_cost = int(declared_cost)
+            if artifact_cost < 0 or (run.artifacts_created + artifact_cost > int(artifact_budget)):
+                from ksadk.harness.subagent import SubAgentExecutionError
+
+                raise SubAgentExecutionError(
+                    "budget_exhausted",
+                    f"execution artifact budget {artifact_budget} exhausted before tool execution",
+                )
+            run.artifacts_created += artifact_cost
+            budgets.check(run, "artifacts", artifact_cost)
+            budgets.spend(run, "artifacts", artifact_cost)
+    call = getattr(tool, "call", None)
+    from ksadk.runtime_context import tool_execution_scope
+
+    with tool_execution_scope(
+        run.state.session_id if run is not None else "",
+        run.handle.run_id if run is not None else "",
+    ):
+        target = call if callable(call) else tool
+        parameters = inspect.signature(target).parameters
+        accepts_call_id = "call_id" in parameters or any(
+            value.kind is inspect.Parameter.VAR_KEYWORD for value in parameters.values()
+        )
+        result = await target(arguments, **({"call_id": call_id} if accepts_call_id else {}))
+    # 缺口 5：工具产出的 Artifact → artifact.created 事件（与子事件同
+    # 缓冲，tool.call.end 之后统一重排并入，seq 单调）。
+    drain = getattr(tool, "drain_artifacts", None)
+    if callable(drain) and run is not None:
+        for artifact in drain():
+            self._pending_child_events.setdefault(run.handle.run_id, []).append(
+                RuntimeEvent.create(
+                    EventType.ARTIFACT_CREATED,
+                    agent_id=run.state.agent_id,
+                    user_id=run.state.user_id,
+                    session_id=run.state.session_id,
+                    invocation_id=run.handle.run_id,
+                    seq_id=0,
+                    payload={
+                        "name": str(artifact.get("name") or ""),
+                        "version": int(artifact.get("version") or 1),
+                        "uri": str(artifact.get("uri") or ""),
+                        "mime": str(artifact.get("mime") or "text/plain"),
+                    },
+                )
+            )
+    return result