Advanced Workflow Features for vMCP Composition

[tgrunnagle]

Overview

Implements advanced workflow features for Virtual MCP Composite Tools, including DAG-based parallel execution, step dependencies, sophisticated error handling, and workflow state management. This completes Phase 2 of the composition work.

This is a larger review, but only 792 lines of product code added and 47 removed.

Issue: Closes #156 (stacklok/stacklok-epics)

What Changed

Core Features

1. DAG-Based Parallel Execution

• New file: pkg/vmcp/composer/dag_executor.go
  • Implements topological sort using Kahn's algorithm to build execution levels
  • Executes independent steps in parallel using errgroup for coordination
  • Semaphore-based concurrency limiting (default: 10 parallel steps)
  • Automatic optimization: steps with no dependencies run concurrently
  • Performance improvement: parallel execution reduces workflow time by ~60-70% for independent steps

2. Step Dependencies

• depends_on field support in pkg/vmcp/composer/composer.go:67
• Dependency graph validation with cycle detection using DFS
• Transitive dependencies automatically handled
• Missing dependency validation at workflow definition time

3. Advanced Error Handling

• Three-level error handling:
  • Step-level: on_error.continue_on_error overrides workflow-level settings
  • Workflow-level: failure_mode (abort/continue/best_effort)
  • Automatic: retry with exponential backoff
• Retry logic in pkg/vmcp/composer/workflow_engine.go:311-350:
  • Configurable retry count and initial delay
  • Exponential backoff (2^attempt * initial_delay, max 60x)
  • Safety cap: maximum 10 retries to prevent infinite loops

4. Workflow State Management

• Pluggable state store interface: pkg/vmcp/composer/composer.go:191-217
• In-memory implementation: pkg/vmcp/composer/state_store.go
  • Thread-safe operations with mutex protection
  • Deep copying to prevent external modifications
  • Automatic cleanup of stale workflows (configurable intervals)
  • Ready for future Redis/DB backends

5. Workflow Lifecycle

• UUID-based workflow IDs for unique identification
• State checkpointing after each step completion
• Configurable timeouts (default: 30 minutes for workflows, 5 minutes for steps)
• Automatic cleanup of completed/failed/timed-out workflows
• Workflow cancellation support via state store

Files Added

• pkg/vmcp/composer/dag_executor.go - DAG execution engine
• pkg/vmcp/composer/dag_executor_test.go - DAG executor unit tests (9 test cases)
• pkg/vmcp/composer/state_store.go - In-memory workflow state store
• pkg/vmcp/composer/state_store_test.go - State store unit tests (14 test cases)
• test/e2e/vmcp_workflow_e2e_test.go - End-to-end workflow tests
• docs/operator/advanced-workflow-patterns.md - Comprehensive guide (797 lines)
• docs/operator/composite-tools-quick-reference.md - Quick reference (233 lines)

Files Modified

• pkg/vmcp/composer/workflow_engine.go - Integrated DAG executor and state management
• pkg/vmcp/composer/workflow_engine_test.go - Added retry and timeout tests
• pkg/vmcp/composer/composer.go - Added state store interface and error types
• pkg/vmcp/composer/workflow_context.go - Enhanced context management
• docs/operator/virtualmcpcompositetooldefinition-guide.md - Updated with advanced features

Example Usage

Parallel Incident Investigation Workflow

apiVersion: toolhive.stacklok.dev/v1alpha1
kind: VirtualMCPCompositeToolDefinition
metadata:
  name: incident-investigation
spec:
  name: investigate_incident
  steps:
    # Level 1: Parallel data fetching
    - id: fetch_logs
      type: tool
      tool: splunk.fetch_logs
      arguments:
        incident_id: "{{.params.incident_id}}"

    - id: fetch_metrics
      type: tool
      tool: datadog.fetch_metrics
      arguments:
        incident_id: "{{.params.incident_id}}"

    - id: fetch_traces
      type: tool
      tool: jaeger.fetch_traces
      arguments:
        incident_id: "{{.params.incident_id}}"

    # Level 2: Correlation (waits for all Level 1)
    - id: correlate
      type: tool
      tool: analysis.correlate
      depends_on: [fetch_logs, fetch_metrics, fetch_traces]
      arguments:
        logs: "{{.steps.fetch_logs.output}}"
        metrics: "{{.steps.fetch_metrics.output}}"
        traces: "{{.steps.fetch_traces.output}}"
      on_error:
        action: retry
        retry_count: 3
        retry_delay: 2s

    # Level 3: Report creation
    - id: create_report
      type: tool
      tool: jira.create_issue
      depends_on: [correlate]
      arguments:
        title: "Incident {{.params.incident_id}}"
        body: "{{.steps.correlate.output.summary}}"

