narratio/internal/app/analyze_checkpoint_recovery_test.go

package app

import (
	"context"
	"errors"
	"os"
	"path/filepath"
	"reflect"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/narratio/internal/adapters/scriptorium"
	"gitea.maximumdirect.net/eric/narratio/internal/config"
	"gitea.maximumdirect.net/eric/narratio/internal/manifest"
	"gitea.maximumdirect.net/eric/narratio/internal/stage"
)

func TestExecuteStagesRetriesFromLastAnalyzeCheckpoint(t *testing.T) {
	cfg := testConfig(t)
	cfg.Pipeline.Scriptorium = &config.ScriptoriumConfig{Artifacts: checkpointRecoveryArtifacts("original")}
	store := &manifest.LocalStore{}
	seed := manifest.New(cfg.Session.SessionID, time.Now().UTC())
	seed.Campaign = cfg.Session.Campaign
	for _, name := range []string{"prepare", "transcribe", "merge", "polish", "normalize", "trim"} {
		seed.MarkStageSucceeded(name, time.Now().UTC(), nil)
	}
	saveBoundedManifest(t, cfg, seed)
	analyze := mustSelectStage(t, "analyze")

	baselineRunner := &analyzeCheckpointRunner{}
	if _, err := executeStages(context.Background(), cfg, []stage.Stage{analyze}, RunOptions{
		Env: &Env{ManifestStore: store, Scriptorium: baselineRunner},
	}); err != nil {
		t.Fatalf("baseline executeStages() error = %v", err)
	}
	baseline, err := store.Load(context.Background(), manifestPathFor(cfg))
	if err != nil {
		t.Fatalf("load baseline manifest: %v", err)
	}
	baseline.MarkStageSucceeded("publish", time.Now().UTC(), nil)
	baseline.MarkStageSucceeded("notify", time.Now().UTC(), nil)
	if err := store.Save(context.Background(), manifestPathFor(cfg), baseline); err != nil {
		t.Fatalf("save delivery baseline: %v", err)
	}
	upstream := cloneStageRecords(baseline, "prepare", "transcribe", "merge", "polish", "normalize", "trim")
	unrelated := baseline.Stages["analyze"].AnalyzeArtifacts["unrelated"]

	cfg.Pipeline.Scriptorium.Artifacts = checkpointRecoveryArtifacts("revised")
	failure := errors.New("injected quest-log failure")
	failingRunner := &analyzeCheckpointRunner{FailPrompt: "quest-log-revised", Err: failure}
	if _, err := executeStages(context.Background(), cfg, []stage.Stage{analyze}, RunOptions{
		Env: &Env{ManifestStore: store, Scriptorium: failingRunner},
	}); !errors.Is(err, failure) {
		t.Fatalf("executeStages() error = %v, want quest-log failure", err)
	}
	if want := []string{"chronicle-revised", "encounters-revised", "quest-log-revised"}; !reflect.DeepEqual(failingRunner.Calls, want) {
		t.Fatalf("failed-run calls = %#v, want %#v", failingRunner.Calls, want)
	}

	failed := loadAnalyzeRecoveryManifest(t, store, cfg)
	failedAnalyze := failed.Stages["analyze"]
	if failedAnalyze.Status != manifest.StatusFailed || len(failedAnalyze.Outputs) != 0 {
		t.Fatalf("aggregate analyze = %#v, want failed without outputs", failedAnalyze)
	}
	assertAnalyzeArtifactStatus(t, failedAnalyze, "chronicle", manifest.AnalyzeArtifactCurrent)
	assertAnalyzeArtifactStatus(t, failedAnalyze, "encounters", manifest.AnalyzeArtifactCurrent)
	assertAnalyzeArtifactStatus(t, failedAnalyze, "quest_log", manifest.AnalyzeArtifactFailed)
	assertAnalyzeArtifactStatus(t, failedAnalyze, "player_handout", manifest.AnalyzeArtifactStale)
	if got := failedAnalyze.AnalyzeArtifacts["unrelated"]; !reflect.DeepEqual(got, unrelated) {
		t.Fatalf("unrelated artifact = %#v, want preserved %#v", got, unrelated)
	}
	for _, name := range []string{"publish", "notify"} {
		if got := failed.Stages[name]; got == nil || got.Status != manifest.StatusStale {
			t.Fatalf("%s stage = %#v, want stale", name, got)
		}
	}
	assertStageRecordsUnchanged(t, failed, upstream)

	retryRunner := &analyzeCheckpointRunner{}
	if _, err := executeStages(context.Background(), cfg, []stage.Stage{analyze}, RunOptions{
		Env: &Env{ManifestStore: store, Scriptorium: retryRunner},
	}); err != nil {
		t.Fatalf("retry executeStages() error = %v", err)
	}
	if want := []string{"quest-log-revised", "player-handout-revised"}; !reflect.DeepEqual(retryRunner.Calls, want) {
		t.Fatalf("retry calls = %#v, want %#v", retryRunner.Calls, want)
	}

	completed := loadAnalyzeRecoveryManifest(t, store, cfg)
	if completed.Stages["analyze"].Status != manifest.StatusSucceeded {
		t.Fatalf("aggregate analyze = %#v, want succeeded", completed.Stages["analyze"])
	}
	for _, key := range []string{"chronicle", "encounters", "player_handout", "quest_log", "unrelated"} {
		assertAnalyzeArtifactStatus(t, completed.Stages["analyze"], key, manifest.AnalyzeArtifactCurrent)
	}
	assertStageRecordsUnchanged(t, completed, upstream)
}

