Applying returnGuidance in orchestrator

Author: khawkins

packages/mcp-workflow/src/common/metadata.ts

@@ -78,12 +78,19 @@ export interface NodeGuidanceData<TResultSchema extends z.ZodObject<z.ZodRawShap
    */
   exampleOutput?: string;
   /**
-   * Optional guidance for the LLM to return to the orchestrator.
-   * When provided, this guidance will replace the default "return to orchestrator"
-   * guidance that the orchestrator normally provides to the LLM, to return the
-   * workflow to the orchestrator after task completion. Make sure this custom
-   * guidance can properly return the workflow to the orchestrator, or the workflow
-   * will likely be broken.
+   * Optional custom guidance for the LLM to return results to the orchestrator.
+   *
+   * When provided, this replaces the orchestrator's default "return to orchestrator"
+   * prompt. The function receives only `workflowStateData` (the runtime session state
+   * that the producer doesn't have at construction time). The producer already owns
+   * `resultSchema` and `exampleOutput` as sibling properties on this same struct,
+   * so they can be captured in the closure if needed.
+   *
+   * Ensure this custom guidance properly instructs the LLM to return the workflow
+   * to the orchestrator, or the workflow will likely be broken.
+   *
+   * @param workflowStateData - The workflow state data to round-trip back to the orchestrator
+   * @returns The return guidance prompt string
    */
   returnGuidance?: (workflowStateData: WorkflowStateData) => string;
 }

packages/mcp-workflow/src/tools/orchestrator/orchestratorTool.ts

@@ -364,22 +364,13 @@ instructions for continuing the workflow.
     nodeGuidanceData: NodeGuidanceData<z.ZodObject<z.ZodRawShape>>,
     workflowStateData: WorkflowStateData
   ): string {
-    const resultSchemaJson = JSON.stringify(
-      zodToJsonSchema(nodeGuidanceData.resultSchema),
-      null,
-      2
-    );
-
-    // Build example section if provided
-    const exampleSection = nodeGuidanceData.exampleOutput
-      ? `
-For example, a properly formatted result should look like:
-
-\`\`\`json
-${nodeGuidanceData.exampleOutput}
-\`\`\`
-`
-      : '';
+    const returnGuidance = nodeGuidanceData.returnGuidance
+      ? nodeGuidanceData.returnGuidance(workflowStateData)
+      : this.defaultReturnGuidance(
+          workflowStateData,
+          nodeGuidanceData.resultSchema,
+          nodeGuidanceData.exampleOutput
+        );
 
     return `
 # ROLE
@@ -391,7 +382,28 @@ you with direct guidance for the current task.
 
 ${nodeGuidanceData.taskGuidance}
 
-# CRITICAL: REQUIRED NEXT STEP
+${returnGuidance}
+`;
+  }
+
+  private defaultReturnGuidance(
+    workflowStateData: WorkflowStateData,
+    resultSchema: z.ZodObject<z.ZodRawShape>,
+    exampleOutput?: string
+  ): string {
+    const resultSchemaJson = JSON.stringify(zodToJsonSchema(resultSchema), null, 2);
+
+    // Build example section if provided
+    const exampleSection = exampleOutput
+      ? `
+For example, a properly formatted result should look like:
+
+\`\`\`json
+${exampleOutput}
+\`\`\`
+`
+      : '';
+    return `# CRITICAL: REQUIRED NEXT STEP
 
 After completing the task above, you **MUST** invoke the \`${this.toolMetadata.toolId}\` tool 
 to continue the workflow. The workflow CANNOT proceed without this tool call.
@@ -424,7 +436,7 @@ Here is an example of the EXACT format your tool call should follow:
 \`\`\`
 Tool: ${this.toolMetadata.toolId}
 Parameters:
-  ${WORKFLOW_PROPERTY_NAMES.userInput}: ${nodeGuidanceData.exampleOutput ? nodeGuidanceData.exampleOutput.replaceAll('\n', '\n    ') : '{ /* your result object here */ }'}
+  ${WORKFLOW_PROPERTY_NAMES.userInput}: ${exampleOutput ? exampleOutput.replaceAll('\n', '\n    ') : '{ /* your result object here */ }'}
   ${WORKFLOW_PROPERTY_NAMES.workflowStateData}: ${JSON.stringify(workflowStateData)}
 \`\`\`
 `;

packages/mcp-workflow/tests/tools/orchestrator/orchestratorTool.test.ts

@@ -578,6 +578,62 @@ Gather user input for platform and projectName.
       expect(prompt).not.toContain('Invoke the following MCP server tool');
     });
 
+    it('should use custom returnGuidance when provided in NodeGuidanceData', async () => {
+      const orchestratorToolId = 'test-return-guidance-orchestrator';
+      const customReturnMarker = '## CUSTOM RETURN GUIDANCE MARKER';
+
+      const workflow = new StateGraph(TestState)
+        .addNode('guidedNode', (_state: State) => {
+          const nodeGuidanceData: NodeGuidanceData<typeof GET_INPUT_WORKFLOW_RESULT_SCHEMA> = {
+            nodeId: 'test-custom-return',
+            taskGuidance: 'Do something interesting.',
+            resultSchema: GET_INPUT_WORKFLOW_RESULT_SCHEMA,
+            exampleOutput: '{ "userUtterance": "example" }',
+            returnGuidance: workflowStateData =>
+              `${customReturnMarker}\nReturn to orchestrator with thread: ${workflowStateData.thread_id}`,
+          };
+          return interrupt(nodeGuidanceData);
+        })
+        .addEdge(START, 'guidedNode')
+        .addEdge('guidedNode', END);
+
+      const config: OrchestratorConfig = {
+        toolId: orchestratorToolId,
+        title: 'Return Guidance Test',
+        description: 'Tests custom returnGuidance in NodeGuidanceData',
+        workflow,
+        stateManager: new WorkflowStateManager({ environment: 'test' }),
+        logger: mockLogger,
+      };
+
+      const orchestrator = new OrchestratorTool(server, config);
+
+      const result = await orchestrator.handleRequest(
+        {
+          userInput: {},
+          workflowStateData: { thread_id: '' },
+        },
+        createMockExtra()
+      );
+
+      expect(result.structuredContent).toBeDefined();
+      const output = result.structuredContent as { orchestrationInstructionsPrompt: string };
+      const prompt = output.orchestrationInstructionsPrompt;
+
+      // Should contain the custom return guidance
+      expect(prompt).toContain(customReturnMarker);
+      expect(prompt).toContain('Return to orchestrator with thread: mmw-');
+
+      // Should still contain the task guidance structure
+      expect(prompt).toContain('# TASK GUIDANCE');
+      expect(prompt).toContain('Do something interesting.');
+
+      // Should NOT contain the default return guidance sections
+      expect(prompt).not.toContain('# CRITICAL: REQUIRED NEXT STEP');
+      expect(prompt).not.toContain('# OUTPUT FORMAT');
+      expect(prompt).not.toContain('# EXAMPLE TOOL CALL');
+    });
+
     it('should use delegate mode (orchestration prompt) when MCPToolInvocationData is used', async () => {
       const workflow = new StateGraph(TestState)
         .addNode('regularInterrupt', (_state: State) => {