diff --git a/tests/unit/test_agent_approval_audit.py b/tests/unit/test_agent_approval_audit.py new file mode 100644 index 000000000..51d0d6bae --- /dev/null +++ b/tests/unit/test_agent_approval_audit.py @@ -0,0 +1,288 @@ +from __future__ import annotations + +import asyncio +import uuid +from contextlib import asynccontextmanager +from datetime import UTC, datetime +from typing import Any, Literal +from unittest.mock import AsyncMock, Mock + +import pytest +from pydantic_ai.tools import ToolApproved, ToolDenied + +from tracecat.agent.session import service as session_service_module +from tracecat.agent.session.service import ( + AgentSessionService, + _emit_approval_audit_events, + _PendingApproval, + _schedule_approval_audit_events, + _should_emit_approval_audit, + _ValidatedContinuation, +) +from tracecat.audit.sanitization import sanitize_audit_metadata +from tracecat.audit.types import AuditEventInput +from tracecat.auth.types import Role +from tracecat.chat.schemas import ApprovalDecision, ContinueRunRequest +from tracecat.contexts import RequestAuditContext + + +def _user_role() -> Role: + return Role( + type="user", + workspace_id=uuid.uuid4(), + organization_id=uuid.uuid4(), + user_id=uuid.uuid4(), + service_id="tracecat-api", + ) + + +def _event() -> AuditEventInput: + return AuditEventInput( + resource_type="agent_approval", + resource_id=uuid.uuid4(), + action="accept", + data={"tool_call_id": "tool-call-1"}, + ) + + +def test_build_approval_audit_events_maps_decisions_without_sensitive_values() -> None: + pending = { + "approve-call": _PendingApproval( + uuid.uuid4(), "approve-call", "tools.ticket.get" + ), + "override-call": _PendingApproval( + uuid.uuid4(), "override-call", "tools.ticket.update" + ), + "deny-call": _PendingApproval(uuid.uuid4(), "deny-call", "core.http_request"), + } + validated = _ValidatedContinuation( + approval_map={ + "approve-call": True, + "override-call": ToolApproved( + override_args={"authorization": "credential-value"} + ), + "deny-call": ToolDenied(message="Needs explicit review"), + }, + decision_metadata={}, + ) + request = ContinueRunRequest( + source="inbox", + decisions=[ + ApprovalDecision(tool_call_id="approve-call", action="approve"), + ApprovalDecision( + tool_call_id="override-call", + action="override", + override_args={"authorization": "credential-value"}, + metadata={"prompt": "sensitive-prompt"}, + ), + ApprovalDecision( + tool_call_id="deny-call", + action="deny", + reason="Needs explicit review", + ), + ], + ) + session_id = uuid.uuid4() + run_id = uuid.uuid4() + dedupe_id = uuid.uuid4() + decided_at = datetime.now(UTC) + + events = AgentSessionService._build_approval_audit_events( + pending_approvals=pending, + validated=validated, + request=request, + session_id=session_id, + run_id=run_id, + dedupe_id=dedupe_id, + decided_at=decided_at, + ) + + assert [event.action for event in events] == ["accept", "accept", "reject"] + assert [event.resource_id for event in events] == [ + pending["approve-call"].approval_id, + pending["override-call"].approval_id, + pending["deny-call"].approval_id, + ] + assert events[1].data is not None + assert events[1].data["arguments_overridden"] is True + assert events[2].data is not None + assert events[2].data["denial_reason"] == "Needs explicit review" + assert "credential-value" not in repr(events) + assert "sensitive-prompt" not in repr(events) + assert all(event.created_at == decided_at for event in events) + + +def test_approval_audit_sanitizer_drops_secret_bearing_denial_reason() -> None: + sanitized = sanitize_audit_metadata( + { + "decision": "deny", + "denial_reason": "Authorization: Bearer secret-token", + "arguments_overridden": False, + } + ) + + assert sanitized == { + "decision": "deny", + "arguments_overridden": False, + } + + +def test_approval_audit_build_is_noop_without_accepted_decisions() -> None: + events = AgentSessionService._build_approval_audit_events( + pending_approvals={}, + validated=_ValidatedContinuation( + approval_map={}, + decision_metadata={}, + ), + request=ContinueRunRequest(source="inbox", decisions=[]), + session_id=uuid.uuid4(), + run_id=uuid.uuid4(), + dedupe_id=uuid.uuid4(), + decided_at=datetime.now(UTC), + ) + + assert events == () + + +@pytest.mark.parametrize( + ("role", "source", "expected"), + [ + (_user_role(), "inbox", True), + (_user_role(), "slack", False), + ( + Role( + type="service", + workspace_id=uuid.uuid4(), + user_id=uuid.uuid4(), + service_id="tracecat-api", + ), + "inbox", + False, + ), + ], +) +def test_approval_audit_only_includes_first_party_user_submissions( + role: Role, + source: Literal["inbox", "slack"], + expected: bool, +) -> None: + assert _should_emit_approval_audit(role=role, source=source) is expected + + +@pytest.mark.anyio +async def test_approval_audit_deduplicates_concurrent_successful_submissions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class RedisDouble: + def __init__(self) -> None: + self.acquired = False + self.lock = asyncio.Lock() + self.calls = 0 + + async def set_if_not_exists( + self, key: str, value: str, *, expire_seconds: int + ) -> bool: + del key, value, expire_seconds + async with self.lock: + self.calls += 1 + if self.acquired: + return False + self.acquired = True + return True + + redis = RedisDouble() + delivered: list[tuple[AuditEventInput, ...]] = [] + + class AuditServiceDouble: + @classmethod + @asynccontextmanager + async def with_session(cls, *, role: Role): + del role + yield cls() + + async def create_events( + self, + events: tuple[AuditEventInput, ...], + *, + request_audit: RequestAuditContext | None, + ) -> None: + del self, request_audit + delivered.append(events) + + monkeypatch.setattr( + session_service_module, "get_redis_client", AsyncMock(return_value=redis) + ) + monkeypatch.setattr(session_service_module, "AuditService", AuditServiceDouble) + kwargs: dict[str, Any] = { + "events": (_event(),), + "role": _user_role(), + "request_audit": RequestAuditContext( + client_ip="192.0.2.1", user_agent="TracecatTest/1.0" + ), + "dedupe_id": uuid.uuid4(), + } + + await asyncio.gather( + _emit_approval_audit_events(**kwargs), + _emit_approval_audit_events(**kwargs), + ) + + assert redis.calls == 2 + assert len(delivered) == 1 + + +@pytest.mark.anyio +async def test_slow_approval_audit_enrichment_does_not_block_scheduler( + monkeypatch: pytest.MonkeyPatch, +) -> None: + gate = asyncio.Event() + started = asyncio.Event() + + async def slow_emit(**kwargs: Any) -> None: + del kwargs + started.set() + await gate.wait() + + monkeypatch.setattr( + session_service_module, "_emit_approval_audit_events", slow_emit + ) + session_service_module._approval_audit_tasks.clear() + + _schedule_approval_audit_events( + events=(_event(),), + role=_user_role(), + request_audit=None, + dedupe_id=uuid.uuid4(), + ) + await started.wait() + + assert len(session_service_module._approval_audit_tasks) == 1 + gate.set() + await asyncio.gather(*session_service_module._approval_audit_tasks) + await asyncio.sleep(0) + assert session_service_module._approval_audit_tasks == set() + + +def test_approval_audit_scheduler_sheds_at_capacity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + emit = AsyncMock() + warning = Mock() + monkeypatch.setattr(session_service_module, "MAX_PENDING_APPROVAL_AUDIT_TASKS", 0) + monkeypatch.setattr(session_service_module, "_emit_approval_audit_events", emit) + monkeypatch.setattr(session_service_module.logger, "warning", warning) + session_service_module._approval_audit_tasks.clear() + + _schedule_approval_audit_events( + events=(_event(),), + role=_user_role(), + request_audit=None, + dedupe_id=uuid.uuid4(), + ) + + emit.assert_not_called() + warning.assert_called_once_with( + "Dropped approval audit batch; pending limit reached", + event_count=1, + max_pending=0, + ) diff --git a/tests/unit/test_audit_service.py b/tests/unit/test_audit_service.py index cd58cbda3..324ae64a9 100644 --- a/tests/unit/test_audit_service.py +++ b/tests/unit/test_audit_service.py @@ -8,7 +8,7 @@ from datetime import datetime from types import SimpleNamespace from typing import Any, cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import orjson @@ -19,7 +19,13 @@ from sqlalchemy.ext.asyncio import AsyncSession from tenacity import wait_none +from tracecat.audit import batch as audit_batch_module from tracecat.audit import service as audit_service_module +from tracecat.audit.batch import ( + AuditBatchDelivery, + AuditBatchEvent, + deliver_audit_batch, +) from tracecat.audit.enums import AuditEventActor, AuditEventStatus from tracecat.audit.logger import ( AuditEventDetails, @@ -31,7 +37,7 @@ _AuditDelivery, _spawn_delivery, ) -from tracecat.audit.types import AuditEvent, AuditMetadata +from tracecat.audit.types import AuditEvent, AuditEventInput, AuditMetadata from tracecat.auth.types import PlatformRole, Role from tracecat.auth.users import UserManager from tracecat.authz.scopes import ADMIN_SCOPES @@ -1132,6 +1138,149 @@ def _delivery(tag: str, url: str = "https://example.com/audit") -> _AuditDeliver ) +@pytest.mark.anyio +async def test_create_events_resolves_webhook_and_actor_once_for_batch( + monkeypatch: pytest.MonkeyPatch, + audit_service: AuditService, +) -> None: + webhook_url = "https://example.com/audit" + webhook = AsyncMock(return_value=webhook_url) + actor = AsyncMock(return_value="actor@example.test") + post_events = AsyncMock() + monkeypatch.setattr(audit_service, "_get_webhook_url", webhook) + monkeypatch.setattr(audit_service, "_get_actor_label", actor) + monkeypatch.setattr(audit_service, "_post_events", post_events) + request_audit = RequestAuditContext( + client_ip="192.0.2.10", user_agent="TracecatAuditTest/1.0" + ) + + await audit_service.create_events( + [ + AuditEventInput( + resource_type="agent_approval", + resource_id=uuid.uuid4(), + action="accept", + data={"tool_call_id": "tool-call-1"}, + ), + AuditEventInput( + resource_type="agent_approval", + resource_id=uuid.uuid4(), + action="reject", + data={"tool_call_id": "tool-call-2"}, + ), + ], + request_audit=request_audit, + ) + + webhook.assert_awaited_once() + actor.assert_awaited_once() + post_events.assert_awaited_once() + assert post_events.await_args is not None + payloads = post_events.await_args.kwargs["payloads"] + assert len(payloads) == 2 + assert {payload.action for payload in payloads} == {"accept", "reject"} + assert all(payload.ip_address == request_audit.client_ip for payload in payloads) + + +@pytest.mark.anyio +async def test_post_events_resolves_shared_configuration_once( + monkeypatch: pytest.MonkeyPatch, + audit_service: AuditService, +) -> None: + custom_headers = AsyncMock(return_value={"X-Audit-Test": "true"}) + custom_payload = AsyncMock(return_value={"schema_version": "1"}) + verify_ssl = AsyncMock(return_value=True) + payload_attribute = AsyncMock(return_value="event") + spawn = Mock() + monkeypatch.setattr(audit_service, "_get_custom_headers", custom_headers) + monkeypatch.setattr(audit_service, "_get_custom_payload", custom_payload) + monkeypatch.setattr(audit_service, "_get_verify_ssl", verify_ssl) + monkeypatch.setattr(audit_service, "_get_payload_attribute", payload_attribute) + monkeypatch.setattr(audit_service_module, "spawn_audit_batch", spawn) + role = audit_service.role + assert role is not None and role.actor_id is not None + payloads = [ + AuditEvent( + actor_type=AuditEventActor.USER, + actor_id=role.actor_id, + resource_type="agent_approval", + action="accept", + status=AuditEventStatus.SUCCESS, + ), + AuditEvent( + actor_type=AuditEventActor.USER, + actor_id=role.actor_id, + resource_type="agent_approval", + action="reject", + status=AuditEventStatus.SUCCESS, + ), + ] + + await audit_service._post_events( + webhook_url="https://example.com/audit", payloads=payloads + ) + + custom_headers.assert_awaited_once() + custom_payload.assert_awaited_once() + verify_ssl.assert_awaited_once() + payload_attribute.assert_awaited_once() + spawn.assert_called_once() + batch = spawn.call_args.args[0] + assert isinstance(batch, AuditBatchDelivery) + assert len(batch.events) == 2 + assert all("event" in item.request_payload for item in batch.events) + + +@pytest.mark.anyio +async def test_deliver_batch_reuses_one_client_and_caps_concurrency( + monkeypatch: pytest.MonkeyPatch, +) -> None: + active = 0 + maximum_active = 0 + post_count = 0 + + async def post(*args: object, **kwargs: object) -> httpx.Response: + nonlocal active, maximum_active, post_count + del args, kwargs + active += 1 + post_count += 1 + maximum_active = max(maximum_active, active) + await asyncio.sleep(0) + active -= 1 + return httpx.Response( + 200, + request=httpx.Request("POST", "https://example.com/audit"), + ) + + client = MagicMock() + client.post = post + client_context = MagicMock() + client_context.__aenter__ = AsyncMock(return_value=client) + client_context.__aexit__ = AsyncMock(return_value=None) + client_factory = Mock(return_value=client_context) + monkeypatch.setattr(audit_batch_module.httpx, "AsyncClient", client_factory) + events = tuple( + AuditBatchEvent( + request_payload={"resource_id": f"event-{index}"}, + resource_type="workflow", + action="update", + ) + for index in range(12) + ) + batch = AuditBatchDelivery( + webhook_url="https://example.com/audit", + events=events, + headers=None, + verify_ssl=True, + ) + + await deliver_audit_batch(batch) + + client_factory.assert_called_once() + assert post_count == len(events) + assert maximum_active <= 4 + + @pytest.fixture(autouse=True) def no_retry_backoff(monkeypatch: pytest.MonkeyPatch) -> None: """Zero the delivery retry backoff so retry-path tests don't sleep.""" diff --git a/tracecat/agent/session/service.py b/tracecat/agent/session/service.py index 566edd3d2..c25095347 100644 --- a/tracecat/agent/session/service.py +++ b/tracecat/agent/session/service.py @@ -96,6 +96,8 @@ from tracecat.artifacts.bindings import ArtifactSideEffect from tracecat.artifacts.schemas import Artifact, ArtifactAdapter, ArtifactType from tracecat.audit.logger import audit_log +from tracecat.audit.service import AuditService +from tracecat.audit.types import AuditAction, AuditEventInput, AuditMetadataValue from tracecat.auth.types import Role from tracecat.authz.scopes import SERVICE_PRINCIPAL_SCOPES from tracecat.cases.prompts import CaseCopilotPrompts @@ -117,6 +119,7 @@ filter_workspace_chat_tools_for_scopes, get_default_tools, ) +from tracecat.contexts import RequestAuditContext, ctx_request_audit from tracecat.db.models import ( APPROVAL_STATUS_ENUM, AgentSession, @@ -155,7 +158,84 @@ AUTO_TITLE_SERVICE_ID = "tracecat-api" APPROVAL_CONTINUATION_DEDUP_TTL_SECONDS = 5 * 60 +APPROVAL_AUDIT_DEDUP_TTL_SECONDS = 24 * 60 * 60 +MAX_PENDING_APPROVAL_AUDIT_TASKS = 64 _background_tasks: set[asyncio.Task[None]] = set() +_approval_audit_tasks: set[asyncio.Task[None]] = set() + + +async def _emit_approval_audit_events( + *, + events: tuple[AuditEventInput, ...], + role: Role, + request_audit: RequestAuditContext | None, + dedupe_id: uuid.UUID, +) -> None: + """Deduplicate and enqueue an approval audit batch outside the request path.""" + try: + try: + redis_client = await get_redis_client() + acquired = await redis_client.set_if_not_exists( + f"agent-approval-audit:{dedupe_id}", + "1", + expire_seconds=APPROVAL_AUDIT_DEDUP_TTL_SECONDS, + ) + if not acquired: + return + except Exception as exc: + # The stable dedupe ID remains in every event so consumers can + # converge the rare duplicate produced while Redis is unavailable. + logger.warning( + "Approval audit dedup unavailable; emitting best-effort", + error_type=type(exc).__name__, + ) + + async with AuditService.with_session(role=role) as audit_service: + await audit_service.create_events(events, request_audit=request_audit) + except Exception as exc: + logger.warning( + "Failed to enqueue approval audit events", + error_type=type(exc).__name__, + event_count=len(events), + ) + + +def _schedule_approval_audit_events( + *, + events: tuple[AuditEventInput, ...], + role: Role, + request_audit: RequestAuditContext | None, + dedupe_id: uuid.UUID, +) -> None: + """Schedule approval audit enrichment without delaying workflow progress.""" + for stranded in [ + task for task in _approval_audit_tasks if task.get_loop().is_closed() + ]: + _approval_audit_tasks.discard(stranded) + if len(_approval_audit_tasks) >= MAX_PENDING_APPROVAL_AUDIT_TASKS: + logger.warning( + "Dropped approval audit batch; pending limit reached", + event_count=len(events), + max_pending=MAX_PENDING_APPROVAL_AUDIT_TASKS, + ) + return + task = asyncio.get_running_loop().create_task( + _emit_approval_audit_events( + events=events, + role=role, + request_audit=request_audit, + dedupe_id=dedupe_id, + ) + ) + _approval_audit_tasks.add(task) + task.add_done_callback(_approval_audit_tasks.discard) + + +def _should_emit_approval_audit( + *, role: Role, source: Literal["inbox", "slack"] +) -> bool: + """Limit first-party approval audits to authenticated UI/API users.""" + return source == "inbox" and role.type == "user" and role.user_id is not None async def _auto_title_session( @@ -348,6 +428,14 @@ class _DecisionRow(NamedTuple): approved_at: datetime +class _PendingApproval(NamedTuple): + """Privacy-bounded fields captured by the existing pending-approval query.""" + + approval_id: uuid.UUID + tool_call_id: str + tool_name: str + + class AgentSessionService(BaseWorkspaceService): """Service for managing agent sessions and history.""" @@ -1333,15 +1421,21 @@ async def has_pending_approvals(self, session_id: uuid.UUID) -> bool: result = await self.session.execute(stmt) return result.scalar_one_or_none() is not None - async def _pending_approval_tool_call_ids(self, session_id: uuid.UUID) -> set[str]: - """Return pending approval tool-call IDs for one session.""" - stmt = select(Approval.tool_call_id).where( + async def _pending_approvals( + self, session_id: uuid.UUID + ) -> dict[str, _PendingApproval]: + """Return audit-safe fields for the session's pending approvals.""" + stmt = select(Approval.id, Approval.tool_call_id, Approval.tool_name).where( Approval.workspace_id == self.workspace_id, Approval.session_id == session_id, Approval.status == ApprovalStatus.PENDING, ) result = await self.session.execute(stmt) - return set(result.scalars().all()) + approvals = ( + _PendingApproval(approval_id, tool_call_id, tool_name) + for approval_id, tool_call_id, tool_name in result.tuples().all() + ) + return {approval.tool_call_id: approval for approval in approvals} async def _settled_approval_decisions( self, session_id: uuid.UUID @@ -2229,6 +2323,67 @@ def _validate_continuation_decisions( decision_metadata=decision_metadata, ) + @staticmethod + def _approval_audit_dedupe_id( + *, + workspace_id: uuid.UUID, + session_id: uuid.UUID, + run_id: uuid.UUID, + stream_id: uuid.UUID, + ) -> uuid.UUID: + """Derive an opaque ID from the stable Temporal approval update identity.""" + update_identity = ( + f"tracecat:set-approvals:{workspace_id}:{session_id}:{run_id}:{stream_id}" + ) + return uuid.uuid5(uuid.NAMESPACE_URL, update_identity) + + @staticmethod + def _build_approval_audit_events( + *, + pending_approvals: Mapping[str, _PendingApproval], + validated: _ValidatedContinuation, + request: ContinueRunRequest, + session_id: uuid.UUID, + run_id: uuid.UUID, + dedupe_id: uuid.UUID, + decided_at: datetime, + ) -> tuple[AuditEventInput, ...]: + """Build one privacy-bounded event for each accepted tool decision.""" + decisions_by_tool_call = { + decision.tool_call_id: decision for decision in request.decisions + } + events: list[AuditEventInput] = [] + for tool_call_id in validated.approval_map: + approval = pending_approvals[tool_call_id] + decision = decisions_by_tool_call[tool_call_id] + action: AuditAction = "reject" if decision.action == "deny" else "accept" + data: dict[str, AuditMetadataValue] = { + "session_id": str(session_id), + "run_id": str(run_id), + "tool_call_id": tool_call_id, + "tool_name": approval.tool_name, + "decision": decision.action, + "source": request.source, + "decision_timestamp": decided_at.isoformat(), + "arguments_overridden": decision.action == "override", + "dedupe_id": str(dedupe_id), + } + if decision.action == "deny" and decision.reason: + # Comments are useful investigation context, but keep the + # webhook payload bounded. The audit sanitizer separately + # drops values containing recognizable PII or credentials. + data["denial_reason"] = decision.reason[:1024] + events.append( + AuditEventInput( + resource_type="agent_approval", + resource_id=approval.approval_id, + action=action, + data=data, + created_at=decided_at, + ) + ) + return tuple(events) + async def _submit_approval_update( self, *, @@ -2306,12 +2461,12 @@ async def _continue_with_approvals( # Idempotency path: resubmissions are normal for partial batches, so # reconcile against settled decisions instead of only pending ones. # A matching replay is a no-op; a contradicting one raises. - pending_tool_call_ids = await self._pending_approval_tool_call_ids(session_id) + pending_approvals = await self._pending_approvals(session_id) settled_decisions = await self._settled_approval_decisions(session_id) validated = self._validate_continuation_decisions( request=request, - pending_tool_call_ids=pending_tool_call_ids, + pending_tool_call_ids=set(pending_approvals), settled_decisions=settled_decisions, ) if not validated.approval_map: @@ -2417,6 +2572,31 @@ async def _continue_with_approvals( await dedup_client.delete(submission_key) raise + if _should_emit_approval_audit(role=self.role, source=source): + decided_at = datetime.now(tz=UTC) + audit_dedupe_id = self._approval_audit_dedupe_id( + workspace_id=self.workspace_id, + session_id=session_id, + run_id=curr_run_id, + stream_id=attempt.stream_id, + ) + audit_events = self._build_approval_audit_events( + pending_approvals=pending_approvals, + validated=validated, + request=request, + session_id=session_id, + run_id=curr_run_id, + dedupe_id=audit_dedupe_id, + decided_at=decided_at, + ) + if audit_events: + _schedule_approval_audit_events( + events=audit_events, + role=self.role, + request_audit=ctx_request_audit.get(), + dedupe_id=audit_dedupe_id, + ) + if not did_resume: await self._apply_submitted_approval_decisions( session_id=session_id, diff --git a/tracecat/audit/batch.py b/tracecat/audit/batch.py new file mode 100644 index 000000000..711bb2efe --- /dev/null +++ b/tracecat/audit/batch.py @@ -0,0 +1,99 @@ +"""Bounded, best-effort delivery for related audit events.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import Any + +import httpx + +from tracecat.audit.types import AuditAction, AuditResourceType +from tracecat.logger import logger + +MAX_PENDING_BATCH_DELIVERIES = 64 +MAX_CONCURRENT_BATCH_POSTS = 4 +BATCH_DELIVERY_DEADLINE_SECONDS = 10.0 + + +@dataclass(frozen=True, slots=True) +class AuditBatchEvent: + """One already-enriched event ready for webhook delivery.""" + + request_payload: dict[str, Any] + resource_type: AuditResourceType + action: AuditAction + + +@dataclass(frozen=True, slots=True) +class AuditBatchDelivery: + """A group of audit events sharing one webhook configuration.""" + + webhook_url: str + events: tuple[AuditBatchEvent, ...] + headers: dict[str, str] | None + verify_ssl: bool + + +# Strong refs to in-flight batches; done callbacks release them. +_batch_delivery_tasks: set[asyncio.Task[None]] = set() + + +def spawn_audit_batch(delivery: AuditBatchDelivery) -> None: + """Post a resolved audit batch on a bounded background task.""" + for stranded in [ + task for task in _batch_delivery_tasks if task.get_loop().is_closed() + ]: + _batch_delivery_tasks.discard(stranded) + if len(_batch_delivery_tasks) >= MAX_PENDING_BATCH_DELIVERIES: + logger.warning( + "Dropped audit webhook batch; pending limit reached", + event_count=len(delivery.events), + max_pending=MAX_PENDING_BATCH_DELIVERIES, + ) + return + task = asyncio.get_running_loop().create_task(deliver_audit_batch(delivery)) + _batch_delivery_tasks.add(task) + task.add_done_callback(_batch_delivery_tasks.discard) + + +async def deliver_audit_batch(delivery: AuditBatchDelivery) -> None: + """Deliver a batch with one client, bounded concurrency, and one deadline.""" + semaphore = asyncio.Semaphore(MAX_CONCURRENT_BATCH_POSTS) + + async def post_event(client: httpx.AsyncClient, event: AuditBatchEvent) -> None: + response: httpx.Response | None = None + try: + async with semaphore: + response = await client.post( + delivery.webhook_url, + json=event.request_payload, + headers=delivery.headers, + ) + response.raise_for_status() + except asyncio.CancelledError: + raise + except Exception as exc: + logger.warning( + "Failed to deliver audit webhook batch event", + error_type=type(exc).__name__, + status_code=response.status_code if response is not None else None, + resource_type=event.resource_type, + action=event.action, + ) + + try: + async with asyncio.timeout(BATCH_DELIVERY_DEADLINE_SECONDS): + async with httpx.AsyncClient( + timeout=BATCH_DELIVERY_DEADLINE_SECONDS, + verify=delivery.verify_ssl, + ) as client: + await asyncio.gather( + *(post_event(client, event) for event in delivery.events) + ) + except TimeoutError: + logger.warning( + "Audit webhook batch exceeded delivery deadline", + event_count=len(delivery.events), + deadline_seconds=BATCH_DELIVERY_DEADLINE_SECONDS, + ) diff --git a/tracecat/audit/sanitization.py b/tracecat/audit/sanitization.py index eaba06bdb..624b77a99 100644 --- a/tracecat/audit/sanitization.py +++ b/tracecat/audit/sanitization.py @@ -12,7 +12,12 @@ "auth_method", "changed_fields", "delete_mode", + "decision", + "decision_timestamp", + "denial_reason", "operation", + "source", + "tool_name", "trigger_type", "workflow_status", } @@ -80,6 +85,8 @@ def _is_allowed_metadata(key: str, value: AuditMetadataValue) -> bool: return isinstance(value, str) or value is None if key.startswith(("is_", "has_", "uses_")): return isinstance(value, bool) + if key == "arguments_overridden": + return isinstance(value, bool) if key.endswith("_count"): return isinstance(value, int) and not isinstance(value, bool) return False diff --git a/tracecat/audit/service.py b/tracecat/audit/service.py index 10c94622a..b2362e161 100644 --- a/tracecat/audit/service.py +++ b/tracecat/audit/service.py @@ -1,19 +1,20 @@ """Fire-and-forget audit webhook delivery with bounded in-memory retries. -Each event is posted by its own asyncio task; transient failures are retried -in process. A retry after a lost response can deliver an exact byte-identical -duplicate. Deliveries are dropped past the pending cap and lost on -process/loop shutdown. Durable, at-least-once delivery arrives with the -ENG-1514 spool. +Single events are posted by their own asyncio tasks with bounded in-process +retries. Explicit batches share one client and a total deadline. A retry after +a lost response can deliver an exact byte-identical duplicate. Deliveries are +dropped past their pending caps and lost on process/loop shutdown. Durable, +at-least-once delivery arrives with the ENG-1514 spool. """ from __future__ import annotations import asyncio import uuid -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass +from datetime import UTC, datetime from typing import Any, Self import httpx @@ -29,11 +30,13 @@ wait_exponential, ) +from tracecat.audit.batch import AuditBatchDelivery, AuditBatchEvent, spawn_audit_batch from tracecat.audit.enums import AuditEventActor, AuditEventStatus from tracecat.audit.sanitization import sanitize_audit_metadata from tracecat.audit.types import ( AuditAction, AuditEvent, + AuditEventInput, AuditMetadata, AuditMetadataValue, AuditResourceType, @@ -41,7 +44,7 @@ ) from tracecat.auth.secrets import get_db_encryption_key from tracecat.auth.types import PlatformRole, Role -from tracecat.contexts import ctx_request_audit, ctx_role +from tracecat.contexts import RequestAuditContext, ctx_request_audit, ctx_role from tracecat.db.engine import ( get_async_session_bypass_rls_context_manager, get_async_session_context_manager, @@ -428,6 +431,47 @@ async def _post_event(self, *, webhook_url: str, payload: AuditEvent) -> None: return _spawn_delivery(delivery) + async def _post_events( + self, *, webhook_url: str, payloads: list[AuditEvent] + ) -> None: + """Resolve shared settings once, then enqueue one bounded delivery batch.""" + try: + custom_headers = await self._get_custom_headers() + custom_payload = await self._get_custom_payload() + verify_ssl = await self._get_verify_ssl() + payload_attribute = await self._get_payload_attribute() + events: list[AuditBatchEvent] = [] + for payload in payloads: + event_payload = payload.model_dump(mode="json") + if custom_payload: + event_payload = {**event_payload, **custom_payload} + request_payload = ( + {payload_attribute: event_payload} + if payload_attribute + else event_payload + ) + events.append( + AuditBatchEvent( + request_payload=request_payload, + resource_type=payload.resource_type, + action=payload.action, + ) + ) + except Exception as exc: + self.logger.warning( + "Failed to resolve audit webhook batch", + error_type=type(exc).__name__, + ) + return + spawn_audit_batch( + AuditBatchDelivery( + webhook_url=webhook_url, + events=tuple(events), + headers=custom_headers, + verify_ssl=verify_ssl, + ) + ) + async def _get_actor_label(self) -> str | None: if self.role is None: return None @@ -476,6 +520,7 @@ def _build_payload( ip_address: str | None, user_agent: str | None, data: dict[str, AuditMetadataValue] | None, + created_at: datetime, ) -> AuditEvent: if self.role is None or self.role.actor_id is None: raise ValueError("Audit payload requires an auditable actor") @@ -502,8 +547,49 @@ def _build_payload( ip_address=ip_address, user_agent=user_agent, data=data, + created_at=created_at, ) + async def create_events( + self, + events: Sequence[AuditEventInput], + *, + request_audit: RequestAuditContext | None = None, + ) -> None: + """Enrich and enqueue multiple events with constant lookup overhead.""" + if not events: + return + if self.role is None or self.role.actor_id is None: + self.logger.debug( + "Skipping audit log batch", + reason="non_auditable_role", + role_type=self.role.type if self.role is not None else None, + ) + return + + webhook_url = await self._get_webhook_url() + if not webhook_url: + self.logger.debug("Skipping audit log batch", reason="webhook_unconfigured") + return + + actor_label = await self._get_actor_label() + payloads = [ + self._build_payload( + resource_type=event.resource_type, + action=event.action, + resource_id=event.resource_id, + status=event.status, + actor_label=actor_label, + ip_address=request_audit.client_ip if request_audit else None, + user_agent=request_audit.user_agent if request_audit else None, + data=sanitize_audit_metadata(event.data), + created_at=event.created_at, + ) + for event in events + ] + await self._post_events(webhook_url=webhook_url, payloads=payloads) + self.logger.debug("Streamed audit event batch", event_count=len(payloads)) + async def create_event( self, *, @@ -579,6 +665,7 @@ async def create_event( else None ), data=sanitize_audit_metadata(data), + created_at=datetime.now(UTC), ) await self._post_event(webhook_url=webhook_url, payload=payload) self.logger.debug( diff --git a/tracecat/audit/types.py b/tracecat/audit/types.py index 6b3178aae..48ff257a1 100644 --- a/tracecat/audit/types.py +++ b/tracecat/audit/types.py @@ -2,6 +2,7 @@ import uuid from collections.abc import Mapping +from dataclasses import dataclass, field from datetime import UTC, datetime from typing import Literal @@ -23,6 +24,7 @@ "reset", "rotate", "accept", + "reject", "revoke", "sign_in", "sync", @@ -54,6 +56,7 @@ "agent_model_access", "agent_preset", "agent_session", + "agent_approval", "organization_domain", "organization_member", "organization_session", @@ -79,6 +82,18 @@ ] +@dataclass(frozen=True, slots=True) +class AuditEventInput: + """One event in a bulk audit delivery request.""" + + resource_type: AuditResourceType + action: AuditAction + resource_id: uuid.UUID | None = None + status: AuditEventStatus = AuditEventStatus.SUCCESS + data: AuditMetadata | None = None + created_at: datetime = field(default_factory=lambda: datetime.now(UTC)) + + class AuditEvent(BaseModel): """A privacy-bounded record of an actor acting on a resource.