gr2gw/gr2gw_test.go

package main

import (
  "bytes"
  "encoding/json"
  "fmt"
  "net/http"
  "net/http/httptest"
  "strings"
  "testing"
  "time"
)

func TestParseSOCKS5URL(t *testing.T) {
  tests := []struct {
    input    string
    expected *SOCKS5Config
  }{
    {"", nil},
    {"127.0.0.1:1080", &SOCKS5Config{Address: "127.0.0.1:1080"}},
    {"socks5://127.0.0.1:9050", &SOCKS5Config{Address: "127.0.0.1:9050"}},
    {"socks5h://user:pass@10.0.0.1:1080", &SOCKS5Config{Address: "10.0.0.1:1080", Username: "user", Password: "pass"}},
  }

  for _, tc := range tests {
    cfg, err := ParseSOCKS5URL(tc.input)
    if err != nil {
      t.Fatalf("unexpected error for %q: %v", tc.input, err)
    }
    if tc.expected == nil {
      if cfg != nil {
        t.Errorf("expected nil config, got %+v", cfg)
      }
      continue
    }
    if cfg.Address != tc.expected.Address || cfg.Username != tc.expected.Username || cfg.Password != tc.expected.Password {
      t.Errorf("for %q, expected %+v, got %+v", tc.input, tc.expected, cfg)
    }
  }
}

func TestChatMessageGetContentString(t *testing.T) {
  m1 := ChatMessage{Role: "user", Content: "hello world"}
  if m1.GetContentString() != "hello world" {
    t.Errorf("expected 'hello world', got %q", m1.GetContentString())
  }

  m2 := ChatMessage{
    Role: "user",
    Content: []interface{}{
      map[string]interface{}{"type": "text", "text": "part 1 "},
      map[string]interface{}{"type": "text", "text": "part 2"},
    },
  }
  if m2.GetContentString() != "part 1 part 2" {
    t.Errorf("expected 'part 1 part 2', got %q", m2.GetContentString())
  }
}

func TestExtractThinking(t *testing.T) {
  content := "<think>Let me calculate 2+2.</think>The answer is 4."
  clean, reasoning := ExtractThinking(content)
  if reasoning != "Let me calculate 2+2." {
    t.Errorf("expected reasoning 'Let me calculate 2+2.', got %q", reasoning)
  }
  if clean != "The answer is 4." {
    t.Errorf("expected clean 'The answer is 4.', got %q", clean)
  }
}

func TestDetectToolCalls(t *testing.T) {
  xmlContent := `<tool_call>
{"name": "get_weather", "arguments": {"city": "Paris"}}
</tool_call>`
  calls, rem, ok := DetectToolCalls(xmlContent)
  if !ok || len(calls) != 1 {
    t.Fatalf("expected 1 tool call, got %d (ok: %v)", len(calls), ok)
  }
  if calls[0].Function.Name != "get_weather" {
    t.Errorf("expected function name get_weather, got %q", calls[0].Function.Name)
  }
  if rem != "" {
    t.Errorf("expected empty remaining content, got %q", rem)
  }

  jsonContent := `{"name": "calculator", "arguments": {"expr": "1+1"}}`
  calls2, rem2, ok2 := DetectToolCalls(jsonContent)
  if !ok2 || len(calls2) != 1 {
    t.Fatalf("expected 1 tool call from JSON, got %d", len(calls2))
  }
  if calls2[0].Function.Name != "calculator" {
    t.Errorf("expected function calculator, got %q", calls2[0].Function.Name)
  }
  if rem2 != "" {
    t.Errorf("expected empty remaining, got %q", rem2)
  }
}

func TestStreamThinkingFilter(t *testing.T) {
  filter := NewStreamThinkingFilter()
  var contentParts []string
  var reasoningParts []string

  onContent := func(s string) { contentParts = append(contentParts, s) }
  onReasoning := func(s string) { reasoningParts = append(reasoningParts, s) }

  chunks := []string{"<thi", "nk>Thinking de", "eply</th", "ink>Here is your answer."}
  for _, c := range chunks {
    filter.Feed(c, onContent, onReasoning)
  }
  filter.Flush(onContent, onReasoning)

  fullReasoning := strings.Join(reasoningParts, "")
  fullContent := strings.Join(contentParts, "")

  if fullReasoning != "Thinking deeply" {
    t.Errorf("expected reasoning 'Thinking deeply', got %q", fullReasoning)
  }
  if fullContent != "Here is your answer." {
    t.Errorf("expected content 'Here is your answer.', got %q", fullContent)
  }
}

func TestStreamToolCallFilter(t *testing.T) {
  filter := NewStreamToolCallFilter()
  var contentParts []string
  var toolCalls []ToolCall

  onContent := func(s string) { contentParts = append(contentParts, s) }
  onToolCall := func(tc ToolCall) { toolCalls = append(toolCalls, tc) }

  chunks := []string{
    "Searching now: ",
    "<tool_c",
    "all>\n{\"name\": \"search_web\", \"arguments\": {\"query\": \"golang\"}}\n</tool_",
    "call>",
    " Done.",
  }

  for _, c := range chunks {
    filter.Feed(c, onContent, onToolCall)
  }
  filter.Flush(onContent, onToolCall)

  if len(toolCalls) != 1 {
    t.Fatalf("expected 1 emitted tool call, got %d", len(toolCalls))
  }
  if toolCalls[0].Function.Name != "search_web" {
    t.Errorf("expected tool name 'search_web', got %q", toolCalls[0].Function.Name)
  }
  fullContent := strings.Join(contentParts, "")
  if fullContent != "Searching now:  Done." {
    t.Errorf("expected 'Searching now:  Done.', got %q", fullContent)
  }
}

func TestMockGradioServerCompletion(t *testing.T) {
  // Setup a mock Gradio server
  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/gradio_api/info" {
      resp := GradioAPIInfoResponse{
        NamedEndpoints: map[string]GradioEndpointInfo{
          "/chat_fn": {
            Parameters: []GradioParamInfo{
              {ParameterName: "message", Component: "Textbox"},
            },
            Returns: []GradioParamInfo{
              {ParameterName: "response", Component: "Json"},
            },
          },
        },
      }
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(resp)
      return
    }

    if r.URL.Path == "/gradio_api/call/chat_fn" {
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(GradioJoinResponse{EventID: "evt_123"})
      return
    }

    if r.URL.Path == "/gradio_api/call/chat_fn/evt_123" {
      w.Header().Set("Content-Type", "text/event-stream")
      flusher, ok := w.(http.Flusher)
      if !ok {
        t.Fatal("expected flusher")
      }
      fmt.Fprintf(w, "event: generating\ndata: [\"Hello \", null]\n\n")
      flusher.Flush()
      fmt.Fprintf(w, "event: generating\ndata: [\"Hello world!\", null]\n\n")
      flusher.Flush()
      fmt.Fprintf(w, "event: complete\ndata: [\"Hello world!\", null]\n\n")
      flusher.Flush()
      return
    }

    http.NotFound(w, r)
  }))
  defer ts.Close()

  gw := NewGradioGateway(ts.URL, "", 10*time.Second)

  // 1. Test Non-streaming request
  reqBody := ChatCompletionRequest{
    Model: "test-model",
    Messages: []ChatMessage{
      {Role: "user", Content: "Hi"},
    },
    Stream: false,
  }
  b, _ := json.Marshal(reqBody)
  httpReq := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b))
  httpReq.Header.Set("Content-Type", "application/json")
  rec := httptest.NewRecorder()

  err := gw.ExecuteChatCompletion(rec, httpReq, reqBody)
  if err != nil {
    t.Fatalf("unexpected completion error: %v", err)
  }

  var resp ChatCompletionResponse
  if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
    t.Fatalf("failed to decode completion response: %v", err)
  }
  if len(resp.Choices) != 1 {
    t.Fatalf("expected 1 choice, got %d", len(resp.Choices))
  }
  if resp.Choices[0].Message.GetContentString() != "Hello world!" {
    t.Errorf("expected 'Hello world!', got %q", resp.Choices[0].Message.GetContentString())
  }

  // 2. Test Streaming request
  reqBodyStream := reqBody
  reqBodyStream.Stream = true
  recStream := httptest.NewRecorder()
  err = gw.ExecuteChatCompletion(recStream, httpReq, reqBodyStream)
  if err != nil {
    t.Fatalf("unexpected streaming error: %v", err)
  }
  streamOutput := recStream.Body.String()
  if !strings.Contains(streamOutput, "data: [DONE]") {
    t.Errorf("expected stream to contain [DONE], got:\n%s", streamOutput)
  }
  if !strings.Contains(streamOutput, "Hello world!") && !strings.Contains(streamOutput, "world!") {
    t.Errorf("expected stream output to contain delta tokens, got:\n%s", streamOutput)
  }
}

func TestParseGradioStreamOutput(t *testing.T) {
  // 1. Standard 1D Gradio array
  frame1 := ParseGradioStreamOutput(`["Hello from 1D", null]`)
  if !frame1.OK || frame1.Content != "Hello from 1D" || frame1.Reasoning != "" || len(frame1.ToolCalls) != 0 {
    t.Errorf("unexpected frame1: %+v", frame1)
  }

  // 2. Hy3 2D array frame with reasoning
  hy3Raw := `[["Hello answer", "Let me think deeply...", [], [{"role": "user", "content": "hi"}]]]`
  frame2 := ParseGradioStreamOutput(hy3Raw)
  if !frame2.OK || frame2.Content != "Hello answer" || frame2.Reasoning != "Let me think deeply..." || len(frame2.ToolCalls) != 0 {
    t.Errorf("unexpected frame2: %+v", frame2)
  }

  // 3. Hy3 2D array frame with tool calls
  hy3ToolRaw := `[["", "Calling weather tool", [{"id": "call_abc", "type": "function", "function": {"name": "get_weather", "arguments": "{\"city\": \"Tokyo\"}"}}], []]]`
  frame3 := ParseGradioStreamOutput(hy3ToolRaw)
  if !frame3.OK || frame3.Content != "" || frame3.Reasoning != "Calling weather tool" || len(frame3.ToolCalls) != 1 {
    t.Fatalf("unexpected frame3: %+v", frame3)
  }
  if frame3.ToolCalls[0].ID != "call_abc" || frame3.ToolCalls[0].Function.Name != "get_weather" {
    t.Errorf("unexpected tool call in frame3: %+v", frame3.ToolCalls[0])
  }

  // 4. Chat pairs
  pairRaw := `[[["user prompt", "assistant answer"]]]`
  frame4 := ParseGradioStreamOutput(pairRaw)
  if !frame4.OK || frame4.Content != "assistant answer" {
    t.Errorf("unexpected frame4: %+v", frame4)
  }

  // 5. Messages array
  msgRaw := `[[{"role": "assistant", "content": "msg answer", "reasoning_content": "msg think"}]]`
  frame5 := ParseGradioStreamOutput(msgRaw)
  if !frame5.OK || frame5.Content != "msg answer" || frame5.Reasoning != "msg think" {
    t.Errorf("unexpected frame5: %+v", frame5)
  }
  // 6. Hy3 with null content and reasoning
  hy3NullContent := `[[null, "Thinking process...", null]]`
  frame6 := ParseGradioStreamOutput(hy3NullContent)
  if !frame6.OK || frame6.Content != "" || frame6.Reasoning != "Thinking process..." {
    t.Errorf("unexpected frame6: %+v", frame6)
  }

  // 7. Hy3 with 3 elements (answer, reasoning, null tool calls)
  hy3ThreeElem := `[["Mocked answer", "Mocked reasoning", null]]`
  frame7 := ParseGradioStreamOutput(hy3ThreeElem)
  if !frame7.OK || frame7.Content != "Mocked answer" || frame7.Reasoning != "Mocked reasoning" || len(frame7.ToolCalls) != 0 {
    t.Errorf("unexpected frame7: %+v", frame7)
  }

  // 8. Gradio 6 Chatbot multi-output format with TextMessage list
  g6Raw := `["", [{"role": "user", "content": [{"text": "weather in Tokyo?", "type": "text"}]}, {"role": "assistant", "content": [{"text": "<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"Tokyo\"}}\n</tool_call>", "type": "text"}]}]]`
  frame8 := ParseGradioStreamOutput(g6Raw)
  if !frame8.OK || len(frame8.ToolCalls) != 1 || frame8.ToolCalls[0].Function.Name != "get_weather" {
    t.Errorf("unexpected frame8: %+v", frame8)
  }
}

