remove adk LivenessProbe and update mcpSessionManager

Author: DmarshalTU

go/internal/controller/translator/agent/adk_api_translator.go

@@ -524,15 +524,6 @@ func (a *adkApiTranslator) buildManifest(
 							TimeoutSeconds:      15,
 							PeriodSeconds:       15,
 						},
-						LivenessProbe: &corev1.Probe{
-							ProbeHandler: corev1.ProbeHandler{
-								HTTPGet: &corev1.HTTPGetAction{Path: "/healthz/mcp", Port: intstr.FromString("http")},
-							},
-							InitialDelaySeconds: 30,
-							PeriodSeconds:       30,
-							TimeoutSeconds:      10,
-							FailureThreshold:    3,
-						},
 						SecurityContext: securityContext,
 						VolumeMounts:    volumeMounts,
 					}},

python/packages/kagent-adk/src/kagent/adk/_mcp_toolset.py

@@ -13,6 +13,7 @@
 logger = logging.getLogger("kagent_adk." + __name__)
 
 _PING_TIMEOUT_SECONDS = 2.0
+_SESSION_REVALIDATE_TIMEOUT_SECONDS = 5.0
 _JSONRPC_METHOD_NOT_FOUND = -32601
 
 
@@ -28,8 +29,17 @@ def _is_server_alive_error(exc: Exception) -> bool:
     return False
 
 
+def _is_session_invalid_error(exc: Exception) -> bool:
+    """Return True if the error indicates the MCP session is no longer valid (e.g. 404)."""
+    parts = [str(exc)]
+    if isinstance(exc, McpError) and exc.error.message:
+        parts.append(exc.error.message)
+    msg = " ".join(parts).lower()
+    return "404" in msg or "session terminated" in msg
+
+
 class KAgentMCPSessionManager(MCPSessionManager):
-    """Session manager that validates cached sessions via ping before reuse.
+    """Session manager that validates cached sessions via ping and list_tools before reuse.
 
     The upstream ``MCPSessionManager`` checks ``_read_stream._closed`` /
     ``_write_stream._closed`` to decide whether a cached session is still
@@ -38,26 +48,42 @@ class KAgentMCPSessionManager(MCPSessionManager):
     stale ``mcp-session-id`` is sent to the new server, which replies
     with HTTP 404 → ``"Session terminated"``.
 
