Design Proposal: Extensions Pipeline for CompleteTaskStep

[mushrafmim]

This proposal introduces an Extensions Pipeline within CompleteTaskStep. The goal is to support custom behaviors—such as authentication/authorization, input schema validation, request filtering, and audit logging—without modifying or cluttering the core task orchestration logic.

---

1. Architectural Design & Lifecycle Phases

To maintain reliability, extensions are executed in two distinct phases relative to TaskDone (the point at which the task's child workflow is resumed in Temporal):

               [CompleteTaskStep(taskID, payload)]
                               │
                               ▼
                     [Load TaskRecord]
                               │
            ┌──────────────────┴──────────────────┐
            │ 1. PRE_RESUME Phase (Blocking)      │
            │  - Run Auth Gate                    │
            │  - Run Schema Validation            │
            │  - Run Sanitization (Mutates Data)  │
            └──────────────────┬──────────────────┘
                               │ (If all succeed)
                               ▼
              [Merge & Save Record to Database]
                               │
                               ▼
                  [Temporal.TaskDone (Resume)]
                               │
            ┌──────────────────┴──────────────────┐
            │ 2. POST_RESUME Phase (Non-Blocking) │
            │  - Run Audit Logging                │
            │  - Send Slack/Email Notifications   │
            │  - Dispatch Webhooks                │
            └─────────────────────────────────────┘

Phase A: PRE_RESUME (Before DB Save & Resume)

• Purpose: Handles validation, authentication, authorization, and data transformation before committing state changes.
• Execution Behavior:
  • Blocking: If any extension in this phase returns an error, execution halts immediately. The payload is not saved to the database, the child workflow is not resumed, and the error is returned to the API caller.
  • Mutable: Extensions can mutate the incoming payload map (e.g. trimming strings or injecting standard parameters).

Phase B: POST_RESUME (After DB Save & Resume)

• Purpose: Handles downstream side-effects (e.g., logging, notification triggers, indexing, external integrations) once execution is guaranteed to progress.
• Execution Behavior:
  • Non-Blocking: Errors returned by post-resume extensions are logged for visibility but do not rollback or abort the completed resume.
  • Immutable: The payload is read-only at this stage because it has already been persisted and passed to Temporal.
  • Asynchronous: Can run in a separate background goroutine to avoid adding latency to the client HTTP response.

---

2. Interface Definition

A single interface with a Phase parameter is recommended. This allows a single extension implementation to orchestrate actions and preserve context across both lifecycle phases if needed.

package extensions

import (
	"context"
	"encoding/json"
	"github.com/OpenNSW/core/taskflow/store"
)

type ExecutionPhase string

const (
	PhasePreResume  ExecutionPhase = "PRE_RESUME"
	PhasePostResume ExecutionPhase = "POST_RESUME"
)

type TaskExtension interface {
	// Execute runs the extension logic.
	// - phase: the current point in the CompleteTaskStep lifecycle.
	// - record: the task record (read-only in POST_RESUME).
	// - payload: incoming payload map (mutable ONLY in PRE_RESUME).
	// - properties: plugin configuration from the subtask template.
	Execute(ctx context.Context, phase ExecutionPhase, record *store.TaskRecord, payload map[string]any, properties json.RawMessage) error
}

---

3. Configuration & Registration

Extensions are configured directly in the subtask template and mapped inside the task orchestrator's registry.

SubTaskTemplate Schema

{
  "id": "submit_payment_details",
  "task_type": "USER_INPUT",
  "output_namespace": "payment_info",
  "extensions": [
    {
      "id": "role-auth-gate",
      "phase": "PRE_RESUME",
      "properties": {
        "required_roles": ["finance_officer"]
      }
    },
    {
      "id": "audit-logger",
      "phase": "POST_RESUME"
    }
  ]
}

Go Configuration Structs

type ExtensionConfig struct {
	ID         string          `json:"id"`
	Phase      ExecutionPhase  `json:"phase"`
	Properties json.RawMessage `json:"properties,omitempty"`
}

type SubTaskTemplate struct {
	ID               string            `json:"id"`
	TaskType         string            `json:"task_type"`
	PluginProperties json.RawMessage   `json:"plugin_properties"`
	OutputNamespace  string            `json:"output_namespace,omitempty"`
	Extensions       []ExtensionConfig `json:"extensions,omitempty"`
}

---

4. Execution Flow in CompleteTaskStep

func (tm *TaskManager) CompleteTaskStep(ctx context.Context, taskID string, payload map[string]any) error {
	record, exists := tm.db.GetTask(ctx, taskID)
	if !exists { return fmt.Errorf("task %s not found", taskID) }
	if record.State == "COMPLETED" { return fmt.Errorf("task %s already completed", taskID) }

	// 1. Run PRE_RESUME Extensions (Blocking, Mutable)
	for _, extCfg := range record.ActiveExtensions {
		if extCfg.Phase != PhasePreResume {
			continue
		}
		ext, registered := tm.extensionsRegistry[extCfg.ID]
		if !registered {
			return fmt.Errorf("extension %q configured but not registered", extCfg.ID)
		}
		if err := ext.Execute(ctx, PhasePreResume, &record, payload, extCfg.Properties); err != nil {
			return fmt.Errorf("pre-resume extension %q failed: %w", extCfg.ID, err) // Aborts execution
		}
	}

	// 2. Persist state
	if len(payload) > 0 && record.ActiveOutputNamespace != "" {
		record.Data[record.ActiveOutputNamespace] = payload
	}
	tm.db.SaveTask(ctx, record)

	// 3. Resume Temporal Child Workflow
	err := tm.taskWorkflowManager.TaskDone(ctx, record.TaskWorkflowID, record.TaskRunID, record.SubTaskNodeID, payload)
	if err != nil {
		return fmt.Errorf("failed to resume task workflow: %w", err)
	}

	// 4. Run POST_RESUME Extensions (Non-Blocking, Immutable, Async)
	go func(rec store.TaskRecord, pay map[string]any) {
		for _, extCfg := range rec.ActiveExtensions {
			if extCfg.Phase != PhasePostResume {
				continue
			}
			ext, registered := tm.extensionsRegistry[extCfg.ID]
			if !registered {
				log.Printf("[TaskManager] ERROR: post-resume extension %q not registered", extCfg.ID)
				continue
			}
			// Run extension in background; errors are logged, not returned to caller
			if err := ext.Execute(context.Background(), PhasePostResume, &rec, pay, extCfg.Properties); err != nil {
				log.Printf("[TaskManager] ERROR: post-resume extension %q failed: %v", extCfg.ID, err)
			}
		}
	}(record, payload)

	return nil
}