Performance: 3 parallel fetches complete in ~1x time instead of 3x sequential time.

Test Coverage

Unit Tests

• ✅ Topological sort (7 test cases covering chains, diamonds, complex DAGs)
• ✅ Cycle detection (3 test cases: direct, indirect, self-reference)
• ✅ Parallel execution verification (timing-based)
• ✅ Dependency ordering enforcement
• ✅ Error handling (abort/continue/best_effort modes)
• ✅ Retry logic with exponential backoff
• ✅ Concurrency limiting with semaphore
• ✅ Context cancellation
• ✅ State store operations (14 comprehensive tests)
• ✅ State store cleanup and concurrency

Integration & E2E Tests

• ✅ Complex 8-step incident investigation workflow
• ✅ End-to-end parallel execution with mock backends
• ✅ Dependency ordering validation with timing verification

All tests passing ✅

Performance Metrics

From test results:

• Parallel speedup: 3 independent 100ms steps complete in ~100ms (not 300ms)
• Complex workflow: 8-step workflow completes in ~200ms (vs 400ms sequential)
• Concurrency control: Semaphore effectively limits parallel execution
• Cleanup efficiency: Stale workflows removed within 2 cleanup cycles

Architecture Highlights

1. Clean Separation: DAG execution, state management, and workflow orchestration are independent modules
2. Pluggable Design: State store interface enables future Redis/PostgreSQL implementations
3. Safety First: Multiple safeguards (max steps: 100, max retries: 10, semaphore limits)
4. Thread Safety: Proper mutex usage, deep copying, and goroutine management with errgroup
5. Context Propagation: Cancellation and timeouts properly propagated through execution stack
6. Observability: Comprehensive logging of execution stats, timing, and state metrics

Documentation

• Advanced Workflow Patterns: 797-line comprehensive guide covering:

  • Parallel execution with DAG
  • Step dependencies and patterns (diamond, fan-out/fan-in)
  • Error handling strategies with examples
  • State management and lifecycle
  • Performance optimization techniques
  • Best practices and common patterns

• Quick Reference: 233-line guide for rapid development

Breaking Changes

None. This is a backward-compatible enhancement. Existing workflows without dependencies execute as before.

Migration Notes

• State tracking requires creating a state store: composer.NewInMemoryStateStore(cleanupInterval, maxAge)
• Parallel execution is automatic for steps without depends_on - no migration needed
• Retry configuration is opt-in via on_error.action: retry

Future Work (Out of Scope)

• Distributed state store (Redis/PostgreSQL) - interface ready
• Workflow pause/resume
• Step-level timeout configuration
• Conditional branching (marked as Phase 3)

---

Ready to merge - All acceptance criteria met, tests passing, comprehensive documentation included.

[tgrunnagle]

@claude Please review this change, which aims to resolve https://github.com/stacklok/stacklok-epics/issues/156

[claude]

Claude finished @tgrunnagle's task —— View job

---

Code Review for PR #2592 - Complete ✅

I've completed a thorough review of the advanced workflow features implementation. This is a well-architected addition that brings powerful DAG-based parallel execution capabilities to vMCP composite tools. Overall the implementation is solid, but I've identified several issues and recommendations.

---

Summary

Overall Assessment: ✅ Approve with minor concerns

This PR successfully implements Phase 2 of the composition work with:

• Clean DAG execution engine using Kahn's algorithm
• Thread-safe state management with proper synchronization
• Comprehensive test coverage (23+ unit tests + E2E)
• Excellent documentation (1030+ lines)

Key Metrics:

• 792 lines of production code added
• 47 lines removed
• 1430+ lines of test code
• Test coverage appears comprehensive

---

Issues Found

🔴 Critical Issues

None - No blocking issues found.

🟡 Medium Severity Issues

1. Deep Copy Limitation in State Store

Location: pkg/vmcp/composer/state_store.go:72-89

