majordomo/agent/finalize_test.go

package agent

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

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

func TestIsWeakFinal(t *testing.T) {
	long := strings.Repeat("As I said, this is the full answer. ", 6) // >120, contains "as i said"
	cases := []struct {
		name string
		in   string
		want bool
	}{
		{"empty", "", true},
		{"whitespace", "   \n\t ", true},
		{"already-answered", "(Already answered above.)", true},
		{"see-above", "see above", true},
		{"as-i-said-short", "As I said, it's 60 minutes.", true},
		{"crisp-number", "42", false},
		{"crisp-yes", "Yes.", false},
		{"crisp-status", "It's down, restarting now.", false},
		{"long-with-as-i-said", long, false}, // >120 chars: not weak despite the phrase
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			if got := isWeakFinal(c.in); got != c.want {
				t.Errorf("isWeakFinal(%q) = %v, want %v", c.in, got, c.want)
			}
		})
	}
}

func TestIsCitationsOnly(t *testing.T) {
	cases := []struct {
		name string
		in   string
		want bool
	}{
		{"sources-md-links", "Sources: [pcprice.watch](https://pcprice.watch/x), [ebay](https://ebay.com/1).", true},
		{"lowercase-bare-url", "sources: see https://example.com/a", true},
		{"bold-label-colon-inside", "**Sources:** [a](https://a), [b](https://b)", true},
		{"bold-label-colon-outside", "**Sources**: [a](https://a), [b](https://b)", true}, // #4: colon after the bold
		{"references-dash", "References — [a](https://a)", true},
		{"citations-label", "Citations: https://x/y", true},
		{"leading-list-marker", "- Sources: [a](https://a)", true},
		{"atx-heading", "## Sources: [a](https://a), [b](https://b)", true}, // #8: heading + its space
		{"further-reading", "Further reading: https://example.com/deep-dive", true},
		{"annotated-multi-source", "Sources: [pcprice.watch](https://a) (tracker), [eBay](https://b) (sold), [bestvaluegpu](https://c) (retail), [resaleprices](https://d) (asking).", true}, // the real #1418 shape
		{"backref-plus-links-is-citations", "References: as noted above, [pcprice.watch](https://pcprice.watch/x).", true},                                                                   // #13: still a citations addendum

		{"empty", "", false},
		{"label-but-no-link", "Source: internal analysis, no URL here", false},
		{"prose-then-sources", "It sells for ~$2,700. Sources: [a](https://a)", false}, // answer first → not a pure addendum
		{"source-led-prose-answer", "Source: According to https://cdc.gov the flu vaccine is 40-60% effective, and the CDC recommends annual vaccination for everyone over six months old.", false}, // #5: prose that merely opens with a label
		{"mentions-sources-midsentence", "The sources of the leak were never confirmed.", false},
		{"link-without-label", "Here is the link you asked for: [a](https://a)", false},
		{"bare-domains-out-of-scope", "Sources: pcprice.watch (used ~$200), ebay.com (sold listings)", false}, // #7: no scheme/link signal
		{"crisp-number", "42", false},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			if got := isCitationsOnly(c.in); got != c.want {
				t.Errorf("isCitationsOnly(%q) = %v, want %v", c.in, got, c.want)
			}
		})
	}
}

func asst(text string, tools ...llm.ToolCall) llm.Message {
	m := llm.Message{Role: llm.RoleAssistant}
	if text != "" {
		m.Parts = []llm.Part{llm.Text(text)}
	}
	m.ToolCalls = tools
	return m
}

