mirror of
https://github.com/CJackHwang/ds2api.git
synced 2026-05-11 19:57:41 +08:00
403 lines
14 KiB
Go
403 lines
14 KiB
Go
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})
|
|
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 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"])
|
|
}
|
|
}
|
|
|
|
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 ds.uploadCalls[0].Filename != "DS2API_HISTORY.txt" || ds.uploadCalls[1].Filename != "DS2API_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, "DS2API_TOOLS.txt") || !strings.Contains(promptText, "TOOL CALL FORMAT") {
|
|
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",
|
|
FinalPrompt: "Continue from the latest state in the attached DS2API_HISTORY.txt context. Available tool descriptions and parameter schemas are attached in DS2API_TOOLS.txt; use only those tools and follow the tool-call format rules in this prompt.",
|
|
PromptTokenText: "# DS2API_HISTORY.txt\n\n=== 1. USER ===\nhello\n\n# DS2API_TOOLS.txt\nAvailable tool 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: "# DS2API_HISTORY.txt\n\n=== 1. USER ===\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 ds.uploadCalls[0].Filename != "DS2API_HISTORY.txt" || ds.uploadCalls[1].Filename != "DS2API_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, "DS2API_TOOLS.txt") {
|
|
t.Fatalf("expected switched payload prompt to retain tools file reference, got %q", promptText)
|
|
}
|
|
}
|