diff --git a/tests/unit/test_agent_activities.py b/tests/unit/test_agent_activities.py index 886f63db2..799c0eaf8 100644 --- a/tests/unit/test_agent_activities.py +++ b/tests/unit/test_agent_activities.py @@ -1215,16 +1215,54 @@ async def test_successful_execution(self, mock_executor_input: AgentExecutorInpu patch( "tracecat.agent.executor.activity.SandboxedAgentExecutor" ) as mock_executor_cls, + patch( + "tracecat.agent.executor.activity.get_pod_deletion_cost_publisher" + ) as get_deletion_cost_publisher, ): mock_activity.heartbeat = MagicMock() mock_executor = MagicMock() mock_executor.run = AsyncMock(return_value=expected_result) mock_executor_cls.return_value = mock_executor + deletion_cost_publisher = MagicMock() + deletion_cost_publisher.increment = AsyncMock() + deletion_cost_publisher.decrement = AsyncMock() + get_deletion_cost_publisher.return_value = deletion_cost_publisher result = await run_agent_activity(mock_executor_input) assert result == expected_result mock_executor_cls.assert_called_once_with(input=mock_executor_input) + deletion_cost_publisher.increment.assert_awaited_once() + deletion_cost_publisher.decrement.assert_awaited_once() + + @pytest.mark.anyio + async def test_decrements_deletion_cost_when_execution_raises( + self, + mock_executor_input: AgentExecutorInput, + ) -> None: + with ( + patch("tracecat.agent.executor.activity.activity") as mock_activity, + patch( + "tracecat.agent.executor.activity.SandboxedAgentExecutor" + ) as mock_executor_cls, + patch( + "tracecat.agent.executor.activity.get_pod_deletion_cost_publisher" + ) as get_deletion_cost_publisher, + ): + mock_activity.heartbeat = MagicMock() + mock_executor = MagicMock() + mock_executor.run = AsyncMock(side_effect=RuntimeError("runtime failed")) + mock_executor_cls.return_value = mock_executor + deletion_cost_publisher = MagicMock() + deletion_cost_publisher.increment = AsyncMock() + deletion_cost_publisher.decrement = AsyncMock() + get_deletion_cost_publisher.return_value = deletion_cost_publisher + + with pytest.raises(RuntimeError, match="runtime failed"): + await run_agent_activity(mock_executor_input) + + deletion_cost_publisher.increment.assert_awaited_once() + deletion_cost_publisher.decrement.assert_awaited_once() @pytest.mark.anyio async def test_emit_session_done_pushes_done_to_active_stream( diff --git a/tests/unit/test_agent_executor_deletion_cost.py b/tests/unit/test_agent_executor_deletion_cost.py new file mode 100644 index 000000000..463ee5d90 --- /dev/null +++ b/tests/unit/test_agent_executor_deletion_cost.py @@ -0,0 +1,260 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Callable, Coroutine +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import httpx +import orjson +import pytest + +from tracecat.agent.executor import deletion_cost + +type AsyncRequestHandler = Callable[ + [httpx.Request], + Coroutine[None, None, httpx.Response], +] + + +def _configure_enabled_publisher( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> tuple[Path, Path]: + token_path = tmp_path / "token" + ca_path = tmp_path / "ca.crt" + token_path.write_text("initial-token") + ca_path.write_text("test-ca") + + monkeypatch.setattr(deletion_cost, "SERVICE_ACCOUNT_TOKEN_PATH", token_path) + monkeypatch.setattr(deletion_cost, "SERVICE_ACCOUNT_CA_PATH", ca_path) + monkeypatch.setattr( + deletion_cost.config, + "TRACECAT__AGENT_EXECUTOR_POD_DELETION_COST_ENABLED", + True, + ) + monkeypatch.setattr( + deletion_cost.config, + "TRACECAT__K8S_POD_NAME", + "agent-executor-abc", + ) + monkeypatch.setattr( + deletion_cost.config, + "TRACECAT__K8S_POD_NAMESPACE", + "tracecat", + ) + monkeypatch.setenv("KUBERNETES_SERVICE_HOST", "10.0.0.1") + monkeypatch.setenv("KUBERNETES_SERVICE_PORT", "6443") + return token_path, ca_path + + +def _install_mock_client( + monkeypatch: pytest.MonkeyPatch, + handler: AsyncRequestHandler, +) -> list[tuple[str, float]]: + real_async_client = httpx.AsyncClient + client_settings: list[tuple[str, float]] = [] + + def client_factory(*, verify: str, timeout: float) -> httpx.AsyncClient: + client_settings.append((verify, timeout)) + return real_async_client( + transport=httpx.MockTransport(handler), + timeout=timeout, + ) + + monkeypatch.setattr(deletion_cost.httpx, "AsyncClient", client_factory) + return client_settings + + +@pytest.mark.parametrize( + "disabled_condition", + [ + "flag_off", + "no_service_host", + "missing_token", + "missing_pod_name", + "missing_pod_namespace", + ], +) +@pytest.mark.anyio +async def test_disabled_publisher_is_noop_without_http( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + disabled_condition: str, +) -> None: + token_path, _ = _configure_enabled_publisher(monkeypatch, tmp_path) + + match disabled_condition: + case "flag_off": + monkeypatch.setattr( + deletion_cost.config, + "TRACECAT__AGENT_EXECUTOR_POD_DELETION_COST_ENABLED", + False, + ) + case "no_service_host": + monkeypatch.delenv("KUBERNETES_SERVICE_HOST") + case "missing_token": + token_path.unlink() + case "missing_pod_name": + monkeypatch.setattr( + deletion_cost.config, + "TRACECAT__K8S_POD_NAME", + None, + ) + case "missing_pod_namespace": + monkeypatch.setattr( + deletion_cost.config, + "TRACECAT__K8S_POD_NAMESPACE", + None, + ) + + def fail_client(**_: object) -> httpx.AsyncClient: + pytest.fail("disabled publisher attempted to create an HTTP client") + + monkeypatch.setattr(deletion_cost.httpx, "AsyncClient", fail_client) + publisher = deletion_cost.PodDeletionCostPublisher() + + await publisher.increment() + await publisher.decrement() + + +@pytest.mark.anyio +async def test_increment_and_decrement_publish_exact_patch( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + token_path, ca_path = _configure_enabled_publisher(monkeypatch, tmp_path) + requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, request=request) + + client_settings = _install_mock_client(monkeypatch, handler) + publisher = deletion_cost.PodDeletionCostPublisher() + + await publisher.increment() + token_path.write_text("rotated-token") + await publisher.decrement() + + assert len(requests) == 2 + assert all( + str(request.url) + == ("https://10.0.0.1:6443/api/v1/namespaces/tracecat/pods/agent-executor-abc") + for request in requests + ) + assert [request.method for request in requests] == ["PATCH", "PATCH"] + assert [request.headers["content-type"] for request in requests] == [ + "application/merge-patch+json", + "application/merge-patch+json", + ] + assert [request.headers["authorization"] for request in requests] == [ + "Bearer initial-token", + "Bearer rotated-token", + ] + assert [orjson.loads(request.content) for request in requests] == [ + { + "metadata": { + "annotations": { + "controller.kubernetes.io/pod-deletion-cost": "1", + } + } + }, + { + "metadata": { + "annotations": { + "controller.kubernetes.io/pod-deletion-cost": "0", + } + } + }, + ] + assert client_settings == [ + (str(ca_path), deletion_cost.PUBLISH_TIMEOUT_SECONDS), + (str(ca_path), deletion_cost.PUBLISH_TIMEOUT_SECONDS), + ] + + +@pytest.mark.anyio +async def test_rapid_increment_and_decrement_coalesce_to_latest_count( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + _configure_enabled_publisher(monkeypatch, tmp_path) + requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, request=request) + + _install_mock_client(monkeypatch, handler) + publisher = deletion_cost.PodDeletionCostPublisher() + + await asyncio.gather( + publisher.increment(), + publisher.decrement(), + ) + + assert len(requests) == 1 + assert orjson.loads(requests[0].content) == { + "metadata": { + "annotations": { + "controller.kubernetes.io/pod-deletion-cost": "0", + } + } + } + + +@pytest.mark.anyio +async def test_three_403_responses_disable_publisher_without_more_requests( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + _configure_enabled_publisher(monkeypatch, tmp_path) + requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(403, request=request) + + _install_mock_client(monkeypatch, handler) + test_logger = SimpleNamespace(debug=Mock(), warning=Mock()) + monkeypatch.setattr(deletion_cost, "logger", test_logger) + publisher = deletion_cost.PodDeletionCostPublisher() + + await publisher.increment() + await publisher.increment() + await publisher.decrement() + await publisher.increment() + + assert len(requests) == 3 + assert test_logger.warning.call_count == 3 + final_warning = test_logger.warning.call_args + assert final_warning.args[0].startswith( + "Disabling Kubernetes pod deletion cost publisher" + ) + assert final_warning.kwargs["failures"] == 3 + assert final_warning.kwargs["status_code"] == 403 + + +@pytest.mark.anyio +async def test_success_resets_consecutive_failure_counter( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + _configure_enabled_publisher(monkeypatch, tmp_path) + status_codes = iter([403, 200, 403, 403, 403]) + requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(next(status_codes), request=request) + + _install_mock_client(monkeypatch, handler) + publisher = deletion_cost.PodDeletionCostPublisher() + + for _ in range(5): + await publisher.increment() + await publisher.decrement() + + assert len(requests) == 5 diff --git a/tracecat/agent/executor/activity.py b/tracecat/agent/executor/activity.py index bc2432542..4b16c5096 100644 --- a/tracecat/agent/executor/activity.py +++ b/tracecat/agent/executor/activity.py @@ -42,6 +42,9 @@ is_stdio_mcp_server, requires_sandbox_internet_access, ) +from tracecat.agent.executor.deletion_cost import ( + get_pod_deletion_cost_publisher, +) from tracecat.agent.executor.loopback import ( LoopbackHandler, LoopbackInput, @@ -1176,34 +1179,41 @@ async def run_agent_activity(input: AgentExecutorInput) -> AgentExecutorResult: Returns: AgentExecutorResult with execution status and terminal output. """ - sandbox_mode = "direct" if TRACECAT__DISABLE_NSJAIL else "nsjail" - activity.heartbeat( - f"Starting agent execution ({sandbox_mode} mode): {input.session_id}" - ) + deletion_cost_publisher = get_pod_deletion_cost_publisher() + try: + await deletion_cost_publisher.increment() + sandbox_mode = "direct" if TRACECAT__DISABLE_NSJAIL else "nsjail" + activity.heartbeat( + f"Starting agent execution ({sandbox_mode} mode): {input.session_id}" + ) - input = await _hydrate_sdk_session_history(input) - - # Stdio MCP servers are spawned directly by the runtime; unlike HTTP - # servers they have no per-call secret resolution hook downstream. The - # configs in ``input.config.mcp_servers`` and each subagent's - # ``config.mcp_servers`` arrive in refs-only shape (no ``env``) — hydrate - # from the DB here so the spawned processes get their credentials. - config = cast(Any, input.config) - config.mcp_servers = await _hydrate_stdio_env(config.mcp_servers, role=input.role) - for subagent in input.subagents: - subagent.config.mcp_servers = await _hydrate_stdio_env( - subagent.config.mcp_servers, role=input.role + input = await _hydrate_sdk_session_history(input) + + # Stdio MCP servers are spawned directly by the runtime; unlike HTTP + # servers they have no per-call secret resolution hook downstream. The + # configs in ``input.config.mcp_servers`` and each subagent's + # ``config.mcp_servers`` arrive in refs-only shape (no ``env``) — hydrate + # from the DB here so the spawned processes get their credentials. + config = cast(Any, input.config) + config.mcp_servers = await _hydrate_stdio_env( + config.mcp_servers, role=input.role ) + for subagent in input.subagents: + subagent.config.mcp_servers = await _hydrate_stdio_env( + subagent.config.mcp_servers, role=input.role + ) - executor = SandboxedAgentExecutor(input=input) - result = await executor.run() + executor = SandboxedAgentExecutor(input=input) + result = await executor.run() - if result.success: - activity.heartbeat(f"Agent execution completed: {input.session_id}") - else: - activity.heartbeat(f"Agent execution failed: {result.error}") + if result.success: + activity.heartbeat(f"Agent execution completed: {input.session_id}") + else: + activity.heartbeat(f"Agent execution failed: {result.error}") - return result + return result + finally: + await deletion_cost_publisher.decrement() async def _hydrate_sdk_session_history(input: AgentExecutorInput) -> AgentExecutorInput: diff --git a/tracecat/agent/executor/deletion_cost.py b/tracecat/agent/executor/deletion_cost.py new file mode 100644 index 000000000..3504f8950 --- /dev/null +++ b/tracecat/agent/executor/deletion_cost.py @@ -0,0 +1,172 @@ +"""Best-effort Kubernetes pod deletion cost publishing.""" + +from __future__ import annotations + +import asyncio +import os +from pathlib import Path +from typing import TypedDict + +import httpx + +from tracecat import config +from tracecat.logger import logger + +SERVICE_ACCOUNT_TOKEN_PATH = Path("/var/run/secrets/kubernetes.io/serviceaccount/token") +SERVICE_ACCOUNT_CA_PATH = Path("/var/run/secrets/kubernetes.io/serviceaccount/ca.crt") +POD_DELETION_COST_ANNOTATION = "controller.kubernetes.io/pod-deletion-cost" +PUBLISH_TIMEOUT_SECONDS = 5.0 +MAX_CONSECUTIVE_FAILURES = 3 + + +class _PodMetadataPatch(TypedDict): + annotations: dict[str, str] + + +class _PodPatch(TypedDict): + metadata: _PodMetadataPatch + + +class PodDeletionCostPublisher: + """Publish this process's in-flight agent turn count to its Kubernetes pod.""" + + def __init__(self) -> None: + self._api_host = os.environ.get("KUBERNETES_SERVICE_HOST") or "" + self._api_port = os.environ.get("KUBERNETES_SERVICE_PORT") or "443" + self._pod_name = (config.TRACECAT__K8S_POD_NAME or "").strip() + self._pod_namespace = (config.TRACECAT__K8S_POD_NAMESPACE or "").strip() + self._enabled = bool( + config.TRACECAT__AGENT_EXECUTOR_POD_DELETION_COST_ENABLED + and self._api_host + and SERVICE_ACCOUNT_TOKEN_PATH.exists() + and self._pod_name + and self._pod_namespace + ) + self._count = 0 + self._last_published_count: int | None = None + self._consecutive_failures = 0 + self._publishing = False + self._lock = asyncio.Lock() + + if not self._enabled: + logger.debug("Kubernetes pod deletion cost publisher disabled") + + async def increment(self) -> None: + """Add one in-flight agent turn and publish the latest count.""" + await self._adjust(1) + + async def decrement(self) -> None: + """Remove one in-flight agent turn and publish the latest count.""" + await self._adjust(-1) + + async def _adjust(self, delta: int) -> None: + if not self._enabled: + return + + async with self._lock: + if not self._enabled: + return + self._count += delta + if self._publishing: + return + self._publishing = True + + try: + # Give changes scheduled in the same event-loop turn a chance to + # coalesce before taking the first snapshot. + await asyncio.sleep(0) + await self._publish_latest() + except BaseException: + async with self._lock: + self._publishing = False + raise + + async def _publish_latest(self) -> None: + while self._enabled: + async with self._lock: + count = self._count + if count == self._last_published_count: + self._publishing = False + return + + published = await self._publish(count) + + async with self._lock: + if published: + self._last_published_count = count + if not self._enabled or self._count == count: + self._publishing = False + return + + async def _publish(self, count: int) -> bool: + url = ( + f"https://{self._api_host}:{self._api_port}/api/v1/namespaces/" + f"{self._pod_namespace}/pods/{self._pod_name}" + ) + patch: _PodPatch = { + "metadata": { + "annotations": { + POD_DELETION_COST_ANNOTATION: str(count), + } + } + } + + try: + token = SERVICE_ACCOUNT_TOKEN_PATH.read_text().strip() + async with httpx.AsyncClient( + verify=str(SERVICE_ACCOUNT_CA_PATH), + timeout=PUBLISH_TIMEOUT_SECONDS, + ) as client: + response = await client.patch( + url, + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/merge-patch+json", + }, + json=patch, + ) + response.raise_for_status() + except httpx.HTTPStatusError as exc: + self._record_failure( + status_code=exc.response.status_code, + error=str(exc), + ) + return False + except (httpx.HTTPError, OSError) as exc: + self._record_failure(status_code=None, error=str(exc)) + return False + + self._consecutive_failures = 0 + return True + + def _record_failure(self, *, status_code: int | None, error: str) -> None: + self._consecutive_failures += 1 + if self._consecutive_failures >= MAX_CONSECUTIVE_FAILURES: + self._enabled = False + logger.warning( + "Disabling Kubernetes pod deletion cost publisher after " + "consecutive failures", + failures=self._consecutive_failures, + status_code=status_code, + error=error, + ) + return + + logger.warning( + "Failed to publish Kubernetes pod deletion cost", + failures=self._consecutive_failures, + status_code=status_code, + error=error, + ) + + +# Lazy singleton - no lifespan required. +_pod_deletion_cost_publisher: PodDeletionCostPublisher | None = None + + +def get_pod_deletion_cost_publisher() -> PodDeletionCostPublisher: + """Get the global pod deletion cost publisher instance.""" + global _pod_deletion_cost_publisher + if _pod_deletion_cost_publisher is None: + _pod_deletion_cost_publisher = PodDeletionCostPublisher() + return _pod_deletion_cost_publisher diff --git a/tracecat/config.py b/tracecat/config.py index a7184e706..8ea30f3cc 100644 --- a/tracecat/config.py +++ b/tracecat/config.py @@ -750,6 +750,17 @@ def _parse_auth_types() -> set[AuthType]: ) """Best-effort readiness sentinel written after the agent executor starts.""" +TRACECAT__AGENT_EXECUTOR_POD_DELETION_COST_ENABLED = env_bool( + "TRACECAT__AGENT_EXECUTOR_POD_DELETION_COST_ENABLED", default=True +) +"""Enable best-effort Kubernetes pod deletion cost publishing.""" + +TRACECAT__K8S_POD_NAME = os.environ.get("TRACECAT__K8S_POD_NAME") or None +"""Kubernetes pod name supplied through the downward API.""" + +TRACECAT__K8S_POD_NAMESPACE = os.environ.get("TRACECAT__K8S_POD_NAMESPACE") or None +"""Kubernetes pod namespace supplied through the downward API.""" + TRACECAT__LITELLM_PORT = int(os.environ.get("TRACECAT__LITELLM_PORT") or 4000) """Bind port for the managed LiteLLM service."""