agent-platform/services/tool-service/app/domain/repositories.py

from datetime import datetime

from core_shared import JSONValue
from sqlalchemy import select
from sqlalchemy.orm import Session

from app.db.models import ToolBinding, ToolCredential, ToolDefinition, ToolConnection


class ToolDefinitionRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def create(
        self,
        *,
        plugin_id: str | None,
        code: str,
        name: str,
        tool_type: str,
        description: str | None) -> ToolDefinition:
        entity = ToolDefinition(
            plugin_id=plugin_id,
            code=code,
            name=name,
            tool_type=tool_type,
            description=description)
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def list_all(self) -> list[ToolDefinition]:
        stmt = (
            select(ToolDefinition)
            .order_by(ToolDefinition.created_time.desc())
        )
        return list(self.db.scalars(stmt))

    def get_by_id(self, *, tool_id: str) -> ToolDefinition | None:
        stmt = (
            select(ToolDefinition)
            .where(ToolDefinition.id == tool_id)
        )
        return self.db.scalar(stmt)

    def save(self, entity: ToolDefinition) -> ToolDefinition:
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def delete(self, entity: ToolDefinition) -> None:
        self.db.delete(entity)
        self.db.commit()


class ToolConnectionRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def create(
        self,
        *,
        tool_id: str,
        input_schema_json: dict[str, JSONValue] | None,
        output_schema_json: dict[str, JSONValue] | None,
        invoke_config_json: dict[str, JSONValue] | None,
        timeout_ms: int | None,
        retry_policy_json: dict[str, JSONValue] | None) -> ToolConnection:
        entity = ToolConnection(
            tool_id=tool_id,
            input_schema_json=input_schema_json,
            output_schema_json=output_schema_json,
            invoke_config_json=invoke_config_json,
            timeout_ms=timeout_ms,
            retry_policy_json=retry_policy_json)
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def list_by_tool(self, *, tool_id: str) -> list[ToolConnection]:
        stmt = (
            select(ToolConnection)
            .where(ToolConnection.tool_id == tool_id)
            .order_by(ToolConnection.created_time.desc())
        )
        return list(self.db.scalars(stmt))

    def list_all(self) -> list[ToolConnection]:
        stmt = (
            select(ToolConnection)
            .order_by(ToolConnection.created_time.desc())
        )
        return list(self.db.scalars(stmt))

    def list_pending_mcp_discovery(
        self,
        *,
        stale_before: datetime,
        limit: int = 20) -> list[ToolConnection]:
        pending: list[ToolConnection] = []
        for connection in self.list_all():
            config = connection.invoke_config_json or {}
            status_payload = config.get("mcp_status")
            if not isinstance(status_payload, dict):
                continue
            status = status_payload.get("status")
            if status == "queued" or (status == "running" and connection.updated_time < stale_before):
                pending.append(connection)
        pending.sort(key=lambda item: item.updated_time)
        return pending[:limit]

    def get_by_id(self, *, tool_connection_id: str) -> ToolConnection | None:
        stmt = (
            select(ToolConnection)
            .where(ToolConnection.id == tool_connection_id)
        )
        return self.db.scalar(stmt)

    def save(self, entity: ToolConnection) -> ToolConnection:
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def delete(self, entity: ToolConnection) -> None:
        self.db.delete(entity)
        self.db.commit()

class ToolBindingRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def create(
        self,
        *,
        app_id: str,
        tool_connection_id: str,
        credential_id: str | None,
        binding_scope: str,
        enabled: bool,
        config_json: dict[str, JSONValue] | None) -> ToolBinding:
        entity = ToolBinding(
            app_id=app_id,
            tool_connection_id=tool_connection_id,
            credential_id=credential_id,
            binding_scope=binding_scope,
            enabled=enabled,
            config_json=config_json)
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def list_by_scope(self, *, app_id: str | None = None) -> list[ToolBinding]:
        stmt = select(ToolBinding)
        if app_id is not None:
            stmt = stmt.where(ToolBinding.app_id == app_id)
        stmt = stmt.order_by(ToolBinding.created_time.desc())
        return list(self.db.scalars(stmt))

    def get_by_id(self, *, binding_id: str) -> ToolBinding | None:
        stmt = (
            select(ToolBinding)
            .where(ToolBinding.id == binding_id)
        )
        return self.db.scalar(stmt)

    def save(self, entity: ToolBinding) -> ToolBinding:
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def delete(self, entity: ToolBinding) -> None:
        self.db.delete(entity)
        self.db.commit()


class ToolCredentialRepository:
    def __init__(self, db: Session) -> None:
        self.db = db

    def create(
        self,
        *,
        name: str,
        credential_type: str,
        encrypted_payload_text: str,
        secret_fingerprint: str,
        encryption_algorithm: str,
        metadata_json: dict[str, JSONValue]) -> ToolCredential:
        entity = ToolCredential(
            name=name,
            credential_type=credential_type,
            encrypted_payload_text=encrypted_payload_text,
            secret_fingerprint=secret_fingerprint,
            encryption_algorithm=encryption_algorithm,
            metadata_json=metadata_json)
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def list_all(self) -> list[ToolCredential]:
        stmt = (
            select(ToolCredential)
            .order_by(ToolCredential.created_time.desc())
        )
        return list(self.db.scalars(stmt))

    def get_by_id(self, *, credential_id: str) -> ToolCredential | None:
        stmt = (
            select(ToolCredential)
            .where(ToolCredential.id == credential_id)
        )
        return self.db.scalar(stmt)

    def save(self, entity: ToolCredential) -> ToolCredential:
        self.db.add(entity)
        self.db.commit()
        self.db.refresh(entity)
        return entity

    def delete(self, entity: ToolCredential) -> None:
        self.db.delete(entity)
        self.db.commit()