diff --git a/src/wf_mcp/admin_surface/handlers/broker.py b/src/wf_mcp/admin_surface/handlers/broker.py index 16b4f6fc..0200d639 100644 --- a/src/wf_mcp/admin_surface/handlers/broker.py +++ b/src/wf_mcp/admin_surface/handlers/broker.py @@ -6,6 +6,13 @@ from wf_api import ( WorkflowSourceAdminApi, WorkflowSourceAdminSurface, ) +from wf_api.models import ( + AdminEventPayload, + ConnectionPayload, + ConnectionStatusPayload, + InspectSourceResult, + ListSourcesResult, +) from wf_mcp.broker.service import WfMcpService from wf_mcp.broker.service.workflow_operation_context import context_from_service from wf_mcp.shared.errors import error_payload @@ -24,11 +31,11 @@ class BrokerAdminHandlers: events=service.events, ) - async def list_connections(self) -> list[dict[str, Any]]: + async def list_connections(self) -> list[ConnectionPayload]: payload = await self.admin.list_connections() return payload["connections"] - async def get_connection_statuses(self) -> list[dict[str, Any]]: + async def get_connection_statuses(self) -> list[ConnectionStatusPayload]: payload = await self.admin.get_connection_statuses() return payload["statuses"] @@ -63,10 +70,10 @@ class BrokerAdminHandlers: *, cursor: str | None = None, limit: int = 50, - ) -> dict[str, Any]: + ) -> ListSourcesResult: return await self.sources.list_sources(cursor=cursor, limit=limit) - async def inspect_source(self, source_id: str) -> dict[str, Any]: + async def inspect_source(self, source_id: str) -> InspectSourceResult: return await self.sources.inspect_source(source_id=source_id) async def read_broker_resource(self, qualified_name: str) -> dict[str, Any]: @@ -97,6 +104,6 @@ class BrokerAdminHandlers: **error_payload(exc), } - async def get_broker_events(self) -> list[dict[str, Any]]: + async def get_broker_events(self) -> list[AdminEventPayload]: payload = await self.admin.list_events() return payload["events"] diff --git a/src/wf_mcp/admin_surface/tools.py b/src/wf_mcp/admin_surface/tools.py index 68191cba..05f3ee0d 100644 --- a/src/wf_mcp/admin_surface/tools.py +++ b/src/wf_mcp/admin_surface/tools.py @@ -4,6 +4,13 @@ from typing import Annotated, Any from pydantic import Field +from wf_api.models import ( + AdminEventPayload, + ConnectionPayload, + ConnectionStatusPayload, + InspectSourceResult, + ListSourcesResult, +) from wf_mcp.broker.service import WfMcpService from .handlers.broker import BrokerAdminHandlers @@ -45,7 +52,7 @@ def register_service_admin_tools( title="List Connections", description="List configured MCP connections known to this server.", ) - async def list_connections() -> list[dict[str, Any]]: + async def list_connections() -> list[ConnectionPayload]: return await handlers.list_connections() @server.tool( @@ -53,7 +60,7 @@ def register_service_admin_tools( title="Get Connection Statuses", description="Show configured MCP connection status and catalog counts.", ) - async def get_connection_statuses() -> list[dict[str, Any]]: + async def get_connection_statuses() -> list[ConnectionStatusPayload]: return await handlers.get_connection_statuses() @server.tool( @@ -106,7 +113,7 @@ def register_service_admin_tools( ), ), ] = 50, - ) -> dict[str, Any]: + ) -> ListSourcesResult: return await handlers.list_sources(cursor=cursor, limit=limit) @server.tool( @@ -124,7 +131,7 @@ def register_service_admin_tools( ) ), ], - ) -> dict[str, Any]: + ) -> InspectSourceResult: return await handlers.inspect_source(source_id) @server.tool( @@ -176,5 +183,5 @@ def register_service_admin_tools( title="Get Events", description="Return locally recorded broker/platform events.", ) - async def get_events() -> list[dict[str, Any]]: + async def get_events() -> list[AdminEventPayload]: return await handlers.get_broker_events() diff --git a/src/wf_mcp/broker/artifact_tools.py b/src/wf_mcp/broker/artifact_tools.py index a822168e..9dbc5c04 100644 --- a/src/wf_mcp/broker/artifact_tools.py +++ b/src/wf_mcp/broker/artifact_tools.py @@ -4,9 +4,19 @@ from typing import Any from mcp.server.fastmcp import FastMCP -from wf_api.models import RawWorkflowPlan +from wf_api.models import ( + DeleteDeploymentResult, + ListArtifactsResult, + ListDeploymentsResult, + RawWorkflowPlan, + SaveArtifactResult, + SaveDeploymentResult, + ValidateDeploymentResult, + WorkflowArtifactPayload, +) from ..workflow_surface import WorkflowSurfaceHandlers +from ..workflow_surface.models import RunDeploymentResult from .service import WfMcpService @@ -15,11 +25,11 @@ def register_artifact_tools(server: FastMCP, service: WfMcpService) -> None: handlers = WorkflowSurfaceHandlers(service) @server.tool() - async def list_workflow_artifacts() -> dict[str, Any]: + async def list_workflow_artifacts() -> ListArtifactsResult: return await handlers.list_artifacts() @server.tool() - async def save_workflow_artifact(artifact: dict[str, Any]) -> dict[str, Any]: + async def save_workflow_artifact(artifact: dict[str, Any]) -> SaveArtifactResult: return await handlers.save_artifact(artifact) @server.tool() @@ -33,7 +43,7 @@ def register_artifact_tools(server: FastMCP, service: WfMcpService) -> None: required_capabilities: dict[str, dict[str, Any]] | None = None, source_bindings: dict[str, str] | None = None, created_from_catalog_version: str | None = None, - ) -> dict[str, Any]: + ) -> SaveArtifactResult: return await handlers.create_artifact_from_plan( artifact_id=artifact_id, version=version, @@ -53,37 +63,45 @@ def register_artifact_tools(server: FastMCP, service: WfMcpService) -> None: async def inspect_workflow_artifact( artifact_id: str, version: int, - ) -> dict[str, Any]: + ) -> WorkflowArtifactPayload: return await handlers.inspect_artifact( artifact_id=artifact_id, version=version, ) @server.tool() - async def list_workflow_deployments() -> dict[str, Any]: + async def list_workflow_deployments() -> ListDeploymentsResult: return await handlers.list_deployments() @server.tool() - async def save_workflow_deployment(deployment: dict[str, Any]) -> dict[str, Any]: + async def save_workflow_deployment( + deployment: dict[str, Any], + ) -> SaveDeploymentResult: return await handlers.save_deployment(deployment) @server.tool() - async def delete_workflow_deployment(deployment_id: str) -> dict[str, Any]: + async def delete_workflow_deployment( + deployment_id: str, + ) -> DeleteDeploymentResult: """Delete one mutable deployment environment binding.""" return await handlers.delete_deployment(deployment_id=deployment_id) @server.tool() - async def validate_workflow_deployment(deployment_id: str) -> dict[str, Any]: + async def validate_workflow_deployment( + deployment_id: str, + ) -> ValidateDeploymentResult: return await handlers.validate_deployment(deployment_id=deployment_id) @server.tool() async def run_workflow_deployment( deployment_id: str, workflow_input: dict[str, Any], - ) -> dict[str, Any]: - return await handlers.run_deployment( - deployment_id=deployment_id, - workflow_input=workflow_input, + ) -> RunDeploymentResult: + return RunDeploymentResult.model_validate( + await handlers.run_deployment( + deployment_id=deployment_id, + workflow_input=workflow_input, + ) ) @server.tool() @@ -91,10 +109,12 @@ def register_artifact_tools(server: FastMCP, service: WfMcpService) -> None: run_id: str, resume_payload: dict[str, Any], resume_outcome: str = "submitted", - ) -> dict[str, Any]: + ) -> RunDeploymentResult: """Resume a durable interrupted deployment run.""" - return await handlers.resume_run( - run_id=run_id, - resume_payload=resume_payload, - resume_outcome=resume_outcome, + return RunDeploymentResult.model_validate( + await handlers.resume_run( + run_id=run_id, + resume_payload=resume_payload, + resume_outcome=resume_outcome, + ) )