-    This subclass adds a lightweight ``send_ping()`` probe after the
-    parent returns a cached session.  If the ping fails the cached
-    session is torn down and a brand-new one is created transparently.
+    This subclass: (1) runs ``send_ping()`` after the parent returns a cached
+    session; (2) then revalidates the session with ``list_tools()`` so that
+    if the server restarted and the session id is invalid (404), we prune
+    the session from the cache and create a new one.
     """
 
     async def create_session(self, headers: dict[str, str] | None = None) -> ClientSession:
         session = await super().create_session(headers)
 
         try:
             await asyncio.wait_for(session.send_ping(), timeout=_PING_TIMEOUT_SECONDS)
-            return session
         except Exception as exc:
             if _is_server_alive_error(exc):
-                return session
-            logger.warning("MCP session failed ping validation, invalidating cached session and creating a fresh one")
-            try:
-                await self.close()
-            except Exception as close_exc:
-                logger.debug("Non-fatal error while closing stale session: %s", close_exc)
-            return await super().create_session(headers)
+                pass
+            else:
+                logger.warning(
+                    "MCP session failed ping validation, invalidating cached session and creating a fresh one"
+                )
+                try:
+                    await self.close()
+                except Exception as close_exc:
+                    logger.debug("Non-fatal error while closing stale session: %s", close_exc)
+                return await super().create_session(headers)
+
+        try:
+            await asyncio.wait_for(session.list_tools(), timeout=_SESSION_REVALIDATE_TIMEOUT_SECONDS)
+            return session
+        except Exception as exc:
+            if _is_session_invalid_error(exc):
+                logger.warning("MCP session invalid (e.g. 404), pruning from cache and creating a fresh one")
+                try:
+                    await self.close()
+                except Exception as close_exc:
+                    logger.debug("Non-fatal error while closing stale session: %s", close_exc)
+                return await super().create_session(headers)
+            raise
 
 
 def _enrich_cancelled_error(error: BaseException) -> asyncio.CancelledError:

python/packages/kagent-adk/tests/unittests/test_mcp_session_recovery.py

@@ -26,6 +26,7 @@ async def test_create_session_returns_session_when_ping_succeeds():
 
     mock_session = AsyncMock()
     mock_session.send_ping = AsyncMock(return_value=None)
+    mock_session.list_tools = AsyncMock(return_value=[])
 
     with patch.object(
         KAgentMCPSessionManager.__bases__[0],
@@ -37,6 +38,7 @@ async def test_create_session_returns_session_when_ping_succeeds():
 
     assert result is mock_session
     mock_session.send_ping.assert_awaited_once()
+    mock_session.list_tools.assert_awaited_once()
     parent_create.assert_awaited_once()
 
 
@@ -49,6 +51,7 @@ async def test_create_session_invalidates_and_retries_when_ping_fails():
 
     fresh_session = AsyncMock()
     fresh_session.send_ping = AsyncMock(return_value=None)
+    fresh_session.list_tools = AsyncMock(return_value=[])
 
     call_count = 0
 
@@ -100,6 +103,7 @@ async def _slow_ping():
 
     fresh_session = AsyncMock()
     fresh_session.send_ping = AsyncMock(return_value=None)
+    fresh_session.list_tools = AsyncMock(return_value=[])
 
     call_count = 0
 
@@ -133,6 +137,7 @@ async def test_create_session_accepts_method_not_found_as_alive():
     mock_session.send_ping = AsyncMock(
         side_effect=McpError(error=ErrorData(code=-32601, message="Method not found: ping"))
     )
+    mock_session.list_tools = AsyncMock(return_value=[])
 
     with patch.object(
         KAgentMCPSessionManager.__bases__[0],
@@ -143,9 +148,105 @@ async def test_create_session_accepts_method_not_found_as_alive():
         result = await mgr.create_session()
 
     assert result is mock_session
+    mock_session.list_tools.assert_awaited_once()
     parent_create.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+async def test_create_session_invalidates_when_list_tools_returns_session_terminated():
+    """After ping passes, list_tools is used to revalidate; 404/session terminated → prune and recreate."""
+    mgr = _make_manager()
+
+    stale_session = AsyncMock()
+    stale_session.send_ping = AsyncMock(return_value=None)
+    stale_session.list_tools = AsyncMock(side_effect=Exception("Session terminated"))
+
+    fresh_session = AsyncMock()
+    fresh_session.send_ping = AsyncMock(return_value=None)
+    fresh_session.list_tools = AsyncMock(return_value=[])
+
+    call_count = 0
+
+    async def _parent_create(headers=None):
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            return stale_session
+        return fresh_session
+
+    with (
+        patch.object(
+            KAgentMCPSessionManager.__bases__[0],
+            "create_session",
+            side_effect=_parent_create,
+        ),
+        patch.object(mgr, "close", new_callable=AsyncMock) as mock_close,
+    ):
+        result = await mgr.create_session()
+
+    assert result is fresh_session
+    mock_close.assert_awaited_once()
+    assert call_count == 2
+    stale_session.list_tools.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_create_session_invalidates_when_list_tools_returns_404():
+    """list_tools returning 404 (session invalid) triggers prune and recreate."""
+    mgr = _make_manager()
+
+    stale_session = AsyncMock()
+    stale_session.send_ping = AsyncMock(return_value=None)
+    stale_session.list_tools = AsyncMock(side_effect=McpError(error=ErrorData(code=-32000, message="404 Not Found")))
+
+    fresh_session = AsyncMock()
+    fresh_session.send_ping = AsyncMock(return_value=None)
+    fresh_session.list_tools = AsyncMock(return_value=[])
+
+    call_count = 0
+
+    async def _parent_create(headers=None):
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            return stale_session
+        return fresh_session
+
+    with (
+        patch.object(
+            KAgentMCPSessionManager.__bases__[0],
+            "create_session",
+            side_effect=_parent_create,
+        ),
+        patch.object(mgr, "close", new_callable=AsyncMock) as mock_close,
+    ):
+        result = await mgr.create_session()
+
+    assert result is fresh_session
+    mock_close.assert_awaited_once()
+    assert call_count == 2
+    stale_session.list_tools.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_create_session_propagates_non_session_error_from_list_tools():
+    """Transient errors (e.g. timeout) from list_tools are propagated, not treated as session invalid."""
+    mgr = _make_manager()
+
+    mock_session = AsyncMock()
+    mock_session.send_ping = AsyncMock(return_value=None)
+    mock_session.list_tools = AsyncMock(side_effect=asyncio.TimeoutError())
+
+    with patch.object(
+        KAgentMCPSessionManager.__bases__[0],
+        "create_session",
+        new_callable=AsyncMock,
+        return_value=mock_session,
+    ):
+        with pytest.raises(asyncio.TimeoutError):
+            await mgr.create_session()
+
+
 @pytest.mark.asyncio
 async def test_create_session_recovers_even_when_close_raises():
     """Recovery must succeed even if close() raises during stale session teardown."""
@@ -155,6 +256,8 @@ async def test_create_session_recovers_even_when_close_raises():
     stale_session.send_ping = AsyncMock(side_effect=Exception("Session terminated"))
 
     fresh_session = AsyncMock()
+    fresh_session.send_ping = AsyncMock(return_value=None)
+    fresh_session.list_tools = AsyncMock(return_value=[])
 
     call_count = 0