fmt

2026-05-20 23:29:53 +07:00
parent bc8dc138c7
commit ecfb64691e
36 changed files with 1191 additions and 1368 deletions
@@ -319,26 +319,22 @@ class WfMcpService:
        statuses: list[dict[str, Any]] = []
        for connection in self.connections.list_all():
            snapshot = self.store.load_catalog(connection.id)
-            statuses.append(
-                {
-                    "connection_id": connection.id,
-                    "server": connection.server,
-                    "account": connection.account,
-                    "enabled": connection.enabled,
-                    "has_snapshot": snapshot is not None,
-                    "fetched_at_epoch_ms": None
-                    if snapshot is None
-                    else snapshot.fetched_at_epoch_ms,
-                    "max_age_seconds": None
-                    if snapshot is None
-                    else snapshot.max_age_seconds,
-                    "node_count": 0 if snapshot is None else len(snapshot.nodes),
-                    "resource_count": 0
-                    if snapshot is None
-                    else len(snapshot.resources),
-                    "prompt_count": 0 if snapshot is None else len(snapshot.prompts),
-                }
-            )
+            statuses.append({
+                "connection_id": connection.id,
+                "server": connection.server,
+                "account": connection.account,
+                "enabled": connection.enabled,
+                "has_snapshot": snapshot is not None,
+                "fetched_at_epoch_ms": None
+                if snapshot is None
+                else snapshot.fetched_at_epoch_ms,
+                "max_age_seconds": None
+                if snapshot is None
+                else snapshot.max_age_seconds,
+                "node_count": 0 if snapshot is None else len(snapshot.nodes),
+                "resource_count": 0 if snapshot is None else len(snapshot.resources),
+                "prompt_count": 0 if snapshot is None else len(snapshot.prompts),
+            })
        return statuses

    def list_resources(
@@ -94,26 +94,20 @@ async def _refresh_all(service, connection_id: str | None) -> list[dict[str, Any
        try:
            await service.refresh_connection_catalog(target_id)
            snapshot = service.get_connection_snapshot(target_id)
-            results.append(
-                {
-                    "connection_id": target_id,
-                    "refreshed": snapshot is not None,
-                    "node_count": 0 if snapshot is None else len(snapshot.nodes),
-                    "resource_count": 0
-                    if snapshot is None
-                    else len(snapshot.resources),
-                    "prompt_count": 0 if snapshot is None else len(snapshot.prompts),
-                }
-            )
+            results.append({
+                "connection_id": target_id,
+                "refreshed": snapshot is not None,
+                "node_count": 0 if snapshot is None else len(snapshot.nodes),
+                "resource_count": 0 if snapshot is None else len(snapshot.resources),
+                "prompt_count": 0 if snapshot is None else len(snapshot.prompts),
+            })
        except Exception as exc:
-            results.append(
-                {
-                    "connection_id": target_id,
-                    "refreshed": False,
-                    "error_type": type(exc).__name__,
-                    "error": str(exc),
-                }
-            )
+            results.append({
+                "connection_id": target_id,
+                "refreshed": False,
+                "error_type": type(exc).__name__,
+                "error": str(exc),
+            })
    return results


@@ -138,18 +132,16 @@ def main(argv: list[str] | None = None) -> int:
    service = _service_from_config(args.config)

    if args.command == "connections":
-        _json_dump(
-            [
-                {
-                    "id": connection.id,
-                    "server": connection.server,
-                    "account": connection.account,
-                    "enabled": connection.enabled,
-                    "metadata": connection.metadata,
-                }
-                for connection in service.connections.list_all()
-            ]
-        )
+        _json_dump([
+            {
+                "id": connection.id,
+                "server": connection.server,
+                "account": connection.account,
+                "enabled": connection.enabled,
+                "metadata": connection.metadata,
+            }
+            for connection in service.connections.list_all()
+        ])
        return 0

    if args.command == "status":
@@ -162,12 +154,10 @@ def main(argv: list[str] | None = None) -> int:

    if args.command == "refresh":
        results = asyncio.run(_refresh_all(service, args.connection_id))
-        _json_dump(
-            {
-                "results": results,
-                "catalog": service.get_catalog().as_payload(),
-            }
-        )
+        _json_dump({
+            "results": results,
+            "catalog": service.get_catalog().as_payload(),
+        })
        if any(not result["refreshed"] for result in results):
            return 1
        return 0
@@ -38,12 +38,10 @@ def connection_to_fastmcp_server_config(
 def broker_config_to_fastmcp_config(config: BrokerConfig) -> MCPConfig:
    """Convert broker config into FastMCP's multi-server config object."""
    validate_proxy_config(config)
-    return MCPConfig.from_dict(
-        {
-            "mcpServers": {
-                connection.id: connection_to_fastmcp_server_config(connection)
-                for connection in config.connections
-                if connection.enabled
-            }
+    return MCPConfig.from_dict({
+        "mcpServers": {
+            connection.id: connection_to_fastmcp_server_config(connection)
+            for connection in config.connections
+            if connection.enabled
        }
-    )
+    })
@@ -46,13 +46,11 @@ class ResourceLinkRewritingTool(Tool):
        rewrite_uri: Callable[[str], str],
    ) -> ResourceLinkRewritingTool:
        """Copy one tool's public schema while replacing only execution."""
-        return cls.model_validate(
-            {
-                **tool.model_dump(),
-                "parent_tool": tool,
-                "rewrite_uri": rewrite_uri,
-            }
-        )
+        return cls.model_validate({
+            **tool.model_dump(),
+            "parent_tool": tool,
+            "rewrite_uri": rewrite_uri,
+        })


 class ResourceLinkNamespace(Transform):
@@ -214,21 +214,19 @@ class WorkflowSurfaceHandlers:
                query=query,
            ):
                continue
-            rows.append(
-                {
-                    "name": name,
-                    "source_id": "workflow",
-                    "kind": "wrapper_artifact",
-                    "artifact_id": artifact.id,
-                    "version": artifact.version,
-                    "title": artifact.title,
-                    "description": artifact.description,
-                    "outcomes": list(artifact.outcomes),
-                    "is_async": True,
-                    "input_fields": _schema_field_names(artifact.input_schema),
-                    "output_fields": _schema_field_names(artifact.output_schema),
-                }
-            )
+            rows.append({
+                "name": name,
+                "source_id": "workflow",
+                "kind": "wrapper_artifact",
+                "artifact_id": artifact.id,
+                "version": artifact.version,
+                "title": artifact.title,
+                "description": artifact.description,
+                "outcomes": list(artifact.outcomes),
+                "is_async": True,
+                "input_fields": _schema_field_names(artifact.input_schema),
+                "output_fields": _schema_field_names(artifact.output_schema),
+            })
        return rows

    def _wrapper_capability_detail(
@@ -447,12 +445,10 @@ class WorkflowSurfaceHandlers:
                },
            )
        )
-        required_sources = sorted(
-            {
-                capability.logical_source
-                for capability in workflow_artifact.required_capability_map().values()
-            }
-        )
+        required_sources = sorted({
+            capability.logical_source
+            for capability in workflow_artifact.required_capability_map().values()
+        })
        return {
            "artifact_id": workflow_artifact.id,
            "version": workflow_artifact.version,
@@ -909,16 +905,14 @@ def _available_sources(service: WfMcpService) -> list[AvailableSource]:
            if (capability_name := _capability_name(spec.name)) is not None
            if (detail := node_spec_details.get(spec.name)) is not None
        }
-        capabilities.update(
-            {
-                capability_name: AvailableCapability(
-                    name=capability_name,
-                    kind="reducer",
-                )
-                for reducer in source.capabilities.reducers.values()
-                if (capability_name := _capability_name(reducer.name)) is not None
-            }
-        )
+        capabilities.update({
+            capability_name: AvailableCapability(
+                name=capability_name,
+                kind="reducer",
+            )
+            for reducer in source.capabilities.reducers.values()
+            if (capability_name := _capability_name(reducer.name)) is not None
+        })
        sources.append(
            AvailableSource(
                id=source.id,
@@ -982,9 +976,9 @@ def _observed_node_specs(service: WfMcpService) -> dict[str, NodeSpecInventory]:
    observed: dict[str, NodeSpecInventory] = {}
    for source in service.capability_sources.values():
        inventory = source.as_inventory()
-        observed.update(
-            {detail.name: detail for detail in inventory.capabilities.node_spec_details}
-        )
+        observed.update({
+            detail.name: detail for detail in inventory.capabilities.node_spec_details
+        })
    return observed


@@ -1063,17 +1057,15 @@ def _artifact_capability_id(artifact: WorkflowArtifact) -> str:

 def _raw_plan_from_artifact(artifact: WorkflowArtifact) -> RawWorkflowPlan:
    """Validate the stored plan shape expected by the broker workflow runner."""
-    return RawWorkflowPlan.model_validate(
-        {
-            "name": _plan_field(artifact, "name"),
-            "input_schema": _plan_field(artifact, "input_schema"),
-            "state_schema": _plan_field(artifact, "state_schema"),
-            "output_schema": _plan_field(artifact, "output_schema"),
-            "start": _plan_field(artifact, "start"),
-            "nodes": _plan_field(artifact, "nodes"),
-            "edges": _plan_field(artifact, "edges"),
-        }
-    )
+    return RawWorkflowPlan.model_validate({
+        "name": _plan_field(artifact, "name"),
+        "input_schema": _plan_field(artifact, "input_schema"),
+        "state_schema": _plan_field(artifact, "state_schema"),
+        "output_schema": _plan_field(artifact, "output_schema"),
+        "start": _plan_field(artifact, "start"),
+        "nodes": _plan_field(artifact, "nodes"),
+        "edges": _plan_field(artifact, "edges"),
+    })


 def _plan_field(artifact: WorkflowArtifact, field_name: str) -> Any: