agent-platform/services/tool-service/app/application/services.py

from core_shared import JSONValue
from core_shared.secrets import EncryptedSecret, SecretCipher

from app.db.models import ToolBinding, ToolCredential, ToolDefinition, ToolVersion
from app.domain.repositories import (
    ToolBindingRepository,
    ToolCredentialRepository,
    ToolDefinitionRepository,
    ToolVersionRepository,
)
from app.schemas.tool import (
    McpConnectData,
    McpConnectRequestDto,
    McpToolDto,
    ToolBindingCreateRequest,
    ToolBindingCreateRequestDto,
    ToolBindingDeleteRequestDto,
    ToolBindingDetailRequestDto,
    ToolBindingUpdateRequestDto,
    ToolCreateRequest,
    ToolCreateRequestDto,
    ToolCredentialCreateRequest,
    ToolCredentialCreateRequestDto,
    ToolCredentialDeleteRequestDto,
    ToolCredentialDetailRequestDto,
    ToolCredentialDto,
    ToolCredentialRevealDto,
    ToolCredentialUpdateRequestDto,
    ToolDeleteRequestDto,
    ToolDetailRequestDto,
    ToolDto,
    ToolUpdateRequestDto,
    ToolVersionCreateRequest,
    ToolVersionCreateRequestDto,
    ToolVersionDetailRequestDto,
    ToolVersionDto,
    ToolVersionUpdateRequestDto,
)


