Merge branch 'main' into jb/python-mcp-parity

Author: jackbatzner

packages/agent-hypervisor/src/hypervisor/security/kill_switch.py

@@ -1,19 +1,23 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-# Public Preview — basic implementation
 """
-Kill Switch — simple hard kill.
+Kill Switch — agent termination with optional handoff.
 
-Public Preview: immediate agent termination, no task transfer.
+Terminates agent processes via registered callbacks and hands off
+in-flight saga steps to a substitute agent when one is available.
 """
 
 from __future__ import annotations
 
+import logging
 import uuid
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from datetime import UTC, datetime
 from enum import Enum
 
+_logger = logging.getLogger(__name__)
+
 
 class KillReason(str, Enum):
     """Why an agent was killed."""
@@ -37,7 +41,7 @@ class HandoffStatus(str, Enum):
 
 @dataclass
 class StepHandoff:
-    """A saga step being handed off (Public Preview: always COMPENSATED)."""
+    """A saga step being handed off to a substitute or compensated."""
 
     step_id: str
     saga_id: str
@@ -58,22 +62,43 @@ class KillResult:
     handoffs: list[StepHandoff] = field(default_factory=list)
     handoff_success_count: int = 0
     compensation_triggered: bool = False
+    terminated: bool = False
     details: str = ""
 
 
 class KillSwitch:
     """
-    Simple hard kill (Public Preview: no handoff, immediate termination).
+    Kill switch with agent process registry and handoff support.
+
+    Agents register termination callbacks via ``register_agent``.  When
+    ``kill`` is called the switch hands in-flight saga steps to a
+    registered substitute (if any) and then invokes the termination
+    callback to stop the agent process.
     """
 
     def __init__(self) -> None:
         self._kill_history: list[KillResult] = []
         self._substitutes: dict[str, list[str]] = {}
+        self._agents: dict[str, Callable[[], None]] = {}
+
+    # ── Agent process registry ─────────────────────────────────────
+
+    def register_agent(
+        self, agent_did: str, process_handle: Callable[[], None]
+    ) -> None:
+        """Register an agent with its termination callback."""
+        self._agents[agent_did] = process_handle
+
+    def unregister_agent(self, agent_did: str) -> None:
+        """Remove an agent from the process registry."""
+        self._agents.pop(agent_did, None)
+
+    # ── Substitute management ──────────────────────────────────────
 
     def register_substitute(
         self, session_id: str, agent_did: str
     ) -> None:
