executus/run/final_guard_test.go

package run

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

	"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
	"gitea.stevedudenhoeffer.com/steve/majordomo/provider/fake"

	"gitea.stevedudenhoeffer.com/steve/executus/tool"
)

// guardFunc adapts a func to FinalGuard, counting calls and capturing the
// snapshot it received.
type guardFunc struct {
	nudge func(info RunInfo, state FinalState) string
	calls int
	info  RunInfo
	state FinalState
}

func (g *guardFunc) FinalNudge(_ context.Context, info RunInfo, state FinalState) string {
	g.calls++
	g.info = info
	g.state = state
	return g.nudge(info, state)
}

// eventRecorder is captureRecorder plus LogEvent capture (for the final_nudge
// audit event).
type eventRecorder struct {
	captureRecorder
	events []string
}

func (r *eventRecorder) LogEvent(eventType string, _ map[string]any) {
	r.events = append(r.events, eventType)
}

// TestFinalGuardNudgeRunsOneExtraRound is the core announce-then-stop fix: the
// model stops with "Done — sent!" (nothing delivered), the guard returns a
// nudge, and the SAME conversation runs one more round — here a tool call then
// a corrected final answer. The guard is consulted exactly once (the nudge
// round's own stop is final), and the audit log records the nudge.
func TestFinalGuardNudgeRunsOneExtraRound(t *testing.T) {
	fp := fake.New("fake")
	fp.Enqueue("m",
		fake.Reply("Done — sent you the file!"),
		// Nudge round: an (unregistered → tolerated error result) delivery call,
		// then the corrected final answer.
		fake.ReplyWith(llm.Response{ToolCalls: []llm.ToolCall{{ID: "c1", Name: "send_attachments", Arguments: []byte(`{}`)}}}),
		fake.Reply("delivered for real"),
	)
	m, _ := fp.Model("m")
	rec := &eventRecorder{}
	guard := &guardFunc{nudge: func(RunInfo, FinalState) string { return "you did not deliver the file" }}
	ex := New(Config{
		Registry: tool.NewRegistry(),
		Models:   func(ctx context.Context, _ string) (context.Context, llm.Model, error) { return ctx, m, nil },
		Ports: Ports{
			Audit:      auditFunc{start: func(RunInfo) RunRecorder { return rec }},
			FinalGuard: guard,
		},
	})

	res := ex.Run(context.Background(),
		RunnableAgent{Name: "g", ModelTier: "m"},
		tool.Invocation{RunID: "r-nudge", CallerID: "c"}, "make me a file")

	if res.Err != nil {
		t.Fatalf("run error: %v", res.Err)
	}
	if res.Output != "delivered for real" {
		t.Fatalf("output = %q, want the nudge round's answer", res.Output)
	}
	if guard.calls != 1 {
		t.Fatalf("guard consulted %d times, want exactly 1 (no re-guard of the nudge round)", guard.calls)
	}
	if guard.info.RunID != "r-nudge" {
		t.Errorf("guard info.RunID = %q", guard.info.RunID)
	}
	if guard.state.Output != "Done — sent you the file!" {
		t.Errorf("guard state.Output = %q, want the pre-nudge final text", guard.state.Output)
	}
	var sawNudgeEvent bool
	for _, ev := range rec.events {
		if ev == "final_nudge" {
			sawNudgeEvent = true
		}
	}
	if !sawNudgeEvent {
		t.Errorf("audit events %v missing final_nudge", rec.events)
	}
	if rec.stats.Output != "delivered for real" {
		t.Errorf("audit close output = %q, want the nudged output", rec.stats.Output)
	}
}

// TestFinalGuardEmptyNudgeNoExtraRound: a guard returning "" adds zero model
// calls — the common path stays one round.
func TestFinalGuardEmptyNudgeNoExtraRound(t *testing.T) {
	fp := fake.New("fake")
	fp.Enqueue("m", fake.Reply("plain answer"))
	m, _ := fp.Model("m")
	guard := &guardFunc{nudge: func(RunInfo, FinalState) string { return "" }}
	ex := New(Config{
		Registry: tool.NewRegistry(),
		Models:   func(ctx context.Context, _ string) (context.Context, llm.Model, error) { return ctx, m, nil },
		Ports:    Ports{FinalGuard: guard},
	})

	res := ex.Run(context.Background(),
		RunnableAgent{ModelTier: "m"}, tool.Invocation{RunID: "r"}, "hi")

	if res.Err != nil {
		t.Fatalf("run error: %v", res.Err)
	}
	if res.Output != "plain answer" {
		t.Fatalf("output = %q", res.Output)
	}
	if guard.calls != 1 {
		t.Errorf("guard consulted %d times, want 1", guard.calls)
	}
	if n := fp.CallCount("m"); n != 1 {
		t.Errorf("model called %d times, want 1 (empty nudge must not re-run)", n)
	}
}

// TestFinalGuardNudgeRoundFailureKeepsResult: the nudge round erroring (model
// failure, deadline, step cap) must NOT damage the original successful result.
func TestFinalGuardNudgeRoundFailureKeepsResult(t *testing.T) {
	fp := fake.New("fake")
	fp.Enqueue("m",
		fake.Reply("original answer"),
		fake.Fail(errors.New("provider fell over")),
	)
	m, _ := fp.Model("m")
	guard := &guardFunc{nudge: func(RunInfo, FinalState) string { return "deliver it" }}
	ex := New(Config{
		Registry: tool.NewRegistry(),
		Models:   func(ctx context.Context, _ string) (context.Context, llm.Model, error) { return ctx, m, nil },
		Ports:    Ports{FinalGuard: guard},
	})

	res := ex.Run(context.Background(),
		RunnableAgent{ModelTier: "m"}, tool.Invocation{RunID: "r"}, "hi")

	if res.Err != nil {
		t.Fatalf("a failed nudge round must not fail the run: %v", res.Err)
	}
	if res.Output != "original answer" {
		t.Fatalf("output = %q, want the original preserved", res.Output)
	}
}

