refactor(workflow): enhance proxy and step execution handling (#2163)

* refactor(workflow): enhance proxy and step execution handling

- Updated the proxy routes to improve client connection management and error handling.
- Refactored the StepDetailPanel to integrate a new StepCodeSection for editing step code.
- Enhanced the WorkflowStepCard to reflect execution status based on completed steps.
- Modified polling logic in usePollingWorkflowExecution to streamline data retrieval.
- Bumped bindings package version to 1.0.7 and added completed_steps tracking in the workflow schema.

* revert

* wip

* wip

* bump mesh

* wips

Author: pedrofrxncx

apps/mesh/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@decocms/mesh",
-  "version": "1.1.0",
+  "version": "1.1.1",
   "description": "MCP Mesh - Self-hostable MCP Gateway for managing AI connections and tools",
   "author": "Deco team",
   "license": "MIT",

apps/mesh/src/api/routes/proxy.ts

@@ -17,30 +17,30 @@ import { getMonitoringConfig } from "@/core/config";
 import { getStableStdioClient } from "@/stdio/stable-transport";
 import {
   ConnectionEntity,
-  isStdioParameters,
   type HttpConnectionParameters,
+  isStdioParameters,
 } from "@/tools/connection/schema";
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js";
 import {
-  CallToolRequestSchema,
-  GetPromptRequestSchema,
-  ListPromptsRequestSchema,
-  ListResourcesRequestSchema,
-  ListResourceTemplatesRequestSchema,
-  ListToolsRequestSchema,
-  ReadResourceRequestSchema,
   type CallToolRequest,
+  CallToolRequestSchema,
   type CallToolResult,
   type GetPromptRequest,
+  GetPromptRequestSchema,
   type GetPromptResult,
+  ListPromptsRequestSchema,
   type ListPromptsResult,
+  ListResourcesRequestSchema,
   type ListResourcesResult,
+  ListResourceTemplatesRequestSchema,
   type ListResourceTemplatesResult,
+  ListToolsRequestSchema,
   type ListToolsResult,
   type ReadResourceRequest,
+  ReadResourceRequestSchema,
   type ReadResourceResult,
   type Tool,
 } from "@modelcontextprotocol/sdk/types.js";
@@ -452,7 +452,7 @@ async function createMCPProxyDoNotUseDirectly(
             throw error;
           } finally {
             // Close client - stdio connections ignore close() via stable-transport
-            await client.close();
+            client.close().catch(console.error);
           }
         },
       );
@@ -475,43 +475,76 @@ async function createMCPProxyDoNotUseDirectly(
     }
 
     // Fall back to client for connections without indexed tools
-    const client = await createClient();
-    return await client.listTools();
+    let client: Awaited<ReturnType<typeof createClient>> | undefined;
+    try {
+      client = await createClient();
+      return await client.listTools();
+    } finally {
+      client?.close().catch(console.error);
+    }
   };
 
   // List resources from downstream connection
   const listResources = async (): Promise<ListResourcesResult> => {
-    const client = await createClient();
-    return await client.listResources();
+    let client: Awaited<ReturnType<typeof createClient>> | undefined;
+    try {
+      client = await createClient();
+      return await client.listResources();
+    } finally {
+      client?.close().catch(console.error);
+    }
   };
 
   // Read a specific resource from downstream connection
   const readResource = async (
     params: ReadResourceRequest["params"],
   ): Promise<ReadResourceResult> => {
-    const client = await createClient();
-    return await client.readResource(params);
+    let client: Awaited<ReturnType<typeof createClient>> | undefined;
+    try {
+      client = await createClient();
+      return await client.readResource(params);
+    } finally {
+      client?.close().catch(console.error);
+    }
   };
 
   // List resource templates from downstream connection
   const listResourceTemplates =
     async (): Promise<ListResourceTemplatesResult> => {
-      const client = await createClient();
-      return await client.listResourceTemplates();
+      let client: Awaited<ReturnType<typeof createClient>> | undefined;
+      try {
+        client = await createClient();
+        return await client.listResourceTemplates();
+      } finally {
+        client?.close().catch(console.error);
+      }
     };
 
   // List prompts from downstream connection
   const listPrompts = async (): Promise<ListPromptsResult> => {
-    const client = await createClient();
-    return await client.listPrompts();
+    let client: Awaited<ReturnType<typeof createClient>> | undefined;
+    try {
+      client = await createClient();
+      return await client.listPrompts();
+    } catch (error) {
+      console.error("[proxy:listPrompts] Error listing prompts:", error);
+      throw error;
+    } finally {
+      client?.close().catch(console.error);
+    }
   };
 
   // Get a specific prompt from downstream connection
   const getPrompt = async (
     params: GetPromptRequest["params"],
   ): Promise<GetPromptResult> => {
-    const client = await createClient();
-    return await client.getPrompt(params);
+    let client: Awaited<ReturnType<typeof createClient>> | undefined;
+    try {
+      client = await createClient();
+      return await client.getPrompt(params);
+    } finally {
+      client?.close().catch(console.error);
+    }
   };
 
   // Call tool using fetch directly for streaming support

apps/mesh/src/web/components/details/workflow/components/step-detail-panel.tsx

@@ -30,7 +30,8 @@ import type { JsonSchema } from "@/web/utils/constants";
 import { MonacoCodeEditor } from "./monaco-editor";
 import type { Step, ToolCallAction } from "@decocms/bindings/workflow";
 import { useMcp } from "@/web/hooks/use-mcp";
-import { usePollingWorkflowExecution } from "../hooks";
+import { useExecutionCompletedStep } from "../hooks";
+import { useState } from "react";
 
 interface StepDetailPanelProps {
   className?: string;
@@ -101,6 +102,7 @@ export function StepDetailPanel({ className }: StepDetailPanelProps) {
       <InputSection step={currentStep} />
       <OutputSection step={currentStep} />
       <TransformCodeSection step={currentStep} />
+      <StepCodeSection step={currentStep} />
     </div>
   );
 }
@@ -203,7 +205,7 @@ function InputSection({ step }: { step: Step }) {
   const isToolStep = "toolName" in step.action;
   const toolName =
     isToolStep && "toolName" in step.action ? step.action.toolName : null;
-
+  const trackingExecutionId = useTrackingExecutionId();
   const { tool } = useGatewayTool(toolName ?? "");
 
   if (!tool || !tool.inputSchema) {
@@ -220,7 +222,7 @@ function InputSection({ step }: { step: Step }) {
     <Accordion
       type="single"
       collapsible
-      defaultValue="input"
+      defaultValue={trackingExecutionId ? "output" : "input"}
       className="border-b border-border shrink-0"
     >
       <AccordionItem value="input" className="border-b-0">
@@ -231,6 +233,7 @@ function InputSection({ step }: { step: Step }) {
         </AccordionTrigger>
         <AccordionContent className="px-5 pt-2">
           <ToolInput
+            key={step.name}
             inputSchema={tool.inputSchema as JsonSchema}
             inputParams={step.input as Record<string, unknown>}
             setInputParams={handleInputChange}
@@ -249,11 +252,11 @@ function InputSection({ step }: { step: Step }) {
 function OutputSection({ step }: { step: Step }) {
   const outputSchema = step.outputSchema;
   const trackingExecutionId = useTrackingExecutionId();
-  const { step_results } = usePollingWorkflowExecution(trackingExecutionId);
-  const stepResult = step_results?.find(
-    (result) => result.step_id === step.name,
+  const { output, error } = useExecutionCompletedStep(
+    trackingExecutionId,
+    step.name,
   );
-  const output = stepResult?.output;
+  const content = output ? output : error ? { error: error } : null;
 
   // Always show the Output section (even if empty)
   const properties =
@@ -283,8 +286,8 @@ function OutputSection({ step }: { step: Step }) {
             <div className="text-sm text-muted-foreground italic">
               No output schema defined
             </div>
-          ) : output ? (
-            <OutputMonacoEditor output={output} />
+          ) : content ? (
+            <OutputMonacoEditor output={content} />
           ) : (
             <div className="space-y-2">
               {propertyEntries.map(([key, propSchema]) => (
@@ -458,6 +461,10 @@ export default async function(input: Input): Promise<Output> {
     });
   };
 
+  if ("code" in step.action && step.action.code) {
+    return null;
+  }
+
   // No transform code → show collapsed with Plus
   if (!hasTransformCode) {
     return (
@@ -503,6 +510,60 @@ export default async function(input: Input): Promise<Output> {
   );
 }
 
+function StepCodeSection({ step }: { step: Step }) {
+  const [isOpen, setIsOpen] = useState(false);
+  const { updateStep } = useWorkflowActions();
+  const code =
+    "code" in step.action && step.action.code ? step.action.code : null;
+  const trackingExecutionId = useTrackingExecutionId();
+  const handleCodeSave = (
+    code: string,
+    outputSchema: Record<string, unknown> | null,
+  ) => {
+    updateStep(step.name, {
+      action: {
+        ...step.action,
+        code: code,
+      },
+      ...(outputSchema ? { outputSchema } : {}),
+    });
+  };
+  if (!code) {
+    return null;
+  }
+  return (
+    <div className="flex-1 flex flex-col min-h-0">
+      <div
+        className="p-5 shrink-0 cursor-pointer hover:bg-accent/50 transition-colors"
+        onClick={() => setIsOpen(!isOpen)}
+      >
+        <div className="flex items-center justify-between">
+          <h3 className="text-sm font-medium text-muted-foreground">
+            Step Code
+          </h3>
+          {isOpen ? (
+            <Minus size={14} className="text-muted-foreground" />
+          ) : (
+            <Plus size={14} className="text-muted-foreground" />
+          )}
+        </div>
+      </div>
+      {isOpen && (
+        <div className="flex-1 min-h-120 h-full">
+          <MonacoCodeEditor
+            onSave={(code, outputSchema) => handleCodeSave(code, outputSchema)}
+            key={`step-code-${step.name}-${trackingExecutionId}`}
+            code={code}
+            language="typescript"
+            height="100%"
+            readOnly={trackingExecutionId !== undefined}
+          />
+        </div>
+      )}
+    </div>
+  );
+}
+
 // Helper function to convert JSON Schema types to TypeScript types
 function jsonSchemaTypeToTS(schema: JsonSchema): string {
   if (Array.isArray(schema.type)) {

apps/mesh/src/web/components/details/workflow/components/workflow-step-card.tsx

@@ -51,6 +51,10 @@ export function WorkflowStepCard({
   const { item: executionItem } =
     usePollingWorkflowExecution(trackingExecutionId);
   const executionCompletedAt = executionItem?.completed_at_epoch_ms;
+  const isCompletedSuccessfully =
+    executionItem?.completed_steps?.success?.includes(step.name) ?? false;
+  const isCompletedWithError =
+    executionItem?.completed_steps?.error?.includes(step.name) ?? false;
 
   const isToolStep = "toolName" in step.action;
   const connectionId =
@@ -62,7 +66,11 @@ export function WorkflowStepCard({
   const hasToolSelected = Boolean(toolName);
   const outputSchemaProperties = getOutputSchemaProperties(step);
 
-  const status = executionStatus?.status;
+  const status = isCompletedSuccessfully
+    ? "success"
+    : isCompletedWithError
+      ? "error"
+      : executionStatus?.status;
   const isTracking = executionStatus !== undefined;
   const hasStatusIndicator = status === "success" || status === "error";
 

apps/mesh/src/web/components/details/workflow/hooks/derived/use-resolved-refs.ts

@@ -3,22 +3,16 @@ import { useTrackingExecutionId } from "../../stores/workflow";
 
 export function useResolvedRefs() {
   const trackingExecutionId = useTrackingExecutionId();
-  const { step_results, item: executionItem } =
+  const { item: executionItem } =
     usePollingWorkflowExecution(trackingExecutionId);
   const resolvedRefs: Record<string, unknown> | undefined =
-    trackingExecutionId && step_results
+    trackingExecutionId && executionItem
       ? (() => {
           const refs: Record<string, unknown> = {};
           // Add workflow input as "input"
           if (executionItem?.input) {
             refs["input"] = executionItem.input;
           }
-          // Add each step's output by step_id
-          for (const result of step_results) {
-            if (result.step_id && result.output !== undefined) {
-              refs[result.step_id as string] = result.output;
-            }
-          }
           return refs;
         })()
       : undefined;