diff --git a/.github/workflows/backend-tests.yml b/.github/workflows/backend-tests.yml index b58d64330..0ee1989e9 100644 --- a/.github/workflows/backend-tests.yml +++ b/.github/workflows/backend-tests.yml @@ -64,6 +64,8 @@ jobs: if: github.event.pull_request.draft == false env: EMBEDDING_MODEL: sentence-transformers/all-MiniLM-L6-v2 + REDIS_APP_URL: redis://localhost:6379/0 + SANDBOX_IMAGE: surfsense/sandbox:dev services: postgres: @@ -79,6 +81,15 @@ jobs: --health-interval 10s --health-timeout 5s --health-retries 5 + redis: + image: redis:7-alpine + ports: + - 6379:6379 + options: >- + --health-cmd "redis-cli ping" + --health-interval 5s + --health-timeout 3s + --health-retries 10 steps: - name: Checkout code diff --git a/surfsense_backend/alembic/versions/179_add_video_presentation_artifact_id.py b/surfsense_backend/alembic/versions/179_add_video_presentation_artifact_id.py index f862aeda9..0b08de32a 100644 --- a/surfsense_backend/alembic/versions/179_add_video_presentation_artifact_id.py +++ b/surfsense_backend/alembic/versions/179_add_video_presentation_artifact_id.py @@ -27,6 +27,4 @@ def upgrade() -> None: def downgrade() -> None: - op.execute( - "ALTER TABLE video_presentations DROP COLUMN IF EXISTS artifact_id" - ) + op.execute("ALTER TABLE video_presentations DROP COLUMN IF EXISTS artifact_id") diff --git a/surfsense_backend/alembic/versions/180_slim_video_presentation_runs.py b/surfsense_backend/alembic/versions/180_slim_video_presentation_runs.py index 1cb8e54e7..7798696b0 100644 --- a/surfsense_backend/alembic/versions/180_slim_video_presentation_runs.py +++ b/surfsense_backend/alembic/versions/180_slim_video_presentation_runs.py @@ -77,7 +77,9 @@ def downgrade() -> None: """ ) op.execute("ALTER TABLE video_presentation_runs DROP COLUMN IF EXISTS error") - op.execute("ALTER TABLE video_presentation_runs ADD COLUMN IF NOT EXISTS slides JSONB") + op.execute( + "ALTER TABLE video_presentation_runs ADD COLUMN IF NOT EXISTS slides JSONB" + ) op.execute( "ALTER TABLE video_presentation_runs ADD COLUMN IF NOT EXISTS scene_codes JSONB" ) diff --git a/surfsense_backend/alembic/versions/181_drop_image_generations.py b/surfsense_backend/alembic/versions/181_drop_image_generations.py index 487301d60..a86f87bac 100644 --- a/surfsense_backend/alembic/versions/181_drop_image_generations.py +++ b/surfsense_backend/alembic/versions/181_drop_image_generations.py @@ -15,9 +15,10 @@ Revises: 180 from collections.abc import Sequence -from alembic import op from sqlalchemy import text +from alembic import op + revision: str = "181" down_revision: str | None = "180" branch_labels: str | Sequence[str] | None = None diff --git a/surfsense_backend/alembic/versions/183_rename_podcasts_to_podcast_runs.py b/surfsense_backend/alembic/versions/183_rename_podcasts_to_podcast_runs.py index 326d0fb8c..242731155 100644 --- a/surfsense_backend/alembic/versions/183_rename_podcasts_to_podcast_runs.py +++ b/surfsense_backend/alembic/versions/183_rename_podcasts_to_podcast_runs.py @@ -66,7 +66,9 @@ def downgrade() -> None: END $$; """ ) - op.execute("ALTER TABLE podcast_runs ADD COLUMN IF NOT EXISTS storage_backend VARCHAR(32)") + op.execute( + "ALTER TABLE podcast_runs ADD COLUMN IF NOT EXISTS storage_backend VARCHAR(32)" + ) op.execute("ALTER TABLE podcast_runs ADD COLUMN IF NOT EXISTS storage_key TEXT") op.execute("ALTER TABLE podcast_runs ADD COLUMN IF NOT EXISTS file_location TEXT") op.execute("ALTER TABLE podcast_runs RENAME TO podcasts") diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/tools/automation/create.py b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/tools/automation/create.py index c00692185..502da73fc 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/tools/automation/create.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/tools/automation/create.py @@ -33,6 +33,7 @@ from app.agents.chat.multi_agent_chat.subagents.shared.hitl.approvals.self_gated from app.auth.context import AuthContext from app.automations.schemas.api import AutomationCreate from app.automations.services.automation import AutomationService +from app.capabilities.core import ActivityDescriptor from app.db import async_session_maker from app.utils.content_utils import extract_text_content @@ -194,6 +195,15 @@ def create_create_automation_tool( logger.exception("create_automation failed") return {"status": "error", "message": f"persistence failed: {exc}"} + create_automation.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Creating automation", + completed_title="Created automation", + category="action", + icon_key="workflow", + kind="create_automation", + ).as_metadata() + } return create_automation diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/tools/update_memory.py b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/tools/update_memory.py index 333d10e4b..803a5a1c3 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/tools/update_memory.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/tools/update_memory.py @@ -9,6 +9,7 @@ from uuid import UUID from langchain_core.tools import tool from sqlalchemy.ext.asyncio import AsyncSession +from app.capabilities.core import ActivityDescriptor from app.db import async_session_maker from app.services.memory import MemoryScope, save_memory @@ -49,6 +50,15 @@ def create_update_memory_tool( logger.exception("Failed to update user memory: %s", e) return {"status": "error", "message": f"Failed to update memory: {e}"} + update_memory.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Updating your memory", + completed_title="Updated your memory", + category="action", + icon_key="brain", + kind="memory.personal", + ).as_metadata() + } return update_memory @@ -84,6 +94,15 @@ def create_update_team_memory_tool( "message": f"Failed to update team memory: {e}", } + update_memory.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Updating team memory", + completed_title="Updated team memory", + category="action", + icon_key="brain", + kind="memory.team", + ).as_metadata() + } return update_memory diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/middleware/middleware.py b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/middleware/middleware.py index e05749e75..99f06aea8 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/middleware/middleware.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/middleware/middleware.py @@ -12,6 +12,7 @@ from app.agents.chat.multi_agent_chat.shared.filesystem_selection import Filesys from app.agents.chat.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) +from app.capabilities.core import ActivityDescriptor from app.sandbox import is_sandbox_enabled from app.utils.perf import get_perf_logger @@ -122,9 +123,27 @@ class SurfSenseFilesystemMiddleware(FilesystemMiddleware): def _create_glob_tool(self) -> BaseTool: tool = super()._create_glob_tool() tool.description = glob_description(self._filesystem_mode).rstrip() + tool.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Finding files", + completed_title="Found files", + category="file", + icon_key="folder-search", + kind="glob", + ).as_metadata() + } return tool def _create_grep_tool(self) -> BaseTool: tool = super()._create_grep_tool() tool.description = grep_description(self._filesystem_mode).rstrip() + tool.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Searching project", + completed_title="Searched project", + category="file", + icon_key="search-code", + kind="grep", + ).as_metadata() + } return tool diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/edit_file/index.py b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/edit_file/index.py index 6ed8a03d7..92df9edc2 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/edit_file/index.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/edit_file/index.py @@ -20,6 +20,7 @@ from app.agents.chat.multi_agent_chat.shared.middleware.filesystem.backends.kb_p from app.agents.chat.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) +from app.capabilities.core import ActivityDescriptor from ...middleware.async_dispatch import run_async_blocking from ...middleware.mode import is_cloud @@ -137,4 +138,13 @@ def create_edit_file_tool(mw: SurfSenseFilesystemMiddleware) -> BaseTool: description=description, func=sync_edit_file, coroutine=async_edit_file, + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Editing file", + completed_title="Edited file", + category="file", + icon_key="file-pen", + kind="edit_file", + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/execute_code/index.py b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/execute_code/index.py index b530c91f2..ad2e9780c 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/execute_code/index.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/execute_code/index.py @@ -10,6 +10,7 @@ from langchain_core.tools import BaseTool, StructuredTool from app.agents.chat.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) +from app.capabilities.core import ActivityDescriptor from ...middleware.async_dispatch import run_async_blocking from .description import select_description @@ -57,4 +58,13 @@ def create_execute_code_tool(mw: SurfSenseFilesystemMiddleware) -> BaseTool: description=description, func=sync_execute_code, coroutine=async_execute_code, + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Running code", + completed_title="Ran code", + category="action", + icon_key="square-code", + kind="execute_code", + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/list_tree/index.py b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/list_tree/index.py index 21bba1fc3..0bb084626 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/list_tree/index.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/list_tree/index.py @@ -15,6 +15,7 @@ from app.agents.chat.multi_agent_chat.shared.middleware.filesystem.backends.kb_p from app.agents.chat.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) +from app.capabilities.core import ActivityDescriptor from ...middleware.async_dispatch import run_async_blocking from ...middleware.path_resolution import resolve_list_target_path @@ -102,4 +103,13 @@ def create_list_tree_tool(mw: SurfSenseFilesystemMiddleware) -> BaseTool: description=description, func=sync_list_tree, coroutine=async_list_tree, + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Reviewing file tree", + completed_title="Reviewed file tree", + category="file", + icon_key="folder-tree", + kind="list_tree", + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/ls/index.py b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/ls/index.py index e45a279d7..9b5f42830 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/ls/index.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/ls/index.py @@ -14,6 +14,7 @@ from app.agents.chat.multi_agent_chat.shared.middleware.filesystem.backends.kb_p from app.agents.chat.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) +from app.capabilities.core import ActivityDescriptor from ...middleware.async_dispatch import run_async_blocking from ...middleware.path_resolution import resolve_list_target_path @@ -97,4 +98,13 @@ def create_ls_tool(mw: SurfSenseFilesystemMiddleware) -> BaseTool: description=description, func=sync_ls, coroutine=async_ls, + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Reviewing folder", + completed_title="Reviewed folder", + category="file", + icon_key="folder-open", + kind="ls", + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/mkdir/index.py b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/mkdir/index.py index 9f5456cd6..4d07b107e 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/mkdir/index.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/mkdir/index.py @@ -17,6 +17,7 @@ from app.agents.chat.multi_agent_chat.shared.middleware.filesystem.backends.git_ from app.agents.chat.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) +from app.capabilities.core import ActivityDescriptor from app.knowledge_store.paths import DOCUMENTS_ROOT from ...middleware.async_dispatch import run_async_blocking @@ -100,4 +101,13 @@ def create_mkdir_tool(mw: SurfSenseFilesystemMiddleware) -> BaseTool: description=description, func=sync_mkdir, coroutine=async_mkdir, + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Creating folder", + completed_title="Created folder", + category="file", + icon_key="folder-plus", + kind="mkdir", + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/move_file/index.py b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/move_file/index.py index 7b3ff56e4..75221c58a 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/move_file/index.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/move_file/index.py @@ -17,6 +17,7 @@ from app.agents.chat.multi_agent_chat.shared.middleware.filesystem.backends.git_ from app.agents.chat.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) +from app.capabilities.core import ActivityDescriptor from ...middleware.async_dispatch import run_async_blocking from ...middleware.mode import is_cloud @@ -99,4 +100,13 @@ def create_move_file_tool(mw: SurfSenseFilesystemMiddleware) -> BaseTool: description=description, func=sync_move_file, coroutine=async_move_file, + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Moving file", + completed_title="Moved file", + category="file", + icon_key="files", + kind="move_file", + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/read_file/index.py b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/read_file/index.py index b085671e3..607578e79 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/read_file/index.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/read_file/index.py @@ -22,6 +22,7 @@ from app.agents.chat.multi_agent_chat.shared.middleware.filesystem.backends.kb_p from app.agents.chat.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) +from app.capabilities.core import ActivityDescriptor from ...middleware.async_dispatch import run_async_blocking from ...middleware.path_resolution import resolve_relative @@ -112,4 +113,13 @@ def create_read_file_tool(mw: SurfSenseFilesystemMiddleware) -> BaseTool: description=description, func=sync_read_file, coroutine=async_read_file, + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Reading file", + completed_title="Read file", + category="file", + icon_key="file-text", + kind="read_file", + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/rm/index.py b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/rm/index.py index facf15725..52d1f105b 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/rm/index.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/rm/index.py @@ -15,6 +15,7 @@ from app.agents.chat.multi_agent_chat.shared.middleware.filesystem.backends.git_ from app.agents.chat.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) +from app.capabilities.core import ActivityDescriptor from ...middleware.async_dispatch import run_async_blocking from ...middleware.mode import is_cloud @@ -66,4 +67,13 @@ def create_rm_tool(mw: SurfSenseFilesystemMiddleware) -> BaseTool: description=description, func=sync_rm, coroutine=async_rm, + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Deleting file", + completed_title="Deleted file", + category="file", + icon_key="file-x", + kind="rm", + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/rmdir/index.py b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/rmdir/index.py index 05b20f184..7f21b4f7a 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/rmdir/index.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/rmdir/index.py @@ -15,6 +15,7 @@ from app.agents.chat.multi_agent_chat.shared.middleware.filesystem.backends.git_ from app.agents.chat.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) +from app.capabilities.core import ActivityDescriptor from ...middleware.async_dispatch import run_async_blocking from ...middleware.mode import is_cloud @@ -66,4 +67,13 @@ def create_rmdir_tool(mw: SurfSenseFilesystemMiddleware) -> BaseTool: description=description, func=sync_rmdir, coroutine=async_rmdir, + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Deleting folder", + completed_title="Deleted folder", + category="file", + icon_key="folder-x", + kind="rmdir", + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/write_file/index.py b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/write_file/index.py index 8415aefab..a55ed43d8 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/write_file/index.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/write_file/index.py @@ -17,6 +17,7 @@ from app.agents.chat.multi_agent_chat.shared.middleware.filesystem.backends.git_ from app.agents.chat.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) +from app.capabilities.core import ActivityDescriptor from ...middleware.async_dispatch import run_async_blocking from ...middleware.mode import is_cloud @@ -86,4 +87,13 @@ def create_write_file_tool(mw: SurfSenseFilesystemMiddleware) -> BaseTool: description=description, func=sync_write_file, coroutine=async_write_file, + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Creating file", + completed_title="Created file", + category="file", + icon_key="file-plus", + kind="write_file", + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/todos.py b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/todos.py index 0316d6e2d..0f25938a7 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/todos.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/shared/middleware/todos.py @@ -6,6 +6,8 @@ from typing import TYPE_CHECKING, Any from langchain.agents.middleware import TodoListMiddleware +from app.capabilities.core import ActivityDescriptor + if TYPE_CHECKING: from collections.abc import Awaitable, Callable @@ -43,7 +45,20 @@ def build_todos_mw(*, system_prompt: str | None = None) -> TodoListMiddleware: - otherwise: append the given custom todo system prompt. """ if system_prompt is None: - return TodoListMiddleware() - if not system_prompt.strip(): - return _ToolOnlyTodoListMiddleware() - return TodoListMiddleware(system_prompt=system_prompt) + middleware = TodoListMiddleware() + elif not system_prompt.strip(): + middleware = _ToolOnlyTodoListMiddleware() + else: + middleware = TodoListMiddleware(system_prompt=system_prompt) + descriptor = ActivityDescriptor( + active_title="Planning work", + completed_title="Planned work", + category="action", + icon_key="list-todo", + kind="write_todos", + lifecycle="phase", + ).as_metadata() + for tool in middleware.tools: + if tool.name == "write_todos": + tool.metadata = {"activity_descriptor": descriptor} + return middleware diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/generate_image.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/generate_image.py index 975734a3d..85086d40e 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/generate_image.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/generate_image.py @@ -18,6 +18,7 @@ from app.agents.chat.multi_agent_chat.subagents.builtins.deliverables.tools.thre resolve_root_thread_id, ) from app.artifacts.media.image.record import record as record_image +from app.capabilities.core import ActivityDescriptor from app.db import ( Model, Workspace, @@ -281,4 +282,13 @@ def create_generate_image_tool( error=err, ) + generate_image.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Creating an image", + completed_title="Created an image", + category="artifact", + icon_key="image", + kind="generate_image", + ).as_metadata() + } return generate_image diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/load_artifact_source.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/load_artifact_source.py index 212a1971c..c5c80668c 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/load_artifact_source.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/load_artifact_source.py @@ -9,6 +9,7 @@ from langchain_core.tools import BaseTool, tool from sqlalchemy import select from app.artifacts.persistence import Artifact, ArtifactFile, ArtifactFileRole +from app.capabilities.core import ActivityDescriptor from app.config import config as app_config from app.db import shielded_async_session from app.file_storage.factory import get_storage_backend @@ -82,4 +83,14 @@ def create_load_artifact_source_tool(*, workspace_id: int) -> BaseTool: ), } + load_artifact_source.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Opening the artifact", + completed_title="Opened the artifact", + category="artifact", + icon_key="file-input", + kind="load_artifact_source", + lifecycle="phase", + ).as_metadata() + } return load_artifact_source diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/podcast.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/podcast.py index ad661091f..c8ab75f96 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/podcast.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/podcast.py @@ -20,6 +20,7 @@ from app.agents.chat.multi_agent_chat.shared.receipts.receipt import make_receip from app.agents.chat.multi_agent_chat.subagents.builtins.deliverables.tools.thread_resolver import ( resolve_root_thread_id, ) +from app.capabilities.core import ActivityDescriptor from app.db import PodcastStatus, shielded_async_session from app.podcasts.generation.brief import propose_brief from app.podcasts.service import PodcastService @@ -142,4 +143,13 @@ def create_generate_podcast_tool( tool_call_id=runtime.tool_call_id, ) + generate_podcast.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Creating the podcast", + completed_title="Created the podcast", + category="artifact", + icon_key="microphone", + kind="generate_podcast", + ).as_metadata() + } return generate_podcast diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/sandbox.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/sandbox.py index 4c3b6cc92..f667e104e 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/sandbox.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/sandbox.py @@ -9,6 +9,7 @@ from typing import Literal from langchain.tools import ToolRuntime from langchain_core.tools import BaseTool, tool +from app.capabilities.core import ActivityDescriptor from app.config import config as app_config from app.sandbox import SandboxSession, get_registry @@ -103,4 +104,23 @@ def create_sandbox_tools(*, workspace_id: int) -> list[BaseTool]: "binary files" ) from exc + execute.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Running command", + completed_title="Ran command", + category="action", + icon_key="terminal", + kind="execute", + ).as_metadata() + } + read_sandbox_file.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Reviewing the artifact", + completed_title="Reviewed the artifact", + category="artifact", + icon_key="file-text", + kind="read_sandbox_file", + lifecycle="phase", + ).as_metadata() + } return [execute, load_artifact_instructions, read_sandbox_file] diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/save_artifact.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/save_artifact.py index 49b5ab02c..2bc0b2566 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/save_artifact.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/save_artifact.py @@ -15,6 +15,7 @@ from app.artifacts import ArtifactFileInput, save_artifact from app.artifacts.source_formats import validate_source_file from app.artifacts.verification.formats.registry import get_format_adapter from app.artifacts.verification.receipt import read_receipt, sha256_bytes +from app.capabilities.core import ActivityDescriptor from app.config import config as app_config from app.db import shielded_async_session from app.sandbox import SandboxSession, get_registry @@ -189,4 +190,13 @@ def create_save_artifact_tool(workspace_id: int): # Keep the public tool name frozen even though the Python symbol avoids # shadowing the service imported above. save_artifact_tool.name = "save_artifact" + save_artifact_tool.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Preparing the file", + completed_title="Presented file", + category="artifact", + icon_key="file-output", + kind="save_artifact", + ).as_metadata() + } return save_artifact_tool diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/verify_artifact.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/verify_artifact.py index 761596b8c..5dfe5b83b 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/verify_artifact.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/verify_artifact.py @@ -6,6 +6,7 @@ from langchain.tools import ToolRuntime from langchain_core.tools import BaseTool, tool from app.artifacts.verification.service import verify_artifact as verify +from app.capabilities.core import ActivityDescriptor from app.db import shielded_async_session from app.sandbox import get_registry from app.services.llm_service import get_vision_llm @@ -50,4 +51,13 @@ def create_verify_artifact_tool(*, workspace_id: int) -> BaseTool: "verification_unavailable": result.unavailable_reason, } + verify_artifact.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Checking the artifact", + completed_title="Checked the artifact", + category="artifact", + icon_key="badge-check", + kind="verify_artifact", + ).as_metadata() + } return verify_artifact diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/video_presentation.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/video_presentation.py index 7558e5a53..2b300aaf7 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/video_presentation.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/video_presentation.py @@ -25,6 +25,7 @@ from app.agents.chat.multi_agent_chat.subagents.builtins.deliverables.deliverabl from app.agents.chat.multi_agent_chat.subagents.builtins.deliverables.tools.thread_resolver import ( resolve_root_thread_id, ) +from app.capabilities.core import ActivityDescriptor from app.db import ( VideoPresentationRun, VideoPresentationStatus, @@ -182,4 +183,13 @@ def create_generate_video_presentation_tool( tool_call_id=runtime.tool_call_id, ) + generate_video_presentation.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Creating the presentation", + completed_title="Created the presentation", + category="artifact", + icon_key="film", + kind="generate_video_presentation", + ).as_metadata() + } return generate_video_presentation diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/knowledge_base/ask_knowledge_base_tool.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/knowledge_base/ask_knowledge_base_tool.py index d2327471e..98290ca09 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/knowledge_base/ask_knowledge_base_tool.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/knowledge_base/ask_knowledge_base_tool.py @@ -15,6 +15,7 @@ from app.agents.chat.multi_agent_chat.subagents.shared.invocation import ( EXCLUDED_STATE_KEYS, subagent_invoke_config, ) +from app.capabilities.core import ActivityDescriptor from .prompts import load_readonly_description @@ -115,4 +116,13 @@ def build_ask_knowledge_base_tool( func=ask_knowledge_base, coroutine=aask_knowledge_base, description=load_readonly_description(), + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Reviewing your sources", + completed_title="Reviewed your sources", + category="research", + icon_key="library", + kind=TOOL_NAME, + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/knowledge_base/tools/search_knowledge_base.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/knowledge_base/tools/search_knowledge_base.py index 8b1e51833..7a51d40fc 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/knowledge_base/tools/search_knowledge_base.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/knowledge_base/tools/search_knowledge_base.py @@ -29,6 +29,7 @@ from app.agents.chat.multi_agent_chat.shared.state.filesystem_state import ( SurfSenseFilesystemState, ) from app.agents.chat.runtime.references import referenced_document_ids +from app.capabilities.core import ActivityDescriptor from app.db import shielded_async_session from app.utils.perf import get_perf_logger @@ -181,4 +182,13 @@ def create_search_knowledge_base_tool( name="search_knowledge_base", description=_TOOL_DESCRIPTION, coroutine=_impl, + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Searching your sources", + completed_title="Searched your sources", + category="research", + icon_key="library", + kind="search_knowledge_base", + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/create_event.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/create_event.py index a4286e1cb..0d1b14e47 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/create_event.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/create_event.py @@ -11,6 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from app.agents.chat.multi_agent_chat.subagents.shared.hitl.approvals.self_gated import ( request_approval, ) +from app.capabilities.core import ActivityDescriptor from app.services.google_calendar import GoogleCalendarToolMetadataService logger = logging.getLogger(__name__) @@ -346,4 +347,14 @@ def create_create_calendar_event_tool( "message": "Something went wrong while creating the event. Please try again.", } + create_calendar_event.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Creating calendar event", + completed_title="Created calendar event", + category="connector", + icon_key="calendar", + integration_key="google_calendar", + kind="create_calendar_event", + ).as_metadata() + } return create_calendar_event diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/delete_event.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/delete_event.py index ac5e2374b..3e6d5bb36 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/delete_event.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/delete_event.py @@ -11,6 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from app.agents.chat.multi_agent_chat.subagents.shared.hitl.approvals.self_gated import ( request_approval, ) +from app.capabilities.core import ActivityDescriptor from app.services.google_calendar import GoogleCalendarToolMetadataService logger = logging.getLogger(__name__) @@ -309,4 +310,14 @@ def create_delete_calendar_event_tool( "message": "Something went wrong while deleting the event. Please try again.", } + delete_calendar_event.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Deleting calendar event", + completed_title="Deleted calendar event", + category="connector", + icon_key="calendar", + integration_key="google_calendar", + kind="delete_calendar_event", + ).as_metadata() + } return delete_calendar_event diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/search_events.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/search_events.py index 7439b2694..69ddd9f79 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/search_events.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/search_events.py @@ -8,6 +8,7 @@ from sqlalchemy.future import select from app.agents.chat.multi_agent_chat.subagents.connectors.google_auth import ( build_credentials as _build_credentials, ) +from app.capabilities.core import ActivityDescriptor from app.db import SearchSourceConnector, SearchSourceConnectorType logger = logging.getLogger(__name__) @@ -164,4 +165,14 @@ def create_search_calendar_events_tool( "message": "Failed to search calendar events. Please try again.", } + search_calendar_events.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Searching calendar", + completed_title="Searched calendar", + category="connector", + icon_key="calendar", + integration_key="google_calendar", + kind="search_calendar_events", + ).as_metadata() + } return search_calendar_events diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/update_event.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/update_event.py index 0ef5be996..1d2de7f3d 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/update_event.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/update_event.py @@ -11,6 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from app.agents.chat.multi_agent_chat.subagents.shared.hitl.approvals.self_gated import ( request_approval, ) +from app.capabilities.core import ActivityDescriptor from app.services.google_calendar import GoogleCalendarToolMetadataService logger = logging.getLogger(__name__) @@ -395,4 +396,14 @@ def create_update_calendar_event_tool( "message": "Something went wrong while updating the event. Please try again.", } + update_calendar_event.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Updating calendar event", + completed_title="Updated calendar event", + category="connector", + icon_key="calendar", + integration_key="google_calendar", + kind="update_calendar_event", + ).as_metadata() + } return update_calendar_event diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/get_connected_accounts.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/get_connected_accounts.py index 292530229..91fef95d2 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/get_connected_accounts.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/get_connected_accounts.py @@ -15,6 +15,7 @@ import logging from langchain_core.tools import BaseTool, StructuredTool from sqlalchemy import select +from app.capabilities.core import ActivityDescriptor from app.services.mcp_oauth.registry import get_service_by_connector_type logger = logging.getLogger(__name__) @@ -96,4 +97,14 @@ def create_get_connected_accounts_tool(*, workspace_id: int) -> BaseTool: name="get_connected_accounts", description=_TOOL_DESCRIPTION, coroutine=_impl, + metadata={ + "activity_descriptor": ActivityDescriptor( + active_title="Checking connected apps", + completed_title="Checked connected apps", + category="connector", + icon_key="search", + kind="get_connected_accounts", + lifecycle="phase", + ).as_metadata() + }, ) diff --git a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/memory/tools/update_memory.py b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/memory/tools/update_memory.py index 3361f13ae..017e28404 100644 --- a/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/memory/tools/update_memory.py +++ b/surfsense_backend/app/agents/chat/multi_agent_chat/subagents/builtins/memory/tools/update_memory.py @@ -9,6 +9,7 @@ from uuid import UUID from langchain_core.tools import tool from sqlalchemy.ext.asyncio import AsyncSession +from app.capabilities.core import ActivityDescriptor from app.services.memory import ( MEMORY_HARD_LIMIT, MEMORY_SOFT_LIMIT, @@ -47,6 +48,15 @@ def create_update_memory_tool( await db_session.rollback() return {"status": "error", "message": f"Failed to update memory: {e}"} + update_memory.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Updating your memory", + completed_title="Updated your memory", + category="action", + icon_key="brain", + kind="memory.personal", + ).as_metadata() + } return update_memory @@ -79,6 +89,15 @@ def create_update_team_memory_tool( "message": f"Failed to update team memory: {e}", } + update_memory.metadata = { + "activity_descriptor": ActivityDescriptor( + active_title="Updating team memory", + completed_title="Updated team memory", + category="action", + icon_key="brain", + kind="memory.team", + ).as_metadata() + } return update_memory diff --git a/surfsense_backend/app/artifacts/media/naming.py b/surfsense_backend/app/artifacts/media/naming.py index b89d6129f..893b5f1f0 100644 --- a/surfsense_backend/app/artifacts/media/naming.py +++ b/surfsense_backend/app/artifacts/media/naming.py @@ -4,7 +4,7 @@ from __future__ import annotations import re -from app.knowledge_store.paths.naming import normalize_filename +from app.knowledge_store.paths import normalize_filename _DOT_RUNS = re.compile(r"\.+") diff --git a/surfsense_backend/app/capabilities/core/types.py b/surfsense_backend/app/capabilities/core/types.py index 55386ce9f..7e83314e5 100644 --- a/surfsense_backend/app/capabilities/core/types.py +++ b/surfsense_backend/app/capabilities/core/types.py @@ -75,6 +75,9 @@ class ActivityDescriptor: category: ActivityCategory icon_key: str integration_key: str | None = None + kind: str | None = None + lifecycle: Literal["invocation", "phase"] = "invocation" + visibility: Literal["show", "hide"] = "show" def as_metadata(self, *, kind: str | None = None) -> dict[str, str]: metadata = { @@ -83,8 +86,12 @@ class ActivityDescriptor: "category": self.category, "icon_key": self.icon_key, } - if kind: - metadata["kind"] = kind + if resolved_kind := kind or self.kind: + metadata["kind"] = resolved_kind + if self.lifecycle != "invocation": + metadata["lifecycle"] = self.lifecycle + if self.visibility != "show": + metadata["visibility"] = self.visibility if self.integration_key: metadata["integration_key"] = self.integration_key return metadata @@ -99,6 +106,9 @@ class ActivityDescriptor: category = value.get("category") icon_key = value.get("icon_key") integration_key = value.get("integration_key") + kind = value.get("kind") + lifecycle = value.get("lifecycle", "invocation") + visibility = value.get("visibility", "show") if not ( isinstance(active, str) and 0 < len(active.strip()) <= 120 @@ -114,6 +124,15 @@ class ActivityDescriptor: and _ACTIVITY_KEY_RE.fullmatch(integration_key.strip()) ) ) + and ( + kind is None + or ( + isinstance(kind, str) + and re.fullmatch(r"[a-z0-9][a-z0-9._-]{0,127}", kind.strip()) + ) + ) + and lifecycle in {"invocation", "phase"} + and visibility in {"show", "hide"} ): return None return cls( @@ -124,6 +143,9 @@ class ActivityDescriptor: integration_key=( integration_key.strip() if isinstance(integration_key, str) else None ), + kind=kind.strip() if isinstance(kind, str) else None, + lifecycle=lifecycle, + visibility=visibility, ) diff --git a/surfsense_backend/app/routes/__init__.py b/surfsense_backend/app/routes/__init__.py index 0cd10f1cb..c9bdfb8ba 100644 --- a/surfsense_backend/app/routes/__init__.py +++ b/surfsense_backend/app/routes/__init__.py @@ -150,4 +150,6 @@ router.include_router(team_memory_router) # Workspace team memory router.include_router(automations_router) # Automations CRUD + run history router.include_router(file_storage_router) # Original file metadata + download router.include_router(build_capabilities_router()) # Scraper-API capability doors (05) -router.include_router(build_authenticated_artifact_router()) # Authenticated artifact generation (dev API) +router.include_router( + build_authenticated_artifact_router() +) # Authenticated artifact generation (dev API) diff --git a/surfsense_backend/app/routes/new_chat_routes.py b/surfsense_backend/app/routes/new_chat_routes.py index 7e475017c..94ac4ad3b 100644 --- a/surfsense_backend/app/routes/new_chat_routes.py +++ b/surfsense_backend/app/routes/new_chat_routes.py @@ -2432,6 +2432,7 @@ async def get_pending_interrupts( """ from app.agents.chat.runtime.checkpointer import get_checkpointer from app.services.new_streaming_service import VercelStreamingService + from app.tasks.chat.persistence import load_assistant_message_for_turn from app.tasks.chat.streaming.helpers.interrupt_inspector import ( pending_interrupt_entries_from_writes, ) @@ -2471,25 +2472,18 @@ async def get_pending_interrupts( payload = {**payload, "interrupt_id": interrupt_id} payloads.append(payload) - # Reattach the card to the paused turn's assistant row. ``turn_id`` on the - # checkpoint mirrors ``NewChatMessage.turn_id``; fall back to the newest - # assistant row (the paused turn is always the head). + # Reattach the card only to the paused turn's assistant row. ``turn_id`` on + # the checkpoint mirrors ``NewChatMessage.turn_id``. metadata = checkpoint_tuple.metadata or {} turn_id = metadata.get("turn_id") if isinstance(metadata, dict) else None - assistant_query = ( - select(NewChatMessage.id) - .filter( - NewChatMessage.thread_id == thread_id, - NewChatMessage.role == NewChatMessageRole.ASSISTANT, - ) - .order_by(NewChatMessage.created_at.desc()) + assistant_message = await load_assistant_message_for_turn( + session, + chat_id=thread_id, + turn_id=turn_id if isinstance(turn_id, str) else None, ) - if turn_id: - assistant_query = assistant_query.filter(NewChatMessage.turn_id == turn_id) - assistant_message_id = (await session.execute(assistant_query.limit(1))).scalar() return PendingInterruptsResponse( - assistant_message_id=assistant_message_id, + assistant_message_id=assistant_message.id if assistant_message else None, pending_interrupts=payloads, ) diff --git a/surfsense_backend/app/tasks/celery_tasks/video_presentation_tasks.py b/surfsense_backend/app/tasks/celery_tasks/video_presentation_tasks.py index 7fedea7e0..e47dcdcee 100644 --- a/surfsense_backend/app/tasks/celery_tasks/video_presentation_tasks.py +++ b/surfsense_backend/app/tasks/celery_tasks/video_presentation_tasks.py @@ -116,9 +116,7 @@ async def _generate_video_presentation( video_pres = result.scalars().first() if not video_pres: - raise ValueError( - f"VideoPresentationRun {video_presentation_id} not found" - ) + raise ValueError(f"VideoPresentationRun {video_presentation_id} not found") try: video_pres.status = VideoPresentationStatus.GENERATING @@ -255,9 +253,7 @@ async def _generate_video_presentation( video_pres.artifact_id, ) await session.commit() - logger.info( - "VideoPresentationRun %s: READY commit complete", video_pres.id - ) + logger.info("VideoPresentationRun %s: READY commit complete", video_pres.id) logger.info(f"Successfully generated video presentation: {video_pres.id}") diff --git a/surfsense_backend/app/tasks/chat/content_builder.py b/surfsense_backend/app/tasks/chat/content_builder.py index fdfdcbfd9..11f807657 100644 --- a/surfsense_backend/app/tasks/chat/content_builder.py +++ b/surfsense_backend/app/tasks/chat/content_builder.py @@ -1,36 +1,9 @@ -"""Server-side mirror of the frontend's assistant-ui ``ContentPart`` projection. +"""Server-owned assistant content projection persisted at the end of a turn. -Background ----------- -The streaming chat task in ``stream_new_chat`` / ``stream_resume_chat`` yields -SSE events that the frontend folds into a ``ContentPartsState`` (see -``surfsense_web/lib/chat/streaming-state.ts`` and the matching pipeline in -``stream-pipeline.ts``). When a turn ends, the frontend calls -``buildContentForPersistence(...)`` and round-trips that ``ContentPart[]`` -JSONB to ``POST /threads/{id}/messages``, which is what was historically -written to ``new_chat_messages.content``. - -After the ghost-thread fix moved persistence server-side, the assistant -row is written by ``finalize_assistant_turn`` in the streaming finally -block. The frontend's later ``appendMessage`` is now a no-op (recovers -via the ``(thread_id, turn_id, role)`` partial unique index added in -migration 141), which means the *server* is now responsible for -producing the rich ``ContentPart[]`` shape the FE expects on history -reload — text + reasoning + tool-call cards (with ``args``, ``argsText``, -``result``, ``langchainToolCallId``) + canonical activity snapshots. - -This module is the in-memory accumulator that mirrors the FE state for -exactly that purpose. The streaming code calls ``on_text_*`` / ``on_reasoning_*`` -/ ``on_tool_*`` / ``on_activity`` / ``on_activity_timing`` / -``mark_interrupted`` at the same call sites it yields the matching -``streaming_service.format_*`` SSE event, so the in-memory ``parts`` list -stays in lockstep with what the FE's pipeline would have produced live. -``snapshot()`` is then taken once in the ``finally`` block and persisted -in a single UPDATE. - -Pure synchronous state — no DB I/O, no async, no flush callbacks. The -streaming code is responsible for driving lifecycle methods; this class -is a thin projection helper. +SSE emit sites update this synchronous accumulator in lockstep with the wire. +``snapshot()`` returns the complete assistant content-part list for one final +database update. Activity lifecycle decisions belong to ``ActivityJournal``; +this class only stores emitted snapshots with terminal precedence. """ from __future__ import annotations @@ -44,9 +17,8 @@ from typing import Any logger = logging.getLogger(__name__) -# Mirrors the FE's filter in ``buildContentForPersistence`` / ``buildContentForUI``: -# only text/reasoning/tool-call parts count as "meaningful". data-activities -# decorates the meaningful parts but never stands alone in a successful turn. +# Only text/reasoning/tool-call parts count as meaningful. Activity data +# decorates those parts but never stands alone in a successful turn. _MEANINGFUL_PART_TYPES: frozenset[str] = frozenset({"text", "reasoning", "tool-call"}) @@ -68,10 +40,10 @@ def _merge_tool_part_metadata( class AssistantContentBuilder: - """Server-side projection of ``surfsense_web/lib/chat/streaming-state.ts``. + """Accumulate canonical assistant content parts beside SSE emission. Output shape (deep copy of ``self.parts`` via ``snapshot()``) strictly - matches the FE ``ContentPart`` union:: + matches the web ``ContentPart`` union:: | { type: "text"; text: string } | { type: "reasoning"; text: string } @@ -82,7 +54,7 @@ class AssistantContentBuilder: | { type: "data-activities"; data: { activities: ActivityData[]; timing: ActivityTimingData } } Order matches the wire order of the SSE events that drive the lifecycle - methods, with two FE-mirrored exceptions: + methods, with one canonical exception: 1. ``data-activities`` is a singleton pinned at index 0. Full snapshots replace entries by id and remain ordered by immutable sequence. @@ -103,9 +75,7 @@ class AssistantContentBuilder: # threads through every ``tool-input-*`` / ``tool-output-*`` event. self._tool_call_idx_by_ui_id: dict[str, int] = {} # Live argsText accumulator (concatenated ``tool-input-delta`` chunks) - # so we can reproduce the FE's ``appendToolInputDelta`` behaviour - # before ``tool-input-available`` overwrites it with the - # pretty-printed final JSON. + # before ``tool-input-available`` replaces it with final formatted JSON. self._args_text_by_ui_id: dict[str, str] = {} # ------------------------------------------------------------------ @@ -393,7 +363,9 @@ class AssistantContentBuilder: break if not replaced: activities.append(new_snapshot) - activities.sort(key=lambda value: value.get("sequence", 0)) + activities.sort( + key=lambda value: (value.get("sequence", 0), value.get("id", "")) + ) self.parts[existing_idx] = { "type": "data-activities", "data": { @@ -419,10 +391,22 @@ class AssistantContentBuilder: self._tool_call_idx_by_ui_id[ui_id] = idx + 1 def on_activity_timing(self, snapshot: dict[str, Any]) -> None: - """Replace the journal's canonical active-time snapshot.""" + """Advance the journal's canonical active-time snapshot monotonically.""" for i, part in enumerate(self.parts): if part.get("type") != "data-activities": continue + current = part.get("data", {}).get("timing") + if isinstance(current, dict): + if current == snapshot or current.get("status") == "completed": + return + current_duration = current.get("activeDurationMs") + next_duration = snapshot.get("activeDurationMs") + if ( + isinstance(current_duration, int) + and isinstance(next_duration, int) + and next_duration < current_duration + ): + return self.parts[i] = { "type": "data-activities", "data": { @@ -454,7 +438,7 @@ class AssistantContentBuilder: # ------------------------------------------------------------------ def mark_interrupted(self) -> None: - """Close any open text/reasoning and flip running tools to aborted. + """Close open text/reasoning and mark unfinished tool parts aborted. Called from the streaming ``finally`` block before ``snapshot()`` so the persisted JSONB reflects a coherent end-state even when the @@ -477,15 +461,6 @@ class AssistantContentBuilder: self._current_reasoning_id = None self._current_reasoning_started_at = None for part in self.parts: - if part.get("type") == "data-activities": - for activity in part.get("data", {}).get("activities", []): - if activity.get("status") == "running": - activity["status"] = "interrupted" - activity["title"] = ( - f"Interrupted: {activity.get('title', 'activity')}" - ) - activity["completedAt"] = datetime.now(UTC).isoformat() - continue if part.get("type") != "tool-call": continue if "result" in part: diff --git a/surfsense_backend/app/tasks/chat/persistence.py b/surfsense_backend/app/tasks/chat/persistence.py index f78849607..a1da3a4fe 100644 --- a/surfsense_backend/app/tasks/chat/persistence.py +++ b/surfsense_backend/app/tasks/chat/persistence.py @@ -57,6 +57,7 @@ from uuid import UUID from sqlalchemy import text as sa_text from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select from app.db import ( @@ -87,15 +88,29 @@ _EMPTY_SHELL_CONTENT: list[dict[str, Any]] = [{"type": "text", "text": ""}] # no tool calls). The streaming layer flips to this when # ``AssistantContentBuilder.is_empty()`` returns True so the persisted # row is at least somewhat self-describing instead of an empty text -# bubble. The FE's ``ContentPart`` union doesn't include ``status`` -# yet, so the history loader will silently drop this part and render -# a blank bubble (matches today's behaviour for empty turns); a follow-up -# FE PR adds the explicit "no response" rendering. +# bubble. The history converter presents this status as readable text. _STATUS_NO_RESPONSE: list[dict[str, Any]] = [ {"type": "status", "text": "(no text response)"} ] +async def load_assistant_message_for_turn( + session: AsyncSession, + *, + chat_id: int, + turn_id: str | None, +) -> NewChatMessage | None: + """Resolve the paused assistant row by checkpoint turn identity.""" + if not turn_id: + return None + query = select(NewChatMessage).where( + NewChatMessage.thread_id == chat_id, + NewChatMessage.role == NewChatMessageRole.ASSISTANT, + NewChatMessage.turn_id == turn_id, + ) + return (await session.execute(query.limit(1))).scalar_one_or_none() + + def _build_user_content( user_query: str, user_image_data_urls: list[str] | None, diff --git a/surfsense_backend/app/tasks/chat/streaming/activity_timing.py b/surfsense_backend/app/tasks/chat/streaming/activity_timing.py index 83126c6a5..384540dcb 100644 --- a/surfsense_backend/app/tasks/chat/streaming/activity_timing.py +++ b/surfsense_backend/app/tasks/chat/streaming/activity_timing.py @@ -12,7 +12,12 @@ _NANOSECONDS_PER_MILLISECOND = 1_000_000 @dataclass class ActivityTimer: - """Accumulate execution time while excluding HITL suspension.""" + """Measure one assistant turn's active wall time. + + The timer starts when the backend accepts a new or resumed turn, includes + model, tool, retry, and final-answer work, and excludes time suspended for + human approval. + """ active_duration_ns: int active_since_ns: int | None @@ -67,6 +72,14 @@ class ActivityTimer: self.status = "completed" return self.snapshot() + def complete_if_running( + self, *, now_ns: int | None = None + ) -> ActivityTimingData | None: + """Complete cleanup work once without changing paused/terminal timers.""" + if self.status != "running": + return None + return self.complete(now_ns=now_ns) + def _stop_segment(self, *, now_ns: int | None) -> None: if self.status != "running" or self.active_since_ns is None: return diff --git a/surfsense_backend/app/tasks/chat/streaming/agent/event_loop.py b/surfsense_backend/app/tasks/chat/streaming/agent/event_loop.py index 4af830152..8564f7ae4 100644 --- a/surfsense_backend/app/tasks/chat/streaming/agent/event_loop.py +++ b/surfsense_backend/app/tasks/chat/streaming/agent/event_loop.py @@ -9,6 +9,7 @@ intent classification, and interrupt detection. from __future__ import annotations from collections.abc import AsyncGenerator +from datetime import UTC, datetime from typing import Any from app.agents.chat.multi_agent_chat.main_agent.middleware.kb_persistence import ( @@ -33,9 +34,6 @@ from app.tasks.chat.streaming.graph_stream.event_stream import stream_output from app.tasks.chat.streaming.helpers.interrupt_inspector import ( all_interrupt_entries, ) -from app.tasks.chat.streaming.relay.activity_completion import ( - iter_complete_open_activity_frames, -) from app.tasks.chat.streaming.relay.activity_sse import ( emit_activity_frame, emit_activity_timing_frame, @@ -55,6 +53,8 @@ async def stream_agent_events( result: StreamResult, step_prefix: str = "turn", initial_activities: list[ActivityData] | None = None, + resume_activity_id_by_tool_call: dict[str, str] | None = None, + resume_tool_call_ids: list[str] | None = None, *, fallback_commit_workspace_id: int | None = None, fallback_commit_created_by_id: str | None = None, @@ -69,6 +69,11 @@ async def stream_agent_events( ``accumulated_text`` and interrupt state. See ``StreamResult`` for the side-channel surface populated by the underlying relay. """ + + async def load_agent_state() -> Any: + return await agent.aget_state(config) + + result.load_agent_state = load_agent_state async for sse in stream_output( agent=agent, config=config, @@ -77,6 +82,8 @@ async def stream_agent_events( result=result, step_prefix=step_prefix, initial_activities=initial_activities, + resume_activity_id_by_tool_call=resume_activity_id_by_tool_call, + resume_tool_call_ids=resume_tool_call_ids, content_builder=content_builder, runtime_context=runtime_context, ): @@ -228,21 +235,12 @@ async def stream_agent_events( ) activity_state = result.activity_state if activity_state is not None: - for activity_id, current in list( - activity_state.activity_snapshot_by_id.items() - ): - if current.get("status") not in {"running", "awaiting_approval"}: - continue - snapshot = activity_state.transition_activity( - activity_id, - status="awaiting_approval", + for snapshot in activity_state.journal.await_approval(): + yield emit_activity_frame( + streaming_service=streaming_service, + content_builder=content_builder, + snapshot=snapshot, ) - if snapshot: - yield emit_activity_frame( - streaming_service=streaming_service, - content_builder=content_builder, - snapshot=snapshot, - ) # One frame per paused subagent so each parallel HITL renders its own # approval card on the wire. Order matches ``state.interrupts``, which # the resume slicer in @@ -253,9 +251,11 @@ async def stream_agent_events( interrupt_value, interrupt_id=interrupt_id ) elif result.activity_state is not None: - for frame in iter_complete_open_activity_frames( - state=result.activity_state, - streaming_service=streaming_service, - content_builder=content_builder, + for snapshot in result.activity_state.journal.complete_open_phases( + completed_at=datetime.now(UTC).isoformat() ): - yield frame + yield emit_activity_frame( + streaming_service=streaming_service, + content_builder=content_builder, + snapshot=snapshot, + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py b/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py index cc272fecc..dc274d48f 100644 --- a/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py +++ b/surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py @@ -106,7 +106,11 @@ from app.tasks.chat.streaming.flows.shared.stream_loop import run_stream_loop from app.tasks.chat.streaming.flows.shared.terminal_error import ( handle_terminal_exception, ) -from app.tasks.chat.streaming.relay.activity_sse import emit_activity_timing_frame +from app.tasks.chat.streaming.relay.activity_sse import ( + emit_activity_timing_frame, + emit_completed_activity_timing_frame, + emit_completed_activity_timing_frame_if_running, +) from app.tasks.chat.streaming.shared.stream_result import StreamResult from app.utils.perf import get_perf_logger, log_system_snapshot @@ -748,10 +752,10 @@ async def stream_new_chat( yield sse return - yield emit_activity_timing_frame( + yield emit_completed_activity_timing_frame( streaming_service=streaming_service, content_builder=stream_result.content_builder, - snapshot=stream_result.activity_timer.complete(), + timer=stream_result.activity_timer, ) async for title_sse in await_pending_title_update( @@ -785,15 +789,14 @@ async def stream_new_chat( yield sse except Exception as exc: - if ( - stream_result.content_builder is not None - and stream_result.activity_timer.status == "running" - ): - yield emit_activity_timing_frame( + if stream_result.content_builder is not None: + completed_timing_frame = emit_completed_activity_timing_frame_if_running( streaming_service=streaming_service, content_builder=stream_result.content_builder, - snapshot=stream_result.activity_timer.complete(), + timer=stream_result.activity_timer, ) + if completed_timing_frame is not None: + yield completed_timing_frame frames, summary = handle_terminal_exception( exc, flow=flow, diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/assistant_shell.py b/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/assistant_shell.py index c90d6ad32..46240c4de 100644 --- a/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/assistant_shell.py +++ b/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/assistant_shell.py @@ -18,61 +18,84 @@ from __future__ import annotations from dataclasses import dataclass from typing import Any, cast -from sqlalchemy.future import select - -from app.db import NewChatMessage, NewChatMessageRole, shielded_async_session +from app.db import shielded_async_session from app.services.streaming.types import ActivityData, ActivityTimingData -from app.tasks.chat.persistence import persist_assistant_shell +from app.tasks.chat.persistence import ( + load_assistant_message_for_turn, + persist_assistant_shell, +) @dataclass(frozen=True) class ResumableActivityJournal: activities: list[ActivityData] timing: ActivityTimingData | None + activity_id_by_tool_call: dict[str, str] + tool_call_ids: list[str] + + +def order_resume_tool_call_ids( + journal: ResumableActivityJournal, + pending_tool_call_ids: list[str], +) -> list[str]: + """Prefer checkpoint interrupt order, then persisted activity sequence.""" + preferred_by_activity = { + journal.activity_id_by_tool_call[tool_call_id]: tool_call_id + for tool_call_id in journal.tool_call_ids + if tool_call_id in journal.activity_id_by_tool_call + } + ordered: list[str] = [] + for pending_id in pending_tool_call_ids: + activity_id = journal.activity_id_by_tool_call.get(pending_id) + preferred_id = preferred_by_activity.get(activity_id) if activity_id else None + if preferred_id and preferred_id not in ordered: + ordered.append(preferred_id) + ordered.extend( + tool_call_id + for tool_call_id in journal.tool_call_ids + if tool_call_id not in ordered + ) + return ordered async def load_resumable_activity_journal( chat_id: int, + *, + turn_id: str | None, ) -> ResumableActivityJournal: - """Load the paused journal rows that the resumed graph may continue.""" + """Load the exact paused assistant row the resumed graph may continue.""" async with shielded_async_session() as session: - contents = ( - ( - await session.execute( - select(NewChatMessage.content) - .where( - NewChatMessage.thread_id == chat_id, - NewChatMessage.role == NewChatMessageRole.ASSISTANT, - ) - .order_by(NewChatMessage.id.desc()) - .limit(20) - ) - ) - .scalars() - .all() + message = await load_assistant_message_for_turn( + session, + chat_id=chat_id, + turn_id=turn_id, ) - - return _resumable_journal_from_messages(contents) + return _resumable_journal_from_content(message.content if message else None) -def _resumable_journal_from_messages(contents: Any) -> ResumableActivityJournal: - if not isinstance(contents, (list, tuple)): - return ResumableActivityJournal([], None) - latest_by_id: dict[str, ActivityData] = {} - timing: ActivityTimingData | None = None - for content in contents: - activities, candidate_timing = _journal_from_content(content) - for activity in activities: - latest_by_id.setdefault(activity["id"], activity) - if timing is None and candidate_timing is not None: - timing = candidate_timing - return ResumableActivityJournal( - activities=[ +def _resumable_journal_from_content(content: Any) -> ResumableActivityJournal: + activities, timing = _journal_from_content(content) + awaiting_activities = sorted( + ( activity - for activity in latest_by_id.values() + for activity in activities if activity["status"] == "awaiting_approval" - ], + ), + key=lambda activity: (activity["sequence"], activity["id"]), + ) + awaiting = {activity["id"]: activity for activity in awaiting_activities} + bindings, preferred_ids = _activity_bindings_from_content( + content, valid_activity_ids=awaiting.keys() + ) + return ResumableActivityJournal( + activities=awaiting_activities, timing=timing if timing and timing["status"] == "paused" else None, + activity_id_by_tool_call=bindings, + tool_call_ids=[ + preferred_ids[activity["id"]] + for activity in awaiting_activities + if activity["id"] in preferred_ids + ], ) @@ -131,6 +154,35 @@ def _is_activity_timing(value: Any) -> bool: ) +def _activity_bindings_from_content( + content: Any, + *, + valid_activity_ids: Any, +) -> tuple[dict[str, str], dict[str, str]]: + if not isinstance(content, list): + return {}, {} + valid_ids = set(valid_activity_ids) + bindings: dict[str, str] = {} + preferred_id_by_activity: dict[str, str] = {} + for part in content: + if not isinstance(part, dict) or part.get("type") != "tool-call": + continue + metadata = part.get("metadata") + activity_id = metadata.get("activityId") if isinstance(metadata, dict) else None + if not isinstance(activity_id, str) or activity_id not in valid_ids: + continue + preferred_id: str | None = None + for key in ("langchainToolCallId", "toolCallId"): + tool_call_id = part.get(key) + if isinstance(tool_call_id, str) and tool_call_id: + bindings[tool_call_id] = activity_id + if preferred_id is None: + preferred_id = tool_call_id + if preferred_id is not None: + preferred_id_by_activity.setdefault(activity_id, preferred_id) + return bindings, preferred_id_by_activity + + async def persist_resume_assistant_shell( *, chat_id: int, diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/orchestrator.py b/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/orchestrator.py index 6a471ff9a..3ba990c69 100644 --- a/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/orchestrator.py +++ b/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/orchestrator.py @@ -41,6 +41,7 @@ from app.tasks.chat.streaming.contract.file_contract import log_file_contract from app.tasks.chat.streaming.errors.emitter import emit_stream_terminal_error from app.tasks.chat.streaming.flows.resume_chat.assistant_shell import ( load_resumable_activity_journal, + order_resume_tool_call_ids, persist_resume_assistant_shell, ) from app.tasks.chat.streaming.flows.resume_chat.resume_routing import ( @@ -91,7 +92,11 @@ from app.tasks.chat.streaming.flows.shared.stream_loop import run_stream_loop from app.tasks.chat.streaming.flows.shared.terminal_error import ( handle_terminal_exception, ) -from app.tasks.chat.streaming.relay.activity_sse import emit_activity_timing_frame +from app.tasks.chat.streaming.relay.activity_sse import ( + emit_activity_timing_frame, + emit_completed_activity_timing_frame, + emit_completed_activity_timing_frame_if_running, +) from app.tasks.chat.streaming.shared.stream_result import StreamResult from app.tasks.chat.streaming.shared.utils import resume_step_prefix from app.utils.perf import get_perf_logger @@ -373,7 +378,19 @@ async def stream_resume_chat( routing = await build_resume_routing( agent, chat_id=chat_id, decisions=decisions ) - resumable_journal = await load_resumable_activity_journal(chat_id) + paused_checkpoint = await checkpointer.aget_tuple( + {"configurable": {"thread_id": str(chat_id)}} + ) + paused_metadata = paused_checkpoint.metadata if paused_checkpoint else {} + paused_turn_id = ( + paused_metadata.get("turn_id") + if isinstance(paused_metadata, dict) + else None + ) + resumable_journal = await load_resumable_activity_journal( + chat_id, + turn_id=paused_turn_id if isinstance(paused_turn_id, str) else None, + ) config = { "configurable": { @@ -537,6 +554,12 @@ async def stream_resume_chat( stream_result=stream_result, step_prefix=resume_step_prefix(stream_result.turn_id), initial_activities=resumable_journal.activities, + resume_activity_id_by_tool_call=( + resumable_journal.activity_id_by_tool_call + ), + resume_tool_call_ids=order_resume_tool_call_ids( + resumable_journal, routing.pending_tool_call_ids + ), fallback_commit_workspace_id=workspace_id, fallback_commit_created_by_id=user_id, fallback_commit_filesystem_mode=( @@ -572,10 +595,10 @@ async def stream_resume_chat( yield sse return - yield emit_activity_timing_frame( + yield emit_completed_activity_timing_frame( streaming_service=streaming_service, content_builder=stream_result.content_builder, - snapshot=stream_result.activity_timer.complete(), + timer=stream_result.activity_timer, ) if premium_reservation is not None and user_id: @@ -595,15 +618,14 @@ async def stream_resume_chat( yield sse except Exception as exc: - if ( - stream_result.content_builder is not None - and stream_result.activity_timer.status == "running" - ): - yield emit_activity_timing_frame( + if stream_result.content_builder is not None: + completed_timing_frame = emit_completed_activity_timing_frame_if_running( streaming_service=streaming_service, content_builder=stream_result.content_builder, - snapshot=stream_result.activity_timer.complete(), + timer=stream_result.activity_timer, ) + if completed_timing_frame is not None: + yield completed_timing_frame frames, summary = handle_terminal_exception( exc, flow="resume", diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/resume_routing.py b/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/resume_routing.py index 8b8b2068c..af520927c 100644 --- a/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/resume_routing.py +++ b/surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/resume_routing.py @@ -27,6 +27,7 @@ class ResumeRoutingPayload: routed_resume_value: dict[str, Any] lg_resume_map: dict[str, Any] + pending_tool_call_ids: list[str] async def build_resume_routing( @@ -75,6 +76,7 @@ async def build_resume_routing( return ResumeRoutingPayload( routed_resume_value={}, lg_resume_map=lg_resume_map, + pending_tool_call_ids=[], ) routed_resume_value = slice_decisions_by_tool_call(decisions, pending) @@ -82,4 +84,5 @@ async def build_resume_routing( return ResumeRoutingPayload( routed_resume_value=routed_resume_value, lg_resume_map=lg_resume_map, + pending_tool_call_ids=[tool_call_id for tool_call_id, _ in pending], ) diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/shared/assistant_finalize.py b/surfsense_backend/app/tasks/chat/streaming/flows/shared/assistant_finalize.py index 5ad9c2c81..18d4e3175 100644 --- a/surfsense_backend/app/tasks/chat/streaming/flows/shared/assistant_finalize.py +++ b/surfsense_backend/app/tasks/chat/streaming/flows/shared/assistant_finalize.py @@ -29,6 +29,7 @@ from app.agents.chat.multi_agent_chat.shared.citations import ( CitationRegistry, normalize_citations, ) +from app.tasks.chat.streaming.helpers.interrupt_inspector import all_interrupt_entries from app.tasks.chat.streaming.shared.stream_result import StreamResult from app.utils.perf import get_perf_logger @@ -92,25 +93,45 @@ async def finalize_assistant_message( builder_stats: dict[str, int] | None = None if stream_result.content_builder is not None: - if stream_result.activity_timer.status == "running": - stream_result.content_builder.on_activity_timing( - stream_result.activity_timer.complete() - ) + paused_for_approval = ( + stream_result.is_interrupted + or stream_result.activity_timer.status == "paused" + ) + if ( + stream_result.activity_timer.status == "running" + and stream_result.load_agent_state is not None + ): + try: + state = await stream_result.load_agent_state() + paused_for_approval = bool(all_interrupt_entries(state)) + except Exception as exc: + _perf_log.warning( + "[%s] unable to inspect checkpoint during timing cleanup: %s", + log_prefix, + exc, + ) + + terminal_timing = ( + stream_result.activity_timer.pause() + if paused_for_approval and stream_result.activity_timer.status == "running" + else stream_result.activity_timer.complete_if_running() + ) + if terminal_timing is not None: + stream_result.content_builder.on_activity_timing(terminal_timing) + if paused_for_approval: + stream_result.is_interrupted = True + activity_state = stream_result.activity_state if activity_state is not None: - interrupted_at = datetime.now(UTC).isoformat() - for activity_id, current in list( - activity_state.activity_snapshot_by_id.items() - ): - if current.get("status") != "running": - continue - snapshot = activity_state.transition_activity( - activity_id, - status="interrupted", - completed_at=interrupted_at, + snapshots = ( + activity_state.journal.await_approval() + if paused_for_approval + else activity_state.journal.interrupt_running( + completed_at=datetime.now(UTC).isoformat() ) - if snapshot: - stream_result.content_builder.on_activity(snapshot) + ) + for snapshot in snapshots: + stream_result.content_builder.on_activity(snapshot) stream_result.content_builder.mark_interrupted() # Snapshot stats BEFORE ``snapshot()`` deepcopies so the perf log # records the actual finalised payload (post-mark_interrupted), not diff --git a/surfsense_backend/app/tasks/chat/streaming/flows/shared/stream_loop.py b/surfsense_backend/app/tasks/chat/streaming/flows/shared/stream_loop.py index 1a10fc7f6..20123e6f9 100644 --- a/surfsense_backend/app/tasks/chat/streaming/flows/shared/stream_loop.py +++ b/surfsense_backend/app/tasks/chat/streaming/flows/shared/stream_loop.py @@ -36,6 +36,8 @@ async def run_stream_loop( stream_result: StreamResult, step_prefix: str = "turn", initial_activities: list[ActivityData] | None = None, + resume_activity_id_by_tool_call: dict[str, str] | None = None, + resume_tool_call_ids: list[str] | None = None, fallback_commit_workspace_id: int | None, fallback_commit_created_by_id: str | None, fallback_commit_filesystem_mode: FilesystemMode, @@ -61,6 +63,8 @@ async def run_stream_loop( result=stream_result, step_prefix=step_prefix, initial_activities=initial_activities, + resume_activity_id_by_tool_call=resume_activity_id_by_tool_call, + resume_tool_call_ids=resume_tool_call_ids, fallback_commit_workspace_id=fallback_commit_workspace_id, fallback_commit_created_by_id=fallback_commit_created_by_id, fallback_commit_filesystem_mode=fallback_commit_filesystem_mode, diff --git a/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py b/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py index 77b872369..7b6077b77 100644 --- a/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py +++ b/surfsense_backend/app/tasks/chat/streaming/graph_stream/event_stream.py @@ -20,11 +20,17 @@ async def stream_output( result: StreamingResult, step_prefix: str = "turn", initial_activities: list[ActivityData] | None = None, + resume_activity_id_by_tool_call: dict[str, str] | None = None, + resume_tool_call_ids: list[str] | None = None, content_builder: Any | None = None, runtime_context: Any = None, ) -> AsyncIterator[str]: """Yield SSE frames from agent ``astream_events`` via ``EventRelay``.""" - state = AgentEventRelayState.for_invocation(initial_activities=initial_activities) + state = AgentEventRelayState.for_invocation( + initial_activities=initial_activities, + resume_activity_id_by_tool_call=resume_activity_id_by_tool_call, + resume_tool_call_ids=resume_tool_call_ids, + ) astream_kwargs: dict[str, Any] = {"config": config, "version": "v2"} if runtime_context is not None: diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py b/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py index ba73ed496..fa78b86b7 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/chat_model_stream.py @@ -3,12 +3,11 @@ from __future__ import annotations from collections.abc import Iterator +from datetime import UTC, datetime from typing import Any from app.tasks.chat.streaming.helpers.chunk_parts import extract_chunk_parts -from app.tasks.chat.streaming.relay.activity_completion import ( - iter_complete_open_activity_frames, -) +from app.tasks.chat.streaming.relay.activity_sse import emit_activity_frame from app.tasks.chat.streaming.relay.state import AgentEventRelayState from app.tasks.chat.streaming.relay.task_span import ensure_pending_task_span_for_lc @@ -39,11 +38,14 @@ def iter_chat_model_stream_frames( content_builder.on_text_end(state.current_text_id) state.current_text_id = None if state.current_reasoning_id is None: - yield from iter_complete_open_activity_frames( - state=state, - streaming_service=streaming_service, - content_builder=content_builder, - ) + for snapshot in state.journal.complete_open_phases( + completed_at=datetime.now(UTC).isoformat() + ): + yield emit_activity_frame( + streaming_service=streaming_service, + content_builder=content_builder, + snapshot=snapshot, + ) state.current_reasoning_id = streaming_service.generate_reasoning_id() yield streaming_service.format_reasoning_start( state.current_reasoning_id @@ -59,18 +61,19 @@ def iter_chat_model_stream_frames( if part_type == "text": if state.current_reasoning_id is not None: - yield streaming_service.format_reasoning_end( - state.current_reasoning_id - ) + yield streaming_service.format_reasoning_end(state.current_reasoning_id) if content_builder is not None: content_builder.on_reasoning_end(state.current_reasoning_id) state.current_reasoning_id = None if state.current_text_id is None: - yield from iter_complete_open_activity_frames( - state=state, - streaming_service=streaming_service, - content_builder=content_builder, - ) + for snapshot in state.journal.complete_open_phases( + completed_at=datetime.now(UTC).isoformat() + ): + yield emit_activity_frame( + streaming_service=streaming_service, + content_builder=content_builder, + snapshot=snapshot, + ) state.current_text_id = streaming_service.generate_text_id() yield streaming_service.format_text_start(state.current_text_id) if content_builder is not None: diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py b/surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py index 4eded93d5..c4deba4b2 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/custom_events.py @@ -49,19 +49,7 @@ def handle_activity_progress( detail = _trusted_progress_detail(data) if not detail: return None - candidates = [ - snapshot - for snapshot in state.activity_snapshot_by_id.values() - if snapshot.get("status") in {"running", "awaiting_approval"} - ] - if not candidates: - return None - current = max(candidates, key=lambda snapshot: snapshot["sequence"]) - snapshot = state.transition_activity( - current["id"], - status="running", - details=[detail], - ) + snapshot = state.journal.update_current_progress(detail) if snapshot is None: return None return emit_activity_frame( diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py index 68fd587c9..4d993dd1a 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_end.py @@ -10,6 +10,7 @@ from typing import Any from langchain_core.messages import ToolMessage from langgraph.types import Command +from app.services.streaming.types import ActivityStatus from app.tasks.chat.streaming.handlers.tools import ( ToolCompletionEmissionContext, iter_tool_completion_emission_frames, @@ -110,30 +111,28 @@ def iter_tool_end_frames( holder["value"] = state.lc_tool_call_id_by_run[run_id] failed = tool_output_has_error(tool_output) - activity_id = state.activity_id_by_run.pop(run_id, None) if run_id else None - tool_metadata = state.tool_activity_metadata(activity_id=activity_id) or {} - if activity_id: - spec = state.activity_spec_by_id.get(activity_id) - raw_status = str(tool_output.get("status") or "").lower() - terminal_status = ( - "cancelled" - if raw_status in {"cancelled", "canceled", "rejected"} - else "error" - if failed - else "completed" + raw_status = str(tool_output.get("status") or "").lower() + terminal_status: ActivityStatus = ( + "cancelled" + if raw_status in {"cancelled", "canceled", "rejected"} + else "error" + if failed + else "completed" + ) + activity_finish = state.journal.finish_tool( + run_id=run_id, + status=terminal_status, + completed_at=completed_at, + ) + tool_metadata = ( + state.tool_activity_metadata(activity_id=activity_finish.activity_id) or {} + ) + if activity_finish.snapshot: + yield emit_activity_frame( + streaming_service=streaming_service, + content_builder=content_builder, + snapshot=activity_finish.snapshot, ) - if spec and (spec.lifecycle == "invocation" or terminal_status != "completed"): - snapshot = state.transition_activity( - activity_id, - status=terminal_status, - completed_at=completed_at, - ) - if snapshot: - yield emit_activity_frame( - streaming_service=streaming_service, - content_builder=content_builder, - snapshot=snapshot, - ) if tool_name == "verify_artifact": state.deliverable_needs_repair = failed diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py index de90d987f..22569085c 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tool_start.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import logging from collections.abc import Iterator from datetime import UTC, datetime from typing import Any @@ -16,6 +17,8 @@ from app.tasks.chat.streaming.relay.activity_sse import emit_activity_frame from app.tasks.chat.streaming.relay.state import AgentEventRelayState from app.tasks.chat.streaming.relay.task_span import open_task_span +logger = logging.getLogger(__name__) + def _safe_integration_metadata( event: dict[str, Any], @@ -36,13 +39,6 @@ def _safe_integration_metadata( return None -def _artifact_instruction_type(tool_name: str, tool_input: Any) -> str | None: - if tool_name != "load_artifact_instructions" or not isinstance(tool_input, dict): - return None - artifact_type = tool_input.get("artifact_type") - return artifact_type if artifact_type in {"pdf", "docx", "pptx", "xlsx"} else None - - def iter_tool_start_frames( event: dict[str, Any], *, @@ -109,7 +105,6 @@ def iter_tool_start_frames( if isinstance(subagent_type, str) and subagent_type.strip(): state.active_subagent_type = subagent_type.strip() - artifact_type = _artifact_instruction_type(tool_name, tool_input) event_metadata = event.get("metadata") trusted_descriptor = ( event_metadata.get("activity_descriptor") @@ -119,57 +114,36 @@ def iter_tool_start_frames( activity = resolve_tool_activity( tool_name, subagent_type=state.active_subagent_type, - artifact_type=artifact_type, repairing_artifact=state.deliverable_needs_repair, trusted_descriptor=trusted_descriptor, ) + if ( + matched_meta is None + and langchain_tool_call_id is None + and activity.visibility != "hide" + ): + langchain_tool_call_id = state.consume_resume_tool_call_id() + if langchain_tool_call_id is None and state.journal.resume_id_by_tool_call: + logger.warning( + "[activity_resume] no persisted tool-call id available " + "for replayed tool name=%s run_id=%s remaining_bindings=%d", + tool_name, + run_id, + len(state.journal.resume_id_by_tool_call), + ) integration = _safe_integration_metadata(event) - activity_id: str | None = None - if activity.visibility != "hide": - scope = state.active_span_id or "root" - phase_key = activity.phase_key if activity.lifecycle == "phase" else None - open_phase = state.open_phase_by_scope.get(scope) - reuse_phase = bool( - phase_key - and open_phase - and open_phase[0] == phase_key - and open_phase[1] not in state.terminal_activity_ids - ) - if open_phase and not reuse_phase: - closed = state.transition_activity( - open_phase[1], status="completed", completed_at=started_at - ) - if closed: - yield emit_activity_frame( - streaming_service=streaming_service, - content_builder=content_builder, - snapshot=closed, - ) - if reuse_phase and open_phase: - activity_id = open_phase[1] - snapshot = state.activity_snapshot_by_id[activity_id] - else: - resumable_ids = state.resumable_activity_ids_by_kind.get(activity.kind) - activity_id = ( - resumable_ids.pop(0) - if resumable_ids - else state.next_activity_id(step_prefix) - ) - previous = state.activity_snapshot_by_id.get(activity_id) - snapshot = activity.snapshot( - activity_id=activity_id, - sequence=(previous["sequence"] if previous else state.activity_counter), - status="running", - started_at=previous["startedAt"] if previous else started_at, - integration=integration - or (previous.get("integration") if previous else None), - ) - state.activity_spec_by_id[activity_id] = activity - state.activity_snapshot_by_id[activity_id] = snapshot - if phase_key: - state.open_phase_by_scope[scope] = (phase_key, activity_id) - if run_id: - state.activity_id_by_run[run_id] = activity_id + activity_start = state.journal.begin_tool( + spec=activity, + run_id=run_id, + step_prefix=step_prefix, + scope=state.active_span_id or "root", + started_at=started_at, + tool_call_id=tool_call_id, + langchain_tool_call_id=langchain_tool_call_id, + integration=integration, + ) + activity_id = activity_start.activity_id + for snapshot in activity_start.snapshots: yield emit_activity_frame( streaming_service=streaming_service, content_builder=content_builder, diff --git a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/activity.py b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/activity.py index 7151cc7ec..109747a8c 100644 --- a/surfsense_backend/app/tasks/chat/streaming/handlers/tools/activity.py +++ b/surfsense_backend/app/tasks/chat/streaming/handlers/tools/activity.py @@ -109,180 +109,18 @@ def _copy( ) -_ACTIVITY_SPECS: dict[str, ActivitySpec] = { - "read_file": _copy("Reading file", "Read file", "file", icon_key="file-text"), - "write_file": _copy("Creating file", "Created file", "file", icon_key="file-plus"), - "edit_file": _copy("Editing file", "Edited file", "file", icon_key="file-pen"), - "move_file": _copy("Moving file", "Moved file", "file", icon_key="files"), - "rm": _copy("Deleting file", "Deleted file", "file", icon_key="file-x"), - "mkdir": _copy("Creating folder", "Created folder", "file", icon_key="folder-plus"), - "rmdir": _copy("Deleting folder", "Deleted folder", "file", icon_key="folder-x"), - "ls": _copy("Reviewing folder", "Reviewed folder", "file", icon_key="folder-open"), - "list_tree": _copy( - "Reviewing file tree", "Reviewed file tree", "file", icon_key="folder-tree" - ), - "glob": _copy("Finding files", "Found files", "file", icon_key="folder-search"), - "grep": _copy( - "Searching project", "Searched project", "file", icon_key="search-code" - ), - "execute": _copy("Running command", "Ran command", "action", icon_key="terminal"), - "execute_code": _copy("Running code", "Ran code", "action", icon_key="square-code"), - "write_todos": _copy( - "Planning work", - "Planned work", - "action", - lifecycle="phase", - icon_key="list-todo", - ), - "load_artifact_source": _copy( - "Opening the artifact", - "Opened the artifact", - "artifact", - lifecycle="phase", - icon_key="file-input", - ), - "read_sandbox_file": _copy( - "Reviewing the artifact", - "Reviewed the artifact", - "artifact", - lifecycle="phase", - icon_key="file-text", - ), - "load_artifact_instructions": _copy( - "Loading artifact instructions", - "Loaded artifact instructions", - "artifact", - "hide", - icon_key="file-input", - ), - "verify_artifact": _copy( - "Checking the artifact", - "Checked the artifact", - "artifact", - icon_key="badge-check", - ), - "save_artifact": _copy( - "Preparing the file", "Presented file", "artifact", icon_key="file-output" - ), - "save_document": _copy( - "Preparing the document", - "Presented document", - "artifact", - icon_key="file-output", - ), - "generate_image": _copy( - "Creating an image", "Created an image", "artifact", icon_key="image" - ), - "display_image": _copy( - "Preparing the image", "Presented image", "artifact", icon_key="image" - ), - "generate_podcast": _copy( - "Creating the podcast", "Created the podcast", "artifact", icon_key="microphone" - ), - "generate_video_presentation": _copy( - "Creating the presentation", - "Created the presentation", - "artifact", - icon_key="film", - ), - "search_knowledge_base": _copy( - "Searching your sources", - "Searched your sources", - "research", - icon_key="library", - ), - "ask_knowledge_base": _copy( - "Reviewing your sources", - "Reviewed your sources", - "research", - icon_key="library", - ), - "scrape_webpage": _copy( - "Reviewing a webpage", - "Reviewed a webpage", - "research", - lifecycle="phase", - icon_key="scan-text", - ), - "link_preview": _copy( - "Reviewing a link", - "Reviewed a link", - "research", - lifecycle="phase", - icon_key="external-link", - ), - "multi_link_preview": _copy( - "Reviewing links", - "Reviewed links", - "research", - lifecycle="phase", - icon_key="external-link", - ), - "create_calendar_event": _copy( - "Creating calendar event", - "Created calendar event", - "connector", - icon_key="calendar", - integration_key="google_calendar", - ), - "update_calendar_event": _copy( - "Updating calendar event", - "Updated calendar event", - "connector", - icon_key="calendar", - integration_key="google_calendar", - ), - "delete_calendar_event": _copy( - "Deleting calendar event", - "Deleted calendar event", - "connector", - icon_key="calendar", - integration_key="google_calendar", - ), - "search_calendar_events": _copy( - "Searching calendar", - "Searched calendar", - "connector", - icon_key="calendar", - integration_key="google_calendar", - ), - "create_automation": _copy( - "Creating automation", "Created automation", "action", icon_key="workflow" - ), - "update_memory": _copy( - "Remembering preference", "Remembered preference", "action", icon_key="brain" - ), - "task": _copy( - "Working with a specialist", - "Worked with a specialist", - "action", - "hide", - icon_key="route", - ), - "get_connected_accounts": _copy( - "Checking connected apps", - "Checked connected apps", - "connector", - lifecycle="phase", - icon_key="search", - ), - "generate_report": _copy( - "Creating report", "Created report", "artifact", "hide", icon_key="file-text" - ), - "generate_resume": _copy( - "Creating resume", "Created resume", "artifact", "hide", icon_key="file-text" - ), - "pwd": _copy( - "Checking folder", "Checked folder", "file", "hide", icon_key="terminal" - ), - "cd": _copy( - "Changing folder", "Changed folder", "file", "hide", icon_key="terminal" - ), - "noop": _copy("Working", "Worked", "action", "hide", icon_key="tool"), - "invalid_tool": _copy( - "Repairing action", "Repaired action", "action", "hide", icon_key="tool" - ), -} +_INTERNAL_TOOL_NAMES = frozenset( + { + "cd", + "generate_report", + "generate_resume", + "invalid_tool", + "load_artifact_instructions", + "noop", + "pwd", + "task", + } +) def _fallback() -> ActivitySpec: @@ -302,14 +140,16 @@ def _activity_from_descriptor(value: object) -> ActivitySpec | None: descriptor = ActivityDescriptor.from_metadata(value) if descriptor is None: return None - kind = value.get("kind") if isinstance(value, dict) else None - if not isinstance(kind, str) or not _ACTIVITY_KIND_RE.fullmatch(kind): + kind = descriptor.kind + if kind is None or not _ACTIVITY_KIND_RE.fullmatch(kind): kind = "connector.action" return _with_kind( _copy( descriptor.active_title, descriptor.completed_title, descriptor.category, + descriptor.visibility, + lifecycle=descriptor.lifecycle, icon_key=descriptor.icon_key, integration_key=descriptor.integration_key, ), @@ -321,12 +161,10 @@ def resolve_tool_activity( tool_name: str, *, subagent_type: str | None, - artifact_type: str | None = None, repairing_artifact: bool = False, trusted_descriptor: dict[str, Any] | None = None, ) -> ActivitySpec: """Resolve display semantics from trusted runtime context, never model copy.""" - del artifact_type if tool_name == "execute" and subagent_type == "deliverables": return _with_kind( _copy( @@ -349,9 +187,11 @@ def resolve_tool_activity( if described is not None: return described - explicit = _ACTIVITY_SPECS.get(tool_name) - if explicit: - return _with_kind(explicit, tool_name) + if tool_name in _INTERNAL_TOOL_NAMES: + return _with_kind( + _copy("Using a tool", "Completed an action", "action", "hide"), + "tool.action", + ) try: capability = get_capability(tool_name) diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/activity_completion.py b/surfsense_backend/app/tasks/chat/streaming/relay/activity_completion.py deleted file mode 100644 index 97d1ab11a..000000000 --- a/surfsense_backend/app/tasks/chat/streaming/relay/activity_completion.py +++ /dev/null @@ -1,31 +0,0 @@ -"""Terminal transitions for open semantic activity phases.""" - -from __future__ import annotations - -from collections.abc import Iterator -from datetime import UTC, datetime -from typing import Any - -from app.tasks.chat.streaming.relay.activity_sse import emit_activity_frame -from app.tasks.chat.streaming.relay.state import AgentEventRelayState - - -def iter_complete_open_activity_frames( - *, - state: AgentEventRelayState, - streaming_service: Any, - content_builder: Any | None, -) -> Iterator[str]: - completed_at = datetime.now(UTC).isoformat() - for _, activity_id in list(state.open_phase_by_scope.values()): - snapshot = state.transition_activity( - activity_id, - status="completed", - completed_at=completed_at, - ) - if snapshot: - yield emit_activity_frame( - streaming_service=streaming_service, - content_builder=content_builder, - snapshot=snapshot, - ) diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/activity_journal.py b/surfsense_backend/app/tasks/chat/streaming/relay/activity_journal.py new file mode 100644 index 000000000..28fca9a4f --- /dev/null +++ b/surfsense_backend/app/tasks/chat/streaming/relay/activity_journal.py @@ -0,0 +1,311 @@ +"""Canonical lifecycle owner for one turn's user-visible activities.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from app.services.streaming.types import ( + ActivityData, + ActivityIntegration, + ActivityStatus, +) +from app.tasks.chat.streaming.handlers.tools.activity import ActivitySpec + +_TERMINAL_STATUSES = {"completed", "error", "cancelled", "interrupted"} +_STATUS_SEVERITY: dict[ActivityStatus, int] = { + "running": 0, + "awaiting_approval": 0, + "completed": 1, + "interrupted": 2, + "cancelled": 3, + "error": 4, +} + + +@dataclass(frozen=True, slots=True) +class ActivityStart: + activity_id: str | None + snapshots: tuple[ActivityData, ...] = () + + +@dataclass(frozen=True, slots=True) +class ActivityFinish: + activity_id: str | None + snapshot: ActivityData | None = None + + +@dataclass +class ActivityJournal: + """Own activity identity, phase reuse, resume binding, and transitions.""" + + counter: int = 0 + id_by_run: dict[str, str] = field(default_factory=dict) + snapshot_by_id: dict[str, ActivityData] = field(default_factory=dict) + spec_by_id: dict[str, ActivitySpec] = field(default_factory=dict) + open_phase_by_scope: dict[str, tuple[str, str]] = field(default_factory=dict) + terminal_ids: set[str] = field(default_factory=set) + resume_id_by_tool_call: dict[str, str] = field(default_factory=dict) + active_runs_by_activity: dict[str, set[str]] = field(default_factory=dict) + deferred_close_at_by_activity: dict[str, str] = field(default_factory=dict) + deferred_outcome_by_activity: dict[str, tuple[ActivityStatus, str]] = field( + default_factory=dict + ) + + @classmethod + def resume( + cls, + *, + activities: list[ActivityData] | None = None, + activity_id_by_tool_call: dict[str, str] | None = None, + ) -> ActivityJournal: + snapshots = { + activity["id"]: activity + for activity in activities or [] + if activity.get("status") == "awaiting_approval" + } + valid_ids = snapshots.keys() + bindings = { + tool_call_id: activity_id + for tool_call_id, activity_id in (activity_id_by_tool_call or {}).items() + if activity_id in valid_ids + } + return cls( + counter=max( + (activity["sequence"] for activity in snapshots.values()), default=0 + ), + snapshot_by_id=snapshots, + resume_id_by_tool_call=bindings, + ) + + def begin_tool( + self, + *, + spec: ActivitySpec, + run_id: str, + step_prefix: str, + scope: str, + started_at: str, + tool_call_id: str, + langchain_tool_call_id: str | None, + integration: ActivityIntegration | None, + ) -> ActivityStart: + if spec.visibility == "hide": + return ActivityStart(None) + + emitted: list[ActivityData] = [] + phase_key = spec.phase_key if spec.lifecycle == "phase" else None + open_phase = self.open_phase_by_scope.get(scope) + reuse_phase = bool( + phase_key + and open_phase + and open_phase[0] == phase_key + and open_phase[1] not in self.terminal_ids + ) + + if open_phase and not reuse_phase: + self.open_phase_by_scope.pop(scope, None) + closed = self._request_phase_close( + open_phase[1], status="completed", completed_at=started_at + ) + if closed is not None: + emitted.append(closed) + + if reuse_phase and open_phase: + activity_id = open_phase[1] + snapshot = self.snapshot_by_id[activity_id] + else: + activity_id = self._consume_resume_id(langchain_tool_call_id, tool_call_id) + if activity_id is None: + activity_id = self._next_id(step_prefix) + previous = self.snapshot_by_id.get(activity_id) + snapshot = spec.snapshot( + activity_id=activity_id, + sequence=previous["sequence"] if previous else self.counter, + status="running", + started_at=previous["startedAt"] if previous else started_at, + integration=integration + or (previous.get("integration") if previous else None), + ) + self.spec_by_id[activity_id] = spec + self.snapshot_by_id[activity_id] = snapshot + if phase_key: + self.open_phase_by_scope[scope] = (phase_key, activity_id) + + if run_id: + self.id_by_run[run_id] = activity_id + self.active_runs_by_activity.setdefault(activity_id, set()).add(run_id) + emitted.append(snapshot) + return ActivityStart(activity_id, tuple(emitted)) + + def finish_tool( + self, + *, + run_id: str, + status: ActivityStatus, + completed_at: str, + ) -> ActivityFinish: + activity_id = self.id_by_run.pop(run_id, None) if run_id else None + if activity_id is None: + return ActivityFinish(None) + spec = self.spec_by_id.get(activity_id) + active_runs = self.active_runs_by_activity.get(activity_id) + if active_runs is not None: + active_runs.discard(run_id) + if not active_runs: + self.active_runs_by_activity.pop(activity_id, None) + if spec is None: + return ActivityFinish(activity_id) + if spec.lifecycle == "phase": + if status != "completed": + self._defer_outcome(activity_id, status, completed_at) + if activity_id in self.active_runs_by_activity: + return ActivityFinish(activity_id) + deferred = self.deferred_outcome_by_activity.pop(activity_id, None) + close_at = self.deferred_close_at_by_activity.pop(activity_id, None) + if deferred is not None: + status, completed_at = deferred + elif close_at is not None: + status, completed_at = "completed", close_at + else: + return ActivityFinish(activity_id) + return ActivityFinish( + activity_id, + self.transition( + activity_id, + status=status, + completed_at=completed_at, + ), + ) + + def transition( + self, + activity_id: str, + *, + status: ActivityStatus, + completed_at: str | None = None, + details: list[str] | None = None, + ) -> ActivityData | None: + current = self.snapshot_by_id.get(activity_id) + spec = self.spec_by_id.get(activity_id) + if current is None or spec is None: + return None + if activity_id in self.terminal_ids: + return current + snapshot = spec.snapshot( + activity_id=activity_id, + sequence=current["sequence"], + status=status, + started_at=current["startedAt"], + completed_at=completed_at, + details=details if details is not None else current.get("details"), + integration=current.get("integration"), + ) + self.snapshot_by_id[activity_id] = snapshot + if status in _TERMINAL_STATUSES: + self.terminal_ids.add(activity_id) + for scope, (_, open_id) in list(self.open_phase_by_scope.items()): + if open_id == activity_id: + self.open_phase_by_scope.pop(scope, None) + return snapshot + + def complete_open_phases(self, *, completed_at: str) -> list[ActivityData]: + snapshots: list[ActivityData] = [] + for scope, (_, activity_id) in list(self.open_phase_by_scope.items()): + self.open_phase_by_scope.pop(scope, None) + snapshot = self._request_phase_close( + activity_id, status="completed", completed_at=completed_at + ) + if snapshot is not None: + snapshots.append(snapshot) + return snapshots + + def await_approval(self) -> list[ActivityData]: + snapshots: list[ActivityData] = [] + for activity_id, current in list(self.snapshot_by_id.items()): + if current.get("status") not in {"running", "awaiting_approval"}: + continue + snapshot = self.transition(activity_id, status="awaiting_approval") + if snapshot is not None: + snapshots.append(snapshot) + return snapshots + + def interrupt_running(self, *, completed_at: str) -> list[ActivityData]: + snapshots: list[ActivityData] = [] + for activity_id, current in list(self.snapshot_by_id.items()): + if current.get("status") != "running": + continue + for run_id in self.active_runs_by_activity.pop(activity_id, set()): + self.id_by_run.pop(run_id, None) + self.deferred_close_at_by_activity.pop(activity_id, None) + self.deferred_outcome_by_activity.pop(activity_id, None) + snapshot = self.transition( + activity_id, + status="interrupted", + completed_at=completed_at, + ) + if snapshot is not None: + snapshots.append(snapshot) + return snapshots + + def update_current_progress(self, detail: str) -> ActivityData | None: + candidates = [ + snapshot + for snapshot in self.snapshot_by_id.values() + if snapshot.get("status") in {"running", "awaiting_approval"} + ] + if not candidates: + return None + current = max(candidates, key=lambda snapshot: snapshot["sequence"]) + return self.transition( + current["id"], + status=current["status"], + details=[detail], + ) + + def _request_phase_close( + self, + activity_id: str, + *, + status: ActivityStatus, + completed_at: str, + ) -> ActivityData | None: + if status != "completed": + self._defer_outcome(activity_id, status, completed_at) + if activity_id in self.active_runs_by_activity: + self.deferred_close_at_by_activity.setdefault(activity_id, completed_at) + return None + deferred = self.deferred_outcome_by_activity.pop(activity_id, None) + self.deferred_close_at_by_activity.pop(activity_id, None) + if deferred is not None: + status, completed_at = deferred + return self.transition(activity_id, status=status, completed_at=completed_at) + + def _defer_outcome( + self, + activity_id: str, + status: ActivityStatus, + completed_at: str, + ) -> None: + current = self.deferred_outcome_by_activity.get(activity_id) + if current is None or _STATUS_SEVERITY[status] > _STATUS_SEVERITY[current[0]]: + self.deferred_outcome_by_activity[activity_id] = (status, completed_at) + + def _next_id(self, step_prefix: str) -> str: + self.counter += 1 + return f"act_{step_prefix}_{self.counter}" + + def _consume_resume_id(self, *tool_call_ids: str | None) -> str | None: + activity_id = next( + ( + self.resume_id_by_tool_call[tool_call_id] + for tool_call_id in tool_call_ids + if tool_call_id and tool_call_id in self.resume_id_by_tool_call + ), + None, + ) + if activity_id is None: + return None + for key, value in list(self.resume_id_by_tool_call.items()): + if value == activity_id: + self.resume_id_by_tool_call.pop(key, None) + return activity_id diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/activity_sse.py b/surfsense_backend/app/tasks/chat/streaming/relay/activity_sse.py index 1ec48ca2c..c110549ac 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/activity_sse.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/activity_sse.py @@ -5,6 +5,7 @@ from __future__ import annotations from typing import Any from app.services.streaming.types import ActivityData, ActivityTimingData +from app.tasks.chat.streaming.activity_timing import ActivityTimer def emit_activity_frame( @@ -27,3 +28,36 @@ def emit_activity_timing_frame( if content_builder is not None: content_builder.on_activity_timing(snapshot) return streaming_service.format_data("activity-timing", snapshot) + + +def emit_completed_activity_timing_frame( + *, + streaming_service: Any, + content_builder: Any | None, + timer: ActivityTimer, + now_ns: int | None = None, +) -> str: + """Strictly complete a successful turn and dual-write its terminal snapshot.""" + return emit_activity_timing_frame( + streaming_service=streaming_service, + content_builder=content_builder, + snapshot=timer.complete(now_ns=now_ns), + ) + + +def emit_completed_activity_timing_frame_if_running( + *, + streaming_service: Any, + content_builder: Any | None, + timer: ActivityTimer, + now_ns: int | None = None, +) -> str | None: + """Complete exception cleanup once without changing paused/terminal timers.""" + snapshot = timer.complete_if_running(now_ns=now_ns) + if snapshot is None: + return None + return emit_activity_timing_frame( + streaming_service=streaming_service, + content_builder=content_builder, + snapshot=snapshot, + ) diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/state.py b/surfsense_backend/app/tasks/chat/streaming/relay/state.py index d8b02b3f7..3b6f22c9d 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/state.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/state.py @@ -2,18 +2,12 @@ from __future__ import annotations +from collections import deque from dataclasses import dataclass, field from typing import Any from app.services.streaming.types import ActivityData -from app.tasks.chat.streaming.handlers.tools.activity import ActivitySpec - -_TERMINAL_ACTIVITY_STATUSES = { - "completed", - "error", - "cancelled", - "interrupted", -} +from app.tasks.chat.streaming.relay.activity_journal import ActivityJournal @dataclass @@ -31,13 +25,7 @@ class AgentEventRelayState: accumulated_text: str = "" current_text_id: str | None = None - activity_counter: int = 0 - activity_id_by_run: dict[str, str] = field(default_factory=dict) - activity_snapshot_by_id: dict[str, ActivityData] = field(default_factory=dict) - activity_spec_by_id: dict[str, ActivitySpec] = field(default_factory=dict) - resumable_activity_ids_by_kind: dict[str, list[str]] = field(default_factory=dict) - open_phase_by_scope: dict[str, tuple[str, str]] = field(default_factory=dict) - terminal_activity_ids: set[str] = field(default_factory=set) + journal: ActivityJournal = field(default_factory=ActivityJournal) active_tool_depth: int = 0 current_reasoning_id: str | None = None pending_tool_call_chunks: list[dict[str, Any]] = field(default_factory=list) @@ -45,6 +33,7 @@ class AgentEventRelayState: file_path_by_run: dict[str, str] = field(default_factory=dict) index_to_meta: dict[int, dict[str, str]] = field(default_factory=dict) ui_tool_call_id_by_run: dict[str, str] = field(default_factory=dict) + resume_tool_call_ids: deque[str] = field(default_factory=deque) current_lc_tool_call_id: dict[str, str | None] = field( default_factory=lambda: {"value": None} ) @@ -91,58 +80,19 @@ class AgentEventRelayState: cls, *, initial_activities: list[ActivityData] | None = None, + resume_activity_id_by_tool_call: dict[str, str] | None = None, + resume_tool_call_ids: list[str] | None = None, ) -> AgentEventRelayState: - snapshots = { - activity["id"]: activity - for activity in initial_activities or [] - if activity.get("status") == "awaiting_approval" - } - resumable: dict[str, list[str]] = {} - for activity in sorted( - snapshots.values(), key=lambda item: (item["sequence"], item["id"]) - ): - resumable.setdefault(activity["kind"], []).append(activity["id"]) return cls( - activity_counter=max( - (activity["sequence"] for activity in snapshots.values()), default=0 + journal=ActivityJournal.resume( + activities=initial_activities, + activity_id_by_tool_call=resume_activity_id_by_tool_call, ), - activity_snapshot_by_id=snapshots, - resumable_activity_ids_by_kind=resumable, + resume_tool_call_ids=deque(resume_tool_call_ids or ()), ) - def next_activity_id(self, step_prefix: str) -> str: - self.activity_counter += 1 - return f"act_{step_prefix}_{self.activity_counter}" - - def transition_activity( - self, - activity_id: str, - *, - status: str, - completed_at: str | None = None, - details: list[str] | None = None, - ) -> ActivityData | None: - current = self.activity_snapshot_by_id.get(activity_id) - spec = self.activity_spec_by_id.get(activity_id) - if current is None or spec is None: - return None - if activity_id in self.terminal_activity_ids: - return current - started_at = current["startedAt"] - integration = current.get("integration") - snapshot = spec.snapshot( - activity_id=activity_id, - sequence=current["sequence"], - status=status, # type: ignore[arg-type] - started_at=started_at, - completed_at=completed_at, - details=details if details is not None else current.get("details"), - integration=integration, + def consume_resume_tool_call_id(self) -> str | None: + """Consume the next persisted call identity for a replayed HITL tool.""" + return ( + self.resume_tool_call_ids.popleft() if self.resume_tool_call_ids else None ) - self.activity_snapshot_by_id[activity_id] = snapshot - if status in _TERMINAL_ACTIVITY_STATUSES: - self.terminal_activity_ids.add(activity_id) - for scope, (_, open_id) in list(self.open_phase_by_scope.items()): - if open_id == activity_id: - self.open_phase_by_scope.pop(scope, None) - return snapshot diff --git a/surfsense_backend/app/tasks/chat/streaming/shared/stream_result.py b/surfsense_backend/app/tasks/chat/streaming/shared/stream_result.py index aafe1d228..f8e6ab8c3 100644 --- a/surfsense_backend/app/tasks/chat/streaming/shared/stream_result.py +++ b/surfsense_backend/app/tasks/chat/streaming/shared/stream_result.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from typing import Any @@ -25,19 +26,18 @@ class StreamResult: commit_gate_passed: bool = True commit_gate_reason: str = "" # Pre-allocated assistant ``new_chat_messages.id`` for this turn, captured by - # ``persist_assistant_shell`` right after the user row is persisted. ``None`` - # for the legacy/anonymous code paths that don't opt in to server-side - # ``ContentPart[]`` projection. + # ``persist_assistant_shell`` right after the user row is persisted. assistant_message_id: int | None = None - # In-memory mirror of the FE's assistant-ui ``ContentPartsState``, populated - # by the lifecycle methods called from the streaming event loop at each - # ``streaming_service.format_*`` yield site. Snapshot in the streaming - # ``finally`` to produce the rich JSONB persisted by - # ``finalize_assistant_turn``. ``repr=False`` keeps the log-on-error path - # (``StreamResult`` is logged in some error branches) from dumping a - # potentially-large parts list. + # Server-side content-part projection populated alongside SSE emission. + # Snapshot in ``finally`` for ``finalize_assistant_turn``. ``repr=False`` + # prevents error logs from dumping a potentially large parts list. content_builder: Any | None = field(default=None, repr=False) activity_state: Any | None = field(default=None, repr=False) + # Reads the authoritative LangGraph checkpoint during disconnect cleanup. + # The in-memory event loop may be cancelled before it observes pending HITL. + load_agent_state: Callable[[], Awaitable[Any]] | None = field( + default=None, repr=False + ) activity_timer: ActivityTimer = field( default_factory=ActivityTimer.start, repr=False ) diff --git a/surfsense_backend/scripts/backfill_podcast_artifacts.py b/surfsense_backend/scripts/backfill_podcast_artifacts.py index b6b3d171c..73203be97 100644 --- a/surfsense_backend/scripts/backfill_podcast_artifacts.py +++ b/surfsense_backend/scripts/backfill_podcast_artifacts.py @@ -194,7 +194,9 @@ async def backfill(*, apply: bool) -> None: files=[ ArtifactFileInput( data=audio, - filename=primary_filename(title, extension="mp3", fallback="podcast"), + filename=primary_filename( + title, extension="mp3", fallback="podcast" + ), mime_type="audio/mpeg", role=ArtifactFileRole.PRIMARY, ) @@ -225,7 +227,9 @@ async def backfill(*, apply: bool) -> None: # before the column landed): stamp the link, don't re-create. strays = [row for row in rows if row.id in artifacts] - print(f"{len(rows)} pending podcast row(s); {len(artifacts)} already converted.") + print( + f"{len(rows)} pending podcast row(s); {len(artifacts)} already converted." + ) if not apply: print(f"Dry run: {len(pending)} row(s) would become Artifacts.") if strays: diff --git a/surfsense_backend/scripts/backfill_video_artifacts.py b/surfsense_backend/scripts/backfill_video_artifacts.py index d8cc52a06..e9db4a86c 100644 --- a/surfsense_backend/scripts/backfill_video_artifacts.py +++ b/surfsense_backend/scripts/backfill_video_artifacts.py @@ -235,7 +235,9 @@ async def backfill(*, apply: bool) -> None: for row in strays: await session.execute( - text("UPDATE video_presentations SET artifact_id = :aid WHERE id = :id"), + text( + "UPDATE video_presentations SET artifact_id = :aid WHERE id = :id" + ), {"aid": artifacts[row.id][0], "id": row.id}, ) if strays: diff --git a/surfsense_backend/tests/integration/artifacts/test_editor_content.py b/surfsense_backend/tests/integration/artifacts/test_editor_content.py index 8662367c2..71dd7241c 100644 --- a/surfsense_backend/tests/integration/artifacts/test_editor_content.py +++ b/surfsense_backend/tests/integration/artifacts/test_editor_content.py @@ -29,12 +29,17 @@ def artifact_backend(monkeypatch, patched_embed_texts): async def test_save_document_rejects_artifact_with_conflict( - db_session, db_workspace, db_user, artifact_backend, monkeypatch + db_session, + db_workspace, + db_user, + artifact_thread, + artifact_backend, + monkeypatch, ): saved = await save_artifact( db_session, workspace_id=db_workspace.id, - thread_id=1, + thread_id=artifact_thread.id, tool_call_id="markdown", title="Current / notes", markdown_representation="# Current notes", diff --git a/surfsense_backend/tests/integration/artifacts/test_office_artifacts.py b/surfsense_backend/tests/integration/artifacts/test_office_artifacts.py index 333d8a778..e2b3304b1 100644 --- a/surfsense_backend/tests/integration/artifacts/test_office_artifacts.py +++ b/surfsense_backend/tests/integration/artifacts/test_office_artifacts.py @@ -46,11 +46,11 @@ def _office_bytes(format_name: str, label: str) -> bytes: return output.getvalue() -def _runtime(format_name: str) -> ToolRuntime: +def _runtime(format_name: str, thread_id: int) -> ToolRuntime: return ToolRuntime( state={}, context=None, - config={"configurable": {"thread_id": f"77::task:{format_name}"}}, + config={"configurable": {"thread_id": f"{thread_id}::task:{format_name}"}}, stream_writer=None, tool_call_id=format_name, store=None, @@ -93,6 +93,7 @@ async def _verify( async def test_office_tool_create_revise_editor_contract_and_purge( db_session, db_workspace, + artifact_thread, patched_embed_texts, monkeypatch, format_name, @@ -141,7 +142,7 @@ async def test_office_tool_create_revise_editor_contract_and_purge( monkeypatch.setattr(load_source_tool, "shielded_async_session", session_context) monkeypatch.setattr(load_source_tool, "get_storage_backend", lambda *_: backend) tool = save_artifact_tool.create_save_artifact_tool(db_workspace.id) - runtime = _runtime(format_name) + runtime = _runtime(format_name, artifact_thread.id) sandbox.files[primary_path] = _office_bytes(format_name, "changed-before-save") rejected = await tool.coroutine( diff --git a/surfsense_backend/tests/integration/artifacts/test_tool.py b/surfsense_backend/tests/integration/artifacts/test_tool.py index 55729c41e..e332a282a 100644 --- a/surfsense_backend/tests/integration/artifacts/test_tool.py +++ b/surfsense_backend/tests/integration/artifacts/test_tool.py @@ -20,11 +20,11 @@ from .test_service import MemoryBackend pytestmark = pytest.mark.integration -def _runtime() -> ToolRuntime: +def _runtime(thread_id: int) -> ToolRuntime: return ToolRuntime( state={}, context=None, - config={"configurable": {"thread_id": "77::task:call-tool"}}, + config={"configurable": {"thread_id": f"{thread_id}::task:call-tool"}}, stream_writer=None, tool_call_id="call-tool", store=None, @@ -32,7 +32,7 @@ def _runtime() -> ToolRuntime: async def test_tool_persists_and_indexes_artifact_document_immediately( - db_session, db_workspace, patched_embed_texts, monkeypatch + db_session, db_workspace, artifact_thread, patched_embed_texts, monkeypatch ): del patched_embed_texts backend = MemoryBackend() @@ -51,7 +51,7 @@ async def test_tool_persists_and_indexes_artifact_document_immediately( command = await tool.coroutine( title="Legacy artifact", markdown_representation="# Legacy artifact\n\nimmediate-search-hit-term", - runtime=_runtime(), + runtime=_runtime(artifact_thread.id), ) payload = json.loads(command.update["messages"][0].content) @@ -80,7 +80,7 @@ async def test_tool_persists_and_indexes_artifact_document_immediately( async def test_load_artifact_source_restores_the_current_source( - db_session, db_workspace, patched_embed_texts, monkeypatch + db_session, db_workspace, artifact_thread, patched_embed_texts, monkeypatch ): del patched_embed_texts backend = MemoryBackend() @@ -91,7 +91,7 @@ async def test_load_artifact_source_restores_the_current_source( saved = await save_artifact( db_session, workspace_id=db_workspace.id, - thread_id=77, + thread_id=artifact_thread.id, tool_call_id="create", title="Restorable", markdown_representation="# Restorable", @@ -128,7 +128,10 @@ async def test_load_artifact_source_restores_the_current_source( tool = load_source_tool.create_load_artifact_source_tool( workspace_id=db_workspace.id ) - loaded = await tool.coroutine(artifact_id=saved.artifact_id, runtime=_runtime()) + loaded = await tool.coroutine( + artifact_id=saved.artifact_id, + runtime=_runtime(artifact_thread.id), + ) expected_path = f"/workspace/artifact-{saved.artifact_id}-out.py" assert loaded == { diff --git a/surfsense_backend/tests/integration/artifacts/test_xlsx_artifacts.py b/surfsense_backend/tests/integration/artifacts/test_xlsx_artifacts.py index 92a5e8676..299b37d91 100644 --- a/surfsense_backend/tests/integration/artifacts/test_xlsx_artifacts.py +++ b/surfsense_backend/tests/integration/artifacts/test_xlsx_artifacts.py @@ -119,7 +119,6 @@ async def test_xlsx_tool_create_revise_without_preview( monkeypatch.setattr( service, "knowledge_store_enabled_for", AsyncMock(return_value=False) ) - monkeypatch.setattr(service, "index_artifact", AsyncMock()) monkeypatch.setattr(save_artifact_tool, "get_registry", get_registry) monkeypatch.setattr(save_artifact_tool, "shielded_async_session", session_context) monkeypatch.setattr(save_artifact_tool.app_config, "SECRET_KEY", SECRET) diff --git a/surfsense_backend/tests/integration/chat/test_resume_activity_journal.py b/surfsense_backend/tests/integration/chat/test_resume_activity_journal.py new file mode 100644 index 000000000..22f2bc76d --- /dev/null +++ b/surfsense_backend/tests/integration/chat/test_resume_activity_journal.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession + +from app.db import ( + ChatVisibility, + NewChatMessage, + NewChatMessageRole, + NewChatThread, + User, + Workspace, +) +from app.services.new_streaming_service import VercelStreamingService +from app.tasks.chat.content_builder import AssistantContentBuilder +from app.tasks.chat.persistence import load_assistant_message_for_turn +from app.tasks.chat.streaming.flows.resume_chat.assistant_shell import ( + _resumable_journal_from_content, + order_resume_tool_call_ids, +) +from app.tasks.chat.streaming.handlers.tool_start import iter_tool_start_frames +from app.tasks.chat.streaming.handlers.tools.activity import resolve_tool_activity +from app.tasks.chat.streaming.relay.state import AgentEventRelayState + +pytestmark = pytest.mark.integration + + +@pytest.mark.asyncio +async def test_old_paused_turn_restores_two_same_kind_tools_by_exact_call_id( + db_session: AsyncSession, + db_user: User, + db_workspace: Workspace, +) -> None: + thread = NewChatThread( + title="Resume identity", + workspace_id=db_workspace.id, + created_by_id=db_user.id, + visibility=ChatVisibility.PRIVATE, + ) + db_session.add(thread) + await db_session.flush() + + spec = resolve_tool_activity("write_file", subagent_type=None) + first = spec.snapshot( + activity_id="act-first", + sequence=1, + status="awaiting_approval", + started_at="2026-01-01T00:00:00+00:00", + ) + second = spec.snapshot( + activity_id="act-second", + sequence=2, + status="awaiting_approval", + started_at="2026-01-01T00:00:01+00:00", + ) + paused = NewChatMessage( + thread_id=thread.id, + role=NewChatMessageRole.ASSISTANT, + turn_id="paused-turn", + content=[ + { + "type": "data-activities", + "data": { + "activities": [first, second], + "timing": {"status": "paused", "activeDurationMs": 1200}, + }, + }, + { + "type": "tool-call", + "toolCallId": "ui-first", + "langchainToolCallId": "lc-first", + "toolName": "write_file", + "metadata": {"activityId": first["id"]}, + }, + { + "type": "tool-call", + "toolCallId": "ui-second", + "langchainToolCallId": "lc-second", + "toolName": "write_file", + "metadata": {"activityId": second["id"]}, + }, + ], + ) + newer = [ + NewChatMessage( + thread_id=thread.id, + role=NewChatMessageRole.ASSISTANT, + turn_id=f"newer-turn-{index}", + content=[{"type": "text", "text": "newer"}], + ) + for index in range(25) + ] + db_session.add_all([paused, *newer]) + await db_session.flush() + + resolved = await load_assistant_message_for_turn( + db_session, + chat_id=thread.id, + turn_id="paused-turn", + ) + assert resolved is not None + journal = _resumable_journal_from_content(resolved.content) + + assert resolved.id == paused.id + assert [activity["id"] for activity in journal.activities] == [ + "act-first", + "act-second", + ] + assert journal.activity_id_by_tool_call == { + "lc-first": "act-first", + "ui-first": "act-first", + "lc-second": "act-second", + "ui-second": "act-second", + } + assert journal.tool_call_ids == ["lc-first", "lc-second"] + assert order_resume_tool_call_ids(journal, ["ui-second", "ui-first"]) == [ + "lc-second", + "lc-first", + ] + + relay_state = AgentEventRelayState.for_invocation( + initial_activities=journal.activities, + resume_activity_id_by_tool_call=journal.activity_id_by_tool_call, + resume_tool_call_ids=journal.tool_call_ids, + ) + builder = AssistantContentBuilder() + for run_id in ("fresh-first", "fresh-second"): + list( + iter_tool_start_frames( + { + "name": "write_file", + "run_id": run_id, + "data": {"input": {"file_path": f"{run_id}.md", "content": "x"}}, + }, + state=relay_state, + streaming_service=VercelStreamingService(), + content_builder=builder, + result=SimpleNamespace(write_attempted=False), + step_prefix="resume", + ) + ) + + assert relay_state.journal.id_by_run == { + "fresh-first": "act-first", + "fresh-second": "act-second", + } + assert not relay_state.resume_tool_call_ids + assert not relay_state.journal.resume_id_by_tool_call diff --git a/surfsense_backend/tests/integration/podcasts/test_render_task.py b/surfsense_backend/tests/integration/podcasts/test_render_task.py index 5c479a531..a71d32682 100644 --- a/surfsense_backend/tests/integration/podcasts/test_render_task.py +++ b/surfsense_backend/tests/integration/podcasts/test_render_task.py @@ -34,7 +34,13 @@ async def _primary_key(db_session, artifact_id: int) -> str: async def test_render_marks_ready_and_records_the_artifact( - db_session, db_workspace, make_podcast, bind_task_session, fake_tts, fake_merge, fake_storage + db_session, + db_workspace, + make_podcast, + bind_task_session, + fake_tts, + fake_merge, + fake_storage, ): podcast = await make_podcast( workspace_id=db_workspace.id, status=PodcastStatus.RENDERING diff --git a/surfsense_backend/tests/unit/artifacts/test_verification_receipt.py b/surfsense_backend/tests/unit/artifacts/test_verification_receipt.py index 806533474..2c96c6ea8 100644 --- a/surfsense_backend/tests/unit/artifacts/test_verification_receipt.py +++ b/surfsense_backend/tests/unit/artifacts/test_verification_receipt.py @@ -65,20 +65,26 @@ async def test_receipts_for_multiple_artifacts_do_not_overwrite_each_other(): await write_receipt(session, docx_receipt, SECRET) await write_receipt(session, pdf_receipt, SECRET) - assert await read_receipt( - session, - SECRET, - workspace_id=WORKSPACE_ID, - primary_path=docx_receipt.primary_path, - now=100, - ) == docx_receipt - assert await read_receipt( - session, - SECRET, - workspace_id=WORKSPACE_ID, - primary_path=pdf_receipt.primary_path, - now=100, - ) == pdf_receipt + assert ( + await read_receipt( + session, + SECRET, + workspace_id=WORKSPACE_ID, + primary_path=docx_receipt.primary_path, + now=100, + ) + == docx_receipt + ) + assert ( + await read_receipt( + session, + SECRET, + workspace_id=WORKSPACE_ID, + primary_path=pdf_receipt.primary_path, + now=100, + ) + == pdf_receipt + ) async def test_receipt_rejects_tampered_payload(): diff --git a/surfsense_backend/tests/unit/capabilities/access/test_agent_tools.py b/surfsense_backend/tests/unit/capabilities/access/test_agent_tools.py index 08ac2b505..d9b15ed33 100644 --- a/surfsense_backend/tests/unit/capabilities/access/test_agent_tools.py +++ b/surfsense_backend/tests/unit/capabilities/access/test_agent_tools.py @@ -139,6 +139,28 @@ async def test_capability_activity_descriptor_reaches_structured_tool_metadata(i } +async def test_activity_descriptor_round_trips_local_lifecycle_policy(): + metadata = ActivityDescriptor( + active_title="Planning work", + completed_title="Planned work", + category="action", + icon_key="list-todo", + kind="write_todos", + lifecycle="phase", + visibility="hide", + ).as_metadata() + + assert ActivityDescriptor.from_metadata(metadata) == ActivityDescriptor( + active_title="Planning work", + completed_title="Planned work", + category="action", + icon_key="list-todo", + kind="write_todos", + lifecycle="phase", + visibility="hide", + ) + + async def test_tool_runs_executor_and_returns_serialized_output(isolate): cap = _capability(name="web.scrape", output=_EchoOutput(echoed="hi there")) tools = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap]) diff --git a/surfsense_backend/tests/unit/scripts/test_backfill_podcast_artifacts.py b/surfsense_backend/tests/unit/scripts/test_backfill_podcast_artifacts.py index 908766281..fa54e241f 100644 --- a/surfsense_backend/tests/unit/scripts/test_backfill_podcast_artifacts.py +++ b/surfsense_backend/tests/unit/scripts/test_backfill_podcast_artifacts.py @@ -12,7 +12,11 @@ pytestmark = pytest.mark.unit def test_tool_parts_matches_only_podcast_tool_calls_with_a_result(): content = [ {"type": "text", "text": "hi"}, - {"type": "tool-call", "toolName": "generate_image", "result": {"artifact_id": 1}}, + { + "type": "tool-call", + "toolName": "generate_image", + "result": {"artifact_id": 1}, + }, {"type": "tool-call", "toolName": "generate_podcast"}, # no result { "type": "tool-call", diff --git a/surfsense_backend/tests/unit/scripts/test_backfill_video_artifacts.py b/surfsense_backend/tests/unit/scripts/test_backfill_video_artifacts.py index 1df566799..02155c0b0 100644 --- a/surfsense_backend/tests/unit/scripts/test_backfill_video_artifacts.py +++ b/surfsense_backend/tests/unit/scripts/test_backfill_video_artifacts.py @@ -12,7 +12,11 @@ pytestmark = pytest.mark.unit def test_tool_parts_matches_only_video_tool_calls_with_a_result(): content = [ {"type": "text", "text": "hi"}, - {"type": "tool-call", "toolName": "generate_image", "result": {"artifact_id": 1}}, + { + "type": "tool-call", + "toolName": "generate_image", + "result": {"artifact_id": 1}, + }, {"type": "tool-call", "toolName": "generate_video_presentation"}, # no result { "type": "tool-call", diff --git a/surfsense_backend/tests/unit/tasks/chat/streaming/test_save_artifact_contract.py b/surfsense_backend/tests/unit/tasks/chat/streaming/test_save_artifact_contract.py index ff126d2e0..da0202657 100644 --- a/surfsense_backend/tests/unit/tasks/chat/streaming/test_save_artifact_contract.py +++ b/surfsense_backend/tests/unit/tasks/chat/streaming/test_save_artifact_contract.py @@ -2,7 +2,17 @@ from app.tasks.chat.streaming.handlers.tools.activity import resolve_tool_activi def test_save_artifact_uses_canonical_activity_presentation(): - activity = resolve_tool_activity("save_artifact", subagent_type=None) + activity = resolve_tool_activity( + "save_artifact", + subagent_type=None, + trusted_descriptor={ + "active_title": "Preparing the file", + "completed_title": "Presented file", + "category": "artifact", + "icon_key": "file-output", + "kind": "save_artifact", + }, + ) assert activity.active_title == "Preparing the file" assert activity.completed_title == "Presented file" diff --git a/surfsense_backend/tests/unit/tasks/chat/test_activity_contract.py b/surfsense_backend/tests/unit/tasks/chat/test_activity_contract.py index 492de454b..8f96371c8 100644 --- a/surfsense_backend/tests/unit/tasks/chat/test_activity_contract.py +++ b/surfsense_backend/tests/unit/tasks/chat/test_activity_contract.py @@ -1,8 +1,12 @@ from __future__ import annotations +import asyncio import json +from pathlib import Path from types import SimpleNamespace +import pytest + from app.agents.chat.multi_agent_chat.shared.tools.mcp.tool import ( _mcp_activity_descriptor, ) @@ -10,19 +14,25 @@ from app.services.new_streaming_service import VercelStreamingService from app.services.streaming.types import ActivityTimingData from app.tasks.chat.content_builder import AssistantContentBuilder from app.tasks.chat.streaming.activity_timing import ActivityTimer +from app.tasks.chat.streaming.agent.event_loop import stream_agent_events from app.tasks.chat.streaming.flows.resume_chat.assistant_shell import ( - _resumable_journal_from_messages, + _resumable_journal_from_content, +) +from app.tasks.chat.streaming.flows.shared.assistant_finalize import ( + finalize_assistant_message, ) from app.tasks.chat.streaming.flows.shared.first_frames import iter_initial_frames from app.tasks.chat.streaming.handlers.custom_events import handle_activity_progress from app.tasks.chat.streaming.handlers.tool_end import iter_tool_end_frames -from app.tasks.chat.streaming.handlers.tool_start import ( - _artifact_instruction_type, - iter_tool_start_frames, -) +from app.tasks.chat.streaming.handlers.tool_start import iter_tool_start_frames from app.tasks.chat.streaming.handlers.tools.activity import resolve_tool_activity -from app.tasks.chat.streaming.relay.activity_sse import emit_activity_timing_frame +from app.tasks.chat.streaming.relay.activity_sse import ( + emit_activity_timing_frame, + emit_completed_activity_timing_frame, + emit_completed_activity_timing_frame_if_running, +) from app.tasks.chat.streaming.relay.state import AgentEventRelayState +from app.tasks.chat.streaming.shared.stream_result import StreamResult def _payload(frame: str) -> dict: @@ -76,18 +86,6 @@ def test_initial_frames_carry_turn_identity_without_timing_copy() -> None: ] turn_info = frames[2]["data"] assert turn_info == {"chat_turn_id": "12:activity-clock"} - assert all(frame["type"] != "data-thinking-step" for frame in frames) - - -def test_artifact_instruction_type_comes_from_the_structured_tool() -> None: - assert ( - _artifact_instruction_type( - "load_artifact_instructions", - {"artifact_type": "pdf"}, - ) - == "pdf" - ) - assert _artifact_instruction_type("execute", {"artifact_type": "pdf"}) is None def test_backend_owns_activity_copy_and_phase_lifecycle() -> None: @@ -188,7 +186,20 @@ def test_backend_owns_activity_copy_and_phase_lifecycle() -> None: verify_frames = [ _payload(frame) for frame in iter_tool_start_frames( - {"name": "verify_artifact", "run_id": "verify-1", "data": {"input": {}}}, + { + "name": "verify_artifact", + "run_id": "verify-1", + "metadata": { + "activity_descriptor": { + "active_title": "Checking the artifact", + "completed_title": "Checked the artifact", + "category": "artifact", + "icon_key": "badge-check", + "kind": "verify_artifact", + } + }, + "data": {"input": {}}, + }, state=state, streaming_service=service, content_builder=builder, @@ -292,7 +303,7 @@ def test_unknown_tools_are_generic_and_internal_tools_are_hidden() -> None: assert all(frame["type"] != "data-activity" for frame in hidden) -def test_backend_assigns_specific_icons_and_safe_fallbacks() -> None: +def test_localized_native_descriptor_inventory_and_safe_fallbacks() -> None: expected_icons = { "read_file": "file-text", "write_file": "file-plus", @@ -312,34 +323,49 @@ def test_backend_assigns_specific_icons_and_safe_fallbacks() -> None: "read_sandbox_file": "file-text", "verify_artifact": "badge-check", "save_artifact": "file-output", - "save_document": "file-output", "generate_image": "image", - "display_image": "image", "generate_podcast": "microphone", "generate_video_presentation": "film", "search_knowledge_base": "library", "ask_knowledge_base": "library", - "scrape_webpage": "scan-text", - "google_search.scrape": "search", - "web.crawl": "scan-text", - "link_preview": "external-link", - "multi_link_preview": "external-link", "create_calendar_event": "calendar", "update_calendar_event": "calendar", "delete_calendar_event": "calendar", "search_calendar_events": "calendar", "create_automation": "workflow", "update_memory": "brain", + "get_connected_accounts": "search", } for tool_name, icon_key in expected_icons.items(): - spec = resolve_tool_activity(tool_name, subagent_type=None) + spec = resolve_tool_activity( + tool_name, + subagent_type=None, + trusted_descriptor={ + "active_title": "Working", + "completed_title": "Worked", + "category": "action", + "icon_key": icon_key, + "kind": tool_name, + }, + ) assert spec.icon_key == icon_key unknown = resolve_tool_activity("dynamic_unknown_tool", subagent_type=None) assert unknown.icon_key == "tool" - service = resolve_tool_activity("youtube.scrape", subagent_type=None) + service = resolve_tool_activity( + "youtube.scrape", + subagent_type=None, + trusted_descriptor={ + "active_title": "Reviewing video", + "completed_title": "Reviewed video", + "category": "research", + "icon_key": "youtube", + "kind": "youtube.scrape", + "integration_key": "youtube", + }, + ) snapshot = service.snapshot( activity_id="act_youtube", sequence=1, @@ -349,6 +375,96 @@ def test_backend_assigns_specific_icons_and_safe_fallbacks() -> None: assert snapshot["integration"] == {"source": "native", "key": "youtube"} +def test_visible_native_tools_declare_descriptors_at_their_definition() -> None: + backend_root = Path(__file__).parents[4] + inventory = { + "app/agents/chat/multi_agent_chat/shared/middleware/filesystem/middleware/middleware.py": { + "glob", + "grep", + }, + "app/agents/chat/multi_agent_chat/shared/middleware/todos.py": {"write_todos"}, + **{ + f"app/agents/chat/multi_agent_chat/shared/middleware/filesystem/tools/{name}/index.py": { + name + } + for name in ( + "edit_file", + "execute_code", + "list_tree", + "ls", + "mkdir", + "move_file", + "read_file", + "rm", + "rmdir", + "write_file", + ) + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/generate_image.py": { + "generate_image" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/load_artifact_source.py": { + "load_artifact_source" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/podcast.py": { + "generate_podcast" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/sandbox.py": { + "read_sandbox_file" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/save_artifact.py": { + "save_artifact" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/verify_artifact.py": { + "verify_artifact" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/deliverables/tools/video_presentation.py": { + "generate_video_presentation" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/knowledge_base/ask_knowledge_base_tool.py": { + "ask_knowledge_base" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/knowledge_base/tools/search_knowledge_base.py": { + "search_knowledge_base" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/create_event.py": { + "create_calendar_event" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/delete_event.py": { + "delete_calendar_event" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/search_events.py": { + "search_calendar_events" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/calendar/update_event.py": { + "update_calendar_event" + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/mcp_discovery/tools/get_connected_accounts.py": { + "get_connected_accounts" + }, + "app/agents/chat/multi_agent_chat/main_agent/tools/automation/create.py": { + "create_automation" + }, + "app/agents/chat/multi_agent_chat/main_agent/tools/update_memory.py": { + "memory.personal", + "memory.team", + }, + "app/agents/chat/multi_agent_chat/subagents/builtins/memory/tools/update_memory.py": { + "memory.personal", + "memory.team", + }, + } + + for relative_path, tool_names in inventory.items(): + source = (backend_root / relative_path).read_text() + assert source.count('"activity_descriptor"') >= len(tool_names), relative_path + for tool_name in tool_names: + assert f'kind="{tool_name}"' in source or ( + f'"{tool_name}"' in source + and ("kind=TOOL_NAME" in source or "kind=tool_name" in source) + ), (relative_path, tool_name) + + def test_unknown_mcp_tool_uses_generic_activity_and_mcp_integration() -> None: frames = [ _payload(frame) @@ -464,7 +580,7 @@ def test_generated_native_tool_keeps_activity_id_through_result_lifecycle() -> N step_prefix="turn", ) ) - activity_id = state.activity_id_by_run["search-1"] + activity_id = state.journal.id_by_run["search-1"] list( iter_tool_end_frames( @@ -523,7 +639,11 @@ def test_resume_reuses_persisted_awaiting_activity_identity() -> None: status="awaiting_approval", started_at="2026-01-01T00:00:00+00:00", ) - state = AgentEventRelayState.for_invocation(initial_activities=[awaiting]) + state = AgentEventRelayState.for_invocation( + initial_activities=[awaiting], + resume_activity_id_by_tool_call={"lc-original-write": awaiting["id"]}, + resume_tool_call_ids=["lc-original-write"], + ) builder = AssistantContentBuilder() result = SimpleNamespace(write_attempted=False) @@ -550,6 +670,9 @@ def test_resume_reuses_persisted_awaiting_activity_identity() -> None: assert resumed["data"]["startedAt"] == "2026-01-01T00:00:00+00:00" tool_part = next(part for part in builder.snapshot() if part["type"] == "tool-call") assert tool_part["metadata"]["activityId"] == "act_original_7" + assert tool_part["langchainToolCallId"] == "lc-original-write" + assert not state.resume_tool_call_ids + assert not state.journal.resume_id_by_tool_call def test_resume_seed_loader_returns_paused_journal() -> None: @@ -568,66 +691,28 @@ def test_resume_seed_loader_returns_paused_journal() -> None: completed_at="2026-01-01T00:00:01+00:00", ) - seed = _resumable_journal_from_messages( + seed = _resumable_journal_from_content( [ - [ - {"type": "data-thinking-steps", "data": {"steps": []}}, - { - "type": "data-activities", - "data": { - "activities": [completed, awaiting], - "timing": {"status": "paused", "activeDurationMs": 2400}, - }, + { + "type": "data-activities", + "data": { + "activities": [completed, awaiting], + "timing": {"status": "paused", "activeDurationMs": 2400}, }, - ] + }, + { + "type": "tool-call", + "toolCallId": "call-write", + "toolName": "write_file", + "metadata": {"activityId": awaiting["id"]}, + }, ] ) assert seed.activities == [awaiting] assert seed.timing == {"status": "paused", "activeDurationMs": 2400} - - -def test_resume_seed_loader_uses_latest_snapshot_across_resume_messages() -> None: - spec = resolve_tool_activity("write_file", subagent_type=None) - awaiting = spec.snapshot( - activity_id="act_shared", - sequence=1, - status="awaiting_approval", - started_at="2026-01-01T00:00:00+00:00", - ) - completed = spec.snapshot( - activity_id="act_shared", - sequence=1, - status="completed", - started_at="2026-01-01T00:00:00+00:00", - completed_at="2026-01-01T00:00:01+00:00", - ) - - seed = _resumable_journal_from_messages( - [ - [ - { - "type": "data-activities", - "data": { - "activities": [completed], - "timing": {"status": "completed", "activeDurationMs": 3100}, - }, - } - ], - [ - { - "type": "data-activities", - "data": { - "activities": [awaiting], - "timing": {"status": "paused", "activeDurationMs": 2000}, - }, - } - ], - ] - ) - - assert seed.activities == [] - assert seed.timing is None + assert seed.activity_id_by_tool_call == {"call-write": awaiting["id"]} + assert seed.tool_call_ids == ["call-write"] def test_activity_timer_excludes_hitl_wait_and_resumes_accumulation() -> None: @@ -652,6 +737,126 @@ def test_activity_timer_excludes_hitl_wait_and_resumes_accumulation() -> None: } +def test_activity_timer_cleanup_completes_only_running_timers() -> None: + running = ActivityTimer.start(now_ns=1_000_000_000) + assert running.complete_if_running(now_ns=3_000_000_000) == { + "status": "completed", + "activeDurationMs": 2000, + } + assert running.complete_if_running(now_ns=9_000_000_000) is None + + paused = ActivityTimer.start(now_ns=1_000_000_000) + paused.pause(now_ns=2_000_000_000) + assert paused.complete_if_running(now_ns=9_000_000_000) is None + assert paused.status == "paused" + + +async def test_disconnect_cleanup_uses_pending_hitl_checkpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from app.tasks.chat import persistence + + persisted: dict = {} + + async def capture_finalize(**kwargs) -> None: + persisted.update(kwargs) + + monkeypatch.setattr(persistence, "finalize_assistant_turn", capture_finalize) + + pending_state = SimpleNamespace( + tasks=[ + SimpleNamespace( + interrupts=( + SimpleNamespace( + id="interrupt-1", + value={"type": "approval", "message": "Approve?"}, + ), + ) + ) + ], + values={}, + ) + + class Agent: + def __init__(self) -> None: + self.state_read_started = asyncio.Event() + self.state_reads = 0 + + async def astream_events(self, *_args, **_kwargs): + return + yield + + async def aget_state(self, _config): + self.state_reads += 1 + if self.state_reads == 1: + self.state_read_started.set() + await asyncio.Event().wait() + return pending_state + + builder = AssistantContentBuilder() + builder.on_activity_timing({"status": "running", "activeDurationMs": 1200}) + result = StreamResult( + turn_id="turn-hitl", + assistant_message_id=42, + content_builder=builder, + activity_timer=ActivityTimer.resume( + {"status": "paused", "activeDurationMs": 1200} + ), + ) + agent = Agent() + + async def consume_stream() -> None: + async for _ in stream_agent_events( + agent=agent, + config={"configurable": {}}, + input_data={}, + streaming_service=VercelStreamingService(), + result=result, + content_builder=builder, + ): + pass + + consumer = asyncio.create_task(consume_stream()) + await agent.state_read_started.wait() + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await consumer + + assert result.activity_timer.status == "running" + + await finalize_assistant_message( + stream_result=result, + chat_id=7, + workspace_id=9, + user_id="user-1", + accumulator=SimpleNamespace(), + log_prefix="test_disconnect", + ) + + journal = next( + part for part in persisted["content"] if part["type"] == "data-activities" + ) + assert result.is_interrupted is True + assert result.activity_timer.status == "paused" + assert journal["data"]["timing"]["status"] == "paused" + + +def test_activity_timer_excludes_multiple_hitl_waits_and_keeps_pause_strict() -> None: + timer = ActivityTimer.start(now_ns=0) + first_pause = timer.pause(now_ns=10_000_000_000) + with pytest.raises(ValueError, match="Only a running activity timer can pause"): + timer.pause(now_ns=20_000_000_000) + + timer = ActivityTimer.resume(first_pause, now_ns=310_000_000_000) + second_pause = timer.pause(now_ns=325_000_000_000) + timer = ActivityTimer.resume(second_pause, now_ns=925_000_000_000) + + assert timer.complete(now_ns=955_000_000_000) == { + "status": "completed", + "activeDurationMs": 55_000, + } + + def test_activity_builder_keeps_timing_and_rows_in_one_journal() -> None: builder = AssistantContentBuilder() builder.on_activity_timing( @@ -668,12 +873,28 @@ def test_activity_builder_keeps_timing_and_rows_in_one_journal() -> None: started_at="2026-01-01T00:00:00+00:00", ) ) + builder.on_activity_timing( + { + "status": "running", + "activeDurationMs": 1500, + } + ) + assert builder.snapshot()[0]["data"]["timing"] == { + "status": "paused", + "activeDurationMs": 2000, + } builder.on_activity_timing( { "status": "completed", "activeDurationMs": 5000, } ) + builder.on_activity_timing( + { + "status": "running", + "activeDurationMs": 6000, + } + ) journal = builder.snapshot()[0] assert journal["type"] == "data-activities" @@ -702,6 +923,53 @@ def test_activity_timing_wire_and_persistence_use_the_same_snapshot() -> None: assert builder.snapshot()[0]["data"]["timing"] == snapshot +def test_completed_timing_frame_is_strict_and_cleanup_is_idempotent() -> None: + service = VercelStreamingService() + builder = AssistantContentBuilder() + running = ActivityTimer.start(now_ns=1_000_000_000) + + frame = emit_completed_activity_timing_frame( + streaming_service=service, + content_builder=builder, + timer=running, + now_ns=3_000_000_000, + ) + + assert frame is not None + assert _payload(frame)["data"] == { + "status": "completed", + "activeDurationMs": 2000, + } + assert ( + emit_completed_activity_timing_frame_if_running( + streaming_service=service, + content_builder=builder, + timer=running, + now_ns=9_000_000_000, + ) + is None + ) + + paused = ActivityTimer.start(now_ns=1_000_000_000) + paused.pause(now_ns=2_000_000_000) + with pytest.raises(ValueError, match="Only a running activity timer can complete"): + emit_completed_activity_timing_frame( + streaming_service=service, + content_builder=builder, + timer=paused, + now_ns=9_000_000_000, + ) + assert ( + emit_completed_activity_timing_frame_if_running( + streaming_service=service, + content_builder=builder, + timer=paused, + now_ns=9_000_000_000, + ) + is None + ) + + def test_custom_progress_uses_allowlisted_details_not_raw_messages() -> None: service = VercelStreamingService() builder = AssistantContentBuilder() @@ -713,8 +981,8 @@ def test_custom_progress_uses_allowlisted_details_not_raw_messages() -> None: status="running", started_at="2026-01-01T00:00:00+00:00", ) - state.activity_spec_by_id[snapshot["id"]] = spec - state.activity_snapshot_by_id[snapshot["id"]] = snapshot + state.journal.spec_by_id[snapshot["id"]] = spec + state.journal.snapshot_by_id[snapshot["id"]] = snapshot frame = handle_activity_progress( { @@ -743,14 +1011,14 @@ def test_activity_state_preserves_terminal_monotonicity() -> None: status="running", started_at="2026-01-01T00:00:00+00:00", ) - state.activity_spec_by_id[running["id"]] = spec - state.activity_snapshot_by_id[running["id"]] = running + state.journal.spec_by_id[running["id"]] = spec + state.journal.snapshot_by_id[running["id"]] = running - awaiting = state.transition_activity(running["id"], status="awaiting_approval") + awaiting = state.journal.transition(running["id"], status="awaiting_approval") assert awaiting and awaiting["status"] == "awaiting_approval" assert "completedAt" not in awaiting - interrupted = state.transition_activity( + interrupted = state.journal.transition( running["id"], status="interrupted", completed_at="2026-01-01T00:01:00+00:00", @@ -758,6 +1026,6 @@ def test_activity_state_preserves_terminal_monotonicity() -> None: assert interrupted and interrupted["status"] == "interrupted" assert interrupted["completedAt"] assert ( - state.transition_activity(running["id"], status="running")["status"] + state.journal.transition(running["id"], status="running")["status"] == "interrupted" ) diff --git a/surfsense_backend/tests/unit/tasks/chat/test_activity_journal.py b/surfsense_backend/tests/unit/tasks/chat/test_activity_journal.py new file mode 100644 index 000000000..8d4a3feb1 --- /dev/null +++ b/surfsense_backend/tests/unit/tasks/chat/test_activity_journal.py @@ -0,0 +1,353 @@ +from app.tasks.chat.streaming.handlers.tools.activity import resolve_tool_activity +from app.tasks.chat.streaming.relay.activity_journal import ActivityJournal + + +def _activity(tool_name: str, *, lifecycle: str = "invocation"): + return resolve_tool_activity( + tool_name, + subagent_type=None, + trusted_descriptor={ + "active_title": "Working", + "completed_title": "Worked", + "category": "action", + "icon_key": "tool", + "kind": tool_name, + "lifecycle": lifecycle, + }, + ) + + +def _awaiting(activity_id: str, sequence: int): + return _activity("write_file").snapshot( + activity_id=activity_id, + sequence=sequence, + status="awaiting_approval", + started_at=f"2026-01-01T00:00:0{sequence}+00:00", + ) + + +def test_resume_binding_distinguishes_same_kind_activities() -> None: + first = _awaiting("act_first", 1) + second = _awaiting("act_second", 2) + journal = ActivityJournal.resume( + activities=[first, second], + activity_id_by_tool_call={ + "tool-call-first": first["id"], + "tool-call-second": second["id"], + }, + ) + spec = _activity("write_file") + + resumed_second = journal.begin_tool( + spec=spec, + run_id="run-second", + step_prefix="resume", + scope="root", + started_at="2026-01-01T00:01:00+00:00", + tool_call_id="tool-call-second", + langchain_tool_call_id=None, + integration=None, + ) + resumed_first = journal.begin_tool( + spec=spec, + run_id="run-first", + step_prefix="resume", + scope="root", + started_at="2026-01-01T00:01:01+00:00", + tool_call_id="tool-call-first", + langchain_tool_call_id=None, + integration=None, + ) + + assert resumed_second.activity_id == second["id"] + assert resumed_second.snapshots[-1]["startedAt"] == second["startedAt"] + assert resumed_first.activity_id == first["id"] + assert resumed_first.snapshots[-1]["startedAt"] == first["startedAt"] + + +def test_resume_prefers_authoritative_langchain_tool_call_id() -> None: + first = _awaiting("act_ui", 1) + second = _awaiting("act_lc", 2) + journal = ActivityJournal.resume( + activities=[first, second], + activity_id_by_tool_call={ + "ui-call": first["id"], + "lc-call": second["id"], + }, + ) + + resumed = journal.begin_tool( + spec=_activity("write_file"), + run_id="run", + step_prefix="resume", + scope="root", + started_at="2026-01-01T00:01:00+00:00", + tool_call_id="ui-call", + langchain_tool_call_id="lc-call", + integration=None, + ) + + assert resumed.activity_id == second["id"] + assert "lc-call" not in journal.resume_id_by_tool_call + + +def test_phase_reuses_identity_until_a_different_phase_starts() -> None: + journal = ActivityJournal() + planning = _activity("write_todos", lifecycle="phase") + research = _activity("web.crawl", lifecycle="phase") + + first = journal.begin_tool( + spec=planning, + run_id="plan-1", + step_prefix="turn", + scope="root", + started_at="2026-01-01T00:00:00+00:00", + tool_call_id="call-plan-1", + langchain_tool_call_id=None, + integration=None, + ) + repeated = journal.begin_tool( + spec=planning, + run_id="plan-2", + step_prefix="turn", + scope="root", + started_at="2026-01-01T00:00:01+00:00", + tool_call_id="call-plan-2", + langchain_tool_call_id=None, + integration=None, + ) + assert ( + journal.finish_tool( + run_id="plan-1", + status="completed", + completed_at="2026-01-01T00:00:01+00:00", + ).snapshot + is None + ) + assert ( + journal.finish_tool( + run_id="plan-2", + status="completed", + completed_at="2026-01-01T00:00:01+00:00", + ).snapshot + is None + ) + next_phase = journal.begin_tool( + spec=research, + run_id="research", + step_prefix="turn", + scope="root", + started_at="2026-01-01T00:00:02+00:00", + tool_call_id="call-research", + langchain_tool_call_id=None, + integration=None, + ) + + assert repeated.activity_id == first.activity_id + assert next_phase.snapshots[0]["id"] == first.activity_id + assert next_phase.snapshots[0]["status"] == "completed" + assert next_phase.activity_id != first.activity_id + + +def test_phase_close_waits_for_all_runs_and_preserves_late_error() -> None: + journal = ActivityJournal() + planning = _activity("write_todos", lifecycle="phase") + research = _activity("web.crawl", lifecycle="phase") + + first = journal.begin_tool( + spec=planning, + run_id="plan-1", + step_prefix="turn", + scope="root", + started_at="2026-01-01T00:00:00+00:00", + tool_call_id="call-plan-1", + langchain_tool_call_id=None, + integration=None, + ) + journal.begin_tool( + spec=planning, + run_id="plan-2", + step_prefix="turn", + scope="root", + started_at="2026-01-01T00:00:01+00:00", + tool_call_id="call-plan-2", + langchain_tool_call_id=None, + integration=None, + ) + next_phase = journal.begin_tool( + spec=research, + run_id="research", + step_prefix="turn", + scope="root", + started_at="2026-01-01T00:00:02+00:00", + tool_call_id="call-research", + langchain_tool_call_id=None, + integration=None, + ) + + assert [snapshot["id"] for snapshot in next_phase.snapshots] == [ + next_phase.activity_id + ] + assert ( + journal.finish_tool( + run_id="plan-2", + status="completed", + completed_at="2026-01-01T00:00:03+00:00", + ).snapshot + is None + ) + failed = journal.finish_tool( + run_id="plan-1", + status="error", + completed_at="2026-01-01T00:00:04+00:00", + ).snapshot + + assert failed is not None + assert failed["id"] == first.activity_id + assert failed["status"] == "error" + assert failed["completedAt"] == "2026-01-01T00:00:04+00:00" + + +def test_phase_outcomes_use_deterministic_severity() -> None: + journal = ActivityJournal() + phase = _activity("write_todos", lifecycle="phase") + for run_id in ("one", "two", "three"): + journal.begin_tool( + spec=phase, + run_id=run_id, + step_prefix="turn", + scope="root", + started_at="2026-01-01T00:00:00+00:00", + tool_call_id=f"call-{run_id}", + langchain_tool_call_id=None, + integration=None, + ) + + assert ( + journal.finish_tool( + run_id="one", + status="interrupted", + completed_at="2026-01-01T00:00:01+00:00", + ).snapshot + is None + ) + assert ( + journal.finish_tool( + run_id="two", + status="cancelled", + completed_at="2026-01-01T00:00:02+00:00", + ).snapshot + is None + ) + final = journal.finish_tool( + run_id="three", + status="error", + completed_at="2026-01-01T00:00:03+00:00", + ).snapshot + + assert final is not None + assert final["status"] == "error" + + +def test_successful_phase_closes_after_its_final_active_run() -> None: + journal = ActivityJournal() + phase = _activity("write_todos", lifecycle="phase") + started = journal.begin_tool( + spec=phase, + run_id="one", + step_prefix="turn", + scope="root", + started_at="2026-01-01T00:00:00+00:00", + tool_call_id="call-one", + langchain_tool_call_id=None, + integration=None, + ) + journal.begin_tool( + spec=phase, + run_id="two", + step_prefix="turn", + scope="root", + started_at="2026-01-01T00:00:01+00:00", + tool_call_id="call-two", + langchain_tool_call_id=None, + integration=None, + ) + + assert journal.complete_open_phases(completed_at="2026-01-01T00:00:02+00:00") == [] + assert ( + journal.finish_tool( + run_id="one", + status="completed", + completed_at="2026-01-01T00:00:03+00:00", + ).snapshot + is None + ) + closed = journal.finish_tool( + run_id="two", + status="completed", + completed_at="2026-01-01T00:00:04+00:00", + ).snapshot + + assert closed is not None + assert closed["id"] == started.activity_id + assert closed["status"] == "completed" + assert closed["completedAt"] == "2026-01-01T00:00:02+00:00" + + +def test_interrupt_running_force_closes_active_phase_runs() -> None: + journal = ActivityJournal() + started = journal.begin_tool( + spec=_activity("write_todos", lifecycle="phase"), + run_id="plan", + step_prefix="turn", + scope="root", + started_at="2026-01-01T00:00:00+00:00", + tool_call_id="call-plan", + langchain_tool_call_id=None, + integration=None, + ) + + interrupted = journal.interrupt_running(completed_at="2026-01-01T00:00:01+00:00") + + assert interrupted[0]["id"] == started.activity_id + assert interrupted[0]["status"] == "interrupted" + assert "plan" not in journal.id_by_run + assert started.activity_id not in journal.active_runs_by_activity + + +def test_terminal_activity_never_regresses() -> None: + journal = ActivityJournal() + started = journal.begin_tool( + spec=_activity("write_file"), + run_id="write", + step_prefix="turn", + scope="root", + started_at="2026-01-01T00:00:00+00:00", + tool_call_id="call-write", + langchain_tool_call_id=None, + integration=None, + ) + assert started.activity_id + + completed = journal.transition( + started.activity_id, + status="completed", + completed_at="2026-01-01T00:00:01+00:00", + ) + stale = journal.transition(started.activity_id, status="running") + + assert completed and completed["status"] == "completed" + assert stale == completed + + +def test_progress_updates_details_without_unpausing_activity() -> None: + awaiting = _awaiting("act_waiting", 1) + journal = ActivityJournal.resume(activities=[awaiting]) + spec = _activity("write_file") + journal.spec_by_id[awaiting["id"]] = spec + + updated = journal.update_current_progress("Reviewing sources (1/2)") + + assert updated is not None + assert updated["status"] == "awaiting_approval" + assert updated["details"] == ["Reviewing sources (1/2)"] diff --git a/surfsense_backend/tests/unit/tasks/chat/test_content_builder.py b/surfsense_backend/tests/unit/tasks/chat/test_content_builder.py index c590503a6..2523abc41 100644 --- a/surfsense_backend/tests/unit/tasks/chat/test_content_builder.py +++ b/surfsense_backend/tests/unit/tasks/chat/test_content_builder.py @@ -418,16 +418,21 @@ class TestActivities: assert snap[0]["data"]["activities"] == [_activity("act-1", 1)] assert snap[1] == {"type": "text", "text": "Hello"} - def test_snapshots_upsert_by_id_and_sort_by_sequence(self): + def test_snapshots_upsert_by_id_and_sort_by_sequence_then_id(self): b = AssistantContentBuilder() - b.on_activity(_activity("act-2", 2)) - b.on_activity(_activity("act-1", 1)) - b.on_activity(_activity("act-1", 1, status="completed", title="Done")) + b.on_activity(_activity("act-c", 2)) + b.on_activity(_activity("act-b", 1)) + b.on_activity(_activity("act-a", 1)) + b.on_activity(_activity("act-a", 1, status="completed", title="Done")) snap = b.snapshot() assert len([p for p in snap if p["type"] == "data-activities"]) == 1 activities = snap[0]["data"]["activities"] - assert [activity["id"] for activity in activities] == ["act-1", "act-2"] + assert [activity["id"] for activity in activities] == [ + "act-a", + "act-b", + "act-c", + ] assert activities[0]["status"] == "completed" def test_terminal_activity_cannot_regress(self): @@ -449,7 +454,7 @@ class TestActivities: class TestMarkInterrupted: - def test_running_activities_interrupt_but_approval_pauses_are_preserved(self): + def test_activity_lifecycle_is_not_invented_by_persistence_builder(self): b = AssistantContentBuilder() b.on_activity(_activity("running", 1)) b.on_activity(_activity("approval", 2, status="awaiting_approval")) @@ -457,8 +462,8 @@ class TestMarkInterrupted: b.mark_interrupted() activities = b.snapshot()[0]["data"]["activities"] - assert activities[0]["status"] == "interrupted" - assert activities[0]["completedAt"] + assert activities[0]["status"] == "running" + assert "completedAt" not in activities[0] assert activities[1]["status"] == "awaiting_approval" assert "completedAt" not in activities[1] @@ -522,6 +527,7 @@ class TestIsEmpty: b.on_activity(_activity("act-1", 1)) assert b.is_empty() + class TestSnapshotSemantics: def test_snapshot_is_deep_copied_so_mutations_do_not_leak(self): b = AssistantContentBuilder() diff --git a/surfsense_web/app/dashboard/[workspace_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[workspace_id]/new-chat/[[...chat_id]]/page.tsx index fa0f5f9a1..842d1f1b5 100644 --- a/surfsense_web/app/dashboard/[workspace_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[workspace_id]/new-chat/[[...chat_id]]/page.tsx @@ -23,7 +23,6 @@ import { mentionedDocumentsAtom, messageDocumentsMapAtom, } from "@/atoms/chat/mentioned-documents.atom"; -import { clearPlanOwnerRegistry } from "@/atoms/chat/plan-state.atom"; import { closeReportPanelAtom } from "@/atoms/chat/report-panel.atom"; import { closeEditorPanelAtom } from "@/atoms/editor/editor-panel.atom"; import { membersAtom } from "@/atoms/members/members-query.atoms"; @@ -275,7 +274,6 @@ export default function NewChatPage() { setMentionedDocuments([]); tokenUsageStore.clear(); setMessageDocumentsMap({}); - clearPlanOwnerRegistry(); closeReportPanel(); closeEditorPanel(); chatStreamStore.clearInactive(nextThreadId); diff --git a/surfsense_web/app/dashboard/[workspace_id]/playground/components/run-detail.tsx b/surfsense_web/app/dashboard/[workspace_id]/playground/components/run-detail.tsx index 4e18f179b..2df72abac 100644 --- a/surfsense_web/app/dashboard/[workspace_id]/playground/components/run-detail.tsx +++ b/surfsense_web/app/dashboard/[workspace_id]/playground/components/run-detail.tsx @@ -69,12 +69,7 @@ export function RunDetail({ if (!run) return null; return ( -
+
{run.error && (
{run.error} diff --git a/surfsense_web/atoms/chat/plan-state.atom.ts b/surfsense_web/atoms/chat/plan-state.atom.ts deleted file mode 100644 index 2436dd300..000000000 --- a/surfsense_web/atoms/chat/plan-state.atom.ts +++ /dev/null @@ -1,224 +0,0 @@ -/** - * Plan State Atom - * - * Tracks the latest state of each plan by title. - * When write_todos is called multiple times with the same title, - * only the FIRST component renders (stays fixed in position), - * subsequent calls just update the shared state. - */ - -import { atom } from "jotai"; - -export interface PlanTodo { - id: string; - content: string; - status: "pending" | "in_progress" | "completed" | "cancelled"; -} - -export interface PlanState { - id: string; - title: string; - todos: PlanTodo[]; - lastUpdated: number; - /** The toolCallId of the first component that rendered this plan */ - ownerToolCallId: string; -} - -/** - * SYNCHRONOUS ownership registry - prevents race conditions - * Only ONE plan allowed per conversation - first plan wins - */ -let firstPlanOwner: { toolCallId: string; title: string } | null = null; - -/** - * Register as owner of a plan SYNCHRONOUSLY - * ONE PLAN PER CONVERSATION: Only the first write_todos call renders. - * All subsequent calls update the state but don't render their own card. - */ -export function registerPlanOwner(title: string, toolCallId: string): boolean { - if (!firstPlanOwner) { - // First plan in this conversation - claim ownership - firstPlanOwner = { toolCallId, title }; - return true; - } - - // Check if we're the owner - return firstPlanOwner.toolCallId === toolCallId; -} - -/** - * Get the canonical title for a plan - * Returns the first plan's title if one exists, otherwise the provided title - */ -export function getCanonicalPlanTitle(title: string): string { - return firstPlanOwner?.title || title; -} - -/** - * Check if a plan already exists in this conversation - */ -export function hasPlan(): boolean { - return firstPlanOwner !== null; -} - -/** - * Get the first plan's info - */ -export function getFirstPlanInfo(): { toolCallId: string; title: string } | null { - return firstPlanOwner; -} - -/** - * Check if a toolCallId is the owner of the plan SYNCHRONOUSLY - */ -export function isPlanOwner(toolCallId: string): boolean { - return !firstPlanOwner || firstPlanOwner.toolCallId === toolCallId; -} - -/** - * Clear ownership registry (call when starting a new chat) - */ -export function clearPlanOwnerRegistry(): void { - firstPlanOwner = null; -} - -/** - * Map of plan title -> latest plan state - * Using title as key since it stays constant across updates - */ -export const planStatesAtom = atom>(new Map()); - -/** - * Input type for updating plan state - */ -export interface UpdatePlanInput { - id: string; - title: string; - todos: PlanTodo[]; - toolCallId: string; -} - -/** - * Helper atom to update a plan state - */ -export const updatePlanStateAtom = atom(null, (get, set, plan: UpdatePlanInput) => { - const states = new Map(get(planStatesAtom)); - - // Register ownership synchronously if not already done - registerPlanOwner(plan.title, plan.toolCallId); - - // Get the actual owner from the first plan - const ownerToolCallId = firstPlanOwner?.toolCallId || plan.toolCallId; - - // Always use the canonical (first) title for the plan key - const canonicalTitle = getCanonicalPlanTitle(plan.title); - - states.set(canonicalTitle, { - id: plan.id, - title: canonicalTitle, - todos: plan.todos, - lastUpdated: Date.now(), - ownerToolCallId, - }); - set(planStatesAtom, states); -}); - -/** - * Helper atom to get the latest plan state by title - */ -export const getPlanStateAtom = atom((get) => { - const states = get(planStatesAtom); - return (title: string) => states.get(title); -}); - -/** - * Helper atom to clear all plan states (useful when starting a new chat) - */ -export const clearPlanStatesAtom = atom(null, (get, set) => { - clearPlanOwnerRegistry(); - set(planStatesAtom, new Map()); -}); - -/** - * Hydrate plan state from persisted message content - * Call this when loading messages from the database to restore plan state - */ -export interface HydratePlanInput { - toolCallId: string; - result: { - id?: string; - title?: string; - todos?: Array<{ - id?: string; - content: string; - status: "pending" | "in_progress" | "completed" | "cancelled"; - }>; - }; -} - -export const hydratePlanStateAtom = atom(null, (get, set, plan: HydratePlanInput) => { - if (!plan.result?.todos || plan.result.todos.length === 0) return; - - const states = new Map(get(planStatesAtom)); - const title = plan.result.title || "Plan"; - - // Register this as the owner if no plan exists yet - registerPlanOwner(title, plan.toolCallId); - - // Get the canonical title - const canonicalTitle = getCanonicalPlanTitle(title); - const ownerToolCallId = firstPlanOwner?.toolCallId || plan.toolCallId; - - // Only set if this is newer or doesn't exist - const existing = states.get(canonicalTitle); - if (!existing) { - states.set(canonicalTitle, { - id: plan.result.id || `plan-${Date.now()}`, - title: canonicalTitle, - todos: plan.result.todos.map((t, i) => ({ - id: t.id || `todo-${i}`, - content: t.content, - status: t.status, - })), - lastUpdated: Date.now(), - ownerToolCallId, - }); - set(planStatesAtom, states); - } -}); - -/** - * Extract write_todos tool call data from message content - * Returns an array of { toolCallId, result } for each write_todos call found - */ -export function extractWriteTodosFromContent(content: unknown): HydratePlanInput[] { - if (!Array.isArray(content)) return []; - - const results: HydratePlanInput[] = []; - - for (const part of content) { - if ( - typeof part === "object" && - part !== null && - "type" in part && - (part as { type: string }).type === "tool-call" && - "toolName" in part && - (part as { toolName: string }).toolName === "write_todos" && - "toolCallId" in part && - "result" in part - ) { - const toolCall = part as { - toolCallId: string; - result: HydratePlanInput["result"]; - }; - if (toolCall.result) { - results.push({ - toolCallId: toolCall.toolCallId, - result: toolCall.result, - }); - } - } - } - - return results; -} diff --git a/surfsense_web/components/citations/run-citation.tsx b/surfsense_web/components/citations/run-citation.tsx index d32902627..21e2cc7f2 100644 --- a/surfsense_web/components/citations/run-citation.tsx +++ b/surfsense_web/components/citations/run-citation.tsx @@ -30,9 +30,7 @@ export const RunCitation: FC<{ runId: string }> = ({ runId }) => { - - - - - - )} - - {actionArray.length > 0 && ( -
- {actionArray.map((action) => ( - - ))} -
- )} - - - ); -}; diff --git a/surfsense_web/components/tool-ui/plan/schema.ts b/surfsense_web/components/tool-ui/plan/schema.ts deleted file mode 100644 index 56a5cd218..000000000 --- a/surfsense_web/components/tool-ui/plan/schema.ts +++ /dev/null @@ -1,171 +0,0 @@ -import { z } from "zod"; - -/** - * Todo item status - */ -export const TodoStatusSchema = z.enum(["pending", "in_progress", "completed", "cancelled"]); -export type TodoStatus = z.infer; - -/** - * Normalize various status string formats to the canonical TodoStatus - * Handles common variations from different sources: - * - Linear: Done, In Progress, Todo, Backlog, Cancelled - * - Jira: To Do, In Progress, Done, In Review, Reopened, Testing + statusCategory - * - ClickUp: Open, In Progress, Complete, Closed, Review - * - GitHub: open, closed - * - Airtable: Any custom field values - */ -export function normalizeStatus(status: unknown): TodoStatus { - if (typeof status !== "string") return "pending"; - - const normalized = status - .toLowerCase() - .trim() - .replace(/[\s_-]+/g, "_"); - - // Completed variations - // Sources: Linear (Done), Jira (Done), ClickUp (Complete, Closed), GitHub (closed) - if ( - normalized === "completed" || - normalized === "complete" || - normalized === "done" || - normalized === "finished" || - normalized === "closed" || - normalized === "resolved" || - normalized === "fixed" || - normalized === "shipped" || - normalized === "released" || - normalized === "merged" - ) { - return "completed"; - } - - // In progress variations - // Sources: Linear (In Progress), Jira (In Progress, In Review, Testing), ClickUp (In Progress, Review) - if ( - normalized === "in_progress" || - normalized === "inprogress" || - normalized === "started" || - normalized === "active" || - normalized === "working" || - normalized === "in_review" || - normalized === "inreview" || - normalized === "review" || - normalized === "reviewing" || - normalized === "testing" || - normalized === "in_testing" || - normalized === "qa" || - normalized === "in_qa" || - normalized === "doing" || - normalized === "wip" || - normalized === "work_in_progress" - ) { - return "in_progress"; - } - - // Cancelled variations - // Sources: Linear (Cancelled), Jira (Won't Fix, Duplicate) - if ( - normalized === "cancelled" || - normalized === "canceled" || - normalized === "dropped" || - normalized === "won't_fix" || - normalized === "wontfix" || - normalized === "wont_fix" || - normalized === "duplicate" || - normalized === "invalid" || - normalized === "rejected" || - normalized === "archived" || - normalized === "removed" || - normalized === "obsolete" - ) { - return "cancelled"; - } - - // Pending variations (default) - // Sources: Linear (Todo, Backlog), Jira (To Do, Reopened), ClickUp (Open), GitHub (open) - // Includes: "pending", "todo", "to_do", "backlog", "open", "new", "triage", "reopened", etc. - return "pending"; -} - -/** - * Single todo item in a plan - * Matches deepagents TodoListMiddleware output: { content, status } - * id is auto-generated if not provided - */ -export const PlanTodoSchema = z.object({ - id: z.string().optional(), - content: z.string(), - status: TodoStatusSchema, -}); - -export type PlanTodo = z.infer; - -/** - * Serializable plan schema for tool results - * Matches deepagents TodoListMiddleware output format - * id/title are auto-generated if not provided - */ -export const SerializablePlanSchema = z.object({ - id: z.string().optional(), - title: z.string().optional(), - todos: z.array(PlanTodoSchema).min(1), - maxVisibleTodos: z.number().optional(), - showProgress: z.boolean().optional(), -}); - -export type SerializablePlan = z.infer; - -/** - * Normalized plan with required fields (after auto-generation) - */ -export interface NormalizedPlan { - id: string; - title: string; - todos: Array<{ id: string; content: string; status: TodoStatus }>; - maxVisibleTodos?: number; - showProgress?: boolean; -} - -/** - * Parse and normalize a plan from tool result - * Auto-generates id/title if not provided (for deepagents compatibility) - */ -export function parseSerializablePlan(data: unknown): NormalizedPlan { - const result = SerializablePlanSchema.safeParse(data); - - if (!result.success) { - console.warn("Invalid plan data:", result.error.issues); - - // Try to extract basic info for fallback - const obj = (data && typeof data === "object" ? data : {}) as Record; - - return { - id: typeof obj.id === "string" ? obj.id : `plan-${Date.now()}`, - title: typeof obj.title === "string" ? obj.title : "Plan", - todos: Array.isArray(obj.todos) - ? obj.todos.map((t: unknown, i: number) => { - const todo = t as Record; - return { - id: typeof todo?.id === "string" ? todo.id : `todo-${i}`, - content: typeof todo?.content === "string" ? todo.content : "Task", - status: normalizeStatus(todo?.status), - }; - }) - : [{ id: "1", content: "No tasks", status: "pending" as const }], - }; - } - - // Normalize: add id/title if missing - return { - id: result.data.id || `plan-${Date.now()}`, - title: result.data.title || "Plan", - todos: result.data.todos.map((t, i) => ({ - id: t.id || `todo-${i}`, - content: t.content, - status: t.status, - })), - maxVisibleTodos: result.data.maxVisibleTodos, - showProgress: result.data.showProgress, - }; -} diff --git a/surfsense_web/components/tool-ui/podcast/player.tsx b/surfsense_web/components/tool-ui/podcast/player.tsx index 1e35fa97e..182e7d9bb 100644 --- a/surfsense_web/components/tool-ui/podcast/player.tsx +++ b/surfsense_web/components/tool-ui/podcast/player.tsx @@ -173,9 +173,7 @@ export function PodcastPlayer({ : Promise.resolve(null), ]); audioBlob = blob; - const parsed = details - ? publicPodcastDetailsSchema.safeParse(details) - : null; + const parsed = details ? publicPodcastDetailsSchema.safeParse(details) : null; lines = (parsed?.success ? (parsed.data.podcast_transcript ?? []) : []).map( (entry, turn) => ({ key: `turn-${turn}`, diff --git a/surfsense_web/components/tool-ui/video-presentation/generate-video-presentation.tsx b/surfsense_web/components/tool-ui/video-presentation/generate-video-presentation.tsx index b928eb41e..2d7a75358 100644 --- a/surfsense_web/components/tool-ui/video-presentation/generate-video-presentation.tsx +++ b/surfsense_web/components/tool-ui/video-presentation/generate-video-presentation.tsx @@ -416,9 +416,7 @@ export function StatusPoller({ shareToken?: string | null; }) { if (artifactId == null && presentationId == null) { - return ( -

Presentation not available

- ); + return

