| package chat |
|
|
| import ( |
| "context" |
| "encoding/json" |
| "net/http" |
| "net/http/httptest" |
| "strings" |
| "testing" |
| "time" |
|
|
| "ds2api/internal/account" |
| "ds2api/internal/auth" |
| "ds2api/internal/config" |
| dsclient "ds2api/internal/deepseek/client" |
| "ds2api/internal/promptcompat" |
| ) |
|
|
| 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}, promptcompat.StandardRequest{}, "test-session-id") |
| if leaseID == "" { |
| t.Fatalf("expected non-empty lease id") |
| } |
| if lease, ok := h.releaseStreamLease(leaseID); !ok { |
| t.Fatalf("expected lease release success") |
| } else if lease.SessionID != "test-session-id" { |
| t.Fatalf("expected released session id, got %q", lease.SessionID) |
| } |
| 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, ".txt") { |
| 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 TestHandleVercelStreamPrepareUsesHalfwidthDSMLToolPrompt(t *testing.T) { |
| t.Setenv("VERCEL", "1") |
| t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "stream-secret") |
|
|
| h := &Handler{ |
| Store: mockOpenAIConfig{}, |
| Auth: streamStatusAuthStub{}, |
| DS: &inlineUploadDSStub{}, |
| } |
|
|
| reqBody, _ := json.Marshal(map[string]any{ |
| "model": "deepseek-v4-flash", |
| "messages": []any{ |
| map[string]any{"role": "user", "content": "search docs"}, |
| }, |
| "tools": []any{ |
| map[string]any{ |
| "type": "function", |
| "function": map[string]any{ |
| "name": "search", |
| "description": "search docs", |
| "parameters": map[string]any{ |
| "type": "object", |
| "properties": map[string]any{ |
| "query": map[string]any{"type": "string"}, |
| }, |
| "required": []any{"query"}, |
| }, |
| }, |
| }, |
| }, |
| "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()) |
| } |
| var body map[string]any |
| if err := json.NewDecoder(rec.Body).Decode(&body); err != nil { |
| t.Fatalf("decode failed: %v", err) |
| } |
| finalPrompt, _ := body["final_prompt"].(string) |
| payload, _ := body["payload"].(map[string]any) |
| payloadPrompt, _ := payload["prompt"].(string) |
| for label, promptText := range map[string]string{"final_prompt": finalPrompt, "payload.prompt": payloadPrompt} { |
| if !strings.Contains(promptText, "<|DSML|tool_calls>") || !strings.Contains(promptText, "Tag punctuation alphabet: ASCII < > / = \" plus the halfwidth pipe |.") { |
| t.Fatalf("expected %s to contain halfwidth DSML tool instructions, got %q", label, promptText) |
| } |
| if strings.Contains(promptText, "\uff5c") || strings.Contains(promptText, "full"+"width vertical bar") { |
| t.Fatalf("expected %s not to contain legacy pipe guidance, got %q", label, promptText) |
| } |
| } |
| toolNames, _ := body["tool_names"].([]any) |
| if len(toolNames) != 1 || toolNames[0] != "search" { |
| t.Fatalf("expected prepared tool names to align with request tools, got %#v", body["tool_names"]) |
| } |
| } |
|
|
| type vercelReleaseAutoDeleteDSStub struct { |
| resp *http.Response |
| deleteCallCount int |
| deletedSessionID string |
| deletedToken string |
| deleteErr error |
| events *[]string |
| } |
|
|
| func (m *vercelReleaseAutoDeleteDSStub) CreateSession(_ context.Context, _ *auth.RequestAuth, _ int) (string, error) { |
| return "session-id", nil |
| } |
|
|
| func (m *vercelReleaseAutoDeleteDSStub) GetPow(_ context.Context, _ *auth.RequestAuth, _ int) (string, error) { |
| return "pow", nil |
| } |
|
|
| func (m *vercelReleaseAutoDeleteDSStub) UploadFile(_ context.Context, _ *auth.RequestAuth, _ dsclient.UploadFileRequest, _ int) (*dsclient.UploadFileResult, error) { |
| return &dsclient.UploadFileResult{ID: "file-id", Filename: "file.txt", Bytes: 1, Status: "uploaded"}, nil |
| } |
|
|
| func (m *vercelReleaseAutoDeleteDSStub) CallCompletion(_ context.Context, _ *auth.RequestAuth, _ map[string]any, _ string, _ int) (*http.Response, error) { |
| return m.resp, nil |
| } |
|
|
| func (m *vercelReleaseAutoDeleteDSStub) DeleteSessionForToken(_ context.Context, token string, sessionID string) (*dsclient.DeleteSessionResult, error) { |
| if m.events != nil { |
| *m.events = append(*m.events, "delete") |
| } |
| m.deleteCallCount++ |
| m.deletedSessionID = sessionID |
| m.deletedToken = token |
| if m.deleteErr != nil { |
| return nil, m.deleteErr |
| } |
| return &dsclient.DeleteSessionResult{SessionID: sessionID, Success: true}, nil |
| } |
|
|
| func (m *vercelReleaseAutoDeleteDSStub) DeleteAllSessionsForToken(_ context.Context, _ string) error { |
| return nil |
| } |
|
|
| type vercelReleaseAuthStub struct { |
| events *[]string |
| } |
|
|
| func (a *vercelReleaseAuthStub) Determine(_ *http.Request) (*auth.RequestAuth, error) { |
| return &auth.RequestAuth{DeepSeekToken: "test-token", AccountID: "test-account"}, nil |
| } |
|
|
| func (a *vercelReleaseAuthStub) DetermineCaller(_ *http.Request) (*auth.RequestAuth, error) { |
| return &auth.RequestAuth{DeepSeekToken: "test-token", AccountID: "test-account"}, nil |
| } |
|
|
| func (a *vercelReleaseAuthStub) Release(_ *auth.RequestAuth) { |
| if a.events != nil { |
| *a.events = append(*a.events, "release") |
| } |
| } |
|
|
| func TestHandleVercelStreamReleaseTriggersAutoDelete(t *testing.T) { |
| t.Setenv("VERCEL", "1") |
| t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "stream-secret") |
|
|
| events := []string{} |
| ds := &vercelReleaseAutoDeleteDSStub{events: &events} |
| h := &Handler{ |
| Store: mockOpenAIConfig{ |
| autoDeleteMode: "single", |
| }, |
| Auth: &vercelReleaseAuthStub{events: &events}, |
| DS: ds, |
| } |
|
|
| leaseID := h.holdStreamLease(&auth.RequestAuth{DeepSeekToken: "test-token", AccountID: "test-account"}, promptcompat.StandardRequest{}, "session-to-delete") |
| if leaseID == "" { |
| t.Fatalf("expected non-empty lease id") |
| } |
|
|
| reqBody := map[string]any{"lease_id": leaseID} |
| reqJSON, _ := json.Marshal(reqBody) |
| req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions?__stream_release=1", strings.NewReader(string(reqJSON))) |
| req.Header.Set("X-Ds2-Internal-Token", "stream-secret") |
| req.Header.Set("Content-Type", "application/json") |
| rec := httptest.NewRecorder() |
|
|
| h.handleVercelStreamRelease(rec, req) |
|
|
| if rec.Code != http.StatusOK { |
| t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) |
| } |
| if ds.deleteCallCount != 1 { |
| t.Fatalf("expected auto delete call count=1, got %d", ds.deleteCallCount) |
| } |
| if ds.deletedSessionID != "session-to-delete" { |
| t.Fatalf("expected deleted session id=session-to-delete, got %q", ds.deletedSessionID) |
| } |
| if got, want := strings.Join(events, ","), "delete,release"; got != want { |
| t.Fatalf("expected auto-delete before auth release, got %s", got) |
| } |
| } |
|
|
| func TestHandleVercelStreamPrepareUploadsToolsSeparately(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": []any{ |
| map[string]any{"role": "user", "content": "search docs"}, |
| }, |
| "tools": []any{ |
| map[string]any{ |
| "type": "function", |
| "function": map[string]any{ |
| "name": "search", |
| "description": "search docs", |
| "parameters": map[string]any{"type": "object"}, |
| }, |
| }, |
| }, |
| "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) != 2 { |
| t.Fatalf("expected history and tools uploads, got %d", len(ds.uploadCalls)) |
| } |
| if strings.Contains(strings.ToLower(ds.uploadCalls[0].Filename), "history") || !strings.HasSuffix(ds.uploadCalls[0].Filename, ".txt") || ds.uploadCalls[1].Filename != "context_tools.txt" { |
| t.Fatalf("unexpected upload filenames: %#v", ds.uploadCalls) |
| } |
| if strings.Contains(string(ds.uploadCalls[0].Data), "Description: search docs") { |
| t.Fatalf("history transcript should not embed tool descriptions, got %q", string(ds.uploadCalls[0].Data)) |
| } |
|
|
| var body map[string]any |
| if err := json.NewDecoder(rec.Body).Decode(&body); err != nil { |
| t.Fatalf("decode failed: %v", err) |
| } |
| finalPrompt, _ := body["final_prompt"].(string) |
| payload, _ := body["payload"].(map[string]any) |
| payloadPrompt, _ := payload["prompt"].(string) |
| for label, promptText := range map[string]string{"final_prompt": finalPrompt, "payload.prompt": payloadPrompt} { |
| if !strings.Contains(promptText, "context_tools.txt") || !strings.Contains(promptText, "TOOL CALL SCHEME") { |
| t.Fatalf("expected %s to reference tools file and retain tool instructions, got %q", label, promptText) |
| } |
| if strings.Contains(promptText, "Description: search docs") { |
| t.Fatalf("expected %s not to inline tool descriptions, got %q", label, promptText) |
| } |
| } |
| refIDs, _ := payload["ref_file_ids"].([]any) |
| if len(refIDs) < 2 || refIDs[0] != "file-inline-1" || refIDs[1] != "file-inline-2" { |
| t.Fatalf("expected history and tools ref ids first, 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()) |
| } |
| } |
|
|
| func TestHandleVercelStreamSwitchReuploadsCurrentInputFile(t *testing.T) { |
| t.Setenv("VERCEL", "1") |
| t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "stream-secret") |
| t.Setenv("DS2API_CONFIG_JSON", `{ |
| "keys":["managed-key"], |
| "accounts":[ |
| {"email":"acc1@test.com","password":"pwd"}, |
| {"email":"acc2@test.com","password":"pwd"} |
| ] |
| }`) |
| store := config.LoadStore() |
| resolver := auth.NewResolver(store, account.NewPool(store), func(_ context.Context, acc config.Account) (string, error) { |
| return "token-" + acc.Identifier(), nil |
| }) |
| authReq := httptest.NewRequest(http.MethodPost, "/", nil) |
| authReq.Header.Set("Authorization", "Bearer managed-key") |
| a, err := resolver.Determine(authReq) |
| if err != nil { |
| t.Fatalf("determine failed: %v", err) |
| } |
| defer resolver.Release(a) |
|
|
| ds := &inlineUploadDSStub{} |
| h := &Handler{ |
| Store: mockOpenAIConfig{currentInputEnabled: true}, |
| Auth: resolver, |
| DS: ds, |
| } |
| stdReq := promptcompat.StandardRequest{ |
| RequestedModel: "deepseek-v4-flash", |
| ResolvedModel: "deepseek-v4-flash", |
| ResponseModel: "deepseek-v4-flash", |
| CurrentInputFilename: "ctx-fixed-1234.txt", |
| FinalPrompt: "已附上下文文件「ctx-fixed-1234.txt」。直接基于该文件完成当前请求，不要重复之前内容。 可用工具与参数说明在 context_tools.txt 中，仅可依据该文件调用工具。", |
| PromptTokenText: "# context_context.txt\n\n=== 1 ===\n[r=1]\nhello\n\n# context_tools.txt\nTool descriptions and parameter schemas for this request.\n\nYou have access to these tools:\n\nTool: search\nDescription: search docs\nParameters: {\"type\":\"object\"}\n", |
| HistoryText: "# context_context.txt\n\n=== 1 ===\n[r=1]\nhello\n", |
| CurrentInputFileApplied: true, |
| CurrentInputFileID: "file-old", |
| CurrentToolsFileID: "file-old-tools", |
| ToolsRaw: []any{ |
| map[string]any{ |
| "type": "function", |
| "function": map[string]any{ |
| "name": "search", |
| "description": "search docs", |
| "parameters": map[string]any{"type": "object"}, |
| }, |
| }, |
| }, |
| RefFileIDs: []string{"file-old", "file-old-tools", "client-file"}, |
| Thinking: true, |
| } |
| leaseID := h.holdStreamLease(a, stdReq, "") |
| req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions?__stream_switch=1", strings.NewReader(`{"lease_id":"`+leaseID+`"}`)) |
| req.Header.Set("X-Ds2-Internal-Token", "stream-secret") |
| rec := httptest.NewRecorder() |
|
|
| h.handleVercelStreamSwitch(rec, req) |
|
|
| if rec.Code != http.StatusOK { |
| t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) |
| } |
| if len(ds.uploadCalls) != 2 { |
| t.Fatalf("expected current input and tools reupload on switched account, got %d", len(ds.uploadCalls)) |
| } |
| if strings.Contains(strings.ToLower(ds.uploadCalls[0].Filename), "history") || !strings.HasSuffix(ds.uploadCalls[0].Filename, ".txt") || ds.uploadCalls[1].Filename != "context_tools.txt" { |
| t.Fatalf("unexpected reupload filenames: %#v", ds.uploadCalls) |
| } |
| var body map[string]any |
| if err := json.NewDecoder(rec.Body).Decode(&body); err != nil { |
| t.Fatalf("decode failed: %v", err) |
| } |
| if body["deepseek_token"] != "token-acc2@test.com" { |
| t.Fatalf("expected switched account token, got %#v", body["deepseek_token"]) |
| } |
| payload, _ := body["payload"].(map[string]any) |
| refIDs, _ := payload["ref_file_ids"].([]any) |
| if len(refIDs) != 3 || refIDs[0] != "file-inline-1" || refIDs[1] != "file-inline-2" || refIDs[2] != "client-file" { |
| t.Fatalf("expected reuploaded current input ref plus client ref, got %#v", payload["ref_file_ids"]) |
| } |
| promptText, _ := payload["prompt"].(string) |
| if !strings.Contains(promptText, "context_tools.txt") { |
| t.Fatalf("expected switched payload prompt to retain tools file reference, got %q", promptText) |
| } |
| } |
|
|