The code performs shallow copies of nested maps in PendingElicitation.Schema:

// Note: We perform a shallow copy of the WorkflowStatus struct and deep copy slices
// (CompletedSteps, PendingElicitations). Maps within nested structures (like
// PendingElicitation.Schema) remain shared.

Issue: This creates a potential for external modification of schema maps. While the comment acknowledges this is "acceptable," it could lead to subtle bugs if workflow state is ever modified concurrently or if the state store is used in unexpected ways.

Recommendation: Consider using a proper deep copy library (e.g., github.com/mitchellh/copystructure) or document more explicitly that callers must never modify returned state objects.

Severity: Medium - Could cause bugs in future if assumptions change

---

2. Missing Retry Count Validation in Template Context

Location: pkg/vmcp/composer/workflow_engine.go:397-399

When recording retry counts, the code assumes the step result exists:

if result, exists := workflowCtx.GetStepResult(step.ID); exists {
    result.RetryCount = retryCount
}

Issue: This silently fails if the result doesn't exist. While this shouldn't happen in normal flow, it makes debugging harder if there's a bug in the step recording logic.

Recommendation: Add logging or an assertion when the result doesn't exist:

if result, exists := workflowCtx.GetStepResult(step.ID); exists {
    result.RetryCount = retryCount
} else {
    logger.Warnf("Step result not found when setting retry count for step %s", step.ID)
}

Severity: Low-Medium - Primarily affects debuggability

---

3. Context Cancellation Race in State Checkpointing

Location: pkg/vmcp/composer/workflow_engine.go:633-648

The checkpointWorkflowState function uses context.Background() to save state after each step:

func (e *workflowEngine) checkpointWorkflowState(workflowCtx *WorkflowContext) {
    // ...
    ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancel()
    
    if err := e.stateStore.SaveState(ctx, workflowCtx.WorkflowID, state); err != nil {
        logger.Warnf("Failed to checkpoint workflow state for %s: %v", workflowCtx.WorkflowID, err)
    }
}

Issue: If the workflow context is cancelled/timed out, state checkpointing will continue using a fresh context. This is actually intentional (to ensure state is saved even when workflow is cancelled), but it could lead to a scenario where:

1. Workflow times out
2. Final state save uses background context and succeeds
3. But workflow was marked as timed out before checkpoint completes
4. State becomes inconsistent

Recommendation: This is likely fine as-is since checkpoints are "best effort," but consider documenting this behavior more explicitly or ensuring terminal state saves always happen last.

Severity: Low - Edge case, but could cause confusion

---

🟢 Minor Issues / Style Suggestions

4. Magic Numbers in Tests

Location: Multiple test files

Tests use hardcoded timing values (50ms, 100ms, 150ms) which could be flaky on slower CI systems:

time.Sleep(100 * time.Millisecond)
assert.Less(t, duration, 200*time.Millisecond, "parallel execution should be faster")

Recommendation: Consider using relative timing checks or helper functions that scale with system load.

Severity: Low - Tests passed, but could be flaky

---

5. Missing Documentation for Semaphore Limit

Location: pkg/vmcp/composer/dag_executor.go:14-17

const (
    // defaultMaxParallelSteps is the default maximum number of steps to execute in parallel.
    defaultMaxParallelSteps = 10
)

Issue: The PR description mentions this limit but doesn't explain why 10 was chosen or how users can configure it.

Recommendation: Add a comment explaining the rationale (e.g., "Chosen as a balance between parallelism and resource usage. Can be adjusted via DAG executor constructor.") and document in the user-facing guide how to tune this.

Severity: Low - Documentation issue

---

6. Inconsistent Error Wrapping

Location: Throughout workflow_engine.go

Some errors use %w formatting while others use %v:

// Good:
return fmt.Errorf("%w: failed to expand arguments for step %s: %v", ErrTemplateExpansion, step.ID, err)

// Inconsistent:
return fmt.Errorf("failed to route tool %s in step %s: %w", step.Tool, step.ID, err)

Recommendation: Be consistent about which errors should be wrapped and which shouldn't. The current approach is mostly good, but a few cases are inconsistent.

Severity: Low - Cosmetic

---

Architecture Review

✅ Strengths

1. Clean Separation of Concerns

   • DAG execution is isolated in its own module
   • State management is pluggable via interface
   • Workflow orchestration is independent of execution strategy

