sm-golem-plugin/classes/services/ai_service_test.go

package services

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"mime"
	"mime/multipart"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"git.golem15.com/golem15/sm-golem-plugin/classes/security"
	"git.golem15.com/golem15/sm-golem-plugin/classes/valueobjects"
	"git.golem15.com/golem15/sm-golem-plugin/internal/pgtest"
	"git.golem15.com/golem15/sm-golem-plugin/models"
	"git.golem15.com/golem15/summercms/modules/fetchguard"
	"git.golem15.com/golem15/summercms/modules/lagoon"
	"gorm.io/gorm"
)

// rt is a scripted provider transport.
type rt func(*http.Request) (*http.Response, error)

func (f rt) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }

func answer(status int, body string) *http.Response {
	return &http.Response{StatusCode: status, Header: http.Header{"Content-Type": {"application/json"}}, Body: io.NopCloser(strings.NewReader(body))}
}

func with(t *testing.T, f func(r *http.Request, body []byte) *http.Response) context.Context {
	t.Helper()
	return fetchguard.WithTransport(context.Background(), rt(func(r *http.Request) (*http.Response, error) {
		var body []byte
		if r.Body != nil {
			body, _ = io.ReadAll(r.Body)
		}
		resp := f(r, body)
		if resp == nil {
			return nil, errors.New("connection refused")
		}
		resp.Request = r
		return resp, nil
	}))
}

// freshModels empties golem15_golem_models and inserts rows in order
// (sort_order is the slice index).
func freshModels(t *testing.T, rows ...models.AiModel) *gorm.DB {
	t.Helper()
	db := pgtest.DB(t)
	if err := db.Exec(`DELETE FROM golem15_golem_models`).Error; err != nil {
		t.Fatal(err)
	}
	for i := range rows {
		rows[i].SortOrder = i
		if rows[i].Adapter == "" {
			rows[i].Adapter = "openai"
		}
		if err := db.Create(&rows[i]).Error; err != nil {
			t.Fatal(err)
		}
	}
	return db
}

func svcFor(db *gorm.DB) *AIService {
	return NewAIService(func() *gorm.DB { return db }, nil)
}

const pong = `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}],"usage":{"total_tokens":3}}`