func TestExecuteStagesJoinsAnalyzeFailureWithTerminalSaveFailure(t *testing.T) {
	cfg := testConfig(t)
	cfg.Pipeline.Scriptorium = &config.ScriptoriumConfig{
		Artifacts: map[string]config.ScriptoriumArtifactConfig{
			"chronicle": {Enabled: true, PromptID: "chronicle", OutputPath: "artifacts/chronicle.md"},
		},
	}
	artifactErr := errors.New("injected analyze failure")
	terminalErr := errors.New("injected terminal save failure")
	store := &failedAnalyzeSaveStore{Store: &manifest.LocalStore{}, Err: terminalErr}
	runner := &analyzeCheckpointRunner{FailPrompt: "chronicle", Err: artifactErr}

	_, err := executeStages(context.Background(), cfg, []stage.Stage{mustSelectStage(t, "analyze")}, RunOptions{
		Env: &Env{ManifestStore: store, Scriptorium: runner},
	})
	if !errors.Is(err, artifactErr) || !errors.Is(err, terminalErr) {
		t.Fatalf("executeStages() error = %v, want analyze and terminal save failures", err)
	}
	if !store.Failed {
		t.Fatal("terminal save failure was not injected")
	}
}

func TestExecuteStagesPersistsAnalyzeCheckpointFailureRecovery(t *testing.T) {
	cfg := testConfig(t)
	cfg.Pipeline.Scriptorium = &config.ScriptoriumConfig{
		Artifacts: map[string]config.ScriptoriumArtifactConfig{
			"chronicle":  {Enabled: true, PromptID: "chronicle", OutputPath: "artifacts/chronicle.md"},
			"encounters": {Enabled: true, PromptID: "encounters", OutputPath: "artifacts/encounters.md"},
			"quest_log":  {Enabled: true, PromptID: "quest-log", OutputPath: "artifacts/quest_log.md"},
		},
	}
	baseStore := &manifest.LocalStore{}
	seed := manifest.New(cfg.Session.SessionID, time.Now().UTC())
	seed.Campaign = cfg.Session.Campaign
	seed.MarkStageSucceeded("publish", time.Now().UTC(), nil)
	seed.MarkStageSucceeded("notify", time.Now().UTC(), nil)
	saveBoundedManifest(t, cfg, seed)
	checkpointErr := errors.New("injected artifact checkpoint failure")
	store := &currentAnalyzeCheckpointFailingStore{
		Store:        baseStore,
		CurrentLimit: 2,
		Err:          checkpointErr,
	}
	runner := &analyzeCheckpointRunner{}

	_, err := executeStages(context.Background(), cfg, []stage.Stage{mustSelectStage(t, "analyze")}, RunOptions{
		Env: &Env{ManifestStore: store, Scriptorium: runner},
	})
	if !errors.Is(err, checkpointErr) {
		t.Fatalf("executeStages() error = %v, want checkpoint failure", err)
	}
	if want := []string{"chronicle", "encounters"}; !reflect.DeepEqual(runner.Calls, want) {
		t.Fatalf("adapter calls = %#v, want %#v", runner.Calls, want)
	}

	persisted := loadAnalyzeRecoveryManifest(t, baseStore, cfg)
	if persisted.Stages["analyze"].Status != manifest.StatusFailed {
		t.Fatalf("aggregate analyze = %#v, want failed", persisted.Stages["analyze"])
	}
	assertAnalyzeArtifactStatus(t, persisted.Stages["analyze"], "chronicle", manifest.AnalyzeArtifactCurrent)
	assertAnalyzeArtifactStatus(t, persisted.Stages["analyze"], "encounters", manifest.AnalyzeArtifactFailed)
	if _, exists := persisted.Stages["analyze"].AnalyzeArtifacts["quest_log"]; exists {
		t.Fatalf("unattempted quest log appeared: %#v", persisted.Stages["analyze"].AnalyzeArtifacts["quest_log"])
	}
	for _, name := range []string{"publish", "notify"} {
		if got := persisted.Stages[name]; got == nil || got.Status != manifest.StatusStale {
			t.Fatalf("%s stage = %#v, want stale", name, got)
		}
	}
}

