feat: align Go/Node DSML tool-call parsing drift tolerance and update API docs

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
CJACK
2026-05-10 16:17:46 +08:00
parent cee8757d14
commit eaeb403fda
32 changed files with 879 additions and 102 deletions

View File

@@ -66,7 +66,7 @@ func (h *Handler) handleNonStreamWithRetry(w http.ResponseWriter, ctx context.Co
writeJSON(w, http.StatusOK, respBody)
}
func (h *Handler) handleStreamWithRetry(w http.ResponseWriter, r *http.Request, a *auth.RequestAuth, resp *http.Response, payload map[string]any, pow, completionID string, sessionIDRef *string, model, finalPrompt string, refFileTokens int, thinkingEnabled, searchEnabled bool, toolNames []string, toolsRaw any, toolChoice promptcompat.ToolChoicePolicy, historySession *chatHistorySession) {
func (h *Handler) handleStreamWithRetry(w http.ResponseWriter, r *http.Request, a *auth.RequestAuth, resp *http.Response, payload map[string]any, pow, completionID string, sessionIDRef *string, stdReq promptcompat.StandardRequest, model, finalPrompt string, refFileTokens int, thinkingEnabled, searchEnabled bool, toolNames []string, toolsRaw any, toolChoice promptcompat.ToolChoicePolicy, historySession *chatHistorySession) {
streamRuntime, initialType, ok := h.prepareChatStreamRuntime(w, resp, completionID, model, finalPrompt, refFileTokens, thinkingEnabled, searchEnabled, toolNames, toolsRaw, toolChoice, historySession)
if !ok {
return
@@ -78,6 +78,8 @@ func (h *Handler) handleStreamWithRetry(w http.ResponseWriter, r *http.Request,
RetryMaxAttempts: emptyOutputRetryMaxAttempts(),
MaxAttempts: 3,
UsagePrompt: finalPrompt,
Request: stdReq,
CurrentInputFile: h.Store,
}, completionruntime.StreamRetryHooks{
ConsumeAttempt: func(currentResp *http.Response, allowDeferEmpty bool) (bool, bool) {
return h.consumeChatStreamAttempt(r, currentResp, streamRuntime, initialType, thinkingEnabled, historySession, allowDeferEmpty)

View File

@@ -33,6 +33,7 @@ type Handler struct {
type streamLease struct {
Auth *auth.RequestAuth
Standard promptcompat.StandardRequest
ExpiresAt time.Time
}

View File

@@ -28,6 +28,10 @@ func (h *Handler) ChatCompletions(w http.ResponseWriter, r *http.Request) {
h.handleVercelStreamPow(w, r)
return
}
if isVercelStreamSwitchRequest(r) {
h.handleVercelStreamSwitch(w, r)
return
}
if isVercelStreamPrepareRequest(r) {
h.handleVercelStreamPrepare(w, r)
return
@@ -114,7 +118,7 @@ func (h *Handler) ChatCompletions(w http.ResponseWriter, r *http.Request) {
}
streamReq := start.Request
refFileTokens := streamReq.RefFileTokens
h.handleStreamWithRetry(w, r, a, start.Response, start.Payload, start.Pow, sessionID, &sessionID, streamReq.ResponseModel, streamReq.PromptTokenText, refFileTokens, streamReq.Thinking, streamReq.Search, streamReq.ToolNames, streamReq.ToolsRaw, streamReq.ToolChoice, historySession)
h.handleStreamWithRetry(w, r, a, start.Response, start.Payload, start.Pow, sessionID, &sessionID, streamReq, streamReq.ResponseModel, streamReq.PromptTokenText, refFileTokens, streamReq.Thinking, streamReq.Search, streamReq.ToolNames, streamReq.ToolsRaw, streamReq.ToolChoice, historySession)
}
func (h *Handler) autoDeleteRemoteSession(ctx context.Context, a *auth.RequestAuth, sessionID string) {

View File

@@ -1,6 +1,7 @@
package chat
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
@@ -8,8 +9,11 @@ import (
"testing"
"time"
"ds2api/internal/account"
"ds2api/internal/auth"
"ds2api/internal/config"
dsclient "ds2api/internal/deepseek/client"
"ds2api/internal/promptcompat"
)
func TestIsVercelStreamPrepareRequest(t *testing.T) {
@@ -206,6 +210,76 @@ func TestHandleVercelStreamPrepareUsesHalfwidthDSMLToolPrompt(t *testing.T) {
}
}
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")
@@ -241,3 +315,88 @@ func TestHandleVercelStreamPrepareMapsCurrentInputFileManagedAuthFailureTo401(t
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)
}
}

View File

@@ -11,6 +11,7 @@ import (
"ds2api/internal/auth"
"ds2api/internal/config"
"ds2api/internal/httpapi/openai/history"
"ds2api/internal/promptcompat"
"ds2api/internal/util"
@@ -96,7 +97,7 @@ func (h *Handler) handleVercelStreamPrepare(w http.ResponseWriter, r *http.Reque
}
payload := stdReq.CompletionPayload(sessionID)
leaseID := h.holdStreamLease(a)
leaseID := h.holdStreamLease(a, stdReq)
if leaseID == "" {
writeOpenAIError(w, http.StatusInternalServerError, "failed to create stream lease")
return
@@ -185,6 +186,80 @@ func (h *Handler) handleVercelStreamPow(w http.ResponseWriter, r *http.Request)
})
}
func (h *Handler) handleVercelStreamSwitch(w http.ResponseWriter, r *http.Request) {
if !config.IsVercel() {
http.NotFound(w, r)
return
}
h.sweepExpiredStreamLeases()
internalSecret := vercelInternalSecret()
internalToken := strings.TrimSpace(r.Header.Get("X-Ds2-Internal-Token"))
if internalSecret == "" || subtle.ConstantTimeCompare([]byte(internalToken), []byte(internalSecret)) != 1 {
writeOpenAIError(w, http.StatusUnauthorized, "unauthorized internal request")
return
}
var req map[string]any
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeOpenAIError(w, http.StatusBadRequest, "invalid json")
return
}
leaseID, _ := req["lease_id"].(string)
leaseID = strings.TrimSpace(leaseID)
if leaseID == "" {
writeOpenAIError(w, http.StatusBadRequest, "lease_id is required")
return
}
lease, ok := h.lookupStreamLease(leaseID)
if !ok || lease.Auth == nil {
writeOpenAIError(w, http.StatusNotFound, "stream lease not found or expired")
return
}
a := lease.Auth
if !a.UseConfigToken || !a.SwitchAccount(r.Context()) {
writeOpenAIErrorWithCode(w, http.StatusTooManyRequests, "Upstream account hit a rate limit and returned reasoning without visible output.", "upstream_empty_output")
return
}
stdReq := lease.Standard
var err error
if stdReq.CurrentInputFileApplied {
stdReq, err = (history.Service{Store: h.Store, DS: h.DS}).ReuploadAppliedCurrentInputFile(r.Context(), a, stdReq)
if err != nil {
status, message := mapCurrentInputFileError(err)
writeOpenAIError(w, status, message)
return
}
}
sessionID, err := h.DS.CreateSession(r.Context(), a, 3)
if err != nil {
writeOpenAIError(w, http.StatusUnauthorized, "Account token is invalid. Please re-login the account in admin.")
return
}
powHeader, err := h.DS.GetPow(r.Context(), a, 3)
if err != nil {
writeOpenAIError(w, http.StatusUnauthorized, "Failed to get PoW (invalid token or unknown error).")
return
}
if strings.TrimSpace(a.DeepSeekToken) == "" {
writeOpenAIError(w, http.StatusUnauthorized, "Account token is invalid. Please re-login the account in admin.")
return
}
h.updateStreamLeaseStandard(leaseID, stdReq)
writeJSON(w, http.StatusOK, map[string]any{
"session_id": sessionID,
"lease_id": leaseID,
"model": stdReq.ResponseModel,
"final_prompt": stdReq.FinalPrompt,
"thinking_enabled": stdReq.Thinking,
"search_enabled": stdReq.Search,
"tool_names": stdReq.ToolNames,
"deepseek_token": a.DeepSeekToken,
"pow_header": powHeader,
"payload": stdReq.CompletionPayload(sessionID),
})
}
func isVercelStreamPrepareRequest(r *http.Request) bool {
if r == nil {
return false
@@ -206,6 +281,13 @@ func isVercelStreamPowRequest(r *http.Request) bool {
return strings.TrimSpace(r.URL.Query().Get("__stream_pow")) == "1"
}
func isVercelStreamSwitchRequest(r *http.Request) bool {
if r == nil {
return false
}
return strings.TrimSpace(r.URL.Query().Get("__stream_switch")) == "1"
}
func vercelInternalSecret() string {
if v := strings.TrimSpace(os.Getenv("DS2API_VERCEL_INTERNAL_SECRET")); v != "" {
return v
@@ -216,10 +298,14 @@ func vercelInternalSecret() string {
return "admin"
}
func (h *Handler) holdStreamLease(a *auth.RequestAuth) string {
func (h *Handler) holdStreamLease(a *auth.RequestAuth, standards ...promptcompat.StandardRequest) string {
if a == nil {
return ""
}
var stdReq promptcompat.StandardRequest
if len(standards) > 0 {
stdReq = standards[0]
}
now := time.Now()
ttl := streamLeaseTTL()
if ttl <= 0 {
@@ -234,6 +320,7 @@ func (h *Handler) holdStreamLease(a *auth.RequestAuth) string {
leaseID := newLeaseID()
h.streamLeases[leaseID] = streamLease{
Auth: a,
Standard: stdReq,
ExpiresAt: now.Add(ttl),
}
h.leaseMu.Unlock()
@@ -241,20 +328,43 @@ func (h *Handler) holdStreamLease(a *auth.RequestAuth) string {
return leaseID
}
func (h *Handler) lookupStreamLeaseAuth(leaseID string) *auth.RequestAuth {
func (h *Handler) lookupStreamLease(leaseID string) (streamLease, bool) {
leaseID = strings.TrimSpace(leaseID)
if leaseID == "" {
return nil
return streamLease{}, false
}
h.leaseMu.Lock()
lease, ok := h.streamLeases[leaseID]
h.leaseMu.Unlock()
if !ok || time.Now().After(lease.ExpiresAt) {
return streamLease{}, false
}
return lease, true
}
func (h *Handler) lookupStreamLeaseAuth(leaseID string) *auth.RequestAuth {
lease, ok := h.lookupStreamLease(leaseID)
if !ok {
return nil
}
return lease.Auth
}
func (h *Handler) updateStreamLeaseStandard(leaseID string, stdReq promptcompat.StandardRequest) {
leaseID = strings.TrimSpace(leaseID)
if leaseID == "" {
return
}
h.leaseMu.Lock()
defer h.leaseMu.Unlock()
lease, ok := h.streamLeases[leaseID]
if !ok {
return
}
lease.Standard = stdReq
h.streamLeases[leaseID] = lease
}
func (h *Handler) releaseStreamLease(leaseID string) bool {
leaseID = strings.TrimSpace(leaseID)
if leaseID == "" {