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

ksadk/plugins/teams/tool_service.py

@@ -0,0 +1,152 @@
+"""Shared trusted Teams tool schemas and domain invocation.
+
+Hosts provide a verified in-process Actor, the same domain transaction, and
+artifact methods appropriate to their storage. This module has no Studio,
+runtime, filesystem, transport or model-provider dependency.
+"""
+
+from .contracts import Actor, TaskCreateInput
+from .errors import TeamsError
+
+
+class TeamsToolMethods:
+    @staticmethod
+    def tool_definitions():
+        def schema(properties, required=()):
+            return {
+                "type": "object",
+                "properties": properties,
+                "required": list(required),
+                "additionalProperties": False,
+            }
+
+        text = {"type": "string"}
+        return [
+            ("team_context", "读取当前团队消息、成员和任务状态。", schema({})),
+            (
+                "team_message",
+                "发布进度，或明确请求某成员执行；wake=false仅留言，不触发执行。",
+                schema(
+                    {"content": text, "targetMemberId": text, "wake": {"type": "boolean"}},
+                    ["content"],
+                ),
+            ),
+            (
+                "team_create_task",
+                "仅Leader：创建有验收标准的任务；指定负责人即自动派发，不要再额外wake。依赖用本轮已有taskId。",
+                TaskCreateInput.model_json_schema(),
+            ),
+            (
+                "team_wait",
+                "仅Leader：持久等待一组任务验收/失败，不阻塞；调用后结束当前回复。",
+                schema(
+                    {"taskIds": {"type": "array", "items": text, "minItems": 1, "maxItems": 64}},
+                    ["taskIds"],
+                ),
+            ),
+            (
+                "team_task_action",
+                "按最新 revision 领取任务；Leader 分配/重试，并仅在leader策略下审核任务结果。",
+                schema(
+                    {
+                        "taskId": text,
+                        "action": {"enum": ["assign", "claim", "retry", "accept", "reject"]},
+                        "expectedRevision": {"type": "integer", "minimum": 1},
+                        "ownerMemberId": text,
+                        "reason": {"type": "string", "maxLength": 2000},
+                    },
+                    ["taskId", "action", "expectedRevision"],
+                ),
+            ),
+            (
+                "team_submit_result",
+                "提交自己当前任务的可核验结果。成功Run结束及验收后才算成功。",
+                schema(
+                    {"result": text, "artifactIds": {"type": "array", "items": text}}, ["result"]
+                ),
+            ),
+            (
+                "team_finish",
+                "仅Leader：所有任务通过验收后提交最终成果，等待群主验收。",
+                schema({"result": text}, ["result"]),
+            ),
+            (
+                "team_publish_artifact",
+                "登记当前独立工作区的实际交付文件，返回有来源和哈希的产物引用。",
+                schema({"path": text, "name": text}, ["path"]),
+            ),
+            (
+                "team_read_artifact",
+                "读取本群已显式共享的不可变交付物，复制到自己工作区并返回来源与摘要。",
+                schema({"artifactId": text}, ["artifactId"]),
+            ),
+        ]
+
+    async def invoke(self, principal, operation, arguments, call_id):
+        # principal is an in-process/opaque-handle Host value, never JSON Actor.
+        actor = principal.get("actor") if isinstance(principal, dict) else None
+        if not isinstance(actor, Actor) or actor.kind != "member":
+            raise TeamsError("invocation_required", "缺少宿主授予的成员调用上下文", status=403)
+        from jsonschema import Draft202012Validator
+        from jsonschema.exceptions import ValidationError
+
+        definition = next((item for item in self.tool_definitions() if item[0] == operation), None)
+        if definition is None:
+            raise TeamsError("tool_unavailable", "当前团队不提供此工具", status=404)
+        try:
+            Draft202012Validator(definition[2]).validate(arguments)
+        except ValidationError as error:
+            raise TeamsError(
+                "invalid_tool_arguments", "协作工具参数不符合声明的结构", status=422
+            ) from error
+        if not isinstance(call_id, str) or not call_id or len(call_id) > 512:
+            raise TeamsError("call_identity_required", "缺少稳定工具调用标识", status=403)
+        with self.domain.store.transaction() as tx:
+            self.domain.authorize(tx, actor, actor.group_id)
+        key = f"tool:{actor.run_id}:{call_id}"
+        if operation == "team_context":
+            from .context import execution_context
+
+            return execution_context(self.domain, actor, principal.get("taskId"))
+        if operation == "team_message":
+            return self.domain.member_message(
+                actor,
+                content=arguments["content"],
+                target_member_id=arguments.get("targetMemberId"),
+                wake=arguments.get("wake", False),
+                source_delivery_id=principal["deliveryId"],
+                key=key,
+            )
+        if operation == "team_create_task":
+            return self.domain.create_task(
+                actor,
+                actor.group_id,
+                actor.team_run_id,
+                TaskCreateInput.model_validate(arguments),
+                key,
+            )
+        if operation == "team_wait":
+            return self.domain.wait_for_tasks(actor, arguments["taskIds"], key)
+        if operation == "team_task_action":
+            return self.domain.task_action(
+                actor,
+                actor.group_id,
+                arguments["taskId"],
+                arguments["action"],
+                arguments["expectedRevision"],
+                key,
+                owner_member_id=arguments.get("ownerMemberId"),
+                reason=arguments.get("reason"),
+            )
+        if operation == "team_submit_result":
+            artifacts = self._artifacts(actor, arguments.get("artifactIds", []))
+            return self.domain.result_candidate(
+                actor, principal["taskId"], arguments["result"], key, artifacts=artifacts
+            )
+        if operation == "team_finish":
+            return self.domain.finish_candidate(actor, arguments["result"], key)
+        if operation == "team_publish_artifact":
+            return self.publish_artifact(actor, arguments["path"], arguments.get("name"), key)
+        if operation == "team_read_artifact":
+            return self.read_artifact(actor, arguments["artifactId"])
+        raise TeamsError("tool_unavailable", "当前团队不提供此工具", status=404)

