chore(deps): bump actions/checkout from 4 to 6

Bumps [actions/checkout](https://github.com/actions/checkout) from 4 to 6.
- [Release notes](https://github.com/actions/checkout/releases)
- [Changelog](https://github.com/actions/checkout/blob/main/CHANGELOG.md)
- [Commits](https://github.com/actions/checkout/compare/v4...v6)

---
updated-dependencies:
- dependency-name: actions/checkout
  dependency-version: '6'
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>

.github/workflows/classic-autogpt-ci.yml

@@ -61,7 +61,7 @@ jobs:
       #   See: https://github.com/actions/runner/issues/598#issuecomment-2011890429
 
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
           submodules: true

.github/workflows/classic-autogpt-docker-cache-clean.yml

@@ -16,7 +16,7 @@ jobs:
         build-type: [release, dev]
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3

.github/workflows/classic-autogpt-docker-ci.yml

@@ -35,7 +35,7 @@ jobs:
         build-type: [release, dev]
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
@@ -101,7 +101,7 @@ jobs:
 
     steps:
       - name: Check out repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
         with:
           submodules: true
 

.github/workflows/classic-autogpt-docker-release.yml

@@ -20,7 +20,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Log in to Docker hub
         uses: docker/login-action@v3

.github/workflows/classic-autogpts-ci.yml

@@ -44,7 +44,7 @@ jobs:
       min-python-version: '3.10'
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
           submodules: true

.github/workflows/classic-benchmark-ci.yml

@@ -42,7 +42,7 @@ jobs:
         working-directory: classic/benchmark
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
           submodules: true
@@ -114,7 +114,7 @@ jobs:
     timeout-minutes: 20
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
           submodules: true

.github/workflows/classic-benchmark_publish_package.yml

@@ -10,7 +10,7 @@ jobs:
       contents: write
     steps:
     - name: Checkout repository
-      uses: actions/checkout@v4
+      uses: actions/checkout@v6
       with:
         submodules: true
         fetch-depth: 0

.github/workflows/classic-forge-ci.yml

@@ -63,7 +63,7 @@ jobs:
       #   See: https://github.com/actions/runner/issues/598#issuecomment-2011890429
 
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
           submodules: true

.github/workflows/classic-frontend-ci.yml

@@ -25,7 +25,7 @@ jobs:
 
     steps:
       - name: Checkout Repo
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Setup Flutter
         uses: subosito/flutter-action@v2

.github/workflows/classic-python-checks.yml

@@ -33,7 +33,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - id: changes-in
         name: Determine affected subprojects
@@ -68,7 +68,7 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
 
@@ -120,7 +120,7 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
 

.github/workflows/pr-overlap-check.yml

@@ -16,7 +16,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0  # Need full history for merge testing
 

autogpt_platform/backend/backend/api/conftest.py

@@ -1,9 +1,4 @@
-"""Common test fixtures for server tests.
-
-Note: Common fixtures like test_user_id, admin_user_id, target_user_id,
-setup_test_user, and setup_admin_user are defined in the parent conftest.py
-(backend/conftest.py) and are available here automatically.
-"""
+"""Common test fixtures for server tests."""
 
 import pytest
 from pytest_snapshot.plugin import Snapshot
@@ -16,6 +11,54 @@ def configured_snapshot(snapshot: Snapshot) -> Snapshot:
     return snapshot
 
 
+@pytest.fixture
+def test_user_id() -> str:
+    """Test user ID fixture."""
+    return "3e53486c-cf57-477e-ba2a-cb02dc828e1a"
+
+
+@pytest.fixture
+def admin_user_id() -> str:
+    """Admin user ID fixture."""
+    return "4e53486c-cf57-477e-ba2a-cb02dc828e1b"
+
+
+@pytest.fixture
+def target_user_id() -> str:
+    """Target user ID fixture."""
+    return "5e53486c-cf57-477e-ba2a-cb02dc828e1c"
+
+
+@pytest.fixture
+async def setup_test_user(test_user_id):
+    """Create test user in database before tests."""
+    from backend.data.user import get_or_create_user
+
+    # Create the test user in the database using JWT token format
+    user_data = {
+        "sub": test_user_id,
+        "email": "test@example.com",
+        "user_metadata": {"name": "Test User"},
+    }
+    await get_or_create_user(user_data)
+    return test_user_id
+
+
+@pytest.fixture
+async def setup_admin_user(admin_user_id):
+    """Create admin user in database before tests."""
+    from backend.data.user import get_or_create_user
+
+    # Create the admin user in the database using JWT token format
+    user_data = {
+        "sub": admin_user_id,
+        "email": "test-admin@example.com",
+        "user_metadata": {"name": "Test Admin"},
+    }
+    await get_or_create_user(user_data)
+    return admin_user_id
+
+
 @pytest.fixture
 def mock_jwt_user(test_user_id):
     """Provide mock JWT payload for regular user testing."""

autogpt_platform/backend/backend/api/external/v1/tools.py

@@ -15,9 +15,9 @@ from prisma.enums import APIKeyPermission
 from pydantic import BaseModel, Field
 
 from backend.api.external.middleware import require_permission
-from backend.copilot.model import ChatSession
-from backend.copilot.tools import find_agent_tool, run_agent_tool
-from backend.copilot.tools.models import ToolResponseBase
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.tools import find_agent_tool, run_agent_tool
+from backend.api.features.chat.tools.models import ToolResponseBase
 from backend.data.auth.base import APIAuthorizationInfo
 
 logger = logging.getLogger(__name__)

autogpt_platform/backend/backend/api/features/chat/completion_consumer.py

@@ -37,10 +37,12 @@ stale pending messages from dead consumers.
 
 import asyncio
 import logging
+import os
 import uuid
 from typing import Any
 
 import orjson
+from prisma import Prisma
 from pydantic import BaseModel
 from redis.exceptions import ResponseError
 
@@ -67,8 +69,8 @@ class OperationCompleteMessage(BaseModel):
 class ChatCompletionConsumer:
     """Consumer for chat operation completion messages from Redis Streams.
 
-    Database operations are handled through the chat_db() accessor, which
-    routes through DatabaseManager RPC when Prisma is not directly connected.
+    This consumer initializes its own Prisma client in start() to ensure
+    database operations work correctly within this async context.
 
     Uses Redis consumer groups to allow multiple platform pods to consume
     messages reliably with automatic redelivery on failure.
@@ -77,6 +79,7 @@ class ChatCompletionConsumer:
     def __init__(self):
         self._consumer_task: asyncio.Task | None = None
         self._running = False
+        self._prisma: Prisma | None = None
         self._consumer_name = f"consumer-{uuid.uuid4().hex[:8]}"
 
     async def start(self) -> None:
@@ -112,6 +115,15 @@ class ChatCompletionConsumer:
             f"Chat completion consumer started (consumer: {self._consumer_name})"
         )
 
+    async def _ensure_prisma(self) -> Prisma:
+        """Lazily initialize Prisma client on first use."""
+        if self._prisma is None:
+            database_url = os.getenv("DATABASE_URL", "postgresql://localhost:5432")
+            self._prisma = Prisma(datasource={"url": database_url})
+            await self._prisma.connect()
+            logger.info("[COMPLETION] Consumer Prisma client connected (lazy init)")
+        return self._prisma
+
     async def stop(self) -> None:
         """Stop the completion consumer."""
         self._running = False
@@ -124,6 +136,11 @@ class ChatCompletionConsumer:
                 pass
             self._consumer_task = None
 
+        if self._prisma:
+            await self._prisma.disconnect()
+            self._prisma = None
+            logger.info("[COMPLETION] Consumer Prisma client disconnected")
+
         logger.info("Chat completion consumer stopped")
 
     async def _consume_messages(self) -> None:
@@ -235,7 +252,7 @@ class ChatCompletionConsumer:
             # XAUTOCLAIM after min_idle_time expires
 
     async def _handle_message(self, body: bytes) -> None:
-        """Handle a completion message."""
+        """Handle a completion message using our own Prisma client."""
         try:
             data = orjson.loads(body)
             message = OperationCompleteMessage(**data)
@@ -285,7 +302,8 @@ class ChatCompletionConsumer:
         message: OperationCompleteMessage,
     ) -> None:
         """Handle successful operation completion."""
-        await process_operation_success(task, message.result)
+        prisma = await self._ensure_prisma()
+        await process_operation_success(task, message.result, prisma)
 
     async def _handle_failure(
         self,
@@ -293,7 +311,8 @@ class ChatCompletionConsumer:
         message: OperationCompleteMessage,
     ) -> None:
         """Handle failed operation completion."""
-        await process_operation_failure(task, message.error)
+        prisma = await self._ensure_prisma()
+        await process_operation_failure(task, message.error, prisma)
 
 
 # Module-level consumer instance

autogpt_platform/backend/backend/api/features/chat/completion_handler.py

@@ -9,8 +9,7 @@ import logging
 from typing import Any
 
 import orjson
-
-from backend.data.db_accessors import chat_db
+from prisma import Prisma
 
 from . import service as chat_service
 from . import stream_registry
@@ -73,40 +72,48 @@ async def _update_tool_message(
     session_id: str,
     tool_call_id: str,
     content: str,
+    prisma_client: Prisma | None,
 ) -> None:
-    """Update tool message in database using the chat_db accessor.
-
-    Routes through DatabaseManager RPC when Prisma is not directly
-    connected (e.g. in the CoPilot Executor microservice).
+    """Update tool message in database.
 
     Args:
         session_id: The session ID
         tool_call_id: The tool call ID to update
         content: The new content for the message
+        prisma_client: Optional Prisma client. If None, uses chat_service.
 
     Raises:
-        ToolMessageUpdateError: If the database update fails.
+        ToolMessageUpdateError: If the database update fails. The caller should
+            handle this to avoid marking the task as completed with inconsistent state.
     """
     try:
-        updated = await chat_db().update_tool_message_content(
-            session_id=session_id,
-            tool_call_id=tool_call_id,
-            new_content=content,
-        )
-        if not updated:
-            raise ToolMessageUpdateError(
-                f"No message found with tool_call_id="
-                f"{tool_call_id} in session {session_id}"
+        if prisma_client:
+            # Use provided Prisma client (for consumer with its own connection)
+            updated_count = await prisma_client.chatmessage.update_many(
+                where={
+                    "sessionId": session_id,
+                    "toolCallId": tool_call_id,
+                },
+                data={"content": content},
+            )
+            # Check if any rows were updated - 0 means message not found
+            if updated_count == 0:
+                raise ToolMessageUpdateError(
+                    f"No message found with tool_call_id={tool_call_id} in session {session_id}"
+                )
+        else:
+            # Use service function (for webhook endpoint)
+            await chat_service._update_pending_operation(
+                session_id=session_id,
+                tool_call_id=tool_call_id,
+                result=content,
             )
     except ToolMessageUpdateError:
         raise
     except Exception as e:
-        logger.error(
-            f"[COMPLETION] Failed to update tool message: {e}",
-            exc_info=True,
-        )
+        logger.error(f"[COMPLETION] Failed to update tool message: {e}", exc_info=True)
         raise ToolMessageUpdateError(
-            f"Failed to update tool message for tool call #{tool_call_id}: {e}"
+            f"Failed to update tool message for tool_call_id={tool_call_id}: {e}"
         ) from e
 
 
@@ -195,6 +202,7 @@ async def _save_agent_from_result(
 async def process_operation_success(
     task: stream_registry.ActiveTask,
     result: dict | str | None,
+    prisma_client: Prisma | None = None,
 ) -> None:
     """Handle successful operation completion.
 
@@ -204,10 +212,12 @@ async def process_operation_success(
     Args:
         task: The active task that completed
         result: The result data from the operation
+        prisma_client: Optional Prisma client for database operations.
+            If None, uses chat_service._update_pending_operation instead.
 
     Raises:
-        ToolMessageUpdateError: If the database update fails. The task
-            will be marked as failed instead of completed.
+        ToolMessageUpdateError: If the database update fails. The task will be
+            marked as failed instead of completed to avoid inconsistent state.
     """
     # For agent generation tools, save the agent to library
     if task.tool_name in AGENT_GENERATION_TOOLS and isinstance(result, dict):
@@ -240,6 +250,7 @@ async def process_operation_success(
             session_id=task.session_id,
             tool_call_id=task.tool_call_id,
             content=result_str,
+            prisma_client=prisma_client,
         )
     except ToolMessageUpdateError:
         # DB update failed - mark task as failed to avoid inconsistent state
@@ -282,15 +293,18 @@ async def process_operation_success(
 async def process_operation_failure(
     task: stream_registry.ActiveTask,
     error: str | None,
+    prisma_client: Prisma | None = None,
 ) -> None:
     """Handle failed operation completion.
 
-    Publishes the error to the stream registry, updates the database
-    with the error response, and marks the task as failed.
+    Publishes the error to the stream registry, updates the database with
+    the error response, and marks the task as failed.
 
     Args:
         task: The active task that failed
         error: The error message from the operation
+        prisma_client: Optional Prisma client for database operations.
+            If None, uses chat_service._update_pending_operation instead.
     """
     error_msg = error or "Operation failed"
 
@@ -311,6 +325,7 @@ async def process_operation_failure(
             session_id=task.session_id,
             tool_call_id=task.tool_call_id,
             content=error_response.model_dump_json(),
+            prisma_client=prisma_client,
         )
     except ToolMessageUpdateError:
         # DB update failed - log but continue with cleanup

autogpt_platform/backend/backend/api/features/chat/config.py

@@ -27,6 +27,7 @@ class ChatConfig(BaseSettings):
     session_ttl: int = Field(default=43200, description="Session TTL in seconds")
 
     # Streaming Configuration
+    stream_timeout: int = Field(default=300, description="Stream timeout in seconds")
     max_retries: int = Field(
         default=3,
         description="Max retries for fallback path (SDK handles retries internally)",
@@ -38,10 +39,8 @@ class ChatConfig(BaseSettings):
 
     # Long-running operation configuration
     long_running_operation_ttl: int = Field(
-        default=3600,
-        description="TTL in seconds for long-running operation deduplication lock "
-        "(1 hour, matches stream_ttl). Prevents duplicate operations if pod dies. "
-        "For longer operations, the stream_registry heartbeat keeps them alive.",
+        default=600,
+        description="TTL in seconds for long-running operation tracking in Redis (safety net if pod dies)",
     )
 
     # Stream registry configuration for SSE reconnection

autogpt_platform/backend/backend/api/features/chat/db.py

@@ -14,27 +14,29 @@ from prisma.types import (
     ChatSessionWhereInput,
 )
 
-from backend.data import db
+from backend.data.db import transaction
 from backend.util.json import SafeJson
 
-from .model import ChatMessage, ChatSession, ChatSessionInfo
-
 logger = logging.getLogger(__name__)
 
 
-async def get_chat_session(session_id: str) -> ChatSession | None:
+async def get_chat_session(session_id: str) -> PrismaChatSession | None:
     """Get a chat session by ID from the database."""
     session = await PrismaChatSession.prisma().find_unique(
         where={"id": session_id},
-        include={"Messages": {"order_by": {"sequence": "asc"}}},
+        include={"Messages": True},
     )
-    return ChatSession.from_db(session) if session else None
+    if session and session.Messages:
+        # Sort messages by sequence in Python - Prisma Python client doesn't support
+        # order_by in include clauses (unlike Prisma JS), so we sort after fetching
+        session.Messages.sort(key=lambda m: m.sequence)
+    return session
 
 
 async def create_chat_session(
     session_id: str,
     user_id: str,
-) -> ChatSessionInfo:
+) -> PrismaChatSession:
     """Create a new chat session in the database."""
     data = ChatSessionCreateInput(
         id=session_id,
@@ -43,8 +45,7 @@ async def create_chat_session(
         successfulAgentRuns=SafeJson({}),
         successfulAgentSchedules=SafeJson({}),
     )
-    prisma_session = await PrismaChatSession.prisma().create(data=data)
-    return ChatSessionInfo.from_db(prisma_session)
+    return await PrismaChatSession.prisma().create(data=data)
 
 
 async def update_chat_session(
@@ -55,7 +56,7 @@ async def update_chat_session(
     total_prompt_tokens: int | None = None,
     total_completion_tokens: int | None = None,
     title: str | None = None,
-) -> ChatSession | None:
+) -> PrismaChatSession | None:
     """Update a chat session's metadata."""
     data: ChatSessionUpdateInput = {"updatedAt": datetime.now(UTC)}
 
@@ -75,9 +76,12 @@ async def update_chat_session(
     session = await PrismaChatSession.prisma().update(
         where={"id": session_id},
         data=data,
-        include={"Messages": {"order_by": {"sequence": "asc"}}},
+        include={"Messages": True},
     )
-    return ChatSession.from_db(session) if session else None
+    if session and session.Messages:
+        # Sort in Python - Prisma Python doesn't support order_by in include clauses
+        session.Messages.sort(key=lambda m: m.sequence)
+    return session
 
 
 async def add_chat_message(
@@ -90,7 +94,7 @@ async def add_chat_message(
     refusal: str | None = None,
     tool_calls: list[dict[str, Any]] | None = None,
     function_call: dict[str, Any] | None = None,
-) -> ChatMessage:
+) -> PrismaChatMessage:
     """Add a message to a chat session."""
     # Build input dict dynamically rather than using ChatMessageCreateInput directly
     # because Prisma's TypedDict validation rejects optional fields set to None.
@@ -125,119 +129,79 @@ async def add_chat_message(
         ),
         PrismaChatMessage.prisma().create(data=cast(ChatMessageCreateInput, data)),
     )
-    return ChatMessage.from_db(message)
+    return message
 
 
 async def add_chat_messages_batch(
     session_id: str,
     messages: list[dict[str, Any]],
     start_sequence: int,
-) -> tuple[list[ChatMessage], int]:
+) -> list[PrismaChatMessage]:
     """Add multiple messages to a chat session in a batch.
 
-    Uses collision detection with retry: tries to create messages starting
-    at start_sequence. If a unique constraint violation occurs (e.g., the
-    streaming loop and long-running callback race), queries MAX(sequence)
-    and retries with the correct next sequence number. This avoids
-    unnecessary upserts and DB queries in the common case (no collision).
-
-    Returns:
-        Tuple of (messages, final_message_count) where final_message_count
-        is the total number of messages in the session after insertion.
-        This allows callers to update their counters even when collision
-        detection adjusts start_sequence.
+    Uses a transaction for atomicity - if any message creation fails,
+    the entire batch is rolled back.
     """
     if not messages:
-        # No messages to add - return current count
-        return [], start_sequence
+        return []
 
-    max_retries = 3
-    for attempt in range(max_retries):
-        try:
-            created_messages = []
-            async with db.transaction() as tx:
-                for i, msg in enumerate(messages):
-                    # Build input dict dynamically rather than using ChatMessageCreateInput
-                    # directly because Prisma's TypedDict validation rejects optional fields
-                    # set to None. We only include fields that have values, then cast.
-                    data: dict[str, Any] = {
-                        "Session": {"connect": {"id": session_id}},
-                        "role": msg["role"],
-                        "sequence": start_sequence + i,
-                    }
+    created_messages = []
 
-                    # Add optional string fields
-                    if msg.get("content") is not None:
-                        data["content"] = msg["content"]
-                    if msg.get("name") is not None:
-                        data["name"] = msg["name"]
-                    if msg.get("tool_call_id") is not None:
-                        data["toolCallId"] = msg["tool_call_id"]
-                    if msg.get("refusal") is not None:
-                        data["refusal"] = msg["refusal"]
+    async with transaction() as tx:
+        for i, msg in enumerate(messages):
+            # Build input dict dynamically rather than using ChatMessageCreateInput
+            # directly because Prisma's TypedDict validation rejects optional fields
+            # set to None. We only include fields that have values, then cast.
+            data: dict[str, Any] = {
+                "Session": {"connect": {"id": session_id}},
+                "role": msg["role"],
+                "sequence": start_sequence + i,
+            }
 
-                    # Add optional JSON fields only when they have values
-                    if msg.get("tool_calls") is not None:
-                        data["toolCalls"] = SafeJson(msg["tool_calls"])
-                    if msg.get("function_call") is not None:
-                        data["functionCall"] = SafeJson(msg["function_call"])
+            # Add optional string fields
+            if msg.get("content") is not None:
+                data["content"] = msg["content"]
+            if msg.get("name") is not None:
+                data["name"] = msg["name"]
+            if msg.get("tool_call_id") is not None:
+                data["toolCallId"] = msg["tool_call_id"]
+            if msg.get("refusal") is not None:
+                data["refusal"] = msg["refusal"]
 
-                    created = await PrismaChatMessage.prisma(tx).create(
-                        data=cast(ChatMessageCreateInput, data)
-                    )
-                    created_messages.append(created)
+            # Add optional JSON fields only when they have values
+            if msg.get("tool_calls") is not None:
+                data["toolCalls"] = SafeJson(msg["tool_calls"])
+            if msg.get("function_call") is not None:
+                data["functionCall"] = SafeJson(msg["function_call"])
 
-                # Update session's updatedAt timestamp within the same transaction.
-                # Note: Token usage (total_prompt_tokens, total_completion_tokens) is updated
-                # separately via update_chat_session() after streaming completes.
-                await PrismaChatSession.prisma(tx).update(
-                    where={"id": session_id},
-                    data={"updatedAt": datetime.now(UTC)},
-                )
-
-            # Return messages and final message count (for shared counter sync)
-            final_count = start_sequence + len(messages)
-            return [ChatMessage.from_db(m) for m in created_messages], final_count
-
-        except Exception as e:
-            # Check if it's a unique constraint violation
-            error_msg = str(e).lower()
-            is_unique_constraint = (
-                "unique constraint" in error_msg or "duplicate key" in error_msg
+            created = await PrismaChatMessage.prisma(tx).create(
+                data=cast(ChatMessageCreateInput, data)
             )
+            created_messages.append(created)
 
-            if is_unique_constraint and attempt < max_retries - 1:
-                # Collision detected - query MAX(sequence)+1 and retry with correct offset
-                logger.info(
-                    f"Collision detected for session {session_id} at sequence "
-                    f"{start_sequence}, querying DB for latest sequence"
-                )
-                start_sequence = await get_next_sequence(session_id)
-                logger.info(
-                    f"Retrying batch insert with start_sequence={start_sequence}"
-                )
-                continue
-            else:
-                # Not a collision or max retries exceeded - propagate error
-                raise
+        # Update session's updatedAt timestamp within the same transaction.
+        # Note: Token usage (total_prompt_tokens, total_completion_tokens) is updated
+        # separately via update_chat_session() after streaming completes.
+        await PrismaChatSession.prisma(tx).update(
+            where={"id": session_id},
+            data={"updatedAt": datetime.now(UTC)},
+        )
 
-    # Should never reach here due to raise in exception handler
-    raise RuntimeError(f"Failed to insert messages after {max_retries} attempts")
+    return created_messages
 
 
 async def get_user_chat_sessions(
     user_id: str,
     limit: int = 50,
     offset: int = 0,
-) -> list[ChatSessionInfo]:
+) -> list[PrismaChatSession]:
     """Get chat sessions for a user, ordered by most recent."""
-    prisma_sessions = await PrismaChatSession.prisma().find_many(
+    return await PrismaChatSession.prisma().find_many(
         where={"userId": user_id},
         order={"updatedAt": "desc"},
         take=limit,
         skip=offset,
     )
-    return [ChatSessionInfo.from_db(s) for s in prisma_sessions]
 
 
 async def get_user_session_count(user_id: str) -> int:
@@ -276,23 +240,10 @@ async def delete_chat_session(session_id: str, user_id: str | None = None) -> bo
         return False
 
 
-async def get_next_sequence(session_id: str) -> int:
-    """Get the next sequence number for a new message in this session.
-
-    Uses MAX(sequence) + 1 for robustness. Returns 0 if no messages exist.
-    More robust than COUNT(*) because it's immune to deleted messages.
-    """
-    result = await db.prisma.query_raw(
-        """
-        SELECT COALESCE(MAX(sequence) + 1, 0) as next_seq
-        FROM "ChatMessage"
-        WHERE "sessionId" = $1
-        """,
-        session_id,
-    )
-    if not result or len(result) == 0:
-        return 0
-    return int(result[0]["next_seq"])
+async def get_chat_session_message_count(session_id: str) -> int:
+    """Get the number of messages in a chat session."""
+    count = await PrismaChatMessage.prisma().count(where={"sessionId": session_id})
+    return count
 
 
 async def update_tool_message_content(

autogpt_platform/backend/backend/api/features/chat/model.py

@@ -2,7 +2,7 @@ import asyncio
 import logging
 import uuid
 from datetime import UTC, datetime
-from typing import Any, Self, cast
+from typing import Any, cast
 from weakref import WeakValueDictionary
 
 from openai.types.chat import (
@@ -23,17 +23,26 @@ from prisma.models import ChatMessage as PrismaChatMessage
 from prisma.models import ChatSession as PrismaChatSession
 from pydantic import BaseModel
 
-from backend.data.db_accessors import chat_db
 from backend.data.redis_client import get_redis_async
 from backend.util import json
 from backend.util.exceptions import DatabaseError, RedisError
 
+from . import db as chat_db
 from .config import ChatConfig
 
 logger = logging.getLogger(__name__)
 config = ChatConfig()
 
 
+def _parse_json_field(value: str | dict | list | None, default: Any = None) -> Any:
+    """Parse a JSON field that may be stored as string or already parsed."""
+    if value is None:
+        return default
+    if isinstance(value, str):
+        return json.loads(value)
+    return value
+
+
 # Redis cache key prefix for chat sessions
 CHAT_SESSION_CACHE_PREFIX = "chat:session:"
 
@@ -43,7 +52,28 @@ def _get_session_cache_key(session_id: str) -> str:
     return f"{CHAT_SESSION_CACHE_PREFIX}{session_id}"
 
 
-# ===================== Chat data models ===================== #
+# Session-level locks to prevent race conditions during concurrent upserts.
+# Uses WeakValueDictionary to automatically garbage collect locks when no longer referenced,
+# preventing unbounded memory growth while maintaining lock semantics for active sessions.
+# Invalidation: Locks are auto-removed by GC when no coroutine holds a reference (after
+# async with lock: completes). Explicit cleanup also occurs in delete_chat_session().
+_session_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary()
+_session_locks_mutex = asyncio.Lock()
+
+
+async def _get_session_lock(session_id: str) -> asyncio.Lock:
+    """Get or create a lock for a specific session to prevent concurrent upserts.
+
+    Uses WeakValueDictionary for automatic cleanup: locks are garbage collected
+    when no coroutine holds a reference to them, preventing memory leaks from
+    unbounded growth of session locks.
+    """
+    async with _session_locks_mutex:
+        lock = _session_locks.get(session_id)
+        if lock is None:
+            lock = asyncio.Lock()
+            _session_locks[session_id] = lock
+        return lock
 
 
 class ChatMessage(BaseModel):
@@ -55,19 +85,6 @@ class ChatMessage(BaseModel):
     tool_calls: list[dict] | None = None
     function_call: dict | None = None
 
-    @staticmethod
-    def from_db(prisma_message: PrismaChatMessage) -> "ChatMessage":
-        """Convert a Prisma ChatMessage to a Pydantic ChatMessage."""
-        return ChatMessage(
-            role=prisma_message.role,
-            content=prisma_message.content,
-            name=prisma_message.name,
-            tool_call_id=prisma_message.toolCallId,
-            refusal=prisma_message.refusal,
-            tool_calls=_parse_json_field(prisma_message.toolCalls),
-            function_call=_parse_json_field(prisma_message.functionCall),
-        )
-
 
 class Usage(BaseModel):
     prompt_tokens: int
@@ -75,10 +92,11 @@ class Usage(BaseModel):
     total_tokens: int
 
 
-class ChatSessionInfo(BaseModel):
+class ChatSession(BaseModel):
     session_id: str
     user_id: str
     title: str | None = None
+    messages: list[ChatMessage]
     usage: list[Usage]
     credentials: dict[str, dict] = {}  # Map of provider -> credential metadata
     started_at: datetime
@@ -86,9 +104,60 @@ class ChatSessionInfo(BaseModel):
     successful_agent_runs: dict[str, int] = {}
     successful_agent_schedules: dict[str, int] = {}
 
-    @classmethod
-    def from_db(cls, prisma_session: PrismaChatSession) -> Self:
-        """Convert Prisma ChatSession to Pydantic ChatSession."""
+    def add_tool_call_to_current_turn(self, tool_call: dict) -> None:
+        """Attach a tool_call to the current turn's assistant message.
+
+        Searches backwards for the most recent assistant message (stopping at
+        any user message boundary). If found, appends the tool_call to it.
+        Otherwise creates a new assistant message with the tool_call.
+        """
+        for msg in reversed(self.messages):
+            if msg.role == "user":
+                break
+            if msg.role == "assistant":
+                if not msg.tool_calls:
+                    msg.tool_calls = []
+                msg.tool_calls.append(tool_call)
+                return
+
+        self.messages.append(
+            ChatMessage(role="assistant", content="", tool_calls=[tool_call])
+        )
+
+    @staticmethod
+    def new(user_id: str) -> "ChatSession":
+        return ChatSession(
+            session_id=str(uuid.uuid4()),
+            user_id=user_id,
+            title=None,
+            messages=[],
+            usage=[],
+            credentials={},
+            started_at=datetime.now(UTC),
+            updated_at=datetime.now(UTC),
+        )
+
+    @staticmethod
+    def from_db(
+        prisma_session: PrismaChatSession,
+        prisma_messages: list[PrismaChatMessage] | None = None,
+    ) -> "ChatSession":
+        """Convert Prisma models to Pydantic ChatSession."""
+        messages = []
+        if prisma_messages:
+            for msg in prisma_messages:
+                messages.append(
+                    ChatMessage(
+                        role=msg.role,
+                        content=msg.content,
+                        name=msg.name,
+                        tool_call_id=msg.toolCallId,
+                        refusal=msg.refusal,
+                        tool_calls=_parse_json_field(msg.toolCalls),
+                        function_call=_parse_json_field(msg.functionCall),
+                    )
+                )
+
         # Parse JSON fields from Prisma
         credentials = _parse_json_field(prisma_session.credentials, default={})
         successful_agent_runs = _parse_json_field(
@@ -110,10 +179,11 @@ class ChatSessionInfo(BaseModel):
                 )
             )
 
-        return cls(
+        return ChatSession(
             session_id=prisma_session.id,
             user_id=prisma_session.userId,
             title=prisma_session.title,
+            messages=messages,
             usage=usage,
             credentials=credentials,
             started_at=prisma_session.createdAt,
@@ -122,55 +192,46 @@ class ChatSessionInfo(BaseModel):
             successful_agent_schedules=successful_agent_schedules,
         )
 
+    @staticmethod
+    def _merge_consecutive_assistant_messages(
+        messages: list[ChatCompletionMessageParam],
+    ) -> list[ChatCompletionMessageParam]:
+        """Merge consecutive assistant messages into single messages.
 
-class ChatSession(ChatSessionInfo):
-    messages: list[ChatMessage]
-
-    @classmethod
-    def new(cls, user_id: str) -> Self:
-        return cls(
-            session_id=str(uuid.uuid4()),
-            user_id=user_id,
-            title=None,
-            messages=[],
-            usage=[],
-            credentials={},
-            started_at=datetime.now(UTC),
-            updated_at=datetime.now(UTC),
-        )
-
-    @classmethod
-    def from_db(cls, prisma_session: PrismaChatSession) -> Self:
-        """Convert Prisma ChatSession to Pydantic ChatSession."""
-        if prisma_session.Messages is None:
-            raise ValueError(
-                f"Prisma session {prisma_session.id} is missing Messages relation"
-            )
-
-        return cls(
-            **ChatSessionInfo.from_db(prisma_session).model_dump(),
-            messages=[ChatMessage.from_db(m) for m in prisma_session.Messages],
-        )
-
-    def add_tool_call_to_current_turn(self, tool_call: dict) -> None:
-        """Attach a tool_call to the current turn's assistant message.
-
-        Searches backwards for the most recent assistant message (stopping at
-        any user message boundary). If found, appends the tool_call to it.
-        Otherwise creates a new assistant message with the tool_call.
+        Long-running tool flows can create split assistant messages: one with
+        text content and another with tool_calls. Anthropic's API requires
+        tool_result blocks to reference a tool_use in the immediately preceding
+        assistant message, so these splits cause 400 errors via OpenRouter.
         """
-        for msg in reversed(self.messages):
-            if msg.role == "user":
-                break
-            if msg.role == "assistant":
-                if not msg.tool_calls:
-                    msg.tool_calls = []
-                msg.tool_calls.append(tool_call)
-                return
+        if len(messages) < 2:
+            return messages
 
-        self.messages.append(
-            ChatMessage(role="assistant", content="", tool_calls=[tool_call])
-        )
+        result: list[ChatCompletionMessageParam] = [messages[0]]
+        for msg in messages[1:]:
+            prev = result[-1]
+            if prev.get("role") != "assistant" or msg.get("role") != "assistant":
+                result.append(msg)
+                continue
+
+            prev = cast(ChatCompletionAssistantMessageParam, prev)
+            curr = cast(ChatCompletionAssistantMessageParam, msg)
+
+            curr_content = curr.get("content") or ""
+            if curr_content:
+                prev_content = prev.get("content") or ""
+                prev["content"] = (
+                    f"{prev_content}\n{curr_content}" if prev_content else curr_content
+                )
+
+            curr_tool_calls = curr.get("tool_calls")
+            if curr_tool_calls:
+                prev_tool_calls = prev.get("tool_calls")
+                prev["tool_calls"] = (
+                    list(prev_tool_calls) + list(curr_tool_calls)
+                    if prev_tool_calls
+                    else list(curr_tool_calls)
+                )
+        return result
 
     def to_openai_messages(self) -> list[ChatCompletionMessageParam]:
         messages = []
@@ -260,70 +321,40 @@ class ChatSession(ChatSessionInfo):
                 )
         return self._merge_consecutive_assistant_messages(messages)
 
-    @staticmethod
-    def _merge_consecutive_assistant_messages(
-        messages: list[ChatCompletionMessageParam],
-    ) -> list[ChatCompletionMessageParam]:
-        """Merge consecutive assistant messages into single messages.
 
-        Long-running tool flows can create split assistant messages: one with
-        text content and another with tool_calls. Anthropic's API requires
-        tool_result blocks to reference a tool_use in the immediately preceding
-        assistant message, so these splits cause 400 errors via OpenRouter.
-        """
-        if len(messages) < 2:
-            return messages
+async def _get_session_from_cache(session_id: str) -> ChatSession | None:
+    """Get a chat session from Redis cache."""
+    redis_key = _get_session_cache_key(session_id)
+    async_redis = await get_redis_async()
+    raw_session: bytes | None = await async_redis.get(redis_key)
 
-        result: list[ChatCompletionMessageParam] = [messages[0]]
-        for msg in messages[1:]:
-            prev = result[-1]
-            if prev.get("role") != "assistant" or msg.get("role") != "assistant":
-                result.append(msg)
-                continue
+    if raw_session is None:
+        return None
 
-            prev = cast(ChatCompletionAssistantMessageParam, prev)
-            curr = cast(ChatCompletionAssistantMessageParam, msg)
-
-            curr_content = curr.get("content") or ""
-            if curr_content:
-                prev_content = prev.get("content") or ""
-                prev["content"] = (
-                    f"{prev_content}\n{curr_content}" if prev_content else curr_content
-                )
-
-            curr_tool_calls = curr.get("tool_calls")
-            if curr_tool_calls:
-                prev_tool_calls = prev.get("tool_calls")
-                prev["tool_calls"] = (
-                    list(prev_tool_calls) + list(curr_tool_calls)
-                    if prev_tool_calls
-                    else list(curr_tool_calls)
-                )
-        return result
+    try:
+        session = ChatSession.model_validate_json(raw_session)
+        logger.info(
+            f"[CACHE] Loaded session {session_id}: {len(session.messages)} messages, "
+            f"last_roles={[m.role for m in session.messages[-3:]]}"  # Last 3 roles
+        )
+        return session
+    except Exception as e:
+        logger.error(f"Failed to deserialize session {session_id}: {e}", exc_info=True)
+        raise RedisError(f"Corrupted session data for {session_id}") from e
 
 
-def _parse_json_field(value: str | dict | list | None, default: Any = None) -> Any:
-    """Parse a JSON field that may be stored as string or already parsed."""
-    if value is None:
-        return default
-    if isinstance(value, str):
-        return json.loads(value)
-    return value
-
-
-# ================ Chat cache + DB operations ================ #
-
-# NOTE: Database calls are automatically routed through DatabaseManager if Prisma is not
-#       connected directly.
-
-
-async def cache_chat_session(session: ChatSession) -> None:
-    """Cache a chat session in Redis (without persisting to the database)."""
+async def _cache_session(session: ChatSession) -> None:
+    """Cache a chat session in Redis."""
     redis_key = _get_session_cache_key(session.session_id)
     async_redis = await get_redis_async()
     await async_redis.setex(redis_key, config.session_ttl, session.model_dump_json())
 
 
+async def cache_chat_session(session: ChatSession) -> None:
+    """Cache a chat session without persisting to the database."""
+    await _cache_session(session)
+
+
 async def invalidate_session_cache(session_id: str) -> None:
     """Invalidate a chat session from Redis cache.
 
@@ -339,6 +370,77 @@ async def invalidate_session_cache(session_id: str) -> None:
         logger.warning(f"Failed to invalidate session cache for {session_id}: {e}")
 
 
+async def _get_session_from_db(session_id: str) -> ChatSession | None:
+    """Get a chat session from the database."""
+    prisma_session = await chat_db.get_chat_session(session_id)
+    if not prisma_session:
+        return None
+
+    messages = prisma_session.Messages
+    logger.debug(
+        f"[DB] Loaded session {session_id}: {len(messages) if messages else 0} messages, "
+        f"roles={[m.role for m in messages[-3:]] if messages else []}"  # Last 3 roles
+    )
+
+    return ChatSession.from_db(prisma_session, messages)
+
+
+async def _save_session_to_db(
+    session: ChatSession, existing_message_count: int
+) -> None:
+    """Save or update a chat session in the database."""
+    # Check if session exists in DB
+    existing = await chat_db.get_chat_session(session.session_id)
+
+    if not existing:
+        # Create new session
+        await chat_db.create_chat_session(
+            session_id=session.session_id,
+            user_id=session.user_id,
+        )
+        existing_message_count = 0
+
+    # Calculate total tokens from usage
+    total_prompt = sum(u.prompt_tokens for u in session.usage)
+    total_completion = sum(u.completion_tokens for u in session.usage)
+
+    # Update session metadata
+    await chat_db.update_chat_session(
+        session_id=session.session_id,
+        credentials=session.credentials,
+        successful_agent_runs=session.successful_agent_runs,
+        successful_agent_schedules=session.successful_agent_schedules,
+        total_prompt_tokens=total_prompt,
+        total_completion_tokens=total_completion,
+    )
+
+    # Add new messages (only those after existing count)
+    new_messages = session.messages[existing_message_count:]
+    if new_messages:
+        messages_data = []
+        for msg in new_messages:
+            messages_data.append(
+                {
+                    "role": msg.role,
+                    "content": msg.content,
+                    "name": msg.name,
+                    "tool_call_id": msg.tool_call_id,
+                    "refusal": msg.refusal,
+                    "tool_calls": msg.tool_calls,
+                    "function_call": msg.function_call,
+                }
+            )
+        logger.debug(
+            f"[DB] Saving {len(new_messages)} messages to session {session.session_id}, "
+            f"roles={[m['role'] for m in messages_data]}"
+        )
+        await chat_db.add_chat_messages_batch(
+            session_id=session.session_id,
+            messages=messages_data,
+            start_sequence=existing_message_count,
+        )
+
+
 async def get_chat_session(
     session_id: str,
     user_id: str | None = None,
@@ -386,73 +488,22 @@ async def get_chat_session(
 
     # Cache the session from DB
     try:
-        await cache_chat_session(session)
-        logger.info(f"Cached session {session_id} from database")
+        await _cache_session(session)
     except Exception as e:
         logger.warning(f"Failed to cache session {session_id}: {e}")
 
     return session
 
 
-async def _get_session_from_cache(session_id: str) -> ChatSession | None:
-    """Get a chat session from Redis cache."""
-    redis_key = _get_session_cache_key(session_id)
-    async_redis = await get_redis_async()
-    raw_session: bytes | None = await async_redis.get(redis_key)
-
-    if raw_session is None:
-        return None
-
-    try:
-        session = ChatSession.model_validate_json(raw_session)
-        logger.info(
-            f"Loading session {session_id} from cache: "
-            f"message_count={len(session.messages)}, "
-            f"roles={[m.role for m in session.messages]}"
-        )
-        return session
-    except Exception as e:
-        logger.error(f"Failed to deserialize session {session_id}: {e}", exc_info=True)
-        raise RedisError(f"Corrupted session data for {session_id}") from e
-
-
-async def _get_session_from_db(session_id: str) -> ChatSession | None:
-    """Get a chat session from the database."""
-    session = await chat_db().get_chat_session(session_id)
-    if not session:
-        return None
-
-    logger.info(
-        f"Loaded session {session_id} from DB: "
-        f"has_messages={bool(session.messages)}, "
-        f"message_count={len(session.messages)}, "
-        f"roles={[m.role for m in session.messages]}"
-    )
-
-    return session
-
-
 async def upsert_chat_session(
     session: ChatSession,
-    *,
-    existing_message_count: int | None = None,
-) -> tuple[ChatSession, int]:
+) -> ChatSession:
     """Update a chat session in both cache and database.
 
     Uses session-level locking to prevent race conditions when concurrent
     operations (e.g., background title update and main stream handler)
     attempt to upsert the same session simultaneously.
 
-    Args:
-        existing_message_count: If provided, skip the DB query to count
-            existing messages. The caller is responsible for tracking this
-            accurately. Useful for incremental saves in a streaming loop
-            where the caller already knows how many messages are persisted.
-
-    Returns:
-        Tuple of (session, final_message_count) where final_message_count is
-        the actual persisted message count after collision detection adjustments.
-
     Raises:
         DatabaseError: If the database write fails. The cache is still updated
             as a best-effort optimization, but the error is propagated to ensure
@@ -464,21 +515,15 @@ async def upsert_chat_session(
 
     async with lock:
         # Get existing message count from DB for incremental saves
-        if existing_message_count is None:
-            existing_message_count = await chat_db().get_next_sequence(
-                session.session_id
-            )
+        existing_message_count = await chat_db.get_chat_session_message_count(
+            session.session_id
+        )
 
         db_error: Exception | None = None
-        final_count = existing_message_count
 
         # Save to database (primary storage)
         try:
-            final_count = await _save_session_to_db(
-                session,
-                existing_message_count,
-                skip_existence_check=existing_message_count > 0,
-            )
+            await _save_session_to_db(session, existing_message_count)
         except Exception as e:
             logger.error(
                 f"Failed to save session {session.session_id} to database: {e}"
@@ -487,7 +532,7 @@ async def upsert_chat_session(
 
         # Save to cache (best-effort, even if DB failed)
         try:
-            await cache_chat_session(session)
+            await _cache_session(session)
         except Exception as e:
             # If DB succeeded but cache failed, raise cache error
             if db_error is None:
@@ -505,82 +550,7 @@ async def upsert_chat_session(
                 f"Failed to persist chat session {session.session_id} to database"
             ) from db_error
 
-        return session, final_count
-
-
-async def _save_session_to_db(
-    session: ChatSession,
-    existing_message_count: int,
-    *,
-    skip_existence_check: bool = False,
-) -> int:
-    """Save or update a chat session in the database.
-
-    Args:
-        skip_existence_check: When True, skip the ``get_chat_session`` query
-            and assume the session row already exists.  Saves one DB round trip
-            for incremental saves during streaming.
-
-    Returns:
-        Final message count after save (accounting for collision detection).
-    """
-    db = chat_db()
-
-    if not skip_existence_check:
-        # Check if session exists in DB
-        existing = await db.get_chat_session(session.session_id)
-
-        if not existing:
-            # Create new session
-            await db.create_chat_session(
-                session_id=session.session_id,
-                user_id=session.user_id,
-            )
-            existing_message_count = 0
-
-    # Calculate total tokens from usage
-    total_prompt = sum(u.prompt_tokens for u in session.usage)
-    total_completion = sum(u.completion_tokens for u in session.usage)
-
-    # Update session metadata
-    await db.update_chat_session(
-        session_id=session.session_id,
-        credentials=session.credentials,
-        successful_agent_runs=session.successful_agent_runs,
-        successful_agent_schedules=session.successful_agent_schedules,
-        total_prompt_tokens=total_prompt,
-        total_completion_tokens=total_completion,
-    )
-
-    # Add new messages (only those after existing count)
-    new_messages = session.messages[existing_message_count:]
-    final_count = existing_message_count
-    if new_messages:
-        messages_data = []
-        for msg in new_messages:
-            messages_data.append(
-                {
-                    "role": msg.role,
-                    "content": msg.content,
-                    "name": msg.name,
-                    "tool_call_id": msg.tool_call_id,
-                    "refusal": msg.refusal,
-                    "tool_calls": msg.tool_calls,
-                    "function_call": msg.function_call,
-                }
-            )
-        logger.info(
-            f"Saving {len(new_messages)} new messages to DB for session {session.session_id}: "
-            f"roles={[m['role'] for m in messages_data]}, "
-            f"start_sequence={existing_message_count}"
-        )
-        _, final_count = await db.add_chat_messages_batch(
-            session_id=session.session_id,
-            messages=messages_data,
-            start_sequence=existing_message_count,
-        )
-
-    return final_count
+        return session
 
 
 async def append_and_save_message(session_id: str, message: ChatMessage) -> ChatSession:
@@ -598,7 +568,9 @@ async def append_and_save_message(session_id: str, message: ChatMessage) -> Chat
             raise ValueError(f"Session {session_id} not found")
 
         session.messages.append(message)
-        existing_message_count = await chat_db().get_next_sequence(session_id)
+        existing_message_count = await chat_db.get_chat_session_message_count(
+            session_id
+        )
 
         try:
             await _save_session_to_db(session, existing_message_count)
@@ -608,7 +580,7 @@ async def append_and_save_message(session_id: str, message: ChatMessage) -> Chat
             ) from e
 
         try:
-            await cache_chat_session(session)
+            await _cache_session(session)
         except Exception as e:
             logger.warning(f"Cache write failed for session {session_id}: {e}")
 
@@ -627,7 +599,7 @@ async def create_chat_session(user_id: str) -> ChatSession:
 
     # Create in database first - fail fast if this fails
     try:
-        await chat_db().create_chat_session(
+        await chat_db.create_chat_session(
             session_id=session.session_id,
             user_id=user_id,
         )
@@ -639,7 +611,7 @@ async def create_chat_session(user_id: str) -> ChatSession:
 
     # Cache the session (best-effort optimization, DB is source of truth)
     try:
-        await cache_chat_session(session)
+        await _cache_session(session)
     except Exception as e:
         logger.warning(f"Failed to cache new session {session.session_id}: {e}")
 
@@ -650,16 +622,20 @@ async def get_user_sessions(
     user_id: str,
     limit: int = 50,
     offset: int = 0,
-) -> tuple[list[ChatSessionInfo], int]:
+) -> tuple[list[ChatSession], int]:
     """Get chat sessions for a user from the database with total count.
 
     Returns:
         A tuple of (sessions, total_count) where total_count is the overall
         number of sessions for the user (not just the current page).
     """
-    db = chat_db()
-    sessions = await db.get_user_chat_sessions(user_id, limit, offset)
-    total_count = await db.get_user_session_count(user_id)
+    prisma_sessions = await chat_db.get_user_chat_sessions(user_id, limit, offset)
+    total_count = await chat_db.get_user_session_count(user_id)
+
+    sessions = []
+    for prisma_session in prisma_sessions:
+        # Convert without messages for listing (lighter weight)
+        sessions.append(ChatSession.from_db(prisma_session, None))
 
     return sessions, total_count
 
@@ -677,7 +653,7 @@ async def delete_chat_session(session_id: str, user_id: str | None = None) -> bo
     """
     # Delete from database first (with optional user_id validation)
     # This confirms ownership before invalidating cache
-    deleted = await chat_db().delete_chat_session(session_id, user_id)
+    deleted = await chat_db.delete_chat_session(session_id, user_id)
 
     if not deleted:
         return False
@@ -712,7 +688,7 @@ async def update_session_title(session_id: str, title: str) -> bool:
         True if updated successfully, False otherwise.
     """
     try:
-        result = await chat_db().update_chat_session(session_id=session_id, title=title)
+        result = await chat_db.update_chat_session(session_id=session_id, title=title)
         if result is None:
             logger.warning(f"Session {session_id} not found for title update")
             return False
@@ -724,7 +700,7 @@ async def update_session_title(session_id: str, title: str) -> bool:
             cached = await _get_session_from_cache(session_id)
             if cached:
                 cached.title = title
-                await cache_chat_session(cached)
+                await _cache_session(cached)
         except Exception as e:
             # Not critical - title will be correct on next full cache refresh
             logger.warning(
@@ -735,29 +711,3 @@ async def update_session_title(session_id: str, title: str) -> bool:
     except Exception as e:
         logger.error(f"Failed to update title for session {session_id}: {e}")
         return False
-
-
-# ==================== Chat session locks ==================== #
-
-_session_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary()
-_session_locks_mutex = asyncio.Lock()
-
-
-async def _get_session_lock(session_id: str) -> asyncio.Lock:
-    """Get or create a lock for a specific session to prevent concurrent upserts.
-
-    This was originally added to solve the specific problem of race conditions between
-    the session title thread and the conversation thread, which always occurs on the
-    same instance as we prevent rapid request sends on the frontend.
-
-    Uses WeakValueDictionary for automatic cleanup: locks are garbage collected
-    when no coroutine holds a reference to them, preventing memory leaks from
-    unbounded growth of session locks. Explicit cleanup also occurs
-    in `delete_chat_session()`.
-    """
-    async with _session_locks_mutex:
-        lock = _session_locks.get(session_id)
-        if lock is None:
-            lock = asyncio.Lock()
-            _session_locks[session_id] = lock
-        return lock

autogpt_platform/backend/backend/api/features/chat/model_test.py

@@ -60,7 +60,7 @@ async def test_chatsession_redis_storage(setup_test_user, test_user_id):
     s = ChatSession.new(user_id=test_user_id)
     s.messages = messages
 
-    s, _ = await upsert_chat_session(s)
+    s = await upsert_chat_session(s)
 
     s2 = await get_chat_session(
         session_id=s.session_id,
@@ -77,7 +77,7 @@ async def test_chatsession_redis_storage_user_id_mismatch(
 
     s = ChatSession.new(user_id=test_user_id)
     s.messages = messages
-    s, _ = await upsert_chat_session(s)
+    s = await upsert_chat_session(s)
 
     s2 = await get_chat_session(s.session_id, "different_user_id")
 
@@ -94,7 +94,7 @@ async def test_chatsession_db_storage(setup_test_user, test_user_id):
     s.messages = messages  # Contains user, assistant, and tool messages
     assert s.session_id is not None, "Session id is not set"
     # Upsert to save to both cache and DB
-    s, _ = await upsert_chat_session(s)
+    s = await upsert_chat_session(s)
 
     # Clear the Redis cache to force DB load
     redis_key = f"chat:session:{s.session_id}"

autogpt_platform/backend/backend/api/features/chat/routes.py

@@ -11,15 +11,14 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Query, Response,
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 
-from backend.copilot import service as chat_service
-from backend.copilot import stream_registry
-from backend.copilot.completion_handler import (
-    process_operation_failure,
-    process_operation_success,
-)
-from backend.copilot.config import ChatConfig
-from backend.copilot.executor.utils import enqueue_cancel_task, enqueue_copilot_task
-from backend.copilot.model import (
+from backend.util.exceptions import NotFoundError
+from backend.util.feature_flag import Flag, is_feature_enabled
+
+from . import service as chat_service
+from . import stream_registry
+from .completion_handler import process_operation_failure, process_operation_success
+from .config import ChatConfig
+from .model import (
     ChatMessage,
     ChatSession,
     append_and_save_message,
@@ -28,8 +27,9 @@ from backend.copilot.model import (
     get_chat_session,
     get_user_sessions,
 )
-from backend.copilot.response_model import StreamError, StreamFinish, StreamHeartbeat
-from backend.copilot.tools.models import (
+from .response_model import StreamError, StreamFinish, StreamHeartbeat, StreamStart
+from .sdk import service as sdk_service
+from .tools.models import (
     AgentDetailsResponse,
     AgentOutputResponse,
     AgentPreviewResponse,
@@ -50,11 +50,9 @@ from backend.copilot.tools.models import (
     OperationPendingResponse,
     OperationStartedResponse,
     SetupRequirementsResponse,
-    SuggestedGoalResponse,
     UnderstandingUpdatedResponse,
 )
-from backend.copilot.tracking import track_user_message
-from backend.util.exceptions import NotFoundError
+from .tracking import track_user_message
 
 config = ChatConfig()
 
@@ -132,14 +130,6 @@ class ListSessionsResponse(BaseModel):
     total: int
 
 
-class CancelTaskResponse(BaseModel):
-    """Response model for the cancel task endpoint."""
-
-    cancelled: bool
-    task_id: str | None = None
-    reason: str | None = None
-
-
 class OperationCompleteRequest(BaseModel):
     """Request model for external completion webhook."""
 
@@ -322,57 +312,6 @@ async def get_session(
     )
 
 
-@router.post(
-    "/sessions/{session_id}/cancel",
-    status_code=200,
-)
-async def cancel_session_task(
-    session_id: str,
-    user_id: Annotated[str | None, Depends(auth.get_user_id)],
-) -> CancelTaskResponse:
-    """Cancel the active streaming task for a session.
-
-    Publishes a cancel event to the executor via RabbitMQ FANOUT, then
-    polls Redis until the task status flips from ``running`` or a timeout
-    (5 s) is reached.  Returns only after the cancellation is confirmed.
-    """
-    await _validate_and_get_session(session_id, user_id)
-
-    active_task, _ = await stream_registry.get_active_task_for_session(
-        session_id, user_id
-    )
-    if not active_task:
-        return CancelTaskResponse(cancelled=False, reason="no_active_task")
-
-    task_id = active_task.task_id
-    await enqueue_cancel_task(task_id)
-    logger.info(
-        f"[CANCEL] Published cancel for task ...{task_id[-8:]} "
-        f"session ...{session_id[-8:]}"
-    )
-
-    # Poll until the executor confirms the task is no longer running.
-    # Keep max_wait below typical reverse-proxy read timeouts.
-    poll_interval = 0.5
-    max_wait = 5.0
-    waited = 0.0
-    while waited < max_wait:
-        await asyncio.sleep(poll_interval)
-        waited += poll_interval
-        task = await stream_registry.get_task(task_id)
-        if task is None or task.status != "running":
-            logger.info(
-                f"[CANCEL] Task ...{task_id[-8:]} confirmed stopped "
-                f"(status={task.status if task else 'gone'}) after {waited:.1f}s"
-            )
-            return CancelTaskResponse(cancelled=True, task_id=task_id)
-
-    logger.warning(f"[CANCEL] Task ...{task_id[-8:]} not confirmed after {max_wait}s")
-    return CancelTaskResponse(
-        cancelled=True, task_id=task_id, reason="cancel_published_not_confirmed"
-    )
-
-
 @router.post(
     "/sessions/{session_id}/stream",
 )
@@ -415,7 +354,7 @@ async def stream_chat_post(
         f"user={user_id}, message_len={len(request.message)}",
         extra={"json_fields": log_meta},
     )
-    await _validate_and_get_session(session_id, user_id)
+    session = await _validate_and_get_session(session_id, user_id)
     logger.info(
         f"[TIMING] session validated in {(time.perf_counter() - stream_start_time) * 1000:.1f}ms",
         extra={
@@ -442,7 +381,7 @@ async def stream_chat_post(
                 message_length=len(request.message),
             )
         logger.info(f"[STREAM] Saving user message to session {session_id}")
-        await append_and_save_message(session_id, message)
+        session = await append_and_save_message(session_id, message)
         logger.info(f"[STREAM] User message saved for session {session_id}")
 
     # Create a task in the stream registry for reconnection support
@@ -469,19 +408,125 @@ async def stream_chat_post(
         },
     )
 
-    await enqueue_copilot_task(
-        task_id=task_id,
-        session_id=session_id,
-        user_id=user_id,
-        operation_id=operation_id,
-        message=request.message,
-        is_user_message=request.is_user_message,
-        context=request.context,
-    )
+    # Background task that runs the AI generation independently of SSE connection
+    async def run_ai_generation():
+        import time as time_module
 
+        gen_start_time = time_module.perf_counter()
+        logger.info(
+            f"[TIMING] run_ai_generation STARTED, task={task_id}, session={session_id}, user={user_id}",
+            extra={"json_fields": log_meta},
+        )
+        first_chunk_time, ttfc = None, None
+        chunk_count = 0
+        try:
+            # Emit a start event with task_id for reconnection
+            start_chunk = StreamStart(messageId=task_id, taskId=task_id)
+            await stream_registry.publish_chunk(task_id, start_chunk)
+            logger.info(
+                f"[TIMING] StreamStart published at {(time_module.perf_counter() - gen_start_time) * 1000:.1f}ms",
+                extra={
+                    "json_fields": {
+                        **log_meta,
+                        "elapsed_ms": (time_module.perf_counter() - gen_start_time)
+                        * 1000,
+                    }
+                },
+            )
+
+            # Choose service based on LaunchDarkly flag (falls back to config default)
+            use_sdk = await is_feature_enabled(
+                Flag.COPILOT_SDK,
+                user_id or "anonymous",
+                default=config.use_claude_agent_sdk,
+            )
+            stream_fn = (
+                sdk_service.stream_chat_completion_sdk
+                if use_sdk
+                else chat_service.stream_chat_completion
+            )
+            logger.info(
+                f"[TIMING] Calling {'sdk' if use_sdk else 'standard'} stream_chat_completion",
+                extra={"json_fields": log_meta},
+            )
+            # Pass message=None since we already added it to the session above
+            async for chunk in stream_fn(
+                session_id,
+                None,  # Message already in session
+                is_user_message=request.is_user_message,
+                user_id=user_id,
+                session=session,  # Pass session with message already added
+                context=request.context,
+            ):
+                # Skip duplicate StreamStart — we already published one above
+                if isinstance(chunk, StreamStart):
+                    continue
+                chunk_count += 1
+                if first_chunk_time is None:
+                    first_chunk_time = time_module.perf_counter()
+                    ttfc = first_chunk_time - gen_start_time
+                    logger.info(
+                        f"[TIMING] FIRST AI CHUNK at {ttfc:.2f}s, type={type(chunk).__name__}",
+                        extra={
+                            "json_fields": {
+                                **log_meta,
+                                "chunk_type": type(chunk).__name__,
+                                "time_to_first_chunk_ms": ttfc * 1000,
+                            }
+                        },
+                    )
+                # Write to Redis (subscribers will receive via XREAD)
+                await stream_registry.publish_chunk(task_id, chunk)
+
+            gen_end_time = time_module.perf_counter()
+            total_time = (gen_end_time - gen_start_time) * 1000
+            logger.info(
+                f"[TIMING] run_ai_generation FINISHED in {total_time / 1000:.1f}s; "
+                f"task={task_id}, session={session_id}, "
+                f"ttfc={ttfc or -1:.2f}s, n_chunks={chunk_count}",
+                extra={
+                    "json_fields": {
+                        **log_meta,
+                        "total_time_ms": total_time,
+                        "time_to_first_chunk_ms": (
+                            ttfc * 1000 if ttfc is not None else None
+                        ),
+                        "n_chunks": chunk_count,
+                    }
+                },
+            )
+            await stream_registry.mark_task_completed(task_id, "completed")
+        except Exception as e:
+            elapsed = time_module.perf_counter() - gen_start_time
+            logger.error(
+                f"[TIMING] run_ai_generation ERROR after {elapsed:.2f}s: {e}",
+                extra={
+                    "json_fields": {
+                        **log_meta,
+                        "elapsed_ms": elapsed * 1000,
+                        "error": str(e),
+                    }
+                },
+            )
+            # Publish a StreamError so the frontend can display an error message
+            try:
+                await stream_registry.publish_chunk(
+                    task_id,
+                    StreamError(
+                        errorText="An error occurred. Please try again.",
+                        code="stream_error",
+                    ),
+                )
+            except Exception:
+                pass  # Best-effort; mark_task_completed will publish StreamFinish
+            await stream_registry.mark_task_completed(task_id, "failed")
+
+    # Start the AI generation in a background task
+    bg_task = asyncio.create_task(run_ai_generation())
+    await stream_registry.set_task_asyncio_task(task_id, bg_task)
     setup_time = (time.perf_counter() - stream_start_time) * 1000
     logger.info(
-        f"[TIMING] Task enqueued to RabbitMQ, setup={setup_time:.1f}ms",
+        f"[TIMING] Background task started, setup={setup_time:.1f}ms",
         extra={"json_fields": {**log_meta, "setup_time_ms": setup_time}},
     )
 
@@ -1044,7 +1089,6 @@ ToolResponseUnion = (
     | AgentPreviewResponse
     | AgentSavedResponse
     | ClarificationNeededResponse
-    | SuggestedGoalResponse
     | BlockListResponse
     | BlockDetailsResponse
     | BlockOutputResponse

autogpt_platform/backend/backend/api/features/chat/sdk/response_adapter.py

@@ -0,0 +1,203 @@
+"""Response adapter for converting Claude Agent SDK messages to Vercel AI SDK format.
+
+This module provides the adapter layer that converts streaming messages from
+the Claude Agent SDK into the Vercel AI SDK UI Stream Protocol format that
+the frontend expects.
+"""
+
+import json
+import logging
+import uuid
+
+from claude_agent_sdk import (
+    AssistantMessage,
+    Message,
+    ResultMessage,
+    SystemMessage,
+    TextBlock,
+    ToolResultBlock,
+    ToolUseBlock,
+    UserMessage,
+)
+
+from backend.api.features.chat.response_model import (
+    StreamBaseResponse,
+    StreamError,
+    StreamFinish,
+    StreamFinishStep,
+    StreamStart,
+    StreamStartStep,
+    StreamTextDelta,
+    StreamTextEnd,
+    StreamTextStart,
+    StreamToolInputAvailable,
+    StreamToolInputStart,
+    StreamToolOutputAvailable,
+)
+from backend.api.features.chat.sdk.tool_adapter import (
+    MCP_TOOL_PREFIX,
+    pop_pending_tool_output,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class SDKResponseAdapter:
+    """Adapter for converting Claude Agent SDK messages to Vercel AI SDK format.
+
+    This class maintains state during a streaming session to properly track
+    text blocks, tool calls, and message lifecycle.
+    """
+
+    def __init__(self, message_id: str | None = None):
+        self.message_id = message_id or str(uuid.uuid4())
+        self.text_block_id = str(uuid.uuid4())
+        self.has_started_text = False
+        self.has_ended_text = False
+        self.current_tool_calls: dict[str, dict[str, str]] = {}
+        self.task_id: str | None = None
+        self.step_open = False
+
+    def set_task_id(self, task_id: str) -> None:
+        """Set the task ID for reconnection support."""
+        self.task_id = task_id
+
+    def convert_message(self, sdk_message: Message) -> list[StreamBaseResponse]:
+        """Convert a single SDK message to Vercel AI SDK format."""
+        responses: list[StreamBaseResponse] = []
+
+        if isinstance(sdk_message, SystemMessage):
+            if sdk_message.subtype == "init":
+                responses.append(
+                    StreamStart(messageId=self.message_id, taskId=self.task_id)
+                )
+                # Open the first step (matches non-SDK: StreamStart then StreamStartStep)
+                responses.append(StreamStartStep())
+                self.step_open = True
+
+        elif isinstance(sdk_message, AssistantMessage):
+            # After tool results, the SDK sends a new AssistantMessage for the
+            # next LLM turn. Open a new step if the previous one was closed.
+            if not self.step_open:
+                responses.append(StreamStartStep())
+                self.step_open = True
+
+            for block in sdk_message.content:
+                if isinstance(block, TextBlock):
+                    if block.text:
+                        self._ensure_text_started(responses)
+                        responses.append(
+                            StreamTextDelta(id=self.text_block_id, delta=block.text)
+                        )
+
+                elif isinstance(block, ToolUseBlock):
+                    self._end_text_if_open(responses)
+
+                    # Strip MCP prefix so frontend sees "find_block"
+                    # instead of "mcp__copilot__find_block".
+                    tool_name = block.name.removeprefix(MCP_TOOL_PREFIX)
+
+                    responses.append(
+                        StreamToolInputStart(toolCallId=block.id, toolName=tool_name)
+                    )
+                    responses.append(
+                        StreamToolInputAvailable(
+                            toolCallId=block.id,
+                            toolName=tool_name,
+                            input=block.input,
+                        )
+                    )
+                    self.current_tool_calls[block.id] = {"name": tool_name}
+
+        elif isinstance(sdk_message, UserMessage):
+            # UserMessage carries tool results back from tool execution.
+            content = sdk_message.content
+            blocks = content if isinstance(content, list) else []
+            for block in blocks:
+                if isinstance(block, ToolResultBlock) and block.tool_use_id:
+                    tool_info = self.current_tool_calls.get(block.tool_use_id, {})
+                    tool_name = tool_info.get("name", "unknown")
+
+                    # Prefer the stashed full output over the SDK's
+                    # (potentially truncated) ToolResultBlock content.
+                    # The SDK truncates large results, writing them to disk,
+                    # which breaks frontend widget parsing.
+                    output = pop_pending_tool_output(tool_name) or (
+                        _extract_tool_output(block.content)
+                    )
+
+                    responses.append(
+                        StreamToolOutputAvailable(
+                            toolCallId=block.tool_use_id,
+                            toolName=tool_name,
+                            output=output,
+                            success=not (block.is_error or False),
+                        )
+                    )
+
+            # Close the current step after tool results — the next
+            # AssistantMessage will open a new step for the continuation.
+            if self.step_open:
+                responses.append(StreamFinishStep())
+                self.step_open = False
+
+        elif isinstance(sdk_message, ResultMessage):
+            self._end_text_if_open(responses)
+            # Close the step before finishing.
+            if self.step_open:
+                responses.append(StreamFinishStep())
+                self.step_open = False
+
+            if sdk_message.subtype == "success":
+                responses.append(StreamFinish())
+            elif sdk_message.subtype in ("error", "error_during_execution"):
+                error_msg = getattr(sdk_message, "result", None) or "Unknown error"
+                responses.append(
+                    StreamError(errorText=str(error_msg), code="sdk_error")
+                )
+                responses.append(StreamFinish())
+            else:
+                logger.warning(
+                    f"Unexpected ResultMessage subtype: {sdk_message.subtype}"
+                )
+                responses.append(StreamFinish())
+
+        else:
+            logger.debug(f"Unhandled SDK message type: {type(sdk_message).__name__}")
+
+        return responses
+
+    def _ensure_text_started(self, responses: list[StreamBaseResponse]) -> None:
+        """Start (or restart) a text block if needed."""
+        if not self.has_started_text or self.has_ended_text:
+            if self.has_ended_text:
+                self.text_block_id = str(uuid.uuid4())
+                self.has_ended_text = False
+            responses.append(StreamTextStart(id=self.text_block_id))
+            self.has_started_text = True
+
+    def _end_text_if_open(self, responses: list[StreamBaseResponse]) -> None:
+        """End the current text block if one is open."""
+        if self.has_started_text and not self.has_ended_text:
+            responses.append(StreamTextEnd(id=self.text_block_id))
+            self.has_ended_text = True
+
+
+def _extract_tool_output(content: str | list[dict[str, str]] | None) -> str:
+    """Extract a string output from a ToolResultBlock's content field."""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts = [item.get("text", "") for item in content if item.get("type") == "text"]
+        if parts:
+            return "".join(parts)
+        try:
+            return json.dumps(content)
+        except (TypeError, ValueError):
+            return str(content)
+    if content is None:
+        return ""
+    try:
+        return json.dumps(content)
+    except (TypeError, ValueError):
+        return str(content)

autogpt_platform/backend/backend/api/features/chat/sdk/response_adapter_test.py

@@ -1,8 +1,5 @@
 """Unit tests for the SDK response adapter."""
 
-import asyncio
-
-import pytest
 from claude_agent_sdk import (
     AssistantMessage,
     ResultMessage,
@@ -13,7 +10,7 @@ from claude_agent_sdk import (
     UserMessage,
 )
 
-from backend.copilot.response_model import (
+from backend.api.features.chat.response_model import (
     StreamBaseResponse,
     StreamError,
     StreamFinish,
@@ -30,10 +27,6 @@ from backend.copilot.response_model import (
 
 from .response_adapter import SDKResponseAdapter
 from .tool_adapter import MCP_TOOL_PREFIX
-from .tool_adapter import _pending_tool_outputs as _pto
-from .tool_adapter import _stash_event
-from .tool_adapter import stash_pending_tool_output as _stash
-from .tool_adapter import wait_for_stash
 
 
 def _adapter() -> SDKResponseAdapter:
@@ -371,310 +364,3 @@ def test_full_conversation_flow():
         "StreamFinishStep",  # step 2 closed
         "StreamFinish",
     ]
-
-
-# -- Flush unresolved tool calls --------------------------------------------
-
-
-def test_flush_unresolved_at_result_message():
-    """Built-in tools (WebSearch) without UserMessage results get flushed at ResultMessage."""
-    adapter = _adapter()
-    all_responses: list[StreamBaseResponse] = []
-
-    # 1. Init
-    all_responses.extend(
-        adapter.convert_message(SystemMessage(subtype="init", data={}))
-    )
-    # 2. Tool use (built-in tool — no MCP prefix)
-    all_responses.extend(
-        adapter.convert_message(
-            AssistantMessage(
-                content=[
-                    ToolUseBlock(id="ws-1", name="WebSearch", input={"query": "test"})
-                ],
-                model="test",
-            )
-        )
-    )
-    # 3. No UserMessage for this tool — go straight to ResultMessage
-    all_responses.extend(
-        adapter.convert_message(
-            ResultMessage(
-                subtype="success",
-                duration_ms=100,
-                duration_api_ms=50,
-                is_error=False,
-                num_turns=1,
-                session_id="s1",
-            )
-        )
-    )
-
-    types = [type(r).__name__ for r in all_responses]
-    assert types == [
-        "StreamStart",
-        "StreamStartStep",
-        "StreamToolInputStart",
-        "StreamToolInputAvailable",
-        "StreamToolOutputAvailable",  # flushed with empty output
-        "StreamFinishStep",  # step closed by flush
-        "StreamFinish",
-    ]
-    # The flushed output should be empty (no stash available)
-    output_event = [
-        r for r in all_responses if isinstance(r, StreamToolOutputAvailable)
-    ][0]
-    assert output_event.toolCallId == "ws-1"
-    assert output_event.toolName == "WebSearch"
-    assert output_event.output == ""
-
-
-def test_flush_unresolved_at_next_assistant_message():
-    """Built-in tools get flushed when the next AssistantMessage arrives."""
-    adapter = _adapter()
-    all_responses: list[StreamBaseResponse] = []
-
-    # 1. Init
-    all_responses.extend(
-        adapter.convert_message(SystemMessage(subtype="init", data={}))
-    )
-    # 2. Tool use (built-in — no UserMessage will come)
-    all_responses.extend(
-        adapter.convert_message(
-            AssistantMessage(
-                content=[
-                    ToolUseBlock(id="ws-1", name="WebSearch", input={"query": "test"})
-                ],
-                model="test",
-            )
-        )
-    )
-    # 3. Next AssistantMessage triggers flush before processing its blocks
-    all_responses.extend(
-        adapter.convert_message(
-            AssistantMessage(
-                content=[TextBlock(text="Here are the results")], model="test"
-            )
-        )
-    )
-
-    types = [type(r).__name__ for r in all_responses]
-    assert types == [
-        "StreamStart",
-        "StreamStartStep",
-        "StreamToolInputStart",
-        "StreamToolInputAvailable",
-        # Flush at next AssistantMessage:
-        "StreamToolOutputAvailable",
-        "StreamFinishStep",  # step closed by flush
-        # New step for continuation text:
-        "StreamStartStep",
-        "StreamTextStart",
-        "StreamTextDelta",
-    ]
-
-
-def test_flush_with_stashed_output():
-    """Stashed output from PostToolUse hook is used when flushing."""
-    adapter = _adapter()
-
-    # Simulate PostToolUse hook stashing output
-    _pto.set({})
-    _stash("WebSearch", "Search result: 5 items found")
-
-    all_responses: list[StreamBaseResponse] = []
-
-    # Tool use
-    all_responses.extend(
-        adapter.convert_message(
-            AssistantMessage(
-                content=[
-                    ToolUseBlock(id="ws-1", name="WebSearch", input={"query": "test"})
-                ],
-                model="test",
-            )
-        )
-    )
-    # ResultMessage triggers flush
-    all_responses.extend(
-        adapter.convert_message(
-            ResultMessage(
-                subtype="success",
-                duration_ms=100,
-                duration_api_ms=50,
-                is_error=False,
-                num_turns=1,
-                session_id="s1",
-            )
-        )
-    )
-
-    output_events = [
-        r for r in all_responses if isinstance(r, StreamToolOutputAvailable)
-    ]
-    assert len(output_events) == 1
-    assert output_events[0].output == "Search result: 5 items found"
-
-    # Cleanup
-    _pto.set({})  # type: ignore[arg-type]
-
-
-# -- wait_for_stash synchronisation tests --
-
-
-@pytest.mark.asyncio
-async def test_wait_for_stash_signaled():
-    """wait_for_stash returns True when stash_pending_tool_output signals."""
-    _pto.set({})
-    event = asyncio.Event()
-    _stash_event.set(event)
-
-    # Simulate a PostToolUse hook that stashes output after a short delay
-    async def delayed_stash():
-        await asyncio.sleep(0.01)
-        _stash("WebSearch", "result data")
-
-    asyncio.create_task(delayed_stash())
-    result = await wait_for_stash(timeout=1.0)
-
-    assert result is True
-    assert _pto.get({}).get("WebSearch") == ["result data"]
-
-    # Cleanup
-    _pto.set({})  # type: ignore[arg-type]
-    _stash_event.set(None)
-
-
-@pytest.mark.asyncio
-async def test_wait_for_stash_timeout():
-    """wait_for_stash returns False on timeout when no stash occurs."""
-    _pto.set({})
-    event = asyncio.Event()
-    _stash_event.set(event)
-
-    result = await wait_for_stash(timeout=0.05)
-    assert result is False
-
-    # Cleanup
-    _pto.set({})  # type: ignore[arg-type]
-    _stash_event.set(None)
-
-
-@pytest.mark.asyncio
-async def test_wait_for_stash_already_stashed():
-    """wait_for_stash picks up a stash that happened just before the wait."""
-    _pto.set({})
-    event = asyncio.Event()
-    _stash_event.set(event)
-
-    # Stash before waiting — simulates hook completing before message arrives
-    _stash("Read", "file contents")
-    # Event is now set; wait_for_stash detects the fast path and returns
-    # immediately without timing out.
-    result = await wait_for_stash(timeout=0.05)
-    assert result is True
-
-    # But the stash itself is populated
-    assert _pto.get({}).get("Read") == ["file contents"]
-
-    # Cleanup
-    _pto.set({})  # type: ignore[arg-type]
-    _stash_event.set(None)
-
-
-# -- Parallel tool call tests --
-
-
-def test_parallel_tool_calls_not_flushed_prematurely():
-    """Parallel tool calls should NOT be flushed when the next AssistantMessage
-    only contains ToolUseBlocks (parallel continuation)."""
-    adapter = SDKResponseAdapter()
-
-    # Init
-    adapter.convert_message(SystemMessage(subtype="init", data={}))
-
-    # First AssistantMessage: tool call #1
-    msg1 = AssistantMessage(
-        content=[ToolUseBlock(id="t1", name="WebSearch", input={"q": "foo"})],
-        model="test",
-    )
-    r1 = adapter.convert_message(msg1)
-    assert any(isinstance(r, StreamToolInputAvailable) for r in r1)
-    assert adapter.has_unresolved_tool_calls
-
-    # Second AssistantMessage: tool call #2 (parallel continuation)
-    msg2 = AssistantMessage(
-        content=[ToolUseBlock(id="t2", name="WebSearch", input={"q": "bar"})],
-        model="test",
-    )
-    r2 = adapter.convert_message(msg2)
-
-    # No flush should have happened — t1 should NOT have StreamToolOutputAvailable
-    output_events = [r for r in r2 if isinstance(r, StreamToolOutputAvailable)]
-    assert len(output_events) == 0, (
-        f"Tool-only AssistantMessage should not flush prior tools, "
-        f"but got {len(output_events)} output events"
-    )
-
-    # Both t1 and t2 should still be unresolved
-    assert "t1" not in adapter.resolved_tool_calls
-    assert "t2" not in adapter.resolved_tool_calls
-
-
-def test_text_assistant_message_flushes_prior_tools():
-    """An AssistantMessage with text (new turn) should flush unresolved tools."""
-    adapter = SDKResponseAdapter()
-
-    # Init
-    adapter.convert_message(SystemMessage(subtype="init", data={}))
-
-    # Tool call
-    msg1 = AssistantMessage(
-        content=[ToolUseBlock(id="t1", name="WebSearch", input={"q": "foo"})],
-        model="test",
-    )
-    adapter.convert_message(msg1)
-    assert adapter.has_unresolved_tool_calls
-
-    # Text AssistantMessage (new turn after tools completed)
-    msg2 = AssistantMessage(
-        content=[TextBlock(text="Here are the results")],
-        model="test",
-    )
-    r2 = adapter.convert_message(msg2)
-
-    # Flush SHOULD have happened — t1 gets empty output
-    output_events = [r for r in r2 if isinstance(r, StreamToolOutputAvailable)]
-    assert len(output_events) == 1
-    assert output_events[0].toolCallId == "t1"
-    assert "t1" in adapter.resolved_tool_calls
-
-
-def test_already_resolved_tool_skipped_in_user_message():
-    """A tool result in UserMessage should be skipped if already resolved by flush."""
-    adapter = SDKResponseAdapter()
-
-    adapter.convert_message(SystemMessage(subtype="init", data={}))
-
-    # Tool call + flush via text message
-    adapter.convert_message(
-        AssistantMessage(
-            content=[ToolUseBlock(id="t1", name="WebSearch", input={})],
-            model="test",
-        )
-    )
-    adapter.convert_message(
-        AssistantMessage(
-            content=[TextBlock(text="Done")],
-            model="test",
-        )
-    )
-    assert "t1" in adapter.resolved_tool_calls
-
-    # Now UserMessage arrives with the real result — should be skipped
-    user_msg = UserMessage(content=[ToolResultBlock(tool_use_id="t1", content="real")])
-    r = adapter.convert_message(user_msg)
-    output_events = [r_ for r_ in r if isinstance(r_, StreamToolOutputAvailable)]
-    assert (
-        len(output_events) == 0
-    ), "Already-resolved tool should not emit duplicate output"

autogpt_platform/backend/backend/api/features/chat/sdk/security_hooks.py

@@ -11,12 +11,11 @@ import re
 from collections.abc import Callable
 from typing import Any, cast
 
-from .tool_adapter import (
+from backend.api.features.chat.sdk.tool_adapter import (
     BLOCKED_TOOLS,
     DANGEROUS_PATTERNS,
     MCP_TOOL_PREFIX,
     WORKSPACE_SCOPED_TOOLS,
-    stash_pending_tool_output,
 )
 
 logger = logging.getLogger(__name__)
@@ -124,20 +123,20 @@ def _validate_user_isolation(
     """Validate that tool calls respect user isolation."""
     # For workspace file tools, ensure path doesn't escape
     if "workspace" in tool_name.lower():
-        # The "path" param is a cloud storage key (e.g. "/ASEAN/report.md")
-        # where a leading "/" is normal.  Only check for ".." traversal.
-        # Filesystem paths (source_path, save_to_path) are validated inside
-        # the tool itself via _validate_ephemeral_path.
         path = tool_input.get("path", "") or tool_input.get("file_path", "")
-        if path and ".." in path:
-            logger.warning(f"Blocked path traversal attempt: {path} by user {user_id}")
-            return {
-                "hookSpecificOutput": {
-                    "hookEventName": "PreToolUse",
-                    "permissionDecision": "deny",
-                    "permissionDecisionReason": "Path traversal not allowed",
+        if path:
+            # Check for path traversal
+            if ".." in path or path.startswith("/"):
+                logger.warning(
+                    f"Blocked path traversal attempt: {path} by user {user_id}"
+                )
+                return {
+                    "hookSpecificOutput": {
+                        "hookEventName": "PreToolUse",
+                        "permissionDecision": "deny",
+                        "permissionDecisionReason": "Path traversal not allowed",
+                    }
                 }
-            }
 
     return {}
 
@@ -188,19 +187,8 @@ def create_security_hooks(
 
             # Rate-limit Task (sub-agent) spawns per session
             if tool_name == "Task":
-                # Block background task execution first — denied calls
-                # should not consume a subtask slot.
-                if tool_input.get("run_in_background"):
-                    logger.info(f"[SDK] Blocked background Task, user={user_id}")
-                    return cast(
-                        SyncHookJSONOutput,
-                        _deny(
-                            "Background task execution is not supported. "
-                            "Run tasks in the foreground instead "
-                            "(remove the run_in_background parameter)."
-                        ),
-                    )
-                if task_spawn_count >= max_subtasks:
+                task_spawn_count += 1
+                if task_spawn_count > max_subtasks:
                     logger.warning(
                         f"[SDK] Task limit reached ({max_subtasks}), user={user_id}"
                     )
@@ -211,7 +199,6 @@ def create_security_hooks(
                             "Please continue in the main conversation."
                         ),
                     )
-                task_spawn_count += 1
 
             # Strip MCP prefix for consistent validation
             is_copilot_tool = tool_name.startswith(MCP_TOOL_PREFIX)
@@ -237,43 +224,10 @@ def create_security_hooks(
             tool_use_id: str | None,
             context: HookContext,
         ) -> SyncHookJSONOutput:
-            """Log successful tool executions and stash SDK built-in tool outputs.
-
-            MCP tools stash their output in ``_execute_tool_sync`` before the
-            SDK can truncate it.  SDK built-in tools (WebSearch, Read, etc.)
-            are executed by the CLI internally — this hook captures their
-            output so the response adapter can forward it to the frontend.
-            """
+            """Log successful tool executions for observability."""
             _ = context
             tool_name = cast(str, input_data.get("tool_name", ""))
-            is_builtin = not tool_name.startswith(MCP_TOOL_PREFIX)
-            logger.info(
-                "[SDK] PostToolUse: %s (builtin=%s, tool_use_id=%s)",
-                tool_name,
-                is_builtin,
-                (tool_use_id or "")[:12],
-            )
-
-            # Stash output for SDK built-in tools so the response adapter can
-            # emit StreamToolOutputAvailable even when the CLI doesn't surface
-            # a separate UserMessage with ToolResultBlock content.
-            if is_builtin:
-                tool_response = input_data.get("tool_response")
-                if tool_response is not None:
-                    resp_preview = str(tool_response)[:100]
-                    logger.info(
-                        "[SDK] Stashing builtin output for %s (%d chars): %s...",
-                        tool_name,
-                        len(str(tool_response)),
-                        resp_preview,
-                    )
-                    stash_pending_tool_output(tool_name, tool_response)
-                else:
-                    logger.warning(
-                        "[SDK] PostToolUse for builtin %s but tool_response is None",
-                        tool_name,
-                    )
-
+            logger.debug(f"[SDK] Tool success: {tool_name}, tool_use_id={tool_use_id}")
             return cast(SyncHookJSONOutput, {})
 
         async def post_tool_failure_hook(

autogpt_platform/backend/backend/api/features/chat/sdk/security_hooks_test.py

@@ -0,0 +1,165 @@
+"""Unit tests for SDK security hooks."""
+
+import os
+
+from .security_hooks import _validate_tool_access, _validate_user_isolation
+
+SDK_CWD = "/tmp/copilot-abc123"
+
+
+def _is_denied(result: dict) -> bool:
+    hook = result.get("hookSpecificOutput", {})
+    return hook.get("permissionDecision") == "deny"
+
+
+# -- Blocked tools -----------------------------------------------------------
+
+
+def test_blocked_tools_denied():
+    for tool in ("bash", "shell", "exec", "terminal", "command"):
+        result = _validate_tool_access(tool, {})
+        assert _is_denied(result), f"{tool} should be blocked"
+
+
+def test_unknown_tool_allowed():
+    result = _validate_tool_access("SomeCustomTool", {})
+    assert result == {}
+
+
+# -- Workspace-scoped tools --------------------------------------------------
+
+
+def test_read_within_workspace_allowed():
+    result = _validate_tool_access(
+        "Read", {"file_path": f"{SDK_CWD}/file.txt"}, sdk_cwd=SDK_CWD
+    )
+    assert result == {}
+
+
+def test_write_within_workspace_allowed():
+    result = _validate_tool_access(
+        "Write", {"file_path": f"{SDK_CWD}/output.json"}, sdk_cwd=SDK_CWD
+    )
+    assert result == {}
+
+
+def test_edit_within_workspace_allowed():
+    result = _validate_tool_access(
+        "Edit", {"file_path": f"{SDK_CWD}/src/main.py"}, sdk_cwd=SDK_CWD
+    )
+    assert result == {}
+
+
+def test_glob_within_workspace_allowed():
+    result = _validate_tool_access("Glob", {"path": f"{SDK_CWD}/src"}, sdk_cwd=SDK_CWD)
+    assert result == {}
+
+
+def test_grep_within_workspace_allowed():
+    result = _validate_tool_access("Grep", {"path": f"{SDK_CWD}/src"}, sdk_cwd=SDK_CWD)
+    assert result == {}
+
+
+def test_read_outside_workspace_denied():
+    result = _validate_tool_access(
+        "Read", {"file_path": "/etc/passwd"}, sdk_cwd=SDK_CWD
+    )
+    assert _is_denied(result)
+
+
+def test_write_outside_workspace_denied():
+    result = _validate_tool_access(
+        "Write", {"file_path": "/home/user/secrets.txt"}, sdk_cwd=SDK_CWD
+    )
+    assert _is_denied(result)
+
+
+def test_traversal_attack_denied():
+    result = _validate_tool_access(
+        "Read",
+        {"file_path": f"{SDK_CWD}/../../etc/passwd"},
+        sdk_cwd=SDK_CWD,
+    )
+    assert _is_denied(result)
+
+
+def test_no_path_allowed():
+    """Glob/Grep without a path argument defaults to cwd — should pass."""
+    result = _validate_tool_access("Glob", {}, sdk_cwd=SDK_CWD)
+    assert result == {}
+
+
+def test_read_no_cwd_denies_absolute():
+    """If no sdk_cwd is set, absolute paths are denied."""
+    result = _validate_tool_access("Read", {"file_path": "/tmp/anything"})
+    assert _is_denied(result)
+
+
+# -- Tool-results directory --------------------------------------------------
+
+
+def test_read_tool_results_allowed():
+    home = os.path.expanduser("~")
+    path = f"{home}/.claude/projects/-tmp-copilot-abc123/tool-results/12345.txt"
+    result = _validate_tool_access("Read", {"file_path": path}, sdk_cwd=SDK_CWD)
+    assert result == {}
+
+
+def test_read_claude_projects_without_tool_results_denied():
+    home = os.path.expanduser("~")
+    path = f"{home}/.claude/projects/-tmp-copilot-abc123/settings.json"
+    result = _validate_tool_access("Read", {"file_path": path}, sdk_cwd=SDK_CWD)
+    assert _is_denied(result)
+
+
+# -- Built-in Bash is blocked (use bash_exec MCP tool instead) ---------------
+
+
+def test_bash_builtin_always_blocked():
+    """SDK built-in Bash is blocked — bash_exec MCP tool with bubblewrap is used instead."""
+    result = _validate_tool_access("Bash", {"command": "echo hello"}, sdk_cwd=SDK_CWD)
+    assert _is_denied(result)
+
+
+# -- Dangerous patterns ------------------------------------------------------
+
+
+def test_dangerous_pattern_blocked():
+    result = _validate_tool_access("SomeTool", {"cmd": "sudo rm -rf /"})
+    assert _is_denied(result)
+
+
+def test_subprocess_pattern_blocked():
+    result = _validate_tool_access("SomeTool", {"code": "subprocess.run(...)"})
+    assert _is_denied(result)
+
+
+# -- User isolation ----------------------------------------------------------
+
+
+def test_workspace_path_traversal_blocked():
+    result = _validate_user_isolation(
+        "workspace_read", {"path": "../../../etc/shadow"}, user_id="user-1"
+    )
+    assert _is_denied(result)
+
+
+def test_workspace_absolute_path_blocked():
+    result = _validate_user_isolation(
+        "workspace_read", {"path": "/etc/passwd"}, user_id="user-1"
+    )
+    assert _is_denied(result)
+
+
+def test_workspace_normal_path_allowed():
+    result = _validate_user_isolation(
+        "workspace_read", {"path": "src/main.py"}, user_id="user-1"
+    )
+    assert result == {}
+
+
+def test_non_workspace_tool_passes_isolation():
+    result = _validate_user_isolation(
+        "find_agent", {"query": "email"}, user_id="user-1"
+    )
+    assert result == {}

autogpt_platform/backend/backend/api/features/chat/sdk/service.py

@@ -0,0 +1,752 @@
+"""Claude Agent SDK service layer for CoPilot chat completions."""
+
+import asyncio
+import json
+import logging
+import os
+import uuid
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass
+from typing import Any
+
+from backend.util.exceptions import NotFoundError
+
+from .. import stream_registry
+from ..config import ChatConfig
+from ..model import (
+    ChatMessage,
+    ChatSession,
+    get_chat_session,
+    update_session_title,
+    upsert_chat_session,
+)
+from ..response_model import (
+    StreamBaseResponse,
+    StreamError,
+    StreamFinish,
+    StreamStart,
+    StreamTextDelta,
+    StreamToolInputAvailable,
+    StreamToolOutputAvailable,
+)
+from ..service import (
+    _build_system_prompt,
+    _execute_long_running_tool_with_streaming,
+    _generate_session_title,
+)
+from ..tools.models import OperationPendingResponse, OperationStartedResponse
+from ..tools.sandbox import WORKSPACE_PREFIX, make_session_path
+from ..tracking import track_user_message
+from .response_adapter import SDKResponseAdapter
+from .security_hooks import create_security_hooks
+from .tool_adapter import (
+    COPILOT_TOOL_NAMES,
+    SDK_DISALLOWED_TOOLS,
+    LongRunningCallback,
+    create_copilot_mcp_server,
+    set_execution_context,
+)
+from .transcript import (
+    download_transcript,
+    read_transcript_file,
+    upload_transcript,
+    validate_transcript,
+    write_transcript_to_tempfile,
+)
+
+logger = logging.getLogger(__name__)
+config = ChatConfig()
+
+# Set to hold background tasks to prevent garbage collection
+_background_tasks: set[asyncio.Task[Any]] = set()
+
+
+@dataclass
+class CapturedTranscript:
+    """Info captured by the SDK Stop hook for stateless --resume."""
+
+    path: str = ""
+    sdk_session_id: str = ""
+
+    @property
+    def available(self) -> bool:
+        return bool(self.path)
+
+
+_SDK_CWD_PREFIX = WORKSPACE_PREFIX
+
+# Appended to the system prompt to inform the agent about available tools.
+# The SDK built-in Bash is NOT available — use mcp__copilot__bash_exec instead,
+# which has kernel-level network isolation (unshare --net).
+_SDK_TOOL_SUPPLEMENT = """
+
+## Tool notes
+
+- The SDK built-in Bash tool is NOT available.  Use the `bash_exec` MCP tool
+  for shell commands — it runs in a network-isolated sandbox.
+- **Shared workspace**: The SDK Read/Write tools and `bash_exec` share the
+  same working directory. Files created by one are readable by the other.
+  These files are **ephemeral** — they exist only for the current session.
+- **Persistent storage**: Use `write_workspace_file` / `read_workspace_file`
+  for files that should persist across sessions (stored in cloud storage).
+- Long-running tools (create_agent, edit_agent, etc.) are handled
+  asynchronously.  You will receive an immediate response; the actual result
+  is delivered to the user via a background stream.
+"""
+
+
+def _build_long_running_callback(user_id: str | None) -> LongRunningCallback:
+    """Build a callback that delegates long-running tools to the non-SDK infrastructure.
+
+    Long-running tools (create_agent, edit_agent, etc.) are delegated to the
+    existing background infrastructure: stream_registry (Redis Streams),
+    database persistence, and SSE reconnection.  This means results survive
+    page refreshes / pod restarts, and the frontend shows the proper loading
+    widget with progress updates.
+
+    The returned callback matches the ``LongRunningCallback`` signature:
+    ``(tool_name, args, session) -> MCP response dict``.
+    """
+
+    async def _callback(
+        tool_name: str, args: dict[str, Any], session: ChatSession
+    ) -> dict[str, Any]:
+        operation_id = str(uuid.uuid4())
+        task_id = str(uuid.uuid4())
+        tool_call_id = f"sdk-{uuid.uuid4().hex[:12]}"
+        session_id = session.session_id
+
+        # --- Build user-friendly messages (matches non-SDK service) ---
+        if tool_name == "create_agent":
+            desc = args.get("description", "")
+            desc_preview = (desc[:100] + "...") if len(desc) > 100 else desc
+            pending_msg = (
+                f"Creating your agent: {desc_preview}"
+                if desc_preview
+                else "Creating agent... This may take a few minutes."
+            )
+            started_msg = (
+                "Agent creation started. You can close this tab - "
+                "check your library in a few minutes."
+            )
+        elif tool_name == "edit_agent":
+            changes = args.get("changes", "")
+            changes_preview = (changes[:100] + "...") if len(changes) > 100 else changes
+            pending_msg = (
+                f"Editing agent: {changes_preview}"
+                if changes_preview
+                else "Editing agent... This may take a few minutes."
+            )
+            started_msg = (
+                "Agent edit started. You can close this tab - "
+                "check your library in a few minutes."
+            )
+        else:
+            pending_msg = f"Running {tool_name}... This may take a few minutes."
+            started_msg = (
+                f"{tool_name} started. You can close this tab - "
+                "check back in a few minutes."
+            )
+
+        # --- Register task in Redis for SSE reconnection ---
+        await stream_registry.create_task(
+            task_id=task_id,
+            session_id=session_id,
+            user_id=user_id,
+            tool_call_id=tool_call_id,
+            tool_name=tool_name,
+            operation_id=operation_id,
+        )
+
+        # --- Save OperationPendingResponse to chat history ---
+        pending_message = ChatMessage(
+            role="tool",
+            content=OperationPendingResponse(
+                message=pending_msg,
+                operation_id=operation_id,
+                tool_name=tool_name,
+            ).model_dump_json(),
+            tool_call_id=tool_call_id,
+        )
+        session.messages.append(pending_message)
+        await upsert_chat_session(session)
+
+        # --- Spawn background task (reuses non-SDK infrastructure) ---
+        bg_task = asyncio.create_task(
+            _execute_long_running_tool_with_streaming(
+                tool_name=tool_name,
+                parameters=args,
+                tool_call_id=tool_call_id,
+                operation_id=operation_id,
+                task_id=task_id,
+                session_id=session_id,
+                user_id=user_id,
+            )
+        )
+        _background_tasks.add(bg_task)
+        bg_task.add_done_callback(_background_tasks.discard)
+        await stream_registry.set_task_asyncio_task(task_id, bg_task)
+
+        logger.info(
+            f"[SDK] Long-running tool {tool_name} delegated to background "
+            f"(operation_id={operation_id}, task_id={task_id})"
+        )
+
+        # --- Return OperationStartedResponse as MCP tool result ---
+        # This flows through SDK → response adapter → frontend, triggering
+        # the loading widget with SSE reconnection support.
+        started_json = OperationStartedResponse(
+            message=started_msg,
+            operation_id=operation_id,
+            tool_name=tool_name,
+            task_id=task_id,
+        ).model_dump_json()
+
+        return {
+            "content": [{"type": "text", "text": started_json}],
+            "isError": False,
+        }
+
+    return _callback
+
+
+def _resolve_sdk_model() -> str | None:
+    """Resolve the model name for the Claude Agent SDK CLI.
+
+    Uses ``config.claude_agent_model`` if set, otherwise derives from
+    ``config.model`` by stripping the OpenRouter provider prefix (e.g.,
+    ``"anthropic/claude-opus-4.6"`` → ``"claude-opus-4.6"``).
+    """
+    if config.claude_agent_model:
+        return config.claude_agent_model
+    model = config.model
+    if "/" in model:
+        return model.split("/", 1)[1]
+    return model
+
+
+def _build_sdk_env() -> dict[str, str]:
+    """Build env vars for the SDK CLI process.
+
+    Routes API calls through OpenRouter (or a custom base_url) using
+    the same ``config.api_key`` / ``config.base_url`` as the non-SDK path.
+    This gives per-call token and cost tracking on the OpenRouter dashboard.
+
+    Only overrides ``ANTHROPIC_API_KEY`` when a valid proxy URL and auth
+    token are both present — otherwise returns an empty dict so the SDK
+    falls back to its default credentials.
+    """
+    env: dict[str, str] = {}
+    if config.api_key and config.base_url:
+        # Strip /v1 suffix — SDK expects the base URL without a version path
+        base = config.base_url.rstrip("/")
+        if base.endswith("/v1"):
+            base = base[:-3]
+        if not base or not base.startswith("http"):
+            # Invalid base_url — don't override SDK defaults
+            return env
+        env["ANTHROPIC_BASE_URL"] = base
+        env["ANTHROPIC_AUTH_TOKEN"] = config.api_key
+        # Must be explicitly empty so the CLI uses AUTH_TOKEN instead
+        env["ANTHROPIC_API_KEY"] = ""
+    return env
+
+
+def _make_sdk_cwd(session_id: str) -> str:
+    """Create a safe, session-specific working directory path.
+
+    Delegates to :func:`~backend.api.features.chat.tools.sandbox.make_session_path`
+    (single source of truth for path sanitization) and adds a defence-in-depth
+    assertion.
+    """
+    cwd = make_session_path(session_id)
+    # Defence-in-depth: normpath + startswith is a CodeQL-recognised sanitizer
+    cwd = os.path.normpath(cwd)
+    if not cwd.startswith(_SDK_CWD_PREFIX):
+        raise ValueError(f"SDK cwd escaped prefix: {cwd}")
+    return cwd
+
+
+def _cleanup_sdk_tool_results(cwd: str) -> None:
+    """Remove SDK tool-result files for a specific session working directory.
+
+    The SDK creates tool-result files under ~/.claude/projects/<encoded-cwd>/tool-results/.
+    We clean only the specific cwd's results to avoid race conditions between
+    concurrent sessions.
+
+    Security: cwd MUST be created by _make_sdk_cwd() which sanitizes session_id.
+    """
+    import shutil
+
+    # Validate cwd is under the expected prefix
+    normalized = os.path.normpath(cwd)
+    if not normalized.startswith(_SDK_CWD_PREFIX):
+        logger.warning(f"[SDK] Rejecting cleanup for path outside workspace: {cwd}")
+        return
+
+    # SDK encodes the cwd path by replacing '/' with '-'
+    encoded_cwd = normalized.replace("/", "-")
+
+    # Construct the project directory path (known-safe home expansion)
+    claude_projects = os.path.expanduser("~/.claude/projects")
+    project_dir = os.path.join(claude_projects, encoded_cwd)
+
+    # Security check 3: Validate project_dir is under ~/.claude/projects
+    project_dir = os.path.normpath(project_dir)
+    if not project_dir.startswith(claude_projects):
+        logger.warning(
+            f"[SDK] Rejecting cleanup for escaped project path: {project_dir}"
+        )
+        return
+
+    results_dir = os.path.join(project_dir, "tool-results")
+    if os.path.isdir(results_dir):
+        for filename in os.listdir(results_dir):
+            file_path = os.path.join(results_dir, filename)
+            try:
+                if os.path.isfile(file_path):
+                    os.remove(file_path)
+            except OSError:
+                pass
+
+    # Also clean up the temp cwd directory itself
+    try:
+        shutil.rmtree(normalized, ignore_errors=True)
+    except OSError:
+        pass
+
+
+async def _compress_conversation_history(
+    session: ChatSession,
+) -> list[ChatMessage]:
+    """Compress prior conversation messages if they exceed the token threshold.
+
+    Uses the shared compress_context() from prompt.py which supports:
+    - LLM summarization of old messages (keeps recent ones intact)
+    - Progressive content truncation as fallback
+    - Middle-out deletion as last resort
+
+    Returns the compressed prior messages (everything except the current message).
+    """
+    prior = session.messages[:-1]
+    if len(prior) < 2:
+        return prior
+
+    from backend.util.prompt import compress_context
+
+    # Convert ChatMessages to dicts for compress_context
+    messages_dict = []
+    for msg in prior:
+        msg_dict: dict[str, Any] = {"role": msg.role}
+        if msg.content:
+            msg_dict["content"] = msg.content
+        if msg.tool_calls:
+            msg_dict["tool_calls"] = msg.tool_calls
+        if msg.tool_call_id:
+            msg_dict["tool_call_id"] = msg.tool_call_id
+        messages_dict.append(msg_dict)
+
+    try:
+        import openai
+
+        async with openai.AsyncOpenAI(
+            api_key=config.api_key, base_url=config.base_url, timeout=30.0
+        ) as client:
+            result = await compress_context(
+                messages=messages_dict,
+                model=config.model,
+                client=client,
+            )
+    except Exception as e:
+        logger.warning(f"[SDK] Context compression with LLM failed: {e}")
+        # Fall back to truncation-only (no LLM summarization)
+        result = await compress_context(
+            messages=messages_dict,
+            model=config.model,
+            client=None,
+        )
+
+    if result.was_compacted:
+        logger.info(
+            f"[SDK] Context compacted: {result.original_token_count} -> "
+            f"{result.token_count} tokens "
+            f"({result.messages_summarized} summarized, "
+            f"{result.messages_dropped} dropped)"
+        )
+        # Convert compressed dicts back to ChatMessages
+        return [
+            ChatMessage(
+                role=m["role"],
+                content=m.get("content"),
+                tool_calls=m.get("tool_calls"),
+                tool_call_id=m.get("tool_call_id"),
+            )
+            for m in result.messages
+        ]
+
+    return prior
+
+
+def _format_conversation_context(messages: list[ChatMessage]) -> str | None:
+    """Format conversation messages into a context prefix for the user message.
+
+    Returns a string like:
+        <conversation_history>
+        User: hello
+        You responded: Hi! How can I help?
+        </conversation_history>
+
+    Returns None if there are no messages to format.
+    """
+    if not messages:
+        return None
+
+    lines: list[str] = []
+    for msg in messages:
+        if not msg.content:
+            continue
+        if msg.role == "user":
+            lines.append(f"User: {msg.content}")
+        elif msg.role == "assistant":
+            lines.append(f"You responded: {msg.content}")
+        # Skip tool messages — they're internal details
+
+    if not lines:
+        return None
+
+    return "<conversation_history>\n" + "\n".join(lines) + "\n</conversation_history>"
+
+
+async def stream_chat_completion_sdk(
+    session_id: str,
+    message: str | None = None,
+    tool_call_response: str | None = None,  # noqa: ARG001
+    is_user_message: bool = True,
+    user_id: str | None = None,
+    retry_count: int = 0,  # noqa: ARG001
+    session: ChatSession | None = None,
+    context: dict[str, str] | None = None,  # noqa: ARG001
+) -> AsyncGenerator[StreamBaseResponse, None]:
+    """Stream chat completion using Claude Agent SDK.
+
+    Drop-in replacement for stream_chat_completion with improved reliability.
+    """
+
+    if session is None:
+        session = await get_chat_session(session_id, user_id)
+
+    if not session:
+        raise NotFoundError(
+            f"Session {session_id} not found. Please create a new session first."
+        )
+
+    if message:
+        session.messages.append(
+            ChatMessage(
+                role="user" if is_user_message else "assistant", content=message
+            )
+        )
+        if is_user_message:
+            track_user_message(
+                user_id=user_id, session_id=session_id, message_length=len(message)
+            )
+
+    session = await upsert_chat_session(session)
+
+    # Generate title for new sessions (first user message)
+    if is_user_message and not session.title:
+        user_messages = [m for m in session.messages if m.role == "user"]
+        if len(user_messages) == 1:
+            first_message = user_messages[0].content or message or ""
+            if first_message:
+                task = asyncio.create_task(
+                    _update_title_async(session_id, first_message, user_id)
+                )
+                _background_tasks.add(task)
+                task.add_done_callback(_background_tasks.discard)
+
+    # Build system prompt (reuses non-SDK path with Langfuse support)
+    has_history = len(session.messages) > 1
+    system_prompt, _ = await _build_system_prompt(
+        user_id, has_conversation_history=has_history
+    )
+    system_prompt += _SDK_TOOL_SUPPLEMENT
+    message_id = str(uuid.uuid4())
+    task_id = str(uuid.uuid4())
+
+    yield StreamStart(messageId=message_id, taskId=task_id)
+
+    stream_completed = False
+    # Initialise sdk_cwd before the try so the finally can reference it
+    # even if _make_sdk_cwd raises (in that case it stays as "").
+    sdk_cwd = ""
+    use_resume = False
+
+    try:
+        # Use a session-specific temp dir to avoid cleanup race conditions
+        # between concurrent sessions.
+        sdk_cwd = _make_sdk_cwd(session_id)
+        os.makedirs(sdk_cwd, exist_ok=True)
+
+        set_execution_context(
+            user_id,
+            session,
+            long_running_callback=_build_long_running_callback(user_id),
+        )
+        try:
+            from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient
+
+            # Fail fast when no API credentials are available at all
+            sdk_env = _build_sdk_env()
+            if not sdk_env and not os.environ.get("ANTHROPIC_API_KEY"):
+                raise RuntimeError(
+                    "No API key configured. Set OPEN_ROUTER_API_KEY "
+                    "(or CHAT_API_KEY) for OpenRouter routing, "
+                    "or ANTHROPIC_API_KEY for direct Anthropic access."
+                )
+
+            mcp_server = create_copilot_mcp_server()
+
+            sdk_model = _resolve_sdk_model()
+
+            # --- Transcript capture via Stop hook ---
+            captured_transcript = CapturedTranscript()
+
+            def _on_stop(transcript_path: str, sdk_session_id: str) -> None:
+                captured_transcript.path = transcript_path
+                captured_transcript.sdk_session_id = sdk_session_id
+
+            security_hooks = create_security_hooks(
+                user_id,
+                sdk_cwd=sdk_cwd,
+                max_subtasks=config.claude_agent_max_subtasks,
+                on_stop=_on_stop if config.claude_agent_use_resume else None,
+            )
+
+            # --- Resume strategy: download transcript from bucket ---
+            resume_file: str | None = None
+            use_resume = False
+
+            if config.claude_agent_use_resume and user_id and len(session.messages) > 1:
+                transcript_content = await download_transcript(user_id, session_id)
+                if transcript_content and validate_transcript(transcript_content):
+                    resume_file = write_transcript_to_tempfile(
+                        transcript_content, session_id, sdk_cwd
+                    )
+                    if resume_file:
+                        use_resume = True
+                        logger.info(
+                            f"[SDK] Using --resume with transcript "
+                            f"({len(transcript_content)} bytes)"
+                        )
+
+            sdk_options_kwargs: dict[str, Any] = {
+                "system_prompt": system_prompt,
+                "mcp_servers": {"copilot": mcp_server},
+                "allowed_tools": COPILOT_TOOL_NAMES,
+                "disallowed_tools": SDK_DISALLOWED_TOOLS,
+                "hooks": security_hooks,
+                "cwd": sdk_cwd,
+                "max_buffer_size": config.claude_agent_max_buffer_size,
+            }
+            if sdk_env:
+                sdk_options_kwargs["model"] = sdk_model
+                sdk_options_kwargs["env"] = sdk_env
+            if use_resume and resume_file:
+                sdk_options_kwargs["resume"] = resume_file
+
+            options = ClaudeAgentOptions(**sdk_options_kwargs)  # type: ignore[arg-type]
+
+            adapter = SDKResponseAdapter(message_id=message_id)
+            adapter.set_task_id(task_id)
+
+            async with ClaudeSDKClient(options=options) as client:
+                current_message = message or ""
+                if not current_message and session.messages:
+                    last_user = [m for m in session.messages if m.role == "user"]
+                    if last_user:
+                        current_message = last_user[-1].content or ""
+
+                if not current_message.strip():
+                    yield StreamError(
+                        errorText="Message cannot be empty.",
+                        code="empty_prompt",
+                    )
+                    yield StreamFinish()
+                    return
+
+                # Build query: with --resume the CLI already has full
+                # context, so we only send the new message.  Without
+                # resume, compress history into a context prefix.
+                query_message = current_message
+                if not use_resume and len(session.messages) > 1:
+                    logger.warning(
+                        f"[SDK] Using compression fallback for session "
+                        f"{session_id} ({len(session.messages)} messages) — "
+                        f"no transcript available for --resume"
+                    )
+                    compressed = await _compress_conversation_history(session)
+                    history_context = _format_conversation_context(compressed)
+                    if history_context:
+                        query_message = (
+                            f"{history_context}\n\n"
+                            f"Now, the user says:\n{current_message}"
+                        )
+
+                logger.info(
+                    f"[SDK] Sending query ({len(session.messages)} msgs in session)"
+                )
+                logger.debug(f"[SDK] Query preview: {current_message[:80]!r}")
+                await client.query(query_message, session_id=session_id)
+
+                assistant_response = ChatMessage(role="assistant", content="")
+                accumulated_tool_calls: list[dict[str, Any]] = []
+                has_appended_assistant = False
+                has_tool_results = False
+
+                async for sdk_msg in client.receive_messages():
+                    logger.debug(
+                        f"[SDK] Received: {type(sdk_msg).__name__} "
+                        f"{getattr(sdk_msg, 'subtype', '')}"
+                    )
+                    for response in adapter.convert_message(sdk_msg):
+                        if isinstance(response, StreamStart):
+                            continue
+
+                        yield response
+
+                        if isinstance(response, StreamTextDelta):
+                            delta = response.delta or ""
+                            # After tool results, start a new assistant
+                            # message for the post-tool text.
+                            if has_tool_results and has_appended_assistant:
+                                assistant_response = ChatMessage(
+                                    role="assistant", content=delta
+                                )
+                                accumulated_tool_calls = []
+                                has_appended_assistant = False
+                                has_tool_results = False
+                                session.messages.append(assistant_response)
+                                has_appended_assistant = True
+                            else:
+                                assistant_response.content = (
+                                    assistant_response.content or ""
+                                ) + delta
+                                if not has_appended_assistant:
+                                    session.messages.append(assistant_response)
+                                    has_appended_assistant = True
+
+                        elif isinstance(response, StreamToolInputAvailable):
+                            accumulated_tool_calls.append(
+                                {
+                                    "id": response.toolCallId,
+                                    "type": "function",
+                                    "function": {
+                                        "name": response.toolName,
+                                        "arguments": json.dumps(response.input or {}),
+                                    },
+                                }
+                            )
+                            assistant_response.tool_calls = accumulated_tool_calls
+                            if not has_appended_assistant:
+                                session.messages.append(assistant_response)
+                                has_appended_assistant = True
+
+                        elif isinstance(response, StreamToolOutputAvailable):
+                            session.messages.append(
+                                ChatMessage(
+                                    role="tool",
+                                    content=(
+                                        response.output
+                                        if isinstance(response.output, str)
+                                        else str(response.output)
+                                    ),
+                                    tool_call_id=response.toolCallId,
+                                )
+                            )
+                            has_tool_results = True
+
+                        elif isinstance(response, StreamFinish):
+                            stream_completed = True
+
+                    if stream_completed:
+                        break
+
+                if (
+                    assistant_response.content or assistant_response.tool_calls
+                ) and not has_appended_assistant:
+                    session.messages.append(assistant_response)
+
+                # --- Capture transcript while CLI is still alive ---
+                # Must happen INSIDE async with: close() sends SIGTERM
+                # which kills the CLI before it can flush the JSONL.
+                if (
+                    config.claude_agent_use_resume
+                    and user_id
+                    and captured_transcript.available
+                ):
+                    # Give CLI time to flush JSONL writes before we read
+                    await asyncio.sleep(0.5)
+                    raw_transcript = read_transcript_file(captured_transcript.path)
+                    if raw_transcript:
+                        task = asyncio.create_task(
+                            _upload_transcript_bg(user_id, session_id, raw_transcript)
+                        )
+                        _background_tasks.add(task)
+                        task.add_done_callback(_background_tasks.discard)
+                    else:
+                        logger.debug("[SDK] Stop hook fired but transcript not usable")
+
+        except ImportError:
+            raise RuntimeError(
+                "claude-agent-sdk is not installed. "
+                "Disable SDK mode (CHAT_USE_CLAUDE_AGENT_SDK=false) "
+                "to use the OpenAI-compatible fallback."
+            )
+
+        await upsert_chat_session(session)
+        logger.debug(
+            f"[SDK] Session {session_id} saved with {len(session.messages)} messages"
+        )
+        if not stream_completed:
+            yield StreamFinish()
+
+    except Exception as e:
+        logger.error(f"[SDK] Error: {e}", exc_info=True)
+        try:
+            await upsert_chat_session(session)
+        except Exception as save_err:
+            logger.error(f"[SDK] Failed to save session on error: {save_err}")
+        yield StreamError(
+            errorText="An error occurred. Please try again.",
+            code="sdk_error",
+        )
+        yield StreamFinish()
+    finally:
+        if sdk_cwd:
+            _cleanup_sdk_tool_results(sdk_cwd)
+
+
+async def _upload_transcript_bg(
+    user_id: str, session_id: str, raw_content: str
+) -> None:
+    """Background task to strip progress entries and upload transcript."""
+    try:
+        await upload_transcript(user_id, session_id, raw_content)
+    except Exception as e:
+        logger.error(f"[SDK] Failed to upload transcript for {session_id}: {e}")
+
+
+async def _update_title_async(
+    session_id: str, message: str, user_id: str | None = None
+) -> None:
+    """Background task to update session title."""
+    try:
+        title = await _generate_session_title(
+            message, user_id=user_id, session_id=session_id
+        )
+        if title:
+            await update_session_title(session_id, title)
+            logger.debug(f"[SDK] Generated title for {session_id}: {title}")
+    except Exception as e:
+        logger.warning(f"[SDK] Failed to update session title: {e}")

autogpt_platform/backend/backend/api/features/chat/sdk/tool_adapter.py

@@ -9,7 +9,6 @@ via a callback provided by the service layer.  This avoids wasteful SDK polling
 and makes results survive page refreshes.
 """
 
-import asyncio
 import itertools
 import json
 import logging
@@ -19,9 +18,9 @@ from collections.abc import Awaitable, Callable
 from contextvars import ContextVar
 from typing import Any
 
-from backend.copilot.model import ChatSession
-from backend.copilot.tools import TOOL_REGISTRY
-from backend.copilot.tools.base import BaseTool
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.tools import TOOL_REGISTRY
+from backend.api.features.chat.tools.base import BaseTool
 
 logger = logging.getLogger(__name__)
 
@@ -42,17 +41,9 @@ _current_session: ContextVar[ChatSession | None] = ContextVar(
 # Stash for MCP tool outputs before the SDK potentially truncates them.
 # Keyed by tool_name → full output string. Consumed (popped) by the
 # response adapter when it builds StreamToolOutputAvailable.
-_pending_tool_outputs: ContextVar[dict[str, list[str]]] = ContextVar(
+_pending_tool_outputs: ContextVar[dict[str, str]] = ContextVar(
     "pending_tool_outputs", default=None  # type: ignore[arg-type]
 )
-# Event signaled whenever stash_pending_tool_output() adds a new entry.
-# Used by the streaming loop to wait for PostToolUse hooks to complete
-# instead of sleeping an arbitrary duration.  The SDK fires hooks via
-# start_soon (fire-and-forget) so the next message can arrive before
-# the hook stashes its output — this event bridges that gap.
-_stash_event: ContextVar[asyncio.Event | None] = ContextVar(
-    "_stash_event", default=None
-)
 
 # Callback type for delegating long-running tools to the non-SDK infrastructure.
 # Args: (tool_name, arguments, session) → MCP-formatted response dict.
@@ -85,7 +76,6 @@ def set_execution_context(
     _current_user_id.set(user_id)
     _current_session.set(session)
     _pending_tool_outputs.set({})
-    _stash_event.set(asyncio.Event())
     _long_running_callback.set(long_running_callback)
 
 
@@ -98,89 +88,19 @@ def get_execution_context() -> tuple[str | None, ChatSession | None]:
 
 
 def pop_pending_tool_output(tool_name: str) -> str | None:
-    """Pop and return the oldest stashed output for *tool_name*.
+    """Pop and return the stashed full output for *tool_name*.
 
     The SDK CLI may truncate large tool results (writing them to disk and
     replacing the content with a file reference). This stash keeps the
     original MCP output so the response adapter can forward it to the
     frontend for proper widget rendering.
 
-    Uses a FIFO queue per tool name so duplicate calls to the same tool
-    in one turn each get their own output.
-
     Returns ``None`` if nothing was stashed for *tool_name*.
     """
     pending = _pending_tool_outputs.get(None)
     if pending is None:
         return None
-    queue = pending.get(tool_name)
-    if not queue:
-        pending.pop(tool_name, None)
-        return None
-    value = queue.pop(0)
-    if not queue:
-        del pending[tool_name]
-    return value
-
-
-def stash_pending_tool_output(tool_name: str, output: Any) -> None:
-    """Stash tool output for later retrieval by the response adapter.
-
-    Used by the PostToolUse hook to capture SDK built-in tool outputs
-    (WebSearch, Read, etc.) that aren't available through the MCP stash
-    mechanism in ``_execute_tool_sync``.
-
-    Appends to a FIFO queue per tool name so multiple calls to the same
-    tool in one turn are all preserved.
-    """
-    pending = _pending_tool_outputs.get(None)
-    if pending is None:
-        return
-    if isinstance(output, str):
-        text = output
-    else:
-        try:
-            text = json.dumps(output)
-        except (TypeError, ValueError):
-            text = str(output)
-    pending.setdefault(tool_name, []).append(text)
-    # Signal any waiters that new output is available.
-    event = _stash_event.get(None)
-    if event is not None:
-        event.set()
-
-
-async def wait_for_stash(timeout: float = 0.5) -> bool:
-    """Wait for a PostToolUse hook to stash tool output.
-
-    The SDK fires PostToolUse hooks asynchronously via ``start_soon()`` —
-    the next message (AssistantMessage/ResultMessage) can arrive before the
-    hook completes and stashes its output.  This function bridges that gap
-    by waiting on the ``_stash_event``, which is signaled by
-    :func:`stash_pending_tool_output`.
-
-    After the event fires, callers should ``await asyncio.sleep(0)`` to
-    give any remaining concurrent hooks a chance to complete.
-
-    Returns ``True`` if a stash signal was received, ``False`` on timeout.
-    The timeout is a safety net — normally the stash happens within
-    microseconds of yielding to the event loop.
-    """
-    event = _stash_event.get(None)
-    if event is None:
-        return False
-    # Fast path: hook already completed before we got here.
-    if event.is_set():
-        event.clear()
-        return True
-    # Slow path: wait for the hook to signal.
-    try:
-        async with asyncio.timeout(timeout):
-            await event.wait()
-        event.clear()
-        return True
-    except TimeoutError:
-        return False
+    return pending.pop(tool_name, None)
 
 
 async def _execute_tool_sync(
@@ -205,63 +125,14 @@ async def _execute_tool_sync(
     # Stash the full output before the SDK potentially truncates it.
     pending = _pending_tool_outputs.get(None)
     if pending is not None:
-        pending.setdefault(base_tool.name, []).append(text)
-
-    content_blocks: list[dict[str, str]] = [{"type": "text", "text": text}]
-
-    # If the tool result contains inline image data, add an MCP image block
-    # so Claude can "see" the image (e.g. read_workspace_file on a small PNG).
-    image_block = _extract_image_block(text)
-    if image_block:
-        content_blocks.append(image_block)
+        pending[base_tool.name] = text
 
     return {
-        "content": content_blocks,
+        "content": [{"type": "text", "text": text}],
         "isError": not result.success,
     }
 
 
-# MIME types that Claude can process as image content blocks.
-_SUPPORTED_IMAGE_TYPES = frozenset(
-    {"image/png", "image/jpeg", "image/gif", "image/webp"}
-)
-
-
-def _extract_image_block(text: str) -> dict[str, str] | None:
-    """Extract an MCP image content block from a tool result JSON string.
-
-    Detects workspace file responses with ``content_base64`` and an image
-    MIME type, returning an MCP-format image block that allows Claude to
-    "see" the image.  Returns ``None`` if the result is not an inline image.
-    """
-    try:
-        data = json.loads(text)
-    except (json.JSONDecodeError, TypeError):
-        return None
-
-    if not isinstance(data, dict):
-        return None
-
-    mime_type = data.get("mime_type", "")
-    base64_content = data.get("content_base64", "")
-
-    # Only inline small images — large ones would exceed Claude's limits.
-    # 32 KB raw ≈ ~43 KB base64.
-    _MAX_IMAGE_BASE64_BYTES = 43_000
-    if (
-        mime_type in _SUPPORTED_IMAGE_TYPES
-        and base64_content
-        and len(base64_content) <= _MAX_IMAGE_BASE64_BYTES
-    ):
-        return {
-            "type": "image",
-            "data": base64_content,
-            "mimeType": mime_type,
-        }
-
-    return None
-
-
 def _mcp_error(message: str) -> dict[str, Any]:
     return {
         "content": [
@@ -440,29 +311,14 @@ def create_copilot_mcp_server():
 # which provides kernel-level network isolation via unshare --net.
 # Task allows spawning sub-agents (rate-limited by security hooks).
 # WebSearch uses Brave Search via Anthropic's API — safe, no SSRF risk.
-# TodoWrite manages the task checklist shown in the UI — no security concern.
-_SDK_BUILTIN_TOOLS = [
-    "Read",
-    "Write",
-    "Edit",
-    "Glob",
-    "Grep",
-    "Task",
-    "WebSearch",
-    "TodoWrite",
-]
+_SDK_BUILTIN_TOOLS = ["Read", "Write", "Edit", "Glob", "Grep", "Task", "WebSearch"]
 
 # SDK built-in tools that must be explicitly blocked.
 # Bash: dangerous — agent uses mcp__copilot__bash_exec with kernel-level
 #   network isolation (unshare --net) instead.
 # WebFetch: SSRF risk — can reach internal network (localhost, 10.x, etc.).
 #   Agent uses the SSRF-protected mcp__copilot__web_fetch tool instead.
-# AskUserQuestion: interactive CLI tool — no terminal in copilot context.
-SDK_DISALLOWED_TOOLS = [
-    "Bash",
-    "WebFetch",
-    "AskUserQuestion",
-]
+SDK_DISALLOWED_TOOLS = ["Bash", "WebFetch"]
 
 # Tools that are blocked entirely in security hooks (defence-in-depth).
 # Includes SDK_DISALLOWED_TOOLS plus common aliases/synonyms.

autogpt_platform/backend/backend/api/features/chat/sdk/transcript.py

@@ -14,8 +14,6 @@ import json
 import logging
 import os
 import re
-import time
-from dataclasses import dataclass
 
 logger = logging.getLogger(__name__)
 
@@ -33,16 +31,6 @@ STRIPPABLE_TYPES = frozenset(
     {"progress", "file-history-snapshot", "queue-operation", "summary", "pr-link"}
 )
 
-
-@dataclass
-class TranscriptDownload:
-    """Result of downloading a transcript with its metadata."""
-
-    content: str
-    message_count: int = 0  # session.messages length when uploaded
-    uploaded_at: float = 0.0  # epoch timestamp of upload
-
-
 # Workspace storage constants — deterministic path from session_id.
 TRANSCRIPT_STORAGE_PREFIX = "chat-transcripts"
 
@@ -131,21 +119,22 @@ def read_transcript_file(transcript_path: str) -> str | None:
             content = f.read()
 
         if not content.strip():
-            logger.debug("[Transcript] File is empty: %s", transcript_path)
+            logger.debug(f"[Transcript] Empty file: {transcript_path}")
             return None
 
         lines = content.strip().split("\n")
-
-        # Validate that the transcript has real conversation content
-        # (not just metadata like queue-operation entries).
-        if not validate_transcript(content):
+        if len(lines) < 3:
+            # Raw files with ≤2 lines are metadata-only
+            # (queue-operation + file-history-snapshot, no conversation).
             logger.debug(
-                "[Transcript] No conversation content (%d lines) in %s",
-                len(lines),
-                transcript_path,
+                f"[Transcript] Too few lines ({len(lines)}): {transcript_path}"
             )
             return None
 
+        # Quick structural validation — parse first and last lines.
+        json.loads(lines[0])
+        json.loads(lines[-1])
+
         logger.info(
             f"[Transcript] Read {len(lines)} lines, "
             f"{len(content)} bytes from {transcript_path}"
@@ -171,41 +160,6 @@ def _sanitize_id(raw_id: str, max_len: int = 36) -> str:
 _SAFE_CWD_PREFIX = os.path.realpath("/tmp/copilot-")
 
 
-def _encode_cwd_for_cli(cwd: str) -> str:
-    """Encode a working directory path the same way the Claude CLI does.
-
-    The CLI replaces all non-alphanumeric characters with ``-``.
-    """
-    return re.sub(r"[^a-zA-Z0-9]", "-", os.path.realpath(cwd))
-
-
-def cleanup_cli_project_dir(sdk_cwd: str) -> None:
-    """Remove the CLI's project directory for a specific working directory.
-
-    The CLI stores session data under ``~/.claude/projects/<encoded_cwd>/``.
-    Each SDK turn uses a unique ``sdk_cwd``, so the project directory is
-    safe to remove entirely after the transcript has been uploaded.
-    """
-    import shutil
-
-    cwd_encoded = _encode_cwd_for_cli(sdk_cwd)
-    config_dir = os.environ.get("CLAUDE_CONFIG_DIR") or os.path.expanduser("~/.claude")
-    projects_base = os.path.realpath(os.path.join(config_dir, "projects"))
-    project_dir = os.path.realpath(os.path.join(projects_base, cwd_encoded))
-
-    if not project_dir.startswith(projects_base + os.sep):
-        logger.warning(
-            f"[Transcript] Cleanup path escaped projects base: {project_dir}"
-        )
-        return
-
-    if os.path.isdir(project_dir):
-        shutil.rmtree(project_dir, ignore_errors=True)
-        logger.debug(f"[Transcript] Cleaned up CLI project dir: {project_dir}")
-    else:
-        logger.debug(f"[Transcript] Project dir not found: {project_dir}")
-
-
 def write_transcript_to_tempfile(
     transcript_content: str,
     session_id: str,
@@ -294,15 +248,6 @@ def _storage_path_parts(user_id: str, session_id: str) -> tuple[str, str, str]:
     )
 
 
-def _meta_storage_path_parts(user_id: str, session_id: str) -> tuple[str, str, str]:
-    """Return (workspace_id, file_id, filename) for a session's transcript metadata."""
-    return (
-        TRANSCRIPT_STORAGE_PREFIX,
-        _sanitize_id(user_id),
-        f"{_sanitize_id(session_id)}.meta.json",
-    )
-
-
 def _build_storage_path(user_id: str, session_id: str, backend: object) -> str:
     """Build the full storage path string that ``retrieve()`` expects.
 
@@ -323,30 +268,21 @@ def _build_storage_path(user_id: str, session_id: str, backend: object) -> str:
         return f"local://{wid}/{fid}/{fname}"
 
 
-async def upload_transcript(
-    user_id: str,
-    session_id: str,
-    content: str,
-    message_count: int = 0,
-) -> None:
+async def upload_transcript(user_id: str, session_id: str, content: str) -> None:
     """Strip progress entries and upload transcript to bucket storage.
 
     Safety: only overwrites when the new (stripped) transcript is larger than
     what is already stored.  Since JSONL is append-only, the latest transcript
     is always the longest.  This prevents a slow/stale background task from
     clobbering a newer upload from a concurrent turn.
-
-    Args:
-        message_count: ``len(session.messages)`` at upload time — used by
-            the next turn to detect staleness and compress only the gap.
     """
     from backend.util.workspace_storage import get_workspace_storage
 
     stripped = strip_progress_entries(content)
     if not validate_transcript(stripped):
         logger.warning(
-            f"[Transcript] Skipping upload — stripped content not valid "
-            f"for session {session_id}"
+            f"[Transcript] Skipping upload — stripped content is not a valid "
+            f"transcript for session {session_id}"
         )
         return
 
@@ -361,8 +297,9 @@ async def upload_transcript(
         existing = await storage.retrieve(path)
         if len(existing) >= new_size:
             logger.info(
-                f"[Transcript] Skipping upload — existing ({len(existing)}B) "
-                f">= new ({new_size}B) for session {session_id}"
+                f"[Transcript] Skipping upload — existing transcript "
+                f"({len(existing)}B) >= new ({new_size}B) for session "
+                f"{session_id}"
             )
             return
     except (FileNotFoundError, Exception):
@@ -374,38 +311,16 @@ async def upload_transcript(
         filename=fname,
         content=encoded,
     )
-
-    # Store metadata alongside the transcript so the next turn can detect
-    # staleness and only compress the gap instead of the full history.
-    # Wrapped in try/except so a metadata write failure doesn't orphan
-    # the already-uploaded transcript — the next turn will just fall back
-    # to full gap fill (msg_count=0).
-    try:
-        meta = {"message_count": message_count, "uploaded_at": time.time()}
-        mwid, mfid, mfname = _meta_storage_path_parts(user_id, session_id)
-        await storage.store(
-            workspace_id=mwid,
-            file_id=mfid,
-            filename=mfname,
-            content=json.dumps(meta).encode("utf-8"),
-        )
-    except Exception as e:
-        logger.warning(f"[Transcript] Failed to write metadata for {session_id}: {e}")
-
     logger.info(
-        f"[Transcript] Uploaded {new_size}B "
-        f"(stripped from {len(content)}B, msg_count={message_count}) "
-        f"for session {session_id}"
+        f"[Transcript] Uploaded {new_size} bytes "
+        f"(stripped from {len(content)}) for session {session_id}"
     )
 
 
-async def download_transcript(
-    user_id: str, session_id: str
-) -> TranscriptDownload | None:
-    """Download transcript and metadata from bucket storage.
+async def download_transcript(user_id: str, session_id: str) -> str | None:
+    """Download transcript from bucket storage.
 
-    Returns a ``TranscriptDownload`` with the JSONL content and the
-    ``message_count`` watermark from the upload, or ``None`` if not found.
+    Returns the JSONL content string, or ``None`` if not found.
     """
     from backend.util.workspace_storage import get_workspace_storage
 
@@ -415,6 +330,10 @@ async def download_transcript(
     try:
         data = await storage.retrieve(path)
         content = data.decode("utf-8")
+        logger.info(
+            f"[Transcript] Downloaded {len(content)} bytes for session {session_id}"
+        )
+        return content
     except FileNotFoundError:
         logger.debug(f"[Transcript] No transcript in storage for {session_id}")
         return None
@@ -422,36 +341,6 @@ async def download_transcript(
         logger.warning(f"[Transcript] Failed to download transcript: {e}")
         return None
 
-    # Try to load metadata (best-effort — old transcripts won't have it)
-    message_count = 0
-    uploaded_at = 0.0
-    try:
-        from backend.util.workspace_storage import GCSWorkspaceStorage
-
-        mwid, mfid, mfname = _meta_storage_path_parts(user_id, session_id)
-        if isinstance(storage, GCSWorkspaceStorage):
-            blob = f"workspaces/{mwid}/{mfid}/{mfname}"
-            meta_path = f"gcs://{storage.bucket_name}/{blob}"
-        else:
-            meta_path = f"local://{mwid}/{mfid}/{mfname}"
-
-        meta_data = await storage.retrieve(meta_path)
-        meta = json.loads(meta_data.decode("utf-8"))
-        message_count = meta.get("message_count", 0)
-        uploaded_at = meta.get("uploaded_at", 0.0)
-    except (FileNotFoundError, json.JSONDecodeError, Exception):
-        pass  # No metadata — treat as unknown (msg_count=0 → always fill gap)
-
-    logger.info(
-        f"[Transcript] Downloaded {len(content)}B "
-        f"(msg_count={message_count}) for session {session_id}"
-    )
-    return TranscriptDownload(
-        content=content,
-        message_count=message_count,
-        uploaded_at=uploaded_at,
-    )
-
 
 async def delete_transcript(user_id: str, session_id: str) -> None:
     """Delete transcript from bucket storage (e.g. after resume failure)."""

autogpt_platform/backend/backend/api/features/chat/service.py

@@ -27,18 +27,20 @@ from openai.types.chat import (
     ChatCompletionToolParam,
 )
 
-from backend.data.db_accessors import chat_db, understanding_db
 from backend.data.redis_client import get_redis_async
-from backend.data.understanding import format_understanding_for_prompt
+from backend.data.understanding import (
+    format_understanding_for_prompt,
+    get_business_understanding,
+)
 from backend.util.exceptions import NotFoundError
 from backend.util.settings import AppEnvironment, Settings
 
+from . import db as chat_db
 from . import stream_registry
 from .config import ChatConfig
 from .model import (
     ChatMessage,
     ChatSession,
-    ChatSessionInfo,
     Usage,
     cache_chat_session,
     get_chat_session,
@@ -118,8 +120,6 @@ Adapt flexibly to the conversation context. Not every interaction requires all s
    - Find reusable components with `find_block`
    - Create custom solutions with `create_agent` if nothing suitable exists
    - Modify existing library agents with `edit_agent`
-   - **When `create_agent` returns `suggested_goal`**: Present the suggestion to the user and ask "Would you like me to proceed with this refined goal?" If they accept, call `create_agent` again with the suggested goal.
-   - **When `create_agent` returns `clarifying_questions`**: After the user answers, call `create_agent` again with the original description AND the answers in the `context` parameter.
 
 5. **Execute**: Run automations immediately, schedule them, or set up webhooks using `run_agent`. Test specific components with `run_block`.
 
@@ -166,11 +166,6 @@ Adapt flexibly to the conversation context. Not every interaction requires all s
 - Use `add_understanding` to capture valuable business context
 - When tool calls fail, try alternative approaches
 
-**Handle Feedback Loops:**
-- When a tool returns a suggested alternative (like a refined goal), present it clearly and ask the user for confirmation before proceeding
-- When clarifying questions are answered, immediately re-call the tool with the accumulated context
-- Don't ask redundant questions if the user has already provided context in the conversation
-
 ## CRITICAL REMINDER
 
 You are NOT a chatbot. You are NOT documentation. You are a partner who helps busy business owners get value quickly by showing proof through working automations. Bias toward action over explanation."""
@@ -268,7 +263,7 @@ async def _build_system_prompt(
     understanding = None
     if user_id:
         try:
-            understanding = await understanding_db().get_business_understanding(user_id)
+            understanding = await get_business_understanding(user_id)
         except Exception as e:
             logger.warning(f"Failed to fetch business understanding: {e}")
             understanding = None
@@ -344,7 +339,7 @@ async def _generate_session_title(
 async def assign_user_to_session(
     session_id: str,
     user_id: str,
-) -> ChatSessionInfo:
+) -> ChatSession:
     """
     Assign a user to a chat session.
     """
@@ -352,8 +347,7 @@ async def assign_user_to_session(
     if not session:
         raise NotFoundError(f"Session {session_id} not found")
     session.user_id = user_id
-    session, _ = await upsert_chat_session(session)
-    return session
+    return await upsert_chat_session(session)
 
 
 async def stream_chat_completion(
@@ -434,16 +428,12 @@ async def stream_chat_completion(
             f"Session {session_id} not found. Please create a new session first."
         )
 
-    # Append the new message to the session if it's not already there
-    new_message_role = "user" if is_user_message else "assistant"
-    if message and (
-        len(session.messages) == 0
-        or not (
-            session.messages[-1].role == new_message_role
-            and session.messages[-1].content == message
+    if message:
+        session.messages.append(
+            ChatMessage(
+                role="user" if is_user_message else "assistant", content=message
+            )
         )
-    ):
-        session.messages.append(ChatMessage(role=new_message_role, content=message))
         logger.info(
             f"Appended message (role={'user' if is_user_message else 'assistant'}), "
             f"new message_count={len(session.messages)}"
@@ -464,7 +454,7 @@ async def stream_chat_completion(
             )
 
     upsert_start = time.monotonic()
-    session, _ = await upsert_chat_session(session)
+    session = await upsert_chat_session(session)
     upsert_time = (time.monotonic() - upsert_start) * 1000
     logger.info(
         f"[TIMING] upsert_chat_session took {upsert_time:.1f}ms",
@@ -690,7 +680,7 @@ async def stream_chat_completion(
                                 f"tool_responses={len(tool_response_messages)}"
                             )
                         if messages_to_save_early or has_appended_streaming_message:
-                            _ = await upsert_chat_session(session)
+                            await upsert_chat_session(session)
                             has_saved_assistant_message = True
 
                     has_yielded_end = True
@@ -729,7 +719,7 @@ async def stream_chat_completion(
             if tool_response_messages:
                 session.messages.extend(tool_response_messages)
             try:
-                _ = await upsert_chat_session(session)
+                await upsert_chat_session(session)
             except Exception as e:
                 logger.warning(
                     f"Failed to save interrupted session {session.session_id}: {e}"
@@ -770,7 +760,7 @@ async def stream_chat_completion(
                 if messages_to_save:
                     session.messages.extend(messages_to_save)
                 if messages_to_save or has_appended_streaming_message:
-                    _ = await upsert_chat_session(session)
+                    await upsert_chat_session(session)
 
             if not has_yielded_error:
                 error_message = str(e)
@@ -854,7 +844,7 @@ async def stream_chat_completion(
             not has_long_running_tool_call
             and (messages_to_save or has_appended_streaming_message)
         ):
-            _ = await upsert_chat_session(session)
+            await upsert_chat_session(session)
     else:
         logger.info(
             "Assistant message already saved when StreamFinish was received, "
@@ -1233,10 +1223,23 @@ async def _stream_chat_chunks(
                     },
                 )
 
-                # Execute all accumulated tool calls in parallel
-                # Events are yielded as they arrive from each concurrent tool
-                async for event in _execute_tool_calls_parallel(tool_calls, session):
-                    yield event
+                # Yield all accumulated tool calls after the stream is complete
+                # This ensures all tool call arguments have been fully received
+                for idx, tool_call in enumerate(tool_calls):
+                    try:
+                        async for tc in _yield_tool_call(tool_calls, idx, session):
+                            yield tc
+                    except (orjson.JSONDecodeError, KeyError, TypeError) as e:
+                        logger.error(
+                            f"Failed to parse tool call {idx}: {e}",
+                            exc_info=True,
+                            extra={"tool_call": tool_call},
+                        )
+                        yield StreamError(
+                            errorText=f"Invalid tool call arguments for tool {tool_call.get('function', {}).get('name', 'unknown')}: {e}",
+                        )
+                        # Re-raise to trigger retry logic in the parent function
+                        raise
 
                 total_time = (time_module.perf_counter() - stream_chunks_start) * 1000
                 logger.info(
@@ -1314,91 +1317,10 @@ async def _stream_chat_chunks(
             return
 
 
-async def _with_optional_lock(
-    lock: asyncio.Lock | None,
-    coro_fn: Any,
-) -> Any:
-    """Run *coro_fn()* under *lock* when provided, otherwise run directly."""
-    if lock:
-        async with lock:
-            return await coro_fn()
-    return await coro_fn()
-
-
-async def _execute_tool_calls_parallel(
-    tool_calls: list[dict[str, Any]],
-    session: ChatSession,
-) -> AsyncGenerator[StreamBaseResponse, None]:
-    """Execute all tool calls concurrently, yielding stream events as they arrive.
-
-    Each tool runs as an ``asyncio.Task``, pushing events into a shared queue.
-    A ``session_lock`` serialises session-state mutations (long-running tool
-    bookkeeping, ``run_agent`` counters).
-    """
-    queue: asyncio.Queue[StreamBaseResponse | None] = asyncio.Queue()
-    session_lock = asyncio.Lock()
-    n_tools = len(tool_calls)
-    retryable_errors: list[Exception] = []
-
-    async def _run_tool(idx: int) -> None:
-        tool_name = tool_calls[idx].get("function", {}).get("name", "unknown")
-        tool_call_id = tool_calls[idx].get("id", f"unknown_{idx}")
-        try:
-            async for event in _yield_tool_call(tool_calls, idx, session, session_lock):
-                await queue.put(event)
-        except (orjson.JSONDecodeError, KeyError, TypeError) as e:
-            logger.error(
-                f"Failed to parse tool call {idx} ({tool_name}): {e}",
-                exc_info=True,
-            )
-            retryable_errors.append(e)
-        except Exception as e:
-            # Infrastructure / setup errors — emit an error output so the
-            # client always sees a terminal event and doesn't hang.
-            logger.error(f"Tool call {idx} ({tool_name}) failed: {e}", exc_info=True)
-            await queue.put(
-                StreamToolOutputAvailable(
-                    toolCallId=tool_call_id,
-                    toolName=tool_name,
-                    output=ErrorResponse(
-                        message=f"Tool execution failed: {e!s}",
-                        error=type(e).__name__,
-                        session_id=session.session_id,
-                    ).model_dump_json(),
-                    success=False,
-                )
-            )
-        finally:
-            await queue.put(None)  # sentinel
-
-    tasks = [asyncio.create_task(_run_tool(idx)) for idx in range(n_tools)]
-    try:
-        finished = 0
-        while finished < n_tools:
-            event = await queue.get()
-            if event is None:
-                finished += 1
-            else:
-                yield event
-        if retryable_errors:
-            if len(retryable_errors) > 1:
-                logger.warning(
-                    f"{len(retryable_errors)} tool calls had retryable errors; "
-                    f"re-raising first to trigger retry"
-                )
-            raise retryable_errors[0]
-    finally:
-        for t in tasks:
-            if not t.done():
-                t.cancel()
-        await asyncio.gather(*tasks, return_exceptions=True)
-
-
 async def _yield_tool_call(
     tool_calls: list[dict[str, Any]],
     yield_idx: int,
     session: ChatSession,
-    session_lock: asyncio.Lock | None = None,
 ) -> AsyncGenerator[StreamBaseResponse, None]:
     """
     Yield a tool call and its execution result.
@@ -1496,7 +1418,8 @@ async def _yield_tool_call(
                 "check back in a few minutes."
             )
 
-        # Track appended message for rollback on failure
+        # Track appended messages for rollback on failure
+        assistant_message: ChatMessage | None = None
         pending_message: ChatMessage | None = None
 
         # Wrap session save and task creation in try-except to release lock on failure
@@ -1511,24 +1434,22 @@ async def _yield_tool_call(
                 operation_id=operation_id,
             )
 
-            # Attach tool_call and save pending result — lock serialises
-            # concurrent session mutations during parallel execution.
-            async def _save_pending() -> None:
-                nonlocal pending_message
-                session.add_tool_call_to_current_turn(tool_calls[yield_idx])
-                pending_message = ChatMessage(
-                    role="tool",
-                    content=OperationPendingResponse(
-                        message=pending_msg,
-                        operation_id=operation_id,
-                        tool_name=tool_name,
-                    ).model_dump_json(),
-                    tool_call_id=tool_call_id,
-                )
-                session.messages.append(pending_message)
-                _ = await upsert_chat_session(session)
+            # Attach the tool_call to the current turn's assistant message
+            # (or create one if this is a tool-only response with no text).
+            session.add_tool_call_to_current_turn(tool_calls[yield_idx])
 
-            await _with_optional_lock(session_lock, _save_pending)
+            # Then save pending tool result
+            pending_message = ChatMessage(
+                role="tool",
+                content=OperationPendingResponse(
+                    message=pending_msg,
+                    operation_id=operation_id,
+                    tool_name=tool_name,
+                ).model_dump_json(),
+                tool_call_id=tool_call_id,
+            )
+            session.messages.append(pending_message)
+            await upsert_chat_session(session)
             logger.info(
                 f"Saved pending operation {operation_id} (task_id={task_id}) "
                 f"for tool {tool_name} in session {session.session_id}"
@@ -1552,25 +1473,27 @@ async def _yield_tool_call(
             # Associate the asyncio task with the stream registry task
             await stream_registry.set_task_asyncio_task(task_id, bg_task)
         except Exception as e:
-            # Roll back appended messages — use identity-based removal so
-            # it works even when other parallel tools have appended after us.
-            async def _rollback() -> None:
-                if pending_message and pending_message in session.messages:
-                    session.messages.remove(pending_message)
-
-            await _with_optional_lock(session_lock, _rollback)
+            # Roll back appended messages to prevent data corruption on subsequent saves
+            if (
+                pending_message
+                and session.messages
+                and session.messages[-1] == pending_message
+            ):
+                session.messages.pop()
+            if (
+                assistant_message
+                and session.messages
+                and session.messages[-1] == assistant_message
+            ):
+                session.messages.pop()
 
             # Release the Redis lock since the background task won't be spawned
             await _mark_operation_completed(tool_call_id)
             # Mark stream registry task as failed if it was created
             try:
-                await stream_registry.mark_task_completed(
-                    task_id,
-                    status="failed",
-                    error_message=f"Failed to setup tool {tool_name}: {e}",
-                )
-            except Exception as mark_err:
-                logger.warning(f"Failed to mark task {task_id} as failed: {mark_err}")
+                await stream_registry.mark_task_completed(task_id, status="failed")
+            except Exception:
+                pass
             logger.error(
                 f"Failed to setup long-running tool {tool_name}: {e}", exc_info=True
             )
@@ -1736,11 +1659,7 @@ async def _execute_long_running_tool_with_streaming(
         session = await get_chat_session(session_id, user_id)
         if not session:
             logger.error(f"Session {session_id} not found for background tool")
-            await stream_registry.mark_task_completed(
-                task_id,
-                status="failed",
-                error_message=f"Session {session_id} not found",
-            )
+            await stream_registry.mark_task_completed(task_id, status="failed")
             return
 
         # Pass operation_id and task_id to the tool for async processing
@@ -1851,7 +1770,7 @@ async def _update_pending_operation(
     This is called by background tasks when long-running operations complete.
     """
     # Update the message in database
-    updated = await chat_db().update_tool_message_content(
+    updated = await chat_db.update_tool_message_content(
         session_id=session_id,
         tool_call_id=tool_call_id,
         new_content=result,
@@ -2020,7 +1939,7 @@ async def _generate_llm_continuation(
             fresh_session.messages.append(assistant_message)
 
             # Save to database (not cache) to persist the response
-            _ = await upsert_chat_session(fresh_session)
+            await upsert_chat_session(fresh_session)
 
             # Invalidate cache so next poll/refresh gets fresh data
             await invalidate_session_cache(session_id)
@@ -2226,7 +2145,7 @@ async def _generate_llm_continuation_with_streaming(
             fresh_session.messages.append(assistant_message)
 
             # Save to database (not cache) to persist the response
-            _ = await upsert_chat_session(fresh_session)
+            await upsert_chat_session(fresh_session)
 
             # Invalidate cache so next poll/refresh gets fresh data
             await invalidate_session_cache(session_id)

autogpt_platform/backend/backend/api/features/chat/service_test.py

@@ -58,7 +58,7 @@ async def test_stream_chat_completion_with_tool_calls(setup_test_user, test_user
         return pytest.skip("OPEN_ROUTER_API_KEY is not set, skipping test")
 
     session = await create_chat_session(test_user_id)
-    session, _ = await upsert_chat_session(session)
+    session = await upsert_chat_session(session)
 
     has_errors = False
     has_ended = False
@@ -104,7 +104,7 @@ async def test_sdk_resume_multi_turn(setup_test_user, test_user_id):
         return pytest.skip("CLAUDE_AGENT_USE_RESUME is not enabled, skipping test")
 
     session = await create_chat_session(test_user_id)
-    session, _ = await upsert_chat_session(session)
+    session = await upsert_chat_session(session)
 
     # --- Turn 1: send a message with a unique keyword ---
     keyword = "ZEPHYR42"
@@ -132,24 +132,18 @@ async def test_sdk_resume_multi_turn(setup_test_user, test_user_id):
     assert not turn1_errors, f"Turn 1 errors: {turn1_errors}"
     assert turn1_text, "Turn 1 produced no text"
 
-    # Wait for background upload task to complete (retry up to 5s).
-    # The CLI may not produce a usable transcript for very short
-    # conversations (only metadata entries) — this is environment-dependent
-    # (CLI version, platform).  When that happens, multi-turn still works
-    # via conversation compression (non-resume path), but we can't test
-    # the --resume round-trip.
+    # Wait for background upload task to complete (retry up to 5s)
     transcript = None
     for _ in range(10):
         await asyncio.sleep(0.5)
         transcript = await download_transcript(test_user_id, session.session_id)
         if transcript:
             break
-    if not transcript:
-        return pytest.skip(
-            "CLI did not produce a usable transcript — "
-            "cannot test --resume round-trip in this environment"
-        )
-    logger.info(f"Turn 1 transcript uploaded: {len(transcript.content)} bytes")
+    assert transcript, (
+        "Transcript was not uploaded to bucket after turn 1 — "
+        "Stop hook may not have fired or transcript was too small"
+    )
+    logger.info(f"Turn 1 transcript uploaded: {len(transcript)} bytes")
 
     # Reload session for turn 2
     session = await get_chat_session(session.session_id, test_user_id)

autogpt_platform/backend/backend/api/features/chat/stream_registry.py

@@ -227,14 +227,7 @@ async def publish_chunk(
         # Only log timing for significant chunks or slow operations
         if (
             chunk_type
-            in (
-                "StreamStart",
-                "StreamFinish",
-                "StreamTextStart",
-                "StreamTextEnd",
-                "StreamToolInputAvailable",
-                "StreamToolOutputAvailable",
-            )
+            in ("StreamStart", "StreamFinish", "StreamTextStart", "StreamTextEnd")
             or total_time > 50
         ):
             logger.info(
@@ -644,8 +637,6 @@ async def _stream_listener(
 async def mark_task_completed(
     task_id: str,
     status: Literal["completed", "failed"] = "completed",
-    *,
-    error_message: str | None = None,
 ) -> bool:
     """Mark a task as completed and publish finish event.
 
@@ -656,10 +647,6 @@ async def mark_task_completed(
     Args:
         task_id: Task ID to mark as completed
         status: Final status ("completed" or "failed")
-        error_message: If provided and status="failed", publish a StreamError
-            before StreamFinish so connected clients see why the task ended.
-            If not provided, no StreamError is published (caller should publish
-            manually if needed to avoid duplicates).
 
     Returns:
         True if task was newly marked completed, False if already completed/failed
@@ -675,17 +662,6 @@ async def mark_task_completed(
         logger.debug(f"Task {task_id} already completed/failed, skipping")
         return False
 
-    # Publish error event before finish so connected clients know WHY the
-    # task ended. Only publish if caller provided an explicit error message
-    # to avoid duplicates with code paths that manually publish StreamError.
-    # This is best-effort — if it fails, the StreamFinish still ensures
-    # listeners clean up.
-    if status == "failed" and error_message:
-        try:
-            await publish_chunk(task_id, StreamError(errorText=error_message))
-        except Exception as e:
-            logger.warning(f"Failed to publish error event for task {task_id}: {e}")
-
     # THEN publish finish event (best-effort - listeners can detect via status polling)
     try:
         await publish_chunk(task_id, StreamFinish())
@@ -838,6 +814,24 @@ async def get_active_task_for_session(
                 if task_user_id and user_id != task_user_id:
                     continue
 
+                # Auto-expire stale tasks that exceeded stream_timeout
+                created_at_str = meta.get("created_at", "")
+                if created_at_str:
+                    try:
+                        created_at = datetime.fromisoformat(created_at_str)
+                        age_seconds = (
+                            datetime.now(timezone.utc) - created_at
+                        ).total_seconds()
+                        if age_seconds > config.stream_timeout:
+                            logger.warning(
+                                f"[TASK_LOOKUP] Auto-expiring stale task {task_id[:8]}... "
+                                f"(age={age_seconds:.0f}s > timeout={config.stream_timeout}s)"
+                            )
+                            await mark_task_completed(task_id, "failed")
+                            continue
+                    except (ValueError, TypeError):
+                        pass
+
                 logger.info(
                     f"[TASK_LOOKUP] Found running task {task_id[:8]}... for session {session_id[:8]}..."
                 )

autogpt_platform/backend/backend/api/features/chat/tools/__init__.py

@@ -3,8 +3,8 @@ from typing import TYPE_CHECKING, Any
 
 from openai.types.chat import ChatCompletionToolParam
 
-from backend.copilot.model import ChatSession
-from backend.copilot.tracking import track_tool_called
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.tracking import track_tool_called
 
 from .add_understanding import AddUnderstandingTool
 from .agent_output import AgentOutputTool
@@ -31,7 +31,7 @@ from .workspace_files import (
 )
 
 if TYPE_CHECKING:
-    from backend.copilot.response_model import StreamToolOutputAvailable
+    from backend.api.features.chat.response_model import StreamToolOutputAvailable
 
 logger = logging.getLogger(__name__)
 

autogpt_platform/backend/backend/api/features/chat/tools/_test_data.py

@@ -6,11 +6,11 @@ import pytest
 from prisma.types import ProfileCreateInput
 from pydantic import SecretStr
 
+from backend.api.features.chat.model import ChatSession
 from backend.api.features.store import db as store_db
 from backend.blocks.firecrawl.scrape import FirecrawlScrapeBlock
 from backend.blocks.io import AgentInputBlock, AgentOutputBlock
 from backend.blocks.llm import AITextGeneratorBlock
-from backend.copilot.model import ChatSession
 from backend.data.db import prisma
 from backend.data.graph import Graph, Link, Node, create_graph
 from backend.data.model import APIKeyCredentials

autogpt_platform/backend/backend/api/features/chat/tools/add_understanding.py

@@ -3,9 +3,11 @@
 import logging
 from typing import Any
 
-from backend.copilot.model import ChatSession
-from backend.data.db_accessors import understanding_db
-from backend.data.understanding import BusinessUnderstandingInput
+from backend.api.features.chat.model import ChatSession
+from backend.data.understanding import (
+    BusinessUnderstandingInput,
+    upsert_business_understanding,
+)
 
 from .base import BaseTool
 from .models import ErrorResponse, ToolResponseBase, UnderstandingUpdatedResponse
@@ -97,9 +99,7 @@ and automations for the user's specific needs."""
         ]
 
         # Upsert with merge
-        understanding = await understanding_db().upsert_business_understanding(
-            user_id, input_data
-        )
+        understanding = await upsert_business_understanding(user_id, input_data)
 
         # Build current understanding summary (filter out empty values)
         current_understanding = {

autogpt_platform/backend/backend/api/features/chat/tools/agent_generator/core.py

@@ -5,8 +5,9 @@ import re
 import uuid
 from typing import Any, NotRequired, TypedDict
 
-from backend.data.db_accessors import graph_db, library_db, store_db
-from backend.data.graph import Graph, Link, Node
+from backend.api.features.library import db as library_db
+from backend.api.features.store import db as store_db
+from backend.data.graph import Graph, Link, Node, get_graph, get_store_listed_graphs
 from backend.util.exceptions import DatabaseError, NotFoundError
 
 from .service import (
@@ -144,9 +145,8 @@ async def get_library_agent_by_id(
     Returns:
         LibraryAgentSummary if found, None otherwise
     """
-    db = library_db()
     try:
-        agent = await db.get_library_agent_by_graph_id(user_id, agent_id)
+        agent = await library_db.get_library_agent_by_graph_id(user_id, agent_id)
         if agent:
             logger.debug(f"Found library agent by graph_id: {agent.name}")
             return LibraryAgentSummary(
@@ -163,7 +163,7 @@ async def get_library_agent_by_id(
         logger.debug(f"Could not fetch library agent by graph_id {agent_id}: {e}")
 
     try:
-        agent = await db.get_library_agent(agent_id, user_id)
+        agent = await library_db.get_library_agent(agent_id, user_id)
         if agent:
             logger.debug(f"Found library agent by library_id: {agent.name}")
             return LibraryAgentSummary(
@@ -215,7 +215,7 @@ async def get_library_agents_for_generation(
         List of LibraryAgentSummary with schemas and recent executions for sub-agent composition
     """
     try:
-        response = await library_db().list_library_agents(
+        response = await library_db.list_library_agents(
             user_id=user_id,
             search_term=search_query,
             page=1,
@@ -272,7 +272,7 @@ async def search_marketplace_agents_for_generation(
         List of LibraryAgentSummary with full input/output schemas
     """
     try:
-        response = await store_db().get_store_agents(
+        response = await store_db.get_store_agents(
             search_query=search_query,
             page=1,
             page_size=max_results,
@@ -286,7 +286,7 @@ async def search_marketplace_agents_for_generation(
             return []
 
         graph_ids = [agent.agent_graph_id for agent in agents_with_graphs]
-        graphs = await graph_db().get_store_listed_graphs(graph_ids)
+        graphs = await get_store_listed_graphs(*graph_ids)
 
         results: list[LibraryAgentSummary] = []
         for agent in agents_with_graphs:
@@ -673,10 +673,9 @@ async def save_agent_to_library(
         Tuple of (created Graph, LibraryAgent)
     """
     graph = json_to_graph(agent_json)
-    db = library_db()
     if is_update:
-        return await db.update_graph_in_library(graph, user_id)
-    return await db.create_graph_in_library(graph, user_id)
+        return await library_db.update_graph_in_library(graph, user_id)
+    return await library_db.create_graph_in_library(graph, user_id)
 
 
 def graph_to_json(graph: Graph) -> dict[str, Any]:
@@ -736,14 +735,12 @@ async def get_agent_as_json(
     Returns:
         Agent as JSON dict or None if not found
     """
-    db = graph_db()
-
-    graph = await db.get_graph(agent_id, version=None, user_id=user_id)
+    graph = await get_graph(agent_id, version=None, user_id=user_id)
 
     if not graph and user_id:
         try:
-            library_agent = await library_db().get_library_agent(agent_id, user_id)
-            graph = await db.get_graph(
+            library_agent = await library_db.get_library_agent(agent_id, user_id)
+            graph = await get_graph(
                 library_agent.graph_id, version=None, user_id=user_id
             )
         except NotFoundError:

autogpt_platform/backend/backend/api/features/chat/tools/agent_output.py

@@ -7,9 +7,10 @@ from typing import Any
 
 from pydantic import BaseModel, field_validator
 
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.library import db as library_db
 from backend.api.features.library.model import LibraryAgent
-from backend.copilot.model import ChatSession
-from backend.data.db_accessors import execution_db, library_db
+from backend.data import execution as execution_db
 from backend.data.execution import ExecutionStatus, GraphExecution, GraphExecutionMeta
 
 from .base import BaseTool
@@ -164,12 +165,10 @@ class AgentOutputTool(BaseTool):
         Resolve agent from provided identifiers.
         Returns (library_agent, error_message).
         """
-        lib_db = library_db()
-
         # Priority 1: Exact library agent ID
         if library_agent_id:
             try:
-                agent = await lib_db.get_library_agent(library_agent_id, user_id)
+                agent = await library_db.get_library_agent(library_agent_id, user_id)
                 return agent, None
             except Exception as e:
                 logger.warning(f"Failed to get library agent by ID: {e}")
@@ -183,7 +182,7 @@ class AgentOutputTool(BaseTool):
                 return None, f"Agent '{store_slug}' not found in marketplace"
 
             # Find in user's library by graph_id
-            agent = await lib_db.get_library_agent_by_graph_id(user_id, graph.id)
+            agent = await library_db.get_library_agent_by_graph_id(user_id, graph.id)
             if not agent:
                 return (
                     None,
@@ -195,7 +194,7 @@ class AgentOutputTool(BaseTool):
         # Priority 3: Fuzzy name search in library
         if agent_name:
             try:
-                response = await lib_db.list_library_agents(
+                response = await library_db.list_library_agents(
                     user_id=user_id,
                     search_term=agent_name,
                     page_size=5,
@@ -229,11 +228,9 @@ class AgentOutputTool(BaseTool):
         Fetch execution(s) based on filters.
         Returns (single_execution, available_executions_meta, error_message).
         """
-        exec_db = execution_db()
-
         # If specific execution_id provided, fetch it directly
         if execution_id:
-            execution = await exec_db.get_graph_execution(
+            execution = await execution_db.get_graph_execution(
                 user_id=user_id,
                 execution_id=execution_id,
                 include_node_executions=False,
@@ -243,7 +240,7 @@ class AgentOutputTool(BaseTool):
             return execution, [], None
 
         # Get completed executions with time filters
-        executions = await exec_db.get_graph_executions(
+        executions = await execution_db.get_graph_executions(
             graph_id=graph_id,
             user_id=user_id,
             statuses=[ExecutionStatus.COMPLETED],
@@ -257,7 +254,7 @@ class AgentOutputTool(BaseTool):
 
         # If only one execution, fetch full details
         if len(executions) == 1:
-            full_execution = await exec_db.get_graph_execution(
+            full_execution = await execution_db.get_graph_execution(
                 user_id=user_id,
                 execution_id=executions[0].id,
                 include_node_executions=False,
@@ -265,7 +262,7 @@ class AgentOutputTool(BaseTool):
             return full_execution, [], None
 
         # Multiple executions - return latest with full details, plus list of available
-        full_execution = await exec_db.get_graph_execution(
+        full_execution = await execution_db.get_graph_execution(
             user_id=user_id,
             execution_id=executions[0].id,
             include_node_executions=False,
@@ -383,7 +380,7 @@ class AgentOutputTool(BaseTool):
             and not input_data.store_slug
         ):
             # Fetch execution directly to get graph_id
-            execution = await execution_db().get_graph_execution(
+            execution = await execution_db.get_graph_execution(
                 user_id=user_id,
                 execution_id=input_data.execution_id,
                 include_node_executions=False,
@@ -395,7 +392,7 @@ class AgentOutputTool(BaseTool):
                 )
 
             # Find library agent by graph_id
-            agent = await library_db().get_library_agent_by_graph_id(
+            agent = await library_db.get_library_agent_by_graph_id(
                 user_id, execution.graph_id
             )
             if not agent:

autogpt_platform/backend/backend/api/features/chat/tools/agent_search.py

@@ -4,7 +4,8 @@ import logging
 import re
 from typing import Literal
 
-from backend.data.db_accessors import library_db, store_db
+from backend.api.features.library import db as library_db
+from backend.api.features.store import db as store_db
 from backend.util.exceptions import DatabaseError, NotFoundError
 
 from .models import (
@@ -44,10 +45,8 @@ async def _get_library_agent_by_id(user_id: str, agent_id: str) -> AgentInfo | N
     Returns:
         AgentInfo if found, None otherwise
     """
-    lib_db = library_db()
-
     try:
-        agent = await lib_db.get_library_agent_by_graph_id(user_id, agent_id)
+        agent = await library_db.get_library_agent_by_graph_id(user_id, agent_id)
         if agent:
             logger.debug(f"Found library agent by graph_id: {agent.name}")
             return AgentInfo(
@@ -72,7 +71,7 @@ async def _get_library_agent_by_id(user_id: str, agent_id: str) -> AgentInfo | N
         )
 
     try:
-        agent = await lib_db.get_library_agent(agent_id, user_id)
+        agent = await library_db.get_library_agent(agent_id, user_id)
         if agent:
             logger.debug(f"Found library agent by library_id: {agent.name}")
             return AgentInfo(
@@ -134,7 +133,7 @@ async def search_agents(
     try:
         if source == "marketplace":
             logger.info(f"Searching marketplace for: {query}")
-            results = await store_db().get_store_agents(search_query=query, page_size=5)
+            results = await store_db.get_store_agents(search_query=query, page_size=5)
             for agent in results.agents:
                 agents.append(
                     AgentInfo(
@@ -160,7 +159,7 @@ async def search_agents(
 
             if not agents:
                 logger.info(f"Searching user library for: {query}")
-                results = await library_db().list_library_agents(
+                results = await library_db.list_library_agents(
                     user_id=user_id,  # type: ignore[arg-type]
                     search_term=query,
                     page_size=10,

autogpt_platform/backend/backend/api/features/chat/tools/base.py

@@ -5,8 +5,8 @@ from typing import Any
 
 from openai.types.chat import ChatCompletionToolParam
 
-from backend.copilot.model import ChatSession
-from backend.copilot.response_model import StreamToolOutputAvailable
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.response_model import StreamToolOutputAvailable
 
 from .models import ErrorResponse, NeedLoginResponse, ToolResponseBase
 

autogpt_platform/backend/backend/api/features/chat/tools/bash_exec.py

@@ -11,11 +11,18 @@ available (e.g. macOS development).
 import logging
 from typing import Any
 
-from backend.copilot.model import ChatSession
-
-from .base import BaseTool
-from .models import BashExecResponse, ErrorResponse, ToolResponseBase
-from .sandbox import get_workspace_dir, has_full_sandbox, run_sandboxed
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.tools.base import BaseTool
+from backend.api.features.chat.tools.models import (
+    BashExecResponse,
+    ErrorResponse,
+    ToolResponseBase,
+)
+from backend.api.features.chat.tools.sandbox import (
+    get_workspace_dir,
+    has_full_sandbox,
+    run_sandboxed,
+)
 
 logger = logging.getLogger(__name__)
 

autogpt_platform/backend/backend/api/features/chat/tools/check_operation_status.py

@@ -3,10 +3,13 @@
 import logging
 from typing import Any
 
-from backend.copilot.model import ChatSession
-
-from .base import BaseTool
-from .models import ErrorResponse, ResponseType, ToolResponseBase
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.tools.base import BaseTool
+from backend.api.features.chat.tools.models import (
+    ErrorResponse,
+    ResponseType,
+    ToolResponseBase,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -75,7 +78,7 @@ class CheckOperationStatusTool(BaseTool):
         session: ChatSession,
         **kwargs,
     ) -> ToolResponseBase:
-        from backend.copilot import stream_registry
+        from backend.api.features.chat import stream_registry
 
         operation_id = (kwargs.get("operation_id") or "").strip()
         task_id = (kwargs.get("task_id") or "").strip()

autogpt_platform/backend/backend/api/features/chat/tools/create_agent.py

@@ -3,7 +3,7 @@
 import logging
 from typing import Any
 
-from backend.copilot.model import ChatSession
+from backend.api.features.chat.model import ChatSession
 
 from .agent_generator import (
     AgentGeneratorNotConfiguredError,
@@ -22,7 +22,6 @@ from .models import (
     ClarificationNeededResponse,
     ClarifyingQuestion,
     ErrorResponse,
-    SuggestedGoalResponse,
     ToolResponseBase,
 )
 
@@ -187,28 +186,26 @@ class CreateAgentTool(BaseTool):
         if decomposition_result.get("type") == "unachievable_goal":
             suggested = decomposition_result.get("suggested_goal", "")
             reason = decomposition_result.get("reason", "")
-            return SuggestedGoalResponse(
+            return ErrorResponse(
                 message=(
-                    f"This goal cannot be accomplished with the available blocks. {reason}"
+                    f"This goal cannot be accomplished with the available blocks. "
+                    f"{reason} "
+                    f"Suggestion: {suggested}"
                 ),
-                suggested_goal=suggested,
-                reason=reason,
-                original_goal=description,
-                goal_type="unachievable",
+                error="unachievable_goal",
+                details={"suggested_goal": suggested, "reason": reason},
                 session_id=session_id,
             )
 
         if decomposition_result.get("type") == "vague_goal":
             suggested = decomposition_result.get("suggested_goal", "")
-            reason = decomposition_result.get(
-                "reason", "The goal needs more specific details"
-            )
-            return SuggestedGoalResponse(
-                message="The goal is too vague to create a specific workflow.",
-                suggested_goal=suggested,
-                reason=reason,
-                original_goal=description,
-                goal_type="vague",
+            return ErrorResponse(
+                message=(
+                    f"The goal is too vague to create a specific workflow. "
+                    f"Suggestion: {suggested}"
+                ),
+                error="vague_goal",
+                details={"suggested_goal": suggested},
                 session_id=session_id,
             )
 

autogpt_platform/backend/backend/api/features/chat/tools/customize_agent.py

@@ -3,9 +3,9 @@
 import logging
 from typing import Any
 
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.store import db as store_db
 from backend.api.features.store.exceptions import AgentNotFoundError
-from backend.copilot.model import ChatSession
-from backend.data.db_accessors import store_db as get_store_db
 
 from .agent_generator import (
     AgentGeneratorNotConfiguredError,
@@ -137,8 +137,6 @@ class CustomizeAgentTool(BaseTool):
 
         creator_username, agent_slug = parts
 
-        store_db = get_store_db()
-
         # Fetch the marketplace agent details
         try:
             agent_details = await store_db.get_store_agent_details(

autogpt_platform/backend/backend/api/features/chat/tools/edit_agent.py

@@ -3,7 +3,7 @@
 import logging
 from typing import Any
 
-from backend.copilot.model import ChatSession
+from backend.api.features.chat.model import ChatSession
 
 from .agent_generator import (
     AgentGeneratorNotConfiguredError,

autogpt_platform/backend/backend/api/features/chat/tools/feature_requests.py

@@ -5,14 +5,9 @@ from typing import Any
 
 from pydantic import SecretStr
 
-from backend.blocks.linear._api import LinearClient
-from backend.copilot.model import ChatSession
-from backend.data.db_accessors import user_db
-from backend.data.model import APIKeyCredentials
-from backend.util.settings import Settings
-
-from .base import BaseTool
-from .models import (
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.tools.base import BaseTool
+from backend.api.features.chat.tools.models import (
     ErrorResponse,
     FeatureRequestCreatedResponse,
     FeatureRequestInfo,
@@ -20,6 +15,10 @@ from .models import (
     NoResultsResponse,
     ToolResponseBase,
 )
+from backend.blocks.linear._api import LinearClient
+from backend.data.model import APIKeyCredentials
+from backend.data.user import get_user_email_by_id
+from backend.util.settings import Settings
 
 logger = logging.getLogger(__name__)
 
@@ -33,6 +32,7 @@ query SearchFeatureRequests($term: String!, $filter: IssueFilter, $first: Int) {
       id
       identifier
       title
+      description
     }
   }
 }
@@ -104,8 +104,8 @@ def _get_linear_config() -> tuple[LinearClient, str, str]:
     Raises RuntimeError if any required setting is missing.
     """
     secrets = _get_settings().secrets
-    if not secrets.copilot_linear_api_key:
-        raise RuntimeError("COPILOT_LINEAR_API_KEY is not configured")
+    if not secrets.linear_api_key:
+        raise RuntimeError("LINEAR_API_KEY is not configured")
     if not secrets.linear_feature_request_project_id:
         raise RuntimeError("LINEAR_FEATURE_REQUEST_PROJECT_ID is not configured")
     if not secrets.linear_feature_request_team_id:
@@ -114,7 +114,7 @@ def _get_linear_config() -> tuple[LinearClient, str, str]:
     credentials = APIKeyCredentials(
         id="system-linear",
         provider="linear",
-        api_key=SecretStr(secrets.copilot_linear_api_key),
+        api_key=SecretStr(secrets.linear_api_key),
         title="System Linear API Key",
     )
     client = LinearClient(credentials=credentials)
@@ -204,6 +204,7 @@ class SearchFeatureRequestsTool(BaseTool):
                     id=node["id"],
                     identifier=node["identifier"],
                     title=node["title"],
+                    description=node.get("description"),
                 )
                 for node in nodes
             ]
@@ -237,11 +238,7 @@ class CreateFeatureRequestTool(BaseTool):
             "Create a new feature request or add a customer need to an existing one. "
             "Always search first with search_feature_requests to avoid duplicates. "
             "If a matching request exists, pass its ID as existing_issue_id to add "
-            "the user's need to it instead of creating a duplicate. "
-            "IMPORTANT: Never include personally identifiable information (PII) in "
-            "the title or description — no names, emails, phone numbers, company "
-            "names, or other identifying details. Write titles and descriptions in "
-            "generic, feature-focused language."
+            "the user's need to it instead of creating a duplicate."
         )
 
     @property
@@ -251,20 +248,11 @@ class CreateFeatureRequestTool(BaseTool):
             "properties": {
                 "title": {
                     "type": "string",
-                    "description": (
-                        "Title for the feature request. Must be generic and "
-                        "feature-focused — do not include any user names, emails, "
-                        "company names, or other PII."
-                    ),
+                    "description": "Title for the feature request.",
                 },
                 "description": {
                     "type": "string",
-                    "description": (
-                        "Detailed description of what the user wants and why. "
-                        "Must not contain any personally identifiable information "
-                        "(PII) — describe the feature need generically without "
-                        "referencing specific users, companies, or contact details."
-                    ),
+                    "description": "Detailed description of what the user wants and why.",
                 },
                 "existing_issue_id": {
                     "type": "string",
@@ -344,9 +332,7 @@ class CreateFeatureRequestTool(BaseTool):
         # Resolve a human-readable name (email) for the Linear customer record.
         # Fall back to user_id if the lookup fails or returns None.
         try:
-            customer_display_name = (
-                await user_db().get_user_email_by_id(user_id) or user_id
-            )
+            customer_display_name = await get_user_email_by_id(user_id) or user_id
         except Exception:
             customer_display_name = user_id
 

autogpt_platform/backend/backend/api/features/chat/tools/feature_requests_test.py

@@ -1,18 +1,22 @@
 """Tests for SearchFeatureRequestsTool and CreateFeatureRequestTool."""
 
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
 
-from ._test_data import make_session
-from .feature_requests import CreateFeatureRequestTool, SearchFeatureRequestsTool
-from .models import (
+from backend.api.features.chat.tools.feature_requests import (
+    CreateFeatureRequestTool,
+    SearchFeatureRequestsTool,
+)
+from backend.api.features.chat.tools.models import (
     ErrorResponse,
     FeatureRequestCreatedResponse,
     FeatureRequestSearchResponse,
     NoResultsResponse,
 )
 
+from ._test_data import make_session
+
 _TEST_USER_ID = "test-user-feature-requests"
 _TEST_USER_EMAIL = "testuser@example.com"
 
@@ -35,7 +39,7 @@ def _mock_linear_config(*, query_return=None, mutate_return=None):
         client.mutate.return_value = mutate_return
     return (
         patch(
-            "backend.copilot.tools.feature_requests._get_linear_config",
+            "backend.api.features.chat.tools.feature_requests._get_linear_config",
             return_value=(client, _FAKE_PROJECT_ID, _FAKE_TEAM_ID),
         ),
         client,
@@ -117,11 +121,13 @@ class TestSearchFeatureRequestsTool:
                 "id": "id-1",
                 "identifier": "FR-1",
                 "title": "Dark mode",
+                "description": "Add dark mode support",
             },
             {
                 "id": "id-2",
                 "identifier": "FR-2",
                 "title": "Dark theme",
+                "description": None,
             },
         ]
         patcher, _ = _mock_linear_config(query_return=_search_response(nodes))
@@ -202,7 +208,7 @@ class TestSearchFeatureRequestsTool:
     async def test_linear_client_init_failure(self):
         session = make_session(user_id=_TEST_USER_ID)
         with patch(
-            "backend.copilot.tools.feature_requests._get_linear_config",
+            "backend.api.features.chat.tools.feature_requests._get_linear_config",
             side_effect=RuntimeError("No API key"),
         ):
             tool = SearchFeatureRequestsTool()
@@ -225,11 +231,10 @@ class TestCreateFeatureRequestTool:
 
     @pytest.fixture(autouse=True)
     def _patch_email_lookup(self):
-        mock_user_db = MagicMock()
-        mock_user_db.get_user_email_by_id = AsyncMock(return_value=_TEST_USER_EMAIL)
         with patch(
-            "backend.copilot.tools.feature_requests.user_db",
-            return_value=mock_user_db,
+            "backend.api.features.chat.tools.feature_requests.get_user_email_by_id",
+            new_callable=AsyncMock,
+            return_value=_TEST_USER_EMAIL,
         ):
             yield
 
@@ -342,7 +347,7 @@ class TestCreateFeatureRequestTool:
     async def test_linear_client_init_failure(self):
         session = make_session(user_id=_TEST_USER_ID)
         with patch(
-            "backend.copilot.tools.feature_requests._get_linear_config",
+            "backend.api.features.chat.tools.feature_requests._get_linear_config",
             side_effect=RuntimeError("No API key"),
         ):
             tool = CreateFeatureRequestTool()

autogpt_platform/backend/backend/api/features/chat/tools/find_agent.py

@@ -2,7 +2,7 @@
 
 from typing import Any
 
-from backend.copilot.model import ChatSession
+from backend.api.features.chat.model import ChatSession
 
 from .agent_search import search_agents
 from .base import BaseTool

autogpt_platform/backend/backend/api/features/chat/tools/find_block.py

@@ -3,18 +3,17 @@ from typing import Any
 
 from prisma.enums import ContentType
 
-from backend.blocks import get_block
-from backend.blocks._base import BlockType
-from backend.copilot.model import ChatSession
-from backend.data.db_accessors import search
-
-from .base import BaseTool, ToolResponseBase
-from .models import (
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.tools.base import BaseTool, ToolResponseBase
+from backend.api.features.chat.tools.models import (
     BlockInfoSummary,
     BlockListResponse,
     ErrorResponse,
     NoResultsResponse,
 )
+from backend.api.features.store.hybrid_search import unified_hybrid_search
+from backend.blocks import get_block
+from backend.blocks._base import BlockType
 
 logger = logging.getLogger(__name__)
 
@@ -108,7 +107,7 @@ class FindBlockTool(BaseTool):
 
         try:
             # Search for blocks using hybrid search
-            results, total = await search().unified_hybrid_search(
+            results, total = await unified_hybrid_search(
                 query=query,
                 content_types=[ContentType.BLOCK],
                 page=1,

autogpt_platform/backend/backend/api/features/chat/tools/find_block_test.py

@@ -4,15 +4,15 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
-from backend.blocks._base import BlockType
-
-from ._test_data import make_session
-from .find_block import (
+from backend.api.features.chat.tools.find_block import (
     COPILOT_EXCLUDED_BLOCK_IDS,
     COPILOT_EXCLUDED_BLOCK_TYPES,
     FindBlockTool,
 )
-from .models import BlockListResponse
+from backend.api.features.chat.tools.models import BlockListResponse
+from backend.blocks._base import BlockType
+
+from ._test_data import make_session
 
 _TEST_USER_ID = "test-user-find-block"
 
@@ -84,17 +84,13 @@ class TestFindBlockFiltering:
                 "standard-block-id": standard_block,
             }.get(block_id)
 
-        mock_search_db = MagicMock()
-        mock_search_db.unified_hybrid_search = AsyncMock(
-            return_value=(search_results, 2)
-        )
-
         with patch(
-            "backend.copilot.tools.find_block.search",
-            return_value=mock_search_db,
+            "backend.api.features.chat.tools.find_block.unified_hybrid_search",
+            new_callable=AsyncMock,
+            return_value=(search_results, 2),
         ):
             with patch(
-                "backend.copilot.tools.find_block.get_block",
+                "backend.api.features.chat.tools.find_block.get_block",
                 side_effect=mock_get_block,
             ):
                 tool = FindBlockTool()
@@ -132,17 +128,13 @@ class TestFindBlockFiltering:
                 "normal-block-id": normal_block,
             }.get(block_id)
 
-        mock_search_db = MagicMock()
-        mock_search_db.unified_hybrid_search = AsyncMock(
-            return_value=(search_results, 2)
-        )
-
         with patch(
-            "backend.copilot.tools.find_block.search",
-            return_value=mock_search_db,
+            "backend.api.features.chat.tools.find_block.unified_hybrid_search",
+            new_callable=AsyncMock,
+            return_value=(search_results, 2),
         ):
             with patch(
-                "backend.copilot.tools.find_block.get_block",
+                "backend.api.features.chat.tools.find_block.get_block",
                 side_effect=mock_get_block,
             ):
                 tool = FindBlockTool()
@@ -361,16 +353,12 @@ class TestFindBlockFiltering:
             for d in block_defs
         }
 
-        mock_search_db = MagicMock()
-        mock_search_db.unified_hybrid_search = AsyncMock(
-            return_value=(search_results, len(search_results))
-        )
-
         with patch(
-            "backend.copilot.tools.find_block.search",
-            return_value=mock_search_db,
+            "backend.api.features.chat.tools.find_block.unified_hybrid_search",
+            new_callable=AsyncMock,
+            return_value=(search_results, len(search_results)),
         ), patch(
-            "backend.copilot.tools.find_block.get_block",
+            "backend.api.features.chat.tools.find_block.get_block",
             side_effect=lambda bid: mock_blocks.get(bid),
         ):
             tool = FindBlockTool()

autogpt_platform/backend/backend/api/features/chat/tools/find_library_agent.py

@@ -2,7 +2,7 @@
 
 from typing import Any
 
-from backend.copilot.model import ChatSession
+from backend.api.features.chat.model import ChatSession
 
 from .agent_search import search_agents
 from .base import BaseTool

autogpt_platform/backend/backend/api/features/chat/tools/get_doc_page.py

@@ -4,10 +4,13 @@ import logging
 from pathlib import Path
 from typing import Any
 
-from backend.copilot.model import ChatSession
-
-from .base import BaseTool
-from .models import DocPageResponse, ErrorResponse, ToolResponseBase
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.tools.base import BaseTool
+from backend.api.features.chat.tools.models import (
+    DocPageResponse,
+    ErrorResponse,
+    ToolResponseBase,
+)
 
 logger = logging.getLogger(__name__)
 

autogpt_platform/backend/backend/api/features/chat/tools/models.py

@@ -2,7 +2,7 @@
 
 from datetime import datetime
 from enum import Enum
-from typing import Any, Literal
+from typing import Any
 
 from pydantic import BaseModel, Field
 
@@ -50,8 +50,6 @@ class ResponseType(str, Enum):
     # Feature request types
     FEATURE_REQUEST_SEARCH = "feature_request_search"
     FEATURE_REQUEST_CREATED = "feature_request_created"
-    # Goal refinement
-    SUGGESTED_GOAL = "suggested_goal"
 
 
 # Base response model
@@ -298,22 +296,6 @@ class ClarificationNeededResponse(ToolResponseBase):
     questions: list[ClarifyingQuestion] = Field(default_factory=list)
 
 
-class SuggestedGoalResponse(ToolResponseBase):
-    """Response when the goal needs refinement with a suggested alternative."""
-
-    type: ResponseType = ResponseType.SUGGESTED_GOAL
-    suggested_goal: str = Field(description="The suggested alternative goal")
-    reason: str = Field(
-        default="", description="Why the original goal needs refinement"
-    )
-    original_goal: str = Field(
-        default="", description="The user's original goal for context"
-    )
-    goal_type: Literal["vague", "unachievable"] = Field(
-        default="vague", description="Type: 'vague' or 'unachievable'"
-    )
-
-
 # Documentation search models
 class DocSearchResult(BaseModel):
     """A single documentation search result."""
@@ -504,6 +486,7 @@ class FeatureRequestInfo(BaseModel):
     id: str
     identifier: str
     title: str
+    description: str | None = None
 
 
 class FeatureRequestSearchResponse(ToolResponseBase):

autogpt_platform/backend/backend/api/features/chat/tools/run_agent.py

@@ -5,12 +5,16 @@ from typing import Any
 
 from pydantic import BaseModel, Field, field_validator
 
-from backend.copilot.config import ChatConfig
-from backend.copilot.model import ChatSession
-from backend.copilot.tracking import track_agent_run_success, track_agent_scheduled
-from backend.data.db_accessors import graph_db, library_db, user_db
+from backend.api.features.chat.config import ChatConfig
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.tracking import (
+    track_agent_run_success,
+    track_agent_scheduled,
+)
+from backend.api.features.library import db as library_db
 from backend.data.graph import GraphModel
 from backend.data.model import CredentialsMetaInput
+from backend.data.user import get_user_by_id
 from backend.executor import utils as execution_utils
 from backend.util.clients import get_scheduler_client
 from backend.util.exceptions import DatabaseError, NotFoundError
@@ -196,7 +200,7 @@ class RunAgentTool(BaseTool):
 
             # Priority: library_agent_id if provided
             if has_library_id:
-                library_agent = await library_db().get_library_agent(
+                library_agent = await library_db.get_library_agent(
                     params.library_agent_id, user_id
                 )
                 if not library_agent:
@@ -205,7 +209,9 @@ class RunAgentTool(BaseTool):
                         session_id=session_id,
                     )
                 # Get the graph from the library agent
-                graph = await graph_db().get_graph(
+                from backend.data.graph import get_graph
+
+                graph = await get_graph(
                     library_agent.graph_id,
                     library_agent.graph_version,
                     user_id=user_id,
@@ -516,7 +522,7 @@ class RunAgentTool(BaseTool):
         library_agent = await get_or_create_library_agent(graph, user_id)
 
         # Get user timezone
-        user = await user_db().get_user_by_id(user_id)
+        user = await get_user_by_id(user_id)
         user_timezone = get_user_timezone_or_utc(user.timezone if user else timezone)
 
         # Create schedule

autogpt_platform/backend/backend/api/features/chat/tools/run_block.py

@@ -7,17 +7,20 @@ from typing import Any
 
 from pydantic_core import PydanticUndefined
 
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.tools.find_block import (
+    COPILOT_EXCLUDED_BLOCK_IDS,
+    COPILOT_EXCLUDED_BLOCK_TYPES,
+)
 from backend.blocks import get_block
 from backend.blocks._base import AnyBlockSchema
-from backend.copilot.model import ChatSession
-from backend.data.db_accessors import workspace_db
 from backend.data.execution import ExecutionContext
 from backend.data.model import CredentialsFieldInfo, CredentialsMetaInput
+from backend.data.workspace import get_or_create_workspace
 from backend.integrations.creds_manager import IntegrationCredentialsManager
 from backend.util.exceptions import BlockError
 
 from .base import BaseTool
-from .find_block import COPILOT_EXCLUDED_BLOCK_IDS, COPILOT_EXCLUDED_BLOCK_TYPES
 from .helpers import get_inputs_from_schema
 from .models import (
     BlockDetails,
@@ -273,7 +276,7 @@ class RunBlockTool(BaseTool):
 
         try:
             # Get or create user's workspace for CoPilot file operations
-            workspace = await workspace_db().get_or_create_workspace(user_id)
+            workspace = await get_or_create_workspace(user_id)
 
             # Generate synthetic IDs for CoPilot context
             # Each chat session is treated as its own agent with one continuous run

autogpt_platform/backend/backend/api/features/chat/tools/run_block_test.py

@@ -4,16 +4,16 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
-from backend.blocks._base import BlockType
-
-from ._test_data import make_session
-from .models import (
+from backend.api.features.chat.tools.models import (
     BlockDetailsResponse,
     BlockOutputResponse,
     ErrorResponse,
     InputValidationErrorResponse,
 )
-from .run_block import RunBlockTool
+from backend.api.features.chat.tools.run_block import RunBlockTool
+from backend.blocks._base import BlockType
+
+from ._test_data import make_session
 
 _TEST_USER_ID = "test-user-run-block"
 
@@ -77,7 +77,7 @@ class TestRunBlockFiltering:
         input_block = make_mock_block("input-block-id", "Input Block", BlockType.INPUT)
 
         with patch(
-            "backend.copilot.tools.run_block.get_block",
+            "backend.api.features.chat.tools.run_block.get_block",
             return_value=input_block,
         ):
             tool = RunBlockTool()
@@ -103,7 +103,7 @@ class TestRunBlockFiltering:
         )
 
         with patch(
-            "backend.copilot.tools.run_block.get_block",
+            "backend.api.features.chat.tools.run_block.get_block",
             return_value=smart_block,
         ):
             tool = RunBlockTool()
@@ -127,7 +127,7 @@ class TestRunBlockFiltering:
         )
 
         with patch(
-            "backend.copilot.tools.run_block.get_block",
+            "backend.api.features.chat.tools.run_block.get_block",
             return_value=standard_block,
         ):
             tool = RunBlockTool()
@@ -183,7 +183,7 @@ class TestRunBlockInputValidation:
         )
 
         with patch(
-            "backend.copilot.tools.run_block.get_block",
+            "backend.api.features.chat.tools.run_block.get_block",
             return_value=mock_block,
         ):
             tool = RunBlockTool()
@@ -222,7 +222,7 @@ class TestRunBlockInputValidation:
         )
 
         with patch(
-            "backend.copilot.tools.run_block.get_block",
+            "backend.api.features.chat.tools.run_block.get_block",
             return_value=mock_block,
         ):
             tool = RunBlockTool()
@@ -263,7 +263,7 @@ class TestRunBlockInputValidation:
         )
 
         with patch(
-            "backend.copilot.tools.run_block.get_block",
+            "backend.api.features.chat.tools.run_block.get_block",
             return_value=mock_block,
         ):
             tool = RunBlockTool()
@@ -302,19 +302,15 @@ class TestRunBlockInputValidation:
 
         mock_block.execute = mock_execute
 
-        mock_workspace_db = MagicMock()
-        mock_workspace_db.get_or_create_workspace = AsyncMock(
-            return_value=MagicMock(id="test-workspace-id")
-        )
-
         with (
             patch(
-                "backend.copilot.tools.run_block.get_block",
+                "backend.api.features.chat.tools.run_block.get_block",
                 return_value=mock_block,
             ),
             patch(
-                "backend.copilot.tools.run_block.workspace_db",
-                return_value=mock_workspace_db,
+                "backend.api.features.chat.tools.run_block.get_or_create_workspace",
+                new_callable=AsyncMock,
+                return_value=MagicMock(id="test-workspace-id"),
             ),
         ):
             tool = RunBlockTool()
@@ -348,7 +344,7 @@ class TestRunBlockInputValidation:
         )
 
         with patch(
-            "backend.copilot.tools.run_block.get_block",
+            "backend.api.features.chat.tools.run_block.get_block",
             return_value=mock_block,
         ):
             tool = RunBlockTool()

autogpt_platform/backend/backend/api/features/chat/tools/sandbox.py

@@ -13,7 +13,6 @@ import logging
 import os
 import platform
 import shutil
-import signal
 
 logger = logging.getLogger(__name__)
 
@@ -246,7 +245,6 @@ async def run_sandboxed(
             stderr=asyncio.subprocess.PIPE,
             cwd=cwd,
             env=safe_env,
-            start_new_session=True,  # Own process group for clean kill
         )
 
         try:
@@ -257,18 +255,7 @@ async def run_sandboxed(
             stderr = stderr_bytes.decode("utf-8", errors="replace")
             return stdout, stderr, proc.returncode or 0, False
         except asyncio.TimeoutError:
-            # Kill entire process group (bwrap + all children).
-            # proc.kill() alone only kills the bwrap parent, leaving
-            # children running until they finish naturally.
-            try:
-                os.killpg(proc.pid, signal.SIGKILL)
-            except ProcessLookupError:
-                pass  # Already exited
-            except OSError as kill_err:
-                logger.warning(
-                    "Failed to kill process group %d: %s", proc.pid, kill_err
-                )
-            # Always reap the subprocess regardless of killpg outcome.
+            proc.kill()
             await proc.communicate()
             return "", f"Execution timed out after {timeout}s", -1, True
 

autogpt_platform/backend/backend/api/features/chat/tools/search_docs.py

@@ -5,17 +5,16 @@ from typing import Any
 
 from prisma.enums import ContentType
 
-from backend.copilot.model import ChatSession
-from backend.data.db_accessors import search
-
-from .base import BaseTool
-from .models import (
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.tools.base import BaseTool
+from backend.api.features.chat.tools.models import (
     DocSearchResult,
     DocSearchResultsResponse,
     ErrorResponse,
     NoResultsResponse,
     ToolResponseBase,
 )
+from backend.api.features.store.hybrid_search import unified_hybrid_search
 
 logger = logging.getLogger(__name__)
 
@@ -118,7 +117,7 @@ class SearchDocsTool(BaseTool):
 
         try:
             # Search using hybrid search for DOCUMENTATION content type only
-            results, total = await search().unified_hybrid_search(
+            results, total = await unified_hybrid_search(
                 query=query,
                 content_types=[ContentType.DOCUMENTATION],
                 page=1,

autogpt_platform/backend/backend/api/features/chat/tools/test_run_block_details.py

@@ -4,13 +4,13 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
+from backend.api.features.chat.tools.models import BlockDetailsResponse
+from backend.api.features.chat.tools.run_block import RunBlockTool
 from backend.blocks._base import BlockType
 from backend.data.model import CredentialsMetaInput
 from backend.integrations.providers import ProviderName
 
 from ._test_data import make_session
-from .models import BlockDetailsResponse
-from .run_block import RunBlockTool
 
 _TEST_USER_ID = "test-user-run-block-details"
 
@@ -61,7 +61,7 @@ async def test_run_block_returns_details_when_no_input_provided():
     )
 
     with patch(
-        "backend.copilot.tools.run_block.get_block",
+        "backend.api.features.chat.tools.run_block.get_block",
         return_value=http_block,
     ):
         # Mock credentials check to return no missing credentials
@@ -120,7 +120,7 @@ async def test_run_block_returns_details_when_only_credentials_provided():
     }
 
     with patch(
-        "backend.copilot.tools.run_block.get_block",
+        "backend.api.features.chat.tools.run_block.get_block",
         return_value=mock,
     ):
         with patch.object(

autogpt_platform/backend/backend/api/features/chat/tools/utils.py

@@ -3,8 +3,9 @@
 import logging
 from typing import Any
 
+from backend.api.features.library import db as library_db
 from backend.api.features.library import model as library_model
-from backend.data.db_accessors import library_db, store_db
+from backend.api.features.store import db as store_db
 from backend.data.graph import GraphModel
 from backend.data.model import (
     Credentials,
@@ -38,14 +39,13 @@ async def fetch_graph_from_store_slug(
     Raises:
         DatabaseError: If there's a database error during lookup.
     """
-    sdb = store_db()
     try:
-        store_agent = await sdb.get_store_agent_details(username, agent_name)
+        store_agent = await store_db.get_store_agent_details(username, agent_name)
     except NotFoundError:
         return None, None
 
     # Get the graph from store listing version
-    graph = await sdb.get_available_graph(
+    graph = await store_db.get_available_graph(
         store_agent.store_listing_version_id, hide_nodes=False
     )
     return graph, store_agent
@@ -210,13 +210,13 @@ async def get_or_create_library_agent(
     Returns:
         LibraryAgent instance
     """
-    existing = await library_db().get_library_agent_by_graph_id(
+    existing = await library_db.get_library_agent_by_graph_id(
         graph_id=graph.id, user_id=user_id
     )
     if existing:
         return existing
 
-    library_agents = await library_db().create_library_agent(
+    library_agents = await library_db.create_library_agent(
         graph=graph,
         user_id=user_id,
         create_library_agents_for_sub_graphs=False,

autogpt_platform/backend/backend/api/features/chat/tools/web_fetch.py

@@ -6,12 +6,15 @@ from typing import Any
 import aiohttp
 import html2text
 
-from backend.copilot.model import ChatSession
+from backend.api.features.chat.model import ChatSession
+from backend.api.features.chat.tools.base import BaseTool
+from backend.api.features.chat.tools.models import (
+    ErrorResponse,
+    ToolResponseBase,
+    WebFetchResponse,
+)
 from backend.util.request import Requests
 
-from .base import BaseTool
-from .models import ErrorResponse, ToolResponseBase, WebFetchResponse
-
 logger = logging.getLogger(__name__)
 
 # Limits

autogpt_platform/backend/backend/api/features/chat/tools/workspace_files.py

@@ -0,0 +1,626 @@
+"""CoPilot tools for workspace file operations."""
+
+import base64
+import logging
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+from backend.api.features.chat.model import ChatSession
+from backend.data.workspace import get_or_create_workspace
+from backend.util.settings import Config
+from backend.util.virus_scanner import scan_content_safe
+from backend.util.workspace import WorkspaceManager
+
+from .base import BaseTool
+from .models import ErrorResponse, ResponseType, ToolResponseBase
+
+logger = logging.getLogger(__name__)
+
+
+class WorkspaceFileInfoData(BaseModel):
+    """Data model for workspace file information (not a response itself)."""
+
+    file_id: str
+    name: str
+    path: str
+    mime_type: str
+    size_bytes: int
+
+
+class WorkspaceFileListResponse(ToolResponseBase):
+    """Response containing list of workspace files."""
+
+    type: ResponseType = ResponseType.WORKSPACE_FILE_LIST
+    files: list[WorkspaceFileInfoData]
+    total_count: int
+
+
+class WorkspaceFileContentResponse(ToolResponseBase):
+    """Response containing workspace file content (legacy, for small text files)."""
+
+    type: ResponseType = ResponseType.WORKSPACE_FILE_CONTENT
+    file_id: str
+    name: str
+    path: str
+    mime_type: str
+    content_base64: str
+
+
+class WorkspaceFileMetadataResponse(ToolResponseBase):
+    """Response containing workspace file metadata and download URL (prevents context bloat)."""
+
+    type: ResponseType = ResponseType.WORKSPACE_FILE_METADATA
+    file_id: str
+    name: str
+    path: str
+    mime_type: str
+    size_bytes: int
+    download_url: str
+    preview: str | None = None  # First 500 chars for text files
+
+
+class WorkspaceWriteResponse(ToolResponseBase):
+    """Response after writing a file to workspace."""
+
+    type: ResponseType = ResponseType.WORKSPACE_FILE_WRITTEN
+    file_id: str
+    name: str
+    path: str
+    size_bytes: int
+
+
+class WorkspaceDeleteResponse(ToolResponseBase):
+    """Response after deleting a file from workspace."""
+
+    type: ResponseType = ResponseType.WORKSPACE_FILE_DELETED
+    file_id: str
+    success: bool
+
+
+class ListWorkspaceFilesTool(BaseTool):
+    """Tool for listing files in user's workspace."""
+
+    @property
+    def name(self) -> str:
+        return "list_workspace_files"
+
+    @property
+    def description(self) -> str:
+        return (
+            "List files in the user's persistent workspace (cloud storage). "
+            "These files survive across sessions. "
+            "For ephemeral session files, use the SDK Read/Glob tools instead. "
+            "Returns file names, paths, sizes, and metadata. "
+            "Optionally filter by path prefix."
+        )
+
+    @property
+    def parameters(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "path_prefix": {
+                    "type": "string",
+                    "description": (
+                        "Optional path prefix to filter files "
+                        "(e.g., '/documents/' to list only files in documents folder). "
+                        "By default, only files from the current session are listed."
+                    ),
+                },
+                "limit": {
+                    "type": "integer",
+                    "description": "Maximum number of files to return (default 50, max 100)",
+                    "minimum": 1,
+                    "maximum": 100,
+                },
+                "include_all_sessions": {
+                    "type": "boolean",
+                    "description": (
+                        "If true, list files from all sessions. "
+                        "Default is false (only current session's files)."
+                    ),
+                },
+            },
+            "required": [],
+        }
+
+    @property
+    def requires_auth(self) -> bool:
+        return True
+
+    async def _execute(
+        self,
+        user_id: str | None,
+        session: ChatSession,
+        **kwargs,
+    ) -> ToolResponseBase:
+        session_id = session.session_id
+
+        if not user_id:
+            return ErrorResponse(
+                message="Authentication required",
+                session_id=session_id,
+            )
+
+        path_prefix: Optional[str] = kwargs.get("path_prefix")
+        limit = min(kwargs.get("limit", 50), 100)
+        include_all_sessions: bool = kwargs.get("include_all_sessions", False)
+
+        try:
+            workspace = await get_or_create_workspace(user_id)
+            # Pass session_id for session-scoped file access
+            manager = WorkspaceManager(user_id, workspace.id, session_id)
+
+            files = await manager.list_files(
+                path=path_prefix,
+                limit=limit,
+                include_all_sessions=include_all_sessions,
+            )
+            total = await manager.get_file_count(
+                path=path_prefix,
+                include_all_sessions=include_all_sessions,
+            )
+
+            file_infos = [
+                WorkspaceFileInfoData(
+                    file_id=f.id,
+                    name=f.name,
+                    path=f.path,
+                    mime_type=f.mimeType,
+                    size_bytes=f.sizeBytes,
+                )
+                for f in files
+            ]
+
+            scope_msg = "all sessions" if include_all_sessions else "current session"
+            return WorkspaceFileListResponse(
+                files=file_infos,
+                total_count=total,
+                message=f"Found {len(files)} files in workspace ({scope_msg})",
+                session_id=session_id,
+            )
+
+        except Exception as e:
+            logger.error(f"Error listing workspace files: {e}", exc_info=True)
+            return ErrorResponse(
+                message=f"Failed to list workspace files: {str(e)}",
+                error=str(e),
+                session_id=session_id,
+            )
+
+
+class ReadWorkspaceFileTool(BaseTool):
+    """Tool for reading file content from workspace."""
+
+    # Size threshold for returning full content vs metadata+URL
+    # Files larger than this return metadata with download URL to prevent context bloat
+    MAX_INLINE_SIZE_BYTES = 32 * 1024  # 32KB
+    # Preview size for text files
+    PREVIEW_SIZE = 500
+
+    @property
+    def name(self) -> str:
+        return "read_workspace_file"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Read a file from the user's persistent workspace (cloud storage). "
+            "These files survive across sessions. "
+            "For ephemeral session files, use the SDK Read tool instead. "
+            "Specify either file_id or path to identify the file. "
+            "For small text files, returns content directly. "
+            "For large or binary files, returns metadata and a download URL. "
+            "Paths are scoped to the current session by default. "
+            "Use /sessions/<session_id>/... for cross-session access."
+        )
+
+    @property
+    def parameters(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "file_id": {
+                    "type": "string",
+                    "description": "The file's unique ID (from list_workspace_files)",
+                },
+                "path": {
+                    "type": "string",
+                    "description": (
+                        "The virtual file path (e.g., '/documents/report.pdf'). "
+                        "Scoped to current session by default."
+                    ),
+                },
+                "force_download_url": {
+                    "type": "boolean",
+                    "description": (
+                        "If true, always return metadata+URL instead of inline content. "
+                        "Default is false (auto-selects based on file size/type)."
+                    ),
+                },
+            },
+            "required": [],  # At least one must be provided
+        }
+
+    @property
+    def requires_auth(self) -> bool:
+        return True
+
+    def _is_text_mime_type(self, mime_type: str) -> bool:
+        """Check if the MIME type is a text-based type."""
+        text_types = [
+            "text/",
+            "application/json",
+            "application/xml",
+            "application/javascript",
+            "application/x-python",
+            "application/x-sh",
+        ]
+        return any(mime_type.startswith(t) for t in text_types)
+
+    async def _execute(
+        self,
+        user_id: str | None,
+        session: ChatSession,
+        **kwargs,
+    ) -> ToolResponseBase:
+        session_id = session.session_id
+
+        if not user_id:
+            return ErrorResponse(
+                message="Authentication required",
+                session_id=session_id,
+            )
+
+        file_id: Optional[str] = kwargs.get("file_id")
+        path: Optional[str] = kwargs.get("path")
+        force_download_url: bool = kwargs.get("force_download_url", False)
+
+        if not file_id and not path:
+            return ErrorResponse(
+                message="Please provide either file_id or path",
+                session_id=session_id,
+            )
+
+        try:
+            workspace = await get_or_create_workspace(user_id)
+            # Pass session_id for session-scoped file access
+            manager = WorkspaceManager(user_id, workspace.id, session_id)
+
+            # Get file info
+            if file_id:
+                file_info = await manager.get_file_info(file_id)
+                if file_info is None:
+                    return ErrorResponse(
+                        message=f"File not found: {file_id}",
+                        session_id=session_id,
+                    )
+                target_file_id = file_id
+            else:
+                # path is guaranteed to be non-None here due to the check above
+                assert path is not None
+                file_info = await manager.get_file_info_by_path(path)
+                if file_info is None:
+                    return ErrorResponse(
+                        message=f"File not found at path: {path}",
+                        session_id=session_id,
+                    )
+                target_file_id = file_info.id
+
+            # Decide whether to return inline content or metadata+URL
+            is_small_file = file_info.sizeBytes <= self.MAX_INLINE_SIZE_BYTES
+            is_text_file = self._is_text_mime_type(file_info.mimeType)
+
+            # Return inline content for small text files (unless force_download_url)
+            if is_small_file and is_text_file and not force_download_url:
+                content = await manager.read_file_by_id(target_file_id)
+                content_b64 = base64.b64encode(content).decode("utf-8")
+
+                return WorkspaceFileContentResponse(
+                    file_id=file_info.id,
+                    name=file_info.name,
+                    path=file_info.path,
+                    mime_type=file_info.mimeType,
+                    content_base64=content_b64,
+                    message=f"Successfully read file: {file_info.name}",
+                    session_id=session_id,
+                )
+
+            # Return metadata + workspace:// reference for large or binary files
+            # This prevents context bloat (100KB file = ~133KB as base64)
+            # Use workspace:// format so frontend urlTransform can add proxy prefix
+            download_url = f"workspace://{target_file_id}"
+
+            # Generate preview for text files
+            preview: str | None = None
+            if is_text_file:
+                try:
+                    content = await manager.read_file_by_id(target_file_id)
+                    preview_text = content[: self.PREVIEW_SIZE].decode(
+                        "utf-8", errors="replace"
+                    )
+                    if len(content) > self.PREVIEW_SIZE:
+                        preview_text += "..."
+                    preview = preview_text
+                except Exception:
+                    pass  # Preview is optional
+
+            return WorkspaceFileMetadataResponse(
+                file_id=file_info.id,
+                name=file_info.name,
+                path=file_info.path,
+                mime_type=file_info.mimeType,
+                size_bytes=file_info.sizeBytes,
+                download_url=download_url,
+                preview=preview,
+                message=f"File: {file_info.name} ({file_info.sizeBytes} bytes). Use download_url to retrieve content.",
+                session_id=session_id,
+            )
+
+        except FileNotFoundError as e:
+            return ErrorResponse(
+                message=str(e),
+                session_id=session_id,
+            )
+        except Exception as e:
+            logger.error(f"Error reading workspace file: {e}", exc_info=True)
+            return ErrorResponse(
+                message=f"Failed to read workspace file: {str(e)}",
+                error=str(e),
+                session_id=session_id,
+            )
+
+
+class WriteWorkspaceFileTool(BaseTool):
+    """Tool for writing files to workspace."""
+
+    @property
+    def name(self) -> str:
+        return "write_workspace_file"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Write or create a file in the user's persistent workspace (cloud storage). "
+            "These files survive across sessions. "
+            "For ephemeral session files, use the SDK Write tool instead. "
+            "Provide the content as a base64-encoded string. "
+            f"Maximum file size is {Config().max_file_size_mb}MB. "
+            "Files are saved to the current session's folder by default. "
+            "Use /sessions/<session_id>/... for cross-session access."
+        )
+
+    @property
+    def parameters(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "filename": {
+                    "type": "string",
+                    "description": "Name for the file (e.g., 'report.pdf')",
+                },
+                "content_base64": {
+                    "type": "string",
+                    "description": "Base64-encoded file content",
+                },
+                "path": {
+                    "type": "string",
+                    "description": (
+                        "Optional virtual path where to save the file "
+                        "(e.g., '/documents/report.pdf'). "
+                        "Defaults to '/{filename}'. Scoped to current session."
+                    ),
+                },
+                "mime_type": {
+                    "type": "string",
+                    "description": (
+                        "Optional MIME type of the file. "
+                        "Auto-detected from filename if not provided."
+                    ),
+                },
+                "overwrite": {
+                    "type": "boolean",
+                    "description": "Whether to overwrite if file exists at path (default: false)",
+                },
+            },
+            "required": ["filename", "content_base64"],
+        }
+
+    @property
+    def requires_auth(self) -> bool:
+        return True
+
+    async def _execute(
+        self,
+        user_id: str | None,
+        session: ChatSession,
+        **kwargs,
+    ) -> ToolResponseBase:
+        session_id = session.session_id
+
+        if not user_id:
+            return ErrorResponse(
+                message="Authentication required",
+                session_id=session_id,
+            )
+
+        filename: str = kwargs.get("filename", "")
+        content_b64: str = kwargs.get("content_base64", "")
+        path: Optional[str] = kwargs.get("path")
+        mime_type: Optional[str] = kwargs.get("mime_type")
+        overwrite: bool = kwargs.get("overwrite", False)
+
+        if not filename:
+            return ErrorResponse(
+                message="Please provide a filename",
+                session_id=session_id,
+            )
+
+        if not content_b64:
+            return ErrorResponse(
+                message="Please provide content_base64",
+                session_id=session_id,
+            )
+
+        # Decode content
+        try:
+            content = base64.b64decode(content_b64)
+        except Exception:
+            return ErrorResponse(
+                message="Invalid base64-encoded content",
+                session_id=session_id,
+            )
+
+        # Check size
+        max_file_size = Config().max_file_size_mb * 1024 * 1024
+        if len(content) > max_file_size:
+            return ErrorResponse(
+                message=f"File too large. Maximum size is {Config().max_file_size_mb}MB",
+                session_id=session_id,
+            )
+
+        try:
+            # Virus scan
+            await scan_content_safe(content, filename=filename)
+
+            workspace = await get_or_create_workspace(user_id)
+            # Pass session_id for session-scoped file access
+            manager = WorkspaceManager(user_id, workspace.id, session_id)
+
+            file_record = await manager.write_file(
+                content=content,
+                filename=filename,
+                path=path,
+                mime_type=mime_type,
+                overwrite=overwrite,
+            )
+
+            return WorkspaceWriteResponse(
+                file_id=file_record.id,
+                name=file_record.name,
+                path=file_record.path,
+                size_bytes=file_record.sizeBytes,
+                message=f"Successfully wrote file: {file_record.name}",
+                session_id=session_id,
+            )
+
+        except ValueError as e:
+            return ErrorResponse(
+                message=str(e),
+                session_id=session_id,
+            )
+        except Exception as e:
+            logger.error(f"Error writing workspace file: {e}", exc_info=True)
+            return ErrorResponse(
+                message=f"Failed to write workspace file: {str(e)}",
+                error=str(e),
+                session_id=session_id,
+            )
+
+
+class DeleteWorkspaceFileTool(BaseTool):
+    """Tool for deleting files from workspace."""
+
+    @property
+    def name(self) -> str:
+        return "delete_workspace_file"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Delete a file from the user's persistent workspace (cloud storage). "
+            "Specify either file_id or path to identify the file. "
+            "Paths are scoped to the current session by default. "
+            "Use /sessions/<session_id>/... for cross-session access."
+        )
+
+    @property
+    def parameters(self) -> dict[str, Any]:
+        return {
+            "type": "object",
+            "properties": {
+                "file_id": {
+                    "type": "string",
+                    "description": "The file's unique ID (from list_workspace_files)",
+                },
+                "path": {
+                    "type": "string",
+                    "description": (
+                        "The virtual file path (e.g., '/documents/report.pdf'). "
+                        "Scoped to current session by default."
+                    ),
+                },
+            },
+            "required": [],  # At least one must be provided
+        }
+
+    @property
+    def requires_auth(self) -> bool:
+        return True
+
+    async def _execute(
+        self,
+        user_id: str | None,
+        session: ChatSession,
+        **kwargs,
+    ) -> ToolResponseBase:
+        session_id = session.session_id
+
+        if not user_id:
+            return ErrorResponse(
+                message="Authentication required",
+                session_id=session_id,
+            )
+
+        file_id: Optional[str] = kwargs.get("file_id")
+        path: Optional[str] = kwargs.get("path")
+
+        if not file_id and not path:
+            return ErrorResponse(
+                message="Please provide either file_id or path",
+                session_id=session_id,
+            )
+
+        try:
+            workspace = await get_or_create_workspace(user_id)
+            # Pass session_id for session-scoped file access
+            manager = WorkspaceManager(user_id, workspace.id, session_id)
+
+            # Determine the file_id to delete
+            target_file_id: str
+            if file_id:
+                target_file_id = file_id
+            else:
+                # path is guaranteed to be non-None here due to the check above
+                assert path is not None
+                file_info = await manager.get_file_info_by_path(path)
+                if file_info is None:
+                    return ErrorResponse(
+                        message=f"File not found at path: {path}",
+                        session_id=session_id,
+                    )
+                target_file_id = file_info.id
+
+            success = await manager.delete_file(target_file_id)
+
+            if not success:
+                return ErrorResponse(
+                    message=f"File not found: {target_file_id}",
+                    session_id=session_id,
+                )
+
+            return WorkspaceDeleteResponse(
+                file_id=target_file_id,
+                success=True,
+                message="File deleted successfully",
+                session_id=session_id,
+            )
+
+        except Exception as e:
+            logger.error(f"Error deleting workspace file: {e}", exc_info=True)
+            return ErrorResponse(
+                message=f"Failed to delete workspace file: {str(e)}",
+                error=str(e),
+                session_id=session_id,
+            )

autogpt_platform/backend/backend/api/features/workspace/routes.py

@@ -11,7 +11,7 @@ import fastapi
 from autogpt_libs.auth.dependencies import get_user_id, requires_user
 from fastapi.responses import Response
 
-from backend.data.workspace import WorkspaceFile, get_workspace, get_workspace_file
+from backend.data.workspace import get_workspace, get_workspace_file
 from backend.util.workspace_storage import get_workspace_storage
 
 
@@ -44,11 +44,11 @@ router = fastapi.APIRouter(
 )
 
 
-def _create_streaming_response(content: bytes, file: WorkspaceFile) -> Response:
+def _create_streaming_response(content: bytes, file) -> Response:
     """Create a streaming response for file content."""
     return Response(
         content=content,
-        media_type=file.mime_type,
+        media_type=file.mimeType,
         headers={
             "Content-Disposition": _sanitize_filename_for_header(file.name),
             "Content-Length": str(len(content)),
@@ -56,7 +56,7 @@ def _create_streaming_response(content: bytes, file: WorkspaceFile) -> Response:
     )
 
 
-async def _create_file_download_response(file: WorkspaceFile) -> Response:
+async def _create_file_download_response(file) -> Response:
     """
     Create a download response for a workspace file.
 
@@ -66,33 +66,33 @@ async def _create_file_download_response(file: WorkspaceFile) -> Response:
     storage = await get_workspace_storage()
 
     # For local storage, stream the file directly
-    if file.storage_path.startswith("local://"):
-        content = await storage.retrieve(file.storage_path)
+    if file.storagePath.startswith("local://"):
+        content = await storage.retrieve(file.storagePath)
         return _create_streaming_response(content, file)
 
     # For GCS, try to redirect to signed URL, fall back to streaming
     try:
-        url = await storage.get_download_url(file.storage_path, expires_in=300)
+        url = await storage.get_download_url(file.storagePath, expires_in=300)
         # If we got back an API path (fallback), stream directly instead
         if url.startswith("/api/"):
-            content = await storage.retrieve(file.storage_path)
+            content = await storage.retrieve(file.storagePath)
             return _create_streaming_response(content, file)
         return fastapi.responses.RedirectResponse(url=url, status_code=302)
     except Exception as e:
         # Log the signed URL failure with context
         logger.error(
             f"Failed to get signed URL for file {file.id} "
-            f"(storagePath={file.storage_path}): {e}",
+            f"(storagePath={file.storagePath}): {e}",
             exc_info=True,
         )
         # Fall back to streaming directly from GCS
         try:
-            content = await storage.retrieve(file.storage_path)
+            content = await storage.retrieve(file.storagePath)
             return _create_streaming_response(content, file)
         except Exception as fallback_error:
             logger.error(
                 f"Fallback streaming also failed for file {file.id} "
-                f"(storagePath={file.storage_path}): {fallback_error}",
+                f"(storagePath={file.storagePath}): {fallback_error}",
                 exc_info=True,
             )
             raise

autogpt_platform/backend/backend/api/rest_api.py

@@ -41,11 +41,11 @@ import backend.data.user
 import backend.integrations.webhooks.utils
 import backend.util.service
 import backend.util.settings
-from backend.blocks.llm import DEFAULT_LLM_MODEL
-from backend.copilot.completion_consumer import (
+from backend.api.features.chat.completion_consumer import (
     start_completion_consumer,
     stop_completion_consumer,
 )
+from backend.blocks.llm import DEFAULT_LLM_MODEL
 from backend.data.model import Credentials
 from backend.integrations.providers import ProviderName
 from backend.monitoring.instrumentation import instrument_fastapi

autogpt_platform/backend/backend/app.py

@@ -38,9 +38,7 @@ def main(**kwargs):
 
     from backend.api.rest_api import AgentServer
     from backend.api.ws_api import WebsocketServer
-    from backend.copilot.executor.manager import CoPilotExecutor
-    from backend.data.db_manager import DatabaseManager
-    from backend.executor import ExecutionManager, Scheduler
+    from backend.executor import DatabaseManager, ExecutionManager, Scheduler
     from backend.notifications import NotificationManager
 
     run_processes(
@@ -50,7 +48,6 @@ def main(**kwargs):
         WebsocketServer(),
         AgentServer(),
         ExecutionManager(),
-        CoPilotExecutor(),
         **kwargs,
     )
 

autogpt_platform/backend/backend/conftest.py

@@ -1,7 +1,6 @@
 import logging
 import os
 
-import pytest
 import pytest_asyncio
 from dotenv import load_dotenv
 
@@ -28,54 +27,6 @@ async def server():
         yield server
 
 
-@pytest.fixture
-def test_user_id() -> str:
-    """Test user ID fixture."""
-    return "3e53486c-cf57-477e-ba2a-cb02dc828e1a"
-
-
-@pytest.fixture
-def admin_user_id() -> str:
-    """Admin user ID fixture."""
-    return "4e53486c-cf57-477e-ba2a-cb02dc828e1b"
-
-
-@pytest.fixture
-def target_user_id() -> str:
-    """Target user ID fixture."""
-    return "5e53486c-cf57-477e-ba2a-cb02dc828e1c"
-
-
-@pytest.fixture
-async def setup_test_user(test_user_id):
-    """Create test user in database before tests."""
-    from backend.data.user import get_or_create_user
-
-    # Create the test user in the database using JWT token format
-    user_data = {
-        "sub": test_user_id,
-        "email": "test@example.com",
-        "user_metadata": {"name": "Test User"},
-    }
-    await get_or_create_user(user_data)
-    return test_user_id
-
-
-@pytest.fixture
-async def setup_admin_user(admin_user_id):
-    """Create admin user in database before tests."""
-    from backend.data.user import get_or_create_user
-
-    # Create the admin user in the database using JWT token format
-    user_data = {
-        "sub": admin_user_id,
-        "email": "test-admin@example.com",
-        "user_metadata": {"name": "Test Admin"},
-    }
-    await get_or_create_user(user_data)
-    return admin_user_id
-
-
 @pytest_asyncio.fixture(scope="session", loop_scope="session", autouse=True)
 async def graph_cleanup(server):
     created_graph_ids = []

autogpt_platform/backend/backend/copilot/__init__.py

@@ -1 +0,0 @@
-

autogpt_platform/backend/backend/copilot/executor/__main__.py

@@ -1,18 +0,0 @@
-"""Entry point for running the CoPilot Executor service.
-
-Usage:
-    python -m backend.copilot.executor
-"""
-
-from backend.app import run_processes
-
-from .manager import CoPilotExecutor
-
-
-def main():
-    """Run the CoPilot Executor service."""
-    run_processes(CoPilotExecutor())
-
-
-if __name__ == "__main__":
-    main()

autogpt_platform/backend/backend/copilot/executor/manager.py

@@ -1,520 +0,0 @@
-"""CoPilot Executor Manager - main service for CoPilot task execution.
-
-This module contains the CoPilotExecutor class that consumes chat tasks from
-RabbitMQ and processes them using a thread pool, following the graph executor pattern.
-"""
-
-import logging
-import os
-import threading
-import time
-import uuid
-from concurrent.futures import Future, ThreadPoolExecutor
-
-from pika.adapters.blocking_connection import BlockingChannel
-from pika.exceptions import AMQPChannelError, AMQPConnectionError
-from pika.spec import Basic, BasicProperties
-from prometheus_client import Gauge, start_http_server
-
-from backend.data import redis_client as redis
-from backend.data.rabbitmq import SyncRabbitMQ
-from backend.executor.cluster_lock import ClusterLock
-from backend.util.decorator import error_logged
-from backend.util.logging import TruncatedLogger
-from backend.util.process import AppProcess
-from backend.util.retry import continuous_retry
-from backend.util.settings import Settings
-
-from .processor import execute_copilot_task, init_worker
-from .utils import (
-    COPILOT_CANCEL_QUEUE_NAME,
-    COPILOT_EXECUTION_QUEUE_NAME,
-    GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS,
-    CancelCoPilotEvent,
-    CoPilotExecutionEntry,
-    create_copilot_queue_config,
-)
-
-logger = TruncatedLogger(logging.getLogger(__name__), prefix="[CoPilotExecutor]")
-settings = Settings()
-
-# Prometheus metrics
-active_tasks_gauge = Gauge(
-    "copilot_executor_active_tasks",
-    "Number of active CoPilot tasks",
-)
-pool_size_gauge = Gauge(
-    "copilot_executor_pool_size",
-    "Maximum number of CoPilot executor workers",
-)
-utilization_gauge = Gauge(
-    "copilot_executor_utilization_ratio",
-    "Ratio of active tasks to pool size",
-)
-
-
-class CoPilotExecutor(AppProcess):
-    """CoPilot Executor service for processing chat generation tasks.
-
-    This service consumes tasks from RabbitMQ, processes them using a thread pool,
-    and publishes results to Redis Streams. It follows the graph executor pattern
-    for reliable message handling and graceful shutdown.
-
-    Key features:
-    - RabbitMQ-based task distribution with manual acknowledgment
-    - Thread pool executor for concurrent task processing
-    - Cluster lock for duplicate prevention across pods
-    - Graceful shutdown with timeout for in-flight tasks
-    - FANOUT exchange for cancellation broadcast
-    """
-
-    def __init__(self):
-        super().__init__()
-        self.pool_size = settings.config.num_copilot_workers
-        self.active_tasks: dict[str, tuple[Future, threading.Event]] = {}
-        self.executor_id = str(uuid.uuid4())
-
-        self._executor = None
-        self._stop_consuming = None
-
-        self._cancel_thread = None
-        self._cancel_client = None
-        self._run_thread = None
-        self._run_client = None
-
-        self._task_locks: dict[str, ClusterLock] = {}
-        self._active_tasks_lock = threading.Lock()
-
-    # ============ Main Entry Points (AppProcess interface) ============ #
-
-    def run(self):
-        """Main service loop - consume from RabbitMQ."""
-        logger.info(f"Pod assigned executor_id: {self.executor_id}")
-        logger.info(f"Spawn max-{self.pool_size} workers...")
-
-        pool_size_gauge.set(self.pool_size)
-        self._update_metrics()
-        start_http_server(settings.config.copilot_executor_port)
-
-        self.cancel_thread.start()
-        self.run_thread.start()
-
-        while True:
-            time.sleep(1e5)
-
-    def cleanup(self):
-        """Graceful shutdown with active execution waiting."""
-        pid = os.getpid()
-        logger.info(f"[cleanup {pid}] Starting graceful shutdown...")
-
-        # Signal the consumer thread to stop
-        try:
-            self.stop_consuming.set()
-            run_channel = self.run_client.get_channel()
-            run_channel.connection.add_callback_threadsafe(
-                lambda: run_channel.stop_consuming()
-            )
-            logger.info(f"[cleanup {pid}] Consumer has been signaled to stop")
-        except Exception as e:
-            logger.error(f"[cleanup {pid}] Error stopping consumer: {e}")
-
-        # Wait for active executions to complete
-        if self.active_tasks:
-            logger.info(
-                f"[cleanup {pid}] Waiting for {len(self.active_tasks)} active tasks to complete (timeout: {GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS}s)..."
-            )
-
-            start_time = time.monotonic()
-            last_refresh = start_time
-            lock_refresh_interval = settings.config.cluster_lock_timeout / 10
-
-            while (
-                self.active_tasks
-                and (time.monotonic() - start_time) < GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS
-            ):
-                self._cleanup_completed_tasks()
-                if not self.active_tasks:
-                    break
-
-                # Refresh cluster locks periodically
-                current_time = time.monotonic()
-                if current_time - last_refresh >= lock_refresh_interval:
-                    for lock in list(self._task_locks.values()):
-                        try:
-                            lock.refresh()
-                        except Exception as e:
-                            logger.warning(
-                                f"[cleanup {pid}] Failed to refresh lock: {e}"
-                            )
-                    last_refresh = current_time
-
-                logger.info(
-                    f"[cleanup {pid}] {len(self.active_tasks)} tasks still active, waiting..."
-                )
-                time.sleep(10.0)
-
-        # Stop message consumers
-        if self._run_thread:
-            self._stop_message_consumers(
-                self._run_thread, self.run_client, "[cleanup][run]"
-            )
-        if self._cancel_thread:
-            self._stop_message_consumers(
-                self._cancel_thread, self.cancel_client, "[cleanup][cancel]"
-            )
-
-        # Clean up worker threads (closes per-loop workspace storage sessions)
-        if self._executor:
-            from .processor import cleanup_worker
-
-            logger.info(f"[cleanup {pid}] Cleaning up workers...")
-            futures = []
-            for _ in range(self._executor._max_workers):
-                futures.append(self._executor.submit(cleanup_worker))
-            for f in futures:
-                try:
-                    f.result(timeout=10)
-                except Exception as e:
-                    logger.warning(f"[cleanup {pid}] Worker cleanup error: {e}")
-
-            logger.info(f"[cleanup {pid}] Shutting down executor...")
-            self._executor.shutdown(wait=False)
-
-        # Release any remaining locks
-        for task_id, lock in list(self._task_locks.items()):
-            try:
-                lock.release()
-                logger.info(f"[cleanup {pid}] Released lock for {task_id}")
-            except Exception as e:
-                logger.error(
-                    f"[cleanup {pid}] Failed to release lock for {task_id}: {e}"
-                )
-
-        logger.info(f"[cleanup {pid}] Graceful shutdown completed")
-
-    # ============ RabbitMQ Consumer Methods ============ #
-
-    @continuous_retry()
-    def _consume_cancel(self):
-        """Consume cancellation messages from FANOUT exchange."""
-        if self.stop_consuming.is_set() and not self.active_tasks:
-            logger.info("Stop reconnecting cancel consumer - service cleaned up")
-            return
-
-        if not self.cancel_client.is_ready:
-            self.cancel_client.disconnect()
-        self.cancel_client.connect()
-
-        # Check again after connect - shutdown may have been requested
-        if self.stop_consuming.is_set() and not self.active_tasks:
-            logger.info("Stop consuming requested during reconnect - disconnecting")
-            self.cancel_client.disconnect()
-            return
-
-        cancel_channel = self.cancel_client.get_channel()
-        cancel_channel.basic_consume(
-            queue=COPILOT_CANCEL_QUEUE_NAME,
-            on_message_callback=self._handle_cancel_message,
-            auto_ack=True,
-        )
-        logger.info("Starting to consume cancel messages...")
-        cancel_channel.start_consuming()
-        if not self.stop_consuming.is_set() or self.active_tasks:
-            raise RuntimeError("Cancel message consumer stopped unexpectedly")
-        logger.info("Cancel message consumer stopped gracefully")
-
-    @continuous_retry()
-    def _consume_run(self):
-        """Consume run messages from DIRECT exchange."""
-        if self.stop_consuming.is_set():
-            logger.info("Stop reconnecting run consumer - service cleaned up")
-            return
-
-        if not self.run_client.is_ready:
-            self.run_client.disconnect()
-        self.run_client.connect()
-
-        # Check again after connect - shutdown may have been requested
-        if self.stop_consuming.is_set():
-            logger.info("Stop consuming requested during reconnect - disconnecting")
-            self.run_client.disconnect()
-            return
-
-        run_channel = self.run_client.get_channel()
-        run_channel.basic_qos(prefetch_count=self.pool_size)
-
-        run_channel.basic_consume(
-            queue=COPILOT_EXECUTION_QUEUE_NAME,
-            on_message_callback=self._handle_run_message,
-            auto_ack=False,
-            consumer_tag="copilot_execution_consumer",
-        )
-        logger.info("Starting to consume run messages...")
-        run_channel.start_consuming()
-        if not self.stop_consuming.is_set():
-            raise RuntimeError("Run message consumer stopped unexpectedly")
-        logger.info("Run message consumer stopped gracefully")
-
-    # ============ Message Handlers ============ #
-
-    @error_logged(swallow=True)
-    def _handle_cancel_message(
-        self,
-        _channel: BlockingChannel,
-        _method: Basic.Deliver,
-        _properties: BasicProperties,
-        body: bytes,
-    ):
-        """Handle cancel message from FANOUT exchange."""
-        request = CancelCoPilotEvent.model_validate_json(body)
-        task_id = request.task_id
-        if not task_id:
-            logger.warning("Cancel message missing 'task_id'")
-            return
-        if task_id not in self.active_tasks:
-            logger.debug(f"Cancel received for {task_id} but not active")
-            return
-
-        _, cancel_event = self.active_tasks[task_id]
-        logger.info(f"Received cancel for {task_id}")
-        if not cancel_event.is_set():
-            cancel_event.set()
-        else:
-            logger.debug(f"Cancel already set for {task_id}")
-
-    def _handle_run_message(
-        self,
-        _channel: BlockingChannel,
-        method: Basic.Deliver,
-        _properties: BasicProperties,
-        body: bytes,
-    ):
-        """Handle run message from DIRECT exchange."""
-        delivery_tag = method.delivery_tag
-        # Capture the channel used at message delivery time to ensure we ack
-        # on the correct channel. Delivery tags are channel-scoped and become
-        # invalid if the channel is recreated after reconnection.
-        delivery_channel = _channel
-
-        def ack_message(reject: bool, requeue: bool):
-            """Acknowledge or reject the message.
-
-            Uses the channel from the original message delivery. If the channel
-            is no longer open (e.g., after reconnection), logs a warning and
-            skips the ack - RabbitMQ will redeliver the message automatically.
-            """
-            try:
-                if not delivery_channel.is_open:
-                    logger.warning(
-                        f"Channel closed, cannot ack delivery_tag={delivery_tag}. "
-                        "Message will be redelivered by RabbitMQ."
-                    )
-                    return
-
-                if reject:
-                    delivery_channel.connection.add_callback_threadsafe(
-                        lambda: delivery_channel.basic_nack(
-                            delivery_tag, requeue=requeue
-                        )
-                    )
-                else:
-                    delivery_channel.connection.add_callback_threadsafe(
-                        lambda: delivery_channel.basic_ack(delivery_tag)
-                    )
-            except (AMQPChannelError, AMQPConnectionError) as e:
-                # Channel/connection errors indicate stale delivery tag - don't retry
-                logger.warning(
-                    f"Cannot ack delivery_tag={delivery_tag} due to channel/connection "
-                    f"error: {e}. Message will be redelivered by RabbitMQ."
-                )
-            except Exception as e:
-                # Other errors might be transient, but log and skip to avoid blocking
-                logger.error(
-                    f"Unexpected error acking delivery_tag={delivery_tag}: {e}"
-                )
-
-        # Check if we're shutting down
-        if self.stop_consuming.is_set():
-            logger.info("Rejecting new task during shutdown")
-            ack_message(reject=True, requeue=True)
-            return
-
-        # Check if we can accept more tasks
-        self._cleanup_completed_tasks()
-        if len(self.active_tasks) >= self.pool_size:
-            ack_message(reject=True, requeue=True)
-            return
-
-        try:
-            entry = CoPilotExecutionEntry.model_validate_json(body)
-        except Exception as e:
-            logger.error(f"Could not parse run message: {e}, body={body}")
-            ack_message(reject=True, requeue=False)
-            return
-
-        task_id = entry.task_id
-
-        # Check for local duplicate - task is already running on this executor
-        if task_id in self.active_tasks:
-            logger.warning(
-                f"Task {task_id} already running locally, rejecting duplicate"
-            )
-            ack_message(reject=True, requeue=False)
-            return
-
-        # Try to acquire cluster-wide lock
-        cluster_lock = ClusterLock(
-            redis=redis.get_redis(),
-            key=f"copilot:task:{task_id}:lock",
-            owner_id=self.executor_id,
-            timeout=settings.config.cluster_lock_timeout,
-        )
-        current_owner = cluster_lock.try_acquire()
-        if current_owner != self.executor_id:
-            if current_owner is not None:
-                logger.warning(f"Task {task_id} already running on pod {current_owner}")
-                ack_message(reject=True, requeue=False)
-            else:
-                logger.warning(
-                    f"Could not acquire lock for {task_id} - Redis unavailable"
-                )
-                ack_message(reject=True, requeue=True)
-            return
-
-        # Execute the task
-        try:
-            self._task_locks[task_id] = cluster_lock
-
-            logger.info(
-                f"Acquired cluster lock for {task_id}, executor_id={self.executor_id}"
-            )
-
-            cancel_event = threading.Event()
-            future = self.executor.submit(
-                execute_copilot_task, entry, cancel_event, cluster_lock
-            )
-            self.active_tasks[task_id] = (future, cancel_event)
-        except Exception as e:
-            logger.warning(f"Failed to setup execution for {task_id}: {e}")
-            cluster_lock.release()
-            if task_id in self._task_locks:
-                del self._task_locks[task_id]
-            ack_message(reject=True, requeue=True)
-            return
-
-        self._update_metrics()
-
-        def on_run_done(f: Future):
-            logger.info(f"Run completed for {task_id}")
-            try:
-                if exec_error := f.exception():
-                    logger.error(f"Execution for {task_id} failed: {exec_error}")
-                    # Don't requeue failed tasks - they've been marked as failed
-                    # in the stream registry. Requeuing would cause infinite retries
-                    # for deterministic failures.
-                    ack_message(reject=True, requeue=False)
-                else:
-                    ack_message(reject=False, requeue=False)
-            except BaseException as e:
-                logger.exception(f"Error in run completion callback: {e}")
-            finally:
-                # Release the cluster lock
-                if task_id in self._task_locks:
-                    logger.info(f"Releasing cluster lock for {task_id}")
-                    self._task_locks[task_id].release()
-                    del self._task_locks[task_id]
-                self._cleanup_completed_tasks()
-
-        future.add_done_callback(on_run_done)
-
-    # ============ Helper Methods ============ #
-
-    def _cleanup_completed_tasks(self) -> list[str]:
-        """Remove completed futures from active_tasks and update metrics."""
-        completed_tasks = []
-        with self._active_tasks_lock:
-            for task_id, (future, _) in list(self.active_tasks.items()):
-                if future.done():
-                    completed_tasks.append(task_id)
-                    self.active_tasks.pop(task_id, None)
-                    logger.info(f"Cleaned up completed task {task_id}")
-
-        self._update_metrics()
-        return completed_tasks
-
-    def _update_metrics(self):
-        """Update Prometheus metrics."""
-        active_count = len(self.active_tasks)
-        active_tasks_gauge.set(active_count)
-        if self.stop_consuming.is_set():
-            utilization_gauge.set(1.0)
-        else:
-            utilization_gauge.set(
-                active_count / self.pool_size if self.pool_size > 0 else 0
-            )
-
-    def _stop_message_consumers(
-        self, thread: threading.Thread, client: SyncRabbitMQ, prefix: str
-    ):
-        """Stop a message consumer thread."""
-        try:
-            channel = client.get_channel()
-            channel.connection.add_callback_threadsafe(lambda: channel.stop_consuming())
-
-            thread.join(timeout=300)
-            if thread.is_alive():
-                logger.error(
-                    f"{prefix} Thread did not finish in time, forcing disconnect"
-                )
-
-            client.disconnect()
-            logger.info(f"{prefix} Client disconnected")
-        except Exception as e:
-            logger.error(f"{prefix} Error disconnecting client: {e}")
-
-    # ============ Lazy-initialized Properties ============ #
-
-    @property
-    def cancel_thread(self) -> threading.Thread:
-        if self._cancel_thread is None:
-            self._cancel_thread = threading.Thread(
-                target=lambda: self._consume_cancel(),
-                daemon=True,
-            )
-        return self._cancel_thread
-
-    @property
-    def run_thread(self) -> threading.Thread:
-        if self._run_thread is None:
-            self._run_thread = threading.Thread(
-                target=lambda: self._consume_run(),
-                daemon=True,
-            )
-        return self._run_thread
-
-    @property
-    def stop_consuming(self) -> threading.Event:
-        if self._stop_consuming is None:
-            self._stop_consuming = threading.Event()
-        return self._stop_consuming
-
-    @property
-    def executor(self) -> ThreadPoolExecutor:
-        if self._executor is None:
-            self._executor = ThreadPoolExecutor(
-                max_workers=self.pool_size,
-                initializer=init_worker,
-            )
-        return self._executor
-
-    @property
-    def cancel_client(self) -> SyncRabbitMQ:
-        if self._cancel_client is None:
-            self._cancel_client = SyncRabbitMQ(create_copilot_queue_config())
-        return self._cancel_client
-
-    @property
-    def run_client(self) -> SyncRabbitMQ:
-        if self._run_client is None:
-            self._run_client = SyncRabbitMQ(create_copilot_queue_config())
-        return self._run_client

autogpt_platform/backend/backend/copilot/executor/processor.py

@@ -1,291 +0,0 @@
-"""CoPilot execution processor - per-worker execution logic.
-
-This module contains the processor class that handles CoPilot task execution
-in a thread-local context, following the graph executor pattern.
-"""
-
-import asyncio
-import logging
-import threading
-import time
-
-from backend.copilot import service as copilot_service
-from backend.copilot import stream_registry
-from backend.copilot.config import ChatConfig
-from backend.copilot.response_model import StreamError, StreamFinish, StreamFinishStep
-from backend.copilot.sdk import service as sdk_service
-from backend.executor.cluster_lock import ClusterLock
-from backend.util.decorator import error_logged
-from backend.util.feature_flag import Flag, is_feature_enabled
-from backend.util.logging import TruncatedLogger, configure_logging
-from backend.util.process import set_service_name
-from backend.util.retry import func_retry
-
-from .utils import CoPilotExecutionEntry, CoPilotLogMetadata
-
-logger = TruncatedLogger(logging.getLogger(__name__), prefix="[CoPilotExecutor]")
-
-
-# ============ Module Entry Points ============ #
-
-# Thread-local storage for processor instances
-_tls = threading.local()
-
-
-def execute_copilot_task(
-    entry: CoPilotExecutionEntry,
-    cancel: threading.Event,
-    cluster_lock: ClusterLock,
-):
-    """Execute a CoPilot task using the thread-local processor.
-
-    This function is the entry point called by the thread pool executor.
-
-    Args:
-        entry: The task payload
-        cancel: Threading event to signal cancellation
-        cluster_lock: Distributed lock for this execution
-    """
-    processor: CoPilotProcessor = _tls.processor
-    return processor.execute(entry, cancel, cluster_lock)
-
-
-def init_worker():
-    """Initialize the processor for the current worker thread.
-
-    This function is called by the thread pool executor when a new worker
-    thread is created. It ensures each worker has its own processor instance.
-    """
-    _tls.processor = CoPilotProcessor()
-    _tls.processor.on_executor_start()
-
-
-def cleanup_worker():
-    """Clean up the processor for the current worker thread.
-
-    Should be called before the worker thread's event loop is destroyed so
-    that event-loop-bound resources (e.g. ``aiohttp.ClientSession``) are
-    closed on the correct loop.
-    """
-    processor: CoPilotProcessor | None = getattr(_tls, "processor", None)
-    if processor is not None:
-        processor.cleanup()
-
-
-# ============ Processor Class ============ #
-
-
-class CoPilotProcessor:
-    """Per-worker execution logic for CoPilot tasks.
-
-    This class is instantiated once per worker thread and handles the execution
-    of CoPilot chat generation tasks. It maintains an async event loop for
-    running the async service code.
-
-    The execution flow:
-        1. CoPilot task is picked from RabbitMQ queue
-        2. Manager submits task to thread pool
-        3. Processor executes the task in its event loop
-        4. Results are published to Redis Streams
-    """
-
-    @func_retry
-    def on_executor_start(self):
-        """Initialize the processor when the worker thread starts.
-
-        This method is called once per worker thread to set up the async event
-        loop and initialize any required resources.
-
-        Database is accessed only through DatabaseManager, so we don't need to connect
-        to Prisma directly.
-        """
-        configure_logging()
-        set_service_name("CoPilotExecutor")
-        self.tid = threading.get_ident()
-        self.execution_loop = asyncio.new_event_loop()
-        self.execution_thread = threading.Thread(
-            target=self.execution_loop.run_forever, daemon=True
-        )
-        self.execution_thread.start()
-
-        logger.info(f"[CoPilotExecutor] Worker {self.tid} started")
-
-    def cleanup(self):
-        """Clean up event-loop-bound resources before the loop is destroyed.
-
-        Shuts down the workspace storage instance that belongs to this
-        worker's event loop, ensuring ``aiohttp.ClientSession.close()``
-        runs on the same loop that created the session.
-        """
-        from backend.util.workspace_storage import shutdown_workspace_storage
-
-        try:
-            future = asyncio.run_coroutine_threadsafe(
-                shutdown_workspace_storage(), self.execution_loop
-            )
-            future.result(timeout=5)
-        except Exception as e:
-            logger.warning(f"[CoPilotExecutor] Worker {self.tid} cleanup error: {e}")
-
-        # Stop the event loop
-        self.execution_loop.call_soon_threadsafe(self.execution_loop.stop)
-        self.execution_thread.join(timeout=5)
-        logger.info(f"[CoPilotExecutor] Worker {self.tid} cleaned up")
-
-    @error_logged(swallow=False)
-    def execute(
-        self,
-        entry: CoPilotExecutionEntry,
-        cancel: threading.Event,
-        cluster_lock: ClusterLock,
-    ):
-        """Execute a CoPilot task.
-
-        This is the main entry point for task execution. It runs the async
-        execution logic in the worker's event loop and handles errors.
-
-        Args:
-            entry: The task payload containing session and message info
-            cancel: Threading event to signal cancellation
-            cluster_lock: Distributed lock to prevent duplicate execution
-        """
-        log = CoPilotLogMetadata(
-            logging.getLogger(__name__),
-            task_id=entry.task_id,
-            session_id=entry.session_id,
-            user_id=entry.user_id,
-        )
-        log.info("Starting execution")
-
-        start_time = time.monotonic()
-
-        try:
-            # Run the async execution in our event loop
-            future = asyncio.run_coroutine_threadsafe(
-                self._execute_async(entry, cancel, cluster_lock, log),
-                self.execution_loop,
-            )
-
-            # Wait for completion, checking cancel periodically
-            while not future.done():
-                try:
-                    future.result(timeout=1.0)
-                except asyncio.TimeoutError:
-                    if cancel.is_set():
-                        log.info("Cancellation requested")
-                        future.cancel()
-                        break
-                    # Refresh cluster lock to maintain ownership
-                    cluster_lock.refresh()
-
-            if not future.cancelled():
-                # Get result to propagate any exceptions
-                future.result()
-
-            elapsed = time.monotonic() - start_time
-            log.info(f"Execution completed in {elapsed:.2f}s")
-
-        except Exception as e:
-            elapsed = time.monotonic() - start_time
-            log.error(f"Execution failed after {elapsed:.2f}s: {e}")
-            # Note: _execute_async already marks the task as failed before re-raising,
-            # so we don't call _mark_task_failed here to avoid duplicate error events.
-            raise
-
-    async def _execute_async(
-        self,
-        entry: CoPilotExecutionEntry,
-        cancel: threading.Event,
-        cluster_lock: ClusterLock,
-        log: CoPilotLogMetadata,
-    ):
-        """Async execution logic for CoPilot task.
-
-        This method calls the existing stream_chat_completion service function
-        and publishes results to the stream registry.
-
-        Args:
-            entry: The task payload
-            cancel: Threading event to signal cancellation
-            cluster_lock: Distributed lock for refresh
-            log: Structured logger for this task
-        """
-        last_refresh = time.monotonic()
-        refresh_interval = 30.0  # Refresh lock every 30 seconds
-
-        try:
-            # Choose service based on LaunchDarkly flag
-            config = ChatConfig()
-            use_sdk = await is_feature_enabled(
-                Flag.COPILOT_SDK,
-                entry.user_id or "anonymous",
-                default=config.use_claude_agent_sdk,
-            )
-            stream_fn = (
-                sdk_service.stream_chat_completion_sdk
-                if use_sdk
-                else copilot_service.stream_chat_completion
-            )
-            log.info(f"Using {'SDK' if use_sdk else 'standard'} service")
-
-            # Stream chat completion and publish chunks to Redis
-            async for chunk in stream_fn(
-                session_id=entry.session_id,
-                message=entry.message if entry.message else None,
-                is_user_message=entry.is_user_message,
-                user_id=entry.user_id,
-                context=entry.context,
-            ):
-                # Check for cancellation
-                if cancel.is_set():
-                    log.info("Cancelled during streaming")
-                    await stream_registry.publish_chunk(
-                        entry.task_id, StreamError(errorText="Operation cancelled")
-                    )
-                    await stream_registry.publish_chunk(
-                        entry.task_id, StreamFinishStep()
-                    )
-                    await stream_registry.publish_chunk(entry.task_id, StreamFinish())
-                    await stream_registry.mark_task_completed(
-                        entry.task_id, status="failed"
-                    )
-                    return
-
-                # Refresh cluster lock periodically
-                current_time = time.monotonic()
-                if current_time - last_refresh >= refresh_interval:
-                    cluster_lock.refresh()
-                    last_refresh = current_time
-
-                # Publish chunk to stream registry
-                await stream_registry.publish_chunk(entry.task_id, chunk)
-
-            # Mark task as completed
-            await stream_registry.mark_task_completed(entry.task_id, status="completed")
-            log.info("Task completed successfully")
-
-        except asyncio.CancelledError:
-            log.info("Task cancelled")
-            await stream_registry.mark_task_completed(
-                entry.task_id,
-                status="failed",
-                error_message="Task was cancelled",
-            )
-            raise
-
-        except Exception as e:
-            log.error(f"Task failed: {e}")
-            await self._mark_task_failed(entry.task_id, str(e))
-            raise
-
-    async def _mark_task_failed(self, task_id: str, error_message: str):
-        """Mark a task as failed and publish error to stream registry."""
-        try:
-            await stream_registry.publish_chunk(
-                task_id, StreamError(errorText=error_message)
-            )
-            await stream_registry.publish_chunk(task_id, StreamFinishStep())
-            await stream_registry.publish_chunk(task_id, StreamFinish())
-            await stream_registry.mark_task_completed(task_id, status="failed")
-        except Exception as e:
-            logger.error(f"Failed to mark task {task_id} as failed: {e}")

autogpt_platform/backend/backend/copilot/executor/utils.py

@@ -1,224 +0,0 @@
-"""RabbitMQ queue configuration for CoPilot executor.
-
-Defines two exchanges and queues following the graph executor pattern:
-- 'copilot_execution' (DIRECT) for chat generation tasks
-- 'copilot_cancel' (FANOUT) for cancellation requests
-"""
-
-import logging
-
-from pydantic import BaseModel
-
-from backend.data.rabbitmq import Exchange, ExchangeType, Queue, RabbitMQConfig
-from backend.util.logging import TruncatedLogger, is_structured_logging_enabled
-
-logger = logging.getLogger(__name__)
-
-
-# ============ Logging Helper ============ #
-
-
-class CoPilotLogMetadata(TruncatedLogger):
-    """Structured logging helper for CoPilot executor.
-
-    In cloud environments (structured logging enabled), uses a simple prefix
-    and passes metadata via json_fields. In local environments, uses a detailed
-    prefix with all metadata key-value pairs for easier debugging.
-
-    Args:
-        logger: The underlying logger instance
-        max_length: Maximum log message length before truncation
-        **kwargs: Metadata key-value pairs (e.g., task_id="abc", session_id="xyz")
-            These are added to json_fields in cloud mode, or to the prefix in local mode.
-    """
-
-    def __init__(
-        self,
-        logger: logging.Logger,
-        max_length: int = 1000,
-        **kwargs: str | None,
-    ):
-        # Filter out None values
-        metadata = {k: v for k, v in kwargs.items() if v is not None}
-        metadata["component"] = "CoPilotExecutor"
-
-        if is_structured_logging_enabled():
-            prefix = "[CoPilotExecutor]"
-        else:
-            # Build prefix from metadata key-value pairs
-            meta_parts = "|".join(
-                f"{k}:{v}" for k, v in metadata.items() if k != "component"
-            )
-            prefix = (
-                f"[CoPilotExecutor|{meta_parts}]" if meta_parts else "[CoPilotExecutor]"
-            )
-
-        super().__init__(
-            logger,
-            max_length=max_length,
-            prefix=prefix,
-            metadata=metadata,
-        )
-
-
-# ============ Exchange and Queue Configuration ============ #
-
-COPILOT_EXECUTION_EXCHANGE = Exchange(
-    name="copilot_execution",
-    type=ExchangeType.DIRECT,
-    durable=True,
-    auto_delete=False,
-)
-COPILOT_EXECUTION_QUEUE_NAME = "copilot_execution_queue"
-COPILOT_EXECUTION_ROUTING_KEY = "copilot.run"
-
-COPILOT_CANCEL_EXCHANGE = Exchange(
-    name="copilot_cancel",
-    type=ExchangeType.FANOUT,
-    durable=True,
-    auto_delete=False,
-)
-COPILOT_CANCEL_QUEUE_NAME = "copilot_cancel_queue"
-
-# CoPilot operations can include extended thinking and agent generation
-# which may take 30+ minutes to complete
-COPILOT_CONSUMER_TIMEOUT_SECONDS = 60 * 60  # 1 hour
-
-# Graceful shutdown timeout - allow in-flight operations to complete
-GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS = 30 * 60  # 30 minutes
-
-
-def create_copilot_queue_config() -> RabbitMQConfig:
-    """Create RabbitMQ configuration for CoPilot executor.
-
-    Defines two exchanges and queues:
-    - 'copilot_execution' (DIRECT) for chat generation tasks
-    - 'copilot_cancel' (FANOUT) for cancellation requests
-
-    Returns:
-        RabbitMQConfig with exchanges and queues defined
-    """
-    run_queue = Queue(
-        name=COPILOT_EXECUTION_QUEUE_NAME,
-        exchange=COPILOT_EXECUTION_EXCHANGE,
-        routing_key=COPILOT_EXECUTION_ROUTING_KEY,
-        durable=True,
-        auto_delete=False,
-        arguments={
-            # Extended consumer timeout for long-running LLM operations
-            # Default 30-minute timeout is insufficient for extended thinking
-            # and agent generation which can take 30+ minutes
-            "x-consumer-timeout": COPILOT_CONSUMER_TIMEOUT_SECONDS
-            * 1000,
-        },
-    )
-    cancel_queue = Queue(
-        name=COPILOT_CANCEL_QUEUE_NAME,
-        exchange=COPILOT_CANCEL_EXCHANGE,
-        routing_key="",  # not used for FANOUT
-        durable=True,
-        auto_delete=False,
-    )
-    return RabbitMQConfig(
-        vhost="/",
-        exchanges=[COPILOT_EXECUTION_EXCHANGE, COPILOT_CANCEL_EXCHANGE],
-        queues=[run_queue, cancel_queue],
-    )
-
-
-# ============ Message Models ============ #
-
-
-class CoPilotExecutionEntry(BaseModel):
-    """Task payload for CoPilot AI generation.
-
-    This model represents a chat generation task to be processed by the executor.
-    """
-
-    task_id: str
-    """Unique identifier for this task (used for stream registry)"""
-
-    session_id: str
-    """Chat session ID"""
-
-    user_id: str | None
-    """User ID (may be None for anonymous users)"""
-
-    operation_id: str
-    """Operation ID for webhook callbacks and completion tracking"""
-
-    message: str
-    """User's message to process"""
-
-    is_user_message: bool = True
-    """Whether the message is from the user (vs system/assistant)"""
-
-    context: dict[str, str] | None = None
-    """Optional context for the message (e.g., {url: str, content: str})"""
-
-
-class CancelCoPilotEvent(BaseModel):
-    """Event to cancel a CoPilot operation."""
-
-    task_id: str
-    """Task ID to cancel"""
-
-
-# ============ Queue Publishing Helpers ============ #
-
-
-async def enqueue_copilot_task(
-    task_id: str,
-    session_id: str,
-    user_id: str | None,
-    operation_id: str,
-    message: str,
-    is_user_message: bool = True,
-    context: dict[str, str] | None = None,
-) -> None:
-    """Enqueue a CoPilot task for processing by the executor service.
-
-    Args:
-        task_id: Unique identifier for this task (used for stream registry)
-        session_id: Chat session ID
-        user_id: User ID (may be None for anonymous users)
-        operation_id: Operation ID for webhook callbacks and completion tracking
-        message: User's message to process
-        is_user_message: Whether the message is from the user (vs system/assistant)
-        context: Optional context for the message (e.g., {url: str, content: str})
-    """
-    from backend.util.clients import get_async_copilot_queue
-
-    entry = CoPilotExecutionEntry(
-        task_id=task_id,
-        session_id=session_id,
-        user_id=user_id,
-        operation_id=operation_id,
-        message=message,
-        is_user_message=is_user_message,
-        context=context,
-    )
-
-    queue_client = await get_async_copilot_queue()
-    await queue_client.publish_message(
-        routing_key=COPILOT_EXECUTION_ROUTING_KEY,
-        message=entry.model_dump_json(),
-        exchange=COPILOT_EXECUTION_EXCHANGE,
-    )
-
-
-async def enqueue_cancel_task(task_id: str) -> None:
-    """Publish a cancel request for a running CoPilot task.
-
-    Sends a ``CancelCoPilotEvent`` to the FANOUT exchange so all executor
-    pods receive the cancellation signal.
-    """
-    from backend.util.clients import get_async_copilot_queue
-
-    event = CancelCoPilotEvent(task_id=task_id)
-    queue_client = await get_async_copilot_queue()
-    await queue_client.publish_message(
-        routing_key="",  # FANOUT ignores routing key
-        message=event.model_dump_json(),
-        exchange=COPILOT_CANCEL_EXCHANGE,
-    )

autogpt_platform/backend/backend/copilot/parallel_tool_calls_test.py

@@ -1,272 +0,0 @@
-"""Tests for parallel tool call execution in CoPilot.
-
-These tests mock _yield_tool_call to avoid importing the full copilot stack
-which requires Prisma, DB connections, etc.
-"""
-
-import asyncio
-import time
-from typing import Any, cast
-
-import pytest
-
-
-@pytest.mark.asyncio
-async def test_parallel_tool_calls_run_concurrently():
-    """Multiple tool calls should complete in ~max(delays), not sum(delays)."""
-    # Import here to allow module-level mocking if needed
-    from backend.copilot.response_model import (
-        StreamToolInputAvailable,
-        StreamToolOutputAvailable,
-    )
-    from backend.copilot.service import _execute_tool_calls_parallel
-
-    n_tools = 3
-    delay_per_tool = 0.2
-    tool_calls = [
-        {
-            "id": f"call_{i}",
-            "type": "function",
-            "function": {"name": f"tool_{i}", "arguments": "{}"},
-        }
-        for i in range(n_tools)
-    ]
-
-    # Minimal session mock
-    class FakeSession:
-        session_id = "test"
-        user_id = "test"
-
-        def __init__(self):
-            self.messages = []
-
-    original_yield = None
-
-    async def fake_yield(tc_list, idx, sess, lock=None):
-        yield StreamToolInputAvailable(
-            toolCallId=tc_list[idx]["id"],
-            toolName=tc_list[idx]["function"]["name"],
-            input={},
-        )
-        await asyncio.sleep(delay_per_tool)
-        yield StreamToolOutputAvailable(
-            toolCallId=tc_list[idx]["id"],
-            toolName=tc_list[idx]["function"]["name"],
-            output="{}",
-        )
-
-    import backend.copilot.service as svc
-
-    original_yield = svc._yield_tool_call
-    svc._yield_tool_call = fake_yield
-    try:
-        start = time.monotonic()
-        events = []
-        async for event in _execute_tool_calls_parallel(
-            tool_calls, cast(Any, FakeSession())
-        ):
-            events.append(event)
-        elapsed = time.monotonic() - start
-    finally:
-        svc._yield_tool_call = original_yield
-
-    assert len(events) == n_tools * 2
-    # Parallel: should take ~delay, not ~n*delay
-    assert elapsed < delay_per_tool * (
-        n_tools - 0.5
-    ), f"Took {elapsed:.2f}s, expected parallel (~{delay_per_tool}s)"
-
-
-@pytest.mark.asyncio
-async def test_single_tool_call_works():
-    """Single tool call should work identically."""
-    from backend.copilot.response_model import (
-        StreamToolInputAvailable,
-        StreamToolOutputAvailable,
-    )
-    from backend.copilot.service import _execute_tool_calls_parallel
-
-    tool_calls = [
-        {
-            "id": "call_0",
-            "type": "function",
-            "function": {"name": "t", "arguments": "{}"},
-        }
-    ]
-
-    class FakeSession:
-        session_id = "test"
-        user_id = "test"
-
-        def __init__(self):
-            self.messages = []
-
-    async def fake_yield(tc_list, idx, sess, lock=None):
-        yield StreamToolInputAvailable(toolCallId="call_0", toolName="t", input={})
-        yield StreamToolOutputAvailable(toolCallId="call_0", toolName="t", output="{}")
-
-    import backend.copilot.service as svc
-
-    orig = svc._yield_tool_call
-    svc._yield_tool_call = fake_yield
-    try:
-        events = [
-            e
-            async for e in _execute_tool_calls_parallel(
-                tool_calls, cast(Any, FakeSession())
-            )
-        ]
-    finally:
-        svc._yield_tool_call = orig
-
-    assert len(events) == 2
-
-
-@pytest.mark.asyncio
-async def test_retryable_error_propagates():
-    """Retryable errors should be raised after all tools finish."""
-    from backend.copilot.response_model import StreamToolOutputAvailable
-    from backend.copilot.service import _execute_tool_calls_parallel
-
-    tool_calls = [
-        {
-            "id": f"call_{i}",
-            "type": "function",
-            "function": {"name": f"t_{i}", "arguments": "{}"},
-        }
-        for i in range(2)
-    ]
-
-    class FakeSession:
-        session_id = "test"
-        user_id = "test"
-
-        def __init__(self):
-            self.messages = []
-
-    async def fake_yield(tc_list, idx, sess, lock=None):
-        if idx == 1:
-            raise KeyError("bad")
-        from backend.copilot.response_model import StreamToolInputAvailable
-
-        yield StreamToolInputAvailable(
-            toolCallId=tc_list[idx]["id"], toolName="t_0", input={}
-        )
-        await asyncio.sleep(0.05)
-        yield StreamToolOutputAvailable(
-            toolCallId=tc_list[idx]["id"], toolName="t_0", output="{}"
-        )
-
-    import backend.copilot.service as svc
-
-    orig = svc._yield_tool_call
-    svc._yield_tool_call = fake_yield
-    try:
-        events = []
-        with pytest.raises(KeyError):
-            async for event in _execute_tool_calls_parallel(
-                tool_calls, cast(Any, FakeSession())
-            ):
-                events.append(event)
-        # First tool's events should still be yielded
-        assert any(isinstance(e, StreamToolOutputAvailable) for e in events)
-    finally:
-        svc._yield_tool_call = orig
-
-
-@pytest.mark.asyncio
-async def test_session_lock_shared():
-    """All parallel tools should receive the same lock instance."""
-    from backend.copilot.response_model import (
-        StreamToolInputAvailable,
-        StreamToolOutputAvailable,
-    )
-    from backend.copilot.service import _execute_tool_calls_parallel
-
-    tool_calls = [
-        {
-            "id": f"call_{i}",
-            "type": "function",
-            "function": {"name": f"t_{i}", "arguments": "{}"},
-        }
-        for i in range(3)
-    ]
-
-    class FakeSession:
-        session_id = "test"
-        user_id = "test"
-
-        def __init__(self):
-            self.messages = []
-
-    observed_locks = []
-
-    async def fake_yield(tc_list, idx, sess, lock=None):
-        observed_locks.append(lock)
-        yield StreamToolInputAvailable(
-            toolCallId=tc_list[idx]["id"], toolName=f"t_{idx}", input={}
-        )
-        yield StreamToolOutputAvailable(
-            toolCallId=tc_list[idx]["id"], toolName=f"t_{idx}", output="{}"
-        )
-
-    import backend.copilot.service as svc
-
-    orig = svc._yield_tool_call
-    svc._yield_tool_call = fake_yield
-    try:
-        async for _ in _execute_tool_calls_parallel(
-            tool_calls, cast(Any, FakeSession())
-        ):
-            pass
-    finally:
-        svc._yield_tool_call = orig
-
-    assert len(observed_locks) == 3
-    assert observed_locks[0] is observed_locks[1] is observed_locks[2]
-    assert isinstance(observed_locks[0], asyncio.Lock)
-
-
-@pytest.mark.asyncio
-async def test_cancellation_cleans_up():
-    """Generator close should cancel in-flight tasks."""
-    from backend.copilot.response_model import StreamToolInputAvailable
-    from backend.copilot.service import _execute_tool_calls_parallel
-
-    tool_calls = [
-        {
-            "id": f"call_{i}",
-            "type": "function",
-            "function": {"name": f"t_{i}", "arguments": "{}"},
-        }
-        for i in range(2)
-    ]
-
-    class FakeSession:
-        session_id = "test"
-        user_id = "test"
-
-        def __init__(self):
-            self.messages = []
-
-    started = asyncio.Event()
-
-    async def fake_yield(tc_list, idx, sess, lock=None):
-        yield StreamToolInputAvailable(
-            toolCallId=tc_list[idx]["id"], toolName=f"t_{idx}", input={}
-        )
-        started.set()
-        await asyncio.sleep(10)  # simulate long-running
-
-    import backend.copilot.service as svc
-
-    orig = svc._yield_tool_call
-    svc._yield_tool_call = fake_yield
-    try:
-        gen = _execute_tool_calls_parallel(tool_calls, cast(Any, FakeSession()))
-        await gen.__anext__()  # get first event
-        await started.wait()
-        await gen.aclose()  # close generator
-    finally:
-        svc._yield_tool_call = orig
-    # If we get here without hanging, cleanup worked

autogpt_platform/backend/backend/copilot/sdk/query_builder_test.py

@@ -1,221 +0,0 @@
-"""Tests for _format_conversation_context and _build_query_message."""
-
-from datetime import UTC, datetime
-
-import pytest
-
-from backend.copilot.model import ChatMessage, ChatSession
-from backend.copilot.sdk.service import (
-    _build_query_message,
-    _format_conversation_context,
-)
-
-# ---------------------------------------------------------------------------
-# _format_conversation_context
-# ---------------------------------------------------------------------------
-
-
-def test_format_empty_list():
-    assert _format_conversation_context([]) is None
-
-
-def test_format_none_content_messages():
-    msgs = [ChatMessage(role="user", content=None)]
-    assert _format_conversation_context(msgs) is None
-
-
-def test_format_user_message():
-    msgs = [ChatMessage(role="user", content="hello")]
-    result = _format_conversation_context(msgs)
-    assert result is not None
-    assert "User: hello" in result
-    assert result.startswith("<conversation_history>")
-    assert result.endswith("</conversation_history>")
-
-
-def test_format_assistant_text():
-    msgs = [ChatMessage(role="assistant", content="hi there")]
-    result = _format_conversation_context(msgs)
-    assert result is not None
-    assert "You responded: hi there" in result
-
-
-def test_format_assistant_tool_calls():
-    msgs = [
-        ChatMessage(
-            role="assistant",
-            content=None,
-            tool_calls=[{"function": {"name": "search", "arguments": '{"q": "test"}'}}],
-        )
-    ]
-    result = _format_conversation_context(msgs)
-    assert result is not None
-    assert 'You called tool: search({"q": "test"})' in result
-
-
-def test_format_tool_result():
-    msgs = [ChatMessage(role="tool", content='{"result": "ok"}')]
-    result = _format_conversation_context(msgs)
-    assert result is not None
-    assert 'Tool result: {"result": "ok"}' in result
-
-
-def test_format_tool_result_none_content():
-    msgs = [ChatMessage(role="tool", content=None)]
-    result = _format_conversation_context(msgs)
-    assert result is not None
-    assert "Tool result: " in result
-
-
-def test_format_full_conversation():
-    msgs = [
-        ChatMessage(role="user", content="find agents"),
-        ChatMessage(
-            role="assistant",
-            content="I'll search for agents.",
-            tool_calls=[
-                {"function": {"name": "find_agents", "arguments": '{"q": "test"}'}}
-            ],
-        ),
-        ChatMessage(role="tool", content='[{"id": "1", "name": "Agent1"}]'),
-        ChatMessage(role="assistant", content="Found Agent1."),
-    ]
-    result = _format_conversation_context(msgs)
-    assert result is not None
-    assert "User: find agents" in result
-    assert "You responded: I'll search for agents." in result
-    assert "You called tool: find_agents" in result
-    assert "Tool result:" in result
-    assert "You responded: Found Agent1." in result
-
-
-# ---------------------------------------------------------------------------
-# _build_query_message
-# ---------------------------------------------------------------------------
-
-
-def _make_session(messages: list[ChatMessage]) -> ChatSession:
-    """Build a minimal ChatSession with the given messages."""
-    now = datetime.now(UTC)
-    return ChatSession(
-        session_id="test-session",
-        user_id="user-1",
-        messages=messages,
-        title="test",
-        usage=[],
-        started_at=now,
-        updated_at=now,
-    )
-
-
-@pytest.mark.asyncio
-async def test_build_query_resume_up_to_date():
-    """With --resume and transcript covers all messages, return raw message."""
-    session = _make_session(
-        [
-            ChatMessage(role="user", content="hello"),
-            ChatMessage(role="assistant", content="hi"),
-            ChatMessage(role="user", content="what's new?"),
-        ]
-    )
-    result = await _build_query_message(
-        "what's new?",
-        session,
-        use_resume=True,
-        transcript_msg_count=2,
-        session_id="test-session",
-    )
-    # transcript_msg_count == msg_count - 1, so no gap
-    assert result == "what's new?"
-
-
-@pytest.mark.asyncio
-async def test_build_query_resume_stale_transcript():
-    """With --resume and stale transcript, gap context is prepended."""
-    session = _make_session(
-        [
-            ChatMessage(role="user", content="turn 1"),
-            ChatMessage(role="assistant", content="reply 1"),
-            ChatMessage(role="user", content="turn 2"),
-            ChatMessage(role="assistant", content="reply 2"),
-            ChatMessage(role="user", content="turn 3"),
-        ]
-    )
-    result = await _build_query_message(
-        "turn 3",
-        session,
-        use_resume=True,
-        transcript_msg_count=2,
-        session_id="test-session",
-    )
-    assert "<conversation_history>" in result
-    assert "turn 2" in result
-    assert "reply 2" in result
-    assert "Now, the user says:\nturn 3" in result
-
-
-@pytest.mark.asyncio
-async def test_build_query_resume_zero_msg_count():
-    """With --resume but transcript_msg_count=0, return raw message."""
-    session = _make_session(
-        [
-            ChatMessage(role="user", content="hello"),
-            ChatMessage(role="assistant", content="hi"),
-            ChatMessage(role="user", content="new msg"),
-        ]
-    )
-    result = await _build_query_message(
-        "new msg",
-        session,
-        use_resume=True,
-        transcript_msg_count=0,
-        session_id="test-session",
-    )
-    assert result == "new msg"
-
-
-@pytest.mark.asyncio
-async def test_build_query_no_resume_single_message():
-    """Without --resume and only 1 message, return raw message."""
-    session = _make_session([ChatMessage(role="user", content="first")])
-    result = await _build_query_message(
-        "first",
-        session,
-        use_resume=False,
-        transcript_msg_count=0,
-        session_id="test-session",
-    )
-    assert result == "first"
-
-
-@pytest.mark.asyncio
-async def test_build_query_no_resume_multi_message(monkeypatch):
-    """Without --resume and multiple messages, compress and prepend."""
-    session = _make_session(
-        [
-            ChatMessage(role="user", content="older question"),
-            ChatMessage(role="assistant", content="older answer"),
-            ChatMessage(role="user", content="new question"),
-        ]
-    )
-
-    # Mock _compress_conversation_history to return the messages as-is
-    async def _mock_compress(sess):
-        return sess.messages[:-1]
-
-    monkeypatch.setattr(
-        "backend.copilot.sdk.service._compress_conversation_history",
-        _mock_compress,
-    )
-
-    result = await _build_query_message(
-        "new question",
-        session,
-        use_resume=False,
-        transcript_msg_count=0,
-        session_id="test-session",
-    )
-    assert "<conversation_history>" in result
-    assert "older question" in result
-    assert "older answer" in result
-    assert "Now, the user says:\nnew question" in result

autogpt_platform/backend/backend/copilot/sdk/response_adapter.py

@@ -1,376 +0,0 @@
-"""Response adapter for converting Claude Agent SDK messages to Vercel AI SDK format.
-
-This module provides the adapter layer that converts streaming messages from
-the Claude Agent SDK into the Vercel AI SDK UI Stream Protocol format that
-the frontend expects.
-"""
-
-import json
-import logging
-import uuid
-
-from claude_agent_sdk import (
-    AssistantMessage,
-    Message,
-    ResultMessage,
-    SystemMessage,
-    TextBlock,
-    ToolResultBlock,
-    ToolUseBlock,
-    UserMessage,
-)
-
-from backend.copilot.response_model import (
-    StreamBaseResponse,
-    StreamError,
-    StreamFinish,
-    StreamFinishStep,
-    StreamStart,
-    StreamStartStep,
-    StreamTextDelta,
-    StreamTextEnd,
-    StreamTextStart,
-    StreamToolInputAvailable,
-    StreamToolInputStart,
-    StreamToolOutputAvailable,
-)
-
-from .tool_adapter import MCP_TOOL_PREFIX, pop_pending_tool_output
-
-logger = logging.getLogger(__name__)
-
-
-class SDKResponseAdapter:
-    """Adapter for converting Claude Agent SDK messages to Vercel AI SDK format.
-
-    This class maintains state during a streaming session to properly track
-    text blocks, tool calls, and message lifecycle.
-    """
-
-    def __init__(self, message_id: str | None = None, session_id: str | None = None):
-        self.message_id = message_id or str(uuid.uuid4())
-        self.session_id = session_id
-        self.text_block_id = str(uuid.uuid4())
-        self.has_started_text = False
-        self.has_ended_text = False
-        self.current_tool_calls: dict[str, dict[str, str]] = {}
-        self.resolved_tool_calls: set[str] = set()
-        self.task_id: str | None = None
-        self.step_open = False
-
-    def set_task_id(self, task_id: str) -> None:
-        """Set the task ID for reconnection support."""
-        self.task_id = task_id
-
-    @property
-    def has_unresolved_tool_calls(self) -> bool:
-        """True when there are tool calls that haven't received output yet."""
-        return bool(self.current_tool_calls.keys() - self.resolved_tool_calls)
-
-    def convert_message(self, sdk_message: Message) -> list[StreamBaseResponse]:
-        """Convert a single SDK message to Vercel AI SDK format."""
-        responses: list[StreamBaseResponse] = []
-
-        if isinstance(sdk_message, SystemMessage):
-            if sdk_message.subtype == "init":
-                responses.append(
-                    StreamStart(messageId=self.message_id, taskId=self.task_id)
-                )
-                # Open the first step (matches non-SDK: StreamStart then StreamStartStep)
-                responses.append(StreamStartStep())
-                self.step_open = True
-
-        elif isinstance(sdk_message, AssistantMessage):
-            # Flush any SDK built-in tool calls that didn't get a UserMessage
-            # result (e.g. WebSearch, Read handled internally by the CLI).
-            # BUT skip flush when this AssistantMessage is a parallel tool
-            # continuation (contains only ToolUseBlocks) — the prior tools
-            # are still executing concurrently and haven't finished yet.
-            is_tool_only = all(isinstance(b, ToolUseBlock) for b in sdk_message.content)
-            if not is_tool_only:
-                self._flush_unresolved_tool_calls(responses)
-
-            # After tool results, the SDK sends a new AssistantMessage for the
-            # next LLM turn. Open a new step if the previous one was closed.
-            if not self.step_open:
-                responses.append(StreamStartStep())
-                self.step_open = True
-
-            for block in sdk_message.content:
-                if isinstance(block, TextBlock):
-                    if block.text:
-                        self._ensure_text_started(responses)
-                        responses.append(
-                            StreamTextDelta(id=self.text_block_id, delta=block.text)
-                        )
-
-                elif isinstance(block, ToolUseBlock):
-                    self._end_text_if_open(responses)
-
-                    # Strip MCP prefix so frontend sees "find_block"
-                    # instead of "mcp__copilot__find_block".
-                    tool_name = block.name.removeprefix(MCP_TOOL_PREFIX)
-
-                    responses.append(
-                        StreamToolInputStart(toolCallId=block.id, toolName=tool_name)
-                    )
-                    responses.append(
-                        StreamToolInputAvailable(
-                            toolCallId=block.id,
-                            toolName=tool_name,
-                            input=block.input,
-                        )
-                    )
-                    self.current_tool_calls[block.id] = {"name": tool_name}
-
-        elif isinstance(sdk_message, UserMessage):
-            # UserMessage carries tool results back from tool execution.
-            content = sdk_message.content
-            blocks = content if isinstance(content, list) else []
-            resolved_in_blocks: set[str] = set()
-
-            sid = (self.session_id or "?")[:12]
-            parent_id_preview = getattr(sdk_message, "parent_tool_use_id", None)
-            logger.info(
-                "[SDK] [%s] UserMessage: %d blocks, content_type=%s, "
-                "parent_tool_use_id=%s",
-                sid,
-                len(blocks),
-                type(content).__name__,
-                parent_id_preview[:12] if parent_id_preview else "None",
-            )
-
-            for block in blocks:
-                if isinstance(block, ToolResultBlock) and block.tool_use_id:
-                    # Skip if already resolved (e.g. by flush) — the real
-                    # result supersedes the empty flush, but re-emitting
-                    # would confuse the frontend's state machine.
-                    if block.tool_use_id in self.resolved_tool_calls:
-                        continue
-                    tool_info = self.current_tool_calls.get(block.tool_use_id, {})
-                    tool_name = tool_info.get("name", "unknown")
-
-                    # Prefer the stashed full output over the SDK's
-                    # (potentially truncated) ToolResultBlock content.
-                    # The SDK truncates large results, writing them to disk,
-                    # which breaks frontend widget parsing.
-                    output = pop_pending_tool_output(tool_name) or (
-                        _extract_tool_output(block.content)
-                    )
-
-                    responses.append(
-                        StreamToolOutputAvailable(
-                            toolCallId=block.tool_use_id,
-                            toolName=tool_name,
-                            output=output,
-                            success=not (block.is_error or False),
-                        )
-                    )
-                    resolved_in_blocks.add(block.tool_use_id)
-
-            # Handle SDK built-in tool results carried via parent_tool_use_id
-            # instead of (or in addition to) ToolResultBlock content.
-            parent_id = sdk_message.parent_tool_use_id
-            if (
-                parent_id
-                and parent_id not in resolved_in_blocks
-                and parent_id not in self.resolved_tool_calls
-            ):
-                tool_info = self.current_tool_calls.get(parent_id, {})
-                tool_name = tool_info.get("name", "unknown")
-
-                # Try stashed output first (from PostToolUse hook),
-                # then tool_use_result dict, then string content.
-                output = pop_pending_tool_output(tool_name)
-                if not output:
-                    tur = sdk_message.tool_use_result
-                    if tur is not None:
-                        output = _extract_tool_use_result(tur)
-                if not output and isinstance(content, str) and content.strip():
-                    output = content.strip()
-
-                if output:
-                    responses.append(
-                        StreamToolOutputAvailable(
-                            toolCallId=parent_id,
-                            toolName=tool_name,
-                            output=output,
-                            success=True,
-                        )
-                    )
-                    resolved_in_blocks.add(parent_id)
-
-            self.resolved_tool_calls.update(resolved_in_blocks)
-
-            # Close the current step after tool results — the next
-            # AssistantMessage will open a new step for the continuation.
-            if self.step_open:
-                responses.append(StreamFinishStep())
-                self.step_open = False
-
-        elif isinstance(sdk_message, ResultMessage):
-            self._flush_unresolved_tool_calls(responses)
-            self._end_text_if_open(responses)
-            # Close the step before finishing.
-            if self.step_open:
-                responses.append(StreamFinishStep())
-                self.step_open = False
-
-            if sdk_message.subtype == "success":
-                responses.append(StreamFinish())
-            elif sdk_message.subtype in ("error", "error_during_execution"):
-                error_msg = sdk_message.result or "Unknown error"
-                responses.append(
-                    StreamError(errorText=str(error_msg), code="sdk_error")
-                )
-                responses.append(StreamFinish())
-            else:
-                logger.warning(
-                    f"Unexpected ResultMessage subtype: {sdk_message.subtype}"
-                )
-                responses.append(StreamFinish())
-
-        else:
-            logger.debug(f"Unhandled SDK message type: {type(sdk_message).__name__}")
-
-        return responses
-
-    def _ensure_text_started(self, responses: list[StreamBaseResponse]) -> None:
-        """Start (or restart) a text block if needed."""
-        if not self.has_started_text or self.has_ended_text:
-            if self.has_ended_text:
-                self.text_block_id = str(uuid.uuid4())
-                self.has_ended_text = False
-            responses.append(StreamTextStart(id=self.text_block_id))
-            self.has_started_text = True
-
-    def _end_text_if_open(self, responses: list[StreamBaseResponse]) -> None:
-        """End the current text block if one is open."""
-        if self.has_started_text and not self.has_ended_text:
-            responses.append(StreamTextEnd(id=self.text_block_id))
-            self.has_ended_text = True
-
-    def _flush_unresolved_tool_calls(self, responses: list[StreamBaseResponse]) -> None:
-        """Emit outputs for tool calls that didn't receive a UserMessage result.
-
-        SDK built-in tools (WebSearch, Read, etc.) may be executed by the CLI
-        internally without surfacing a separate ``UserMessage`` with
-        ``ToolResultBlock`` content.  The ``PostToolUse`` hook stashes their
-        output, which we pop and emit here before the next ``AssistantMessage``
-        starts.
-        """
-        unresolved = [
-            (tid, info.get("name", "unknown"))
-            for tid, info in self.current_tool_calls.items()
-            if tid not in self.resolved_tool_calls
-        ]
-        sid = (self.session_id or "?")[:12]
-        if not unresolved:
-            logger.info(
-                "[SDK] [%s] Flush called but all %d tool(s) already resolved",
-                sid,
-                len(self.current_tool_calls),
-            )
-            return
-        logger.info(
-            "[SDK] [%s] Flushing %d unresolved tool call(s): %s",
-            sid,
-            len(unresolved),
-            ", ".join(f"{name}({tid[:12]})" for tid, name in unresolved),
-        )
-
-        flushed = False
-        for tool_id, tool_name in unresolved:
-            output = pop_pending_tool_output(tool_name)
-            if output is not None:
-                responses.append(
-                    StreamToolOutputAvailable(
-                        toolCallId=tool_id,
-                        toolName=tool_name,
-                        output=output,
-                        success=True,
-                    )
-                )
-                self.resolved_tool_calls.add(tool_id)
-                flushed = True
-                logger.info(
-                    "[SDK] [%s] Flushed stashed output for %s " "(call %s, %d chars)",
-                    sid,
-                    tool_name,
-                    tool_id[:12],
-                    len(output),
-                )
-            else:
-                # No output available — emit an empty output so the frontend
-                # transitions the tool from input-available to output-available
-                # (stops the spinner).
-                responses.append(
-                    StreamToolOutputAvailable(
-                        toolCallId=tool_id,
-                        toolName=tool_name,
-                        output="",
-                        success=True,
-                    )
-                )
-                self.resolved_tool_calls.add(tool_id)
-                flushed = True
-                logger.warning(
-                    "[SDK] [%s] Flushed EMPTY output for unresolved tool %s "
-                    "(call %s) — stash was empty (likely SDK hook race "
-                    "condition: PostToolUse hook hadn't completed before "
-                    "flush was triggered)",
-                    sid,
-                    tool_name,
-                    tool_id[:12],
-                )
-
-        if flushed and self.step_open:
-            responses.append(StreamFinishStep())
-            self.step_open = False
-
-
-def _extract_tool_output(content: str | list[dict[str, str]] | None) -> str:
-    """Extract a string output from a ToolResultBlock's content field."""
-    if isinstance(content, str):
-        return content
-    if isinstance(content, list):
-        parts = [item.get("text", "") for item in content if item.get("type") == "text"]
-        if parts:
-            return "".join(parts)
-        try:
-            return json.dumps(content)
-        except (TypeError, ValueError):
-            return str(content)
-    if content is None:
-        return ""
-    try:
-        return json.dumps(content)
-    except (TypeError, ValueError):
-        return str(content)
-
-
-def _extract_tool_use_result(result: object) -> str:
-    """Extract a string from a UserMessage's ``tool_use_result`` dict.
-
-    SDK built-in tools may store their result in ``tool_use_result``
-    instead of (or in addition to) ``ToolResultBlock`` content blocks.
-    """
-    if isinstance(result, str):
-        return result
-    if isinstance(result, dict):
-        # Try common result keys
-        for key in ("content", "text", "output", "stdout", "result"):
-            val = result.get(key)
-            if isinstance(val, str) and val:
-                return val
-        # Fall back to JSON serialization of the whole dict
-        try:
-            return json.dumps(result)
-        except (TypeError, ValueError):
-            return str(result)
-    if result is None:
-        return ""
-    try:
-        return json.dumps(result)
-    except (TypeError, ValueError):
-        return str(result)

autogpt_platform/backend/backend/copilot/sdk/sdk_compat_test.py

@@ -1,194 +0,0 @@
-"""SDK compatibility tests — verify the claude-agent-sdk public API surface we depend on.
-
-Instead of pinning to a narrow version range, these tests verify that the
-installed SDK exposes every class, function, attribute, and method the copilot
-integration relies on.  If an SDK upgrade removes or renames something these
-tests will catch it immediately.
-"""
-
-import inspect
-
-import pytest
-
-# ---------------------------------------------------------------------------
-# Public types & factories
-# ---------------------------------------------------------------------------
-
-
-def test_sdk_exports_client_and_options():
-    from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient
-
-    assert inspect.isclass(ClaudeSDKClient)
-    assert inspect.isclass(ClaudeAgentOptions)
-
-
-def test_sdk_exports_message_types():
-    from claude_agent_sdk import (
-        AssistantMessage,
-        Message,
-        ResultMessage,
-        SystemMessage,
-        UserMessage,
-    )
-
-    for cls in (AssistantMessage, ResultMessage, SystemMessage, UserMessage):
-        assert inspect.isclass(cls), f"{cls.__name__} is not a class"
-    # Message is a Union type alias, just verify it's importable
-    assert Message is not None
-
-
-def test_sdk_exports_content_block_types():
-    from claude_agent_sdk import TextBlock, ToolResultBlock, ToolUseBlock
-
-    for cls in (TextBlock, ToolResultBlock, ToolUseBlock):
-        assert inspect.isclass(cls), f"{cls.__name__} is not a class"
-
-
-def test_sdk_exports_mcp_helpers():
-    from claude_agent_sdk import create_sdk_mcp_server, tool
-
-    assert callable(create_sdk_mcp_server)
-    assert callable(tool)
-
-
-# ---------------------------------------------------------------------------
-# ClaudeSDKClient interface
-# ---------------------------------------------------------------------------
-
-
-def test_client_has_required_methods():
-    from claude_agent_sdk import ClaudeSDKClient
-
-    required = ["connect", "disconnect", "query", "receive_messages"]
-    for name in required:
-        attr = getattr(ClaudeSDKClient, name, None)
-        assert attr is not None, f"ClaudeSDKClient.{name} missing"
-        assert callable(attr), f"ClaudeSDKClient.{name} is not callable"
-
-
-def test_client_supports_async_context_manager():
-    from claude_agent_sdk import ClaudeSDKClient
-
-    assert hasattr(ClaudeSDKClient, "__aenter__")
-    assert hasattr(ClaudeSDKClient, "__aexit__")
-
-
-# ---------------------------------------------------------------------------
-# ClaudeAgentOptions fields
-# ---------------------------------------------------------------------------
-
-
-def test_agent_options_accepts_required_fields():
-    """Verify ClaudeAgentOptions accepts all kwargs our code passes."""
-    from claude_agent_sdk import ClaudeAgentOptions
-
-    opts = ClaudeAgentOptions(
-        system_prompt="test",
-        cwd="/tmp",
-    )
-    assert opts.system_prompt == "test"
-    assert opts.cwd == "/tmp"
-
-
-def test_agent_options_accepts_all_our_fields():
-    """Comprehensive check of every field we use in service.py."""
-    from claude_agent_sdk import ClaudeAgentOptions
-
-    fields_we_use = [
-        "system_prompt",
-        "mcp_servers",
-        "allowed_tools",
-        "disallowed_tools",
-        "hooks",
-        "cwd",
-        "model",
-        "env",
-        "resume",
-        "max_buffer_size",
-    ]
-    sig = inspect.signature(ClaudeAgentOptions)
-    for field in fields_we_use:
-        assert field in sig.parameters, (
-            f"ClaudeAgentOptions no longer accepts '{field}' — "
-            f"available params: {list(sig.parameters.keys())}"
-        )
-
-
-# ---------------------------------------------------------------------------
-# Message attributes
-# ---------------------------------------------------------------------------
-
-
-def test_assistant_message_has_content_and_model():
-    from claude_agent_sdk import AssistantMessage, TextBlock
-
-    msg = AssistantMessage(content=[TextBlock(text="hi")], model="test")
-    assert hasattr(msg, "content")
-    assert hasattr(msg, "model")
-
-
-def test_result_message_has_required_attrs():
-    from claude_agent_sdk import ResultMessage
-
-    msg = ResultMessage(
-        subtype="success",
-        duration_ms=100,
-        duration_api_ms=50,
-        is_error=False,
-        num_turns=1,
-        session_id="s1",
-    )
-    assert msg.subtype == "success"
-    assert hasattr(msg, "result")
-
-
-def test_system_message_has_subtype_and_data():
-    from claude_agent_sdk import SystemMessage
-
-    msg = SystemMessage(subtype="init", data={})
-    assert msg.subtype == "init"
-    assert msg.data == {}
-
-
-def test_user_message_has_parent_tool_use_id():
-    from claude_agent_sdk import UserMessage
-
-    msg = UserMessage(content="test")
-    assert hasattr(msg, "parent_tool_use_id")
-    assert hasattr(msg, "tool_use_result")
-
-
-def test_tool_use_block_has_id_name_input():
-    from claude_agent_sdk import ToolUseBlock
-
-    block = ToolUseBlock(id="t1", name="test", input={"key": "val"})
-    assert block.id == "t1"
-    assert block.name == "test"
-    assert block.input == {"key": "val"}
-
-
-def test_tool_result_block_has_required_attrs():
-    from claude_agent_sdk import ToolResultBlock
-
-    block = ToolResultBlock(tool_use_id="t1", content="result")
-    assert block.tool_use_id == "t1"
-    assert block.content == "result"
-    assert hasattr(block, "is_error")
-
-
-# ---------------------------------------------------------------------------
-# Hook types
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize(
-    "hook_event",
-    ["PreToolUse", "PostToolUse", "Stop"],
-)
-def test_sdk_exports_hook_event_type(hook_event: str):
-    """Verify HookEvent literal includes the events our security_hooks use."""
-    from claude_agent_sdk.types import HookEvent
-
-    # HookEvent is a Literal type — check that our events are valid values.
-    # We can't easily inspect Literal at runtime, so just verify the type exists.
-    assert HookEvent is not None

autogpt_platform/backend/backend/copilot/sdk/security_hooks_test.py

@@ -1,335 +0,0 @@
-"""Tests for SDK security hooks — workspace paths, tool access, and deny messages.
-
-These are pure unit tests with no external dependencies (no SDK, no DB, no server).
-They validate that the security hooks correctly block unauthorized paths,
-tool access, and dangerous input patterns.
-"""
-
-import os
-
-import pytest
-
-from .security_hooks import _validate_tool_access, _validate_user_isolation
-from .service import _is_tool_error_or_denial
-
-SDK_CWD = "/tmp/copilot-abc123"
-
-
-def _sdk_available() -> bool:
-    try:
-        import claude_agent_sdk  # noqa: F401
-
-        return True
-    except ImportError:
-        return False
-
-
-def _is_denied(result: dict) -> bool:
-    hook = result.get("hookSpecificOutput", {})
-    return hook.get("permissionDecision") == "deny"
-
-
-def _reason(result: dict) -> str:
-    return result.get("hookSpecificOutput", {}).get("permissionDecisionReason", "")
-
-
-# -- Blocked tools -----------------------------------------------------------
-
-
-def test_blocked_tools_denied():
-    for tool in ("bash", "shell", "exec", "terminal", "command"):
-        result = _validate_tool_access(tool, {})
-        assert _is_denied(result), f"{tool} should be blocked"
-
-
-def test_unknown_tool_allowed():
-    result = _validate_tool_access("SomeCustomTool", {})
-    assert result == {}
-
-
-# -- Workspace-scoped tools --------------------------------------------------
-
-
-def test_read_within_workspace_allowed():
-    result = _validate_tool_access(
-        "Read", {"file_path": f"{SDK_CWD}/file.txt"}, sdk_cwd=SDK_CWD
-    )
-    assert result == {}
-
-
-def test_write_within_workspace_allowed():
-    result = _validate_tool_access(
-        "Write", {"file_path": f"{SDK_CWD}/output.json"}, sdk_cwd=SDK_CWD
-    )
-    assert result == {}
-
-
-def test_edit_within_workspace_allowed():
-    result = _validate_tool_access(
-        "Edit", {"file_path": f"{SDK_CWD}/src/main.py"}, sdk_cwd=SDK_CWD
-    )
-    assert result == {}
-
-
-def test_glob_within_workspace_allowed():
-    result = _validate_tool_access("Glob", {"path": f"{SDK_CWD}/src"}, sdk_cwd=SDK_CWD)
-    assert result == {}
-
-
-def test_grep_within_workspace_allowed():
-    result = _validate_tool_access("Grep", {"path": f"{SDK_CWD}/src"}, sdk_cwd=SDK_CWD)
-    assert result == {}
-
-
-def test_read_outside_workspace_denied():
-    result = _validate_tool_access(
-        "Read", {"file_path": "/etc/passwd"}, sdk_cwd=SDK_CWD
-    )
-    assert _is_denied(result)
-
-
-def test_write_outside_workspace_denied():
-    result = _validate_tool_access(
-        "Write", {"file_path": "/home/user/secrets.txt"}, sdk_cwd=SDK_CWD
-    )
-    assert _is_denied(result)
-
-
-def test_traversal_attack_denied():
-    result = _validate_tool_access(
-        "Read",
-        {"file_path": f"{SDK_CWD}/../../etc/passwd"},
-        sdk_cwd=SDK_CWD,
-    )
-    assert _is_denied(result)
-
-
-def test_no_path_allowed():
-    """Glob/Grep without a path argument defaults to cwd — should pass."""
-    result = _validate_tool_access("Glob", {}, sdk_cwd=SDK_CWD)
-    assert result == {}
-
-
-def test_read_no_cwd_denies_absolute():
-    """If no sdk_cwd is set, absolute paths are denied."""
-    result = _validate_tool_access("Read", {"file_path": "/tmp/anything"})
-    assert _is_denied(result)
-
-
-# -- Tool-results directory --------------------------------------------------
-
-
-def test_read_tool_results_allowed():
-    home = os.path.expanduser("~")
-    path = f"{home}/.claude/projects/-tmp-copilot-abc123/tool-results/12345.txt"
-    result = _validate_tool_access("Read", {"file_path": path}, sdk_cwd=SDK_CWD)
-    assert result == {}
-
-
-def test_read_claude_projects_without_tool_results_denied():
-    home = os.path.expanduser("~")
-    path = f"{home}/.claude/projects/-tmp-copilot-abc123/settings.json"
-    result = _validate_tool_access("Read", {"file_path": path}, sdk_cwd=SDK_CWD)
-    assert _is_denied(result)
-
-
-# -- Built-in Bash is blocked (use bash_exec MCP tool instead) ---------------
-
-
-def test_bash_builtin_always_blocked():
-    """SDK built-in Bash is blocked — bash_exec MCP tool with bubblewrap is used instead."""
-    result = _validate_tool_access("Bash", {"command": "echo hello"}, sdk_cwd=SDK_CWD)
-    assert _is_denied(result)
-
-
-# -- Dangerous patterns ------------------------------------------------------
-
-
-def test_dangerous_pattern_blocked():
-    result = _validate_tool_access("SomeTool", {"cmd": "sudo rm -rf /"})
-    assert _is_denied(result)
-
-
-def test_subprocess_pattern_blocked():
-    result = _validate_tool_access("SomeTool", {"code": "subprocess.run(...)"})
-    assert _is_denied(result)
-
-
-# -- User isolation ----------------------------------------------------------
-
-
-def test_workspace_path_traversal_blocked():
-    result = _validate_user_isolation(
-        "workspace_read", {"path": "../../../etc/shadow"}, user_id="user-1"
-    )
-    assert _is_denied(result)
-
-
-def test_workspace_absolute_path_allowed():
-    """Workspace 'path' is a cloud storage key — leading '/' is normal."""
-    result = _validate_user_isolation(
-        "workspace_read", {"path": "/ASEAN/report.md"}, user_id="user-1"
-    )
-    assert result == {}
-
-
-def test_workspace_normal_path_allowed():
-    result = _validate_user_isolation(
-        "workspace_read", {"path": "src/main.py"}, user_id="user-1"
-    )
-    assert result == {}
-
-
-def test_non_workspace_tool_passes_isolation():
-    result = _validate_user_isolation(
-        "find_agent", {"query": "email"}, user_id="user-1"
-    )
-    assert result == {}
-
-
-# -- Deny message quality ----------------------------------------------------
-
-
-def test_blocked_tool_message_clarity():
-    """Deny messages must include [SECURITY] and 'cannot be bypassed'."""
-    reason = _reason(_validate_tool_access("bash", {}))
-    assert "[SECURITY]" in reason
-    assert "cannot be bypassed" in reason
-
-
-def test_bash_builtin_blocked_message_clarity():
-    reason = _reason(_validate_tool_access("Bash", {"command": "echo hello"}))
-    assert "[SECURITY]" in reason
-    assert "cannot be bypassed" in reason
-
-
-# -- Task sub-agent hooks (require SDK) --------------------------------------
-
-
-@pytest.fixture()
-def _hooks():
-    """Create security hooks and return the PreToolUse handler."""
-    from .security_hooks import create_security_hooks
-
-    hooks = create_security_hooks(user_id="u1", sdk_cwd=SDK_CWD, max_subtasks=2)
-    pre = hooks["PreToolUse"][0].hooks[0]
-    return pre
-
-
-@pytest.mark.skipif(not _sdk_available(), reason="claude_agent_sdk not installed")
-@pytest.mark.asyncio
-async def test_task_background_blocked(_hooks):
-    """Task with run_in_background=true must be denied."""
-    result = await _hooks(
-        {"tool_name": "Task", "tool_input": {"run_in_background": True, "prompt": "x"}},
-        tool_use_id=None,
-        context={},
-    )
-    assert _is_denied(result)
-    assert "foreground" in _reason(result).lower()
-
-
-@pytest.mark.skipif(not _sdk_available(), reason="claude_agent_sdk not installed")
-@pytest.mark.asyncio
-async def test_task_foreground_allowed(_hooks):
-    """Task without run_in_background should be allowed."""
-    result = await _hooks(
-        {"tool_name": "Task", "tool_input": {"prompt": "do stuff"}},
-        tool_use_id=None,
-        context={},
-    )
-    assert not _is_denied(result)
-
-
-@pytest.mark.skipif(not _sdk_available(), reason="claude_agent_sdk not installed")
-@pytest.mark.asyncio
-async def test_task_limit_enforced(_hooks):
-    """Task spawns beyond max_subtasks should be denied."""
-    # First two should pass
-    for _ in range(2):
-        result = await _hooks(
-            {"tool_name": "Task", "tool_input": {"prompt": "ok"}},
-            tool_use_id=None,
-            context={},
-        )
-        assert not _is_denied(result)
-
-    # Third should be denied (limit=2)
-    result = await _hooks(
-        {"tool_name": "Task", "tool_input": {"prompt": "over limit"}},
-        tool_use_id=None,
-        context={},
-    )
-    assert _is_denied(result)
-    assert "Maximum" in _reason(result)
-
-
-# -- _is_tool_error_or_denial ------------------------------------------------
-
-
-class TestIsToolErrorOrDenial:
-    def test_none_content(self):
-        assert _is_tool_error_or_denial(None) is False
-
-    def test_empty_content(self):
-        assert _is_tool_error_or_denial("") is False
-
-    def test_benign_output(self):
-        assert _is_tool_error_or_denial("All good, no issues.") is False
-
-    def test_security_marker(self):
-        assert _is_tool_error_or_denial("[SECURITY] Tool access blocked") is True
-
-    def test_cannot_be_bypassed(self):
-        assert _is_tool_error_or_denial("This restriction cannot be bypassed.") is True
-
-    def test_not_allowed(self):
-        assert _is_tool_error_or_denial("Operation not allowed in sandbox") is True
-
-    def test_background_task_denial(self):
-        assert (
-            _is_tool_error_or_denial(
-                "Background task execution is not supported. "
-                "Run tasks in the foreground instead."
-            )
-            is True
-        )
-
-    def test_subtask_limit_denial(self):
-        assert (
-            _is_tool_error_or_denial(
-                "Maximum 2 sub-tasks per session. Please continue in the main conversation."
-            )
-            is True
-        )
-
-    def test_denied_marker(self):
-        assert (
-            _is_tool_error_or_denial("Access denied: insufficient privileges") is True
-        )
-
-    def test_blocked_marker(self):
-        assert _is_tool_error_or_denial("Request blocked by security policy") is True
-
-    def test_failed_marker(self):
-        assert _is_tool_error_or_denial("Failed to execute tool: timeout") is True
-
-    def test_mcp_iserror(self):
-        assert _is_tool_error_or_denial('{"isError": true, "content": []}') is True
-
-    def test_benign_error_in_value(self):
-        """Content like '0 errors found' should not trigger — 'error' was removed."""
-        assert _is_tool_error_or_denial("0 errors found") is False
-
-    def test_benign_permission_field(self):
-        """Schema descriptions mentioning 'permission' should not trigger."""
-        assert (
-            _is_tool_error_or_denial(
-                '{"fields": [{"name": "permission_level", "type": "int"}]}'
-            )
-            is False
-        )
-
-    def test_benign_not_found_in_listing(self):
-        """File listing containing 'not found' in filenames should not trigger."""
-        assert _is_tool_error_or_denial("readme.md\nfile-not-found-handler.py") is False

autogpt_platform/backend/backend/copilot/tools/create_agent_test.py

@@ -1,142 +0,0 @@
-"""Tests for CreateAgentTool response types."""
-
-from unittest.mock import AsyncMock, patch
-
-import pytest
-
-from backend.copilot.tools.create_agent import CreateAgentTool
-from backend.copilot.tools.models import (
-    ClarificationNeededResponse,
-    ErrorResponse,
-    SuggestedGoalResponse,
-)
-
-from ._test_data import make_session
-
-_TEST_USER_ID = "test-user-create-agent"
-
-
-@pytest.fixture
-def tool():
-    return CreateAgentTool()
-
-
-@pytest.fixture
-def session():
-    return make_session(_TEST_USER_ID)
-
-
-@pytest.mark.asyncio
-async def test_missing_description_returns_error(tool, session):
-    """Missing description returns ErrorResponse."""
-    result = await tool._execute(user_id=_TEST_USER_ID, session=session, description="")
-    assert isinstance(result, ErrorResponse)
-    assert result.error == "Missing description parameter"
-
-
-@pytest.mark.asyncio
-async def test_vague_goal_returns_suggested_goal_response(tool, session):
-    """vague_goal decomposition result returns SuggestedGoalResponse, not ErrorResponse."""
-    vague_result = {
-        "type": "vague_goal",
-        "suggested_goal": "Monitor Twitter mentions for a specific keyword and send a daily digest email",
-    }
-
-    with (
-        patch(
-            "backend.copilot.tools.create_agent.get_all_relevant_agents_for_generation",
-            new_callable=AsyncMock,
-            return_value=[],
-        ),
-        patch(
-            "backend.copilot.tools.create_agent.decompose_goal",
-            new_callable=AsyncMock,
-            return_value=vague_result,
-        ),
-    ):
-        result = await tool._execute(
-            user_id=_TEST_USER_ID,
-            session=session,
-            description="monitor social media",
-        )
-
-    assert isinstance(result, SuggestedGoalResponse)
-    assert result.goal_type == "vague"
-    assert result.suggested_goal == vague_result["suggested_goal"]
-    assert result.original_goal == "monitor social media"
-    assert result.reason == "The goal needs more specific details"
-    assert not isinstance(result, ErrorResponse)
-
-
-@pytest.mark.asyncio
-async def test_unachievable_goal_returns_suggested_goal_response(tool, session):
-    """unachievable_goal decomposition result returns SuggestedGoalResponse, not ErrorResponse."""
-    unachievable_result = {
-        "type": "unachievable_goal",
-        "suggested_goal": "Summarize the latest news articles on a topic and send them by email",
-        "reason": "There are no blocks for mind-reading.",
-    }
-
-    with (
-        patch(
-            "backend.copilot.tools.create_agent.get_all_relevant_agents_for_generation",
-            new_callable=AsyncMock,
-            return_value=[],
-        ),
-        patch(
-            "backend.copilot.tools.create_agent.decompose_goal",
-            new_callable=AsyncMock,
-            return_value=unachievable_result,
-        ),
-    ):
-        result = await tool._execute(
-            user_id=_TEST_USER_ID,
-            session=session,
-            description="read my mind",
-        )
-
-    assert isinstance(result, SuggestedGoalResponse)
-    assert result.goal_type == "unachievable"
-    assert result.suggested_goal == unachievable_result["suggested_goal"]
-    assert result.original_goal == "read my mind"
-    assert result.reason == unachievable_result["reason"]
-    assert not isinstance(result, ErrorResponse)
-
-
-@pytest.mark.asyncio
-async def test_clarifying_questions_returns_clarification_needed_response(
-    tool, session
-):
-    """clarifying_questions decomposition result returns ClarificationNeededResponse."""
-    clarifying_result = {
-        "type": "clarifying_questions",
-        "questions": [
-            {
-                "question": "What platform should be monitored?",
-                "keyword": "platform",
-                "example": "Twitter, Reddit",
-            }
-        ],
-    }
-
-    with (
-        patch(
-            "backend.copilot.tools.create_agent.get_all_relevant_agents_for_generation",
-            new_callable=AsyncMock,
-            return_value=[],
-        ),
-        patch(
-            "backend.copilot.tools.create_agent.decompose_goal",
-            new_callable=AsyncMock,
-            return_value=clarifying_result,
-        ),
-    ):
-        result = await tool._execute(
-            user_id=_TEST_USER_ID,
-            session=session,
-            description="monitor social media and alert me",
-        )
-
-    assert isinstance(result, ClarificationNeededResponse)
-    assert len(result.questions) == 1
-    assert result.questions[0].keyword == "platform"

autogpt_platform/backend/backend/copilot/tools/workspace_files.py

@@ -1,789 +0,0 @@
-"""CoPilot tools for workspace file operations."""
-
-import base64
-import logging
-import os
-from typing import Any, Optional
-
-from pydantic import BaseModel
-
-from backend.copilot.model import ChatSession
-from backend.copilot.tools.sandbox import make_session_path
-from backend.data.db_accessors import workspace_db
-from backend.util.settings import Config
-from backend.util.virus_scanner import scan_content_safe
-from backend.util.workspace import WorkspaceManager
-
-from .base import BaseTool
-from .models import ErrorResponse, ResponseType, ToolResponseBase
-
-logger = logging.getLogger(__name__)
-
-
-def _resolve_write_content(
-    content_text: str | None,
-    content_b64: str | None,
-    source_path: str | None,
-    session_id: str,
-) -> bytes | ErrorResponse:
-    """Resolve file content from exactly one of three input sources.
-
-    Returns the raw bytes on success, or an ``ErrorResponse`` on validation
-    failure (wrong number of sources, invalid path, file not found, etc.).
-    """
-    # Normalise empty strings to None so counting and dispatch stay in sync.
-    if content_text is not None and content_text == "":
-        content_text = None
-    if content_b64 is not None and content_b64 == "":
-        content_b64 = None
-    if source_path is not None and source_path == "":
-        source_path = None
-
-    sources_provided = sum(
-        x is not None for x in [content_text, content_b64, source_path]
-    )
-    if sources_provided == 0:
-        return ErrorResponse(
-            message="Please provide one of: content, content_base64, or source_path",
-            session_id=session_id,
-        )
-    if sources_provided > 1:
-        return ErrorResponse(
-            message="Provide only one of: content, content_base64, or source_path",
-            session_id=session_id,
-        )
-
-    if source_path is not None:
-        validated = _validate_ephemeral_path(
-            source_path, param_name="source_path", session_id=session_id
-        )
-        if isinstance(validated, ErrorResponse):
-            return validated
-        try:
-            with open(validated, "rb") as f:
-                return f.read()
-        except FileNotFoundError:
-            return ErrorResponse(
-                message=f"Source file not found: {source_path}",
-                session_id=session_id,
-            )
-        except Exception as e:
-            return ErrorResponse(
-                message=f"Failed to read source file: {e}",
-                session_id=session_id,
-            )
-
-    if content_b64 is not None:
-        try:
-            return base64.b64decode(content_b64)
-        except Exception:
-            return ErrorResponse(
-                message=(
-                    "Invalid base64 encoding in content_base64. "
-                    "Please encode the file content with standard base64, "
-                    "or use the 'content' parameter for plain text, "
-                    "or 'source_path' to copy from the working directory."
-                ),
-                session_id=session_id,
-            )
-
-    assert content_text is not None
-    return content_text.encode("utf-8")
-
-
-def _validate_ephemeral_path(
-    path: str, *, param_name: str, session_id: str
-) -> ErrorResponse | str:
-    """Validate that *path* is inside the session's ephemeral directory.
-
-    Uses the session-specific directory (``make_session_path(session_id)``)
-    rather than the bare prefix, so ``/tmp/copilot-evil/...`` is rejected.
-
-    Returns the resolved real path on success, or an ``ErrorResponse`` when the
-    path escapes the session directory.
-    """
-    session_dir = os.path.realpath(make_session_path(session_id)) + os.sep
-    real = os.path.realpath(path)
-    if not real.startswith(session_dir):
-        return ErrorResponse(
-            message=(
-                f"{param_name} must be within the ephemeral working "
-                f"directory ({make_session_path(session_id)})"
-            ),
-            session_id=session_id,
-        )
-    return real
-
-
-_TEXT_MIME_PREFIXES = (
-    "text/",
-    "application/json",
-    "application/xml",
-    "application/javascript",
-    "application/x-python",
-    "application/x-sh",
-)
-
-_IMAGE_MIME_TYPES = {"image/png", "image/jpeg", "image/gif", "image/webp"}
-
-
-def _is_text_mime(mime_type: str) -> bool:
-    return any(mime_type.startswith(t) for t in _TEXT_MIME_PREFIXES)
-
-
-async def _get_manager(user_id: str, session_id: str) -> WorkspaceManager:
-    """Create a session-scoped WorkspaceManager."""
-    workspace = await workspace_db().get_or_create_workspace(user_id)
-    return WorkspaceManager(user_id, workspace.id, session_id)
-
-
-async def _resolve_file(
-    manager: WorkspaceManager,
-    file_id: str | None,
-    path: str | None,
-    session_id: str,
-) -> tuple[str, Any] | ErrorResponse:
-    """Resolve a file by file_id or path.
-
-    Returns ``(target_file_id, file_info)`` on success, or an
-    ``ErrorResponse`` if the file was not found.
-    """
-    if file_id:
-        file_info = await manager.get_file_info(file_id)
-        if file_info is None:
-            return ErrorResponse(
-                message=f"File not found: {file_id}", session_id=session_id
-            )
-        return file_id, file_info
-
-    assert path is not None
-    file_info = await manager.get_file_info_by_path(path)
-    if file_info is None:
-        return ErrorResponse(
-            message=f"File not found at path: {path}", session_id=session_id
-        )
-    return file_info.id, file_info
-
-
-class WorkspaceFileInfoData(BaseModel):
-    """Data model for workspace file information (not a response itself)."""
-
-    file_id: str
-    name: str
-    path: str
-    mime_type: str
-    size_bytes: int
-
-
-class WorkspaceFileListResponse(ToolResponseBase):
-    """Response containing list of workspace files."""
-
-    type: ResponseType = ResponseType.WORKSPACE_FILE_LIST
-    files: list[WorkspaceFileInfoData]
-    total_count: int
-
-
-class WorkspaceFileContentResponse(ToolResponseBase):
-    """Response containing workspace file content (legacy, for small text files)."""
-
-    type: ResponseType = ResponseType.WORKSPACE_FILE_CONTENT
-    file_id: str
-    name: str
-    path: str
-    mime_type: str
-    content_base64: str
-
-
-class WorkspaceFileMetadataResponse(ToolResponseBase):
-    """Response containing workspace file metadata and download URL (prevents context bloat)."""
-
-    type: ResponseType = ResponseType.WORKSPACE_FILE_METADATA
-    file_id: str
-    name: str
-    path: str
-    mime_type: str
-    size_bytes: int
-    download_url: str
-    preview: str | None = None  # First 500 chars for text files
-
-
-class WorkspaceWriteResponse(ToolResponseBase):
-    """Response after writing a file to workspace."""
-
-    type: ResponseType = ResponseType.WORKSPACE_FILE_WRITTEN
-    file_id: str
-    name: str
-    path: str
-    size_bytes: int
-    source: str | None = None  # "content", "base64", or "copied from <path>"
-    content_preview: str | None = None  # First 200 chars for text files
-
-
-class WorkspaceDeleteResponse(ToolResponseBase):
-    """Response after deleting a file from workspace."""
-
-    type: ResponseType = ResponseType.WORKSPACE_FILE_DELETED
-    file_id: str
-    success: bool
-
-
-class ListWorkspaceFilesTool(BaseTool):
-    """Tool for listing files in user's workspace."""
-
-    @property
-    def name(self) -> str:
-        return "list_workspace_files"
-
-    @property
-    def description(self) -> str:
-        return (
-            "List files in the user's persistent workspace (cloud storage). "
-            "These files survive across sessions. "
-            "For ephemeral session files, use the SDK Read/Glob tools instead. "
-            "Returns file names, paths, sizes, and metadata. "
-            "Optionally filter by path prefix."
-        )
-
-    @property
-    def parameters(self) -> dict[str, Any]:
-        return {
-            "type": "object",
-            "properties": {
-                "path_prefix": {
-                    "type": "string",
-                    "description": (
-                        "Optional path prefix to filter files "
-                        "(e.g., '/documents/' to list only files in documents folder). "
-                        "By default, only files from the current session are listed."
-                    ),
-                },
-                "limit": {
-                    "type": "integer",
-                    "description": "Maximum number of files to return (default 50, max 100)",
-                    "minimum": 1,
-                    "maximum": 100,
-                },
-                "include_all_sessions": {
-                    "type": "boolean",
-                    "description": (
-                        "If true, list files from all sessions. "
-                        "Default is false (only current session's files)."
-                    ),
-                },
-            },
-            "required": [],
-        }
-
-    @property
-    def requires_auth(self) -> bool:
-        return True
-
-    async def _execute(
-        self,
-        user_id: str | None,
-        session: ChatSession,
-        **kwargs,
-    ) -> ToolResponseBase:
-        session_id = session.session_id
-        if not user_id:
-            return ErrorResponse(
-                message="Authentication required", session_id=session_id
-            )
-
-        path_prefix: Optional[str] = kwargs.get("path_prefix")
-        limit = min(kwargs.get("limit", 50), 100)
-        include_all_sessions: bool = kwargs.get("include_all_sessions", False)
-
-        try:
-            manager = await _get_manager(user_id, session_id)
-            files = await manager.list_files(
-                path=path_prefix, limit=limit, include_all_sessions=include_all_sessions
-            )
-            total = await manager.get_file_count(
-                path=path_prefix, include_all_sessions=include_all_sessions
-            )
-            file_infos = [
-                WorkspaceFileInfoData(
-                    file_id=f.id,
-                    name=f.name,
-                    path=f.path,
-                    mime_type=f.mime_type,
-                    size_bytes=f.size_bytes,
-                )
-                for f in files
-            ]
-            scope = "all sessions" if include_all_sessions else "current session"
-            total_size = sum(f.size_bytes for f in file_infos)
-
-            # Build a human-readable summary so the agent can relay details.
-            lines = [f"Found {len(files)} file(s) in workspace ({scope}):"]
-            for f in file_infos:
-                lines.append(f"  - {f.path} ({f.size_bytes:,} bytes, {f.mime_type})")
-            if total > len(files):
-                lines.append(f"  ... and {total - len(files)} more")
-            lines.append(f"Total size: {total_size:,} bytes")
-
-            return WorkspaceFileListResponse(
-                files=file_infos,
-                total_count=total,
-                message="\n".join(lines),
-                session_id=session_id,
-            )
-        except Exception as e:
-            logger.error(f"Error listing workspace files: {e}", exc_info=True)
-            return ErrorResponse(
-                message=f"Failed to list workspace files: {e}",
-                error=str(e),
-                session_id=session_id,
-            )
-
-
-class ReadWorkspaceFileTool(BaseTool):
-    """Tool for reading file content from workspace."""
-
-    MAX_INLINE_SIZE_BYTES = 32 * 1024  # 32KB
-    PREVIEW_SIZE = 500
-
-    @property
-    def name(self) -> str:
-        return "read_workspace_file"
-
-    @property
-    def description(self) -> str:
-        return (
-            "Read a file from the user's persistent workspace (cloud storage). "
-            "These files survive across sessions. "
-            "For ephemeral session files, use the SDK Read tool instead. "
-            "Specify either file_id or path to identify the file. "
-            "For small text files, returns content directly. "
-            "For large or binary files, returns metadata and a download URL. "
-            "Optionally use 'save_to_path' to copy the file to the ephemeral "
-            "working directory for processing with bash_exec or SDK tools. "
-            "Paths are scoped to the current session by default. "
-            "Use /sessions/<session_id>/... for cross-session access."
-        )
-
-    @property
-    def parameters(self) -> dict[str, Any]:
-        return {
-            "type": "object",
-            "properties": {
-                "file_id": {
-                    "type": "string",
-                    "description": "The file's unique ID (from list_workspace_files)",
-                },
-                "path": {
-                    "type": "string",
-                    "description": (
-                        "The virtual file path (e.g., '/documents/report.pdf'). "
-                        "Scoped to current session by default."
-                    ),
-                },
-                "save_to_path": {
-                    "type": "string",
-                    "description": (
-                        "If provided, save the file to this path in the ephemeral "
-                        "working directory (e.g., '/tmp/copilot-.../data.csv') "
-                        "so it can be processed with bash_exec or SDK tools. "
-                        "The file content is still returned in the response."
-                    ),
-                },
-                "force_download_url": {
-                    "type": "boolean",
-                    "description": (
-                        "If true, always return metadata+URL instead of inline content. "
-                        "Default is false (auto-selects based on file size/type)."
-                    ),
-                },
-            },
-            "required": [],  # At least one must be provided
-        }
-
-    @property
-    def requires_auth(self) -> bool:
-        return True
-
-    async def _execute(
-        self,
-        user_id: str | None,
-        session: ChatSession,
-        **kwargs,
-    ) -> ToolResponseBase:
-        session_id = session.session_id
-        if not user_id:
-            return ErrorResponse(
-                message="Authentication required", session_id=session_id
-            )
-
-        file_id: Optional[str] = kwargs.get("file_id")
-        path: Optional[str] = kwargs.get("path")
-        save_to_path: Optional[str] = kwargs.get("save_to_path")
-        force_download_url: bool = kwargs.get("force_download_url", False)
-
-        if not file_id and not path:
-            return ErrorResponse(
-                message="Please provide either file_id or path", session_id=session_id
-            )
-
-        # Validate and resolve save_to_path (use sanitized real path).
-        if save_to_path:
-            validated_save = _validate_ephemeral_path(
-                save_to_path, param_name="save_to_path", session_id=session_id
-            )
-            if isinstance(validated_save, ErrorResponse):
-                return validated_save
-            save_to_path = validated_save
-
-        try:
-            manager = await _get_manager(user_id, session_id)
-            resolved = await _resolve_file(manager, file_id, path, session_id)
-            if isinstance(resolved, ErrorResponse):
-                return resolved
-            target_file_id, file_info = resolved
-
-            # If save_to_path, read + save; cache bytes for possible inline reuse.
-            cached_content: bytes | None = None
-            if save_to_path:
-                cached_content = await manager.read_file_by_id(target_file_id)
-                dir_path = os.path.dirname(save_to_path)
-                if dir_path:
-                    os.makedirs(dir_path, exist_ok=True)
-                with open(save_to_path, "wb") as f:
-                    f.write(cached_content)
-
-            is_small = file_info.size_bytes <= self.MAX_INLINE_SIZE_BYTES
-            is_text = _is_text_mime(file_info.mime_type)
-            is_image = file_info.mime_type in _IMAGE_MIME_TYPES
-
-            # Inline content for small text/image files
-            if is_small and (is_text or is_image) and not force_download_url:
-                content = cached_content or await manager.read_file_by_id(
-                    target_file_id
-                )
-                msg = (
-                    f"Read {file_info.name} from workspace:{file_info.path} "
-                    f"({file_info.size_bytes:,} bytes, {file_info.mime_type})"
-                )
-                if save_to_path:
-                    msg += f" — also saved to {save_to_path}"
-                return WorkspaceFileContentResponse(
-                    file_id=file_info.id,
-                    name=file_info.name,
-                    path=file_info.path,
-                    mime_type=file_info.mime_type,
-                    content_base64=base64.b64encode(content).decode("utf-8"),
-                    message=msg,
-                    session_id=session_id,
-                )
-
-            # Metadata + download URL for large/binary files
-            preview: str | None = None
-            if is_text:
-                try:
-                    raw = cached_content or await manager.read_file_by_id(
-                        target_file_id
-                    )
-                    preview = raw[: self.PREVIEW_SIZE].decode("utf-8", errors="replace")
-                    if len(raw) > self.PREVIEW_SIZE:
-                        preview += "..."
-                except Exception:
-                    pass
-
-            msg = (
-                f"File: {file_info.name} at workspace:{file_info.path} "
-                f"({file_info.size_bytes:,} bytes, {file_info.mime_type})"
-            )
-            if save_to_path:
-                msg += f" — saved to {save_to_path}"
-            else:
-                msg += (
-                    " — use read_workspace_file with this file_id to retrieve content"
-                )
-            return WorkspaceFileMetadataResponse(
-                file_id=file_info.id,
-                name=file_info.name,
-                path=file_info.path,
-                mime_type=file_info.mime_type,
-                size_bytes=file_info.size_bytes,
-                download_url=f"workspace://{target_file_id}",
-                preview=preview,
-                message=msg,
-                session_id=session_id,
-            )
-        except FileNotFoundError as e:
-            return ErrorResponse(message=str(e), session_id=session_id)
-        except Exception as e:
-            logger.error(f"Error reading workspace file: {e}", exc_info=True)
-            return ErrorResponse(
-                message=f"Failed to read workspace file: {e}",
-                error=str(e),
-                session_id=session_id,
-            )
-
-
-class WriteWorkspaceFileTool(BaseTool):
-    """Tool for writing files to workspace."""
-
-    @property
-    def name(self) -> str:
-        return "write_workspace_file"
-
-    @property
-    def description(self) -> str:
-        return (
-            "Write or create a file in the user's persistent workspace (cloud storage). "
-            "These files survive across sessions. "
-            "For ephemeral session files, use the SDK Write tool instead. "
-            "Provide content as plain text via 'content', OR base64-encoded via "
-            "'content_base64', OR copy a file from the ephemeral working directory "
-            "via 'source_path'. Exactly one of these three is required. "
-            f"Maximum file size is {Config().max_file_size_mb}MB. "
-            "Files are saved to the current session's folder by default. "
-            "Use /sessions/<session_id>/... for cross-session access."
-        )
-
-    @property
-    def parameters(self) -> dict[str, Any]:
-        return {
-            "type": "object",
-            "properties": {
-                "filename": {
-                    "type": "string",
-                    "description": "Name for the file (e.g., 'report.pdf')",
-                },
-                "content": {
-                    "type": "string",
-                    "description": (
-                        "Plain text content to write. Use this for text files "
-                        "(code, configs, documents, etc.). "
-                        "Mutually exclusive with content_base64 and source_path."
-                    ),
-                },
-                "content_base64": {
-                    "type": "string",
-                    "description": (
-                        "Base64-encoded file content. Use this for binary files "
-                        "(images, PDFs, etc.). "
-                        "Mutually exclusive with content and source_path."
-                    ),
-                },
-                "source_path": {
-                    "type": "string",
-                    "description": (
-                        "Path to a file in the ephemeral working directory to "
-                        "copy to workspace (e.g., '/tmp/copilot-.../output.csv'). "
-                        "Use this to persist files created by bash_exec or SDK Write. "
-                        "Mutually exclusive with content and content_base64."
-                    ),
-                },
-                "path": {
-                    "type": "string",
-                    "description": (
-                        "Optional virtual path where to save the file "
-                        "(e.g., '/documents/report.pdf'). "
-                        "Defaults to '/{filename}'. Scoped to current session."
-                    ),
-                },
-                "mime_type": {
-                    "type": "string",
-                    "description": (
-                        "Optional MIME type of the file. "
-                        "Auto-detected from filename if not provided."
-                    ),
-                },
-                "overwrite": {
-                    "type": "boolean",
-                    "description": "Whether to overwrite if file exists at path (default: false)",
-                },
-            },
-            "required": ["filename"],
-        }
-
-    @property
-    def requires_auth(self) -> bool:
-        return True
-
-    async def _execute(
-        self,
-        user_id: str | None,
-        session: ChatSession,
-        **kwargs,
-    ) -> ToolResponseBase:
-        session_id = session.session_id
-        if not user_id:
-            return ErrorResponse(
-                message="Authentication required", session_id=session_id
-            )
-
-        filename: str = kwargs.get("filename", "")
-        if not filename:
-            return ErrorResponse(
-                message="Please provide a filename", session_id=session_id
-            )
-
-        source_path_arg: str | None = kwargs.get("source_path")
-        content_text: str | None = kwargs.get("content")
-        content_b64: str | None = kwargs.get("content_base64")
-
-        resolved = _resolve_write_content(
-            content_text,
-            content_b64,
-            source_path_arg,
-            session_id,
-        )
-        if isinstance(resolved, ErrorResponse):
-            return resolved
-        content: bytes = resolved
-
-        max_size = Config().max_file_size_mb * 1024 * 1024
-        if len(content) > max_size:
-            return ErrorResponse(
-                message=f"File too large. Maximum size is {Config().max_file_size_mb}MB",
-                session_id=session_id,
-            )
-
-        try:
-            await scan_content_safe(content, filename=filename)
-            manager = await _get_manager(user_id, session_id)
-            rec = await manager.write_file(
-                content=content,
-                filename=filename,
-                path=kwargs.get("path"),
-                mime_type=kwargs.get("mime_type"),
-                overwrite=kwargs.get("overwrite", False),
-            )
-
-            # Build informative source label and message.
-            if source_path_arg:
-                source = f"copied from {source_path_arg}"
-                msg = (
-                    f"Copied {source_path_arg} → workspace:{rec.path} "
-                    f"({rec.size_bytes:,} bytes)"
-                )
-            elif content_b64:
-                source = "base64"
-                msg = (
-                    f"Wrote {rec.name} to workspace ({rec.size_bytes:,} bytes, "
-                    f"decoded from base64)"
-                )
-            else:
-                source = "content"
-                msg = f"Wrote {rec.name} to workspace ({rec.size_bytes:,} bytes)"
-
-            # Include a short preview for text content.
-            preview: str | None = None
-            if _is_text_mime(rec.mime_type):
-                try:
-                    preview = content[:200].decode("utf-8", errors="replace")
-                    if len(content) > 200:
-                        preview += "..."
-                except Exception:
-                    pass
-
-            return WorkspaceWriteResponse(
-                file_id=rec.id,
-                name=rec.name,
-                path=rec.path,
-                size_bytes=rec.size_bytes,
-                source=source,
-                content_preview=preview,
-                message=msg,
-                session_id=session_id,
-            )
-        except ValueError as e:
-            return ErrorResponse(message=str(e), session_id=session_id)
-        except Exception as e:
-            logger.error(f"Error writing workspace file: {e}", exc_info=True)
-            return ErrorResponse(
-                message=f"Failed to write workspace file: {e}",
-                error=str(e),
-                session_id=session_id,
-            )
-
-
-class DeleteWorkspaceFileTool(BaseTool):
-    """Tool for deleting files from workspace."""
-
-    @property
-    def name(self) -> str:
-        return "delete_workspace_file"
-
-    @property
-    def description(self) -> str:
-        return (
-            "Delete a file from the user's persistent workspace (cloud storage). "
-            "Specify either file_id or path to identify the file. "
-            "Paths are scoped to the current session by default. "
-            "Use /sessions/<session_id>/... for cross-session access."
-        )
-
-    @property
-    def parameters(self) -> dict[str, Any]:
-        return {
-            "type": "object",
-            "properties": {
-                "file_id": {
-                    "type": "string",
-                    "description": "The file's unique ID (from list_workspace_files)",
-                },
-                "path": {
-                    "type": "string",
-                    "description": (
-                        "The virtual file path (e.g., '/documents/report.pdf'). "
-                        "Scoped to current session by default."
-                    ),
-                },
-            },
-            "required": [],  # At least one must be provided
-        }
-
-    @property
-    def requires_auth(self) -> bool:
-        return True
-
-    async def _execute(
-        self,
-        user_id: str | None,
-        session: ChatSession,
-        **kwargs,
-    ) -> ToolResponseBase:
-        session_id = session.session_id
-        if not user_id:
-            return ErrorResponse(
-                message="Authentication required", session_id=session_id
-            )
-
-        file_id: Optional[str] = kwargs.get("file_id")
-        path: Optional[str] = kwargs.get("path")
-        if not file_id and not path:
-            return ErrorResponse(
-                message="Please provide either file_id or path", session_id=session_id
-            )
-
-        try:
-            manager = await _get_manager(user_id, session_id)
-            resolved = await _resolve_file(manager, file_id, path, session_id)
-            if isinstance(resolved, ErrorResponse):
-                return resolved
-            target_file_id, file_info = resolved
-
-            if not await manager.delete_file(target_file_id):
-                return ErrorResponse(
-                    message=f"File not found: {target_file_id}", session_id=session_id
-                )
-            return WorkspaceDeleteResponse(
-                file_id=target_file_id,
-                success=True,
-                message=(
-                    f"Deleted {file_info.name} from workspace:{file_info.path} "
-                    f"({file_info.size_bytes:,} bytes)"
-                ),
-                session_id=session_id,
-            )
-        except Exception as e:
-            logger.error(f"Error deleting workspace file: {e}", exc_info=True)
-            return ErrorResponse(
-                message=f"Failed to delete workspace file: {e}",
-                error=str(e),
-                session_id=session_id,
-            )

autogpt_platform/backend/backend/copilot/tools/workspace_files_test.py

@@ -1,267 +0,0 @@
-"""Tests for workspace file tool helpers and path validation."""
-
-import base64
-import os
-
-import pytest
-
-from backend.copilot.tools._test_data import make_session, setup_test_data
-from backend.copilot.tools.workspace_files import (
-    DeleteWorkspaceFileTool,
-    ListWorkspaceFilesTool,
-    ReadWorkspaceFileTool,
-    WorkspaceDeleteResponse,
-    WorkspaceFileListResponse,
-    WorkspaceWriteResponse,
-    WriteWorkspaceFileTool,
-    _resolve_write_content,
-    _validate_ephemeral_path,
-)
-
-# Re-export so pytest discovers the session-scoped fixture
-setup_test_data = setup_test_data
-
-# We need to mock make_session_path to return a known temp dir for tests.
-# The real one uses WORKSPACE_PREFIX = "/tmp/copilot-"
-
-
-@pytest.fixture
-def ephemeral_dir(tmp_path, monkeypatch):
-    """Create a temp dir that acts as the ephemeral session directory."""
-    session_dir = tmp_path / "copilot-test-session"
-    session_dir.mkdir()
-
-    monkeypatch.setattr(
-        "backend.copilot.tools.workspace_files.make_session_path",
-        lambda session_id: str(session_dir),
-    )
-    return session_dir
-
-
-# ---------------------------------------------------------------------------
-# _validate_ephemeral_path
-# ---------------------------------------------------------------------------
-
-
-class TestValidateEphemeralPath:
-    def test_valid_path(self, ephemeral_dir):
-        target = ephemeral_dir / "file.txt"
-        target.touch()
-        result = _validate_ephemeral_path(
-            str(target), param_name="test", session_id="s1"
-        )
-        assert isinstance(result, str)
-        assert result == os.path.realpath(str(target))
-
-    def test_path_traversal_rejected(self, ephemeral_dir):
-        evil_path = str(ephemeral_dir / ".." / "etc" / "passwd")
-        result = _validate_ephemeral_path(evil_path, param_name="test", session_id="s1")
-        # Should return ErrorResponse
-        from backend.copilot.tools.models import ErrorResponse
-
-        assert isinstance(result, ErrorResponse)
-
-    def test_different_session_rejected(self, ephemeral_dir, tmp_path):
-        other_dir = tmp_path / "copilot-evil-session"
-        other_dir.mkdir()
-        target = other_dir / "steal.txt"
-        target.touch()
-        result = _validate_ephemeral_path(
-            str(target), param_name="test", session_id="s1"
-        )
-        from backend.copilot.tools.models import ErrorResponse
-
-        assert isinstance(result, ErrorResponse)
-
-    def test_symlink_escape_rejected(self, ephemeral_dir, tmp_path):
-        """Symlink inside session dir pointing outside should be rejected."""
-        outside_file = tmp_path / "secret.txt"
-        outside_file.write_text("secret")
-        symlink = ephemeral_dir / "link.txt"
-        symlink.symlink_to(outside_file)
-        result = _validate_ephemeral_path(
-            str(symlink), param_name="test", session_id="s1"
-        )
-        from backend.copilot.tools.models import ErrorResponse
-
-        assert isinstance(result, ErrorResponse)
-
-    def test_nested_path_valid(self, ephemeral_dir):
-        nested = ephemeral_dir / "subdir" / "deep"
-        nested.mkdir(parents=True)
-        target = nested / "data.csv"
-        target.touch()
-        result = _validate_ephemeral_path(
-            str(target), param_name="test", session_id="s1"
-        )
-        assert isinstance(result, str)
-
-
-# ---------------------------------------------------------------------------
-# _resolve_write_content
-# ---------------------------------------------------------------------------
-
-
-class TestResolveWriteContent:
-    def test_no_sources_returns_error(self):
-        from backend.copilot.tools.models import ErrorResponse
-
-        result = _resolve_write_content(None, None, None, "s1")
-        assert isinstance(result, ErrorResponse)
-
-    def test_multiple_sources_returns_error(self):
-        from backend.copilot.tools.models import ErrorResponse
-
-        result = _resolve_write_content("text", "b64data", None, "s1")
-        assert isinstance(result, ErrorResponse)
-
-    def test_plain_text_content(self):
-        result = _resolve_write_content("hello world", None, None, "s1")
-        assert result == b"hello world"
-
-    def test_base64_content(self):
-        raw = b"binary data"
-        b64 = base64.b64encode(raw).decode()
-        result = _resolve_write_content(None, b64, None, "s1")
-        assert result == raw
-
-    def test_invalid_base64_returns_error(self):
-        from backend.copilot.tools.models import ErrorResponse
-
-        result = _resolve_write_content(None, "not-valid-b64!!!", None, "s1")
-        assert isinstance(result, ErrorResponse)
-        assert "base64" in result.message.lower()
-
-    def test_source_path(self, ephemeral_dir):
-        target = ephemeral_dir / "input.txt"
-        target.write_bytes(b"file content")
-        result = _resolve_write_content(None, None, str(target), "s1")
-        assert result == b"file content"
-
-    def test_source_path_not_found(self, ephemeral_dir):
-        from backend.copilot.tools.models import ErrorResponse
-
-        missing = str(ephemeral_dir / "nope.txt")
-        result = _resolve_write_content(None, None, missing, "s1")
-        assert isinstance(result, ErrorResponse)
-
-    def test_source_path_outside_ephemeral(self, ephemeral_dir, tmp_path):
-        from backend.copilot.tools.models import ErrorResponse
-
-        outside = tmp_path / "outside.txt"
-        outside.write_text("nope")
-        result = _resolve_write_content(None, None, str(outside), "s1")
-        assert isinstance(result, ErrorResponse)
-
-    def test_empty_string_sources_treated_as_none(self):
-        from backend.copilot.tools.models import ErrorResponse
-
-        # All empty strings → same as no sources
-        result = _resolve_write_content("", "", "", "s1")
-        assert isinstance(result, ErrorResponse)
-
-    def test_empty_string_source_path_with_text(self):
-        # source_path="" should be normalised to None, so only content counts
-        result = _resolve_write_content("hello", "", "", "s1")
-        assert result == b"hello"
-
-
-# ---------------------------------------------------------------------------
-# E2E: workspace file tool round-trip (write → list → read → delete)
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.asyncio(loop_scope="session")
-async def test_workspace_file_round_trip(setup_test_data):
-    """E2E: write a file, list it, read it back (with save_to_path), then delete it."""
-    user = setup_test_data["user"]
-    session = make_session(user.id)
-    session_id = session.session_id
-
-    # ---- Write ----
-    write_tool = WriteWorkspaceFileTool()
-    write_resp = await write_tool._execute(
-        user_id=user.id,
-        session=session,
-        filename="test_round_trip.txt",
-        content="Hello from e2e test!",
-    )
-    assert isinstance(write_resp, WorkspaceWriteResponse), write_resp.message
-    file_id = write_resp.file_id
-
-    # ---- List ----
-    list_tool = ListWorkspaceFilesTool()
-    list_resp = await list_tool._execute(user_id=user.id, session=session)
-    assert isinstance(list_resp, WorkspaceFileListResponse), list_resp.message
-    assert any(f.file_id == file_id for f in list_resp.files)
-
-    # ---- Read (inline) ----
-    read_tool = ReadWorkspaceFileTool()
-    read_resp = await read_tool._execute(
-        user_id=user.id, session=session, file_id=file_id
-    )
-    from backend.copilot.tools.workspace_files import WorkspaceFileContentResponse
-
-    assert isinstance(read_resp, WorkspaceFileContentResponse), read_resp.message
-    decoded = base64.b64decode(read_resp.content_base64).decode()
-    assert decoded == "Hello from e2e test!"
-
-    # ---- Read with save_to_path ----
-    from backend.copilot.tools.sandbox import make_session_path
-
-    ephemeral_dir = make_session_path(session_id)
-    os.makedirs(ephemeral_dir, exist_ok=True)
-    save_path = os.path.join(ephemeral_dir, "saved_copy.txt")
-
-    read_resp2 = await read_tool._execute(
-        user_id=user.id, session=session, file_id=file_id, save_to_path=save_path
-    )
-    assert not isinstance(read_resp2, type(None))
-    assert os.path.exists(save_path)
-    with open(save_path) as f:
-        assert f.read() == "Hello from e2e test!"
-
-    # ---- Delete ----
-    delete_tool = DeleteWorkspaceFileTool()
-    del_resp = await delete_tool._execute(
-        user_id=user.id, session=session, file_id=file_id
-    )
-    assert isinstance(del_resp, WorkspaceDeleteResponse), del_resp.message
-    assert del_resp.success is True
-
-    # Verify file is gone
-    list_resp2 = await list_tool._execute(user_id=user.id, session=session)
-    assert isinstance(list_resp2, WorkspaceFileListResponse)
-    assert not any(f.file_id == file_id for f in list_resp2.files)
-
-
-@pytest.mark.asyncio(loop_scope="session")
-async def test_write_workspace_file_source_path(setup_test_data):
-    """E2E: write a file from ephemeral source_path to workspace."""
-    user = setup_test_data["user"]
-    session = make_session(user.id)
-    session_id = session.session_id
-
-    # Create a file in the ephemeral dir
-    from backend.copilot.tools.sandbox import make_session_path
-
-    ephemeral_dir = make_session_path(session_id)
-    os.makedirs(ephemeral_dir, exist_ok=True)
-    source = os.path.join(ephemeral_dir, "generated_output.csv")
-    with open(source, "w") as f:
-        f.write("col1,col2\n1,2\n")
-
-    write_tool = WriteWorkspaceFileTool()
-    write_resp = await write_tool._execute(
-        user_id=user.id,
-        session=session,
-        filename="output.csv",
-        source_path=source,
-    )
-    assert isinstance(write_resp, WorkspaceWriteResponse), write_resp.message
-
-    # Clean up
-    delete_tool = DeleteWorkspaceFileTool()
-    await delete_tool._execute(
-        user_id=user.id, session=session, file_id=write_resp.file_id
-    )

autogpt_platform/backend/backend/data/db_accessors.py

@@ -1,118 +0,0 @@
-from backend.data import db
-
-
-def chat_db():
-    if db.is_connected():
-        from backend.copilot import db as _chat_db
-
-        chat_db = _chat_db
-    else:
-        from backend.util.clients import get_database_manager_async_client
-
-        chat_db = get_database_manager_async_client()
-
-    return chat_db
-
-
-def graph_db():
-    if db.is_connected():
-        from backend.data import graph as _graph_db
-
-        graph_db = _graph_db
-    else:
-        from backend.util.clients import get_database_manager_async_client
-
-        graph_db = get_database_manager_async_client()
-
-    return graph_db
-
-
-def library_db():
-    if db.is_connected():
-        from backend.api.features.library import db as _library_db
-
-        library_db = _library_db
-    else:
-        from backend.util.clients import get_database_manager_async_client
-
-        library_db = get_database_manager_async_client()
-
-    return library_db
-
-
-def store_db():
-    if db.is_connected():
-        from backend.api.features.store import db as _store_db
-
-        store_db = _store_db
-    else:
-        from backend.util.clients import get_database_manager_async_client
-
-        store_db = get_database_manager_async_client()
-
-    return store_db
-
-
-def search():
-    if db.is_connected():
-        from backend.api.features.store import hybrid_search as _search
-
-        search = _search
-    else:
-        from backend.util.clients import get_database_manager_async_client
-
-        search = get_database_manager_async_client()
-
-    return search
-
-
-def execution_db():
-    if db.is_connected():
-        from backend.data import execution as _execution_db
-
-        execution_db = _execution_db
-    else:
-        from backend.util.clients import get_database_manager_async_client
-
-        execution_db = get_database_manager_async_client()
-
-    return execution_db
-
-
-def user_db():
-    if db.is_connected():
-        from backend.data import user as _user_db
-
-        user_db = _user_db
-    else:
-        from backend.util.clients import get_database_manager_async_client
-
-        user_db = get_database_manager_async_client()
-
-    return user_db
-
-
-def understanding_db():
-    if db.is_connected():
-        from backend.data import understanding as _understanding_db
-
-        understanding_db = _understanding_db
-    else:
-        from backend.util.clients import get_database_manager_async_client
-
-        understanding_db = get_database_manager_async_client()
-
-    return understanding_db
-
-
-def workspace_db():
-    if db.is_connected():
-        from backend.data import workspace as _workspace_db
-
-        workspace_db = _workspace_db
-    else:
-        from backend.util.clients import get_database_manager_async_client
-
-        workspace_db = get_database_manager_async_client()
-
-    return workspace_db

autogpt_platform/backend/backend/data/graph.py

@@ -1147,14 +1147,14 @@ async def get_graph(
     return GraphModel.from_db(graph, for_export)
 
 
-async def get_store_listed_graphs(graph_ids: list[str]) -> dict[str, GraphModel]:
+async def get_store_listed_graphs(*graph_ids: str) -> dict[str, GraphModel]:
     """Batch-fetch multiple store-listed graphs by their IDs.
 
     Only returns graphs that have approved store listings (publicly available).
     Does not require permission checks since store-listed graphs are public.
 
     Args:
-        graph_ids: List of graph IDs to fetch
+        *graph_ids: Variable number of graph IDs to fetch
 
     Returns:
         Dict mapping graph_id to GraphModel for graphs with approved store listings

autogpt_platform/backend/backend/data/workspace.py

@@ -8,7 +8,6 @@ import logging
 from datetime import datetime, timezone
 from typing import Optional
 
-import pydantic
 from prisma.models import UserWorkspace, UserWorkspaceFile
 from prisma.types import UserWorkspaceFileWhereInput
 
@@ -17,61 +16,7 @@ from backend.util.json import SafeJson
 logger = logging.getLogger(__name__)
 
 
-class Workspace(pydantic.BaseModel):
-    """Pydantic model for UserWorkspace, safe for RPC transport."""
-
-    id: str
-    user_id: str
-    created_at: datetime
-    updated_at: datetime
-
-    @staticmethod
-    def from_db(workspace: "UserWorkspace") -> "Workspace":
-        return Workspace(
-            id=workspace.id,
-            user_id=workspace.userId,
-            created_at=workspace.createdAt,
-            updated_at=workspace.updatedAt,
-        )
-
-
-class WorkspaceFile(pydantic.BaseModel):
-    """Pydantic model for UserWorkspaceFile, safe for RPC transport."""
-
-    id: str
-    workspace_id: str
-    created_at: datetime
-    updated_at: datetime
-    name: str
-    path: str
-    storage_path: str
-    mime_type: str
-    size_bytes: int
-    checksum: Optional[str] = None
-    is_deleted: bool = False
-    deleted_at: Optional[datetime] = None
-    metadata: dict = pydantic.Field(default_factory=dict)
-
-    @staticmethod
-    def from_db(file: "UserWorkspaceFile") -> "WorkspaceFile":
-        return WorkspaceFile(
-            id=file.id,
-            workspace_id=file.workspaceId,
-            created_at=file.createdAt,
-            updated_at=file.updatedAt,
-            name=file.name,
-            path=file.path,
-            storage_path=file.storagePath,
-            mime_type=file.mimeType,
-            size_bytes=file.sizeBytes,
-            checksum=file.checksum,
-            is_deleted=file.isDeleted,
-            deleted_at=file.deletedAt,
-            metadata=file.metadata if isinstance(file.metadata, dict) else {},
-        )
-
-
-async def get_or_create_workspace(user_id: str) -> Workspace:
+async def get_or_create_workspace(user_id: str) -> UserWorkspace:
     """
     Get user's workspace, creating one if it doesn't exist.
 
@@ -82,7 +27,7 @@ async def get_or_create_workspace(user_id: str) -> Workspace:
         user_id: The user's ID
 
     Returns:
-        Workspace instance
+        UserWorkspace instance
     """
     workspace = await UserWorkspace.prisma().upsert(
         where={"userId": user_id},
@@ -92,10 +37,10 @@ async def get_or_create_workspace(user_id: str) -> Workspace:
         },
     )
 
-    return Workspace.from_db(workspace)
+    return workspace
 
 
-async def get_workspace(user_id: str) -> Optional[Workspace]:
+async def get_workspace(user_id: str) -> Optional[UserWorkspace]:
     """
     Get user's workspace if it exists.
 
@@ -103,10 +48,9 @@ async def get_workspace(user_id: str) -> Optional[Workspace]:
         user_id: The user's ID
 
     Returns:
-        Workspace instance or None
+        UserWorkspace instance or None
     """
-    workspace = await UserWorkspace.prisma().find_unique(where={"userId": user_id})
-    return Workspace.from_db(workspace) if workspace else None
+    return await UserWorkspace.prisma().find_unique(where={"userId": user_id})
 
 
 async def create_workspace_file(
@@ -119,7 +63,7 @@ async def create_workspace_file(
     size_bytes: int,
     checksum: Optional[str] = None,
     metadata: Optional[dict] = None,
-) -> WorkspaceFile:
+) -> UserWorkspaceFile:
     """
     Create a new workspace file record.
 
@@ -135,7 +79,7 @@ async def create_workspace_file(
         metadata: Optional additional metadata
 
     Returns:
-        Created WorkspaceFile instance
+        Created UserWorkspaceFile instance
     """
     # Normalize path to start with /
     if not path.startswith("/"):
@@ -159,37 +103,34 @@ async def create_workspace_file(
         f"Created workspace file {file.id} at path {path} "
         f"in workspace {workspace_id}"
     )
-    return WorkspaceFile.from_db(file)
+    return file
 
 
 async def get_workspace_file(
     file_id: str,
-    workspace_id: str,
-) -> Optional[WorkspaceFile]:
+    workspace_id: Optional[str] = None,
+) -> Optional[UserWorkspaceFile]:
     """
     Get a workspace file by ID.
 
     Args:
         file_id: The file ID
-        workspace_id: Workspace ID for scoping (required)
+        workspace_id: Optional workspace ID for validation
 
     Returns:
-        WorkspaceFile instance or None
+        UserWorkspaceFile instance or None
     """
-    where_clause: UserWorkspaceFileWhereInput = {
-        "id": file_id,
-        "isDeleted": False,
-        "workspaceId": workspace_id,
-    }
+    where_clause: dict = {"id": file_id, "isDeleted": False}
+    if workspace_id:
+        where_clause["workspaceId"] = workspace_id
 
-    file = await UserWorkspaceFile.prisma().find_first(where=where_clause)
-    return WorkspaceFile.from_db(file) if file else None
+    return await UserWorkspaceFile.prisma().find_first(where=where_clause)
 
 
 async def get_workspace_file_by_path(
     workspace_id: str,
     path: str,
-) -> Optional[WorkspaceFile]:
+) -> Optional[UserWorkspaceFile]:
     """
     Get a workspace file by its virtual path.
 
@@ -198,20 +139,19 @@ async def get_workspace_file_by_path(
         path: Virtual path
 
     Returns:
-        WorkspaceFile instance or None
+        UserWorkspaceFile instance or None
     """
     # Normalize path
     if not path.startswith("/"):
         path = f"/{path}"
 
-    file = await UserWorkspaceFile.prisma().find_first(
+    return await UserWorkspaceFile.prisma().find_first(
         where={
             "workspaceId": workspace_id,
             "path": path,
             "isDeleted": False,
         }
     )
-    return WorkspaceFile.from_db(file) if file else None
 
 
 async def list_workspace_files(
@@ -220,7 +160,7 @@ async def list_workspace_files(
     include_deleted: bool = False,
     limit: Optional[int] = None,
     offset: int = 0,
-) -> list[WorkspaceFile]:
+) -> list[UserWorkspaceFile]:
     """
     List files in a workspace.
 
@@ -232,7 +172,7 @@ async def list_workspace_files(
         offset: Number of files to skip
 
     Returns:
-        List of WorkspaceFile instances
+        List of UserWorkspaceFile instances
     """
     where_clause: UserWorkspaceFileWhereInput = {"workspaceId": workspace_id}
 
@@ -245,13 +185,12 @@ async def list_workspace_files(
             path_prefix = f"/{path_prefix}"
         where_clause["path"] = {"startswith": path_prefix}
 
-    files = await UserWorkspaceFile.prisma().find_many(
+    return await UserWorkspaceFile.prisma().find_many(
         where=where_clause,
         order={"createdAt": "desc"},
         take=limit,
         skip=offset,
     )
-    return [WorkspaceFile.from_db(f) for f in files]
 
 
 async def count_workspace_files(
@@ -270,7 +209,7 @@ async def count_workspace_files(
     Returns:
         Number of files
     """
-    where_clause: UserWorkspaceFileWhereInput = {"workspaceId": workspace_id}
+    where_clause: dict = {"workspaceId": workspace_id}
     if not include_deleted:
         where_clause["isDeleted"] = False
 
@@ -285,8 +224,8 @@ async def count_workspace_files(
 
 async def soft_delete_workspace_file(
     file_id: str,
-    workspace_id: str,
-) -> Optional[WorkspaceFile]:
+    workspace_id: Optional[str] = None,
+) -> Optional[UserWorkspaceFile]:
     """
     Soft-delete a workspace file.
 
@@ -295,10 +234,10 @@ async def soft_delete_workspace_file(
 
     Args:
         file_id: The file ID
-        workspace_id: Workspace ID for scoping (required)
+        workspace_id: Optional workspace ID for validation
 
     Returns:
-        Updated WorkspaceFile instance or None if not found
+        Updated UserWorkspaceFile instance or None if not found
     """
     # First verify the file exists and belongs to workspace
     file = await get_workspace_file(file_id, workspace_id)
@@ -320,7 +259,7 @@ async def soft_delete_workspace_file(
     )
 
     logger.info(f"Soft-deleted workspace file {file_id}")
-    return WorkspaceFile.from_db(updated) if updated else None
+    return updated
 
 
 async def get_workspace_total_size(workspace_id: str) -> int:
@@ -334,4 +273,4 @@ async def get_workspace_total_size(workspace_id: str) -> int:
         Total size in bytes
     """
     files = await list_workspace_files(workspace_id)
-    return sum(file.size_bytes for file in files)
+    return sum(file.sizeBytes for file in files)

autogpt_platform/backend/backend/db.py

@@ -1,5 +1,5 @@
 from backend.app import run_processes
-from backend.data.db_manager import DatabaseManager
+from backend.executor import DatabaseManager
 
 
 def main():

autogpt_platform/backend/backend/executor/__init__.py

@@ -1,7 +1,11 @@
+from .database import DatabaseManager, DatabaseManagerAsyncClient, DatabaseManagerClient
 from .manager import ExecutionManager
 from .scheduler import Scheduler
 
 __all__ = [
+    "DatabaseManager",
+    "DatabaseManagerClient",
+    "DatabaseManagerAsyncClient",
     "ExecutionManager",
     "Scheduler",
 ]

autogpt_platform/backend/backend/executor/activity_status_generator.py

@@ -22,7 +22,7 @@ from backend.util.settings import Settings
 from backend.util.truncate import truncate
 
 if TYPE_CHECKING:
-    from backend.data.db_manager import DatabaseManagerAsyncClient
+    from backend.executor import DatabaseManagerAsyncClient
 
 logger = logging.getLogger(__name__)
 

autogpt_platform/backend/backend/executor/automod/manager.py

@@ -4,7 +4,7 @@ import logging
 from typing import TYPE_CHECKING, Any, Literal
 
 if TYPE_CHECKING:
-    from backend.data.db_manager import DatabaseManagerAsyncClient
+    from backend.executor import DatabaseManagerAsyncClient
 
 from pydantic import ValidationError
 

autogpt_platform/backend/backend/executor/cluster_lock.py

@@ -1,7 +1,6 @@
 """Redis-based distributed locking for cluster coordination."""
 
 import logging
-import threading
 import time
 from typing import TYPE_CHECKING
 
@@ -20,7 +19,6 @@ class ClusterLock:
         self.owner_id = owner_id
         self.timeout = timeout
         self._last_refresh = 0.0
-        self._refresh_lock = threading.Lock()
 
     def try_acquire(self) -> str | None:
         """Try to acquire the lock.
@@ -33,8 +31,7 @@ class ClusterLock:
         try:
             success = self.redis.set(self.key, self.owner_id, nx=True, ex=self.timeout)
             if success:
-                with self._refresh_lock:
-                    self._last_refresh = time.time()
+                self._last_refresh = time.time()
                 return self.owner_id  # Successfully acquired
 
             # Failed to acquire, get current owner
@@ -60,27 +57,23 @@ class ClusterLock:
         Rate limited to at most once every timeout/10 seconds (minimum 1 second).
         During rate limiting, still verifies lock existence but skips TTL extension.
         Setting _last_refresh to 0 bypasses rate limiting for testing.
-
-        Thread-safe: uses _refresh_lock to protect _last_refresh access.
         """
         # Calculate refresh interval: max(timeout // 10, 1)
         refresh_interval = max(self.timeout // 10, 1)
         current_time = time.time()
 
-        # Check if we're within the rate limit period (thread-safe read)
+        # Check if we're within the rate limit period
         # _last_refresh == 0 forces a refresh (bypasses rate limiting for testing)
-        with self._refresh_lock:
-            last_refresh = self._last_refresh
         is_rate_limited = (
-            last_refresh > 0 and (current_time - last_refresh) < refresh_interval
+            self._last_refresh > 0
+            and (current_time - self._last_refresh) < refresh_interval
         )
 
         try:
             # Always verify lock existence, even during rate limiting
             current_value = self.redis.get(self.key)
             if not current_value:
-                with self._refresh_lock:
-                    self._last_refresh = 0
+                self._last_refresh = 0
                 return False
 
             stored_owner = (
@@ -89,8 +82,7 @@ class ClusterLock:
                 else str(current_value)
             )
             if stored_owner != self.owner_id:
-                with self._refresh_lock:
-                    self._last_refresh = 0
+                self._last_refresh = 0
                 return False
 
             # If rate limited, return True but don't update TTL or timestamp
@@ -99,30 +91,25 @@ class ClusterLock:
 
             # Perform actual refresh
             if self.redis.expire(self.key, self.timeout):
-                with self._refresh_lock:
-                    self._last_refresh = current_time
+                self._last_refresh = current_time
                 return True
 
-            with self._refresh_lock:
-                self._last_refresh = 0
+            self._last_refresh = 0
             return False
 
         except Exception as e:
             logger.error(f"ClusterLock.refresh failed for key {self.key}: {e}")
-            with self._refresh_lock:
-                self._last_refresh = 0
+            self._last_refresh = 0
             return False
 
     def release(self):
         """Release the lock."""
-        with self._refresh_lock:
-            if self._last_refresh == 0:
-                return
+        if self._last_refresh == 0:
+            return
 
         try:
             self.redis.delete(self.key)
         except Exception:
             pass
 
-        with self._refresh_lock:
-            self._last_refresh = 0.0
+        self._last_refresh = 0.0

autogpt_platform/backend/backend/executor/database.py

@@ -4,26 +4,14 @@ from typing import TYPE_CHECKING, Callable, Concatenate, ParamSpec, TypeVar, cas
 
 from backend.api.features.library.db import (
     add_store_agent_to_library,
-    create_graph_in_library,
-    create_library_agent,
-    get_library_agent,
-    get_library_agent_by_graph_id,
     list_library_agents,
-    update_graph_in_library,
-)
-from backend.api.features.store.db import (
-    get_agent,
-    get_available_graph,
-    get_store_agent_details,
-    get_store_agents,
 )
+from backend.api.features.store.db import get_store_agent_details, get_store_agents
 from backend.api.features.store.embeddings import (
     backfill_missing_embeddings,
     cleanup_orphaned_embeddings,
     get_embedding_stats,
 )
-from backend.api.features.store.hybrid_search import unified_hybrid_search
-from backend.copilot import db as chat_db
 from backend.data import db
 from backend.data.analytics import (
     get_accuracy_trends_and_alerts,
@@ -60,7 +48,6 @@ from backend.data.graph import (
     get_graph_metadata,
     get_graph_settings,
     get_node,
-    get_store_listed_graphs,
     validate_graph_execution_permissions,
 )
 from backend.data.human_review import (
@@ -80,10 +67,6 @@ from backend.data.notifications import (
     remove_notifications_from_batch,
 )
 from backend.data.onboarding import increment_onboarding_runs
-from backend.data.understanding import (
-    get_business_understanding,
-    upsert_business_understanding,
-)
 from backend.data.user import (
     get_active_user_ids_in_timerange,
     get_user_by_id,
@@ -93,15 +76,6 @@ from backend.data.user import (
     get_user_notification_preference,
     update_user_integrations,
 )
-from backend.data.workspace import (
-    count_workspace_files,
-    create_workspace_file,
-    get_or_create_workspace,
-    get_workspace_file,
-    get_workspace_file_by_path,
-    list_workspace_files,
-    soft_delete_workspace_file,
-)
 from backend.util.service import (
     AppService,
     AppServiceClient,
@@ -133,13 +107,6 @@ async def _get_credits(user_id: str) -> int:
 
 
 class DatabaseManager(AppService):
-    """Database connection pooling service.
-
-    This service connects to the Prisma engine and exposes database
-    operations via RPC endpoints. It acts as a centralized connection pool
-    for all services that need database access.
-    """
-
     @asynccontextmanager
     async def lifespan(self, app: "FastAPI"):
         async with super().lifespan(app):
@@ -175,15 +142,11 @@ class DatabaseManager(AppService):
     def _(
         f: Callable[P, R], name: str | None = None
     ) -> Callable[Concatenate[object, P], R]:
-        """
-        Exposes a function as an RPC endpoint, and adds a virtual `self` param
-        to the function's type so it can be bound as a method.
-        """
         if name is not None:
             f.__name__ = name
         return cast(Callable[Concatenate[object, P], R], expose(f))
 
-    # ============ Graph Executions ============ #
+    # Executions
     get_child_graph_executions = _(get_child_graph_executions)
     get_graph_executions = _(get_graph_executions)
     get_graph_executions_count = _(get_graph_executions_count)
@@ -207,37 +170,36 @@ class DatabaseManager(AppService):
     get_frequently_executed_graphs = _(get_frequently_executed_graphs)
     get_marketplace_graphs_for_monitoring = _(get_marketplace_graphs_for_monitoring)
 
-    # ============ Graphs ============ #
+    # Graphs
     get_node = _(get_node)
     get_graph = _(get_graph)
     get_connected_output_nodes = _(get_connected_output_nodes)
     get_graph_metadata = _(get_graph_metadata)
     get_graph_settings = _(get_graph_settings)
-    get_store_listed_graphs = _(get_store_listed_graphs)
 
-    # ============ Credits ============ #
+    # Credits
     spend_credits = _(_spend_credits, name="spend_credits")
     get_credits = _(_get_credits, name="get_credits")
 
-    # ============ User + Integrations ============ #
-    get_user_by_id = _(get_user_by_id)
+    # User + User Metadata + User Integrations
     get_user_integrations = _(get_user_integrations)
     update_user_integrations = _(update_user_integrations)
 
-    # ============ User Comms ============ #
+    # User Comms - async
     get_active_user_ids_in_timerange = _(get_active_user_ids_in_timerange)
+    get_user_by_id = _(get_user_by_id)
     get_user_email_by_id = _(get_user_email_by_id)
     get_user_email_verification = _(get_user_email_verification)
     get_user_notification_preference = _(get_user_notification_preference)
 
-    # ============ Human In The Loop ============ #
+    # Human In The Loop
     cancel_pending_reviews_for_execution = _(cancel_pending_reviews_for_execution)
     check_approval = _(check_approval)
     get_or_create_human_review = _(get_or_create_human_review)
     has_pending_reviews_for_graph_exec = _(has_pending_reviews_for_graph_exec)
     update_review_processed_status = _(update_review_processed_status)
 
-    # ============ Notifications ============ #
+    # Notifications - async
     clear_all_user_notification_batches = _(clear_all_user_notification_batches)
     create_or_add_to_user_notification_batch = _(
         create_or_add_to_user_notification_batch
@@ -250,62 +212,29 @@ class DatabaseManager(AppService):
         get_user_notification_oldest_message_in_batch
     )
 
-    # ============ Library ============ #
+    # Library
     list_library_agents = _(list_library_agents)
     add_store_agent_to_library = _(add_store_agent_to_library)
-    create_graph_in_library = _(create_graph_in_library)
-    create_library_agent = _(create_library_agent)
-    get_library_agent = _(get_library_agent)
-    get_library_agent_by_graph_id = _(get_library_agent_by_graph_id)
-    update_graph_in_library = _(update_graph_in_library)
     validate_graph_execution_permissions = _(validate_graph_execution_permissions)
 
-    # ============ Onboarding ============ #
+    # Onboarding
     increment_onboarding_runs = _(increment_onboarding_runs)
 
-    # ============ OAuth ============ #
+    # OAuth
     cleanup_expired_oauth_tokens = _(cleanup_expired_oauth_tokens)
 
-    # ============ Store ============ #
+    # Store
     get_store_agents = _(get_store_agents)
     get_store_agent_details = _(get_store_agent_details)
-    get_agent = _(get_agent)
-    get_available_graph = _(get_available_graph)
 
-    # ============ Search ============ #
+    # Store Embeddings
     get_embedding_stats = _(get_embedding_stats)
     backfill_missing_embeddings = _(backfill_missing_embeddings)
     cleanup_orphaned_embeddings = _(cleanup_orphaned_embeddings)
-    unified_hybrid_search = _(unified_hybrid_search)
 
-    # ============ Summary Data ============ #
+    # Summary data - async
     get_user_execution_summary_data = _(get_user_execution_summary_data)
 
-    # ============ Workspace ============ #
-    count_workspace_files = _(count_workspace_files)
-    create_workspace_file = _(create_workspace_file)
-    get_or_create_workspace = _(get_or_create_workspace)
-    get_workspace_file = _(get_workspace_file)
-    get_workspace_file_by_path = _(get_workspace_file_by_path)
-    list_workspace_files = _(list_workspace_files)
-    soft_delete_workspace_file = _(soft_delete_workspace_file)
-
-    # ============ Understanding ============ #
-    get_business_understanding = _(get_business_understanding)
-    upsert_business_understanding = _(upsert_business_understanding)
-
-    # ============ CoPilot Chat Sessions ============ #
-    get_chat_session = _(chat_db.get_chat_session)
-    create_chat_session = _(chat_db.create_chat_session)
-    update_chat_session = _(chat_db.update_chat_session)
-    add_chat_message = _(chat_db.add_chat_message)
-    add_chat_messages_batch = _(chat_db.add_chat_messages_batch)
-    get_user_chat_sessions = _(chat_db.get_user_chat_sessions)
-    get_user_session_count = _(chat_db.get_user_session_count)
-    delete_chat_session = _(chat_db.delete_chat_session)
-    get_next_sequence = _(chat_db.get_next_sequence)
-    update_tool_message_content = _(chat_db.update_tool_message_content)
-
 
 class DatabaseManagerClient(AppServiceClient):
     d = DatabaseManager
@@ -367,50 +296,43 @@ class DatabaseManagerAsyncClient(AppServiceClient):
     def get_service_type(cls):
         return DatabaseManager
 
-    # ============ Graph Executions ============ #
     create_graph_execution = d.create_graph_execution
     get_child_graph_executions = d.get_child_graph_executions
     get_connected_output_nodes = d.get_connected_output_nodes
     get_latest_node_execution = d.get_latest_node_execution
-    get_graph_execution = d.get_graph_execution
-    get_graph_execution_meta = d.get_graph_execution_meta
-    get_graph_executions = d.get_graph_executions
-    get_node_execution = d.get_node_execution
-    get_node_executions = d.get_node_executions
-    update_graph_execution_stats = d.update_graph_execution_stats
-    update_node_execution_status = d.update_node_execution_status
-    update_node_execution_status_batch = d.update_node_execution_status_batch
-    upsert_execution_input = d.upsert_execution_input
-    upsert_execution_output = d.upsert_execution_output
-    get_execution_outputs_by_node_exec_id = d.get_execution_outputs_by_node_exec_id
-    get_execution_kv_data = d.get_execution_kv_data
-    set_execution_kv_data = d.set_execution_kv_data
-
-    # ============ Graphs ============ #
     get_graph = d.get_graph
     get_graph_metadata = d.get_graph_metadata
     get_graph_settings = d.get_graph_settings
+    get_graph_execution = d.get_graph_execution
+    get_graph_execution_meta = d.get_graph_execution_meta
     get_node = d.get_node
-    get_store_listed_graphs = d.get_store_listed_graphs
-
-    # ============ User + Integrations ============ #
+    get_node_execution = d.get_node_execution
+    get_node_executions = d.get_node_executions
     get_user_by_id = d.get_user_by_id
     get_user_integrations = d.get_user_integrations
+    upsert_execution_input = d.upsert_execution_input
+    upsert_execution_output = d.upsert_execution_output
+    get_execution_outputs_by_node_exec_id = d.get_execution_outputs_by_node_exec_id
+    update_graph_execution_stats = d.update_graph_execution_stats
+    update_node_execution_status = d.update_node_execution_status
+    update_node_execution_status_batch = d.update_node_execution_status_batch
     update_user_integrations = d.update_user_integrations
+    get_execution_kv_data = d.get_execution_kv_data
+    set_execution_kv_data = d.set_execution_kv_data
 
-    # ============ Human In The Loop ============ #
+    # Human In The Loop
     cancel_pending_reviews_for_execution = d.cancel_pending_reviews_for_execution
     check_approval = d.check_approval
     get_or_create_human_review = d.get_or_create_human_review
     update_review_processed_status = d.update_review_processed_status
 
-    # ============ User Comms ============ #
+    # User Comms
     get_active_user_ids_in_timerange = d.get_active_user_ids_in_timerange
     get_user_email_by_id = d.get_user_email_by_id
     get_user_email_verification = d.get_user_email_verification
     get_user_notification_preference = d.get_user_notification_preference
 
-    # ============ Notifications ============ #
+    # Notifications
     clear_all_user_notification_batches = d.clear_all_user_notification_batches
     create_or_add_to_user_notification_batch = (
         d.create_or_add_to_user_notification_batch
@@ -423,55 +345,20 @@ class DatabaseManagerAsyncClient(AppServiceClient):
         d.get_user_notification_oldest_message_in_batch
     )
 
-    # ============ Library ============ #
+    # Library
     list_library_agents = d.list_library_agents
     add_store_agent_to_library = d.add_store_agent_to_library
-    create_graph_in_library = d.create_graph_in_library
-    create_library_agent = d.create_library_agent
-    get_library_agent = d.get_library_agent
-    get_library_agent_by_graph_id = d.get_library_agent_by_graph_id
-    update_graph_in_library = d.update_graph_in_library
     validate_graph_execution_permissions = d.validate_graph_execution_permissions
 
-    # ============ Onboarding ============ #
+    # Onboarding
     increment_onboarding_runs = d.increment_onboarding_runs
 
-    # ============ OAuth ============ #
+    # OAuth
     cleanup_expired_oauth_tokens = d.cleanup_expired_oauth_tokens
 
-    # ============ Store ============ #
+    # Store
     get_store_agents = d.get_store_agents
     get_store_agent_details = d.get_store_agent_details
-    get_agent = d.get_agent
-    get_available_graph = d.get_available_graph
 
-    # ============ Search ============ #
-    unified_hybrid_search = d.unified_hybrid_search
-
-    # ============ Summary Data ============ #
+    # Summary data
     get_user_execution_summary_data = d.get_user_execution_summary_data
-
-    # ============ Workspace ============ #
-    count_workspace_files = d.count_workspace_files
-    create_workspace_file = d.create_workspace_file
-    get_or_create_workspace = d.get_or_create_workspace
-    get_workspace_file = d.get_workspace_file
-    get_workspace_file_by_path = d.get_workspace_file_by_path
-    list_workspace_files = d.list_workspace_files
-    soft_delete_workspace_file = d.soft_delete_workspace_file
-
-    # ============ Understanding ============ #
-    get_business_understanding = d.get_business_understanding
-    upsert_business_understanding = d.upsert_business_understanding
-
-    # ============ CoPilot Chat Sessions ============ #
-    get_chat_session = d.get_chat_session
-    create_chat_session = d.create_chat_session
-    update_chat_session = d.update_chat_session
-    add_chat_message = d.add_chat_message
-    add_chat_messages_batch = d.add_chat_messages_batch
-    get_user_chat_sessions = d.get_user_chat_sessions
-    get_user_session_count = d.get_user_session_count
-    delete_chat_session = d.delete_chat_session
-    get_next_sequence = d.get_next_sequence
-    update_tool_message_content = d.update_tool_message_content

autogpt_platform/backend/backend/executor/manager.py

@@ -93,10 +93,7 @@ from .utils import (
 )
 
 if TYPE_CHECKING:
-    from backend.data.db_manager import (
-        DatabaseManagerAsyncClient,
-        DatabaseManagerClient,
-    )
+    from backend.executor import DatabaseManagerAsyncClient, DatabaseManagerClient
 
 
 _logger = logging.getLogger(__name__)

autogpt_platform/backend/backend/util/clients.py

@@ -13,15 +13,12 @@ if TYPE_CHECKING:
     from openai import AsyncOpenAI
     from supabase import AClient, Client
 
-    from backend.data.db_manager import (
-        DatabaseManagerAsyncClient,
-        DatabaseManagerClient,
-    )
     from backend.data.execution import (
         AsyncRedisExecutionEventBus,
         RedisExecutionEventBus,
     )
     from backend.data.rabbitmq import AsyncRabbitMQ, SyncRabbitMQ
+    from backend.executor import DatabaseManagerAsyncClient, DatabaseManagerClient
     from backend.executor.scheduler import SchedulerClient
     from backend.integrations.credentials_store import IntegrationCredentialsStore
     from backend.notifications.notifications import NotificationManagerClient
@@ -30,7 +27,7 @@ if TYPE_CHECKING:
 @thread_cached
 def get_database_manager_client() -> "DatabaseManagerClient":
     """Get a thread-cached DatabaseManagerClient with request retry enabled."""
-    from backend.data.db_manager import DatabaseManagerClient
+    from backend.executor import DatabaseManagerClient
     from backend.util.service import get_service_client
 
     return get_service_client(DatabaseManagerClient, request_retry=True)
@@ -41,7 +38,7 @@ def get_database_manager_async_client(
     should_retry: bool = True,
 ) -> "DatabaseManagerAsyncClient":
     """Get a thread-cached DatabaseManagerAsyncClient with request retry enabled."""
-    from backend.data.db_manager import DatabaseManagerAsyncClient
+    from backend.executor import DatabaseManagerAsyncClient
     from backend.util.service import get_service_client
 
     return get_service_client(DatabaseManagerAsyncClient, request_retry=should_retry)
@@ -109,20 +106,6 @@ async def get_async_execution_queue() -> "AsyncRabbitMQ":
     return client
 
 
-# ============ CoPilot Queue Helpers ============ #
-
-
-@thread_cached
-async def get_async_copilot_queue() -> "AsyncRabbitMQ":
-    """Get a thread-cached AsyncRabbitMQ CoPilot queue client."""
-    from backend.copilot.executor.utils import create_copilot_queue_config
-    from backend.data.rabbitmq import AsyncRabbitMQ
-
-    client = AsyncRabbitMQ(create_copilot_queue_config())
-    await client.connect()
-    return client
-
-
 # ============ Integration Credentials Store ============ #
 
 

autogpt_platform/backend/backend/util/file.py

@@ -383,7 +383,7 @@ async def store_media_file(
                     else:
                         info = await workspace_manager.get_file_info(ws.file_ref)
                     if info:
-                        return MediaFileType(f"{file}#{info.mime_type}")
+                        return MediaFileType(f"{file}#{info.mimeType}")
                 except Exception:
                     pass
             return MediaFileType(file)
@@ -397,7 +397,7 @@ async def store_media_file(
             filename=filename,
             overwrite=True,
         )
-        return MediaFileType(f"workspace://{file_record.id}#{file_record.mime_type}")
+        return MediaFileType(f"workspace://{file_record.id}#{file_record.mimeType}")
 
     else:
         raise ValueError(f"Invalid return_format: {return_format}")

autogpt_platform/backend/backend/util/prompt.py

@@ -535,18 +535,14 @@ async def _summarize_messages_llm(
             {
                 "role": "system",
                 "content": (
-                    "Create a factual summary of the conversation so far. "
+                    "Create a detailed summary of the conversation so far. "
                     "This summary will be used as context when continuing the conversation.\n\n"
-                    "CRITICAL: Only include information that is EXPLICITLY present in the "
-                    "conversation. Do NOT fabricate, infer, or invent any details. "
-                    "If a section has no relevant content in the conversation, skip it entirely.\n\n"
                     "Before writing the summary, analyze each message chronologically to identify:\n"
                     "- User requests and their explicit goals\n"
-                    "- Actions taken and key decisions made\n"
+                    "- Your approach and key decisions made\n"
                     "- Technical specifics (file names, tool outputs, function signatures)\n"
                     "- Errors encountered and resolutions applied\n\n"
-                    "Include ONLY the sections below that have relevant content "
-                    "(skip sections with nothing to report):\n\n"
+                    "You MUST include ALL of the following sections:\n\n"
                     "## 1. Primary Request and Intent\n"
                     "The user's explicit goals and what they are trying to accomplish.\n\n"
                     "## 2. Key Technical Concepts\n"
@@ -554,14 +550,19 @@ async def _summarize_messages_llm(
                     "## 3. Files and Resources Involved\n"
                     "Specific files examined or modified, with relevant snippets and identifiers.\n\n"
                     "## 4. Errors and Fixes\n"
-                    "Problems encountered, error messages, and their resolutions.\n\n"
-                    "## 5. All User Messages\n"
-                    "A complete list of all user inputs (excluding tool outputs) "
-                    "to preserve their exact requests.\n\n"
-                    "## 6. Pending Tasks\n"
+                    "Problems encountered, error messages, and their resolutions. "
+                    "Include any user feedback on fixes.\n\n"
+                    "## 5. Problem Solving\n"
+                    "Issues that have been resolved and how they were addressed.\n\n"
+                    "## 6. All User Messages\n"
+                    "A complete list of all user inputs (excluding tool outputs) to preserve their exact requests.\n\n"
+                    "## 7. Pending Tasks\n"
                     "Work items the user explicitly requested that have not yet been completed.\n\n"
-                    "## 7. Current State\n"
-                    "What was happening most recently in the conversation."
+                    "## 8. Current Work\n"
+                    "Precise description of what was being worked on most recently, including relevant context.\n\n"
+                    "## 9. Next Steps\n"
+                    "What should happen next, aligned with the user's most recent requests. "
+                    "Include verbatim quotes of recent instructions if relevant."
                 ),
             },
             {"role": "user", "content": f"Summarize:\n\n{conversation_text}"},