vojo/apps/ai-bot/bot_test.go

package main

import "testing"

const botID = "@ai:vojo.chat"

func msg(body, formatted string, userIDs []string, withMentions bool) *MessageContent {
	mc := &MessageContent{Body: body, FormattedBody: formatted}
	if withMentions {
		mc.Mentions = &Mentions{UserIDs: userIDs}
	}
	return mc
}

func TestMentionsBot(t *testing.T) {
	cases := []struct {
		name       string
		mc         *MessageContent
		replyIsBot bool
		want       bool
	}{
		{"explicit user_ids mention", msg("hi", "", []string{botID}, true), false, true},
		{"empty m.mentions {} (F29)", msg("hi ai", "", nil, true), false, false},
		{"someone else mentioned", msg("hi", "", []string{"@alice:vojo.chat"}, true), false, false},
		{"typed @ai no pill no mentions (F30)", msg("hey @ai what's up", "", nil, false), false, false},
		{"pill href in formatted_body", msg("hi", `<a href="https://matrix.to/#/@ai:vojo.chat">Vojo AI</a>`, nil, false), false, true},
		{"pill href %40 encoded", msg("hi", `<a href="https://matrix.to/#/%40ai:vojo.chat">Vojo AI</a>`, nil, false), false, true},
		{"reply to bot's message", msg("thanks", "", nil, true), true, true},
		{"plain message, not a DM", msg("just chatting", "", nil, true), false, false},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			if got := mentionsBot(c.mc, botID, c.replyIsBot); got != c.want {
				t.Fatalf("mentionsBot = %v, want %v", got, c.want)
			}
		})
	}
}

func TestIsDM(t *testing.T) {
	cases := []struct {
		name            string
		joined, invited int
		known           bool
		want            bool
	}{
		{"2 joined", 2, 0, true, true},
		{"1 joined + 1 invited (fresh DM)", 1, 1, true, true},
		{"2 joined + 1 invited NOT a 1:1 (F3)", 2, 1, true, false},
		{"3 joined group", 3, 0, true, false},
		{"counts unknown", 2, 0, false, false},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			m := &roomMeta{joined: c.joined, invited: c.invited, countsKnown: c.known}
			if got := m.isDM(); got != c.want {
				t.Fatalf("isDM = %v, want %v", got, c.want)
			}
		})
	}
}

func TestStripReplyFallback(t *testing.T) {
	in := "> <@alice:vojo.chat> secret third-party text\n> more quote\n\n@ai answer me"
	if got := stripReplyFallback(in); got != "@ai answer me" {
		t.Fatalf("stripReplyFallback = %q", got)
	}
	if got := stripReplyFallback("  plain  "); got != "plain" {
		t.Fatalf("plain trim = %q", got)
	}
}

func TestStripBotMention(t *testing.T) {
	cases := []struct{ in, want string }{
		// The headline regression: the full-mxid pill fallback cinny writes must not reach
		// the search query (it made the grounding provider search for "vojo.chat").
		{"@ai:vojo.chat мессенджер макс удалили из эппстора?", "мессенджер макс удалили из эппстора?"},
		// Bare "@localpart" fallback some clients write, with trailing address punctuation.
		{"@ai, какая погода в Москве", "какая погода в Москве"},
		// Mention mid-message is still removed (it is never user content).
		{"скажи @ai:vojo.chat кто выиграл", "скажи кто выиграл"},
		// No mention → unchanged (DMs, where the bot isn't addressed by name).
		{"кто выиграл вчера", "кто выиграл вчера"},
		// The product name in a real question must survive (we never strip the display name).
		{"@ai:vojo.chat что умеет Vojo AI", "что умеет Vojo AI"},
		// A longer handle that merely contains the localpart is kept.
		{"@ai:vojo.chat пинг @aibot", "пинг @aibot"},
	}
	for _, c := range cases {
		if got := stripBotMention(c.in, botID); got != c.want {
			t.Errorf("stripBotMention(%q) = %q, want %q", c.in, got, c.want)
		}
	}
}

func TestComputeUSD(t *testing.T) {
	const model = "grok-test"
	cfg := &Config{XAIModel: model, Prices: map[string]ModelPrice{
		model: {InputPerM: 1.25, CachedPerM: 0.20, OutputPerM: 2.50},
	}}
	u := Usage{PromptTokens: 1_000_000, CachedTokens: 400_000, CompletionTokens: 1_000_000}
	// nonCached 600k*1.25 + cached 400k*0.20 + out 1M*2.50 = 0.75 + 0.08 + 2.50
	got := computeUSD(model, u, cfg)
	want := 0.75 + 0.08 + 2.50
	if diff := got - want; diff > 1e-9 || diff < -1e-9 {
		t.Fatalf("computeUSD = %v, want %v", got, want)
	}
	// An unknown model falls back to the default model's price (never $0, which would
	// blind the ceiling).
	if got := computeUSD("unknown-model", u, cfg); got != want {
		t.Fatalf("unknown-model fallback = %v, want default %v", got, want)
	}
}

func TestBuildContextGroupDropsThirdParties(t *testing.T) {
	history := []bufferedMsg{
		{sender: "@alice:vojo.chat", body: "third-party chatter", isBot: false},
		{sender: botID, body: "previous bot reply", isBot: true},
		{sender: "@bob:vojo.chat", body: "more third-party", isBot: false},
	}
	got := buildContext("SYS", history, false /* group */, "what is 2+2?", 20, 8000)

	// system first, trigger last, and NO third-party user content in between.
	if got[0].Role != "system" || got[0].Content != "SYS" {
		t.Fatalf("first message must be system prompt, got %+v", got[0])
	}
	last := got[len(got)-1]
	if last.Role != "user" || last.Content != "what is 2+2?" {
		t.Fatalf("last message must be the trigger, got %+v", last)
	}
	for _, m := range got {
		if m.Content == "third-party chatter" || m.Content == "more third-party" {
			t.Fatalf("group context leaked third-party content: %+v", got)
		}
	}
	// the bot's own prior reply is kept as an assistant turn
	foundAssistant := false
	for _, m := range got {
		if m.Role == "assistant" && m.Content == "previous bot reply" {
			foundAssistant = true
		}
	}
	if !foundAssistant {
		t.Fatalf("group context should keep the bot's own prior reply: %+v", got)
	}
}

func TestBuildContextDMIncludesPeer(t *testing.T) {
	history := []bufferedMsg{
		{sender: "@peer:vojo.chat", body: "earlier peer line", isBot: false},
		{sender: botID, body: "earlier bot line", isBot: true},
	}
	got := buildContext("SYS", history, true /* DM */, "follow up", 20, 8000)
	var sawPeer, sawBot bool
	for _, m := range got {
		if m.Role == "user" && m.Content == "earlier peer line" {
			sawPeer = true
		}
		if m.Role == "assistant" && m.Content == "earlier bot line" {
			sawBot = true
		}
	}
	if !sawPeer || !sawBot {
		t.Fatalf("DM context should include peer + bot history: %+v", got)
	}
}