ksadk/plugins/teams/transport.py

@@ -11,7 +11,7 @@ import httpx
 
 from .errors import TeamsError
 
-Headers = Callable[[str, str, str], Awaitable[dict[str, str]]]
+Headers = Callable[[str, str, str | bytes], Awaitable[dict[str, str]]]
 
 
 class TeamsHTTPClient:
@@ -41,8 +41,13 @@ class TeamsHTTPClient:
             trust_env=False,
         )
 
-    async def request_headers(self, method: str, path: str, body: str = ""):
-        result = {"Accept": "application/json", "Content-Type": "application/json"}
+    async def request_headers(self, method: str, path: str, body: str | bytes = ""):
+        result = {
+            "Accept": "application/json",
+            "Content-Type": "application/octet-stream"
+            if isinstance(body, bytes)
+            else "application/json",
+        }
         if self.access_token:
             result["Authorization"] = "Bearer " + self.access_token
         if self.node_token:

ksadk/resource_runtime/skill_runtime.py

@@ -4,6 +4,7 @@ from pathlib import Path
 
 from pydantic import SecretStr
 
+from ksadk.common.cloud_headers import DEFAULT_TRAFFIC_SOURCE
 from ksadk.resource_runtime.snapshots import ConnectionTarget, ResourceSnapshot
 from ksadk.sandbox.e2b_connection import ExplicitE2BConnection
 from ksadk.skills.runtime.backends.e2b import E2BSkillRuntimeBackend
