VLCode-Lite/related/workflow/bugfix_test.go

package workflow_test

import (
	"context"
	"encoding/json"
	"testing"

	"workflow"
)

// ---------------------------------------------------------------------------
// BUG-1: StepOutput string shorthand — "out": "$var" must deserialise to
//         {"$var": "=_result"} (spec §5.2.4)
// ---------------------------------------------------------------------------

// TestOutStringShorthandUnmarshal verifies that a JSON string value for "out"
// is accepted and expanded to the canonical map form.
func TestOutStringShorthandUnmarshal(t *testing.T) {
	raw := `{"id":"LLM_Answer","out":"$answer","next":"Stop_End"}`
	var step workflow.Step
	if err := json.Unmarshal([]byte(raw), &step); err != nil {
		t.Fatalf("unexpected unmarshal error: %v", err)
	}
	if len(step.Out) != 1 {
		t.Fatalf("expected Out to have 1 entry, got %d", len(step.Out))
	}
	if step.Out["$answer"] != "=_result" {
		t.Errorf("expected Out[\"$answer\"] == \"=_result\", got %q", step.Out["$answer"])
	}
}

// TestOutStringShorthandEndToEnd verifies that a workflow whose "out" field uses
// the string shorthand actually runs and stores _result into the target variable.
func TestOutStringShorthandEndToEnd(t *testing.T) {
	wf := &workflow.Workflow{
		Version: "3.13",
		Name:    "BUG1 Out Shorthand E2E",
		Registry: workflow.Registry{
			Vars: []string{"$answer(STRING)"},
		},
		Steps: []workflow.Step{
			{
				ID: "LLM_Answer",
				In: workflow.StepInput{
					"messages": []interface{}{
						map[string]interface{}{"role": "user", "content": "hi"},
					},
				},
				Out:  workflow.StepOutput{"$answer": "=_result"}, // canonical form
				Next: "Stop_End",
			},
			{ID: "Stop_End"},
		},
	}

	// Replicate the shorthand by round-tripping through JSON
	jsonBytes, err := json.Marshal(wf.Steps[0])
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}
	// Replace the canonical out object with its string shorthand
	var raw map[string]interface{}
	if err := json.Unmarshal(jsonBytes, &raw); err != nil {
		t.Fatalf("unmarshal to raw: %v", err)
	}
	raw["out"] = "$answer"
	patchedBytes, _ := json.Marshal(raw)
	var patchedStep workflow.Step
	if err := json.Unmarshal(patchedBytes, &patchedStep); err != nil {
		t.Fatalf("unmarshal patched step: %v", err)
	}
	wf.Steps[0] = patchedStep

	adapters := createTestAdapters()
	llm := adapters.LLM.(*workflow.DefaultLLMAdapter)
	llm.SetHandler(func(ctx context.Context, params map[string]interface{}, stream chan<- string) (map[string]interface{}, error) {
		return map[string]interface{}{
			"content":       "hello world",
			"finish_reason": "stop",
			"model":         "test-model",
		}, nil
	})

	eng, err := workflow.NewEngine(wf)
	if err != nil {
		t.Fatalf("NewEngine: %v", err)
	}
	result, err := eng.Execute(context.Background(), nil, adapters)
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	for range result.RunEventStream {
	}

	vars := result.Context.Variables
	if vars["$answer"] != "hello world" {
		t.Errorf("expected $answer == \"hello world\", got %v", vars["$answer"])
	}
}

// ---------------------------------------------------------------------------
// BUG-2: Branch → Stop_* must not trigger a second workflow_done
// ---------------------------------------------------------------------------

// makeBug2Workflow builds the workflow described in the bug report:
//
//	Branch_Go
//	  case true  → Stop_A   (the branch that fires)
//	  ELSE       → Stop_B   (never reached)
//	next         → Stop_C   (must NOT be reached after Stop_A fires)
func makeBug2Workflow() *workflow.Workflow {
	return &workflow.Workflow{
		Version: "3.13",
		Name:    "BUG2 Double WorkflowDone",
		Registry: workflow.Registry{
			Vars: []string{"$x(STRING)"},
		},
		Steps: []workflow.Step{
			{
				ID: "Branch_Go",
				Cases: [][]string{
					{"true", "Stop_A"},
					{"ELSE", "Stop_B"},
				},
				Next: "Stop_C",
			},
			{ID: "Stop_A"},
			{ID: "Stop_B"},
			{ID: "Stop_C"},
		},
	}
}

// TestBranchStopNoDoubleWorkflowDone verifies that when a Branch case directly
// targets a Stop_* node exactly one workflow_done event is emitted and the
// Branch's next pointer is never followed.
func TestBranchStopNoDoubleWorkflowDone(t *testing.T) {
	wf := makeBug2Workflow()

	eng, err := workflow.NewEngine(wf)
	if err != nil {
		t.Fatalf("NewEngine: %v", err)
	}
	result, err := eng.Execute(context.Background(), nil, createTestAdapters())
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}

	var events []workflow.RunEvent
	for ev := range result.RunEventStream {
		events = append(events, ev)
	}

	// Count workflow_done events — must be exactly 1
	doneCount := 0
	var doneStopID string
	for _, ev := range events {
		if ev.Type == workflow.RunEventWorkflowDone {
			doneCount++
			if sid, ok := ev.Payload["stop_id"].(string); ok {
				doneStopID = sid
			}
		}
	}
	if doneCount != 1 {
		t.Errorf("expected exactly 1 workflow_done event, got %d", doneCount)
		for _, ev := range events {
			t.Logf("  seq=%d type=%s step_id=%v payload=%v", ev.Seq, ev.Type, ev.StepID, ev.Payload)
		}
	}
	if doneStopID != "Stop_A" {
		t.Errorf("expected workflow_done stop_id=Stop_A, got %q", doneStopID)
	}

	// Stop_C must never have started
	for _, ev := range events {
		if ev.Type == workflow.RunEventStepStart && ev.StepID != nil && *ev.StepID == "Stop_C" {
			t.Errorf("Stop_C was started — it should have been suppressed after Stop_A fired")
		}
	}
}