455 lines
16 KiB
Go
455 lines
16 KiB
Go
package openai
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
edgeservice "iop/apps/edge/internal/service"
|
|
"iop/packages/go/config"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
// fakeRunResultWithTimeout implements edgeservice.RunResult with a directly
|
|
// configurable WaitTimeout, so tests can exercise the run-timeout cancel path
|
|
// without waiting out the real DefaultTimeoutSec-derived duration.
|
|
type fakeRunResultWithTimeout struct {
|
|
dispatch edgeservice.RunDispatch
|
|
stream edgeservice.RunStream
|
|
waitTimeout time.Duration
|
|
}
|
|
|
|
func (f *fakeRunResultWithTimeout) Dispatch() edgeservice.RunDispatch { return f.dispatch }
|
|
func (f *fakeRunResultWithTimeout) Stream() edgeservice.RunStream { return f.stream }
|
|
func (f *fakeRunResultWithTimeout) Close() {}
|
|
func (f *fakeRunResultWithTimeout) WaitTimeout() time.Duration { return f.waitTimeout }
|
|
|
|
func TestCollectRunResultTimesOut(t *testing.T) {
|
|
handle := &edgeservice.RunHandle{
|
|
RunDispatch: edgeservice.RunDispatch{TimeoutSec: 1},
|
|
RunStream: edgeservice.RunStream{
|
|
Events: make(chan *iop.RunEvent),
|
|
NodeEvents: make(chan *iop.EdgeNodeEvent),
|
|
},
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), 1500*time.Millisecond)
|
|
defer cancel()
|
|
_, _, _, _, _, _, err := collectRunResult(ctx, handle.Stream(), handle.WaitTimeout())
|
|
if err == nil {
|
|
t.Fatal("expected timeout error")
|
|
}
|
|
}
|
|
func TestChatCompletionContextCancelSendsCancelRun(t *testing.T) {
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent)}
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
Adapter: "ollama",
|
|
Target: "llama-fixed",
|
|
SessionID: "cline",
|
|
TimeoutSec: 15,
|
|
}, fake, nil)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"client-model",
|
|
"messages":[{"role":"user","content":"hi"}]
|
|
}`)).WithContext(ctx)
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.handleChatCompletions(w, req)
|
|
|
|
if w.Code != http.StatusRequestTimeout {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
calls := fake.cancelCallsSnapshot()
|
|
if len(calls) != 1 {
|
|
t.Fatalf("expected 1 CancelRun call, got %d: %+v", len(calls), calls)
|
|
}
|
|
if calls[0].RunID != "run-test" || calls[0].Adapter != "ollama" || calls[0].Target != "llama-fixed" || calls[0].SessionID != "cline" {
|
|
t.Fatalf("unexpected cancel request: %+v", calls[0])
|
|
}
|
|
}
|
|
func TestResponsesContextCancelSendsCancelRun(t *testing.T) {
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent)}
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
Adapter: "ollama",
|
|
Target: "llama-fixed",
|
|
SessionID: "cline",
|
|
TimeoutSec: 15,
|
|
}, fake, nil)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
|
"model":"client-model",
|
|
"input":"say hello"
|
|
}`)).WithContext(ctx)
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.handleResponses(w, req)
|
|
|
|
if w.Code != http.StatusRequestTimeout {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
calls := fake.cancelCallsSnapshot()
|
|
if len(calls) != 1 {
|
|
t.Fatalf("expected 1 CancelRun call, got %d: %+v", len(calls), calls)
|
|
}
|
|
if calls[0].RunID != "run-test" || calls[0].Adapter != "ollama" || calls[0].Target != "llama-fixed" || calls[0].SessionID != "cline" {
|
|
t.Fatalf("unexpected cancel request: %+v", calls[0])
|
|
}
|
|
}
|
|
func TestStreamChatCompletionContextCancelSendsCancelRun(t *testing.T) {
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent)}
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
Adapter: "ollama",
|
|
Target: "llama-fixed",
|
|
SessionID: "cline",
|
|
TimeoutSec: 15,
|
|
}, fake, nil)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"client-model",
|
|
"stream":true,
|
|
"messages":[{"role":"user","content":"hi"}]
|
|
}`)).WithContext(ctx)
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.handleChatCompletions(w, req)
|
|
|
|
calls := fake.cancelCallsSnapshot()
|
|
if len(calls) != 1 {
|
|
t.Fatalf("expected 1 CancelRun call, got %d: %+v", len(calls), calls)
|
|
}
|
|
if calls[0].RunID != "run-test" || calls[0].Adapter != "ollama" || calls[0].Target != "llama-fixed" || calls[0].SessionID != "cline" {
|
|
t.Fatalf("unexpected cancel request: %+v", calls[0])
|
|
}
|
|
}
|
|
func TestStreamChatCompletionTimeoutSendsCancelRun(t *testing.T) {
|
|
fake := &fakeRunService{}
|
|
srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, fake, nil)
|
|
|
|
handle := &fakeRunResultWithTimeout{
|
|
dispatch: edgeservice.RunDispatch{
|
|
RunID: "run-timeout",
|
|
NodeID: "node-1",
|
|
Adapter: "ollama",
|
|
Target: "llama3",
|
|
SessionID: "cline",
|
|
},
|
|
stream: edgeservice.RunStream{
|
|
Events: make(chan *iop.RunEvent),
|
|
NodeEvents: make(chan *iop.EdgeNodeEvent),
|
|
},
|
|
waitTimeout: 50 * time.Millisecond,
|
|
}
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.streamChatCompletion(w, newTestChatDispatchContext(req, chatCompletionRequest{Model: "llama3"}), handle)
|
|
|
|
calls := fake.cancelCallsSnapshot()
|
|
if len(calls) != 1 {
|
|
t.Fatalf("expected 1 CancelRun call, got %d: %+v", len(calls), calls)
|
|
}
|
|
if calls[0].RunID != "run-timeout" || calls[0].NodeRef != "node-1" {
|
|
t.Fatalf("unexpected cancel request: %+v", calls[0])
|
|
}
|
|
if !strings.Contains(w.Body.String(), "run timed out") {
|
|
t.Fatalf("expected timeout SSE error, got %s", w.Body.String())
|
|
}
|
|
}
|
|
func TestChatCompletionCompleteDoesNotSendCancelRun(t *testing.T) {
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
|
|
fake.events <- &iop.RunEvent{Type: "delta", Delta: "hello"}
|
|
fake.events <- &iop.RunEvent{Type: "complete"}
|
|
|
|
srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama", Target: "llama-fixed"}, fake, nil)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"client-model",
|
|
"messages":[{"role":"user","content":"hi"}]
|
|
}`))
|
|
w := httptest.NewRecorder()
|
|
|
|
srv.handleChatCompletions(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if calls := fake.cancelCallsSnapshot(); len(calls) != 0 {
|
|
t.Fatalf("expected no CancelRun calls for a normal completion, got %d: %+v", len(calls), calls)
|
|
}
|
|
}
|
|
func TestModelsExposesCatalogRouteModels(t *testing.T) {
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
ModelRoutes: []config.OpenAIRouteEntry{
|
|
{Model: "model-a", Adapter: "ollama", Target: "llama3"},
|
|
{Model: "model-b", Adapter: "vllm", Target: "qwen"},
|
|
},
|
|
}, &fakeRunService{}, nil)
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
|
w := httptest.NewRecorder()
|
|
srv.handleModels(w, req)
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d", w.Code)
|
|
}
|
|
body := w.Body.String()
|
|
if !strings.Contains(body, `"id":"model-a"`) || !strings.Contains(body, `"id":"model-b"`) {
|
|
t.Fatalf("expected catalog model IDs, got %s", body)
|
|
}
|
|
}
|
|
func TestModelsSkipsRoutesWithEmptyTarget(t *testing.T) {
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
ModelRoutes: []config.OpenAIRouteEntry{
|
|
{Model: "model-with-target", Adapter: "ollama", Target: "llama3"},
|
|
{Model: "model-no-target", Adapter: "ollama", Target: ""},
|
|
},
|
|
}, &fakeRunService{}, nil)
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
|
w := httptest.NewRecorder()
|
|
srv.handleModels(w, req)
|
|
body := w.Body.String()
|
|
if !strings.Contains(body, `"id":"model-with-target"`) {
|
|
t.Fatalf("expected model-with-target in response, got %s", body)
|
|
}
|
|
if strings.Contains(body, "model-no-target") {
|
|
t.Fatalf("model-no-target should be skipped (no target), got %s", body)
|
|
}
|
|
}
|
|
func TestModelsRouteCatalogWinsOverModelsField(t *testing.T) {
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
Models: []string{"legacy-model"},
|
|
ModelRoutes: []config.OpenAIRouteEntry{
|
|
{Model: "catalog-model", Adapter: "ollama", Target: "llama3"},
|
|
},
|
|
}, &fakeRunService{}, nil)
|
|
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
|
w := httptest.NewRecorder()
|
|
srv.handleModels(w, req)
|
|
body := w.Body.String()
|
|
if !strings.Contains(body, `"id":"catalog-model"`) {
|
|
t.Fatalf("expected catalog model in response, got %s", body)
|
|
}
|
|
if strings.Contains(body, "legacy-model") {
|
|
t.Fatalf("legacy model should be hidden when catalog is set, got %s", body)
|
|
}
|
|
}
|
|
func TestChatCompletionsRouteCatalogDispatches(t *testing.T) {
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
|
|
fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
|
|
fake.events <- &iop.RunEvent{Type: "complete"}
|
|
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
ModelRoutes: []config.OpenAIRouteEntry{
|
|
{Model: "model-a", Adapter: "ollama", Target: "llama3"},
|
|
{Model: "model-b", Adapter: "vllm", Target: "qwen"},
|
|
},
|
|
}, fake, nil)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"model-a",
|
|
"messages":[{"role":"user","content":"hi"}]
|
|
}`))
|
|
w := httptest.NewRecorder()
|
|
srv.handleChatCompletions(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if fake.req.Adapter != "ollama" || fake.req.Target != "llama3" {
|
|
t.Fatalf("dispatch mismatch: adapter=%q target=%q", fake.req.Adapter, fake.req.Target)
|
|
}
|
|
}
|
|
func TestChatCompletionsRouteCatalogDispatchesModelB(t *testing.T) {
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
|
|
fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
|
|
fake.events <- &iop.RunEvent{Type: "complete"}
|
|
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
ModelRoutes: []config.OpenAIRouteEntry{
|
|
{Model: "model-a", Adapter: "ollama", Target: "llama3"},
|
|
{Model: "model-b", Adapter: "ollama", Target: "qwen", MaxQueue: 5, QueueTimeoutMS: 2000},
|
|
},
|
|
}, fake, nil)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"metadata":{},
|
|
"model":"model-b",
|
|
"messages":[{"role":"user","content":"hi"}]
|
|
}`))
|
|
w := httptest.NewRecorder()
|
|
srv.handleChatCompletions(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if fake.req.Adapter != "ollama" || fake.req.Target != "qwen" {
|
|
t.Fatalf("dispatch mismatch: adapter=%q target=%q", fake.req.Adapter, fake.req.Target)
|
|
}
|
|
if fake.req.ModelGroupKey != "model-b" {
|
|
t.Fatalf("model group key: got %q, want model-b", fake.req.ModelGroupKey)
|
|
}
|
|
if fake.req.MaxQueue != 5 || fake.req.QueueTimeoutMS != 2000 {
|
|
t.Fatalf("queue policy mismatch: MaxQueue=%d, QueueTimeoutMS=%d", fake.req.MaxQueue, fake.req.QueueTimeoutMS)
|
|
}
|
|
}
|
|
func TestChatCompletionsCatalogMissFallsToLegacyTarget(t *testing.T) {
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
|
|
fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
|
|
fake.events <- &iop.RunEvent{Type: "complete"}
|
|
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
Adapter: "ollama",
|
|
Target: "legacy-target",
|
|
ModelRoutes: []config.OpenAIRouteEntry{
|
|
{Model: "model-a", Adapter: "ollama", Target: "llama3"},
|
|
},
|
|
}, fake, nil)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"unknown-model",
|
|
"messages":[{"role":"user","content":"hi"}]
|
|
}`))
|
|
w := httptest.NewRecorder()
|
|
srv.handleChatCompletions(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if fake.req.Target != "legacy-target" {
|
|
t.Fatalf("expected legacy-target fallback, got %q", fake.req.Target)
|
|
}
|
|
}
|
|
func TestChatCompletionsEmptyModelReturns400(t *testing.T) {
|
|
srv := NewServer(config.EdgeOpenAIConf{Adapter: "ollama"}, &fakeRunService{}, nil)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"",
|
|
"messages":[{"role":"user","content":"hi"}]
|
|
}`))
|
|
w := httptest.NewRecorder()
|
|
srv.handleChatCompletions(w, req)
|
|
if w.Code != http.StatusBadRequest {
|
|
t.Fatalf("expected 400 for empty model, got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
}
|
|
func TestResponsesRouteCatalogDispatchesRoute(t *testing.T) {
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
|
|
fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
|
|
fake.events <- &iop.RunEvent{Type: "complete"}
|
|
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
ModelRoutes: []config.OpenAIRouteEntry{
|
|
{Model: "model-a", Adapter: "ollama", Target: "qwen"},
|
|
},
|
|
}, fake, nil)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
|
"model":"model-a",
|
|
"input":"say hello"
|
|
}`))
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if fake.req.Adapter != "ollama" || fake.req.Target != "qwen" {
|
|
t.Fatalf("dispatch mismatch: adapter=%q target=%q", fake.req.Adapter, fake.req.Target)
|
|
}
|
|
if fake.req.ModelGroupKey != "model-a" {
|
|
t.Fatalf("model group key: got %q, want model-a", fake.req.ModelGroupKey)
|
|
}
|
|
}
|
|
func TestResponsesRouteCatalogDispatchesQueuePolicy(t *testing.T) {
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
|
|
fake.events <- &iop.RunEvent{Type: "delta", Delta: "ok"}
|
|
fake.events <- &iop.RunEvent{Type: "complete"}
|
|
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
ModelRoutes: []config.OpenAIRouteEntry{
|
|
{Model: "model-a", Adapter: "ollama", Target: "route-target", MaxQueue: 3, QueueTimeoutMS: 1500},
|
|
},
|
|
}, fake, nil)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
|
"model":"model-a",
|
|
"input":"test"
|
|
}`))
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if fake.req.Target != "route-target" {
|
|
t.Fatalf("expected route-target, got %q", fake.req.Target)
|
|
}
|
|
if fake.req.ModelGroupKey != "model-a" {
|
|
t.Fatalf("model group key: got %q, want model-a", fake.req.ModelGroupKey)
|
|
}
|
|
if fake.req.MaxQueue != 3 || fake.req.QueueTimeoutMS != 1500 {
|
|
t.Fatalf("queue policy mismatch: MaxQueue=%d, QueueTimeoutMS=%d", fake.req.MaxQueue, fake.req.QueueTimeoutMS)
|
|
}
|
|
}
|
|
|
|
func TestResolveRouteDispatchPreservesWorkspaceRequired(t *testing.T) {
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
ModelRoutes: []config.OpenAIRouteEntry{
|
|
{Model: "codex", Adapter: "cli", Target: "codex", WorkspaceRequired: true},
|
|
{Model: "llama3", Adapter: "ollama", Target: "llama3:8b"},
|
|
},
|
|
}, &fakeRunService{}, nil)
|
|
|
|
dispatch, ok := srv.resolveRouteDispatch("codex")
|
|
if !ok {
|
|
t.Fatal("expected dispatch to succeed for codex route")
|
|
}
|
|
if !dispatch.WorkspaceRequired {
|
|
t.Fatalf("expected workspace_required=true for codex route, got false")
|
|
}
|
|
|
|
dispatch, ok = srv.resolveRouteDispatch("llama3")
|
|
if !ok {
|
|
t.Fatal("expected dispatch to succeed for llama3 route")
|
|
}
|
|
if dispatch.WorkspaceRequired {
|
|
t.Fatalf("expected workspace_required=false for llama3 route, got true")
|
|
}
|
|
}
|
|
func TestResolveRouteDispatchFallbackWorkspaceRequiredFalse(t *testing.T) {
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
Adapter: "ollama",
|
|
Target: "llama3",
|
|
ModelRoutes: []config.OpenAIRouteEntry{
|
|
{Model: "codex", Adapter: "cli", Target: "codex", WorkspaceRequired: true},
|
|
},
|
|
}, &fakeRunService{}, nil)
|
|
|
|
dispatch, ok := srv.resolveRouteDispatch("unknown-model")
|
|
if !ok {
|
|
t.Fatal("expected fallback dispatch to succeed")
|
|
}
|
|
if dispatch.WorkspaceRequired {
|
|
t.Fatalf("expected workspace_required=false for fallback route, got true")
|
|
}
|
|
}
|
|
func TestCollectRunResultFailsWhenEventStreamCloses(t *testing.T) {
|
|
events := make(chan *iop.RunEvent)
|
|
close(events)
|
|
handle := &edgeservice.RunHandle{
|
|
RunDispatch: edgeservice.RunDispatch{TimeoutSec: 60},
|
|
RunStream: edgeservice.RunStream{
|
|
Events: events,
|
|
NodeEvents: make(chan *iop.EdgeNodeEvent),
|
|
},
|
|
}
|
|
|
|
_, _, _, _, _, _, err := collectRunResult(context.Background(), handle.Stream(), handle.WaitTimeout())
|
|
if err == nil {
|
|
t.Fatal("expected closed stream error")
|
|
}
|
|
if !strings.Contains(err.Error(), "run stream closed") {
|
|
t.Fatalf("expected run stream closed error, got %v", err)
|
|
}
|
|
}
|