func TestFinalOutput(t *testing.T) {
	cite := []llm.ToolCall{{ID: "c1", Name: "cite", Arguments: json.RawMessage(`{}`)}}
	longAnswer := strings.TrimSpace(strings.Repeat("Free group calls are capped at sixty minutes. ", 6)) // >200
	// A sources/citations-only terminal — the glm-5.2 "cite" shape behind mort
	// issue #1418: the prose answer was front-loaded into the tool-call turn and
	// the terminal turn carried only the citations.
	sources := "Sources: [pcprice.watch](https://pcprice.watch/x), [ebay](https://www.ebay.com/itm/1)."
	answerWithSources := longAnswer + "\n\n" + sources
	// A concise (>80, <200 byte) front-loaded answer + a long citations terminal:
	// the ratio arm can't be met against the long terminal, so citations mode
	// must fall back to the floor.
	conciseAnswer := "It sells for about $2,700 used on eBay, typically $2,400 to $2,900 depending on condition and bundle."
	longSources := "Sources: [pcprice.watch](https://pcprice.watch/gpu/rtx5090) (tracker), [ebay](https://www.ebay.com/sch/rtx5090) (sold), [newegg](https://newegg.com/rtx5090) (retail), [pcpartpicker](https://pcpartpicker.com/rtx5090) (history)."
	// A substantive answer that merely OPENS with "Source:" (not a bare list).
	sourceLedAnswer := "Source: https://nvd.nist.gov/vuln/detail/CVE-2024-1234 — this is the authoritative NVD entry for the vulnerability, rated CVSS 9.8 critical."
	// A >=200-byte planning preamble (must be skipped during recovery).
	longPreamble := "I'll look up the current eBay sold listings, then cross-reference the pcprice.watch tracker and a couple of retail sources, compare the medians across all of them, and put together a clear price range for you before I give the final number."

	tests := []struct {
		name     string
		msgs     []llm.Message
		terminal string
		want     string
	}{
		{
			name: "front-loaded answer recovered over back-ref closer",
			msgs: []llm.Message{
				llm.UserText("q?"),
				asst(longAnswer, cite...),
				llm.ToolResultsMessage(llm.ToolResult{ID: "c1", Name: "cite", Content: "ok"}),
				asst("(Already answered above.)"),
			},
			terminal: "(Already answered above.)",
			want:     longAnswer,
		},
		{
			name: "empty terminal recovers prior substantive answer",
			msgs: []llm.Message{
				llm.UserText("q?"),
				asst(longAnswer, cite...),
				llm.ToolResultsMessage(llm.ToolResult{ID: "c1", Name: "cite", Content: "ok"}),
				asst(""),
			},
			terminal: "",
			want:     longAnswer,
		},
		{
			name: "healthy terminal answer is unchanged",
			msgs: []llm.Message{
				llm.UserText("q?"),
				asst("Let me check.", cite...),
				llm.ToolResultsMessage(llm.ToolResult{ID: "c1", Name: "cite", Content: "ok"}),
				asst(longAnswer),
			},
			terminal: longAnswer,
			want:     longAnswer,
		},
		{
			name: "short crisp answer not overridden by a short preamble prior",
			msgs: []llm.Message{
				llm.UserText("is it up?"),
				asst("Let me check the server status.", cite...),
				llm.ToolResultsMessage(llm.ToolResult{ID: "c1", Name: "cite", Content: "ok"}),
				asst("It's down, restarting now."),
			},
			terminal: "It's down, restarting now.", // not weak → returned as-is
			want:     "It's down, restarting now.",
		},
		{
			name: "weak terminal but only a preamble prior: no recovery",
			msgs: []llm.Message{
				llm.UserText("q?"),
				asst("Let me look that up for you.", cite...),
				llm.ToolResultsMessage(llm.ToolResult{ID: "c1", Name: "cite", Content: "ok"}),
				asst("(see above)"),
			},
			terminal: "(see above)",
			want:     "(see above)", // preamble excluded; falls back to terminal
		},
		{
			name: "citations-only terminal recovers front-loaded answer and keeps sources",
			msgs: []llm.Message{
				llm.UserText("q?"),
				asst(longAnswer, cite...),
				llm.ToolResultsMessage(llm.ToolResult{ID: "c1", Name: "cite", Content: "ok"}),
				asst(sources),
			},
			terminal: sources,
			want:     answerWithSources, // answer recovered, citations appended
		},
		{
			name: "citations-only terminal but only a preamble prior: keeps the sources",
			msgs: []llm.Message{
				llm.UserText("q?"),
				asst("Let me gather the sources.", cite...),
				llm.ToolResultsMessage(llm.ToolResult{ID: "c1", Name: "cite", Content: "ok"}),
				asst(sources),
			},
			terminal: sources,
			want:     sources, // nothing substantive to recover → keep the addendum
		},
		{
			name: "citations already in the recovered answer are not duplicated",
			msgs: []llm.Message{
				llm.UserText("q?"),
				asst(answerWithSources, cite...),
				llm.ToolResultsMessage(llm.ToolResult{ID: "c1", Name: "cite", Content: "ok"}),
				asst(sources),
			},
			terminal: sources,
			want:     answerWithSources, // recovered turn already carries the sources
		},
		{
			// #1418 persisted for CONCISE answers: a <200-char front-loaded
			// answer must still be recovered against a long citations terminal
			// (the ratio arm is skipped in citations mode).
			name: "concise front-loaded answer recovered against a long citations terminal",
			msgs: []llm.Message{
				llm.UserText("q?"),
				asst(conciseAnswer, cite...),
				llm.ToolResultsMessage(llm.ToolResult{ID: "c1", Name: "cite", Content: "ok"}),
				asst(longSources),
			},
			terminal: longSources,
			want:     conciseAnswer + "\n\n" + longSources,
		},
		{
			// A substantive answer that merely OPENS with "Source:" and cites a
			// URL mid-sentence is NOT a citations addendum — return it verbatim,
			// never prepend the prior planning turn.
			name: "source-led substantive answer is not hijacked by a prior turn",
			msgs: []llm.Message{
				llm.UserText("what's the authoritative URL?"),
				asst("I'll look up the CVE in the NVD database, cross-reference the vendor advisory, and confirm the canonical URL before I answer.", cite...),
				llm.ToolResultsMessage(llm.ToolResult{ID: "c1", Name: "cite", Content: "ok"}),
				asst(sourceLedAnswer),
			},
			terminal: sourceLedAnswer,
			want:     sourceLedAnswer,
		},
		{
			// A long (>=200) planning preamble must be skipped during recovery
			// (preamble check applies before the length shortcut); the older
			// real answer is recovered instead.
			name: "long preamble is skipped; older real answer recovered",
			msgs: []llm.Message{
				llm.UserText("q?"),
				asst(conciseAnswer, cite...),
				llm.ToolResultsMessage(llm.ToolResult{ID: "c1", Name: "cite", Content: "ok"}),
				asst(longPreamble, cite...),
				llm.ToolResultsMessage(llm.ToolResult{ID: "c2", Name: "cite", Content: "ok"}),
				asst(sources),
			},
			terminal: sources,
			want:     conciseAnswer + "\n\n" + sources,
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			if got := finalOutput(tc.msgs, tc.terminal); got != tc.want {
				t.Errorf("finalOutput = %q, want %q", got, tc.want)
			}
		})
	}
}

