package vllm import ( "bytes" "context" "encoding/json" "errors" "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 TestVllmExecuteParsesReasoningAndCachedInputTokens(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "text/event-stream") _, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":9,"completion_tokens":6,"prompt_tokens_details":{"cached_tokens":2},"completion_tokens_details":{"reasoning_tokens":5}}}`) _, _ = fmt.Fprintf(w, "data: [DONE]\n\n") })) defer server.Close() adapter := New(config.VllmConf{Endpoint: server.URL}, zap.NewNop()) sink := &fakeSink{} if err := adapter.Execute(context.Background(), runtime.ExecutionSpec{ RunID: "run-1", Target: "llama-3", Input: map[string]any{"prompt": "hi"}, }, sink); err != nil { t.Fatalf("Execute failed: %v", err) } events := sink.all() complete := events[len(events)-1] if complete.Type != runtime.EventTypeComplete || complete.Usage == nil { t.Fatalf("expected complete event with usage, got %+v", complete) } u := complete.Usage if u.InputTokens != 9 || u.OutputTokens != 6 || u.CachedInputTokens != 2 || u.ReasoningTokens != 5 { t.Fatalf("usage: got in=%d out=%d cached=%d reasoning=%d, want 9/6/2/5", u.InputTokens, u.OutputTokens, u.CachedInputTokens, u.ReasoningTokens) } } 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":"\n\n[\"git status\"]\n\n"}}]}`) _, _ = 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), "") || !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...) } // failingTunnelSink emulates the Edge socket write failing after failAfter // successful frames, as when the Edge relay aborts an abandoned tunnel. type failingTunnelSink struct { mu sync.Mutex frames []runtime.ProviderTunnelFrame failAfter int } func (s *failingTunnelSink) EmitTunnelFrame(_ context.Context, frame runtime.ProviderTunnelFrame) error { s.mu.Lock() defer s.mu.Unlock() if len(s.frames) >= s.failAfter { return errors.New("edge socket write failed") } s.frames = append(s.frames, frame) return nil } func (s *failingTunnelSink) all() []runtime.ProviderTunnelFrame { s.mu.Lock() defer s.mu.Unlock() return append([]runtime.ProviderTunnelFrame(nil), s.frames...) } // assertOrderedTunnelFrames verifies S09 frame semantics: contiguous sequences // starting at 0 (no drop/reorder) and exactly one terminal END/ERROR frame, // which must be the last frame emitted. func assertOrderedTunnelFrames(t *testing.T, frames []runtime.ProviderTunnelFrame) { t.Helper() if len(frames) == 0 { t.Fatal("no tunnel frames emitted") } terminals := 0 for i, f := range frames { if f.Sequence != int64(i) { t.Errorf("frame %d: sequence %d, want %d (ordered, lossless)", i, f.Sequence, i) } switch f.Kind { case runtime.ProviderTunnelFrameKindEnd, runtime.ProviderTunnelFrameKindError: terminals++ if i != len(frames)-1 { t.Errorf("terminal %s frame at index %d must be the last of %d frames", f.Kind, i, len(frames)) } } } if terminals != 1 { t.Errorf("expected exactly one terminal END/ERROR frame, got %d", terminals) } } 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) } assertOrderedTunnelFrames(t, frames) } 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) } assertOrderedTunnelFrames(t, frames) if frames[0].Kind != runtime.ProviderTunnelFrameKindResponseStart { t.Errorf("expected first frame RESPONSE_START, got %s", frames[0].Kind) } for i, f := range frames[1 : len(frames)-1] { if f.Kind != runtime.ProviderTunnelFrameKindBody { t.Errorf("mid frame %d: expected BODY before the terminal ERROR, got %s", i+1, f.Kind) } } lastFrame := frames[len(frames)-1] if lastFrame.Kind != runtime.ProviderTunnelFrameKindError { t.Errorf("expected last frame to be ERROR, got %s", lastFrame.Kind) } if lastFrame.Error == "" { t.Error("expected non-empty error message in frame") } if !strings.Contains(lastFrame.Error, "context canceled") { t.Errorf("expected cancellation error in terminal frame, got %q", lastFrame.Error) } } 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) } assertOrderedTunnelFrames(t, frames) lastFrame := frames[len(frames)-1] if lastFrame.Kind != runtime.ProviderTunnelFrameKindError { t.Errorf("expected last frame to be ERROR, got %s", lastFrame.Kind) } if !strings.Contains(lastFrame.Error, "read response body") && !strings.Contains(lastFrame.Error, "EOF") { t.Errorf("expected body read error msg, got %q", lastFrame.Error) } } // TestVllmTunnelProvider_SinkWriteFailureStopsUpstreamRelay covers the // Edge-write-failure leg of S09: when EmitTunnelFrame fails mid-stream the // adapter stops relaying immediately, returns the emit error without appending // further frames, and the upstream provider request is torn down. func TestVllmTunnelProvider_SinkWriteFailureStopsUpstreamRelay(t *testing.T) { handlerObservedTeardown := 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(handlerObservedTeardown) return } if ok { flusher.Flush() } select { case <-r.Context().Done(): close(handlerObservedTeardown) return case <-time.After(10 * time.Millisecond): } } })) defer server.Close() adapter := New(config.VllmConf{ Endpoint: server.URL, }, zap.NewNop()) // RESPONSE_START and the first BODY frame succeed; the next emit fails as // if the Edge socket write broke after the caller went away. sink := &failingTunnelSink{failAfter: 2} 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(context.Background(), req, sink) }() var tunnelErr error select { case tunnelErr = <-errCh: case <-time.After(2 * time.Second): t.Fatal("timeout waiting for TunnelProvider to return after sink write failure") } if tunnelErr == nil || !strings.Contains(tunnelErr.Error(), "edge socket write failed") { t.Fatalf("expected the sink emit error to propagate, got %v", tunnelErr) } select { case <-handlerObservedTeardown: case <-time.After(2 * time.Second): t.Fatal("provider handler did not observe upstream request teardown") } frames := sink.all() if len(frames) != 2 { t.Fatalf("expected relay to stop at 2 frames after emit failure, got %d: %+v", len(frames), frames) } if frames[0].Kind != runtime.ProviderTunnelFrameKindResponseStart { t.Errorf("expected first frame RESPONSE_START, got %s", frames[0].Kind) } if frames[1].Kind != runtime.ProviderTunnelFrameKindBody { t.Errorf("expected second frame BODY, got %s", frames[1].Kind) } for i, f := range frames { if f.Sequence != int64(i) { t.Errorf("frame %d: sequence %d, want %d", i, f.Sequence, i) } } } 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) } assertOrderedTunnelFrames(t, frames) }