func TestHunyuan3BuildPayload(t *testing.T) {
  gw := &GradioGateway{}
  disc := &SpaceDiscovery{
    TotalInputs:        9,
    MessageIndex:       0,
    SystemIndex:        1,
    HistoryIndex:       2,
    ThinkLevelIndex:    3,
    TempIndex:          4,
    MaxTokensIndex:     5,
    TopPIndex:          6,
    FunctionsJSONIndex: 8,
    IsHunyuan3:         true,
  }

  temp := 0.2
  req := ChatCompletionRequest{
    Model:           "hy3",
    ReasoningEffort: "low",
    Temperature:     &temp,
    Tools: []Tool{
      {
        Type: "function",
        Function: map[string]interface{}{
          "name": "calc",
        },
      },
    },
    Messages: []ChatMessage{
      {Role: "system", Content: "Be helpful"},
      {Role: "user", Content: "2+2"},
      {
        Role:             "assistant",
        ReasoningContent: "Thinking...",
        ToolCalls: []ToolCall{
          {ID: "c1", Type: "function", Function: ToolCallFunction{Name: "calc", Arguments: `{"expr":"2+2"}`}},
        },
      },
      {Role: "tool", ToolCallID: "c1", Content: "4"},
    },
  }

  data, err := gw.BuildGradioPayload(disc, req)
  if err != nil {
    t.Fatalf("BuildGradioPayload failed: %v", err)
  }

  if len(data) != 9 {
    t.Fatalf("expected 9 payload items, got %d", len(data))
  }

  // Message parameter (0): should be tool result prompt matching hygate behavior
  if msg, ok := data[0].(string); !ok || msg != "Tool result for c1: 4" {
    t.Errorf("expected 'Tool result for c1: 4', got %v", data[0])
  }

  // System parameter (1)
  if sys, ok := data[1].(string); !ok || sys != "Be helpful" {
    t.Errorf("expected 'Be helpful', got %v", data[1])
  }

  // History parameter (2): should contain prior turns (user, assistant with tool calls)
  hist, ok := data[2].([]map[string]interface{})
  if !ok {
    t.Fatalf("expected history slice of maps, got %T", data[2])
  }
  if len(hist) != 2 {
    t.Fatalf("expected 2 history items (user, assistant), got %d", len(hist))
  }
  if hist[0]["role"] != "user" || hist[0]["content"] != "2+2" {
    t.Errorf("unexpected user history entry: %+v", hist[0])
  }
  if hist[1]["role"] != "assistant" || len(hist[1]["tool_calls"].([]ToolCall)) != 1 {
    t.Errorf("unexpected assistant history entry: %+v", hist[1])
  }

  // ThinkLevel parameter (3)
  if data[3] != "low" {
    t.Errorf("expected think_level 'low', got %v", data[3])
  }

  // Temp parameter (4)
  if data[4] != 0.2 {
    t.Errorf("expected temp 0.2, got %v", data[4])
  }

  // FunctionsJSON parameter (8)
  fnStr, ok := data[8].(string)
  if !ok || !strings.Contains(fnStr, "calc") {
    t.Errorf("expected functions_json_str to contain 'calc', got %v", data[8])
  }

  // Test multi-tool turn: 2 tool messages at the end
  reqMulti := req
  reqMulti.Messages = append(reqMulti.Messages, ChatMessage{Role: "tool", Name: "fetch", Content: "done"})
  dataMulti, err := gw.BuildGradioPayload(disc, reqMulti)
  if err != nil {
    t.Fatalf("BuildGradioPayload failed on multi-tool: %v", err)
  }
  // The last tool message is data[0]
  if msg, ok := dataMulti[0].(string); !ok || msg != "Tool result for fetch: done" {
    t.Errorf("expected 'Tool result for fetch: done', got %v", dataMulti[0])
  }
  // The first tool message is in history
  histMulti, _ := dataMulti[2].([]map[string]interface{})
  if len(histMulti) != 3 {
    t.Fatalf("expected 3 history items (user, assistant, tool 1), got %d", len(histMulti))
  }
  if histMulti[2]["role"] != "tool" || histMulti[2]["content"] != "4" {
    t.Errorf("unexpected tool 1 in history: %+v", histMulti[2])
  }
}

func TestHunyuan3MockServerCompletion(t *testing.T) {
  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/gradio_api/info" {
      resp := GradioAPIInfoResponse{
        NamedEndpoints: map[string]GradioEndpointInfo{
          "/chat": {
            Parameters: []GradioParamInfo{
              {ParameterName: "message"},
              {ParameterName: "system_prompt"},
              {ParameterName: "history"},
              {ParameterName: "think_level"},
              {ParameterName: "temperature"},
              {ParameterName: "max_tokens"},
              {ParameterName: "top_p"},
              {ParameterName: "preserved_thinking"},
              {ParameterName: "functions_json_str"},
            },
          },
        },
      }
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(resp)
      return
    }

    if r.URL.Path == "/gradio_api/call/chat" {
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(GradioJoinResponse{EventID: "evt_hy3"})
      return
    }

    if r.URL.Path == "/gradio_api/call/chat/evt_hy3" {
      w.Header().Set("Content-Type", "text/event-stream")
      flusher, ok := w.(http.Flusher)
      if !ok {
        t.Fatal("expected flusher")
      }
      // Frame 1: Reasoning delta
      fmt.Fprintf(w, "event: generating\ndata: [[\"\", \"Reasoning part 1 \", [], []]]\n\n")
      flusher.Flush()
      // Frame 2: Tool call initiated
      fmt.Fprintf(w, "event: generating\ndata: [[\"\", \"Reasoning part 1 and 2\", [{\"id\": \"call_hy3\", \"type\": \"function\", \"function\": {\"name\": \"search\", \"arguments\": \"{\\\"q\\\": \\\"tencent\\\"}\"}}], []]]\n\n")
      flusher.Flush()
      // Frame 3: Completion
      fmt.Fprintf(w, "event: complete\ndata: [[\"\", \"Reasoning part 1 and 2\", [{\"id\": \"call_hy3\", \"type\": \"function\", \"function\": {\"name\": \"search\", \"arguments\": \"{\\\"q\\\": \\\"tencent\\\"}\"}}], []]]\n\n")
      flusher.Flush()
      return
    }

    http.NotFound(w, r)
  }))
  defer ts.Close()

  gw := NewGradioGateway(ts.URL, "", 10*time.Second)

  // 1. Non-streaming tool call test
  reqBody := ChatCompletionRequest{
    Model: "hy3",
    Messages: []ChatMessage{
      {Role: "user", Content: "search for tencent"},
    },
    Stream: false,
  }
  b, _ := json.Marshal(reqBody)
  httpReq := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b))
  httpReq.Header.Set("Content-Type", "application/json")
  rec := httptest.NewRecorder()

  err := gw.ExecuteChatCompletion(rec, httpReq, reqBody)
  if err != nil {
    t.Fatalf("unexpected completion error: %v", err)
  }

  var resp ChatCompletionResponse
  if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
    t.Fatalf("failed to decode response: %v", err)
  }

  if resp.Choices[0].FinishReason != "tool_calls" {
    t.Errorf("expected finish_reason 'tool_calls', got %q", resp.Choices[0].FinishReason)
  }
  if resp.Choices[0].Message.ReasoningContent != "Reasoning part 1 and 2" {
    t.Errorf("expected native reasoning, got %q", resp.Choices[0].Message.ReasoningContent)
  }
  if len(resp.Choices[0].Message.ToolCalls) != 1 {
    t.Fatalf("expected 1 tool call, got %d", len(resp.Choices[0].Message.ToolCalls))
  }
  if resp.Choices[0].Message.ToolCalls[0].Function.Name != "search" {
    t.Errorf("expected function 'search', got %q", resp.Choices[0].Message.ToolCalls[0].Function.Name)
  }

  // 2. Streaming tool call test
  reqBodyStream := reqBody
  reqBodyStream.Stream = true
  recStream := httptest.NewRecorder()
  err = gw.ExecuteChatCompletion(recStream, httpReq, reqBodyStream)
  if err != nil {
    t.Fatalf("unexpected streaming error: %v", err)
  }

  streamOut := recStream.Body.String()
  if !strings.Contains(streamOut, "reasoning_content") {
    t.Errorf("expected stream to contain reasoning_content, got:\n%s", streamOut)
  }
  if !strings.Contains(streamOut, "tool_calls") {
    t.Errorf("expected stream to contain tool_calls, got:\n%s", streamOut)
  }
  if !strings.Contains(streamOut, "call_hy3") {
    t.Errorf("expected stream to contain tool call ID call_hy3, got:\n%s", streamOut)
  }
  if !strings.Contains(streamOut, "\"finish_reason\":\"tool_calls\"") {
    t.Errorf("expected stream finish_reason tool_calls, got:\n%s", streamOut)
  }
}