class ToolApplicationService:
    def __init__(
        self,
        tool_definition_repository: ToolDefinitionRepository,
        tool_version_repository: ToolVersionRepository,
        tool_binding_repository: ToolBindingRepository,
        tool_credential_repository: ToolCredentialRepository,
        secret_cipher: SecretCipher) -> None:
        self.tool_definition_repository = tool_definition_repository
        self.tool_version_repository = tool_version_repository
        self.tool_binding_repository = tool_binding_repository
        self.tool_credential_repository = tool_credential_repository
        self.secret_cipher = secret_cipher

    def create_tool_definition(self, payload: ToolCreateRequest) -> ToolDefinition:
        code = payload.code or self._build_tool_code(payload.name)
        return self.tool_definition_repository.create(
            plugin_id=payload.plugin_id,
            code=code,
            name=payload.name,
            tool_type=payload.tool_type,
            description=payload.description)

    def list_tool_definitions(self) -> list[ToolDefinition]:
        return self.tool_definition_repository.list_all()

    def create_tool_definition_from_contract(
        self,
        payload: ToolCreateRequestDto) -> ToolDefinition:
        return self.create_tool_definition(
            ToolCreateRequest(
                plugin_id=payload.pluginId,
                name=payload.name,
                tool_type=payload.toolType,
                description=payload.description))

    def delete_tool_definition_from_contract(self, payload: ToolDeleteRequestDto) -> bool:
        entity = self.tool_definition_repository.get_by_id(tool_id=payload.toolId)
        if entity is None:
            return False
        self.tool_definition_repository.delete(entity)
        return True

    def get_tool_definition_from_contract(
        self,
        payload: ToolDetailRequestDto) -> ToolDefinition | None:
        return self.tool_definition_repository.get_by_id(tool_id=payload.toolId)

    def update_tool_definition_from_contract(
        self,
        payload: ToolUpdateRequestDto) -> ToolDefinition | None:
        entity = self.tool_definition_repository.get_by_id(tool_id=payload.toolId)
        if entity is None:
            return None
        if payload.name is not None:
            entity.name = payload.name
            entity.code = self._build_tool_code(payload.name)
        if payload.toolType is not None:
            entity.tool_type = payload.toolType
        if payload.description is not None:
            entity.description = payload.description
        if payload.pluginId is not None:
            entity.plugin_id = payload.pluginId
        return self.tool_definition_repository.save(entity)

    def create_tool_version(self, payload: ToolVersionCreateRequest) -> ToolVersion:
        return self.tool_version_repository.create(
            tool_id=payload.tool_id,
            input_schema_json=payload.input_schema_json,
            output_schema_json=payload.output_schema_json,
            invoke_config_json=payload.invoke_config_json,
            timeout_ms=payload.timeout_ms,
            retry_policy_json=payload.retry_policy_json)

    def list_tool_versions(self, tool_id: str | None = None) -> list[ToolVersion]:
        if tool_id is None:
            return self.tool_version_repository.list_all()
        return self.tool_version_repository.list_by_tool(tool_id=tool_id)

    def create_tool_version_from_contract(
        self,
        payload: ToolVersionCreateRequestDto) -> ToolVersion:
        return self.create_tool_version(
            ToolVersionCreateRequest(
                tool_id=payload.toolId,
                input_schema_json=payload.inputSchema,
                output_schema_json=payload.outputSchema,
                invoke_config_json=payload.invokeConfig,
                timeout_ms=payload.timeoutMs,
                retry_policy_json=payload.retryPolicy))

    def delete_tool_version(self, *, connection_id: str) -> bool:
        entity = self.tool_version_repository.get_by_id(tool_version_id=connection_id)
        if entity is None:
            return False
        self.tool_version_repository.delete(entity)
        return True

    def get_tool_version_from_contract(
        self,
        payload: ToolVersionDetailRequestDto) -> ToolVersion | None:
        return self.tool_version_repository.get_by_id(
            tool_version_id=payload.connectionId)

    def update_tool_version_from_contract(
        self,
        payload: ToolVersionUpdateRequestDto) -> ToolVersion | None:
        entity = self.tool_version_repository.get_by_id(
            tool_version_id=payload.connectionId)
        if entity is None:
            return None
        if payload.inputSchema is not None:
            entity.input_schema_json = payload.inputSchema
        if payload.outputSchema is not None:
            entity.output_schema_json = payload.outputSchema
        if payload.invokeConfig is not None:
            entity.invoke_config_json = payload.invokeConfig
        if payload.timeoutMs is not None:
            entity.timeout_ms = payload.timeoutMs
        if payload.retryPolicy is not None:
            entity.retry_policy_json = payload.retryPolicy
        return self.tool_version_repository.save(entity)

    def connect_mcp_server(self, payload: McpConnectRequestDto) -> McpConnectData:
        server_name, config = self._normalize_mcp_config(payload)
        discovered_tools = self._extract_mcp_tools(config)
        if discovered_tools:
            config["mcp_tools"] = [
                tool.model_dump(mode="json", by_alias=False)
                for tool in discovered_tools
            ]
        tool = self.create_tool_definition(
            ToolCreateRequest(
                name=payload.name or server_name,
                tool_type="mcp",
                description=f"MCP SSE server: {config.get('url', '')}"))
        version = self.create_tool_version(
            ToolVersionCreateRequest(
                tool_id=tool.id,
                input_schema_json={},
                output_schema_json={},
                invoke_config_json=config,
                timeout_ms=self._timeout_ms(config),
                retry_policy_json={"max_attempts": 1}))
        return McpConnectData(
            tool=ToolDto.from_entity(tool),
            connection=ToolVersionDto.from_entity(version),
            discoveredTools=discovered_tools)

    def create_tool_binding(self, payload: ToolBindingCreateRequest) -> ToolBinding:
        if payload.credential_id is not None:
            credential = self.tool_credential_repository.get_by_id(
                credential_id=payload.credential_id)
            if credential is None:
                raise ValueError(f"tool credential not found: {payload.credential_id}")
        return self.tool_binding_repository.create(
            app_id=payload.app_id,
            tool_version_id=payload.tool_version_id,
            credential_id=payload.credential_id,
            binding_scope=payload.binding_scope,
            enabled=payload.enabled,
            config_json=payload.config_json)

    def create_tool_binding_from_contract(
        self,
        payload: ToolBindingCreateRequestDto) -> ToolBinding:
        return self.create_tool_binding(
            ToolBindingCreateRequest(
                app_id=payload.appId,
                tool_version_id=payload.toolVersionId,
                credential_id=payload.credentialId,
                binding_scope=payload.bindingScope,
                enabled=True,
                config_json=payload.configJson))

    def list_tool_bindings(self, app_id: str | None = None) -> list[ToolBinding]:
        return self.tool_binding_repository.list_by_scope(app_id=app_id)

    def delete_tool_binding(self, payload: ToolBindingDeleteRequestDto) -> bool:
        entity = self.tool_binding_repository.get_by_id(binding_id=payload.bindingId)
        if entity is None:
            return False
        self.tool_binding_repository.delete(entity)
        return True

    def get_tool_binding_from_contract(
        self,
        payload: ToolBindingDetailRequestDto) -> ToolBinding | None:
        return self.tool_binding_repository.get_by_id(binding_id=payload.bindingId)

    def update_tool_binding_from_contract(
        self,
        payload: ToolBindingUpdateRequestDto) -> ToolBinding | None:
        entity = self.tool_binding_repository.get_by_id(binding_id=payload.bindingId)
        if entity is None:
            return None
        if payload.credentialId is not None:
            entity.credential_id = payload.credentialId
        if payload.bindingScope is not None:
            entity.binding_scope = payload.bindingScope
        if payload.configJson is not None:
            entity.config_json = payload.configJson
        return self.tool_binding_repository.save(entity)

    def create_tool_credential(self, payload: ToolCredentialCreateRequest) -> ToolCredential:
        encrypted = self.secret_cipher.encrypt_json(payload.secret_json)
        return self.tool_credential_repository.create(
            name=payload.name,
            credential_type=payload.credential_type,
            encrypted_payload_text=encrypted.ciphertext,
            secret_fingerprint=encrypted.fingerprint,
            encryption_algorithm=encrypted.algorithm,
            metadata_json=payload.metadata_json)

    def create_tool_credential_from_contract(
        self,
        payload: ToolCredentialCreateRequestDto) -> ToolCredential:
        return self.create_tool_credential(
            ToolCredentialCreateRequest(
                name=payload.name,
                credential_type=payload.credentialType,
                secret_json=payload.secretJson,
                metadata_json=payload.metadataJson))

    def list_tool_credentials(self) -> list[ToolCredential]:
        return self.tool_credential_repository.list_all()

    def delete_tool_credential(self, payload: ToolCredentialDeleteRequestDto) -> bool:
        entity = self.tool_credential_repository.get_by_id(
            credential_id=payload.credentialId)
        if entity is None:
            return False
        self.tool_credential_repository.delete(entity)
        return True

    def get_tool_credential_from_contract(
        self,
        payload: ToolCredentialDetailRequestDto) -> ToolCredential | None:
        return self.tool_credential_repository.get_by_id(
            credential_id=payload.credentialId)

    def update_tool_credential_from_contract(
        self,
        payload: ToolCredentialUpdateRequestDto) -> ToolCredential | None:
        entity = self.tool_credential_repository.get_by_id(
            credential_id=payload.credentialId)
        if entity is None:
            return None
        if payload.name is not None:
            entity.name = payload.name
        if payload.metadataJson is not None:
            entity.metadata_json = payload.metadataJson
        return self.tool_credential_repository.save(entity)

    def reveal_tool_credential(
        self,
        *,
        credential_id: str) -> tuple[ToolCredential, dict[str, JSONValue]] | None:
        credential = self.tool_credential_repository.get_by_id(
            credential_id=credential_id)
        if credential is None:
            return None
        secret_json = self.secret_cipher.decrypt_json(
            EncryptedSecret(
                ciphertext=credential.encrypted_payload_text,
                fingerprint=credential.secret_fingerprint,
                algorithm=credential.encryption_algorithm)
        )
        return credential, secret_json

    def reveal_tool_credential_from_contract(
        self,
        *,
        credential_id: str) -> ToolCredentialRevealDto | None:
        result = self.reveal_tool_credential(credential_id=credential_id)
        if result is None:
            return None
        credential, secret_json = result
        return ToolCredentialRevealDto(
            credential=ToolCredentialDto.from_entity(credential),
            secretJson=secret_json)

    def get_tool_binding_detail(
        self,
        *,
        binding_id: str) -> tuple[ToolBinding, ToolVersion, ToolDefinition] | None:
        binding = self.tool_binding_repository.get_by_id(binding_id=binding_id)
        if binding is None:
            return None

        tool_version = self.tool_version_repository.get_by_id(
            tool_version_id=binding.tool_version_id)
        if tool_version is None:
            return None

        tool_definition = self.tool_definition_repository.get_by_id(
            tool_id=tool_version.tool_id)
        if tool_definition is None:
            return None

        return binding, tool_version, tool_definition

    def _build_tool_code(self, name: str) -> str:
        base = "".join(
            char.lower() if char.isalnum() else "_"
            for char in name
        ).strip("_") or "tool"
        return base[:64]

    def _normalize_mcp_config(
        self,
        payload: McpConnectRequestDto) -> tuple[str, dict[str, JSONValue]]:
        config = payload.config
        if len(config) == 1:
            server_name, value = next(iter(config.items()))
            if isinstance(value, dict):
                normalized = {str(key): item for key, item in value.items()}
                normalized.setdefault("server_name", server_name)
                normalized.setdefault("transport", "sse")
                return payload.name or server_name, normalized

        server_name_value = config.get("server_name") or config.get("serverName")
        server_name = str(server_name_value or payload.name or "mcp_server")
        normalized = {str(key): value for key, value in config.items()}
        normalized.setdefault("server_name", server_name)
        normalized.setdefault("transport", "sse")
        return server_name, normalized

    def _extract_mcp_tools(self, config: dict[str, JSONValue]) -> list[McpToolDto]:
        raw_tools = config.get("mcp_tools") or config.get("tools")
        if not isinstance(raw_tools, list):
            return []
        tools: list[McpToolDto] = []
        for item in raw_tools:
            if not isinstance(item, dict):
                continue
            name = item.get("name")
            if not isinstance(name, str) or not name:
                continue
            description = item.get("description")
            input_schema = item.get("inputSchema") or item.get("input_schema")
            tools.append(
                McpToolDto(
                    name=name,
                    description=description if isinstance(description, str) else None,
                    inputSchema=input_schema if isinstance(input_schema, dict) else None))
        return tools

    def _timeout_ms(self, config: dict[str, JSONValue]) -> int | None:
        timeout = config.get("timeout") or config.get("timeout_seconds")
        if isinstance(timeout, int | float):
            return int(timeout * 1000)
        return None