@@ -16,6 +17,7 @@ def create_bound_skill_runtime(
     current_connection: ConnectionTarget,
     api_key: SecretStr,
     artifact_directory: Path | None = None,
+    traffic_source: str = DEFAULT_TRAFFIC_SOURCE,
 ) -> E2BSkillRuntimeBackend:
     """Host revalidates account/grant first; only secret values may rotate here.
 
@@ -32,6 +34,7 @@ def create_bound_skill_runtime(
         api_url=target.connection.endpoint,
         domain=target.domain,
         api_key=api_key,
+        traffic_source=traffic_source,
     )
     return E2BSkillRuntimeBackend(
         template_id=target.template_id,

ksadk/resource_runtime/socket_server.py

@@ -11,6 +11,9 @@ from pathlib import Path
 from ksadk.resource_runtime.broker import ResourceBroker
 from ksadk.resource_runtime.ipc import MAX_FRAME_BYTES, FrameError, decode_body, encode_frame
 
+# macOS/BSD sockaddr_un has 104 bytes, including the terminating NUL.
+_MAX_SOCKET_PATH_BYTES = 103
+
 
 class ResourceSocketServer:
     def __init__(self, broker: ResourceBroker, *, max_connections: int = 64):
@@ -33,6 +36,14 @@ class ResourceSocketServer:
         self._directory = directory
         self.path = Path(directory.name) / "broker.sock"
         try:
+            if len(os.fsencode(self.path)) > _MAX_SOCKET_PATH_BYTES:
+                # Artifact/workspace TMPDIRs can exceed Unix socket limits.
+                # Keep their preference when usable; only this ephemeral IPC
+                # endpoint falls back to an atomic, independently private path.
+                directory.cleanup()
+                directory = tempfile.TemporaryDirectory(prefix="ksr-", dir="/tmp")
+                self._directory = directory
+                self.path = Path(directory.name) / "broker.sock"
             os.chmod(directory.name, 0o700)
             # Private parent directory prevents access during the bind/chmod gap.
             self._server = await asyncio.start_unix_server(self._accept, path=str(self.path))

ksadk/resource_runtime/worker.py

@@ -11,9 +11,10 @@ from typing import ClassVar
 from urllib.parse import urlsplit
 
 import httpx
-from pydantic import ConfigDict, Field, SecretStr, model_validator
+from pydantic import ConfigDict, Field, SecretStr, field_validator, model_validator
 from requests.exceptions import RequestException
 
+from ksadk.common.cloud_headers import DEFAULT_TRAFFIC_SOURCE, resolve_traffic_source
 from ksadk.knowledge_base.client import KnowledgeBaseClient
 from ksadk.memory.adk.backends.sdk_ltm_backend import SdkLTMBackend
 from ksadk.memory.models import MemorySearchRequest
@@ -70,10 +71,13 @@ class ExplicitWorkerBinding(PluginContractModel):
 class KnowledgeWorkerBinding(ExplicitWorkerBinding):
     resource_kind: ClassVar[str] = "knowledge-base"
 
-    def create_service(self) -> BoundKnowledgeService:
+    def create_service(
+        self, *, traffic_source: str = DEFAULT_TRAFFIC_SOURCE
+    ) -> BoundKnowledgeService:
         endpoint = urlsplit(self.endpoint)
         client = KnowledgeBaseClient(
             dataset_id=self.config.binding.resource.id,
+            traffic_source=traffic_source,
             region=self.config.binding.resource.region,
             endpoint=endpoint.netloc,
             scheme=endpoint.scheme,
@@ -87,10 +91,13 @@ class KnowledgeWorkerBinding(ExplicitWorkerBinding):
 class MemoryWorkerBinding(ExplicitWorkerBinding):
     resource_kind: ClassVar[str] = "memory-instance"
 
-    def create_provider(self, identity: InvocationIdentity) -> AicpMemoryProvider:
+    def create_provider(
+        self, identity: InvocationIdentity, *, traffic_source: str = DEFAULT_TRAFFIC_SOURCE
+    ) -> AicpMemoryProvider:
         endpoint = urlsplit(self.endpoint)
         backend = SdkLTMBackend(
             index="bound-resource",
+            traffic_source=traffic_source,
             memory_collection_id=self.config.binding.resource.id,
             namespace=self.config.binding.resource.id,
             agent_id=identity.agent_id,
@@ -135,6 +142,7 @@ class DiscoverySkillWorkerBinding(ExplicitWorkerBinding):
 
     def create_provider(
         self, scope: ResourceScope, *, execution_target=None,
+        traffic_source: str = DEFAULT_TRAFFIC_SOURCE,
     ) -> DiscoverySkillService:
         namespace = discovery_scope_key(scope, self.selection_run_ref)
         client = SkillServiceClient(
@@ -143,6 +151,7 @@ class DiscoverySkillWorkerBinding(ExplicitWorkerBinding):
             secret_key=self.secret_key.get_secret_value(),
             region=self.config.binding.resource.region,
             allow_env_fallback=False,
+            traffic_source=traffic_source,
             timeout=15,
         )
         return DiscoverySkillService(
@@ -221,6 +230,9 @@ class PinnedSkillWorkerBinding(PluginContractModel):
 
 class WorkerInitialization(PluginContractModel):
     model_config = ConfigDict(hide_input_in_errors=True)
+    # Resolve in the trusted parent and serialize on stdin; never inherit the
+    # parent's environment or mutate global process state in a shared host.
+    traffic_source: str = Field(default_factory=resolve_traffic_source)
     resource_snapshot: ResourceSnapshot
     bindings: tuple[
         KnowledgeWorkerBinding
@@ -231,6 +243,11 @@ class WorkerInitialization(PluginContractModel):
     ]
     scopes: tuple[ResourceScope, ...]
 
+    @field_validator("traffic_source")
+    @classmethod
+    def validate_traffic_source(cls, value):
+        return resolve_traffic_source(value)
+
     @model_validator(mode="after")
     def single_activation(self) -> WorkerInitialization:
         if not self.bindings or not self.scopes:
@@ -337,13 +354,19 @@ def main() -> int:
             binding.config.binding.id: (
                 binding.create_provider(
                     scopes[binding.config.binding.id],
+                    traffic_source=initialization.traffic_source,
                     execution_target=initialization.resource_snapshot.binding(
                         binding.config.binding.id
                     ).skill_execution,
                 )
                 if isinstance(binding, DiscoverySkillWorkerBinding)
-                else binding.create_provider(scopes[binding.config.binding.id].identity)
+                else binding.create_provider(
+                    scopes[binding.config.binding.id].identity,
+                    traffic_source=initialization.traffic_source,
+                )
                 if isinstance(binding, MemoryWorkerBinding)
+                else binding.create_service(traffic_source=initialization.traffic_source)
+                if isinstance(binding, KnowledgeWorkerBinding)
                 else binding.create_service()
             )
             for binding in initialization.bindings
@@ -362,6 +385,7 @@ def main() -> int:
                     binding.config.binding.id,
                     current_connection=binding.execution_connection,
                     api_key=binding.execution_api_key,
+                    traffic_source=initialization.traffic_source,
                     artifact_directory=root.resolve(),
                 )
                 skill_runtimes[binding.config.binding.id].preflight()

ksadk/runners/_adk_a2a_stream.py

@@ -0,0 +1,122 @@
+"""Preserve wire identity before ADK's converters discard stream semantics."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from functools import partial, wraps
+from typing import Any
+
+A2A_STREAM_METADATA = "ksadk_a2a_stream"
+
+
+def _wrap_converter(converter: Any, *, artifact_update: bool) -> Any:
+    if getattr(converter, "_ksadk_stream_identity", False):
+        return converter
+
+    @wraps(converter)
+    def convert(wire: Any, *args: Any, **kwargs: Any) -> Any:
+        event = converter(wire, *args, **kwargs)
+        if event is None:
+            return None
+        from a2a.types import TaskState
+
+        artifacts = [wire.artifact] if artifact_update else list(wire.artifacts)
+        observations = []
+        for artifact in artifacts:
+            converted = event
+            if not artifact_update and len(artifacts) > 1:
+                single = deepcopy(wire)
+                del single.artifacts[:]
+                single.artifacts.extend([artifact])
+                converted = converter(single, *args, **kwargs)
+            parts = getattr(getattr(converted, "content", None), "parts", None) or ()
+            terminal = not artifact_update and wire.status.state == TaskState.TASK_STATE_COMPLETED
+            observations.append(
+                {
+                    "task_id": wire.task_id if artifact_update else wire.id,
+                    "artifact_id": artifact.artifact_id,
+                    "append": bool(wire.append) if artifact_update else False,
+                    "last_chunk": bool(wire.last_chunk) if artifact_update else terminal,
+                    "snapshot": terminal,
+                    "text": {
+                        lane: "".join(
+                            part.text
+                            for part in parts
+                            if part.text and bool(part.thought) == (lane == "reasoning")
+                        )
+                        for lane in ("message", "reasoning")
+                    },
+                }
+            )
+        if observations:
+            event.custom_metadata = dict(event.custom_metadata or {})
+            event.custom_metadata[A2A_STREAM_METADATA] = observations
+        return event
+
+    convert._ksadk_stream_identity = True
+    return convert
+
+
+async def _reconcile_terminal_status(
+    ctx: Any, response: Any, event: Any, *, task_converter: Any, part_converter: Any
+) -> Any:
+    """A completed status message summarizes the Task, not another LLM output.
+
+    Use the actual accumulated Task artifacts supplied by ADK's A2A client.
+    No association is inferred from the summary's author or text.
+    """
+    from a2a.types import TaskState, TaskStatusUpdateEvent
+
+    if isinstance(response, tuple):
+        task, update = response
+        if (
+            isinstance(update, TaskStatusUpdateEvent)
+            and update.status.state == TaskState.TASK_STATE_COMPLETED
+            and task.artifacts
+        ):
+            converted = task_converter(task, event.author, ctx, part_converter)
+            event.custom_metadata = dict(event.custom_metadata or {})
+            event.custom_metadata[A2A_STREAM_METADATA] = converted.custom_metadata[
+                A2A_STREAM_METADATA
+            ]
+    return event
+
+
+def preserve_a2a_stream_identity(agent: Any) -> None:
+    """Decorate configured converters once, including nested native agents."""
+    seen: set[int] = set()
+
+    def visit(node: Any) -> None:
+        if id(node) in seen:
+            return
+        seen.add(id(node))
+        config = getattr(node, "_config", None)
+        if callable(getattr(config, "a2a_artifact_update_converter", None)) and callable(
+            getattr(config, "a2a_task_converter", None)
+        ):
+            config.a2a_artifact_update_converter = _wrap_converter(
+                config.a2a_artifact_update_converter, artifact_update=True
+            )
+            config.a2a_task_converter = _wrap_converter(
+                config.a2a_task_converter, artifact_update=False
+            )
+            from google.adk.a2a.agent.config import RequestInterceptor
+
+            if not any(
+                getattr(hook.after_request, "_ksadk_terminal_reconcile", False)
+                for hook in (config.request_interceptors or ())
+            ):
+                reconcile = partial(
+                    _reconcile_terminal_status,
+                    task_converter=config.a2a_task_converter,
+                    part_converter=config.a2a_part_converter,
+                )
+                reconcile._ksadk_terminal_reconcile = True
+                config.request_interceptors = [
+                    *(config.request_interceptors or ()),
+                    RequestInterceptor(after_request=reconcile),
+                ]
+        for child in getattr(node, "sub_agents", ()):
+            visit(child)
+
+    visit(agent)