backend/internal/server/chat_test.go

package server

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"

	"code.littlelan.cn/CarrotAssistant/backend/internal/config"
	"code.littlelan.cn/CarrotAssistant/backend/internal/security"
	"code.littlelan.cn/CarrotAssistant/backend/internal/skill"
	"code.littlelan.cn/CarrotAssistant/backend/internal/store"
)

// newTestDB opens an in-memory-ish SQLite (temp file) and returns the gorm DB
// plus a cleanup function.
func newTestDB(t *testing.T) (*gorm.DB, func()) {
	t.Helper()
	dir := t.TempDir()
	db, err := store.Open(dir + "/test.db")
	if err != nil {
		t.Fatalf("open db: %v", err)
	}
	return db, func() {}
}

// newTestDeps builds a Deps wired to a mock LLM server, with crypto + skills
// dir under the test temp directory.
func newTestDeps(t *testing.T, mockLLMURL string) (Deps, *store.App, *skill.Skill, string, func()) {
	t.Helper()
	db, dbCleanup := newTestDB(t)
	dir := t.TempDir()
	crypto, err := security.NewCrypto("0123456789abcdef0123456789abcdef")
	if err != nil {
		t.Fatalf("crypto: %v", err)
	}

	// A model config pointing at the mock server.
	mc := store.ModelConfig{
		Name: "mock", BaseURL: mockLLMURL, APIKeyEncrypted: mustEnc(crypto, "sk-mock"),
		APIKeyPreview: "mock", DefaultModel: "mock-model", SupportsTools: true,
	}
	if err := db.Create(&mc).Error; err != nil {
		t.Fatalf("create mc: %v", err)
	}

	// An app that uses it, with a known token so the chat client can auth.
	const appToken = "ca_live_testtoken"
	app := &store.App{
		Slug: "demo", Name: "Demo", TokenHash: security.HashToken(appToken),
		ModelConfigID: &mc.ID, Status: "active",
	}
	if err := db.Create(app).Error; err != nil {
		t.Fatalf("create app: %v", err)
	}

	// A published skill with one read tool and one mutating tool that
	// requires confirmation.
	sk := &skill.Skill{
		Name:        "demo-assistant",
		Description: "demo",
		AppSlug:     "demo",
		Permissions: skill.Permissions{
			DefaultMode:         "readonly",
			RequireConfirmation: []string{"delete_post"},
		},
		Tools: []skill.Tool{
			{
				Name: "search", Description: "search",
				Parameters: skill.MakeParameters([]byte(`{"type":"object","properties":{"q":{"type":"string"}}}`)),
				Endpoint:   skill.Endpoint{Method: "GET", URL: "https://example.invalid/search", Query: map[string]string{"q": "{{q}}"}},
			},
			{
				Name: "delete_post", Description: "delete",
				Parameters: skill.MakeParameters([]byte(`{"type":"object","properties":{"id":{"type":"integer"}}}`)),
				Endpoint:   skill.Endpoint{Method: "DELETE", URL: "https://example.invalid/posts/{{id}}"},
			},
		},
		SystemPrompt: "you are a demo assistant",
	}
	mgr := skill.NewManager(dir)
	if _, err := mgr.Write("demo", "demo-assistant", sk); err != nil {
		t.Fatalf("write skill: %v", err)
	}
	row := store.Skill{
		AppID: app.ID, Slug: "demo-assistant", Name: sk.Name, FilePath: mgr.RelPath("demo", "demo-assistant"),
		Status: "published", Source: "manual",
	}
	if err := db.Create(&row).Error; err != nil {
		t.Fatalf("create skill: %v", err)
	}

	cfg := &config.Config{JWTSecret: "test-jwt"}
	deps := Deps{DB: db, Config: cfg, Crypto: crypto, SkillsDir: dir}
	return deps, app, sk, appToken, dbCleanup
}

func mustEnc(c *security.Crypto, s string) string {
	out, err := c.Encrypt(s)
	if err != nil {
		panic(err)
	}
	return out
}

// mockLLM is a tiny SSE chat-completion server whose response sequence is
// driven by a script: a list of chunk-producers invoked in order, one per
// LLM call. This lets a test script a tool-call turn followed by a final
// answer turn.
type mockLLM struct {
	turns [][]string // each turn is a list of SSE chunks (raw JSON delta lines)
	idx   int
}

func (m *mockLLM) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Read the request body to find how many turns we've served (we just
	// advance the script on each call).
	body, _ := io.ReadAll(r.Body)
	_ = body
	if m.idx >= len(m.turns) {
		// No more scripted turns: emit a bare stop so the loop ends cleanly.
		m.idx++
		writeSSEChunk(w, map[string]any{"choices": []map[string]any{{"finish_reason": "stop", "delta": map[string]any{}}}})
		fmt.Fprint(w, "data: [DONE]\n\n")
		return
	}
	chunks := m.turns[m.idx]
	m.idx++
	w.Header().Set("Content-Type", "text/event-stream")
	for _, c := range chunks {
		fmt.Fprintf(w, "data: %s\n\n", c)
	}
	fmt.Fprint(w, "data: [DONE]\n\n")
	if f, ok := w.(http.Flusher); ok {
		f.Flush()
	}
}