func TestUniversalToolCallingTransformMessages(t *testing.T) {
  req := ChatCompletionRequest{
    Tools: []Tool{
      {
        Type: "function",
        Function: map[string]interface{}{
          "name":        "get_weather",
          "description": "Get current weather",
        },
      },
    },
    Messages: []ChatMessage{
      {Role: "user", Content: "What is the weather in Tokyo and Paris?"},
      {
        Role: "assistant",
        ToolCalls: []ToolCall{
          {ID: "call_tokyo", Type: "function", Function: ToolCallFunction{Name: "get_weather", Arguments: `{"city":"Tokyo"}`}},
          {ID: "call_paris", Type: "function", Function: ToolCallFunction{Name: "get_weather", Arguments: `{"city":"Paris"}`}},
        },
      },
      {Role: "tool", ToolCallID: "call_tokyo", Content: `{"temp": 20}`},
      {Role: "tool", ToolCallID: "call_paris", Content: `{"temp": 15}`},
    },
  }

  processed, toolInstruction, hasSystem := TransformMessages(req)
  if !hasSystem {
    t.Errorf("expected hasSystem to be true after injecting tool instructions")
  }
  if toolInstruction == "" {
    t.Errorf("expected non-empty toolInstruction")
  }

  // Expect:
  // [0] System message with tool instructions
  // [1] User message: "What is the weather in Tokyo and Paris?"
  // [2] Assistant message with <tool_call> blocks
  // [3] User message with coalesced <tool_response> blocks
  if len(processed) != 4 {
    t.Fatalf("expected 4 processed messages, got %d", len(processed))
  }

  if processed[0].Role != "system" || !strings.Contains(processed[0].GetContentString(), "API router") {
    t.Errorf("unexpected message 0: %+v", processed[0])
  }

  if processed[1].Role != "user" || processed[1].GetContentString() != "What is the weather in Tokyo and Paris?" {
    t.Errorf("unexpected message 1: %+v", processed[1])
  }

  if processed[2].Role != "assistant" || !strings.Contains(processed[2].GetContentString(), "get_weather") {
    t.Errorf("unexpected message 2: %+v", processed[2])
  }

  respContent := processed[3].GetContentString()
  if processed[3].Role != "user" {
    t.Errorf("expected coalesced message 3 to have role user, got %q", processed[3].Role)
  }
  if !strings.Contains(respContent, `{"name": "get_weather", "content": {"temp": 20}}`) {
    t.Errorf("expected resolved function name get_weather for tokyo, got:\n%s", respContent)
  }
  if !strings.Contains(respContent, `{"name": "get_weather", "content": {"temp": 15}}`) {
    t.Errorf("expected resolved function name get_weather for paris, got:\n%s", respContent)
  }
  if !strings.Contains(respContent, "Please answer the user's request based on the tool results.") {
    t.Errorf("expected continuation prompt in coalesced message, got:\n%s", respContent)
  }
}

func TestUniversalToolCallDetectionVariants(t *testing.T) {
  // 1. Array of tool calls inside <tool_calls> tag
  multiXML := `<tool_calls>
[
  {"name": "get_weather", "arguments": {"city": "Tokyo"}},
  {"name": "get_weather", "arguments": {"city": "Paris"}}
]
</tool_calls>`
  calls1, rem1, ok1 := DetectToolCalls(multiXML)
  if !ok1 || len(calls1) != 2 {
    t.Fatalf("expected 2 tool calls from <tool_calls>, got %d", len(calls1))
  }
  if calls1[0].Function.Name != "get_weather" || calls1[1].Function.Name != "get_weather" {
    t.Errorf("unexpected function names: %+v", calls1)
  }
  if rem1 != "" {
    t.Errorf("expected empty remaining, got %q", rem1)
  }

  // 2. <function_call> tag
  fnCallXML := `Some preamble before call.
<function_call>
{"name": "search", "arguments": {"q": "golang"}}
</function_call>
Some postamble.`
  calls2, rem2, ok2 := DetectToolCalls(fnCallXML)
  if !ok2 || len(calls2) != 1 {
    t.Fatalf("expected 1 tool call from <function_call>, got %d", len(calls2))
  }
  if calls2[0].Function.Name != "search" {
    t.Errorf("expected function search, got %q", calls2[0].Function.Name)
  }
  if strings.Contains(rem2, "function_call") {
    t.Errorf("expected tag stripped from remaining, got %q", rem2)
  }
  if !strings.Contains(rem2, "Some preamble") || !strings.Contains(rem2, "Some postamble") {
    t.Errorf("expected surrounding text preserved in remaining, got %q", rem2)
  }

  // 3. [TOOL_CALLS] bracket syntax
  bracketXML := `[TOOL_CALLS]
{"name": "calculate", "arguments": {"x": 42}}
[/TOOL_CALLS]`
  calls3, rem3, ok3 := DetectToolCalls(bracketXML)
  if !ok3 || len(calls3) != 1 {
    t.Fatalf("expected 1 tool call from [TOOL_CALLS], got %d", len(calls3))
  }
  if calls3[0].Function.Name != "calculate" {
    t.Errorf("expected function calculate, got %q", calls3[0].Function.Name)
  }
  if rem3 != "" {
    t.Errorf("expected empty remaining, got %q", rem3)
  }

  // 4. Raw JSON array without tags
  rawArray := `[{"name": "f1", "arguments": {}}, {"name": "f2", "arguments": {}}]`
  calls4, rem4, ok4 := DetectToolCalls(rawArray)
  if !ok4 || len(calls4) != 2 {
    t.Fatalf("expected 2 calls from raw array, got %d", len(calls4))
  }
  if calls4[0].Function.Name != "f1" || calls4[1].Function.Name != "f2" {
    t.Errorf("unexpected names from raw array: %+v", calls4)
  }
  if rem4 != "" {
    t.Errorf("expected empty remaining, got %q", rem4)
  }
}

func TestUniversalStreamToolCallFilterVariants(t *testing.T) {
  filter := NewStreamToolCallFilter()
  var contentParts []string
  var toolCalls []ToolCall

  onContent := func(s string) { contentParts = append(contentParts, s) }
  onToolCall := func(tc ToolCall) { toolCalls = append(toolCalls, tc) }

  // Stream using [TOOL_CALLS] across multiple chunk boundaries
  chunks := []string{
    "Preamble text: ",
    "[TOOL_",
    "CALLS]\n{\"name\": \"browse\", \"arguments\": {\"url\": \"example.com\"}}\n[/TOOL_",
    "CALLS]",
    " Completed.",
  }

  for _, c := range chunks {
    filter.Feed(c, onContent, onToolCall)
  }
  filter.Flush(onContent, onToolCall)

  if len(toolCalls) != 1 {
    t.Fatalf("expected 1 tool call from stream filter, got %d", len(toolCalls))
  }
  if toolCalls[0].Function.Name != "browse" {
    t.Errorf("expected function browse, got %q", toolCalls[0].Function.Name)
  }
  if !filter.emittedCall {
    t.Errorf("expected emittedCall to be true")
  }
  fullContent := strings.Join(contentParts, "")
  if strings.Contains(fullContent, "TOOL_CALLS") {
    t.Errorf("tag leaked into stream content: %q", fullContent)
  }
  if fullContent != "Preamble text:  Completed." {
    t.Errorf("unexpected streamed content: %q", fullContent)
  }
}

func TestBuildGradioPayloadGenericSpaces(t *testing.T) {
  gw := &GradioGateway{}

  req := ChatCompletionRequest{
    Tools: []Tool{
      {Type: "function", Function: map[string]interface{}{"name": "lookup"}},
    },
    Messages: []ChatMessage{
      {Role: "user", Content: "What is 10+10?"},
      {
        Role: "assistant",
        ToolCalls: []ToolCall{
          {ID: "c1", Type: "function", Function: ToolCallFunction{Name: "lookup", Arguments: `{"q":"10+10"}`}},
        },
      },
      {Role: "tool", ToolCallID: "c1", Content: `{"result": 20}`},
    },
  }

  // 1. Space with native system prompt input (SystemIndex: 0, MessageIndex: 1, HistoryIndex: 2)
  discWithSystem := NewDefaultSpaceDiscovery("https://space-1.hf.space")
  discWithSystem.TotalInputs = 3
  discWithSystem.SystemIndex = 0
  discWithSystem.MessageIndex = 1
  discWithSystem.HistoryIndex = 2
  discWithSystem.HistoryFormat = "pairs"

  data1, err := gw.BuildGradioPayload(discWithSystem, req)
  if err != nil {
    t.Fatalf("failed to build payload 1: %v", err)
  }
  sysStr, ok := data1[0].(string)
  if !ok || !strings.Contains(sysStr, "API router") {
    t.Errorf("expected system prompt at index 0, got %v", data1[0])
  }
  msgStr, ok := data1[1].(string)
  if !ok || !strings.Contains(msgStr, "Please answer the user's request based on the tool result.") {
    t.Errorf("expected coalesced tool prompt at index 1, got %v", data1[1])
  }
  pairs1, ok := data1[2].([][]string)
  if !ok || len(pairs1) != 1 {
    t.Fatalf("expected 1 history pair at index 2, got %T (%v)", data1[2], data1[2])
  }
  if pairs1[0][0] != "What is 10+10?" || !strings.Contains(pairs1[0][1], "lookup") {
    t.Errorf("unexpected history pair: %+v", pairs1[0])
  }

  // 2. Space without system prompt (SystemIndex: -1, MessageIndex: 0, HistoryIndex: 1)
  discNoSystem := NewDefaultSpaceDiscovery("https://space-2.hf.space")
  discNoSystem.TotalInputs = 2
  discNoSystem.SystemIndex = -1
  discNoSystem.MessageIndex = 0
  discNoSystem.HistoryIndex = 1
  discNoSystem.HistoryFormat = "pairs"

  data2, err := gw.BuildGradioPayload(discNoSystem, req)
  if err != nil {
    t.Fatalf("failed to build payload 2: %v", err)
  }
  pairs2, ok := data2[1].([][]string)
  if !ok || len(pairs2) != 1 {
    t.Fatalf("expected 1 history pair at index 1, got %T (%v)", data2[1], data2[1])
  }
  // Instructions prepended to the first user turn:
  if !strings.Contains(pairs2[0][0], "API router") || !strings.Contains(pairs2[0][0], "What is 10+10?") {
    t.Errorf("expected system instructions prepended to first pair user message, got: %q", pairs2[0][0])
  }
  if !strings.Contains(pairs2[0][1], "lookup") {
    t.Errorf("expected assistant tool call in pair bot turn, got: %q", pairs2[0][1])
  }

  // 3. Single-textbox space (SystemIndex: -1, MessageIndex: 0, HistoryIndex: -1)
  discSingleInput := NewDefaultSpaceDiscovery("https://space-3.hf.space")
  discSingleInput.TotalInputs = 1
  discSingleInput.SystemIndex = -1
  discSingleInput.MessageIndex = 0
  discSingleInput.HistoryIndex = -1

  data3, err := gw.BuildGradioPayload(discSingleInput, req)
  if err != nil {
    t.Fatalf("failed to build payload 3: %v", err)
  }
  transcript, ok := data3[0].(string)
  if !ok {
    t.Fatalf("expected string transcript, got %T", data3[0])
  }
  if !strings.Contains(transcript, "# Instructions") || !strings.Contains(transcript, "User: What is 10+10?") || !strings.Contains(transcript, "Assistant: <tool_call>") || !strings.Contains(transcript, "# Current Request") {
    t.Errorf("unexpected single-input transcript: %s", transcript)
  }
}

