iop/apps/node/internal/adapters/vllm/vllm_test.go

package vllm

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

	"go.uber.org/zap"

	"iop/apps/node/internal/runtime"
	"iop/packages/go/config"
)

type fakeSink struct {
	mu     sync.Mutex
	events []runtime.RuntimeEvent
}

func (s *fakeSink) Emit(_ context.Context, event runtime.RuntimeEvent) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.events = append(s.events, event)
	return nil
}

func (s *fakeSink) all() []runtime.RuntimeEvent {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]runtime.RuntimeEvent(nil), s.events...)
}

func TestVllmCapabilitiesQueryModels(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/models" {
			t.Fatalf("unexpected path %s", r.URL.Path)
		}
		_, _ = w.Write([]byte(`{"object":"list","data":[{"id":"model-a"},{"id":"model-b"}]}`))
	}))
	defer server.Close()

	adapter := New(config.VllmConf{
		Endpoint:         server.URL,
		Capacity:         10,
		MaxQueue:         20,
		QueueTimeoutMS:   5000,
		RequestTimeoutMS: 10000,
	}, zap.NewNop())
	caps, err := adapter.Capabilities(context.Background())
	if err != nil {
		t.Fatalf("Capabilities failed: %v", err)
	}
	if got := strings.Join(caps.Targets, ","); got != "model-a,model-b" {
		t.Fatalf("targets: got %q", got)
	}
	if caps.MaxConcurrency != 10 {
		t.Fatalf("expected MaxConcurrency 10, got %d", caps.MaxConcurrency)
	}
	if caps.MaxQueue != 20 {
		t.Fatalf("expected MaxQueue 20, got %d", caps.MaxQueue)
	}
	if caps.QueueTimeoutMS != 5000 {
		t.Fatalf("expected QueueTimeoutMS 5000, got %d", caps.QueueTimeoutMS)
	}
	if caps.RequestTimeoutMS != 10000 {
		t.Fatalf("expected RequestTimeoutMS 10000, got %d", caps.RequestTimeoutMS)
	}

	// Verify default capacity fallback
	adapterDefault := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop())
	capsDefault, err := adapterDefault.Capabilities(context.Background())
	if err != nil {
		t.Fatalf("Capabilities failed: %v", err)
	}
	if capsDefault.MaxConcurrency != 8 {
		t.Fatalf("expected default MaxConcurrency 8, got %d", capsDefault.MaxConcurrency)
	}
}

func TestVllmExecuteStreamsDeltas(t *testing.T) {
	var gotModel string
	var gotMessages int
	var gotStream bool
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/chat/completions" {
			t.Fatalf("unexpected path %s", r.URL.Path)
		}
		var req struct {
			Model    string `json:"model"`
			Messages []any  `json:"messages"`
			Stream   bool   `json:"stream"`
		}
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			t.Fatalf("decode request: %v", err)
		}
		gotModel = req.Model
		gotMessages = len(req.Messages)
		gotStream = req.Stream
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"hello "}}]}`)
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"world"}}]}`)
		_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
	defer server.Close()

	adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop())
	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-1",
		Target: "llama-3",
		Input:  map[string]any{"prompt": "say hello"},
	}, sink)
	if err != nil {
		t.Fatalf("Execute failed: %v", err)
	}
	if gotModel != "llama-3" {
		t.Fatalf("model: got %q", gotModel)
	}
	if gotMessages != 1 {
		t.Fatalf("messages: got %d", gotMessages)
	}
	if !gotStream {
		t.Fatal("expected stream=true")
	}

	events := sink.all()
	if len(events) != 4 {
		t.Fatalf("expected 4 events (start+delta+delta+complete), got %d: %+v", len(events), events)
	}
	if events[0].Type != runtime.EventTypeStart {
		t.Fatalf("expected start event, got %+v", events[0])
	}
	if events[1].Delta+events[2].Delta != "hello world" {
		t.Fatalf("unexpected deltas: %q + %q", events[1].Delta, events[2].Delta)
	}
	if events[3].Type != runtime.EventTypeComplete {
		t.Fatalf("expected complete event, got %+v", events[3])
	}
}

