core/server/sandbox/workflow_module.go

package sandbox

// workflow_module.go — v0.30.2 CS1
//
// The workflow module lets extensions read workflow definitions and
// stage data, and programmatically advance or reject workflow stages.
//
// Starlark API:
//   wf = workflow.get_definition(workflow_id)  # returns dict
//   data = workflow.get_stage_data(channel_id) # returns dict
//   workflow.advance(channel_id)               # advance to next stage
//   workflow.reject(channel_id, reason)        # reject current stage

import (
	"context"
	"encoding/json"
	"fmt"

	"go.starlark.net/starlark"
	"go.starlark.net/starlarkstruct"

	"chat-switchboard/store"
)

// BuildWorkflowModule creates the "workflow" Starlark module for a package.
// Requires the workflow.access permission.
func BuildWorkflowModule(ctx context.Context, stores store.Stores) *starlarkstruct.Module {
	return MakeModule("workflow", starlark.StringDict{
		"get_definition": starlark.NewBuiltin("workflow.get_definition", workflowGetDef(ctx, stores)),
		"get_stage_data": starlark.NewBuiltin("workflow.get_stage_data", workflowGetStageData(ctx, stores)),
		"advance":        starlark.NewBuiltin("workflow.advance", workflowAdvance(ctx, stores)),
		"reject":         starlark.NewBuiltin("workflow.reject", workflowReject(ctx, stores)),
	})
}

func workflowGetDef(ctx context.Context, stores store.Stores) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
	return func(_ *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var workflowID string
		if err := starlark.UnpackPositionalArgs("workflow.get_definition", args, kwargs, 1, &workflowID); err != nil {
			return nil, err
		}
		wf, err := stores.Workflows.GetByID(ctx, workflowID)
		if err != nil {
			return starlark.None, fmt.Errorf("workflow.get_definition: %w", err)
		}

		stages, _ := stores.Workflows.ListStages(ctx, workflowID)

		// Build stages list
		stageList := make([]starlark.Value, 0, len(stages))
		for _, s := range stages {
			d := starlark.NewDict(8)
			d.SetKey(starlark.String("id"), starlark.String(s.ID))
			d.SetKey(starlark.String("name"), starlark.String(s.Name))
			d.SetKey(starlark.String("ordinal"), starlark.MakeInt(s.Ordinal))
			d.SetKey(starlark.String("stage_mode"), starlark.String(s.StageMode))
			d.SetKey(starlark.String("history_mode"), starlark.String(s.HistoryMode))
			d.SetKey(starlark.String("auto_transition"), starlark.Bool(s.AutoTransition))
			if s.PersonaID != nil {
				d.SetKey(starlark.String("persona_id"), starlark.String(*s.PersonaID))
			}
			if s.AssignmentTeamID != nil {
				d.SetKey(starlark.String("assignment_team_id"), starlark.String(*s.AssignmentTeamID))
			}
			if s.SurfacePkgID != nil {
				d.SetKey(starlark.String("surface_pkg_id"), starlark.String(*s.SurfacePkgID))
			}
			stageList = append(stageList, d)
		}

		result := starlark.NewDict(8)
		result.SetKey(starlark.String("id"), starlark.String(wf.ID))
		result.SetKey(starlark.String("name"), starlark.String(wf.Name))
		result.SetKey(starlark.String("slug"), starlark.String(wf.Slug))
		result.SetKey(starlark.String("entry_mode"), starlark.String(wf.EntryMode))
		result.SetKey(starlark.String("is_active"), starlark.Bool(wf.IsActive))
		result.SetKey(starlark.String("version"), starlark.MakeInt(wf.Version))
		result.SetKey(starlark.String("stages"), starlark.NewList(stageList))
		return result, nil
	}
}

func workflowGetStageData(ctx context.Context, stores store.Stores) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
	return func(_ *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var channelID string
		if err := starlark.UnpackPositionalArgs("workflow.get_stage_data", args, kwargs, 1, &channelID); err != nil {
			return nil, err
		}

		ws, err := stores.Channels.GetWorkflowStatus(ctx, channelID)
		if err != nil || ws == nil {
			return starlark.None, fmt.Errorf("workflow.get_stage_data: channel not found or not a workflow")
		}

		result := starlark.NewDict(8)
		result.SetKey(starlark.String("current_stage"), starlark.MakeInt(ws.CurrentStage))
		result.SetKey(starlark.String("status"), starlark.String(ws.Status))

		if ws.StageData != nil {
			var data map[string]any
			if json.Unmarshal(ws.StageData, &data) == nil {
				for k, v := range data {
					sv, err := goToStarlark(v)
					if err == nil {
						result.SetKey(starlark.String(k), sv)
					}
				}
			}
		}

		return result, nil
	}
}

func workflowAdvance(ctx context.Context, stores store.Stores) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
	return func(_ *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var channelID string
		if err := starlark.UnpackPositionalArgs("workflow.advance", args, kwargs, 1, &channelID); err != nil {
			return nil, err
		}

		ws, err := stores.Channels.GetWorkflowStatus(ctx, channelID)
		if err != nil || ws == nil || ws.WorkflowID == nil {
			return starlark.None, fmt.Errorf("workflow.advance: channel not found or not a workflow")
		}

		stages, err := stores.Workflows.ListStages(ctx, *ws.WorkflowID)
		if err != nil {
			return starlark.None, fmt.Errorf("workflow.advance: %w", err)
		}

		nextStage := ws.CurrentStage + 1
		if nextStage >= len(stages) {
			// Complete the workflow
			if err := stores.Channels.CompleteWorkflow(ctx, channelID, ws.CurrentStage, ws.StageData); err != nil {
				return starlark.None, fmt.Errorf("workflow.advance: %w", err)
			}
			return starlark.String("completed"), nil
		}

		if err := stores.Channels.AdvanceWorkflowStage(ctx, channelID, nextStage, ws.StageData); err != nil {
			return starlark.None, fmt.Errorf("workflow.advance: %w", err)
		}
		return starlark.String("advanced"), nil
	}
}

func workflowReject(ctx context.Context, stores store.Stores) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
	return func(_ *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var channelID, reason string
		if err := starlark.UnpackPositionalArgs("workflow.reject", args, kwargs, 2, &channelID, &reason); err != nil {
			return nil, err
		}

		ws, err := stores.Channels.GetWorkflowStatus(ctx, channelID)
		if err != nil || ws == nil {
			return starlark.None, fmt.Errorf("workflow.reject: channel not found or not a workflow")
		}

		if ws.CurrentStage <= 0 {
			return starlark.None, fmt.Errorf("workflow.reject: already at stage 0")
		}

		prevStage := ws.CurrentStage - 1
		if err := stores.Channels.RejectWorkflowToStage(ctx, channelID, prevStage); err != nil {
			return starlark.None, fmt.Errorf("workflow.reject: %w", err)
		}
		return starlark.String("rejected"), nil
	}
}