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

ksadk/plugins/providers/harness.py

@@ -6,10 +6,13 @@ delegates every turn to ``HarnessRuntimeAdapter`` through
 history therefore stay on the existing conversation pipeline; this module does
 not create another event stream or transcript store.
 """
+
 from __future__ import annotations
 
+import asyncio
 import json
-from collections.abc import Mapping, Sequence
+from collections.abc import Awaitable, Mapping, Sequence
+from contextlib import AsyncExitStack
 from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import Any, Protocol, cast, runtime_checkable
@@ -20,6 +23,7 @@ from ksadk.harness.runtime import HarnessRuntimeAdapter
 from ksadk.plugins.bundle import ResolvedPluginBundle
 from ksadk.plugins.contracts import CompositionProfile, PluginManifest
 from ksadk.plugins.host import PluginExecutionContext, PluginHostError
+from ksadk.plugins.providers.mcp_projection import project_mcp_capabilities
 from ksadk.runtime import RuntimeExecutor, RuntimeLaunchContext, RuntimeRegistry, StartRequest
 from ksadk.runtime.conversation_execution import invoke_runtime_conversation_once
 from ksadk.sessions import create_session_service
@@ -59,23 +63,15 @@ class HarnessTurnRequest:
         if isinstance(value, cls):
             return value
         if not isinstance(value, Mapping):
-            raise PluginHostError(
-                "harness_input_invalid", "Harness input must be an object"
-            )
+            raise PluginHostError("harness_input_invalid", "Harness input must be an object")
         user_id = str(value.get("user_id") or value.get("userId") or "").strip()
         if not user_id:
-            raise PluginHostError(
-                "harness_input_invalid", "Harness input requires user_id"
-            )
+            raise PluginHostError("harness_input_invalid", "Harness input requires user_id")
         raw_messages = value.get("messages")
         if raw_messages is None and value.get("input") is not None:
             raw_messages = [{"role": "user", "content": value.get("input")}]
-        if not isinstance(raw_messages, Sequence) or isinstance(
-            raw_messages, (str, bytes)
-        ):
-            raise PluginHostError(
-                "harness_input_invalid", "Harness input requires messages"
-            )
+        if not isinstance(raw_messages, Sequence) or isinstance(raw_messages, (str, bytes)):
+            raise PluginHostError("harness_input_invalid", "Harness input requires messages")
         messages: list[Mapping[str, Any]] = []
         for index, message in enumerate(raw_messages):
             if not isinstance(message, Mapping):
@@ -94,24 +90,18 @@ class HarnessTurnRequest:
             raise PluginHostError(
                 "harness_input_invalid", "Harness turn must end with a user message"
             )
-        session_id = str(
-            value.get("session_id") or value.get("sessionId") or ""
-        ).strip()
+        session_id = str(value.get("session_id") or value.get("sessionId") or "").strip()
         model = str(value.get("model") or "").strip()
         metadata = value.get("request_metadata") or value.get("requestMetadata")
         if metadata is not None and not isinstance(metadata, Mapping):
-            raise PluginHostError(
-                "harness_input_invalid", "request_metadata must be an object"
-            )
+            raise PluginHostError("harness_input_invalid", "request_metadata must be an object")
         return cls(
             user_id=user_id,
             session_id=session_id or None,
             messages=tuple(messages),
             model=model or None,
             request_metadata=dict(metadata) if metadata is not None else None,
-            invocation_id=str(
-                value.get("invocation_id") or value.get("invocationId") or ""
-            ).strip()
+            invocation_id=str(value.get("invocation_id") or value.get("invocationId") or "").strip()
             or None,
         )
 
@@ -129,14 +119,17 @@ class HarnessTurnResult:
 class HarnessMCPSource(Protocol):
     def harness_mcp_specs(
         self, bundle: ResolvedPluginBundle
-    ) -> Sequence[McpToolSpec]: ...
+    ) -> Sequence[McpToolSpec] | Awaitable[Sequence[McpToolSpec]]: ...
+
+
+@runtime_checkable
+class HarnessMCPActivationSource(Protocol):
+    async def release_harness_mcp_specs(self, specs: Sequence[McpToolSpec]) -> None: ...
 
 
 @runtime_checkable
 class HarnessSkillSource(Protocol):
-    def harness_skill(
-        self, bundle: ResolvedPluginBundle
-    ) -> HarnessSkillContribution: ...
+    def harness_skill(self, bundle: ResolvedPluginBundle) -> HarnessSkillContribution: ...
 
 
 @runtime_checkable
@@ -269,97 +262,91 @@ class KsADKHarnessProviderRuntime:
         capabilities: PluginExecutionContext,
     ) -> "KsADKHarnessActivation":
         if not self._ready or self._disposed:
-            raise PluginHostError(
-                "harness_provider_unavailable", "Harness provider is not ready"
+            raise PluginHostError("harness_provider_unavailable", "Harness provider is not ready")
+
+        async with AsyncExitStack() as mcp_cleanup:
+            projection = await project_mcp_capabilities(capabilities, bundle)
+            mcp_cleanup.push_async_callback(projection.aclose)
+            mcp_specs = list(projection.specs)
+            mcp_owners = list(projection.owners)
+
+            skills: list[HarnessSkillContribution] = []
+            skill_owners: list[str] = []
+            for binding in capabilities.all("skill.source/v1"):
+                if not isinstance(binding.runtime, HarnessSkillSource):
+                    raise PluginHostError(
+                        "harness_skill_incompatible",
+                        f"plugin {binding.plugin_id} cannot project Harness instructions",
+                    )
+                contribution = binding.runtime.harness_skill(bundle)
+                if not contribution.name.strip() or not contribution.instructions.strip():
+                    raise PluginHostError(
+                        "harness_skill_invalid",
+                        f"plugin {binding.plugin_id} returned an empty Skill contribution",
+                    )
+                skills.append(contribution)
+                skill_owners.append(binding.plugin_id)
+
+            context_sources: list[HarnessContextSource] = []
+            context_owners: list[str] = []
+            for binding in capabilities.all("context.contributor/v1"):
+                if not isinstance(binding.runtime, HarnessContextSource):
+                    raise PluginHostError(
+                        "harness_context_incompatible",
+                        f"plugin {binding.plugin_id} cannot contribute Harness context",
+                    )
+                context_sources.append(binding.runtime)
+                context_owners.append(binding.plugin_id)
+
+            execution = bundle.resolved_agent_spec.get("execution")
+            strategy = (
+                str(execution.get("strategy") or "direct").strip()
+                if isinstance(execution, Mapping)
+                else "direct"
             )
-
-        mcp_specs: list[McpToolSpec] = []
-        mcp_owners: list[str] = []
-        for binding in capabilities.all("mcp.connector/v1"):
-            if not isinstance(binding.runtime, HarnessMCPSource):
-                raise PluginHostError(
-                    "harness_mcp_incompatible",
-                    f"plugin {binding.plugin_id} cannot project Harness MCP config",
-                )
-            mcp_specs.extend(binding.runtime.harness_mcp_specs(bundle))
-            mcp_owners.append(binding.plugin_id)
-
-        skills: list[HarnessSkillContribution] = []
-        skill_owners: list[str] = []
-        for binding in capabilities.all("skill.source/v1"):
-            if not isinstance(binding.runtime, HarnessSkillSource):
-                raise PluginHostError(
-                    "harness_skill_incompatible",
-                    f"plugin {binding.plugin_id} cannot project Harness instructions",
-                )
-            contribution = binding.runtime.harness_skill(bundle)
-            if not contribution.name.strip() or not contribution.instructions.strip():
+            if strategy != "direct":
                 raise PluginHostError(
-                    "harness_skill_invalid",
-                    f"plugin {binding.plugin_id} returned an empty Skill contribution",
+                    "harness_execution_strategy_unsupported",
+                    f"KsADK Harness Provider does not support {strategy!r}",
                 )
-            skills.append(contribution)
-            skill_owners.append(binding.plugin_id)
 
-        context_sources: list[HarnessContextSource] = []
-        context_owners: list[str] = []
-        for binding in capabilities.all("context.contributor/v1"):
-            if not isinstance(binding.runtime, HarnessContextSource):
-                raise PluginHostError(
-                    "harness_context_incompatible",
-                    f"plugin {binding.plugin_id} cannot contribute Harness context",
+            model, prompt = _bundle_model_and_prompt(bundle)
+            if skills:
+                prompt = _append_prompt_sections(
+                    prompt,
+                    [f"Skill {item.name}:\n{item.instructions}" for item in skills],
                 )
-            context_sources.append(binding.runtime)
-            context_owners.append(binding.plugin_id)
-
-        execution = bundle.resolved_agent_spec.get("execution")
-        strategy = (
-            str(execution.get("strategy") or "direct").strip()
-            if isinstance(execution, Mapping)
-            else "direct"
-        )
-        if strategy != "direct":
-            raise PluginHostError(
-                "harness_execution_strategy_unsupported",
-                f"KsADK Harness Provider does not support {strategy!r}",
+            config = HarnessConfig(
+                model=model,
+                prompt=prompt,
+                mcp_tools=tuple(mcp_specs),
+                sandbox=SandboxPolicy(read_only=True),
+                runtime="yaml",
             )
-
-        model, prompt = _bundle_model_and_prompt(bundle)
-        if skills:
-            prompt = _append_prompt_sections(
-                prompt,
-                [f"Skill {item.name}:\n{item.instructions}" for item in skills],
+            inventory = HarnessProviderInventory(
+                provider=self._plugin_id,
+                execution_strategy=strategy,
+                model=model,
+                history_owner="canonical_session_service",
+                mcp_servers=tuple(mcp_owners),
+                skills=tuple(item.name for item in skills),
+                context_contributors=tuple(context_owners),
+            )
+            self._last_inventory = inventory
+            workspace_root = bundle.root / "runtime"
+            if not workspace_root.is_dir():
+                workspace_root = bundle.root
+            return KsADKHarnessActivation(
+                bundle=bundle,
+                config=config,
+                agent_name=bundle.manifest.agent_id,
+                workspace_root=workspace_root,
+                reasoner=self._reasoner,
+                context_sources=tuple(context_sources),
+                session_service=self._session_service,
+                inventory=inventory,
+                mcp_cleanup=mcp_cleanup.pop_all(),
             )
-        config = HarnessConfig(
-            model=model,
-            prompt=prompt,
-            mcp_tools=tuple(mcp_specs),
-            sandbox=SandboxPolicy(read_only=True),
-            runtime="yaml",
-        )
-        inventory = HarnessProviderInventory(
-            provider=self._plugin_id,
-            execution_strategy=strategy,
-            model=model,
-            history_owner="canonical_session_service",
-            mcp_servers=tuple(mcp_owners),
-            skills=tuple(item.name for item in skills),
-            context_contributors=tuple(context_owners),
-        )
-        self._last_inventory = inventory
-        workspace_root = bundle.root / "runtime"
-        if not workspace_root.is_dir():
-            workspace_root = bundle.root
-        return KsADKHarnessActivation(
-            bundle=bundle,
-            config=config,
-            agent_name=bundle.manifest.agent_id,
-            workspace_root=workspace_root,
-            reasoner=self._reasoner,
-            context_sources=tuple(context_sources),
-            session_service=self._session_service,
-            inventory=inventory,
-        )
 
 
 class KsADKHarnessProviderFactory:
@@ -412,6 +399,7 @@ class KsADKHarnessActivation:
         context_sources: tuple[HarnessContextSource, ...],
         session_service: BaseSessionService,
         inventory: HarnessProviderInventory,
+        mcp_cleanup: AsyncExitStack | None = None,
     ) -> None:
         self._bundle = bundle
         self._config = config
@@ -421,6 +409,8 @@ class KsADKHarnessActivation:
         self._context_sources = context_sources
         self._session_service = session_service
         self._inventory = inventory
+        self._mcp_cleanup = mcp_cleanup
+        self._mcp_cleanup_lock = asyncio.Lock()
         self._ready = False
         self._disposed = False
         self._executors: list[RuntimeExecutor] = []
@@ -504,9 +494,18 @@ class KsADKHarnessActivation:
     async def drain(self) -> None:
         self._ready = False
 
+    async def abort(self) -> None:
+        """Revoke external MCP credentials before waiting on a stuck turn."""
+
+        await self._release_mcp_credentials()
+
     async def dispose(self) -> None:
         self._ready = False
         first_error: BaseException | None = None
+        try:
+            await self._release_mcp_credentials()
+        except BaseException as error:  # cleanup must continue
+            first_error = error
         if self._kernel_adapter is not None:
             try:
                 await self._kernel_adapter.close_all()
@@ -524,6 +523,13 @@ class KsADKHarnessActivation:
         if first_error is not None:
             raise first_error
 
+    async def _release_mcp_credentials(self) -> None:
+        async with self._mcp_cleanup_lock:
+            cleanup = self._mcp_cleanup
+            if cleanup is not None:
+                await cleanup.aclose()
+                self._mcp_cleanup = None
+
 
 def _build_direct_backend(
     config: HarnessConfig,
@@ -597,9 +603,7 @@ def _chat_history(history: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]:
         if role not in {"user", "assistant", "tool"}:
             continue
         content = item.get("content")
-        if isinstance(content, Sequence) and not isinstance(
-            content, (str, bytes, bytearray)
-        ):
+        if isinstance(content, Sequence) and not isinstance(content, (str, bytes, bytearray)):
             segments = [
                 str(part.get("text") or "")
                 for part in content
@@ -617,6 +621,7 @@ def _chat_history(history: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]:
 __all__ = [
     "HarnessContextSource",
     "HarnessMCPSource",
+    "HarnessMCPActivationSource",
     "HarnessProviderInventory",
     "HarnessSkillContribution",
     "HarnessSkillSource",

ksadk/plugins/providers/mcp_projection.py

@@ -0,0 +1,150 @@
+"""Provider-neutral projection of PluginHost MCP connector capabilities."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Sequence
+from contextlib import AsyncExitStack
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+from ksadk.harness.config import McpToolSpec
+from ksadk.plugins.bundle import ResolvedPluginBundle
+from ksadk.plugins.host import PluginExecutionContext, PluginHostError
+
+
+@runtime_checkable
+class ActivationMCPSource(Protocol):
+    """MCP source whose lease is bound to one PluginHost activation."""
+
+    def activation_mcp_specs(
+        self,
+        bundle: ResolvedPluginBundle,
+        *,
+        activation_key: str,
+    ) -> Sequence[McpToolSpec] | Awaitable[Sequence[McpToolSpec]]: ...
+
+    async def release_activation_mcp_specs(
+        self,
+        activation_key: str,
+        specs: Sequence[McpToolSpec],
+    ) -> None: ...
+
+
+@runtime_checkable
+class StaticMCPSource(Protocol):
+    """Compatibility seam implemented by existing MCP capability plugins."""
+
+    def harness_mcp_specs(
+        self,
+        bundle: ResolvedPluginBundle,
+    ) -> Sequence[McpToolSpec] | Awaitable[Sequence[McpToolSpec]]: ...
+
+
+@runtime_checkable
+class ReleasableStaticMCPSource(Protocol):
+    async def release_harness_mcp_specs(self, specs: Sequence[McpToolSpec]) -> None: ...
+
+
+@dataclass
+class MCPProjectionLease:
+    """Projected connector specs and their activation-owned cleanup."""
+
+    specs: tuple[McpToolSpec, ...]
+    owners: tuple[str, ...]
+    _cleanup: AsyncExitStack
+    _close_task: asyncio.Task[None] | None = None
+
+    async def aclose(self) -> None:
+        if self._close_task is None:
+            self._close_task = asyncio.create_task(self._cleanup.aclose())
+        await _finish_owned_task(self._close_task)
+
+
+async def _finish_owned_task(task: asyncio.Task[None]) -> None:
+    """Finish a security cleanup before propagating caller cancellation."""
+
+    interrupted = False
+    while not task.done():
+        try:
+            await asyncio.shield(task)
+        except asyncio.CancelledError:
+            interrupted = True
+    task.result()
+    if interrupted:
+        raise asyncio.CancelledError
+
+
+async def project_mcp_capabilities(
+    capabilities: PluginExecutionContext,
+    bundle: ResolvedPluginBundle,
+    *,
+    activation_only: bool = False,
+) -> MCPProjectionLease:
+    """Activate every MCP capability and retain cleanup for the provider."""
+
+    cleanup = AsyncExitStack()
+    specs: list[McpToolSpec] = []
+    owners: list[str] = []
+    names: set[str] = set()
+    try:
+        for binding in capabilities.all("mcp.connector/v1"):
+            runtime = binding.runtime
+            if isinstance(runtime, ActivationMCPSource):
+                projected = runtime.activation_mcp_specs(
+                    bundle,
+                    activation_key=capabilities.activation_key,
+                )
+                if isinstance(projected, Awaitable):
+                    projected = await projected
+                projected_specs = tuple(projected)
+                cleanup.push_async_callback(
+                    runtime.release_activation_mcp_specs,
+                    capabilities.activation_key,
+                    projected_specs,
+                )
+            elif isinstance(runtime, StaticMCPSource):
+                if activation_only:
+                    continue
+                projected = runtime.harness_mcp_specs(bundle)
+                if isinstance(projected, Awaitable):
+                    projected = await projected
+                projected_specs = tuple(projected)
+                if isinstance(runtime, ReleasableStaticMCPSource):
+                    cleanup.push_async_callback(
+                        runtime.release_harness_mcp_specs,
+                        projected_specs,
+                    )
+            else:
+                raise PluginHostError(
+                    "mcp_capability_incompatible",
+                    f"plugin {binding.plugin_id} cannot project MCP connector config",
+                )
+            for spec in projected_specs:
+                if not isinstance(spec, McpToolSpec) or not spec.name.strip():
+                    raise PluginHostError(
+                        "mcp_capability_invalid",
+                        f"plugin {binding.plugin_id} returned an invalid MCP connector",
+                    )
+                if spec.name in names:
+                    raise PluginHostError(
+                        "mcp_capability_ambiguous",
+                        f"more than one MCP connector is named {spec.name!r}",
+                    )
+                names.add(spec.name)
+                specs.append(spec)
+            owners.append(binding.plugin_id)
+    except BaseException:
+        close = asyncio.create_task(cleanup.aclose())
+        await _finish_owned_task(close)
+        raise
+    return MCPProjectionLease(tuple(specs), tuple(owners), cleanup)
+
+
+__all__ = [
+    "ActivationMCPSource",
+    "MCPProjectionLease",
+    "ReleasableStaticMCPSource",
+    "StaticMCPSource",
+    "project_mcp_capabilities",
+]