-        """Register a substitute (Public Preview: no-op, handoff not supported)."""
+        """Register a substitute agent for a session."""
         self._substitutes.setdefault(session_id, []).append(agent_did)
 
     def unregister_substitute(
@@ -83,6 +108,8 @@ def unregister_substitute(
         if agent_did in subs:
             subs.remove(agent_did)
 
+    # ── Kill ───────────────────────────────────────────────────────
+
     def kill(
         self,
         agent_did: str,
@@ -91,35 +118,73 @@ def kill(
         in_flight_steps: list[dict] | None = None,
         details: str = "",
     ) -> KillResult:
-        """Kill an agent immediately (Public Preview: no handoff)."""
+        """Kill an agent, handing off in-flight steps to a substitute if available."""
         in_flight = in_flight_steps or []
 
-        handoffs = [
-            StepHandoff(
-                step_id=step_info.get("step_id", ""),
-                saga_id=step_info.get("saga_id", ""),
-                from_agent=agent_did,
-                status=HandoffStatus.COMPENSATED,
+        # Attempt to find a substitute for handoff
+        substitute = self._find_substitute(session_id, agent_did)
+
+        handoffs: list[StepHandoff] = []
+        handoff_success_count = 0
+        for step_info in in_flight:
+            if substitute is not None:
+                handoffs.append(
+                    StepHandoff(
+                        step_id=step_info.get("step_id", ""),
+                        saga_id=step_info.get("saga_id", ""),
+                        from_agent=agent_did,
+                        to_agent=substitute,
+                        status=HandoffStatus.HANDED_OFF,
+                    )
+                )
+                handoff_success_count += 1
+            else:
+                handoffs.append(
+                    StepHandoff(
+                        step_id=step_info.get("step_id", ""),
+                        saga_id=step_info.get("saga_id", ""),
+                        from_agent=agent_did,
+                        status=HandoffStatus.COMPENSATED,
+                    )
+                )
+
+        # Terminate the agent process
+        terminated = False
+        callback = self._agents.get(agent_did)
+        if callback is not None:
+            callback()
+            terminated = True
+        else:
+            _logger.warning(
+                "No termination callback registered for agent %s",
+                agent_did,
             )
-            for step_info in in_flight
-        ]
 
         result = KillResult(
             agent_did=agent_did,
             session_id=session_id,
             reason=reason,
             handoffs=handoffs,
-            handoff_success_count=0,
-            compensation_triggered=len(handoffs) > 0,
+            handoff_success_count=handoff_success_count,
+            compensation_triggered=any(
+                h.status == HandoffStatus.COMPENSATED for h in handoffs
+            ),
+            terminated=terminated,
             details=details,
         )
         self._kill_history.append(result)
         self.unregister_substitute(session_id, agent_did)
+        self.unregister_agent(agent_did)
         return result
 
     def _find_substitute(
         self, session_id: str, exclude_did: str
     ) -> str | None:
+        """Find a registered substitute for the session, excluding the given agent."""
+        subs = self._substitutes.get(session_id, [])
+        for sub in subs:
+            if sub != exclude_did:
+                return sub
         return None
 
     @property
@@ -132,4 +197,4 @@ def total_kills(self) -> int:
 
     @property
     def total_handoffs(self) -> int:
-        return 0
+        return sum(r.handoff_success_count for r in self._kill_history)

packages/agent-hypervisor/tests/test_kill_switch.py

@@ -144,14 +144,93 @@ def test_kill_unregisters_agent_substitute(self):
         ks.kill("agent-1", "sess-1", KillReason.MANUAL)
         assert "agent-1" not in ks._substitutes.get("sess-1", [])
 
-    def test_find_substitute_always_none(self):
-        """Public preview: no substitute finding."""
+    def test_find_substitute_returns_registered(self):
+        """Substitute lookup returns a registered agent."""
         ks = KillSwitch()
         ks.register_substitute("s1", "backup-agent")
+        assert ks._find_substitute("s1", "agent-1") == "backup-agent"
+
+    def test_find_substitute_excludes_killed_agent(self):
+        ks = KillSwitch()
+        ks.register_substitute("s1", "agent-1")
+        assert ks._find_substitute("s1", "agent-1") is None
+
+    def test_find_substitute_no_session(self):
+        ks = KillSwitch()
         assert ks._find_substitute("s1", "agent-1") is None
 
     def test_session_timeout_reason(self):
         ks = KillSwitch()
         result = ks.kill("a1", "s1", KillReason.SESSION_TIMEOUT)
         assert result.reason == KillReason.SESSION_TIMEOUT
         assert ks.total_kills == 1
+
+    # ── Agent process registry ─────────────────────────────────────
+
+    def test_register_and_unregister_agent(self):
+        ks = KillSwitch()
+        ks.register_agent("agent-1", lambda: None)
+        assert "agent-1" in ks._agents
+        ks.unregister_agent("agent-1")
+        assert "agent-1" not in ks._agents
+
+    def test_unregister_nonexistent_agent(self):
+        ks = KillSwitch()
+        ks.unregister_agent("nonexistent")  # Should not raise
+
+    def test_kill_unregisters_agent(self):
+        ks = KillSwitch()
+        ks.register_agent("agent-1", lambda: None)
+        ks.kill("agent-1", "sess-1", KillReason.MANUAL)
+        assert "agent-1" not in ks._agents
+
+    # ── Termination callback ───────────────────────────────────────
+
+    def test_kill_with_registered_callback_terminates(self):
+        """Registered termination callback is called and terminated=True."""
+        ks = KillSwitch()
+        terminated_agents: list[str] = []
+        ks.register_agent("agent-1", lambda: terminated_agents.append("agent-1"))
+        result = ks.kill("agent-1", "sess-1", KillReason.MANUAL)
+        assert result.terminated is True
+        assert "agent-1" in terminated_agents
+
+    def test_kill_without_registered_callback(self):
+        """Kill without registered callback sets terminated=False."""
+        ks = KillSwitch()
+        result = ks.kill("agent-1", "sess-1", KillReason.MANUAL)
+        assert result.terminated is False
+
+    # ── Handoff vs compensation ────────────────────────────────────
+
+    def test_kill_with_substitute_hands_off_steps(self):
+        """Steps are handed off to substitute when available."""
+        ks = KillSwitch()
+        ks.register_substitute("sess-1", "backup-agent")
+        steps = [
+            {"step_id": "s1", "saga_id": "saga-1"},
+            {"step_id": "s2", "saga_id": "saga-1"},
+        ]
+        result = ks.kill(
+            "agent-1", "sess-1", KillReason.RING_BREACH, in_flight_steps=steps
+        )
+        assert result.handoff_success_count == 2
+        assert all(h.status == HandoffStatus.HANDED_OFF for h in result.handoffs)
+        assert all(h.to_agent == "backup-agent" for h in result.handoffs)
+        assert result.compensation_triggered is False
+
+    def test_kill_without_substitute_compensates_steps(self):
+        """Steps are compensated when no substitute is available."""
+        ks = KillSwitch()
+        steps = [{"step_id": "s1", "saga_id": "saga-1"}]
+        result = ks.kill("agent-1", "sess-1", KillReason.MANUAL, in_flight_steps=steps)
+        assert result.handoff_success_count == 0
+        assert all(h.status == HandoffStatus.COMPENSATED for h in result.handoffs)
+        assert result.compensation_triggered is True
+
+    def test_total_handoffs_counts_successes(self):
+        """total_handoffs sums handoff_success_count across kills."""
+        ks = KillSwitch()
+        ks.register_substitute("s1", "backup")
+        ks.kill("a1", "s1", KillReason.MANUAL, [{"step_id": "s1", "saga_id": "sg1"}])
+        assert ks.total_handoffs == 1

packages/agent-hypervisor/tests/unit/test_session_security.py

@@ -276,10 +276,10 @@ def test_kill_with_handoff(self):
                 {"step_id": "step-1", "saga_id": "saga-1"},
             ],
         )
-        # Public Preview: no handoff, always compensated
-        assert result.handoff_success_count == 0
-        assert result.handoffs[0].status == HandoffStatus.COMPENSATED
-        assert result.compensation_triggered
+        assert result.handoff_success_count == 1
+        assert result.handoffs[0].status == HandoffStatus.HANDED_OFF
+        assert result.handoffs[0].to_agent == "backup-agent"
+        assert not result.compensation_triggered
 
     def test_kill_without_substitute(self):
         ks = KillSwitch()
@@ -324,8 +324,7 @@ def test_total_handoffs(self):
         ks = KillSwitch()
         ks.register_substitute("s1", "backup")
         ks.kill("a1", "s1", KillReason.MANUAL, [{"step_id": "s1", "saga_id": "sg1"}])
-        # Public Preview: no handoffs
-        assert ks.total_handoffs == 0
+        assert ks.total_handoffs == 1
 
     def test_unregister_substitute(self):
         ks = KillSwitch()