func TestGenericSpaceMockServerToolCalling(t *testing.T) {
  var lastReceivedData []interface{}

  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/gradio_api/info" {
      resp := GradioAPIInfoResponse{
        NamedEndpoints: map[string]GradioEndpointInfo{
          "/chat_fn": {
            Parameters: []GradioParamInfo{
              {ParameterName: "system_prompt"},
              {ParameterName: "message"},
              {ParameterName: "history"},
            },
          },
        },
      }
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(resp)
      return
    }

    if r.URL.Path == "/gradio_api/call/chat_fn" {
      var body map[string]interface{}
      json.NewDecoder(r.Body).Decode(&body)
      if dataSlice, ok := body["data"].([]interface{}); ok {
        lastReceivedData = dataSlice
      }
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(GradioJoinResponse{EventID: "evt_generic"})
      return
    }

    if r.URL.Path == "/gradio_api/call/chat_fn/evt_generic" {
      w.Header().Set("Content-Type", "text/event-stream")
      flusher, ok := w.(http.Flusher)
      if !ok {
        t.Fatal("expected flusher")
      }
      msgStr := ""
      if len(lastReceivedData) > 1 {
        msgStr, _ = lastReceivedData[1].(string)
      }

      if strings.Contains(msgStr, "<tool_response>") {
        // Turn 2: answer
        fmt.Fprintf(w, "event: generating\ndata: [\"The weather in Tokyo is 20 C.\", null]\n\n")
        flusher.Flush()
        fmt.Fprintf(w, "event: complete\ndata: [\"The weather in Tokyo is 20 C.\", null]\n\n")
        flusher.Flush()
      } else {
        // Turn 1: tool call
        fmt.Fprintf(w, "event: generating\ndata: [\"<tool_call>\\n{\\\"name\\\": \\\"get_weather\\\", \\\"arguments\\\": {\\\"city\\\": \\\"Tokyo\\\"}}\\n</tool_call>\", null]\n\n")
        flusher.Flush()
        fmt.Fprintf(w, "event: complete\ndata: [\"<tool_call>\\n{\\\"name\\\": \\\"get_weather\\\", \\\"arguments\\\": {\\\"city\\\": \\\"Tokyo\\\"}}\\n</tool_call>\", null]\n\n")
        flusher.Flush()
      }
      return
    }

    http.NotFound(w, r)
  }))
  defer ts.Close()

  gw := NewGradioGateway(ts.URL, "", 10*time.Second)

  // Turn 1: User question with tools
  req1 := ChatCompletionRequest{
    Model: "generic-bot",
    Tools: []Tool{
      {Type: "function", Function: map[string]interface{}{"name": "get_weather"}},
    },
    Messages: []ChatMessage{
      {Role: "user", Content: "Weather in Tokyo?"},
    },
    Stream: false,
  }

  b1, _ := json.Marshal(req1)
  httpReq1 := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b1))
  rec1 := httptest.NewRecorder()

  err := gw.ExecuteChatCompletion(rec1, httpReq1, req1)
  if err != nil {
    t.Fatalf("Turn 1 execution failed: %v", err)
  }

  var resp1 ChatCompletionResponse
  if err := json.NewDecoder(rec1.Body).Decode(&resp1); err != nil {
    t.Fatalf("Turn 1 decode failed: %v", err)
  }
  if resp1.Choices[0].FinishReason != "tool_calls" {
    t.Fatalf("expected finish_reason 'tool_calls', got %q", resp1.Choices[0].FinishReason)
  }
  if len(resp1.Choices[0].Message.ToolCalls) != 1 {
    t.Fatalf("expected 1 tool call, got %d", len(resp1.Choices[0].Message.ToolCalls))
  }
  tc := resp1.Choices[0].Message.ToolCalls[0]
  if tc.Function.Name != "get_weather" {
    t.Fatalf("expected function name get_weather, got %q", tc.Function.Name)
  }

  // Turn 2: Send tool response
  req2 := ChatCompletionRequest{
    Model: "generic-bot",
    Tools: []Tool{
      {Type: "function", Function: map[string]interface{}{"name": "get_weather"}},
    },
    Messages: []ChatMessage{
      {Role: "user", Content: "Weather in Tokyo?"},
      resp1.Choices[0].Message,
      {Role: "tool", ToolCallID: tc.ID, Content: `{"temp": 20}`},
    },
    Stream: false,
  }

  b2, _ := json.Marshal(req2)
  httpReq2 := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b2))
  rec2 := httptest.NewRecorder()

  err = gw.ExecuteChatCompletion(rec2, httpReq2, req2)
  if err != nil {
    t.Fatalf("Turn 2 execution failed: %v", err)
  }

  var resp2 ChatCompletionResponse
  if err := json.NewDecoder(rec2.Body).Decode(&resp2); err != nil {
    t.Fatalf("Turn 2 decode failed: %v", err)
  }
  if resp2.Choices[0].FinishReason != "stop" {
    t.Errorf("expected finish_reason 'stop', got %q", resp2.Choices[0].FinishReason)
  }
  if resp2.Choices[0].Message.GetContentString() != "The weather in Tokyo is 20 C." {
    t.Errorf("expected final answer, got %q", resp2.Choices[0].Message.GetContentString())
  }
}

func TestExtractGradioErrorMessage(t *testing.T) {
  cases := []struct {
    input    string
    expected string
  }{
    {
      input:    `{"error": "Client error '402 Payment Required'"}`,
      expected: "Client error '402 Payment Required'",
    },
    {
      input:    `{"message": "Rate limit exceeded"}`,
      expected: "Rate limit exceeded",
    },
    {
      input:    `{"error": null, "title": "Validation Error"}`,
      expected: "Validation Error",
    },
    {
      input:    `{"error": null}`,
      expected: "internal space error (check Gradio inputs/types)",
    },
    {
      input:    `raw server failure`,
      expected: "raw server failure",
    },
  }

  for _, c := range cases {
    got := extractGradioErrorMessage(c.input)
    if got != c.expected {
      t.Errorf("extractGradioErrorMessage(%q) = %q, expected %q", c.input, got, c.expected)
    }
  }
}

func TestStreamerLazyStartAndErrors(t *testing.T) {
  rec := httptest.NewRecorder()
  s := NewStreamer(rec, nil, "cmpl-1", 12345, "test-model")
  if s.started {
    t.Errorf("expected streamer to start as not started")
  }

  // Ensure Role/headers only sent on first write
  s.Content("Hello")
  if !s.started {
    t.Errorf("expected streamer to be started after Content")
  }
  s.Finish("stop")
  s.Done()

  body := rec.Body.String()
  if !strings.Contains(body, `"role":"assistant"`) {
    t.Errorf("expected role chunk in body: %s", body)
  }
  if !strings.Contains(body, `"content":"Hello"`) {
    t.Errorf("expected content chunk in body: %s", body)
  }
  if !strings.Contains(body, "[DONE]") {
    t.Errorf("expected [DONE] in body: %s", body)
  }
}

func TestUpstreamGradioErrorPropagation(t *testing.T) {
  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/gradio_api/info" {
      resp := GradioAPIInfoResponse{
        NamedEndpoints: map[string]GradioEndpointInfo{
          "/chat_fn": {
            Parameters: []GradioParamInfo{
              {ParameterName: "message", Component: "Textbox"},
            },
          },
        },
      }
      json.NewEncoder(w).Encode(resp)
      return
    }
    if r.URL.Path == "/gradio_api/call/chat_fn" {
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(GradioJoinResponse{EventID: "err-event"})
      return
    }
    if r.URL.Path == "/gradio_api/call/chat_fn/err-event" {
      w.Header().Set("Content-Type", "text/event-stream")
      flusher := w.(http.Flusher)
      fmt.Fprintf(w, "event: error\ndata: {\"error\": \"Quota exceeded: 402 Payment Required\"}\n\n")
      flusher.Flush()
      return
    }
    http.NotFound(w, r)
  }))
  defer ts.Close()

  gw := NewGradioGateway(ts.URL, "", 10*time.Second)

  // 1. Non-streaming error should return error with upstream message
  recNonStream := httptest.NewRecorder()
  httpReqNonStream := httptest.NewRequest("POST", "/v1/chat/completions", nil)
  err := gw.ExecuteChatCompletion(recNonStream, httpReqNonStream, ChatCompletionRequest{
    Messages: []ChatMessage{{Role: "user", Content: "Hello"}},
    Stream:   false,
  })
  if err == nil {
    t.Fatalf("expected error from non-streaming upstream failure, got nil")
  }
  if !strings.Contains(err.Error(), "402 Payment Required") {
    t.Errorf("expected error to mention 402 Payment Required, got %v", err)
  }

  // 2. Streaming error before start should return error with upstream message
  recStream := httptest.NewRecorder()
  httpReqStream := httptest.NewRequest("POST", "/v1/chat/completions", nil)
  err = gw.ExecuteChatCompletion(recStream, httpReqStream, ChatCompletionRequest{
    Messages: []ChatMessage{{Role: "user", Content: "Hello"}},
    Stream:   true,
  })
  if err == nil {
    t.Fatalf("expected error from streaming upstream failure, got nil")
  }
  if !strings.Contains(err.Error(), "402 Payment Required") {
    t.Errorf("expected error to mention 402 Payment Required, got %v", err)
  }
  // And nothing should have been written to body
  if strings.Contains(recStream.Body.String(), "[DONE]") {
    t.Errorf("did not expect [DONE] on upstream error: %s", recStream.Body.String())
  }
}