// TestAIServiceFailures pins AIService's PHP failure messages and the
// request each call sends.
func TestAIServiceFailures(t *testing.T) {
	user := &ModelConfig{Adapter: "openai", APIKey: "sk-user-test", Model: "gpt-4o-mini"}
	prompt := valueobjects.NewPrompt().AddUserMessage("ping")

	t.Run("no model, no key", func(t *testing.T) {
		s := NewAIService(nil, nil)
		if r := s.Send(context.Background(), prompt, nil); r.Success || r.Error != "No AI model configured. Please add a model in Settings > AI." {
			t.Fatalf("no model = %+v", r)
		}
		if r := s.Send(context.Background(), prompt, &ModelConfig{Adapter: "openai"}); r.Error != "API key is not configured for the selected model." {
			t.Fatalf("no key = %+v", r)
		}
		if r := s.Ask(context.Background(), "hi", ""); r.Error != "No AI model configured. Please add a model in Settings > AI." {
			t.Fatalf("ask = %+v", r)
		}
	})

	t.Run("connection and answers", func(t *testing.T) {
		s := NewAIService(nil, nil)
		r := s.Send(with(t, func(*http.Request, []byte) *http.Response { return nil }), prompt, user)
		if r.Error != "Failed to connect to AI service." || r.Err == nil {
			t.Fatalf("refused = %+v", r)
		}
		r = s.Send(with(t, func(*http.Request, []byte) *http.Response { return answer(502, "<html>Bad gateway</html>") }), prompt, user)
		if r.Error != "Invalid response from AI service: Syntax error" || r.Raw["raw"] != "<html>Bad gateway</html>" {
			t.Fatalf("not json = %+v", r)
		}
		r = s.Send(with(t, func(*http.Request, []byte) *http.Response { return answer(200, "{\"x\":\"\xff\"}") }), prompt, user)
		if r.Error != "Invalid response from AI service: Malformed UTF-8 characters, possibly incorrectly encoded" {
			t.Fatalf("bad utf-8 = %+v", r)
		}
		r = s.Send(with(t, func(*http.Request, []byte) *http.Response { return answer(200, "") }), prompt, user)
		if r.Error != "Invalid response from AI service: Syntax error" {
			t.Fatalf("empty = %+v", r)
		}
		// The status is not judged: a 401 body is parsed.
		r = s.Send(with(t, func(*http.Request, []byte) *http.Response {
			return answer(401, `{"error":{"message":"Incorrect API key provided"}}`)
		}), prompt, user)
		if r.Error != "Incorrect API key provided" {
			t.Fatalf("401 = %+v", r)
		}
	})

	t.Run("model system prompt", func(t *testing.T) {
		s := NewAIService(nil, nil)
		cfg := *user
		cfg.SystemPrompt = "Model default."
		var got []string
		ctx := with(t, func(r *http.Request, body []byte) *http.Response {
			var m struct {
				Messages []struct {
					Role    string `json:"role"`
					Content string `json:"content"`
				} `json:"messages"`
			}
			_ = json.Unmarshal(body, &m)
			got = append(got, m.Messages[0].Role+":"+m.Messages[0].Content)
			return answer(200, pong)
		})
		s.Send(ctx, prompt, &cfg)
		s.Send(ctx, prompt.WithSystem("Caller."), &cfg)
		if strings.Join(got, "|") != "system:Model default.|system:Caller." {
			t.Fatalf("system prompts = %v", got)
		}
	})

	t.Run("settings lookups", func(t *testing.T) {
		db := freshModels(t,
			models.AiModel{Name: "chat", Model: "gpt-4o", APIKey: lagoon.NewEncrypted("sk-admin-chat"), IsEnabled: true, IsDefault: true},
			models.AiModel{Name: "claude", Adapter: "anthropic", Model: "claude-sonnet-4-6", APIKey: lagoon.NewEncrypted("sk-admin-claude"), IsEnabled: true},
		)
		s := svcFor(db)
		var url string
		ctx := with(t, func(r *http.Request, _ []byte) *http.Response {
			url = r.URL.String()
			if strings.Contains(url, "anthropic") {
				return answer(200, `{"content":[{"type":"text","text":"hi"}]}`)
			}
			return answer(200, pong)
		})
		if r := s.Ask(ctx, "hi", "be brief"); !r.Success || url != "https://api.openai.com/v1/chat/completions" {
			t.Fatalf("ask with the default model = %+v %s", r, url)
		}
		if r := s.SendWithModel(ctx, prompt, "claude"); !r.Success || url != "https://api.anthropic.com/v1/messages" {
			t.Fatalf("by name = %+v %s", r, url)
		}
		if r := s.SendWithModel(ctx, prompt, "nope"); r.Error != "Model 'nope' not found in settings." {
			t.Fatalf("unknown name = %+v", r)
		}
		if r := s.SendToImageModel(ctx, prompt); r.Error != "No image generation model configured." {
			t.Fatalf("no image model = %+v", r)
		}
		if r := s.SendToVisionModel(ctx, prompt); r.Error != "No vision model configured." {
			t.Fatalf("no vision model = %+v", r)
		}
		if r := s.GenerateImage(ctx, "a cat", "", ""); r.Error != `No image generation model configured. Enable a model with "generates_images" in Settings > AI.` {
			t.Fatalf("generate without model = %+v", r)
		}
		if r := s.SendFilePath(ctx, "/nonexistent", "assistants", "", "", ""); r.Error != "No file model configured." {
			t.Fatalf("no file model = %+v", r)
		}
		// A named model without a files endpoint refuses the upload before
		// anything is read or sent.
		if r := s.SendFilePath(ctx, "/nonexistent", "assistants", "chat", "", ""); r.Error != "Model 'chat' does not support file uploads." {
			t.Fatalf("file upload refused = %+v", r)
		}
		if r := s.SendFile(ctx, strings.NewReader("x"), "a.txt", "text/plain", "missing"); r.Error != "Model 'missing' does not support file uploads." {
			t.Fatalf("missing model upload = %+v", r)
		}
	})

	t.Run("file upload", func(t *testing.T) {
		db := freshModels(t,
			models.AiModel{Name: "files", Model: "gpt-4o", APIKey: lagoon.NewEncrypted("sk-admin-files"), IsEnabled: true, HasFilesEndpoint: true, BaseURL: ptr("https://llm.example.test/v1")},
		)
		s := svcFor(db)
		path := filepath.Join(t.TempDir(), "invoice.pdf")
		if err := os.WriteFile(path, []byte("%PDF-1.4 parity"), 0o600); err != nil {
			t.Fatal(err)
		}
		ctx := with(t, func(r *http.Request, body []byte) *http.Response {
			if r.URL.String() != "https://llm.example.test/v1/files" || r.Header.Get("Authorization") != "Bearer sk-admin-files" {
				t.Errorf("upload request %s %v", r.URL, r.Header)
			}
			mt, params, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
			if err != nil || mt != "multipart/form-data" {
				t.Errorf("content type %q", r.Header.Get("Content-Type"))
				return answer(400, "{}")
			}
			mr := multipart.NewReader(strings.NewReader(string(body)), params["boundary"])
			var parts []string
			for {
				p, err := mr.NextPart()
				if err != nil {
					break
				}
				data, _ := io.ReadAll(p)
				parts = append(parts, p.FormName()+"|"+p.FileName()+"|"+p.Header.Get("Content-Type")+"|"+string(data))
			}
			if strings.Join(parts, ";") != "purpose|||assistants;file|invoice.pdf|application/pdf|%PDF-1.4 parity" {
				t.Errorf("parts = %v", parts)
			}
			return answer(200, `{"id":"file-abc123","object":"file"}`)
		})
		if r := s.SendFilePath(ctx, path, "assistants", "files", "application/pdf", ""); !r.Success || r.Content != "file-abc123" {
			t.Fatalf("upload = %+v", r)
		}
		if r := s.SendFilePath(ctx, filepath.Join(t.TempDir(), "gone.pdf"), "assistants", "", "", ""); r.Error != "Internal server error" {
			t.Fatalf("missing file = %+v", r)
		}
		noID := with(t, func(*http.Request, []byte) *http.Response { return answer(200, `{"object":"file"}`) })
		if r := s.SendFile(noID, strings.NewReader("x"), "a.txt", "text/plain", ""); r.Error != "No file ID in response." {
			t.Fatalf("no id = %+v", r)
		}
		apiErr := with(t, func(*http.Request, []byte) *http.Response {
			return answer(400, `{"error":{"message":"purpose is invalid"}}`)
		})
		if r := s.SendFile(apiErr, strings.NewReader("x"), "a.txt", "text/plain", ""); r.Error != "purpose is invalid" {
			t.Fatalf("api error = %+v", r)
		}
		down := with(t, func(*http.Request, []byte) *http.Response { return nil })
		if r := s.SendFile(down, strings.NewReader("x"), "a.txt", "text/plain", ""); r.Error != "Failed to connect to files API at: https://llm.example.test/v1/files" {
			t.Fatalf("down = %+v", r)
		}
	})

	t.Run("generate image", func(t *testing.T) {
		db := freshModels(t,
			models.AiModel{Name: "dalle", Model: "dall-e-3", APIKey: lagoon.NewEncrypted("sk-admin-dalle"), IsEnabled: true, GeneratesImages: true},
		)
		s := svcFor(db)
		returned := "https://oaidalleapiprodscus.blob.core.windows.net/private/img-1.png"
		ctx := with(t, func(r *http.Request, body []byte) *http.Response {
			var m map[string]any
			_ = json.Unmarshal(body, &m)
			if r.URL.String() != "https://api.openai.com/v1/images/generations" || m["model"] != "dall-e-3" || m["prompt"] != "a cat" || m["n"] != float64(1) || m["size"] != "1792x1024" || m["quality"] != "standard" {
				t.Errorf("image request %s %v", r.URL, m)
			}
			return answer(200, `{"data":[{"url":"`+returned+`","revised_prompt":"a fluffy cat"}]}`)
		})
		r := s.GenerateImage(ctx, "a cat", "", "")
		if !r.Success || r.Content != returned || r.Raw["revised_prompt"] != "a fluffy cat" {
			t.Fatalf("generate = %+v", r)
		}
		// The returned URL is checked before anyone fetches it.
		prev := security.DefaultResolver
		security.DefaultResolver = staticResolver{"oaidalleapiprodscus.blob.core.windows.net": "20.60.40.1", "evil.example.com": "203.0.114.9"}
		t.Cleanup(func() { security.DefaultResolver = prev })
		if err := security.AssertSafeURL(context.Background(), r.Content, security.DefaultAllowedHosts); err != nil {
			t.Fatalf("returned URL refused: %v", err)
		}
		evil := with(t, func(*http.Request, []byte) *http.Response {
			return answer(200, `{"data":[{"url":"https://evil.example.com/x.png"}]}`)
		})
		r = s.GenerateImage(evil, "a cat", "1024x1024", "hd")
		if err := security.AssertSafeURL(context.Background(), r.Content, security.DefaultAllowedHosts); err == nil || err.Error() != "Host not in allowlist: evil.example.com" {
			t.Fatalf("a foreign returned URL must fail the guard: %v", err)
		}
		if r := s.GenerateImage(with(t, func(*http.Request, []byte) *http.Response { return answer(200, `{"data":[]}`) }), "x", "", ""); r.Error != "No image URL in response." {
			t.Fatalf("no url = %+v", r)
		}
		if r := s.GenerateImage(with(t, func(*http.Request, []byte) *http.Response { return answer(200, `nope`) }), "x", "", ""); r.Error != "Invalid response: Syntax error" {
			t.Fatalf("bad json = %+v", r)
		}
		if r := s.GenerateImage(with(t, func(*http.Request, []byte) *http.Response {
			return answer(400, `{"error":{"message":"content policy"}}`)
		}), "x", "", ""); r.Error != "content policy" {
			t.Fatalf("api error = %+v", r)
		}
		if r := s.GenerateImage(with(t, func(*http.Request, []byte) *http.Response { return nil }), "x", "", ""); r.Error != "Failed to connect to image generation service." {
			t.Fatalf("down = %+v", r)
		}
		freshModels(t, models.AiModel{Name: "claude", Adapter: "anthropic", Model: "claude", APIKey: lagoon.NewEncrypted("k"), IsEnabled: true, GeneratesImages: true})
		if r := s.GenerateImage(ctx, "x", "", ""); r.Error != "The configured adapter does not support image generation." {
			t.Fatalf("anthropic image = %+v", r)
		}
	})
}