func TestVllmExecutePassesToolsAndPreservesNativeToolCalls(t *testing.T) {
	var gotToolChoice any
	var gotTools []any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req map[string]any
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			t.Fatalf("decode request: %v", err)
		}
		gotToolChoice = req["tool_choice"]
		gotTools, _ = req["tools"].([]any)
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"run_commands","arguments":"{\"commands\":"}}]}}]}`)
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"[\"git status\"]}"}}]},"finish_reason":"tool_calls"}]}`)
		_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
	defer server.Close()

	adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop())
	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-tools",
		Target: "llama-3",
		Input: map[string]any{
			"prompt":      "status",
			"tools":       []any{map[string]any{"type": "function"}},
			"tool_choice": "auto",
		},
	}, sink)
	if err != nil {
		t.Fatalf("Execute failed: %v", err)
	}
	if gotToolChoice != "auto" || len(gotTools) != 1 {
		t.Fatalf("tools/tool_choice not passed: tools=%+v choice=%+v", gotTools, gotToolChoice)
	}
	events := sink.all()
	complete := events[len(events)-1]
	if complete.Metadata["finish_reason"] != "tool_calls" {
		t.Fatalf("finish_reason: %+v", complete.Metadata)
	}
	var toolCalls []map[string]any
	if err := json.Unmarshal([]byte(complete.Metadata[runtimeMetadataOpenAIToolCalls]), &toolCalls); err != nil {
		t.Fatalf("tool_calls metadata JSON: %v", err)
	}
	fn := toolCalls[0]["function"].(map[string]any)
	if fn["arguments"] != `{"commands":["git status"]}` {
		t.Fatalf("arguments: %+v", fn["arguments"])
	}
}

func TestVllmExecuteRetriesSingleToolWhenAutoUnsupported(t *testing.T) {
	attempts := 0
	var retryBody map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		attempts++
		if attempts == 1 {
			w.WriteHeader(http.StatusBadRequest)
			_, _ = w.Write([]byte(`{"error":{"message":"\"auto\" tool choice requires --enable-auto-tool-choice and --tool-call-parser to be set"}}`))
			return
		}
		if err := json.NewDecoder(r.Body).Decode(&retryBody); err != nil {
			t.Fatalf("decode retry: %v", err)
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"ok"}}]}`)
		_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
	defer server.Close()

	adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop())
	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-retry",
		Target: "llama-3",
		Input: map[string]any{
			"prompt": "status",
			"tools": []any{map[string]any{
				"type": "function",
				"function": map[string]any{
					"name": "run_commands",
				},
			}},
		},
	}, sink)
	if err != nil {
		t.Fatalf("Execute failed: %v", err)
	}
	if attempts != 2 {
		t.Fatalf("attempts: got %d, want 2", attempts)
	}
	choice, ok := retryBody["tool_choice"].(map[string]any)
	if !ok {
		t.Fatalf("retry tool_choice missing: %+v", retryBody)
	}
	fn := choice["function"].(map[string]any)
	if choice["type"] != "function" || fn["name"] != "run_commands" {
		t.Fatalf("retry tool_choice: %+v", choice)
	}
}

func TestVllmExecuteFallsBackToTextToolsWhenNativeToolsUnsupported(t *testing.T) {
	attempts := 0
	var fallbackBody map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		attempts++
		switch attempts {
		case 1:
			w.WriteHeader(http.StatusBadRequest)
			_, _ = w.Write([]byte(`{"error":{"message":"\"auto\" tool choice requires --enable-auto-tool-choice and --tool-call-parser to be set"}}`))
			return
		case 2:
			w.WriteHeader(http.StatusBadRequest)
			_, _ = w.Write([]byte(`{"error":{"message":"tool_choice=\"function=ChatCompletionNamedFunction(name='run_commands') type='function'\" requires --tool-call-parser to be set"}}`))
			return
		}
		if err := json.NewDecoder(r.Body).Decode(&fallbackBody); err != nil {
			t.Fatalf("decode fallback: %v", err)
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"<tool_call>\n<function=run_commands>\n<parameter=commands>[\"git status\"]</parameter>\n</function>\n</tool_call>"}}]}`)
		_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
	defer server.Close()

	adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop())
	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-text-fallback",
		Target: "llama-3",
		Input: map[string]any{
			"messages": []any{
				map[string]any{"role": "system", "content": "Existing Cline system prompt."},
				map[string]any{"role": "user", "content": "status"},
			},
			"tools": []any{map[string]any{
				"type": "function",
				"function": map[string]any{
					"name": "run_commands",
					"parameters": map[string]any{
						"type": "object",
						"properties": map[string]any{
							"commands": map[string]any{"type": "array", "items": map[string]any{"type": "string"}},
						},
					},
				},
			}},
		},
	}, sink)
	if err != nil {
		t.Fatalf("Execute failed: %v", err)
	}
	if attempts != 3 {
		t.Fatalf("attempts: got %d, want 3", attempts)
	}
	if _, ok := fallbackBody["tools"]; ok {
		t.Fatalf("fallback must omit tools: %+v", fallbackBody["tools"])
	}
	if _, ok := fallbackBody["tool_choice"]; ok {
		t.Fatalf("fallback must omit tool_choice: %+v", fallbackBody["tool_choice"])
	}
	messages := fallbackBody["messages"].([]any)
	first := messages[0].(map[string]any)
	if first["role"] != "system" || !strings.Contains(first["content"].(string), "<tool_call>") || !strings.Contains(first["content"].(string), "run_commands") {
		t.Fatalf("fallback system instruction missing tool format/name: %+v", first)
	}
	if !strings.Contains(first["content"].(string), "client workspace root") || !strings.Contains(first["content"].(string), "Do not prepend cd") {
		t.Fatalf("fallback system instruction missing workspace-root guidance: %+v", first)
	}
	if !strings.Contains(first["content"].(string), "Existing Cline system prompt.") {
		t.Fatalf("fallback system instruction did not preserve existing system content: %+v", first)
	}
	if len(messages) != 2 || messages[1].(map[string]any)["role"] != "user" {
		t.Fatalf("fallback should keep a single leading system message: %+v", messages)
	}
	events := sink.all()
	complete := events[len(events)-1]
	if complete.Metadata[runtimeMetadataOpenAITextToolFallback] != "true" {
		t.Fatalf("fallback metadata missing: %+v", complete.Metadata)
	}
}