// writeSSEChunk marshals one chunk object and writes it as a data line.
func writeSSEChunk(w http.ResponseWriter, chunk map[string]any) {
	b, _ := json.Marshal(chunk)
	fmt.Fprintf(w, "data: %s\n\n", b)
}

// TestChatToolCallFlow drives a chat request against a mock LLM that first
// requests the search tool (a read), then produces a final answer. Asserts
// the SSE stream emits the expected event sequence.
func TestChatToolCallFlow(t *testing.T) {
	mock := &mockLLM{
		turns: [][]string{
			// Turn 1: model asks to call search(q="hi").
			{
				chunkJSON("search", "call_1", "", `{"q":"hi"}`, ""),
				`{"choices":[{"finish_reason":"tool_calls","delta":{}}]}`,
			},
			// Turn 2: model produces the final text answer.
			{
				`{"choices":[{"finish_reason":"","delta":{"content":"found "}}]}`,
				`{"choices":[{"finish_reason":"","delta":{"content":"2 posts"}}]}`,
				`{"choices":[{"finish_reason":"stop","delta":{}}]}`,
			},
		},
	}
	mockSrv := httptest.NewServer(mock)
	defer mockSrv.Close()

	deps, _, _, appToken, cleanup := newTestDeps(t, mockSrv.URL)
	defer cleanup()

	// Also stand up a fake "target site" so the tool call resolves.
	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`["post1","post2"]`))
	}))
	defer target.Close()
	// Patch the search tool URL to point at the fake target. (We do this by
	// editing the published skill file in place.)
	patchSkillURL(t, deps, "demo", "demo-assistant", "search", target.URL+"/search")

	gin.SetMode(gin.TestMode)
	r := gin.New()
	deps.mountChatRoutes(r)

	// POST /api/v1/chat and collect the SSE stream.
	body := bytes.NewBufferString(`{"message":"hi"}`)
	req := httptest.NewRequest("POST", "/api/v1/chat", body)
	req.Header.Set("X-App-Token", appToken)
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()
	r.ServeHTTP(rec, req)

	events := parseSSE(rec.Body.String())
	kinds := eventKinds(events)
	wantSeq := []string{"session", "tool_call", "tool_result", "token", "token", "done"}
	if !containsSeq(kinds, wantSeq) {
		t.Fatalf("event kinds = %v, want sequence containing %v", kinds, wantSeq)
	}
}

// patchSkillURL rewrites a tool's endpoint URL in the published skill file
// for the test, so tool calls hit the in-process fake target.
func patchSkillURL(t *testing.T, deps Deps, appSlug, skillSlug, toolName, newURL string) {
	t.Helper()
	mgr := skill.NewManager(deps.SkillsDir)
	var row store.Skill
	deps.DB.Where("slug = ? AND app_id IN (SELECT id FROM apps WHERE slug = ?)", skillSlug, appSlug).First(&row)
	sk, err := mgr.Load(&row)
	if err != nil {
		t.Fatalf("load skill: %v", err)
	}
	for i := range sk.Tools {
		if sk.Tools[i].Name == toolName {
			sk.Tools[i].Endpoint.URL = newURL
		}
	}
	if _, err := mgr.Write(appSlug, skillSlug, sk); err != nil {
		t.Fatalf("rewrite skill: %v", err)
	}
}

// chunkJSON builds a delta chunk for a tool call (function arguments arrive
// as a JSON string). When content is non-empty it's a content delta instead.
func chunkJSON(name, id, content, argsJSON, _ string) string {
	if content != "" {
		b, _ := json.Marshal(map[string]any{"choices": []map[string]any{{
			"delta": map[string]any{"content": content},
		}}})
		return string(b)
	}
	// tool-call delta
	b, _ := json.Marshal(map[string]any{"choices": []map[string]any{{
		"delta": map[string]any{
			"tool_calls": []map[string]any{{
				"index": 0, "id": id,
				"function": map[string]any{"name": name, "arguments": argsJSON},
			}},
		},
	}}})
	return string(b)
}

// parseSSE turns a raw SSE body into a list of (event, data) pairs.
func parseSSE(raw string) []sseEvent {
	var out []sseEvent
	var cur sseEvent
	sc := bufio.NewScanner(strings.NewReader(raw))
	sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
	for sc.Scan() {
		ln := sc.Text()
		if ln == "" {
			if cur.event != "" || cur.data != "" {
				out = append(out, cur)
			}
			cur = sseEvent{}
			continue
		}
		if strings.HasPrefix(ln, "event: ") {
			cur.event = strings.TrimPrefix(ln, "event: ")
		} else if strings.HasPrefix(ln, "data: ") {
			cur.data += strings.TrimPrefix(ln, "data: ")
		}
	}
	if cur.event != "" || cur.data != "" {
		out = append(out, cur)
	}
	return out
}

type sseEvent struct {
	event string
	data  string
}

func eventKinds(es []sseEvent) []string {
	out := make([]string, 0, len(es))
	for _, e := range es {
		out = append(out, e.event)
	}
	return out
}

// containsSeq reports whether want appears as a contiguous subsequence of got.
func containsSeq(got, want []string) bool {
	if len(want) > len(got) {
		return false
	}
	for i := 0; i+len(want) <= len(got); i++ {
		match := true
		for j := range want {
			if got[i+j] != want[j] {
				match = false
				break
			}
		}
		if match {
			return true
		}
	}
	return false
}