type staticResolver map[string]string

func (s staticResolver) LookupIPAddr(_ context.Context, host string) ([]net.IPAddr, error) {
	ip, ok := s[host]
	if !ok {
		return nil, &net.DNSError{Err: "no such host", Name: host, IsNotFound: true}
	}
	return []net.IPAddr{{IP: net.ParseIP(ip)}}, nil
}

func ptr(s string) *string { return &s }

// TestAIServiceStream delivers the content deltas of a server-sent event
// stream in order, as {"content":…} chunks, and stops at the end event.
func TestAIServiceStream(t *testing.T) {
	s := NewAIService(nil, nil)
	stream := func(status int, body string, check func(r *http.Request, payload map[string]any)) context.Context {
		return fetchguard.WithTransport(context.Background(), rt(func(r *http.Request) (*http.Response, error) {
			raw, _ := io.ReadAll(r.Body)
			var m map[string]any
			_ = json.Unmarshal(raw, &m)
			check(r, m)
			return &http.Response{StatusCode: status, Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: io.NopCloser(strings.NewReader(body)), Request: r}, nil
		}))
	}
	var chunks []string
	collect := func(c string) { chunks = append(chunks, c) }
	openai := "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n" +
		"data: {\"choices\":[{\"delta\":{\"content\":\"Hel\"}}]}\n\n" +
		"data: {\"choices\":[{\"delta\":{\"content\":\"lo \\\"you\\\"\"}}]}\r\n\r\n" +
		"data: [DONE]\n\n" +
		"data: {\"choices\":[{\"delta\":{\"content\":\"after done\"}}]}\n"
	err := s.SendStream(stream(200, openai, func(r *http.Request, m map[string]any) {
		if m["stream"] != true || r.Header.Get("Accept") != "text/event-stream" || r.Header.Get("Content-Type") != "application/json" {
			t.Errorf("stream request %v %v", r.Header, m)
		}
	}), valueobjects.NewPrompt().AddUserMessage("hi"), collect, &ModelConfig{Adapter: "openai", APIKey: "k", Model: "gpt-4o"})
	if err != nil {
		t.Fatal(err)
	}
	if strings.Join(chunks, "") != `{"content":"Hel"}{"content":"lo \"you\""}` {
		t.Fatalf("openai chunks = %v", chunks)
	}

	chunks = nil
	anthropic := "event: message_start\ndata: {\"type\":\"message_start\"}\n\n" +
		"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"Cze\"}}\n\n" +
		"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"ść\"}}\n\n" +
		"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n" +
		"data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"late\"}}\n"
	err = s.SendStream(stream(200, anthropic, func(r *http.Request, m map[string]any) {
		if r.Header.Get("X-Api-Key") != "k" || m["stream"] != true {
			t.Errorf("anthropic stream request %v", r.Header)
		}
	}), valueobjects.NewPrompt().AddUserMessage("hi"), collect, &ModelConfig{Adapter: "anthropic", APIKey: "k", Model: "claude"})
	if err != nil {
		t.Fatal(err)
	}
	if strings.Join(chunks, "") != `{"content":"Cze"}{"content":"ść"}` {
		t.Fatalf("anthropic chunks = %v", chunks)
	}

	if err := s.SendStream(stream(401, `{"error":{"message":"bad key"}}`, func(*http.Request, map[string]any) {}), valueobjects.NewPrompt(), collect, &ModelConfig{APIKey: "k"}); err == nil || err.Error() != "HTTP 401 from AI service" {
		t.Fatalf("401 = %v", err)
	}
	if err := s.SendStream(context.Background(), valueobjects.NewPrompt(), collect, nil); err == nil || err.Error() != "No AI model configured. Please add a model in Settings > AI." {
		t.Fatalf("no model = %v", err)
	}
	if err := s.SendStream(context.Background(), valueobjects.NewPrompt(), collect, &ModelConfig{}); err == nil || err.Error() != "API key is not configured for the selected model." {
		t.Fatalf("no key = %v", err)
	}
}