// TestFinalGuardNotConsultedOnRunError: a failed run is never nudged — the
// guard exists for successful-but-undelivered finishes only.
func TestFinalGuardNotConsultedOnRunError(t *testing.T) {
	fp := fake.New("fake")
	fp.Enqueue("m", fake.Fail(errors.New("boom")))
	m, _ := fp.Model("m")
	guard := &guardFunc{nudge: func(RunInfo, FinalState) string { return "nudge" }}
	ex := New(Config{
		Registry: tool.NewRegistry(),
		Models:   func(ctx context.Context, _ string) (context.Context, llm.Model, error) { return ctx, m, nil },
		Ports:    Ports{FinalGuard: guard},
	})

	res := ex.Run(context.Background(),
		RunnableAgent{ModelTier: "m"}, tool.Invocation{RunID: "r"}, "hi")

	if res.Err == nil {
		t.Fatal("expected a run error")
	}
	if guard.calls != 0 {
		t.Errorf("guard consulted %d times on a failed run, want 0", guard.calls)
	}
}

// TestFinalGuardPanicIsolated: a panicking guard degrades to "no nudge" — the
// successful result survives untouched.
func TestFinalGuardPanicIsolated(t *testing.T) {
	fp := fake.New("fake")
	fp.Enqueue("m", fake.Reply("safe answer"))
	m, _ := fp.Model("m")
	ex := New(Config{
		Registry: tool.NewRegistry(),
		Models:   func(ctx context.Context, _ string) (context.Context, llm.Model, error) { return ctx, m, nil },
		Ports:    Ports{FinalGuard: panicGuard{}},
	})

	res := ex.Run(context.Background(),
		RunnableAgent{ModelTier: "m"}, tool.Invocation{RunID: "r"}, "hi")

	if res.Err != nil {
		t.Fatalf("guard panic must not fail the run: %v", res.Err)
	}
	if res.Output != "safe answer" {
		t.Fatalf("output = %q", res.Output)
	}
}

type panicGuard struct{}

func (panicGuard) FinalNudge(context.Context, RunInfo, FinalState) string {
	panic("guard bug")
}

// TestFinalGuardReceivesToolNamesAndStagedIDs: the guard's snapshot carries the
// run's tool names (via ExtraTools here) and the input-staged file ids, so a
// host can skip runs that can't deliver and exclude the user's own attachments.
func TestFinalGuardReceivesToolNamesAndStagedIDs(t *testing.T) {
	fp := fake.New("fake")
	fp.Enqueue("m", fake.Reply("done"))
	m, _ := fp.Model("m")
	guard := &guardFunc{nudge: func(RunInfo, FinalState) string { return "" }}
	ex := New(Config{
		Registry: tool.NewRegistry(),
		Models:   func(ctx context.Context, _ string) (context.Context, llm.Model, error) { return ctx, m, nil },
		Ports: Ports{
			FinalGuard: guard,
			InputFiles: guardStager(func(name string) string { return "staged-" + name }),
		},
	})

	res := ex.Run(context.Background(),
		RunnableAgent{ModelTier: "m"},
		tool.Invocation{
			RunID: "r",
			ExtraTools: []llm.Tool{{
				Name:        "send_attachments",
				Description: "deliver",
				Handler:     func(context.Context, json.RawMessage) (any, error) { return "ok", nil },
			}},
			InputFiles: []tool.InputFile{{Name: "input.wav", MimeType: "audio/wave", Data: []byte("x")}},
		}, "hi")

	if res.Err != nil {
		t.Fatalf("run error: %v", res.Err)
	}
	var sawTool bool
	for _, n := range guard.state.ToolNames {
		if n == "send_attachments" {
			sawTool = true
		}
	}
	if !sawTool {
		t.Errorf("guard state.ToolNames = %v, want send_attachments present", guard.state.ToolNames)
	}
	if len(guard.state.StagedFileIDs) != 1 || guard.state.StagedFileIDs[0] != "staged-input.wav" {
		t.Errorf("guard state.StagedFileIDs = %v, want [staged-input.wav]", guard.state.StagedFileIDs)
	}
}

// guardStager maps an input file's name to a deterministic file id.
type guardStager func(name string) string

func (f guardStager) StageInputFile(_ context.Context, _, _, name, _ string, _ []byte) (string, error) {
	return f(name), nil
}

// TestFinalGuardSkipsMultiPhaseRuns: a phase pipeline's output contract is the
// pipeline's, not a delivery surface's — the guard must not be consulted.
func TestFinalGuardSkipsMultiPhaseRuns(t *testing.T) {
	fp := fake.New("fake")
	fp.Enqueue("m", fake.Reply("phase-out"))
	m, _ := fp.Model("m")
	guard := &guardFunc{nudge: func(RunInfo, FinalState) string { return "nudge" }}
	ex := New(Config{
		Registry: tool.NewRegistry(),
		Models:   func(ctx context.Context, _ string) (context.Context, llm.Model, error) { return ctx, m, nil },
		Ports:    Ports{FinalGuard: guard},
	})

	res := ex.Run(context.Background(),
		RunnableAgent{ModelTier: "m", Phases: []Phase{{Name: "a", SystemPrompt: "do a"}}},
		tool.Invocation{RunID: "r"}, "hi")

	if res.Err != nil {
		t.Fatalf("run error: %v", res.Err)
	}
	if guard.calls != 0 {
		t.Errorf("guard consulted %d times on a phase run, want 0", guard.calls)
	}
}