type analyzeCheckpointRunner struct {
	Calls      []string
	FailPrompt string
	Err        error
}

type failedAnalyzeSaveStore struct {
	manifest.Store
	Err    error
	Failed bool
}

type currentAnalyzeCheckpointFailingStore struct {
	manifest.Store
	CurrentLimit int
	Err          error
	Failed       bool
}

func (s *currentAnalyzeCheckpointFailingStore) Save(ctx context.Context, path string, m *manifest.Manifest) error {
	if !s.Failed {
		current := 0
		if record := m.Stages["analyze"]; record != nil && record.Status == manifest.StatusRunning {
			for _, artifact := range record.AnalyzeArtifacts {
				if artifact.Status == manifest.AnalyzeArtifactCurrent {
					current++
				}
			}
		}
		if current >= s.CurrentLimit {
			s.Failed = true
			return s.Err
		}
	}
	return s.Store.Save(ctx, path, m)
}

func (s *failedAnalyzeSaveStore) Save(ctx context.Context, path string, m *manifest.Manifest) error {
	if record := m.Stages["analyze"]; !s.Failed && record != nil && record.Status == manifest.StatusFailed {
		s.Failed = true
		return s.Err
	}
	return s.Store.Save(ctx, path, m)
}

func (r *analyzeCheckpointRunner) RunArtifact(
	ctx context.Context,
	req scriptorium.RunArtifactRequest,
) (scriptorium.ArtifactResult, error) {
	if err := ctx.Err(); err != nil {
		return scriptorium.ArtifactResult{}, err
	}
	r.Calls = append(r.Calls, req.PromptID)
	if req.PromptID == r.FailPrompt {
		return scriptorium.ArtifactResult{}, r.Err
	}
	if err := os.MkdirAll(filepath.Dir(req.OutputPath), 0o755); err != nil {
		return scriptorium.ArtifactResult{}, err
	}
	if err := os.WriteFile(req.OutputPath, []byte("generated "+req.PromptID+"\n"), 0o644); err != nil {
		return scriptorium.ArtifactResult{}, err
	}
	return scriptorium.ArtifactResult{
		OutputPath:  req.OutputPath,
		CommandMode: scriptorium.CommandModeRun,
		PromptID:    req.PromptID,
		ProfileID:   req.ProfileID,
	}, nil
}