// TestDefaultModelQuirk pins Settings::getDefaultModel: the first enabled
// default model, else the model at position 0 only when it is enabled.
func TestDefaultModelQuirk(t *testing.T) {
	ctx := context.Background()
	key := lagoon.NewEncrypted("sk-admin")
	db := freshModels(t,
		models.AiModel{Name: "zero", Model: "m0", APIKey: key},
		models.AiModel{Name: "one", Model: "m1", APIKey: key, IsEnabled: true, AcceptsImages: true},
	)
	def, err := DefaultModel(ctx, db)
	if err != nil || def != nil {
		t.Fatalf("model 0 disabled, model 1 enabled without is_default: default = %+v, %v; want none", def, err)
	}
	all, _ := Models(ctx, db)
	if len(all) != 1 || all[0].Name != "one" || all[0].Position != 1 || !all[0].Trusted || all[0].APIKey != "sk-admin" {
		t.Fatalf("models = %+v", all)
	}
	if v, _ := VisionModel(ctx, db); v == nil || v.Name != "one" || !v.Trusted {
		t.Fatalf("vision = %+v", v)
	}

	db = freshModels(t,
		models.AiModel{Name: "zero", Model: "m0", APIKey: key},
		models.AiModel{Name: "one", Model: "m1", APIKey: key, IsEnabled: true, IsDefault: true},
	)
	if def, _ := DefaultModel(ctx, db); def == nil || def.Name != "one" {
		t.Fatalf("explicit default = %+v", def)
	}
	db = freshModels(t,
		models.AiModel{Name: "zero", Model: "m0", APIKey: key, IsEnabled: true},
		models.AiModel{Name: "one", Model: "m1", APIKey: key, IsEnabled: true},
	)
	if def, _ := DefaultModel(ctx, db); def == nil || def.Name != "zero" {
		t.Fatalf("position 0 fallback = %+v", def)
	}
	// A disabled default never counts.
	db = freshModels(t,
		models.AiModel{Name: "zero", Model: "m0", APIKey: key, IsEnabled: true},
		models.AiModel{Name: "one", Model: "m1", APIKey: key, IsDefault: true},
	)
	if def, _ := DefaultModel(ctx, db); def == nil || def.Name != "zero" {
		t.Fatalf("disabled default = %+v", def)
	}
	db = freshModels(t,
		models.AiModel{Name: "a", Model: "m", APIKey: key, IsEnabled: true, HasFilesEndpoint: true, BaseURL: ptr("https://x.test/v1"), SystemPrompt: ptr("sys"), MaxCompletionTokens: func() *int { n := 2048; return &n }()},
		models.AiModel{Name: "b", Model: "m", APIKey: key, IsEnabled: true, GeneratesImages: true},
	)
	if m, _ := FileModel(ctx, db); m == nil || m.Name != "a" || m.BaseURL != "https://x.test/v1" || m.SystemPrompt != "sys" || m.MaxCompletionTokens != 2048 {
		t.Fatalf("file model = %+v", m)
	}
	if m, _ := ImageModel(ctx, db); m == nil || m.Name != "b" {
		t.Fatalf("image model = %+v", m)
	}
	if m, _ := ModelByName(ctx, db, "b"); m == nil || m.Name != "b" {
		t.Fatalf("by name = %+v", m)
	}
	if m, _ := Models(ctx, nil); m != nil {
		t.Fatal("no database means no models")
	}
}