func citeToolbox(t *testing.T) *llm.Toolbox {
	t.Helper()
	return llm.NewToolbox("sources", llm.Tool{
		Name:        "cite",
		Description: "Record a citation.",
		Parameters:  json.RawMessage(`{"type":"object","properties":{}}`),
		Handler: func(_ context.Context, _ json.RawMessage) (any, error) {
			return map[string]bool{"ok": true}, nil
		},
	})
}

// TestRun_RecoversFrontLoadedAnswer reproduces the glm-5.2 shape end-to-end: a
// turn carrying the full answer text AND a tool call, then a degenerate
// terminal turn. The recovered answer must be delivered with no extra model
// call (zero-cost recovery from the transcript).
func TestRun_RecoversFrontLoadedAnswer(t *testing.T) {
	longAnswer := strings.TrimSpace(strings.Repeat("Free group calls are capped at sixty minutes. ", 6))
	fp := fake.New("fp")
	fp.Enqueue("test-model",
		fake.ReplyWith(llm.Response{
			Parts:        []llm.Part{llm.Text(longAnswer)},
			ToolCalls:    []llm.ToolCall{{ID: "c1", Name: "cite", Arguments: json.RawMessage(`{}`)}},
			FinishReason: llm.FinishToolCalls,
			Usage:        llm.Usage{InputTokens: 10, OutputTokens: 5},
		}),
		fake.Reply("(Already answered above.)"),
	)

	a := New(newModel(t, fp), "sys", WithToolbox(citeToolbox(t)))
	res, err := a.Run(context.Background(), "is there a meet time limit?")
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Output != longAnswer {
		t.Errorf("Output = %q, want recovered front-loaded answer", res.Output)
	}
	if n := len(fp.Calls()); n != 2 {
		t.Errorf("model calls = %d, want 2 (no extra nudge turn)", n)
	}
}

// TestRun_HealthyTerminalUnchanged guards against regressing the normal case:
// a deferred answer in the terminal turn is delivered verbatim.
func TestRun_HealthyTerminalUnchanged(t *testing.T) {
	fp := fake.New("fp")
	fp.Enqueue("test-model",
		toolCallReply("c1", "cite", `{}`),
		fake.Reply("The limit is 60 minutes for free group calls."),
	)
	a := New(newModel(t, fp), "sys", WithToolbox(citeToolbox(t)))
	res, err := a.Run(context.Background(), "q?")
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Output != "The limit is 60 minutes for free group calls." {
		t.Errorf("Output = %q, want terminal answer unchanged", res.Output)
	}
}

// TestRun_RecoversFrontLoadedAnswerWithCitations reproduces mort issue #1418
// end-to-end: the model front-loads the prose answer into the tool-call turn
// and closes with a sources-only terminal turn. The delivered output must be
// the recovered answer with the citations appended (not the bare sources list),
// with no extra model call.
func TestRun_RecoversFrontLoadedAnswerWithCitations(t *testing.T) {
	longAnswer := strings.TrimSpace(strings.Repeat("Free group calls are capped at sixty minutes. ", 6))
	sources := "Sources: [docs](https://example.com/docs), [pricing](https://example.com/pricing)."
	fp := fake.New("fp")
	fp.Enqueue("test-model",
		fake.ReplyWith(llm.Response{
			Parts:        []llm.Part{llm.Text(longAnswer)},
			ToolCalls:    []llm.ToolCall{{ID: "c1", Name: "cite", Arguments: json.RawMessage(`{}`)}},
			FinishReason: llm.FinishToolCalls,
			Usage:        llm.Usage{InputTokens: 10, OutputTokens: 5},
		}),
		fake.Reply(sources),
	)

	a := New(newModel(t, fp), "sys", WithToolbox(citeToolbox(t)))
	res, err := a.Run(context.Background(), "is there a meet time limit?")
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	want := longAnswer + "\n\n" + sources
	if res.Output != want {
		t.Errorf("Output = %q, want recovered answer + citations %q", res.Output, want)
	}
	if n := len(fp.Calls()); n != 2 {
		t.Errorf("model calls = %d, want 2 (no extra nudge turn)", n)
	}
}