Thomas
5e23659d99
* feat(assistant): optimize MCP and custom widget authoring * fix(assistant): isolate server-safe JSX analyzer * refactor(assistant): generate MCP tool staging * fix(assistant): preserve widget placement follow-up * fix(ci): allow fast gate to finish * fix(custom-widgets): address review findings * fix(custom-widgets): split preview evidence module * style: format assistant custom widget files * fix(custom-widgets): sync formatted skill schema
281 lines
11 KiB
Go
281 lines
11 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestSanitizeProviderPayload(t *testing.T) {
|
|
payload := map[string]any{
|
|
"model": homarrProviderModelID,
|
|
"models": []string{"attacker/model"},
|
|
"provider": map[string]any{"order": []string{"attacker"}},
|
|
"route": "fallback",
|
|
"plugins": []string{"web"},
|
|
"transforms": []string{"middle-out"},
|
|
"max_tokens": 1_000_000,
|
|
"max_completion_tokens": 1_000_000,
|
|
"n": 100,
|
|
"parallel_tool_calls": true,
|
|
"extra_body": map[string]any{"provider": map[string]any{"order": []string{"attacker"}}},
|
|
"extra_headers": map[string]any{"Authorization": "Bearer attacker"},
|
|
"metadata": map[string]any{"user": "private-user-id"},
|
|
"audio": map[string]any{"format": "wav"},
|
|
"modalities": []string{"text", "audio"},
|
|
"logprobs": true,
|
|
"top_logprobs": 20,
|
|
"prediction": map[string]any{"content": strings.Repeat("x", 100)},
|
|
"service_tier": "priority",
|
|
"reasoning": map[string]any{"effort": "high", "max_tokens": 1_000_000},
|
|
"user": "private-user-id",
|
|
"tools": []any{
|
|
map[string]any{"type": "function", "function": map[string]any{"name": "board_list"}},
|
|
map[string]any{
|
|
"type": "openrouter:web_search",
|
|
"parameters": map[string]any{"max_results": 1000, "max_uses": 1000},
|
|
},
|
|
map[string]any{"type": "openrouter:web_search"},
|
|
},
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
}
|
|
if err := sanitizeProviderPayload(payload, "mock/team-selected-model"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for _, field := range clientControlledRoutingFields {
|
|
if field == "provider" {
|
|
continue
|
|
}
|
|
if _, exists := payload[field]; exists {
|
|
t.Fatalf("client-controlled routing field %q was forwarded", field)
|
|
}
|
|
}
|
|
for _, field := range clientControlledCostFields {
|
|
if _, exists := payload[field]; exists {
|
|
t.Fatalf("client-controlled cost field %q was forwarded", field)
|
|
}
|
|
}
|
|
if payload["model"] != "mock/team-selected-model" {
|
|
t.Fatalf("unexpected upstream model: %v", payload["model"])
|
|
}
|
|
if payload["max_completion_tokens"] != maxChatOutputTokens || payload["n"] != 1 || payload["parallel_tool_calls"] != false {
|
|
t.Fatalf("unsafe generation controls were forwarded: %#v", payload)
|
|
}
|
|
if _, exists := payload["max_tokens"]; exists {
|
|
t.Fatal("deprecated client max_tokens was forwarded")
|
|
}
|
|
if _, exists := payload["user"]; exists {
|
|
t.Fatal("client user identifier was forwarded")
|
|
}
|
|
if _, exists := payload["metadata"]; exists {
|
|
t.Fatal("client metadata was forwarded")
|
|
}
|
|
reasoning := payload["reasoning"].(map[string]any)
|
|
if _, exists := reasoning["max_tokens"]; exists || reasoning["effort"] != "high" {
|
|
t.Fatalf("reasoning budget was not bounded safely: %#v", reasoning)
|
|
}
|
|
tools := payload["tools"].([]any)
|
|
if len(tools) != 2 {
|
|
t.Fatalf("unexpected sanitized tool count: %d", len(tools))
|
|
}
|
|
webSearch := tools[1].(map[string]any)
|
|
parameters := webSearch["parameters"].(map[string]any)
|
|
if parameters["max_results"] != maxWebSearchResults ||
|
|
parameters["max_uses"] != maxWebSearchUses ||
|
|
parameters["max_total_results"] != maxWebSearchTotal ||
|
|
parameters["max_characters"] != maxWebSearchCharacters ||
|
|
parameters["search_context_size"] != webSearchContextSize {
|
|
t.Fatalf("web search limits were not enforced: %#v", parameters)
|
|
}
|
|
usage, err := json.Marshal(payload["usage"])
|
|
if err != nil || string(usage) != `{"include":true}` {
|
|
t.Fatalf("unexpected usage options: %s, %v", usage, err)
|
|
}
|
|
privacy := payload["provider"].(map[string]any)
|
|
if privacy["zdr"] != true || privacy["data_collection"] != "deny" {
|
|
t.Fatalf("upstream privacy controls were not enforced: %#v", privacy)
|
|
}
|
|
}
|
|
|
|
func TestSanitizeProviderPayloadRejectsUnsupportedServerTools(t *testing.T) {
|
|
payload := map[string]any{"tools": []any{map[string]any{"type": "openrouter:computer"}}}
|
|
if err := sanitizeProviderPayload(payload, "mock/team-selected-model"); err == nil {
|
|
t.Fatal("expected unsupported server tool to be rejected")
|
|
}
|
|
}
|
|
|
|
func TestSafeUpstreamStatus(t *testing.T) {
|
|
for _, status := range []int{401, 403} {
|
|
if actual := safeUpstreamStatus(status); actual != 502 {
|
|
t.Fatalf("expected upstream %d to become 502, got %d", status, actual)
|
|
}
|
|
}
|
|
if actual := safeUpstreamStatus(429); actual != 429 {
|
|
t.Fatalf("expected upstream 429 to be preserved, got %d", actual)
|
|
}
|
|
}
|
|
|
|
func TestValidateProviderInput(t *testing.T) {
|
|
payload := map[string]any{
|
|
"messages": []any{map[string]any{
|
|
"role": "user",
|
|
"content": []any{map[string]any{
|
|
"type": "image_url",
|
|
"image_url": map[string]any{"url": "data:image/png;base64," + strings.Repeat("a", 1_300_000)},
|
|
}},
|
|
}},
|
|
}
|
|
if err := validateProviderInput(payload); err != nil {
|
|
t.Fatalf("expected a supported Homarr image request, got %v", err)
|
|
}
|
|
payload["messages"] = []any{map[string]any{"role": "user", "content": strings.Repeat("x", maxChatTextBytes+1)}}
|
|
if err := validateProviderInput(payload); !errors.Is(err, errInputTooLarge) {
|
|
t.Fatalf("expected oversized text input to be rejected, got %v", err)
|
|
}
|
|
payload["messages"] = []any{map[string]any{
|
|
"role": "user",
|
|
"content": []any{map[string]any{
|
|
"type": "image_url",
|
|
"image_url": map[string]any{"url": "data:image/png;base64," + strings.Repeat("a", maxChatImageDataBytes)},
|
|
}},
|
|
}}
|
|
if err := validateProviderInput(payload); !errors.Is(err, errTooManyImages) {
|
|
t.Fatalf("expected oversized image input to be rejected, got %v", err)
|
|
}
|
|
payload["messages"] = []any{map[string]any{
|
|
"role": "user",
|
|
"content": []any{map[string]any{
|
|
"type": "image_url",
|
|
"image_url": map[string]any{"url": "https://attacker.example/expensive-image.png"},
|
|
}},
|
|
}}
|
|
if err := validateProviderInput(payload); !errors.Is(err, errRemoteImage) {
|
|
t.Fatalf("expected a remote image URL to be rejected, got %v", err)
|
|
}
|
|
payload["messages"] = []any{map[string]any{
|
|
"role": "user",
|
|
"content": []any{
|
|
map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,YQ=="}},
|
|
map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,YQ=="}},
|
|
map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,YQ=="}},
|
|
map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,YQ=="}},
|
|
map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,YQ=="}},
|
|
},
|
|
}}
|
|
if err := validateProviderInput(payload); err != nil {
|
|
t.Fatalf("expected five uploaded images to be accepted, got %v", err)
|
|
}
|
|
payload["messages"].([]any)[0].(map[string]any)["content"] = append(
|
|
payload["messages"].([]any)[0].(map[string]any)["content"].([]any),
|
|
map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,YQ=="}},
|
|
)
|
|
if err := validateProviderInput(payload); !errors.Is(err, errTooManyImages) {
|
|
t.Fatalf("expected six uploaded images to be rejected, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestProviderInFlightLimits(t *testing.T) {
|
|
provider := &homarrProvider{inFlightByUser: make(map[string]int)}
|
|
if !provider.acquireRequest("one") {
|
|
t.Fatal("expected the first user request to be admitted")
|
|
}
|
|
if !provider.acquireRequest("one") {
|
|
t.Fatal("expected the second user request to be admitted")
|
|
}
|
|
if provider.acquireRequest("one") {
|
|
t.Fatal("expected the third user request to be rejected")
|
|
}
|
|
provider.releaseRequest("one")
|
|
provider.releaseRequest("one")
|
|
for index := range maxGlobalInFlight {
|
|
if !provider.acquireRequest(string(rune(index + 1))) {
|
|
t.Fatalf("expected global request %d to be admitted", index)
|
|
}
|
|
}
|
|
if provider.acquireRequest("overflow") {
|
|
t.Fatal("expected the global in-flight limit to reject overflow")
|
|
}
|
|
}
|
|
|
|
func TestCompletionModelAlias(t *testing.T) {
|
|
body, err := aliasCompletionBody([]byte(`{"id":"one","model":"private/model","choices":[]}`))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if strings.Contains(string(body), "private/model") || !strings.Contains(string(body), `"model":"homarr/model"`) {
|
|
t.Fatalf("private upstream model leaked in JSON: %s", body)
|
|
}
|
|
var stream strings.Builder
|
|
_, err = copyAliasedSSE(
|
|
&stream,
|
|
strings.NewReader("data: {\"id\":\"one\",\"model\":\"private/model\",\"choices\":[]}\n\ndata: [DONE]\n\n"),
|
|
1024,
|
|
)
|
|
if err != nil || strings.Contains(stream.String(), "private/model") ||
|
|
!strings.Contains(stream.String(), `"model":"homarr/model"`) {
|
|
t.Fatalf("private upstream model leaked in SSE: %s, %v", stream.String(), err)
|
|
}
|
|
}
|
|
|
|
func TestProviderEnvironment(t *testing.T) {
|
|
t.Setenv("OPENROUTER_API_KEY", "test-key")
|
|
t.Setenv("HOMARR_AI_OPENROUTER_BASE_URL", "https://router.example/v1/")
|
|
t.Setenv("HOMARR_AI_OPENROUTER_MODEL", "mock/team-selected-model")
|
|
t.Setenv("HOMARR_AI_DAILY_REQUEST_LIMIT", "75")
|
|
t.Setenv("HOMARR_AI_GLOBAL_DAILY_REQUEST_LIMIT", "900")
|
|
provider, err := newHomarrProviderFromEnv()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if provider.apiKey != "test-key" || provider.baseURL != "https://router.example/v1" ||
|
|
provider.modelID != "mock/team-selected-model" || provider.dailyLimit != 75 || provider.globalDailyLimit != 900 {
|
|
t.Fatalf("unexpected provider config: %#v", provider)
|
|
}
|
|
|
|
t.Setenv("HOMARR_AI_DAILY_REQUEST_LIMIT", "0")
|
|
if _, err := newHomarrProviderFromEnv(); err == nil {
|
|
t.Fatal("expected invalid daily limit to fail")
|
|
}
|
|
|
|
t.Setenv("HOMARR_AI_DAILY_REQUEST_LIMIT", "50")
|
|
for _, unsafeURL := range []string{
|
|
"https://user:password@router.example/v1",
|
|
"https://router.example/v1?key=secret",
|
|
"https://router.example/v1#fragment",
|
|
} {
|
|
t.Setenv("HOMARR_AI_OPENROUTER_BASE_URL", unsafeURL)
|
|
if _, err := newHomarrProviderFromEnv(); err == nil {
|
|
t.Fatalf("expected unsafe upstream URL %q to fail", unsafeURL)
|
|
}
|
|
}
|
|
t.Setenv("HOMARR_AI_OPENROUTER_BASE_URL", "http://router.example/v1")
|
|
if _, err := newHomarrProviderFromEnv(); err == nil {
|
|
t.Fatal("expected an insecure upstream to fail without an explicit development opt-in")
|
|
}
|
|
}
|
|
|
|
func TestReadBoundedBody(t *testing.T) {
|
|
body, err := readBoundedBody(strings.NewReader("1234"), 4)
|
|
if err != nil || string(body) != "1234" {
|
|
t.Fatalf("unexpected bounded body: %q, %v", body, err)
|
|
}
|
|
if _, err := readBoundedBody(strings.NewReader("12345"), 4); err != errProviderResponseTooLarge {
|
|
t.Fatalf("expected response size error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestCopyBounded(t *testing.T) {
|
|
var target strings.Builder
|
|
written, err := copyBounded(&target, strings.NewReader("1234"), 4)
|
|
if err != nil || written != 4 || target.String() != "1234" {
|
|
t.Fatalf("unexpected bounded copy: %q, %d, %v", target.String(), written, err)
|
|
}
|
|
|
|
target.Reset()
|
|
written, err = copyBounded(&target, strings.NewReader("12345"), 4)
|
|
if err != errProviderResponseTooLarge || written != 4 || target.String() != "1234" {
|
|
t.Fatalf("expected bounded copy failure, got %q, %d, %v", target.String(), written, err)
|
|
}
|
|
}
|