feat(temporal): encrypt durable agent histories

Author: daryllimyt

packages/tracecat-ee/tracecat_ee/agent/approvals/service.py

@@ -33,6 +33,7 @@
 from tracecat.identifiers.workflow import exec_id_to_parts
 from tracecat.logger import logger
 from tracecat.service import BaseWorkspaceService
+from tracecat.temporal.visibility import tokenize_visibility_value
 from tracecat.workflow.executions.enums import TemporalSearchAttr
 from tracecat_ee.agent.activities import (
     ApplyApprovalResultsActivityInputs,
@@ -328,7 +329,9 @@ async def list_session_history(
                 session_id=session_id,
             )
             return []
-        alias = build_agent_alias(session.parent_workflow_id, session.action_ref)
+        alias = tokenize_visibility_value(
+            build_agent_alias(session.parent_workflow_id, session.action_ref)
+        )
 
         query = " AND ".join(
             [

tests/unit/test_approvals_service.py

@@ -1,16 +1,22 @@
 """Tests for ApprovalService CRUD operations."""
 
 import uuid
+from collections.abc import AsyncIterator
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, patch
 
 import pytest
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession
 from tracecat_ee.agent.approvals.schemas import ApprovalCreate, ApprovalUpdate
 from tracecat_ee.agent.approvals.service import ApprovalService, SessionInfo
 
+from tracecat.agent.aliases import build_agent_alias
 from tracecat.agent.approvals.enums import ApprovalStatus
 from tracecat.auth.types import Role
 from tracecat.db.models import AgentSession, User
+from tracecat.identifiers import WorkflowID
+from tracecat.temporal.visibility import tokenize_visibility_value
 
 pytestmark = pytest.mark.usefixtures("db")
 
@@ -264,6 +270,47 @@ async def test_update_approvals_empty_dict(
         updated = await approvals_service.update_approvals({})
         assert updated == []
 
+    async def test_list_session_history_uses_tokenized_alias_query(
+        self,
+        approvals_service: ApprovalService,
+    ) -> None:
+        session_id = uuid.uuid4()
+        parent_workflow_id = WorkflowID(str(uuid.uuid4()))
+        action_ref = "tools.some_action"
+        expected_alias = tokenize_visibility_value(
+            build_agent_alias(parent_workflow_id, action_ref)
+        )
+        captured_query: str | None = None
+
+        class FakeTemporalClient:
+            async def list_workflows(self, **kwargs: str) -> AsyncIterator[object]:
+                nonlocal captured_query
+                captured_query = kwargs["query"]
+                if False:
+                    yield object()
+
+        with (
+            patch.object(
+                approvals_service,
+                "get_session",
+                AsyncMock(
+                    return_value=SimpleNamespace(
+                        parent_workflow_id=parent_workflow_id,
+                        action_ref=action_ref,
+                    )
+                ),
+            ),
+            patch(
+                "tracecat_ee.agent.approvals.service.get_temporal_client",
+                AsyncMock(return_value=FakeTemporalClient()),
+            ),
+        ):
+            history = await approvals_service.list_session_history(session_id)
+
+        assert history == []
+        assert captured_query is not None
+        assert expected_alias in captured_query
+
     async def test_delete_approval(
         self,
         approvals_service: ApprovalService,

tracecat/agent/session/service.py

@@ -55,6 +55,7 @@
 )
 from tracecat.chat.service import ChatService
 from tracecat.chat.tools import get_default_tools
+from tracecat.contexts import with_temporal_workspace_id
 from tracecat.db.models import AgentSession, AgentSessionHistory, Approval, Chat
 from tracecat.dsl.client import get_temporal_client
 from tracecat.dsl.common import RETRY_POLICIES
@@ -891,15 +892,16 @@ async def run_turn(
                 task_queue=config.TRACECAT__AGENT_QUEUE,
             )
 
-            await client.start_workflow(
-                DurableAgentWorkflow.run,
-                workflow_args,
-                id=str(workflow_id),
-                task_queue=config.TRACECAT__AGENT_QUEUE,
-                execution_timeout=timedelta(hours=1),
-                retry_policy=RETRY_POLICIES["workflow:fail_fast"],
-                search_attributes=self._build_direct_agent_search_attributes(),
-            )
+            with with_temporal_workspace_id(self.role.workspace_id):
+                await client.start_workflow(
+                    DurableAgentWorkflow.run,
+                    workflow_args,
+                    id=str(workflow_id),
+                    task_queue=config.TRACECAT__AGENT_QUEUE,
+                    execution_timeout=timedelta(hours=1),
+                    retry_policy=RETRY_POLICIES["workflow:fail_fast"],
+                    search_attributes=self._build_direct_agent_search_attributes(),
+                )
 
         # Return ChatResponse with session_id for streaming
         stream_url = f"/api/agent/sessions/{session_id}/stream"
@@ -1029,14 +1031,15 @@ async def _continue_with_approvals(
         )
 
         try:
-            await handle.execute_update(
-                DurableAgentWorkflow.set_approvals,
-                WorkflowApprovalSubmission(
-                    approvals=approval_map,
-                    approved_by=self.role.user_id,
-                    decision_metadata=decision_metadata or None,
-                ),
-            )
+            with with_temporal_workspace_id(self.role.workspace_id):
+                await handle.execute_update(
+                    DurableAgentWorkflow.set_approvals,
+                    WorkflowApprovalSubmission(
+                        approvals=approval_map,
+                        approved_by=self.role.user_id,
+                        decision_metadata=decision_metadata or None,
+                    ),
+                )
         except Exception:
             # Allow retriable failures to be resubmitted by clearing dedup marker.
             if dedup_client is not None and dedup_key is not None: