package chat import ( "encoding/json" "net/http" "net/http/httptest" "strings" "testing" "time" "ds2api/internal/auth" dsclient "ds2api/internal/deepseek/client" ) func TestIsVercelStreamPrepareRequest(t *testing.T) { req := httptest.NewRequest("POST", "/v1/chat/completions?__stream_prepare=1", nil) if !isVercelStreamPrepareRequest(req) { t.Fatalf("expected prepare request to be detected") } req2 := httptest.NewRequest("POST", "/v1/chat/completions", nil) if isVercelStreamPrepareRequest(req2) { t.Fatalf("expected non-prepare request") } } func TestIsVercelStreamReleaseRequest(t *testing.T) { req := httptest.NewRequest("POST", "/v1/chat/completions?__stream_release=1", nil) if !isVercelStreamReleaseRequest(req) { t.Fatalf("expected release request to be detected") } req2 := httptest.NewRequest("POST", "/v1/chat/completions", nil) if isVercelStreamReleaseRequest(req2) { t.Fatalf("expected non-release request") } } func TestVercelInternalSecret(t *testing.T) { t.Run("prefer explicit secret", func(t *testing.T) { t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "stream-secret") t.Setenv("DS2API_ADMIN_KEY", "admin-fallback") if got := vercelInternalSecret(); got != "stream-secret" { t.Fatalf("expected explicit secret, got %q", got) } }) t.Run("fallback to admin key", func(t *testing.T) { t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "") t.Setenv("DS2API_ADMIN_KEY", "admin-fallback") if got := vercelInternalSecret(); got != "admin-fallback" { t.Fatalf("expected admin key fallback, got %q", got) } }) t.Run("default admin when env missing", func(t *testing.T) { t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "") t.Setenv("DS2API_ADMIN_KEY", "") if got := vercelInternalSecret(); got != "admin" { t.Fatalf("expected default admin fallback, got %q", got) } }) } func TestStreamLeaseLifecycle(t *testing.T) { h := &Handler{} leaseID := h.holdStreamLease(&auth.RequestAuth{UseConfigToken: false}) if leaseID == "" { t.Fatalf("expected non-empty lease id") } if ok := h.releaseStreamLease(leaseID); !ok { t.Fatalf("expected lease release success") } if ok := h.releaseStreamLease(leaseID); ok { t.Fatalf("expected duplicate release to fail") } } func TestStreamLeaseTTL(t *testing.T) { t.Setenv("DS2API_VERCEL_STREAM_LEASE_TTL_SECONDS", "120") if got := streamLeaseTTL(); got != 120*time.Second { t.Fatalf("expected ttl=120s, got %v", got) } t.Setenv("DS2API_VERCEL_STREAM_LEASE_TTL_SECONDS", "invalid") if got := streamLeaseTTL(); got != 15*time.Minute { t.Fatalf("expected default ttl on invalid value, got %v", got) } } func TestHandleVercelStreamPrepareAppliesCurrentInputFile(t *testing.T) { t.Setenv("VERCEL", "1") t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "stream-secret") ds := &inlineUploadDSStub{} h := &Handler{ Store: mockOpenAIConfig{ currentInputEnabled: true, }, Auth: streamStatusAuthStub{}, DS: ds, } reqBody, _ := json.Marshal(map[string]any{ "model": "deepseek-v4-flash", "messages": historySplitTestMessages(), "stream": true, }) req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions?__stream_prepare=1", strings.NewReader(string(reqBody))) req.Header.Set("Authorization", "Bearer direct-token") req.Header.Set("Content-Type", "application/json") req.Header.Set("X-Ds2-Internal-Token", "stream-secret") rec := httptest.NewRecorder() h.handleVercelStreamPrepare(rec, req) if rec.Code != http.StatusOK { t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) } if len(ds.uploadCalls) != 1 { t.Fatalf("expected 1 current input upload, got %d", len(ds.uploadCalls)) } var body map[string]any if err := json.NewDecoder(rec.Body).Decode(&body); err != nil { t.Fatalf("decode failed: %v", err) } payload, _ := body["payload"].(map[string]any) if payload == nil { t.Fatalf("expected payload object, got %#v", body["payload"]) } promptText, _ := payload["prompt"].(string) if !strings.Contains(promptText, "Continue from the latest state in the attached DS2API_HISTORY.txt context.") { t.Fatalf("expected continuation prompt, got %s", promptText) } if strings.Contains(promptText, "first user turn") || strings.Contains(promptText, "latest user turn") { t.Fatalf("expected original turns hidden from prompt, got %s", promptText) } refIDs, _ := payload["ref_file_ids"].([]any) if len(refIDs) == 0 || refIDs[0] != "file-inline-1" { t.Fatalf("expected uploaded history file first in ref_file_ids, got %#v", payload["ref_file_ids"]) } } func TestHandleVercelStreamPrepareMapsCurrentInputFileManagedAuthFailureTo401(t *testing.T) { t.Setenv("VERCEL", "1") t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "stream-secret") ds := &inlineUploadDSStub{ uploadErr: &dsclient.RequestFailure{Op: "upload file", Kind: dsclient.FailureManagedUnauthorized, Message: "expired token"}, } h := &Handler{ Store: mockOpenAIConfig{ currentInputEnabled: true, }, Auth: streamStatusManagedAuthStub{}, DS: ds, } reqBody, _ := json.Marshal(map[string]any{ "model": "deepseek-v4-flash", "messages": historySplitTestMessages(), "stream": true, }) req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions?__stream_prepare=1", strings.NewReader(string(reqBody))) req.Header.Set("Authorization", "Bearer managed-key") req.Header.Set("Content-Type", "application/json") req.Header.Set("X-Ds2-Internal-Token", "stream-secret") rec := httptest.NewRecorder() h.handleVercelStreamPrepare(rec, req) if rec.Code != http.StatusUnauthorized { t.Fatalf("expected 401, got %d body=%s", rec.Code, rec.Body.String()) } if !strings.Contains(rec.Body.String(), "Please re-login the account in admin") { t.Fatalf("expected managed auth error message, got %s", rec.Body.String()) } }