// TestAdminModelTrusted: an operator's model reaches a plain-http
// loopback endpoint (fetchguard TrustedMode); a user's config to a
// loopback endpoint is refused by the dial guard.
func TestAdminModelTrusted(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/chat/completions" {
			http.NotFound(w, r)
			return
		}
		_, _ = io.WriteString(w, pong)
	}))
	defer srv.Close()
	tls := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.WriteString(w, pong)
	}))
	defer tls.Close()
	s := NewAIService(nil, nil)
	prompt := valueobjects.NewPrompt().AddUserMessage("ping")

	admin := &ModelConfig{Adapter: "openai", APIKey: "sk-admin", Model: "local", BaseURL: srv.URL + "/v1", Trusted: true}
	if r := s.Send(context.Background(), prompt, admin); !r.Success || r.Content != "pong" {
		t.Fatalf("admin model on the local network = %+v", r)
	}

	user := &ModelConfig{Adapter: "openai", APIKey: "sk-user", Model: "local", BaseURL: tls.URL + "/v1"}
	r := s.Send(context.Background(), prompt, user)
	var fe *fetchguard.Error
	if r.Success || r.Error != "Failed to connect to AI service." || !errors.As(r.Err, &fe) || fe.Reason != fetchguard.ReasonPrivateIP {
		t.Fatalf("user config to loopback = %+v (%v)", r, r.Err)
	}
	user.BaseURL = srv.URL + "/v1"
	r = s.Send(context.Background(), prompt, user)
	if !errors.As(r.Err, &fe) || fe.Reason != fetchguard.ReasonScheme {
		t.Fatalf("user config over http = %+v (%v)", r, r.Err)
	}
}