feat(agent): publish pod deletion cost from in-flight turns

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(

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

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: