From ced3780787b018ccd19ee96a26ad4c5b8c02ae29 Mon Sep 17 00:00:00 2001 From: NVIDIAN Date: Sun, 12 Apr 2026 22:13:53 -0700 Subject: [PATCH] refactor: migrate mcp_server from marshal_with/api.model to Pydantic BaseModel (#34935) Co-authored-by: ai-hpc Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- api/controllers/console/app/mcp_server.py | 67 ++++++++++++------ .../console/app/test_mcp_server_response.py | 70 +++++++++++++++++++ 2 files changed, 117 insertions(+), 20 deletions(-) create mode 100644 api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py diff --git a/api/controllers/console/app/mcp_server.py b/api/controllers/console/app/mcp_server.py index 412fc8795a..9454d28bcf 100644 --- a/api/controllers/console/app/mcp_server.py +++ b/api/controllers/console/app/mcp_server.py @@ -1,23 +1,27 @@ import json +from datetime import datetime +from typing import Any -from flask_restx import Resource, marshal_with -from pydantic import BaseModel, Field +from flask_restx import Resource +from pydantic import BaseModel, Field, field_validator from sqlalchemy import select from werkzeug.exceptions import NotFound +from controllers.common.schema import register_schema_models from controllers.console import console_ns from controllers.console.app.wraps import get_app_model from controllers.console.wraps import account_initialization_required, edit_permission_required, setup_required from extensions.ext_database import db -from fields.app_fields import app_server_fields +from fields.base import ResponseModel from libs.login import current_account_with_tenant, login_required from models.enums import AppMCPServerStatus from models.model import AppMCPServer -DEFAULT_REF_TEMPLATE_SWAGGER_2_0 = "#/definitions/{model}" -# Register model for flask_restx to avoid dict type issues in Swagger -app_server_model = console_ns.model("AppServer", app_server_fields) +def _to_timestamp(value: datetime | int | None) -> int | None: + if isinstance(value, datetime): + return int(value.timestamp()) + return value class MCPServerCreatePayload(BaseModel): @@ -32,8 +36,33 @@ class MCPServerUpdatePayload(BaseModel): status: str | None = Field(default=None, description="Server status") -for model in (MCPServerCreatePayload, MCPServerUpdatePayload): - console_ns.schema_model(model.__name__, model.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0)) +class AppMCPServerResponse(ResponseModel): + id: str + name: str + server_code: str + description: str + status: str + parameters: dict[str, Any] | list[Any] | str + created_at: int | None = None + updated_at: int | None = None + + @field_validator("parameters", mode="before") + @classmethod + def _parse_json_string(cls, value: Any) -> Any: + if isinstance(value, str): + try: + return json.loads(value) + except (json.JSONDecodeError, TypeError): + return value + return value + + @field_validator("created_at", "updated_at", mode="before") + @classmethod + def _normalize_timestamp(cls, value: datetime | int | None) -> int | None: + return _to_timestamp(value) + + +register_schema_models(console_ns, MCPServerCreatePayload, MCPServerUpdatePayload, AppMCPServerResponse) @console_ns.route("/apps//server") @@ -41,27 +70,27 @@ class AppMCPServerController(Resource): @console_ns.doc("get_app_mcp_server") @console_ns.doc(description="Get MCP server configuration for an application") @console_ns.doc(params={"app_id": "Application ID"}) - @console_ns.response(200, "MCP server configuration retrieved successfully", app_server_model) + @console_ns.response(200, "Server configuration", console_ns.models[AppMCPServerResponse.__name__]) @login_required @account_initialization_required @setup_required @get_app_model - @marshal_with(app_server_model) def get(self, app_model): server = db.session.scalar(select(AppMCPServer).where(AppMCPServer.app_id == app_model.id).limit(1)) - return server + if server is None: + return {} + return AppMCPServerResponse.model_validate(server, from_attributes=True).model_dump(mode="json") @console_ns.doc("create_app_mcp_server") @console_ns.doc(description="Create MCP server configuration for an application") @console_ns.doc(params={"app_id": "Application ID"}) @console_ns.expect(console_ns.models[MCPServerCreatePayload.__name__]) - @console_ns.response(201, "MCP server configuration created successfully", app_server_model) + @console_ns.response(200, "Server created", console_ns.models[AppMCPServerResponse.__name__]) @console_ns.response(403, "Insufficient permissions") @account_initialization_required @get_app_model @login_required @setup_required - @marshal_with(app_server_model) @edit_permission_required def post(self, app_model): _, current_tenant_id = current_account_with_tenant() @@ -82,20 +111,19 @@ class AppMCPServerController(Resource): ) db.session.add(server) db.session.commit() - return server + return AppMCPServerResponse.model_validate(server, from_attributes=True).model_dump(mode="json") @console_ns.doc("update_app_mcp_server") @console_ns.doc(description="Update MCP server configuration for an application") @console_ns.doc(params={"app_id": "Application ID"}) @console_ns.expect(console_ns.models[MCPServerUpdatePayload.__name__]) - @console_ns.response(200, "MCP server configuration updated successfully", app_server_model) + @console_ns.response(200, "Server updated", console_ns.models[AppMCPServerResponse.__name__]) @console_ns.response(403, "Insufficient permissions") @console_ns.response(404, "Server not found") @get_app_model @login_required @setup_required @account_initialization_required - @marshal_with(app_server_model) @edit_permission_required def put(self, app_model): payload = MCPServerUpdatePayload.model_validate(console_ns.payload or {}) @@ -118,7 +146,7 @@ class AppMCPServerController(Resource): except ValueError: raise ValueError("Invalid status") db.session.commit() - return server + return AppMCPServerResponse.model_validate(server, from_attributes=True).model_dump(mode="json") @console_ns.route("/apps//server/refresh") @@ -126,13 +154,12 @@ class AppMCPServerRefreshController(Resource): @console_ns.doc("refresh_app_mcp_server") @console_ns.doc(description="Refresh MCP server configuration and regenerate server code") @console_ns.doc(params={"server_id": "Server ID"}) - @console_ns.response(200, "MCP server refreshed successfully", app_server_model) + @console_ns.response(200, "Server refreshed", console_ns.models[AppMCPServerResponse.__name__]) @console_ns.response(403, "Insufficient permissions") @console_ns.response(404, "Server not found") @setup_required @login_required @account_initialization_required - @marshal_with(app_server_model) @edit_permission_required def get(self, server_id): _, current_tenant_id = current_account_with_tenant() @@ -145,4 +172,4 @@ class AppMCPServerRefreshController(Resource): raise NotFound() server.server_code = AppMCPServer.generate_server_code(16) db.session.commit() - return server + return AppMCPServerResponse.model_validate(server, from_attributes=True).model_dump(mode="json") diff --git a/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py b/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py new file mode 100644 index 0000000000..baac4cd4e0 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/app/test_mcp_server_response.py @@ -0,0 +1,70 @@ +import datetime + +from controllers.console.app.mcp_server import AppMCPServerResponse + + +class TestAppMCPServerResponse: + def test_parameters_json_string_parsed(self): + data = { + "id": "s1", + "name": "test", + "server_code": "code", + "description": "desc", + "status": "active", + "parameters": '{"key": "value"}', + } + resp = AppMCPServerResponse.model_validate(data) + assert resp.parameters == {"key": "value"} + + def test_parameters_invalid_json_returns_original(self): + data = { + "id": "s1", + "name": "test", + "server_code": "code", + "description": "desc", + "status": "active", + "parameters": "not-valid-json", + } + resp = AppMCPServerResponse.model_validate(data) + assert resp.parameters == "not-valid-json" + + def test_parameters_dict_passthrough(self): + data = { + "id": "s1", + "name": "test", + "server_code": "code", + "description": "desc", + "status": "active", + "parameters": {"already": "parsed"}, + } + resp = AppMCPServerResponse.model_validate(data) + assert resp.parameters == {"already": "parsed"} + + def test_timestamps_normalized(self): + dt = datetime.datetime(2024, 1, 1, 0, 0, 0, tzinfo=datetime.UTC) + data = { + "id": "s1", + "name": "test", + "server_code": "code", + "description": "desc", + "status": "active", + "parameters": {}, + "created_at": dt, + "updated_at": dt, + } + resp = AppMCPServerResponse.model_validate(data) + assert resp.created_at == int(dt.timestamp()) + assert resp.updated_at == int(dt.timestamp()) + + def test_timestamps_none(self): + data = { + "id": "s1", + "name": "test", + "server_code": "code", + "description": "desc", + "status": "active", + "parameters": {}, + } + resp = AppMCPServerResponse.model_validate(data) + assert resp.created_at is None + assert resp.updated_at is None