Merge pull request #157 from khawkins/workflow_customization

Making various workflow entities configurable

Author: khawkins

docs/9_mcp_workflow_engine_extraction/design.md

@@ -1,6 +1,6 @@
 {
   "name": "@salesforce/magen-mcp-workflow",
-  "version": "0.0.4",
+  "version": "0.0.5",
   "type": "module",
   "files": [
     "dist",

package-lock.json

@@ -77,6 +77,22 @@ export interface NodeGuidanceData<TResultSchema extends z.ZodObject<z.ZodRawShap
    * LLM compliance with the expected structure.
    */
   exampleOutput?: string;
+  /**
+   * 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/package.json

@@ -68,6 +68,7 @@ export { type BaseGraphConfig, type WorkflowRunnableConfig } from './common/grap
 export {
   BaseNode,
   AbstractToolNode,
+  UserInputExtractionNode,
   createGetUserInputNode,
   createUserInputExtractionNode,
   type GetUserInputNodeOptions,
@@ -78,7 +79,7 @@ export {
 export { CheckPropertiesFulfilledRouter } from './routers/index.js';
 
 // Base Service Classes
-export { AbstractService } from './services/index.js';
+export { AbstractService, InputExtractionService } from './services/index.js';
 
 // Checkpointing Infrastructure
 export {
@@ -93,6 +94,7 @@ export {
 export {
   OrchestratorTool,
   createOrchestratorToolMetadata,
+  type DefaultOrchestratorInputSchema,
   type OrchestratorConfig,
   type OrchestratorInput,
   type OrchestratorOutput,

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

@@ -204,7 +204,7 @@ export class InputExtractionService
    * @param propertiesToExtract - Array of properties to extract
    * @returns The guidance prompt string
    */
-  private generateTaskGuidance(
+  protected generateTaskGuidance(
     userUtterance: unknown,
     propertiesToExtract: Array<{ propertyName: string; description: string }>
   ): string {

packages/mcp-workflow/src/index.ts

@@ -5,14 +5,22 @@
  * For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/MIT
  */
 
+import z from 'zod';
 import { StateGraph } from '@langchain/langgraph';
 import { Logger } from '../../logging/logger.js';
 import { WorkflowStateManager } from '../../checkpointing/workflowStateManager.js';
+import type { DefaultOrchestratorInputSchema } from './metadata.js';
 
 /**
  * Orchestrator configuration interface
  *
- * Example usage:
+ * @template TInputSchema - The Zod input schema type for the orchestrator MCP tool.
+ *   Defaults to the standard ORCHESTRATOR_INPUT_SCHEMA. When providing a custom schema,
+ *   the OrchestratorTool subclass MUST also override extractUserInput() and
+ *   extractWorkflowStateData() to map the custom schema's properties to the
+ *   orchestrator's semantic needs.
+ *
+ * Example usage (default schema):
  * ```
  * const MyWorkflowState = Annotation.Root({ messages: Annotation<string[]> });
  * const workflow = new StateGraph(MyWorkflowState)
@@ -27,8 +35,26 @@ import { WorkflowStateManager } from '../../checkpointing/workflowStateManager.j
  *   workflow,
  * };
  * ```
+ *
+ * Example usage (custom schema):
+ * ```
+ * const MY_CUSTOM_SCHEMA = z.object({
+ *   payload: z.unknown().optional(),
+ *   sessionState: z.object({ thread_id: z.string() }).default({ thread_id: '' }),
+ * });
+ *
+ * const config: OrchestratorConfig<typeof MY_CUSTOM_SCHEMA> = {
+ *   toolId: 'my-orchestrator',
+ *   title: 'My Orchestrator',
+ *   description: 'Orchestrates my workflow',
+ *   workflow,
+ *   inputSchema: MY_CUSTOM_SCHEMA,
+ * };
+ * ```
  */
-export interface OrchestratorConfig {
+export interface OrchestratorConfig<
+  TInputSchema extends z.ZodObject<z.ZodRawShape> = DefaultOrchestratorInputSchema,
+> {
   /** Unique tool identifier for MCP registration */
   toolId: string;
 
@@ -48,6 +74,18 @@ export interface OrchestratorConfig {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   workflow: StateGraph<any, any, any, any, any, any, any, any>;
 
+  /**
+   * Custom Zod input schema for the orchestrator MCP tool.
+   *
+   * Optional - defaults to ORCHESTRATOR_INPUT_SCHEMA, which provides the standard
+   * `userInput` and `workflowStateData` properties.
+   *
+   * When providing a custom schema, the OrchestratorTool subclass MUST also override
+   * `extractUserInput()` and `extractWorkflowStateData()` to map the custom schema's
+   * properties to the orchestrator's semantic needs.
+   */
+  inputSchema?: TInputSchema;
+
   /**
    * Workflow state manager for checkpointing and persistence
    *

packages/mcp-workflow/src/services/inputExtractionService.ts

@@ -7,6 +7,7 @@
 
 export { type OrchestratorConfig } from './config.js';
 export {
+  type DefaultOrchestratorInputSchema,
   type OrchestratorInput,
   type OrchestratorOutput,
   type OrchestratorToolMetadata,

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

@@ -74,31 +74,49 @@ export const ORCHESTRATOR_OUTPUT_SCHEMA = z.object({
 
 export type OrchestratorOutput = z.infer<typeof ORCHESTRATOR_OUTPUT_SCHEMA>;
 
+/**
+ * Type alias for the default orchestrator input schema type.
+ *
+ * Used as the default generic parameter for OrchestratorConfig, OrchestratorToolMetadata,
+ * and OrchestratorTool to avoid circular imports between config.ts and metadata.ts.
+ */
+export type DefaultOrchestratorInputSchema = typeof ORCHESTRATOR_INPUT_SCHEMA;
+
 /**
  * Orchestrator tool metadata type
  * The metadata for the orchestrator tool (inputs/outputs)
+ *
+ * @template TInputSchema - The Zod input schema type. Defaults to ORCHESTRATOR_INPUT_SCHEMA.
  */
-export type OrchestratorToolMetadata = ToolMetadata<
-  typeof ORCHESTRATOR_INPUT_SCHEMA,
-  typeof ORCHESTRATOR_OUTPUT_SCHEMA
->;
+export type OrchestratorToolMetadata<
+  TInputSchema extends z.ZodObject<z.ZodRawShape> = DefaultOrchestratorInputSchema,
+> = ToolMetadata<TInputSchema, typeof ORCHESTRATOR_OUTPUT_SCHEMA>;
 
 /**
- * Factory function to create orchestrator tool metadata from configuration
+ * Factory function to create orchestrator tool metadata from configuration.
  * Takes the consumer-provided config and creates the tool metadata with
- * standardized input/output schemas.
+ * the appropriate input/output schemas.
+ *
+ * When a custom inputSchema is provided in the config, it is used instead of
+ * the default ORCHESTRATOR_INPUT_SCHEMA.
  *
+ * @template TInputSchema - The Zod input schema type. Defaults to ORCHESTRATOR_INPUT_SCHEMA.
  * @param config - The orchestrator configuration
- * @returns Tool metadata with the specified toolId, title, description
+ * @returns Tool metadata with the specified toolId, title, description, and schemas
  */
-export function createOrchestratorToolMetadata(
-  config: OrchestratorConfig
-): OrchestratorToolMetadata {
+export function createOrchestratorToolMetadata<
+  TInputSchema extends z.ZodObject<z.ZodRawShape> = DefaultOrchestratorInputSchema,
+>(config: OrchestratorConfig<TInputSchema>): OrchestratorToolMetadata<TInputSchema> {
+  // The cast is necessary because TypeScript cannot prove that when config.inputSchema
+  // is undefined, TInputSchema must be DefaultOrchestratorInputSchema (i.e. the generic
+  // default and the runtime default independently track the same fallback). The cast is
+  // safe: when inputSchema is omitted, TInputSchema defaults to typeof ORCHESTRATOR_INPUT_SCHEMA.
+  const effectiveInputSchema = (config.inputSchema ?? ORCHESTRATOR_INPUT_SCHEMA) as TInputSchema;
   return {
     toolId: config.toolId,
     title: config.title,
     description: config.description,
-    inputSchema: ORCHESTRATOR_INPUT_SCHEMA,
+    inputSchema: effectiveInputSchema,
     outputSchema: ORCHESTRATOR_OUTPUT_SCHEMA,
   };
 }