func TestExtractFailedGeneration(t *testing.T) {
  cases := []struct {
    name     string
    input    string
    expected string
  }{
    {
      name:     "json wrapped",
      input:    `{"error": {"code": 400, "failed_generation": "{\"name\": \"test_fn\", \"arguments\": {\"a\": 1}}"}}`,
      expected: `{"name": "test_fn", "arguments": {"a": 1}}`,
    },
    {
      name:     "python repr with single quotes and escaped quotes and newlines",
      input:    `upstream Gradio error: {'message': 'Tool choice is none, but model called a tool', 'type': 'invalid_request_error', 'code': 'tool_use_failed', 'failed_generation': '{\"name\": \"repo_browser.print_tree\", \"arguments\": {\"path\": \"\", \"depth\": 2}\\n}', 'status_code': 400}`,
      expected: "{\"name\": \"repo_browser.print_tree\", \"arguments\": {\"path\": \"\", \"depth\": 2}\n}",
    },
    {
      name:     "python repr with standard single quotes",
      input:    `{'code': 'tool_use_failed', 'failed_generation': '{"name": "calc", "arguments": {"x": 5}}'}`,
      expected: `{"name": "calc", "arguments": {"x": 5}}`,
    },
    {
      name:     "raw object failed_generation",
      input:    `{"failed_generation": {"name": "calc", "arguments": {"x": 5}}}`,
      expected: `{"name": "calc", "arguments": {"x": 5}}`,
    },
  }

  for _, c := range cases {
    t.Run(c.name, func(t *testing.T) {
      got, ok := extractFailedGeneration(c.input)
      if !ok {
        t.Fatalf("extractFailedGeneration failed to find failed_generation in %q", c.input)
      }
      tcs, _, has := DetectToolCalls(got)
      if !has || len(tcs) == 0 {
        t.Fatalf("DetectToolCalls failed on extracted generation %q", got)
      }
    })
  }
}

func TestToolUseFailedRecovery(t *testing.T) {
  errPayload := `upstream Gradio error: {'message': 'Tool choice is none, but model called a tool', 'type': 'invalid_request_error', 'code': 'tool_use_failed', 'failed_generation': '{\"name\": \"repo_browser.print_tree\", \"arguments\": {\"path\": \"\", \"depth\": 2}\\n}', 'status_code': 400}`

  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/gradio_api/info" {
      resp := GradioAPIInfoResponse{
        NamedEndpoints: map[string]GradioEndpointInfo{
          "/chat_fn": {
            Parameters: []GradioParamInfo{
              {ParameterName: "message", Component: "Textbox"},
            },
          },
        },
      }
      json.NewEncoder(w).Encode(resp)
      return
    }
    if r.URL.Path == "/gradio_api/call/chat_fn" {
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(GradioJoinResponse{EventID: "evt_fail"})
      return
    }
    if r.URL.Path == "/gradio_api/call/chat_fn/evt_fail" {
      w.Header().Set("Content-Type", "text/event-stream")
      flusher := w.(http.Flusher)
      fmt.Fprintf(w, "event: error\ndata: %s\n\n", errPayload)
      flusher.Flush()
      return
    }
    http.NotFound(w, r)
  }))
  defer ts.Close()

  gw := NewGradioGateway(ts.URL, "", 10*time.Second)

  // 1. Non-streaming test
  recNonStream := httptest.NewRecorder()
  httpReqNonStream := httptest.NewRequest("POST", "/v1/chat/completions", nil)
  err := gw.ExecuteChatCompletion(recNonStream, httpReqNonStream, ChatCompletionRequest{
    Model:    "test-model",
    Messages: []ChatMessage{{Role: "user", Content: "List files"}},
    Stream:   false,
  })
  if err != nil {
    t.Fatalf("expected non-streaming to succeed by recovering tool call, got error: %v", err)
  }
  var nonStreamResp ChatCompletionResponse
  if err := json.NewDecoder(recNonStream.Body).Decode(&nonStreamResp); err != nil {
    t.Fatalf("failed to decode response: %v", err)
  }
  if nonStreamResp.Choices[0].FinishReason != "tool_calls" {
    t.Fatalf("expected finish_reason 'tool_calls', got %q", nonStreamResp.Choices[0].FinishReason)
  }
  if len(nonStreamResp.Choices[0].Message.ToolCalls) != 1 {
    t.Fatalf("expected 1 tool call, got %d", len(nonStreamResp.Choices[0].Message.ToolCalls))
  }
  tc := nonStreamResp.Choices[0].Message.ToolCalls[0]
  if tc.Function.Name != "repo_browser.print_tree" {
    t.Errorf("expected function repo_browser.print_tree, got %q", tc.Function.Name)
  }

  // 2. Streaming test
  recStream := httptest.NewRecorder()
  httpReqStream := httptest.NewRequest("POST", "/v1/chat/completions", nil)
  err = gw.ExecuteChatCompletion(recStream, httpReqStream, ChatCompletionRequest{
    Model:    "test-model",
    Messages: []ChatMessage{{Role: "user", Content: "List files"}},
    Stream:   true,
  })
  if err != nil {
    t.Fatalf("expected streaming to succeed by recovering tool call, got error: %v", err)
  }
  bodyStream := recStream.Body.String()
  if !strings.Contains(bodyStream, "repo_browser.print_tree") {
    t.Errorf("expected streaming body to contain repo_browser.print_tree, got: %s", bodyStream)
  }
  if !strings.Contains(bodyStream, `"finish_reason":"tool_calls"`) {
    t.Errorf("expected streaming body to contain finish_reason tool_calls, got: %s", bodyStream)
  }
  if !strings.Contains(bodyStream, "[DONE]") {
    t.Errorf("expected streaming body to contain [DONE], got: %s", bodyStream)
  }
}

func TestToolUseFailedImmediateCallRecovery(t *testing.T) {
  errPayload := `{"error": {"code": 400, "message": "Tool choice is none, but model called a tool", "failed_generation": "{\"name\": \"calculator\", \"arguments\": {\"expr\": \"2+2\"}}"}}`

  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/gradio_api/info" {
      resp := GradioAPIInfoResponse{
        NamedEndpoints: map[string]GradioEndpointInfo{
          "/chat_fn": {
            Parameters: []GradioParamInfo{
              {ParameterName: "message", Component: "Textbox"},
            },
          },
        },
      }
      json.NewEncoder(w).Encode(resp)
      return
    }
    if strings.HasPrefix(r.URL.Path, "/gradio_api/call/chat_fn") || strings.HasPrefix(r.URL.Path, "/call/chat_fn") {
      w.Header().Set("Content-Type", "application/json")
      w.WriteHeader(http.StatusBadRequest)
      w.Write([]byte(errPayload))
      return
    }
    http.NotFound(w, r)
  }))
  defer ts.Close()

  gw := NewGradioGateway(ts.URL, "", 10*time.Second)

  rec := httptest.NewRecorder()
  httpReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
  err := gw.ExecuteChatCompletion(rec, httpReq, ChatCompletionRequest{
    Model:    "test-model",
    Messages: []ChatMessage{{Role: "user", Content: "Calculate"}},
    Stream:   false,
  })
  if err != nil {
    t.Fatalf("expected call-level error recovery to succeed, got: %v", err)
  }
  var resp ChatCompletionResponse
  if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
    t.Fatalf("failed to decode response: %v", err)
  }
  if resp.Choices[0].FinishReason != "tool_calls" {
    t.Fatalf("expected finish_reason 'tool_calls', got %q", resp.Choices[0].FinishReason)
  }
  if len(resp.Choices[0].Message.ToolCalls) != 1 {
    t.Fatalf("expected 1 tool call, got %d", len(resp.Choices[0].Message.ToolCalls))
  }
  if resp.Choices[0].Message.ToolCalls[0].Function.Name != "calculator" {
    t.Errorf("expected calculator function, got %q", resp.Choices[0].Message.ToolCalls[0].Function.Name)
  }
}

func TestSystemPromptAugmentationWithoutNativeToolCalling(t *testing.T) {
  gw := NewGradioGateway("https://generic-space.hf.space", "", 10*time.Second)

  tools := []Tool{
    {
      Type: "function",
      Function: map[string]interface{}{
        "name":        "search_docs",
        "description": "Search local documentation",
        "parameters": map[string]interface{}{
          "type": "object",
          "properties": map[string]interface{}{
            "query": map[string]interface{}{"type": "string"},
          },
          "required": []string{"query"},
        },
      },
    },
  }

  // Case 1: Space with SystemIndex and space default prompt, client sends no system message
  discWithSys := NewDefaultSpaceDiscovery("https://generic-space.hf.space")
  discWithSys.TotalInputs = 2
  discWithSys.MessageIndex = 0
  discWithSys.SystemIndex = 1
  discWithSys.DefaultSystemPrompt = "You are a specialized documentation bot."

  req1 := ChatCompletionRequest{
    Tools: tools,
    Messages: []ChatMessage{
      {Role: "user", Content: "How to configure SSL?"},
    },
  }

  payload1, err := gw.BuildGradioPayload(discWithSys, req1)
  if err != nil {
    t.Fatalf("BuildGradioPayload failed: %v", err)
  }
  sysStr, ok := payload1[1].(string)
  if !ok {
    t.Fatalf("expected string at SystemIndex 1, got %T", payload1[1])
  }
  if !strings.Contains(sysStr, "You are a specialized documentation bot.") {
    t.Errorf("expected default system prompt to be retained, got: %s", sysStr)
  }
  if !strings.Contains(sysStr, "API router") || !strings.Contains(sysStr, "search_docs") {
    t.Errorf("expected tool calling instructions and function name in system prompt, got: %s", sysStr)
  }

  // Case 2: Space with SystemIndex, client sends their own system message
  req2 := ChatCompletionRequest{
    Tools: tools,
    Messages: []ChatMessage{
      {Role: "system", Content: "You are an expert developer assistant."},
      {Role: "user", Content: "How to configure SSL?"},
    },
  }
  payload2, err := gw.BuildGradioPayload(discWithSys, req2)
  if err != nil {
    t.Fatalf("BuildGradioPayload failed: %v", err)
  }
  sysStr2, ok := payload2[1].(string)
  if !ok {
    t.Fatalf("expected string at SystemIndex 1, got %T", payload2[1])
  }
  if !strings.Contains(sysStr2, "You are an expert developer assistant.") {
    t.Errorf("expected client system message, got: %s", sysStr2)
  }
  if !strings.Contains(sysStr2, "search_docs") {
    t.Errorf("expected search_docs tool instruction, got: %s", sysStr2)
  }

  // Case 3: Space without SystemIndex, single input textbox
  discSingle := NewDefaultSpaceDiscovery("https://single-input.hf.space")
  discSingle.TotalInputs = 1
  discSingle.MessageIndex = 0
  discSingle.SystemIndex = -1
  discSingle.HistoryIndex = -1

  payload3, err := gw.BuildGradioPayload(discSingle, req1)
  if err != nil {
    t.Fatalf("BuildGradioPayload failed: %v", err)
  }
  singleMsg, ok := payload3[0].(string)
  if !ok {
    t.Fatalf("expected string at index 0, got %T", payload3[0])
  }
  if !strings.Contains(singleMsg, "API router") || !strings.Contains(singleMsg, "How to configure SSL?") {
    t.Errorf("expected single message to contain augmented instructions and user query, got: %s", singleMsg)
  }
}

func TestToolChoiceHandling(t *testing.T) {
  tools := []Tool{
    {
      Type: "function",
      Function: map[string]interface{}{
        "name":        "calculator",
        "description": "Evaluate math expression",
      },
    },
  }

  // 1. ToolChoice: "none" -> no tool instruction generated
  instrNone := BuildToolInstruction(tools, "none")
  if instrNone != "" {
    t.Errorf("expected empty instruction for tool_choice: none, got: %s", instrNone)
  }

  // 2. ToolChoice: "required" -> mandatory directive
  instrReq := BuildToolInstruction(tools, "required")
  if !strings.Contains(instrReq, "You MUST call one of the available tools") {
    t.Errorf("expected required directive in instruction, got: %s", instrReq)
  }

  // 3. ToolChoice: specific function
  instrFn := BuildToolInstruction(tools, map[string]interface{}{
    "type": "function",
    "function": map[string]interface{}{
      "name": "calculator",
    },
  })
  if !strings.Contains(instrFn, "You MUST call the calculator tool") {
    t.Errorf("expected specific function directive in instruction, got: %s", instrFn)
  }
}

func TestConfiguredModelName(t *testing.T) {
  disc := &SpaceDiscovery{
    PrimaryModel: "hy3",
    Models:       []string{"hy3", "hunyuan3"},
  }

  // Without override
  ConfiguredModelName = ""
  list1 := disc.GetModelList()
  if len(list1) < 2 || list1[0].ID != "hy3" {
    t.Errorf("expected default first model hy3, got %+v", list1)
  }

  // With override
  ConfiguredModelName = "my-custom-hy3"
  defer func() { ConfiguredModelName = "" }()

  list2 := disc.GetModelList()
  if len(list2) < 3 || list2[0].ID != "my-custom-hy3" {
    t.Errorf("expected first model to be my-custom-hy3, got %+v", list2)
  }
}

func TestMultimodalStateSpaceMockServerCompletion(t *testing.T) {
  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/config" {
      cfg := GradioConfigResponse{
        Version: "5.29.0",
        Dependencies: []GradioDependency{
          {
            ID:      6,
            APIName: "chat",
            Inputs:  []int{12, 16, 20, 21, 22, 23, 24, 25, 26},
            Outputs: []int{14, 16},
          },
        },
        Components: []GradioComponent{
          {ID: 12, Type: "multimodaltextbox", Props: map[string]interface{}{"label": "Message"}},
          {ID: 16, Type: "state"},
          {ID: 20, Type: "radio", Props: map[string]interface{}{"label": "Model Type", "value": "Chat"}},
          {ID: 21, Type: "checkbox", Props: map[string]interface{}{"label": "Use Internet", "value": false}},
          {ID: 22, Type: "slider", Props: map[string]interface{}{"label": "Max Tokens", "value": 32768}},
          {ID: 23, Type: "slider", Props: map[string]interface{}{"label": "Temperature", "value": 0.8}},
          {ID: 24, Type: "slider", Props: map[string]interface{}{"label": "Top P", "value": 0.95}},
          {ID: 25, Type: "checkbox", Props: map[string]interface{}{"label": "stream", "value": true}},
          {ID: 26, Type: "textbox", Props: map[string]interface{}{"label": "user", "value": "null"}},
        },
      }
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(cfg)
      return
    }

    if r.URL.Path == "/gradio_api/info" {
      info := GradioAPIInfoResponse{
        NamedEndpoints: map[string]GradioEndpointInfo{
          "/chat": {
            Parameters: []GradioParamInfo{
              {ParameterName: "param_0", Label: "Message", Component: "Multimodaltextbox"},
              {ParameterName: "param_2", Label: "Model Type", Component: "Radio", ParameterDefault: "Chat"},
              {ParameterName: "param_3", Label: "Use Internet", Component: "Checkbox", ParameterDefault: false},
              {ParameterName: "param_4", Label: "Max Tokens", Component: "Slider", ParameterDefault: 32768},
              {ParameterName: "param_5", Label: "Temperature", Component: "Slider", ParameterDefault: 0.8},
              {ParameterName: "param_6", Label: "Top P", Component: "Slider", ParameterDefault: 0.95},
              {ParameterName: "param_7", Label: "stream", Component: "Checkbox", ParameterDefault: true},
              {ParameterName: "param_8", Label: "user", Component: "Textbox", ParameterDefault: "null"},
            },
          },
        },
      }
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(info)
      return
    }

    if r.URL.Path == "/gradio_api/call/chat" {
      var body struct {
        Data []interface{} `json:"data"`
      }
      if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
        http.Error(w, err.Error(), http.StatusBadRequest)
        return
      }
      if len(body.Data) != 9 {
        http.Error(w, fmt.Sprintf("expected 9 inputs, got %d", len(body.Data)), http.StatusBadRequest)
        return
      }
      // Check multimodal dict at index 0
      mmDict, ok := body.Data[0].(map[string]interface{})
      if !ok {
        http.Error(w, "input 0 must be multimodal dict", http.StatusBadRequest)
        return
      }
      txt, _ := mmDict["text"].(string)
      evt := "evt_normal"
      if strings.Contains(txt, "get_weather") {
        evt = "evt_tool"
      }
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(GradioJoinResponse{EventID: evt})
      return
    }

    if r.URL.Path == "/gradio_api/call/chat/evt_normal" {
      w.Header().Set("Content-Type", "text/event-stream")
      flusher, _ := w.(http.Flusher)
      fmt.Fprintf(w, "event: complete\ndata: [\"Hello from multimodal space!\", null]\n\n")
      flusher.Flush()
      return
    }

    if r.URL.Path == "/gradio_api/call/chat/evt_tool" {
      w.Header().Set("Content-Type", "text/event-stream")
      flusher, _ := w.(http.Flusher)
      fmt.Fprintf(w, "event: complete\ndata: [\"```json\\n[{\\\"name\\\": \\\"get_weather\\\", \\\"arguments\\\": {\\\"location\\\": \\\"Tokyo\\\"}}]\\n```\", null]\n\n")
      flusher.Flush()
      return
    }

    http.NotFound(w, r)
  }))
  defer ts.Close()

  gw := NewGradioGateway(ts.URL, "", 10*time.Second)

  // 1. Verify Space Inspection
  disc, err := InspectSpace(gw.client, ts.URL, DefaultUserAgent)
  if err != nil {
    t.Fatalf("InspectSpace failed: %v", err)
  }
  if disc.TotalInputs != 9 {
    t.Errorf("expected TotalInputs 9, got %d", disc.TotalInputs)
  }
  if disc.MessageIndex != 0 {
    t.Errorf("expected MessageIndex 0, got %d", disc.MessageIndex)
  }
  if !disc.MessageIsMultimodal {
    t.Errorf("expected MessageIsMultimodal true")
  }
  if disc.TempIndex != 5 {
    t.Errorf("expected TempIndex 5, got %d", disc.TempIndex)
  }
  if disc.MaxTokensIndex != 4 {
    t.Errorf("expected MaxTokensIndex 4, got %d", disc.MaxTokensIndex)
  }
  if disc.TopPIndex != 6 {
    t.Errorf("expected TopPIndex 6, got %d", disc.TopPIndex)
  }
  if disc.StreamIndex != 7 {
    t.Errorf("expected StreamIndex 7, got %d", disc.StreamIndex)
  }

  // 2. Normal Chat Completion
  req1 := ChatCompletionRequest{
    Model: "askcyph",
    Messages: []ChatMessage{
      {Role: "user", Content: "Hello world"},
    },
    Stream: false,
  }
  b1, _ := json.Marshal(req1)
  httpReq1 := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b1))
  rec1 := httptest.NewRecorder()

  err = gw.ExecuteChatCompletion(rec1, httpReq1, req1)
  if err != nil {
    t.Fatalf("ExecuteChatCompletion failed: %v", err)
  }
  var resp1 ChatCompletionResponse
  if err := json.NewDecoder(rec1.Body).Decode(&resp1); err != nil {
    t.Fatalf("failed to decode response: %v", err)
  }
  if resp1.Choices[0].Message.Content != "Hello from multimodal space!" {
    t.Errorf("unexpected content: %v", resp1.Choices[0].Message.Content)
  }

  // 3. Tool Calling Chat Completion
  req2 := ChatCompletionRequest{
    Model: "askcyph",
    Messages: []ChatMessage{
      {Role: "user", Content: "Weather in Tokyo?"},
    },
    Tools: []Tool{
      {
        Type: "function",
        Function: map[string]interface{}{
          "name":        "get_weather",
          "description": "Get current weather",
          "parameters": map[string]interface{}{
            "type": "object",
            "properties": map[string]interface{}{
              "location": map[string]interface{}{"type": "string"},
            },
            "required": []string{"location"},
          },
        },
      },
    },
    Stream: false,
  }
  b2, _ := json.Marshal(req2)
  httpReq2 := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b2))
  rec2 := httptest.NewRecorder()

  err = gw.ExecuteChatCompletion(rec2, httpReq2, req2)
  if err != nil {
    t.Fatalf("ExecuteChatCompletion tool calling failed: %v", err)
  }
  var resp2 ChatCompletionResponse
  if err := json.NewDecoder(rec2.Body).Decode(&resp2); err != nil {
    t.Fatalf("failed to decode response: %v", err)
  }
  if resp2.Choices[0].FinishReason != "tool_calls" {
    t.Fatalf("expected finish_reason 'tool_calls', got %q", resp2.Choices[0].FinishReason)
  }
  if len(resp2.Choices[0].Message.ToolCalls) != 1 {
    t.Fatalf("expected 1 tool call, got %d", len(resp2.Choices[0].Message.ToolCalls))
  }
  if resp2.Choices[0].Message.ToolCalls[0].Function.Name != "get_weather" {
    t.Errorf("expected get_weather, got %s", resp2.Choices[0].Message.ToolCalls[0].Function.Name)
  }
}

// TestGradio3DirectPredictProtocol verifies discovery and chat completion against a Gradio 3 space
// where /gradio_api/info returns 404 and the protocol resolves to /run/predict with tuple pairs history.
func TestGradio3DirectPredictProtocol(t *testing.T) {
  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/gradio_api/info" || r.URL.Path == "/info" {
      http.NotFound(w, r)
      return
    }

    if r.URL.Path == "/config" {
      cfg := GradioConfigResponse{
        Version: "3.41.2",
        Mode:    "chat_interface",
        Title:   "Legacy Gradio 3 Chat",
        Components: []GradioComponent{
          {ID: 1, Type: "textbox", Props: map[string]interface{}{"label": "Input"}},
          {ID: 2, Type: "chatbot", Props: map[string]interface{}{"label": "Chatbot"}},
        },
        Dependencies: []GradioDependency{
          {
            ID:      0,
            Inputs:  []int{1, 2},
            Outputs: []int{2},
          },
        },
      }
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(cfg)
      return
    }

    if r.URL.Path == "/run/predict" {
      var body struct {
        Data        []interface{} `json:"data"`
        FnIndex     int           `json:"fn_index"`
        SessionHash string        `json:"session_hash"`
      }
      if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
        http.Error(w, err.Error(), http.StatusBadRequest)
        return
      }
      if body.FnIndex != 0 {
        http.Error(w, fmt.Sprintf("expected fn_index 0, got %d", body.FnIndex), http.StatusBadRequest)
        return
      }
      msg, _ := body.Data[0].(string)
      reply := "Echo from Gradio 3: " + msg
      respData := map[string]interface{}{
        "data": []interface{}{
          [][]string{
            {msg, reply},
          },
        },
        "is_generating": false,
      }
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(respData)
      return
    }

    http.NotFound(w, r)
  }))
  defer ts.Close()

  gw := NewGradioGateway(ts.URL, "", 10*time.Second)
  disc := gw.GetDiscovery(ts.URL, DefaultUserAgent)

  if disc.GradioVersion != "3.41.2" {
    t.Errorf("expected GradioVersion 3.41.2, got %s", disc.GradioVersion)
  }
  if disc.Protocol != "predict" {
    t.Errorf("expected Protocol predict, got %s", disc.Protocol)
  }
  if disc.Flavor != "ChatInterface" {
    t.Errorf("expected Flavor ChatInterface, got %s", disc.Flavor)
  }
  if disc.HistoryFormat != "pairs" {
    t.Errorf("expected HistoryFormat pairs, got %s", disc.HistoryFormat)
  }
  if disc.FnIndex != 0 {
    t.Errorf("expected FnIndex 0, got %d", disc.FnIndex)
  }

  // Non-streaming completion
  req := ChatCompletionRequest{
    Messages: []ChatMessage{
      {Role: "user", Content: "Hello Gradio 3!"},
    },
    Stream: false,
  }
  b, _ := json.Marshal(req)
  httpReq := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b))
  rec := httptest.NewRecorder()

  if err := gw.ExecuteChatCompletion(rec, httpReq, req); err != nil {
    t.Fatalf("ExecuteChatCompletion failed on Gradio 3: %v", err)
  }

  var res ChatCompletionResponse
  if err := json.NewDecoder(rec.Body).Decode(&res); err != nil {
    t.Fatalf("failed to decode response: %v", err)
  }
  if res.Choices[0].Message.Content != "Echo from Gradio 3: Hello Gradio 3!" {
    t.Errorf("unexpected completion content: %v", res.Choices[0].Message.Content)
  }
}

// TestEndpointScoringDisambiguation verifies that chat generation endpoints are selected over
// UI utility, reset, voting, and feedback endpoints.
func TestEndpointScoringDisambiguation(t *testing.T) {
  compMap := map[int]GradioComponent{
    1: {ID: 1, Type: "textbox", Props: map[string]interface{}{"label": "Message"}},
    2: {ID: 2, Type: "chatbot", Props: map[string]interface{}{"label": "Chat"}},
    3: {ID: 3, Type: "state", Props: map[string]interface{}{"label": "State"}},
  }

  chatDep := GradioDependency{
    ID:      0,
    Inputs:  []int{1, 2, 3},
    Outputs: []int{2},
    Types:   GradioDependencyTypes{Generator: true},
  }
  clearDep := GradioDependency{
    ID:      1,
    Inputs:  []int{2},
    Outputs: []int{2},
  }
  voteDep := GradioDependency{
    ID:      2,
    Inputs:  []int{2},
    Outputs: []int{},
  }

  chatScore := ScoreCandidateEndpoint("chat", nil, &chatDep, compMap)
  clearScore := ScoreCandidateEndpoint("clear", nil, &clearDep, compMap)
  voteScore := ScoreCandidateEndpoint("vote", nil, &voteDep, compMap)
  resetScore := ScoreCandidateEndpoint("reset_all", nil, &clearDep, compMap)

  if chatScore <= 0 {
    t.Errorf("expected positive chat score, got %d", chatScore)
  }
  if clearScore >= chatScore {
    t.Errorf("expected chat score > clear score, got chat=%d clear=%d", chatScore, clearScore)
  }
  if voteScore >= chatScore {
    t.Errorf("expected chat score > vote score, got chat=%d vote=%d", chatScore, voteScore)
  }
  if resetScore >= chatScore {
    t.Errorf("expected chat score > reset score, got chat=%d reset=%d", chatScore, resetScore)
  }
}

// TestToolCallingPictureResolution verifies that tool calling mode is accurately resolved
// based on component topology.
func TestToolCallingPictureResolution(t *testing.T) {
  // Case A: Native slot
  discNative := NewDefaultSpaceDiscovery("https://tencent-hy3.hf.space")
  discNative.FunctionsJSONIndex = 8
  discNative.SystemIndex = 2
  discNative.HistoryIndex = 1
  discNative.ToolCallMode = "native_slot"
  if discNative.ToolCallMode != "native_slot" {
    t.Errorf("expected native_slot, got %s", discNative.ToolCallMode)
  }

  // Case B: Dedicated system prompt slot
  discSys := NewDefaultSpaceDiscovery("https://custom-chat.hf.space")
  discSys.FunctionsJSONIndex = -1
  discSys.SystemIndex = 2
  discSys.HistoryIndex = 1
  discSys.ToolCallMode = "prompt_augmented_system"
  if discSys.ToolCallMode != "prompt_augmented_system" {
    t.Errorf("expected prompt_augmented_system, got %s", discSys.ToolCallMode)
  }

  // Case C: Conversation history (first turn)
  discFirst := NewDefaultSpaceDiscovery("https://chat-only.hf.space")
  discFirst.FunctionsJSONIndex = -1
  discFirst.SystemIndex = -1
  discFirst.HistoryIndex = 1
  discFirst.ToolCallMode = "prompt_augmented_first_turn"
  if discFirst.ToolCallMode != "prompt_augmented_first_turn" {
    t.Errorf("expected prompt_augmented_first_turn, got %s", discFirst.ToolCallMode)
  }

  // Case D: Single prompt input
  discSingle := NewDefaultSpaceDiscovery("https://single-prompt.hf.space")
  discSingle.FunctionsJSONIndex = -1
  discSingle.SystemIndex = -1
  discSingle.HistoryIndex = -1
  discSingle.ToolCallMode = "prompt_augmented_single_prompt"
  if discSingle.ToolCallMode != "prompt_augmented_single_prompt" {
    t.Errorf("expected prompt_augmented_single_prompt, got %s", discSingle.ToolCallMode)
  }
}

// TestCallToPredictProtocolFallback verifies that if a space reports /call support in /info
// but /call returns 404 at runtime, gr2gw gracefully falls back to /run/predict.
func TestCallToPredictProtocolFallback(t *testing.T) {
  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/gradio_api/info" {
      info := GradioAPIInfoResponse{
        NamedEndpoints: map[string]GradioEndpointInfo{
          "/chat": {
            Parameters: []GradioParamInfo{
              {ParameterName: "prompt", Label: "Prompt", Component: "Textbox"},
            },
          },
        },
      }
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(info)
      return
    }

    if r.URL.Path == "/config" {
      cfg := GradioConfigResponse{
        Version: "4.20.0",
        Mode:    "interface",
        Components: []GradioComponent{
          {ID: 1, Type: "textbox", Props: map[string]interface{}{"label": "Prompt"}},
        },
        Dependencies: []GradioDependency{
          {ID: 0, APIName: "/chat", Inputs: []int{1}, Outputs: []int{1}},
        },
      }
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(cfg)
      return
    }

    // /call/chat returns 404 (endpoint disabled or unsupported)
    if strings.HasPrefix(r.URL.Path, "/gradio_api/call/") || strings.HasPrefix(r.URL.Path, "/call/") {
      http.NotFound(w, r)
      return
    }

    // Fallback /run/predict works
    if r.URL.Path == "/run/predict" || r.URL.Path == "/gradio_api/run/predict" {
      var body struct {
        Data []interface{} `json:"data"`
      }
      json.NewDecoder(r.Body).Decode(&body)
      msg, _ := body.Data[0].(string)
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(map[string]interface{}{
        "data":          []interface{}{"Fallback response for: " + msg},
        "is_generating": false,
      })
      return
    }

    http.NotFound(w, r)
  }))
  defer ts.Close()

  gw := NewGradioGateway(ts.URL, "", 10*time.Second)

  req := ChatCompletionRequest{
    Messages: []ChatMessage{
      {Role: "user", Content: "Testing fallback"},
    },
    Stream: false,
  }
  b, _ := json.Marshal(req)
  httpReq := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b))
  rec := httptest.NewRecorder()

  if err := gw.ExecuteChatCompletion(rec, httpReq, req); err != nil {
    t.Fatalf("ExecuteChatCompletion fallback failed: %v", err)
  }

  var res ChatCompletionResponse
  if err := json.NewDecoder(rec.Body).Decode(&res); err != nil {
    t.Fatalf("failed to decode response: %v", err)
  }
  if res.Choices[0].Message.Content != "Fallback response for: Testing fallback" {
    t.Errorf("unexpected content: %v", res.Choices[0].Message.Content)
  }
}