func TestVllmExecuteUsesMessagesInput(t *testing.T) {
	var gotMessages []map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req struct {
			Messages []map[string]any `json:"messages"`
		}
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			t.Fatalf("decode request: %v", err)
		}
		gotMessages = req.Messages
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"ok"}}]}`)
		_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
	defer server.Close()

	adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop())
	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-2",
		Target: "llama-3",
		Input: map[string]any{
			"messages": []any{
				map[string]any{"role": "system", "content": "You are helpful."},
				map[string]any{"role": "user", "content": "hi"},
			},
		},
	}, sink)
	if err != nil {
		t.Fatalf("Execute failed: %v", err)
	}
	if len(gotMessages) != 2 {
		t.Fatalf("messages: got %d", len(gotMessages))
	}
	if gotMessages[0]["role"] != "system" || gotMessages[1]["role"] != "user" {
		t.Fatalf("unexpected messages: %+v", gotMessages)
	}
}

func TestVllmExecuteEmitsErrorForHTTPFailure(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		http.Error(w, "service unavailable", http.StatusServiceUnavailable)
	}))
	defer server.Close()

	adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop())
	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-err",
		Target: "llama-3",
		Input:  map[string]any{"prompt": "hi"},
	}, sink)
	if err == nil {
		t.Fatal("expected error")
	}
	events := sink.all()
	if len(events) < 2 || events[1].Type != runtime.EventTypeError {
		t.Fatalf("expected error event after start, got %+v", events)
	}
	if !strings.Contains(events[1].Error, "503") {
		t.Fatalf("expected status in error, got %q", events[1].Error)
	}
}

func TestVllmInstanceKey(t *testing.T) {
	adapter := New(config.VllmConf{Endpoint: "http://localhost:8000"}, zap.NewNop(), "vllm-gpu")
	caps, _ := adapter.Capabilities(context.Background())
	if caps.InstanceKey != "vllm-gpu" {
		t.Fatalf("InstanceKey: got %q, want vllm-gpu", caps.InstanceKey)
	}
	if caps.AdapterName != Name {
		t.Fatalf("AdapterName: got %q, want %q", caps.AdapterName, Name)
	}
}

func TestVllmProbeProviderAvailability(t *testing.T) {
	t.Run("200_ok_and_target_hit", func(t *testing.T) {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.URL.Path != "/v1/models" {
				t.Fatalf("unexpected path %s", r.URL.Path)
			}
			_, _ = w.Write([]byte(`{"object":"list","data":[{"id":"model-a"},{"id":"model-b"}]}`))
		}))
		defer server.Close()

		adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop())
		res, err := adapter.ProbeProvider(context.Background(), "model-a")
		if err != nil {
			t.Fatalf("ProbeProvider failed: %v", err)
		}
		if res.Status != runtime.ProviderStatusAvailable {
			t.Errorf("expected Status available, got %s", res.Status)
		}
		if len(res.Targets) != 2 || res.Targets[0] != "model-a" {
			t.Errorf("unexpected Targets: %+v", res.Targets)
		}
	})

	t.Run("200_ok_and_target_miss", func(t *testing.T) {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			_, _ = w.Write([]byte(`{"object":"list","data":[{"id":"model-a"}]}`))
		}))
		defer server.Close()

		adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop())
		res, err := adapter.ProbeProvider(context.Background(), "model-b")
		if err != nil {
			t.Fatalf("ProbeProvider failed: %v", err)
		}
		if res.Status != runtime.ProviderStatusUnavailable {
			t.Errorf("expected Status unavailable, got %s", res.Status)
		}
		if !strings.Contains(res.Detail, "not found") {
			t.Errorf("expected 'not found' in detail, got %s", res.Detail)
		}
	})

	t.Run("500_internal_error", func(t *testing.T) {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusInternalServerError)
		}))
		defer server.Close()

		adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop())
		res, err := adapter.ProbeProvider(context.Background(), "model-a")
		if err != nil {
			t.Fatalf("ProbeProvider failed: %v", err)
		}
		if res.Status != runtime.ProviderStatusUnavailable {
			t.Errorf("expected Status unavailable, got %s", res.Status)
		}
	})

	t.Run("empty_target", func(t *testing.T) {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			_, _ = w.Write([]byte(`{"object":"list","data":[{"id":"model-a"}]}`))
		}))
		defer server.Close()

		adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop())
		res, err := adapter.ProbeProvider(context.Background(), "")
		if err != nil {
			t.Fatalf("ProbeProvider failed: %v", err)
		}
		if res.Status != runtime.ProviderStatusAvailable {
			t.Errorf("expected Status available, got %s", res.Status)
		}
	})
}

type fakeTunnelSink struct {
	mu     sync.Mutex
	frames []runtime.ProviderTunnelFrame
}

func (s *fakeTunnelSink) EmitTunnelFrame(_ context.Context, frame runtime.ProviderTunnelFrame) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.frames = append(s.frames, frame)
	return nil
}

func (s *fakeTunnelSink) all() []runtime.ProviderTunnelFrame {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]runtime.ProviderTunnelFrame(nil), s.frames...)
}

func TestVllmTunnelProvider(t *testing.T) {
	expectedBody := "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\ndata: [DONE]\n\n"

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			t.Errorf("expected POST method, got %s", r.Method)
		}
		if r.URL.Path != "/v1/chat/completions" {
			t.Errorf("expected path /v1/chat/completions, got %s", r.URL.Path)
		}

		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("X-Custom-Header", "custom-value")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(expectedBody))
	}))
	defer server.Close()

	adapter := New(config.VllmConf{
		Endpoint: server.URL,
	}, zap.NewNop())

	sink := &fakeTunnelSink{}
	req := runtime.ProviderTunnelRequest{
		RunID:    "run-1",
		TunnelID: "tunnel-1",
		Method:   "POST",
		Path:     "/v1/chat/completions",
		Body:     []byte(`{"model":"qwen3.6:35b","prompt":"hi"}`),
	}

	err := adapter.TunnelProvider(context.Background(), req, sink)
	if err != nil {
		t.Fatalf("TunnelProvider failed: %v", err)
	}

	frames := sink.all()
	if len(frames) < 3 {
		t.Fatalf("expected at least 3 frames, got %d", len(frames))
	}

	startFrame := frames[0]
	if startFrame.Kind != runtime.ProviderTunnelFrameKindResponseStart {
		t.Errorf("expected RESPONSE_START, got %s", startFrame.Kind)
	}
	if startFrame.StatusCode != http.StatusOK {
		t.Errorf("expected 200 OK, got %d", startFrame.StatusCode)
	}
	if startFrame.Headers["X-Custom-Header"] != "custom-value" {
		t.Errorf("expected Custom Header, got %v", startFrame.Headers)
	}

	var bodyBuffer bytes.Buffer
	var lastSeq int64 = 0
	for _, f := range frames[1 : len(frames)-1] {
		if f.Kind != runtime.ProviderTunnelFrameKindBody {
			t.Errorf("expected BODY kind, got %s", f.Kind)
		}
		if f.Sequence != lastSeq+1 {
			t.Errorf("expected seq %d, got %d", lastSeq+1, f.Sequence)
		}
		bodyBuffer.Write(f.Body)
		lastSeq = f.Sequence
	}
	if bodyBuffer.String() != expectedBody {
		t.Errorf("body mismatch: got %q, want %q", bodyBuffer.String(), expectedBody)
	}

	endFrame := frames[len(frames)-1]
	if endFrame.Kind != runtime.ProviderTunnelFrameKindEnd {
		t.Errorf("expected END, got %s", endFrame.Kind)
	}
	if !endFrame.End {
		t.Errorf("expected End true, got %v", endFrame.End)
	}
	if endFrame.Sequence != lastSeq+1 {
		t.Errorf("expected seq %d, got %d", lastSeq+1, endFrame.Sequence)
	}
}

func TestVllmTunnelProvider_Cancel(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	handlerObservedCancel := make(chan struct{})

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)

		flusher, ok := w.(http.Flusher)
		for {
			_, err := w.Write([]byte("data: chunk\n\n"))
			if err != nil {
				close(handlerObservedCancel)
				return
			}
			if ok {
				flusher.Flush()
			}
			select {
			case <-r.Context().Done():
				close(handlerObservedCancel)
				return
			case <-time.After(10 * time.Millisecond):
			}
		}
	}))
	defer server.Close()

	adapter := New(config.VllmConf{
		Endpoint: server.URL,
	}, zap.NewNop())

	sink := &fakeTunnelSink{}
	req := runtime.ProviderTunnelRequest{
		RunID:    "run-1",
		TunnelID: "tunnel-1",
		Method:   "POST",
		Path:     "/v1/chat/completions",
		Body:     []byte(`{"prompt":"hi"}`),
	}

	errCh := make(chan error, 1)
	go func() {
		errCh <- adapter.TunnelProvider(ctx, req, sink)
	}()

	time.Sleep(100 * time.Millisecond)
	cancel()

	select {
	case err := <-errCh:
		if err == nil {
			t.Error("expected error from TunnelProvider on cancellation, got nil")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for TunnelProvider to return")
	}

	select {
	case <-handlerObservedCancel:
	case <-time.After(2 * time.Second):
		t.Fatal("handler did not observe context cancel")
	}

	frames := sink.all()
	if len(frames) < 2 {
		t.Fatalf("expected at least 2 frames, got %d: %+v", len(frames), frames)
	}

	lastFrame := frames[len(frames)-1]
	if lastFrame.Kind != runtime.ProviderTunnelFrameKindError {
		t.Errorf("expected last frame to be ERROR, got %s", lastFrame.Kind)
	}
	expectedSeq := int64(len(frames) - 1)
	if lastFrame.Sequence != expectedSeq {
		t.Errorf("expected error frame sequence %d, got %d", expectedSeq, lastFrame.Sequence)
	}
	if lastFrame.Error == "" {
		t.Error("expected non-empty error message in frame")
	}
}

func TestVllmTunnelProvider_BodyReadError(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hj, ok := w.(http.Hijacker)
		if !ok {
			t.Fatal("webserver doesn't support hijacking")
		}
		conn, _, err := hj.Hijack()
		if err != nil {
			t.Fatalf("hijack failed: %v", err)
		}
		defer conn.Close()

		_, _ = conn.Write([]byte("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: 100\r\n\r\n"))
		_, _ = conn.Write([]byte("data: chunk1\n\n"))
	}))
	defer server.Close()

	adapter := New(config.VllmConf{
		Endpoint: server.URL,
	}, zap.NewNop())

	sink := &fakeTunnelSink{}
	req := runtime.ProviderTunnelRequest{
		RunID:    "run-1",
		TunnelID: "tunnel-1",
		Method:   "POST",
		Path:     "/v1/chat/completions",
		Body:     []byte(`{"prompt":"hi"}`),
	}

	err := adapter.TunnelProvider(context.Background(), req, sink)
	if err == nil {
		t.Fatal("expected error from body read failure, got nil")
	}

	frames := sink.all()
	if len(frames) < 2 {
		t.Fatalf("expected at least 2 frames, got %d: %+v", len(frames), frames)
	}

	lastFrame := frames[len(frames)-1]
	if lastFrame.Kind != runtime.ProviderTunnelFrameKindError {
		t.Errorf("expected last frame to be ERROR, got %s", lastFrame.Kind)
	}
	expectedSeq := int64(len(frames) - 1)
	if lastFrame.Sequence != expectedSeq {
		t.Errorf("expected error frame sequence %d, got %d", expectedSeq, lastFrame.Sequence)
	}
	if !strings.Contains(lastFrame.Error, "read response body") && !strings.Contains(lastFrame.Error, "EOF") {
		t.Errorf("expected body read error msg, got %q", lastFrame.Error)
	}
}

func TestVllmTunnelProvider_GemmaContentToolCallRawBytes(t *testing.T) {
	expectedBody := "data: {\"choices\":[{\"delta\":{\"content\":\"Gemma response \"}}]}\n\n" +
		"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_gemma_1\",\"type\":\"function\",\"function\":{\"name\":\"run_commands\",\"arguments\":\"{\\\"commands\\\": \"}}]}}]}\n\n" +
		"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"[\\\"git status\\\"]}\"}}]}}]}\n\n" +
		"data: [DONE]\n\n"

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			t.Errorf("expected POST method, got %s", r.Method)
		}
		if r.URL.Path != "/v1/chat/completions" {
			t.Errorf("expected path /v1/chat/completions, got %s", r.URL.Path)
		}

		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(expectedBody))
	}))
	defer server.Close()

	adapter := New(config.VllmConf{
		Endpoint: server.URL,
	}, zap.NewNop())

	sink := &fakeTunnelSink{}
	req := runtime.ProviderTunnelRequest{
		RunID:    "run-gemma-1",
		TunnelID: "tunnel-gemma-1",
		Method:   "POST",
		Path:     "/v1/chat/completions",
		Body:     []byte(`{"model":"gemma-2b","prompt":"hi"}`),
	}

	err := adapter.TunnelProvider(context.Background(), req, sink)
	if err != nil {
		t.Fatalf("TunnelProvider failed: %v", err)
	}

	frames := sink.all()
	if len(frames) < 3 {
		t.Fatalf("expected at least 3 frames, got %d", len(frames))
	}

	startFrame := frames[0]
	if startFrame.Kind != runtime.ProviderTunnelFrameKindResponseStart {
		t.Errorf("expected RESPONSE_START, got %s", startFrame.Kind)
	}
	if startFrame.StatusCode != http.StatusOK {
		t.Errorf("expected 200 OK, got %d", startFrame.StatusCode)
	}

	var bodyBuffer bytes.Buffer
	var lastSeq int64 = 0
	for _, f := range frames[1 : len(frames)-1] {
		if f.Kind != runtime.ProviderTunnelFrameKindBody {
			t.Errorf("expected BODY kind, got %s", f.Kind)
		}
		if f.Sequence != lastSeq+1 {
			t.Errorf("expected seq %d, got %d", lastSeq+1, f.Sequence)
		}
		bodyBuffer.Write(f.Body)
		lastSeq = f.Sequence
	}
	if bodyBuffer.String() != expectedBody {
		t.Errorf("body mismatch:\ngot:  %q\nwant: %q", bodyBuffer.String(), expectedBody)
	}

	endFrame := frames[len(frames)-1]
	if endFrame.Kind != runtime.ProviderTunnelFrameKindEnd {
		t.Errorf("expected END, got %s", endFrame.Kind)
	}
	if !endFrame.End {
		t.Errorf("expected End true, got %v", endFrame.End)
	}
	if endFrame.Sequence != lastSeq+1 {
		t.Errorf("expected seq %d, got %d", lastSeq+1, endFrame.Sequence)
	}
}