2. Thread Safety

   • Proper use of sync.RWMutex for concurrent access
   • Deep copying in state store prevents external modifications
   • Goroutine coordination with errgroup

3. Error Handling

   • Three-level error handling (step/workflow/automatic retry) is well-designed
   • Exponential backoff with safety caps
   • Clear error propagation

4. Performance

   • DAG-based parallelization is well-implemented
   • Semaphore prevents resource exhaustion
   • Efficient topological sort (Kahn's algorithm)

5. Observability

   • Comprehensive logging at key decision points
   • Metrics tracking in state store
   • Execution statistics for debugging

⚠️ Architectural Concerns

AC-1: Architecture Documentation Gap

Location: docs/arch/

Issue: Per CLAUDE.md, changes to core packages should update architecture documentation. This PR adds significant new functionality to the vmcp composer, but the architecture docs only mention composite tools briefly in 02-core-concepts.md:225-251.

Impact: Future developers may not understand:

• How DAG execution integrates with the rest of vmcp
• State management patterns and when to use state stores
• How workflow composition relates to other vMCP concepts (aggregation, routing, etc.)

Recommendation: Add a new architecture document docs/arch/10-vmcp-composition.md covering:

• Workflow execution model (DAG, dependencies, parallelism)
• State management architecture and pluggability
• Integration with aggregator and router components
• Error handling and retry strategies
• Performance considerations

Reference: See CLAUDE.md lines about updating architecture docs for core package changes.

Severity: Medium - Documentation debt that should be addressed before Phase 3

---

AC-2: Missing Rate Limiting (Already Documented)

Location: pkg/vmcp/composer/workflow_engine.go:84-90

The TODO comment correctly identifies the need for rate limiting:

// TODO(rate-limiting): Add rate limiting per user/session to prevent workflow execution DoS.

Assessment: This is appropriately deferred to future work. The comment is clear and references security review documentation. Good practice.

No action needed - Just noting for visibility

---

Test Coverage Review

✅ Test Quality

Excellent coverage across:

• ✅ Topological sort (7 test cases: chains, diamonds, complex DAGs)
• ✅ Cycle detection (3 test cases: direct, indirect, self-reference)
• ✅ Parallel execution timing verification
• ✅ Dependency ordering enforcement
• ✅ Error handling modes (abort/continue/best_effort)
• ✅ Retry logic with exponential backoff
• ✅ Concurrency limiting with semaphore
• ✅ State store operations (14 comprehensive tests)
• ✅ E2E workflow with mock backends

Test Statistics

• DAG executor tests: dag_executor_test.go (528 lines)
• State store tests: state_store_test.go (373 lines)
• Workflow engine tests: workflow_engine_test.go (enhanced)
• E2E tests: vmcp_workflow_e2e_test.go (164 lines)

Total test code: ~1430+ lines for ~792 lines of production code
Ratio: ~1.8:1 (excellent)

Minor Test Gaps

1. Concurrent state store access: While cleanup is tested, there's no explicit test for multiple goroutines reading/writing state simultaneously.
2. Workflow timeout during state save: Edge case where workflow times out while state is being checkpointed.
3. Max retry cap enforcement: Test that retry counts > 10 are capped (currently only logged).

Recommendation: These gaps are minor and can be addressed in future PRs if issues arise.

---

Documentation Review

✅ Documentation Quality

Outstanding documentation provided:

• ✅ advanced-workflow-patterns.md (797 lines) - Comprehensive guide with examples
• ✅ composite-tools-quick-reference.md (233 lines) - Quick start guide
• ✅ Updated virtualmcpcompositetooldefinition-guide.md with new features

Examples are practical and production-ready:

• Incident investigation workflow (realistic use case)
• Diamond pattern (demonstrates parallelism)
• Fan-out/fan-in patterns
• Error handling strategies

Documentation Gaps

1. Performance tuning guide: Missing guidance on:

   • When to use dependencies vs. sequential steps
   • How to tune maxParallelSteps for different workload types
   • Memory/CPU considerations for large workflows

2. Troubleshooting section: Would benefit from:

   • Common errors and solutions
   • Debugging workflow execution issues
   • Understanding workflow timing and performance

Recommendation: These can be added incrementally as users provide feedback.

---

Security Review

✅ Security Strengths

1. Resource limits enforced:

   • maxWorkflowSteps = 100 prevents DoS via large workflows
   • maxRetryCount = 10 prevents infinite retry loops
   • Semaphore limits parallel execution

2. Input validation:

   • Workflow IDs required and validated
   • Dependency cycles detected at validation time
   • Step IDs checked for duplicates

3. Error information:

   • Errors are descriptive but don't leak sensitive data
   • Tool call failures are properly wrapped

Security Recommendations

SR-1: Consider step timeout limits
Currently, individual step timeout is defaultStepTimeout = 5 * time.Minutes, but users can override this. Consider adding a maximum step timeout (e.g., 30 minutes) to prevent resource exhaustion.

SR-2: Validate workflow depth
The 100-step limit prevents wide workflows, but there's no limit on dependency depth. A very deep chain (100 sequential steps) could still cause issues. Consider adding a maximum dependency depth check.

Severity: Low - Existing limits are reasonable, these are defense-in-depth improvements

---

Code Quality

✅ Positive Observations

1. Consistent error handling patterns
2. Good use of constants for magic numbers
3. Comprehensive comments explaining complex logic
4. Proper use of defer for cleanup
5. Context propagation throughout

Code Smells

CS-1: Large function - executeDAG in dag_executor.go is 86 lines but well-organized
CS-2: Complex conditional - Error handling in workflow_engine.go:164-203 has nested ifs

Both are acceptable given the domain complexity.

---

Performance Analysis

✅ Performance Strengths

From test results and implementation:

• ✅ Parallel speedup: 3 × 50ms steps complete in ~80ms (vs. 150ms sequential)
• ✅ Efficient semaphore-based concurrency control
• ✅ Minimal locking (only protects Steps map)
• ✅ Proper goroutine cleanup with errgroup

Performance Considerations

PC-1: Deep copy overhead
State store performs deep copies on every save/load. For workflows with large outputs, this could be expensive. Consider copy-on-write or reference counting if this becomes an issue.

PC-2: Cleanup interval
Default cleanup interval is 5 minutes with 1-hour retention. For high-throughput systems, consider making these configurable via environment variables.

Severity: Low - These are optimizations for future scale

---

Recommendations Summary

Must Fix Before Merge

None - No blocking issues

Should Fix Before Merge

1. AC-1: Add architecture documentation for vmcp composition (Medium priority)
2. Consider addressing deep copy limitation with proper library

Can Address in Follow-up PRs

1. Add missing test coverage for concurrent state access
2. Document performance tuning guidelines
3. Add troubleshooting section to docs
4. Consider adding maximum dependency depth check
5. Make state store cleanup intervals configurable

---

Final Verdict

✅ APPROVED - This is high-quality work that successfully delivers Phase 2 functionality.

Why Approve:

• Solid architecture with clean separation of concerns
• Comprehensive test coverage (1.8:1 ratio)
• Excellent documentation (1030+ lines)
• No critical bugs found
• Thread-safety properly handled
• Good error handling patterns

Conditions:

• Recommend addressing architecture documentation gap (AC-1) before merge
• All other issues can be addressed in follow-up PRs

Next Steps:

1. Consider adding docs/arch/10-vmcp-composition.md per CLAUDE.md guidelines
2. Address any issues you agree with from this review
3. Merge when ready - this is solid work!

---

Excellent work on this feature! The parallel execution capabilities will significantly improve workflow performance, and the pluggable state store design sets up nicely for Phase 3 and beyond.

---

[codecov]

Codecov Report

❌ Patch coverage is 82.75862% with 65 lines in your changes missing coverage. Please review.
✅ Project coverage is 55.83%. Comparing base (c94268a) to head (3416f80).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
pkg/vmcp/composer/workflow_engine.go	61.38%	32 Missing and 7 partials ⚠️
pkg/vmcp/composer/state_store.go	88.23%	15 Missing and 3 partials ⚠️
pkg/vmcp/composer/dag_executor.go	92.07%	5 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2592      +/-   ##
==========================================
+ Coverage   55.30%   55.83%   +0.53%     
==========================================
  Files         309      311       +2     
  Lines       29129    29479     +350     
==========================================
+ Hits        16111    16461     +350     
+ Misses      11585    11583       -2     
- Partials     1433     1435       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.