// TestEndpointScoringUserHandlerDepPenalization verifies that user submission helper
// endpoints with shared input/output textboxes are penalized while streaming generators
// and bot endpoints like Chat_Message are favored.
func TestEndpointScoringUserHandlerDepPenalization(t *testing.T) {
  compMap := map[int]GradioComponent{
    5: {ID: 5, Type: "state", Props: map[string]interface{}{"label": "State"}},
    6: {ID: 6, Type: "chatbot", Props: map[string]interface{}{"label": "Chatbot"}},
    8: {ID: 8, Type: "textbox", Props: map[string]interface{}{"label": "Message"}},
    27: {ID: 27, Type: "chatbot", Props: map[string]interface{}{"label": "Chatbot 2"}},
    29: {ID: 29, Type: "textbox", Props: map[string]interface{}{"label": "Link"}},
    30: {ID: 30, Type: "textbox", Props: map[string]interface{}{"label": "User Message"}},
  }

  // user: inputs [8, 6], outputs [8, 6] (clears textbox 8)
  userDep := GradioDependency{
    ID:      2,
    Inputs:  []int{8, 6},
    Outputs: []int{8, 6},
    Types:   GradioDependencyTypes{Generator: false},
  }

  // user2: inputs [30, 27, 29], outputs [30, 27, 29] (clears textboxes 30 and 29)
  user2Dep := GradioDependency{
    ID:      23,
    Inputs:  []int{30, 27, 29},
    Outputs: []int{30, 27, 29},
    Types:   GradioDependencyTypes{Generator: false},
  }

  // Chat_Message: inputs [6, 5], outputs [6, 5], generator: true
  chatMsgDep := GradioDependency{
    ID:      3,
    Inputs:  []int{6, 5},
    Outputs: []int{6, 5},
    Types:   GradioDependencyTypes{Generator: true},
  }

  scoreUser := ScoreCandidateEndpoint("/user", nil, &userDep, compMap)
  scoreUser2 := ScoreCandidateEndpoint("/user2", nil, &user2Dep, compMap)
  scoreChatMsg := ScoreCandidateEndpoint("/Chat_Message", nil, &chatMsgDep, compMap)

  if scoreChatMsg <= 0 {
    t.Errorf("expected positive score for Chat_Message, got %d", scoreChatMsg)
  }
  if scoreChatMsg <= scoreUser {
    t.Errorf("expected Chat_Message score > user score, got chatMsg=%d user=%d", scoreChatMsg, scoreUser)
  }
  if scoreChatMsg <= scoreUser2 {
    t.Errorf("expected Chat_Message score > user2 score, got chatMsg=%d user2=%d", scoreChatMsg, scoreUser2)
  }
}

// TestHistoryOnlyEndpointPayloadBuilding verifies that endpoints with no separate
// message textbox (MessageIndex == -1, HistoryIndex == 0) correctly append the user prompt
// into the history array without overwriting other slots.
func TestHistoryOnlyEndpointPayloadBuilding(t *testing.T) {
  gw := NewGradioGateway("https://test-space.hf.space", "", 10*time.Second)

  // 1. Gradio 6 messages format
  discG6 := NewDefaultSpaceDiscovery("https://test-space.hf.space")
  discG6.TotalInputs = 2
  discG6.MessageIndex = -1
  discG6.HistoryIndex = 0
  discG6.HistoryFormat = "gradio_messages"

  req := ChatCompletionRequest{
    Messages: []ChatMessage{
      {Role: "user", Content: "Hello world"},
    },
  }

  data, err := gw.BuildGradioPayload(discG6, req)
  if err != nil {
    t.Fatalf("BuildGradioPayload failed: %v", err)
  }
  if len(data) != 2 {
    t.Fatalf("expected 2 inputs, got %d", len(data))
  }

  gMsgs, ok := data[0].([]map[string]interface{})
  if !ok {
    t.Fatalf("expected []map[string]interface{} for gradio_messages, got %T", data[0])
  }
  if len(gMsgs) != 1 {
    t.Fatalf("expected 1 message in history, got %d", len(gMsgs))
  }
  if gMsgs[0]["role"] != "user" {
    t.Errorf("expected user role, got %v", gMsgs[0]["role"])
  }
  contents, ok := gMsgs[0]["content"].([]map[string]string)
  if !ok || len(contents) != 1 || contents[0]["text"] != "Hello world" {
    t.Errorf("unexpected content structure: %v", gMsgs[0]["content"])
  }
  if data[1] != nil {
    t.Errorf("expected slot 1 (state) to remain nil, got %v", data[1])
  }

  // 2. Pairs format with MessageIndex == -1
  discPairs := NewDefaultSpaceDiscovery("https://test-space.hf.space")
  discPairs.TotalInputs = 1
  discPairs.MessageIndex = -1
  discPairs.HistoryIndex = 0
  discPairs.HistoryFormat = "pairs"

  dataPairs, err := gw.BuildGradioPayload(discPairs, req)
  if err != nil {
    t.Fatalf("BuildGradioPayload failed: %v", err)
  }
  pairs, ok := dataPairs[0].([][]string)
  if !ok || len(pairs) != 1 {
    t.Fatalf("expected 1 pair in history, got %T (%v)", dataPairs[0], dataPairs[0])
  }
  if pairs[0][0] != "Hello world" {
    t.Errorf("expected user prompt in pair[0], got %v", pairs[0][0])
  }
}

// TestInspectSpaceBlocksChatbotStateResolution verifies InspectSpace selects /Chat_Message
// when presented with a Blocks layout containing user handlers and generator functions.
func TestInspectSpaceBlocksChatbotStateResolution(t *testing.T) {
  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/gradio_api/info" {
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(map[string]interface{}{
        "named_endpoints": map[string]interface{}{
          "/user": map[string]interface{}{
            "parameters": []map[string]interface{}{
              {"parameter_name": "user_message", "component": "Textbox", "label": "msg"},
              {"parameter_name": "history", "component": "Chatbot", "label": "chat"},
            },
          },
          "/Chat_Message": map[string]interface{}{
            "parameters": []map[string]interface{}{
              {
                "parameter_name": "history",
                "component":      "Chatbot",
                "label":          "chat",
                "type":           map[string]interface{}{"title": "ChatbotDataMessages"},
                "python_type":    map[string]interface{}{"type": "dict(text: str)"},
              },
            },
          },
        },
      })
      return
    }

    if r.URL.Path == "/config" {
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(map[string]interface{}{
        "version": "6.20.0",
        "components": []map[string]interface{}{
          {"id": 5, "type": "state", "props": map[string]interface{}{"label": "State"}},
          {"id": 6, "type": "chatbot", "props": map[string]interface{}{"label": "Chatbot"}},
          {"id": 8, "type": "textbox", "props": map[string]interface{}{"label": "Message"}},
        },
        "dependencies": []map[string]interface{}{
          {
            "id":       2,
            "api_name": "user",
            "inputs":   []int{8, 6},
            "outputs":  []int{8, 6},
            "types":    map[string]interface{}{"generator": false},
          },
          {
            "id":       3,
            "api_name": "Chat_Message",
            "inputs":   []int{6, 5},
            "outputs":  []int{6, 5},
            "types":    map[string]interface{}{"generator": true},
          },
        },
      })
      return
    }

    http.NotFound(w, r)
  }))
  defer ts.Close()

  client := &http.Client{Timeout: 5 * time.Second}
  disc, err := InspectSpace(client, ts.URL, "test-agent")
  if err != nil {
    t.Fatalf("InspectSpace failed: %v", err)
  }

  if disc.Endpoint != "/Chat_Message" {
    t.Errorf("expected resolved endpoint /Chat_Message, got %s", disc.Endpoint)
  }
  if disc.FnIndex != 3 {
    t.Errorf("expected FnIndex 3, got %d", disc.FnIndex)
  }
  if disc.MessageIndex != -1 {
    t.Errorf("expected MessageIndex -1, got %d", disc.MessageIndex)
  }
  if disc.HistoryIndex != 0 {
    t.Errorf("expected HistoryIndex 0, got %d", disc.HistoryIndex)
  }
  if disc.HistoryFormat != "gradio_messages" {
    t.Errorf("expected HistoryFormat gradio_messages, got %s", disc.HistoryFormat)
  }
  // TotalInputs should be trimmed to 1 to match the exposed API parameters and avoid padding unexposed state
  if disc.TotalInputs != 1 {
    t.Errorf("expected TotalInputs 1 (unexposed trailing state trimmed), got %d", disc.TotalInputs)
  }
}

// TestExecuteCallCompletionSessionHashAndTrailingStateTrim verifies that call requests
// send a valid session_hash and trim unexposed trailing state inputs from the payload.
func TestExecuteCallCompletionSessionHashAndTrailingStateTrim(t *testing.T) {
  var receivedSessionHash string
  var receivedDataLen int

  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path == "/gradio_api/info" {
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(map[string]interface{}{
        "named_endpoints": map[string]interface{}{
          "/Chat_Message": map[string]interface{}{
            "parameters": []map[string]interface{}{
              {
                "parameter_name": "history",
                "component":      "Chatbot",
                "type":           map[string]interface{}{"title": "ChatbotDataMessages"},
                "python_type":    map[string]interface{}{"type": "dict(text: str)"},
              },
            },
          },
        },
      })
      return
    }

    if r.URL.Path == "/config" {
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(map[string]interface{}{
        "version": "6.20.0",
        "components": []map[string]interface{}{
          {"id": 5, "type": "state"},
          {"id": 6, "type": "chatbot"},
        },
        "dependencies": []map[string]interface{}{
          {
            "id":       3,
            "api_name": "Chat_Message",
            "inputs":   []int{6, 5},
            "outputs":  []int{6, 5},
            "types":    map[string]interface{}{"generator": true},
          },
        },
      })
      return
    }

    if r.URL.Path == "/gradio_api/call/Chat_Message" {
      var body struct {
        Data        []interface{} `json:"data"`
        SessionHash string        `json:"session_hash"`
      }
      if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
        http.Error(w, err.Error(), http.StatusBadRequest)
        return
      }
      receivedSessionHash = body.SessionHash
      receivedDataLen = len(body.Data)
      w.Header().Set("Content-Type", "application/json")
      json.NewEncoder(w).Encode(map[string]interface{}{
        "event_id": "evt-12345",
      })
      return
    }

    if r.URL.Path == "/gradio_api/call/Chat_Message/evt-12345" {
      w.Header().Set("Content-Type", "text/event-stream")
      w.Header().Set("Cache-Control", "no-cache")
      w.WriteHeader(http.StatusOK)
      flusher, _ := w.(http.Flusher)
      fmt.Fprintf(w, "event: generating\ndata: [[{\"role\":\"user\",\"content\":\"hello\"},{\"role\":\"assistant\",\"content\":\"hi\"}]]\n\n")
      if flusher != nil {
        flusher.Flush()
      }
      fmt.Fprintf(w, "event: complete\ndata: [[{\"role\":\"user\",\"content\":\"hello\"},{\"role\":\"assistant\",\"content\":\"hi\"}]]\n\n")
      if flusher != nil {
        flusher.Flush()
      }
      return
    }

    http.NotFound(w, r)
  }))
  defer ts.Close()

  gw := NewGradioGateway(ts.URL, "", 10*time.Second)
  req := ChatCompletionRequest{
    Model: "test-model",
    Messages: []ChatMessage{
      {Role: "user", Content: "hello"},
    },
  }
  httpReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
  rec := httptest.NewRecorder()

  err := gw.ExecuteChatCompletion(rec, httpReq, req)
  if err != nil {
    t.Fatalf("ExecuteChatCompletion failed: %v", err)
  }

  if rec.Code != http.StatusOK {
    t.Fatalf("expected status 200, got %d: %s", rec.Code, rec.Body.String())
  }
  if receivedSessionHash == "" {
    t.Errorf("expected non-empty session_hash in call payload")
  }
  if receivedDataLen != 1 {
    t.Errorf("expected call data length 1 (unexposed state trimmed), got %d", receivedDataLen)
  }
}