Presentation not available

; } return ( ; - -/** - * Loading state component - */ -function WriteTodosLoading() { - return ( -
-
- - Creating plan... -
-
- ); -} - -/** - * WriteTodos Tool UI Component - * - * Displays the agent's planning/todo list with a beautiful UI. - * Uses deepagents TodoListMiddleware output directly: { todos: [{ content, status }] } - * - * FIXED POSITION: When multiple write_todos calls happen in a conversation, - * only the FIRST component renders. Subsequent updates just update the - * shared state, and the first component reads from it. - */ -export const WriteTodosToolUI = ({ - args, - result, - status, - toolCallId, -}: ToolCallMessagePartProps) => { - const updatePlanState = useSetAtom(updatePlanStateAtom); - const planStates = useAtomValue(planStatesAtom); - - // Check if the THREAD is running - const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); - - // Use result if available, otherwise args (for streaming) - const data = result || args; - const hasTodos = data?.todos && data.todos.length > 0; - - // Fixed title for all plans in conversation - const planTitle = "Plan"; - - // SYNCHRONOUS ownership check - const isOwner = useMemo(() => { - return registerPlanOwner(planTitle, toolCallId); - }, [planTitle, toolCallId]); - - // Get canonical title - const canonicalTitle = useMemo(() => getCanonicalPlanTitle(planTitle), [planTitle]); - - // Register/update the plan state - useEffect(() => { - if (hasTodos) { - const normalizedPlan = parseSerializablePlan({ todos: data.todos }); - updatePlanState({ - id: normalizedPlan.id, - title: canonicalTitle, - todos: normalizedPlan.todos, - toolCallId, - }); - } - }, [data, hasTodos, canonicalTitle, updatePlanState, toolCallId]); - - // Get the current plan state - const currentPlanState = planStates.get(canonicalTitle); - - // If we're NOT the owner, render nothing - if (!isOwner) { - return null; - } - - // Loading state - if (status.type === "running" || status.type === "requires-action") { - if (hasTodos) { - const plan = parseSerializablePlan({ todos: data.todos }); - return ( -
- - - -
- ); - } - return ; - } - - // Incomplete/cancelled state - if (status.type === "incomplete") { - if (currentPlanState || hasTodos) { - const plan = currentPlanState || parseSerializablePlan({ todos: data?.todos || [] }); - return ( -
- - - -
- ); - } - return null; - } - - // Success - render the plan - const planToRender = - currentPlanState || (hasTodos ? parseSerializablePlan({ todos: data.todos }) : null); - if (!planToRender) { - return ; - } - - return ( -
- - - -
- ); -}; - -export { type WriteTodosData, WriteTodosSchema }; diff --git a/surfsense_web/features/artifacts-library/ui/artifact-card.tsx b/surfsense_web/features/artifacts-library/ui/artifact-card.tsx index 2eca218d0..d75aec54e 100644 --- a/surfsense_web/features/artifacts-library/ui/artifact-card.tsx +++ b/surfsense_web/features/artifacts-library/ui/artifact-card.tsx @@ -32,10 +32,7 @@ export function ArtifactCard({ {statusLabel ? ( <> -