DocsGPT/application/api/answer/services/continuation_service.py

"""Service for saving and restoring tool-call continuation state.

When a stream pauses (tool needs approval or client-side execution),
the full execution state is persisted to Postgres so the client can
resume later by sending tool_actions.
"""

import logging
from typing import Any, Dict, List, Optional

from application.storage.db.base_repository import looks_like_uuid
from application.storage.db.repositories.conversations import ConversationsRepository
from application.storage.db.repositories.pending_tool_state import (
    PendingToolStateRepository,
)
from application.storage.db.serialization import coerce_pg_native as _make_serializable
from application.storage.db.session import db_readonly, db_session

logger = logging.getLogger(__name__)

# TTL for pending states — auto-cleaned after this period
PENDING_STATE_TTL_SECONDS = 30 * 60  # 30 minutes

# Re-export so the existing tests at tests/api/answer/services/test_continuation_service_pg.py
# can keep importing ``_make_serializable`` from here.
__all__ = ["_make_serializable", "ContinuationService", "PENDING_STATE_TTL_SECONDS"]


class ContinuationService:
    """Manages pending tool-call state in Postgres."""

    def __init__(self):
        # No-op constructor retained for call-site compatibility. State
        # lives in Postgres now; each operation opens its own short-lived
        # session rather than holding a connection on the service.
        pass

    def save_state(
        self,
        conversation_id: str,
        user: str,
        messages: List[Dict],
        pending_tool_calls: List[Dict],
        tools_dict: Dict,
        tool_schemas: List[Dict],
        agent_config: Dict,
        client_tools: Optional[List[Dict]] = None,
    ) -> str:
        """Save execution state for later continuation.

        ``conversation_id`` may be a Postgres UUID or the legacy Mongo
        ``ObjectId`` string — the latter is resolved via
        ``conversations.legacy_mongo_id`` to find the matching row.

        Args:
            conversation_id: The conversation this state belongs to.
            user: Owner user ID.
            messages: Full messages array at the pause point.
            pending_tool_calls: Tool calls awaiting client action.
            tools_dict: Serializable tools configuration dict.
            tool_schemas: LLM-formatted tool schemas (agent.tools).
            agent_config: Config needed to recreate the agent on resume.
            client_tools: Client-provided tool schemas for client-side execution.

        Returns:
            The string ID (conversation_id as provided) of the saved state.
        """
        with db_session() as conn:
            conv = ConversationsRepository(conn).get_by_legacy_id(conversation_id)
            if conv is not None:
                pg_conv_id = conv["id"]
            elif looks_like_uuid(conversation_id):
                pg_conv_id = conversation_id
            else:
                # Unresolvable legacy ObjectId — downstream ``CAST AS uuid``
                # would raise and poison the save. Surface the mismatch so
                # the caller can decide (the stream loop in routes/base.py
                # already wraps this in try/except).
                raise ValueError(
                    f"Cannot save continuation state: conversation_id "
                    f"{conversation_id!r} is neither a PG UUID nor a "
                    f"backfilled legacy Mongo id."
                )
            PendingToolStateRepository(conn).save_state(
                pg_conv_id,
                user,
                messages=_make_serializable(messages),
                pending_tool_calls=_make_serializable(pending_tool_calls),
                tools_dict=_make_serializable(tools_dict),
                tool_schemas=_make_serializable(tool_schemas),
                agent_config=_make_serializable(agent_config),
                client_tools=_make_serializable(client_tools) if client_tools else None,
            )

        logger.info(
            f"Saved continuation state for conversation {conversation_id} "
            f"with {len(pending_tool_calls)} pending tool call(s)"
        )
        return conversation_id

    def load_state(
        self, conversation_id: str, user: str
    ) -> Optional[Dict[str, Any]]:
        """Load pending continuation state.

        Returns:
            The state dict, or None if no pending state exists.
        """
        with db_readonly() as conn:
            conv = ConversationsRepository(conn).get_by_legacy_id(conversation_id)
            if conv is not None:
                pg_conv_id = conv["id"]
            elif looks_like_uuid(conversation_id):
                pg_conv_id = conversation_id
            else:
                # Unresolvable legacy ObjectId → no state can exist for it.
                return None
            doc = PendingToolStateRepository(conn).load_state(pg_conv_id, user)
        if not doc:
            return None
        return doc

    def delete_state(self, conversation_id: str, user: str) -> bool:
        """Delete pending state after successful resumption.

        Returns:
            True if a row was deleted.
        """
        with db_session() as conn:
            conv = ConversationsRepository(conn).get_by_legacy_id(conversation_id)
            if conv is not None:
                pg_conv_id = conv["id"]
            elif looks_like_uuid(conversation_id):
                pg_conv_id = conversation_id
            else:
                # Unresolvable legacy ObjectId → nothing to delete.
                return False
            deleted = PendingToolStateRepository(conn).delete_state(pg_conv_id, user)
        if deleted:
            logger.info(
                f"Deleted continuation state for conversation {conversation_id}"
            )
        return deleted

    def mark_resuming(self, conversation_id: str, user: str) -> bool:
        """Flip the pending row to ``resuming`` so a crashed resume can be retried."""
        with db_session() as conn:
            conv = ConversationsRepository(conn).get_by_legacy_id(conversation_id)
            if conv is not None:
                pg_conv_id = conv["id"]
            elif looks_like_uuid(conversation_id):
                pg_conv_id = conversation_id
            else:
                return False
            flipped = PendingToolStateRepository(conn).mark_resuming(
                pg_conv_id, user
            )
        if flipped:
            logger.info(
                f"Marked continuation state as resuming for conversation "
                f"{conversation_id}"
            )
        return flipped