func (*analyzeCheckpointRunner) RenderArtifact(
	context.Context,
	scriptorium.RenderArtifactRequest,
) (scriptorium.ArtifactResult, error) {
	return scriptorium.ArtifactResult{}, errors.New("unexpected render")
}

func checkpointRecoveryArtifacts(revision string) map[string]config.ScriptoriumArtifactConfig {
	return map[string]config.ScriptoriumArtifactConfig{
		"chronicle": {
			Enabled: true, PromptID: "chronicle-" + revision, OutputPath: "artifacts/chronicle.md",
		},
		"encounters": {
			Enabled: true, DependsOn: []string{"chronicle"}, PromptID: "encounters-" + revision,
			OutputPath: "artifacts/encounters.md",
			Inputs: map[string]config.ScriptoriumInputConfig{
				"chronicle": {Source: "narratio.artifact.chronicle", Required: true},
			},
		},
		"quest_log": {
			Enabled: true, DependsOn: []string{"encounters"}, PromptID: "quest-log-" + revision,
			OutputPath: "artifacts/quest_log.md",
			Inputs: map[string]config.ScriptoriumInputConfig{
				"encounters": {Source: "narratio.artifact.encounters", Required: true},
			},
		},
		"player_handout": {
			Enabled: true, DependsOn: []string{"quest_log"}, PromptID: "player-handout-" + revision,
			OutputPath: "artifacts/player_handout.md",
			Inputs: map[string]config.ScriptoriumInputConfig{
				"quest_log": {Source: "narratio.artifact.quest_log", Required: true},
			},
		},
		"unrelated": {
			Enabled: true, PromptID: "unrelated", OutputPath: "artifacts/unrelated.md",
		},
	}
}

func mustSelectStage(t *testing.T, name string) stage.Stage {
	t.Helper()
	selected, err := stage.Select(name)
	if err != nil {
		t.Fatalf("Select(%q) error = %v", name, err)
	}
	return selected
}

func loadAnalyzeRecoveryManifest(t *testing.T, store manifest.Store, cfg *config.Config) *manifest.Manifest {
	t.Helper()
	m, err := store.Load(context.Background(), manifestPathFor(cfg))
	if err != nil {
		t.Fatalf("load session manifest: %v", err)
	}
	return m
}

func assertAnalyzeArtifactStatus(
	t *testing.T,
	record *manifest.StageRecord,
	key string,
	want manifest.AnalyzeArtifactStatus,
) {
	t.Helper()
	got, exists := record.AnalyzeArtifacts[key]
	if !exists || got.Status != want {
		t.Fatalf("artifact %q = %#v, want status %q", key, got, want)
	}
	if want != manifest.AnalyzeArtifactCurrent && got.Output != nil {
		t.Fatalf("artifact %q retained output authority: %#v", key, got)
	}
}

func cloneStageRecords(m *manifest.Manifest, names ...string) map[string]manifest.StageRecord {
	clones := make(map[string]manifest.StageRecord, len(names))
	for _, name := range names {
		clones[name] = *m.Stages[name]
	}
	return clones
}

func assertStageRecordsUnchanged(t *testing.T, m *manifest.Manifest, want map[string]manifest.StageRecord) {
	t.Helper()
	for name, expected := range want {
		if got := m.Stages[name]; got == nil || !reflect.DeepEqual(*got, expected) {
			t.Fatalf("upstream stage %q = %#v, want unchanged %#v", name, got, expected)
		}
	}
}