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

@@ -114,7 +114,7 @@ func ExecuteNonStreamStartedWithRetry(ctx context.Context, ds DeepSeekCaller, a
turn, outErr := collectAttempt(currentResp, stdReq, usagePrompt, opts)
if outErr != nil {
if canRetryOnAlternateAccount(ctx, a, outErr, opts.RetryEnabled, &accountSwitchAttempted) {
switched, switchErr := startStandardCompletionOnAlternateAccount(ctx, ds, a, stdReq, maxAttempts)
switched, switchErr := startStandardCompletionOnAlternateAccount(ctx, ds, a, stdReq, opts, maxAttempts)
if switchErr != nil {
return NonStreamResult{SessionID: sessionID, Payload: payload, Attempts: attempts}, switchErr
}
@@ -154,7 +154,7 @@ func ExecuteNonStreamStartedWithRetry(ctx context.Context, ds DeepSeekCaller, a
}
if !opts.RetryEnabled || !assistantturn.ShouldRetryEmptyOutput(turn, attempts, retryMax) {
if canRetryOnAlternateAccount(ctx, a, turn.Error, opts.RetryEnabled, &accountSwitchAttempted) {
switched, switchErr := startStandardCompletionOnAlternateAccount(ctx, ds, a, stdReq, maxAttempts)
switched, switchErr := startStandardCompletionOnAlternateAccount(ctx, ds, a, stdReq, opts, maxAttempts)
if switchErr != nil {
return NonStreamResult{SessionID: sessionID, Payload: payload, Turn: turn, Attempts: attempts}, switchErr
}
@@ -205,7 +205,12 @@ func canRetryOnAlternateAccount(ctx context.Context, a *auth.RequestAuth, outErr
return a.SwitchAccount(ctx)
}
func startStandardCompletionOnAlternateAccount(ctx context.Context, ds DeepSeekCaller, a *auth.RequestAuth, stdReq promptcompat.StandardRequest, maxAttempts int) (StartResult, *assistantturn.OutputError) {
func startStandardCompletionOnAlternateAccount(ctx context.Context, ds DeepSeekCaller, a *auth.RequestAuth, stdReq promptcompat.StandardRequest, opts Options, maxAttempts int) (StartResult, *assistantturn.OutputError) {
var prepErr *assistantturn.OutputError
stdReq, prepErr = reuploadCurrentInputFileForAccount(ctx, ds, a, stdReq, opts)
if prepErr != nil {
return StartResult{Request: stdReq}, prepErr
}
sessionID, err := ds.CreateSession(ctx, a, maxAttempts)
if err != nil {
return StartResult{}, authOutputError(a)
@@ -222,6 +227,18 @@ func startStandardCompletionOnAlternateAccount(ctx context.Context, ds DeepSeekC
return StartResult{SessionID: sessionID, Payload: payload, Pow: pow, Response: resp, Request: stdReq}, nil
}
func reuploadCurrentInputFileForAccount(ctx context.Context, ds DeepSeekCaller, a *auth.RequestAuth, stdReq promptcompat.StandardRequest, opts Options) (promptcompat.StandardRequest, *assistantturn.OutputError) {
if opts.CurrentInputFile == nil || !stdReq.CurrentInputFileApplied {
return stdReq, nil
}
out, err := (history.Service{Store: opts.CurrentInputFile, DS: ds}).ReuploadAppliedCurrentInputFile(ctx, a, stdReq)
if err != nil {
status, message := history.MapError(err)
return out, &assistantturn.OutputError{Status: status, Message: message, Code: "error"}
}
return out, nil
}
func collectAttempt(resp *http.Response, stdReq promptcompat.StandardRequest, usagePrompt string, opts Options) (assistantturn.Turn, *assistantturn.OutputError) {
defer func() {
if err := resp.Body.Close(); err != nil {

View File

@@ -38,8 +38,11 @@ func (f *fakeDeepSeekCaller) GetPow(context.Context, *auth.RequestAuth, int) (st
return "pow", nil
}
func (f *fakeDeepSeekCaller) UploadFile(_ context.Context, _ *auth.RequestAuth, req dsclient.UploadFileRequest, _ int) (*dsclient.UploadFileResult, error) {
func (f *fakeDeepSeekCaller) UploadFile(_ context.Context, a *auth.RequestAuth, req dsclient.UploadFileRequest, _ int) (*dsclient.UploadFileResult, error) {
f.uploads = append(f.uploads, req)
if a != nil && a.AccountID != "" {
return &dsclient.UploadFileResult{ID: "file-runtime-" + a.AccountID}, nil
}
return &dsclient.UploadFileResult{ID: "file-runtime-1"}, nil
}
@@ -162,6 +165,66 @@ func TestExecuteNonStreamWithRetrySwitchesManagedAccountBeforeFinal429(t *testin
}
}
func TestExecuteNonStreamWithRetryReuploadsCurrentInputFileAfterAccountSwitch(t *testing.T) {
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
})
req, _ := http.NewRequest(http.MethodPost, "/", nil)
req.Header.Set("Authorization", "Bearer managed-key")
a, err := resolver.Determine(req)
if err != nil {
t.Fatalf("determine failed: %v", err)
}
defer resolver.Release(a)
ds := &fakeDeepSeekCaller{
sessionByAccount: true,
responses: []*http.Response{
sseHTTPResponse(http.StatusOK, `data: {"response_message_id":11,"p":"response/thinking_content","v":"first empty"}`),
sseHTTPResponse(http.StatusOK, `data: {"response_message_id":12,"p":"response/thinking_content","v":"retry empty"}`),
sseHTTPResponse(http.StatusOK, `data: {"response_message_id":21,"p":"response/content","v":"ok from second account"}`),
},
}
stdReq := promptcompat.StandardRequest{
Surface: "test",
RequestedModel: "deepseek-v4-flash",
ResolvedModel: "deepseek-v4-flash",
ResponseModel: "deepseek-v4-flash",
Messages: []any{
map[string]any{"role": "user", "content": "large current input"},
},
PromptTokenText: "large current input",
FinalPrompt: "large current input",
Thinking: true,
}
result, outErr := ExecuteNonStreamWithRetry(context.Background(), ds, a, stdReq, Options{
RetryEnabled: true,
CurrentInputFile: currentInputRuntimeConfig{},
})
if outErr != nil {
t.Fatalf("unexpected output error after account switch retry: %#v", outErr)
}
if result.Turn.Text != "ok from second account" {
t.Fatalf("text mismatch after switch retry: %q", result.Turn.Text)
}
if len(ds.uploads) != 2 {
t.Fatalf("expected current input file uploaded once per account, got %d", len(ds.uploads))
}
refIDs, _ := ds.payloads[2]["ref_file_ids"].([]any)
if len(refIDs) != 1 || refIDs[0] != "file-runtime-acc2@test.com" {
t.Fatalf("expected switched account ref_file_ids to use reuploaded file, got %#v", ds.payloads[2]["ref_file_ids"])
}
}
func TestExecuteNonStreamWithRetryUsesParentMessageForEmptyRetry(t *testing.T) {
ds := &fakeDeepSeekCaller{responses: []*http.Response{
sseHTTPResponse(http.StatusOK, `data: {"response_message_id":77,"p":"response/thinking_content","v":"plan"}`),

View File

@@ -9,7 +9,9 @@ import (
"ds2api/internal/assistantturn"
"ds2api/internal/auth"
"ds2api/internal/config"
"ds2api/internal/httpapi/openai/history"
"ds2api/internal/httpapi/openai/shared"
"ds2api/internal/promptcompat"
)
type StreamRetryOptions struct {
@@ -19,6 +21,8 @@ type StreamRetryOptions struct {
RetryMaxAttempts int
MaxAttempts int
UsagePrompt string
Request promptcompat.StandardRequest
CurrentInputFile history.CurrentInputConfigReader
}
type StreamRetryHooks struct {
@@ -71,7 +75,7 @@ func ExecuteStreamWithRetry(ctx context.Context, ds DeepSeekCaller, a *auth.Requ
if attempts >= retryMax {
if canRetryOnAlternateAccount(ctx, a, &assistantturn.OutputError{Status: http.StatusTooManyRequests}, opts.RetryEnabled, &accountSwitchAttempted) {
switched, switchErr := startPayloadCompletionOnAlternateAccount(ctx, ds, a, payload, maxAttempts)
switched, switchErr := startPayloadCompletionOnAlternateAccount(ctx, ds, a, payload, opts, maxAttempts)
if switchErr != nil {
if hooks.OnRetryFailure != nil {
hooks.OnRetryFailure(switchErr.Status, switchErr.Message, switchErr.Code)
@@ -142,7 +146,7 @@ func ExecuteStreamWithRetry(ctx context.Context, ds DeepSeekCaller, a *auth.Requ
}
}
func startPayloadCompletionOnAlternateAccount(ctx context.Context, ds DeepSeekCaller, a *auth.RequestAuth, payload map[string]any, maxAttempts int) (StartResult, *assistantturn.OutputError) {
func startPayloadCompletionOnAlternateAccount(ctx context.Context, ds DeepSeekCaller, a *auth.RequestAuth, payload map[string]any, opts StreamRetryOptions, maxAttempts int) (StartResult, *assistantturn.OutputError) {
sessionID, err := ds.CreateSession(ctx, a, maxAttempts)
if err != nil {
return StartResult{}, authOutputError(a)
@@ -152,6 +156,13 @@ func startPayloadCompletionOnAlternateAccount(ctx context.Context, ds DeepSeekCa
return StartResult{SessionID: sessionID}, &assistantturn.OutputError{Status: http.StatusUnauthorized, Message: "Failed to get PoW (invalid token or unknown error).", Code: "error"}
}
nextPayload := clonePayload(payload)
if opts.CurrentInputFile != nil && opts.Request.CurrentInputFileApplied {
stdReq, prepErr := reuploadCurrentInputFileForAccount(ctx, ds, a, opts.Request, Options{CurrentInputFile: opts.CurrentInputFile})
if prepErr != nil {
return StartResult{SessionID: sessionID}, prepErr
}
nextPayload = stdReq.CompletionPayload(sessionID)
}
nextPayload["chat_session_id"] = sessionID
delete(nextPayload, "parent_message_id")
resp, err := ds.CallCompletion(ctx, a, nextPayload, pow, maxAttempts)

View File

@@ -5,9 +5,7 @@ import (
"context"
dsprotocol "ds2api/internal/deepseek/protocol"
"encoding/json"
"errors"
"net/http"
"time"
"ds2api/internal/auth"
"ds2api/internal/config"
@@ -15,39 +13,33 @@ import (
)
func (c *Client) CallCompletion(ctx context.Context, a *auth.RequestAuth, payload map[string]any, powResp string, maxAttempts int) (*http.Response, error) {
if maxAttempts <= 0 {
maxAttempts = c.maxRetries
}
_ = maxAttempts
clients := c.requestClientsForAuth(ctx, a)
headers := c.authHeaders(a.DeepSeekToken)
headers["x-ds-pow-response"] = powResp
captureSession := c.capture.Start("deepseek_completion", dsprotocol.DeepSeekCompletionURL, a.AccountID, payload)
attempts := 0
for attempts < maxAttempts {
resp, err := c.streamPost(ctx, clients.stream, dsprotocol.DeepSeekCompletionURL, headers, payload)
if err != nil {
attempts++
time.Sleep(time.Second)
continue
}
if resp.StatusCode == http.StatusOK {
if captureSession != nil {
resp.Body = captureSession.WrapBody(resp.Body, resp.StatusCode)
}
resp = c.wrapCompletionWithAutoContinue(ctx, a, payload, powResp, resp)
return resp, nil
}
if captureSession != nil {
resp.Body = captureSession.WrapBody(resp.Body, resp.StatusCode)
}
_ = resp.Body.Close()
attempts++
time.Sleep(time.Second)
resp, err := c.streamPostOnce(ctx, clients.stream, dsprotocol.DeepSeekCompletionURL, headers, payload)
if err != nil {
return nil, err
}
return nil, errors.New("completion failed")
if captureSession != nil {
resp.Body = captureSession.WrapBody(resp.Body, resp.StatusCode)
}
if resp.StatusCode == http.StatusOK {
resp = c.wrapCompletionWithAutoContinue(ctx, a, payload, powResp, resp)
}
return resp, nil
}
func (c *Client) streamPost(ctx context.Context, doer trans.Doer, url string, headers map[string]string, payload any) (*http.Response, error) {
return c.streamPostWithFallback(ctx, doer, url, headers, payload, true)
}
func (c *Client) streamPostOnce(ctx context.Context, doer trans.Doer, url string, headers map[string]string, payload any) (*http.Response, error) {
return c.streamPostWithFallback(ctx, doer, url, headers, payload, false)
}
func (c *Client) streamPostWithFallback(ctx context.Context, doer trans.Doer, url string, headers map[string]string, payload any, allowFallback bool) (*http.Response, error) {
b, err := json.Marshal(payload)
if err != nil {
return nil, err
@@ -63,15 +55,18 @@ func (c *Client) streamPost(ctx context.Context, doer trans.Doer, url string, he
}
resp, err := doer.Do(req)
if err != nil {
config.Logger.Warn("[deepseek] fingerprint stream request failed, fallback to std transport", "url", url, "error", err)
req2, reqErr := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b))
if reqErr != nil {
return nil, reqErr
if allowFallback {
config.Logger.Warn("[deepseek] fingerprint stream request failed, fallback to std transport", "url", url, "error", err)
req2, reqErr := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b))
if reqErr != nil {
return nil, reqErr
}
for k, v := range headers {
req2.Header.Set(k, v)
}
return clients.fallbackS.Do(req2)
}
for k, v := range headers {
req2.Header.Set(k, v)
}
return clients.fallbackS.Do(req2)
return nil, err
}
return resp, nil
}

View File

@@ -0,0 +1,36 @@
package client
import (
"context"
"errors"
"net/http"
"testing"
"ds2api/internal/auth"
)
func TestCallCompletionDoesNotFallbackForNonIdempotentCompletion(t *testing.T) {
var fallbackCalled bool
client := &Client{
stream: doerFunc(func(*http.Request) (*http.Response, error) {
return nil, errors.New("ambiguous completion write failure")
}),
fallbackS: &http.Client{Transport: roundTripperFunc(func(*http.Request) (*http.Response, error) {
fallbackCalled = true
return &http.Response{StatusCode: http.StatusOK}, nil
})},
}
_, err := client.CallCompletion(
context.Background(),
&auth.RequestAuth{DeepSeekToken: "token"},
map[string]any{"prompt": "hello"},
"pow",
3,
)
if err == nil {
t.Fatal("expected completion error")
}
if fallbackCalled {
t.Fatal("completion fallback should not be called for a non-idempotent request")
}
}

View File

@@ -95,11 +95,7 @@ func (c *Client) UploadFile(ctx context.Context, a *auth.RequestAuth, req Upload
resp, err := c.doUpload(ctx, clients.regular, clients.fallback, dsprotocol.DeepSeekUploadFileURL, headers, body)
if err != nil {
config.Logger.Warn("[upload_file] request error", "error", err, "account", a.AccountID, "filename", filename)
powHeader = ""
lastFailureKind = FailureUnknown
lastFailureMessage = err.Error()
attempts++
continue
return nil, err
}
if captureSession != nil {
resp.Body = captureSession.WrapBody(resp.Body, resp.StatusCode)
@@ -201,7 +197,7 @@ func escapeMultipartFilename(filename string) string {
return filename
}
func (c *Client) doUpload(ctx context.Context, doer trans.Doer, fallback trans.Doer, url string, headers map[string]string, body []byte) (*http.Response, error) {
func (c *Client) doUpload(ctx context.Context, doer trans.Doer, _ trans.Doer, url string, headers map[string]string, body []byte) (*http.Response, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if err != nil {
return nil, err
@@ -213,15 +209,7 @@ func (c *Client) doUpload(ctx context.Context, doer trans.Doer, fallback trans.D
if err == nil {
return resp, nil
}
config.Logger.Warn("[deepseek] fingerprint upload request failed, fallback to std transport", "url", url, "error", err)
req2, reqErr := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
if reqErr != nil {
return nil, reqErr
}
for k, v := range headers {
req2.Header.Set(k, v)
}
return fallback.Do(req2)
return nil, err
}
func extractUploadFileResult(resp map[string]any) *UploadFileResult {

View File

@@ -6,6 +6,7 @@ import (
"encoding/base64"
"encoding/hex"
"encoding/json"
"errors"
"io"
"net/http"
"strings"
@@ -39,6 +40,31 @@ func TestBuildUploadMultipartBodyOmitsPurposeAndIncludesFilePart(t *testing.T) {
}
}
func TestDoUploadDoesNotFallbackForNonIdempotentUpload(t *testing.T) {
var fallbackCalled bool
client := &Client{}
_, err := client.doUpload(
context.Background(),
doerFunc(func(req *http.Request) (*http.Response, error) {
_, _ = io.ReadAll(req.Body)
return nil, errors.New("ambiguous upload write failure")
}),
doerFunc(func(*http.Request) (*http.Response, error) {
fallbackCalled = true
return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader("{}"))}, nil
}),
dsprotocol.DeepSeekUploadFileURL,
map[string]string{"Content-Type": "multipart/form-data"},
[]byte("body"),
)
if err == nil {
t.Fatal("expected upload error")
}
if fallbackCalled {
t.Fatal("upload fallback should not be called for a non-idempotent request")
}
}
func TestExtractUploadFileResultSupportsNestedShapes(t *testing.T) {
got := extractUploadFileResult(map[string]any{
"data": map[string]any{

View File

@@ -145,7 +145,7 @@ func (h *Handler) handleClaudeDirectStream(w http.ResponseWriter, r *http.Reques
return
}
streamReq := start.Request
h.handleClaudeStreamRealtimeWithRetry(w, r, a, start.Response, start.Payload, start.Pow, streamReq.ResponseModel, streamReq.Messages, streamReq.Thinking, streamReq.Search, streamReq.ToolNames, streamReq.ToolsRaw, streamReq.PromptTokenText, historySession)
h.handleClaudeStreamRealtimeWithRetry(w, r, a, start.Response, start.Payload, start.Pow, streamReq, streamReq.ResponseModel, streamReq.Messages, streamReq.Thinking, streamReq.Search, streamReq.ToolNames, streamReq.ToolsRaw, streamReq.PromptTokenText, historySession)
}
func (h *Handler) proxyViaOpenAI(w http.ResponseWriter, r *http.Request, store ConfigReader) bool {
@@ -361,7 +361,7 @@ func (h *Handler) handleClaudeStreamRealtime(w http.ResponseWriter, r *http.Requ
})
}
func (h *Handler) handleClaudeStreamRealtimeWithRetry(w http.ResponseWriter, r *http.Request, a *auth.RequestAuth, resp *http.Response, payload map[string]any, pow, model string, messages []any, thinkingEnabled, searchEnabled bool, toolNames []string, toolsRaw any, promptTokenText string, historySession *responsehistory.Session) {
func (h *Handler) handleClaudeStreamRealtimeWithRetry(w http.ResponseWriter, r *http.Request, a *auth.RequestAuth, resp *http.Response, payload map[string]any, pow string, stdReq promptcompat.StandardRequest, model string, messages []any, thinkingEnabled, searchEnabled bool, toolNames []string, toolsRaw any, promptTokenText string, historySession *responsehistory.Session) {
if resp.StatusCode != http.StatusOK {
defer func() { _ = resp.Body.Close() }()
body, _ := io.ReadAll(resp.Body)
@@ -399,11 +399,13 @@ func (h *Handler) handleClaudeStreamRealtimeWithRetry(w http.ResponseWriter, r *
streamRuntime.sendMessageStart()
completionruntime.ExecuteStreamWithRetry(r.Context(), h.DS, a, resp, payload, pow, completionruntime.StreamRetryOptions{
Surface: "claude.messages",
Stream: true,
RetryEnabled: true,
MaxAttempts: 3,
UsagePrompt: promptTokenText,
Surface: "claude.messages",
Stream: true,
RetryEnabled: true,
MaxAttempts: 3,
UsagePrompt: promptTokenText,
Request: stdReq,
CurrentInputFile: h.Store,
}, completionruntime.StreamRetryHooks{
ConsumeAttempt: func(currentResp *http.Response, allowDeferEmpty bool) (bool, bool) {
return h.consumeClaudeStreamAttempt(r, currentResp, streamRuntime, thinkingEnabled, allowDeferEmpty)

View File

@@ -137,7 +137,7 @@ func (h *Handler) handleGeminiDirectStream(w http.ResponseWriter, r *http.Reques
return
}
streamReq := start.Request
h.handleStreamGenerateContentWithRetry(w, r, a, start.Response, start.Payload, start.Pow, streamReq.ResponseModel, streamReq.PromptTokenText, streamReq.Thinking, streamReq.Search, streamReq.ToolNames, streamReq.ToolsRaw, historySession)
h.handleStreamGenerateContentWithRetry(w, r, a, start.Response, start.Payload, start.Pow, streamReq, streamReq.ResponseModel, streamReq.PromptTokenText, streamReq.Thinking, streamReq.Search, streamReq.ToolNames, streamReq.ToolsRaw, historySession)
}
func (h *Handler) proxyViaOpenAI(w http.ResponseWriter, r *http.Request, stream bool) bool {

View File

@@ -12,6 +12,7 @@ import (
"ds2api/internal/auth"
"ds2api/internal/completionruntime"
dsprotocol "ds2api/internal/deepseek/protocol"
"ds2api/internal/promptcompat"
"ds2api/internal/responsehistory"
"ds2api/internal/sse"
streamengine "ds2api/internal/stream"
@@ -87,7 +88,7 @@ type geminiStreamRuntime struct {
history *responsehistory.Session
}
func (h *Handler) handleStreamGenerateContentWithRetry(w http.ResponseWriter, r *http.Request, a *auth.RequestAuth, resp *http.Response, payload map[string]any, pow, model, finalPrompt string, thinkingEnabled, searchEnabled bool, toolNames []string, toolsRaw any, historySession *responsehistory.Session) {
func (h *Handler) handleStreamGenerateContentWithRetry(w http.ResponseWriter, r *http.Request, a *auth.RequestAuth, resp *http.Response, payload map[string]any, pow string, stdReq promptcompat.StandardRequest, model, finalPrompt string, thinkingEnabled, searchEnabled bool, toolNames []string, toolsRaw any, historySession *responsehistory.Session) {
if resp.StatusCode != http.StatusOK {
defer func() { _ = resp.Body.Close() }()
body, _ := io.ReadAll(resp.Body)
@@ -108,11 +109,13 @@ func (h *Handler) handleStreamGenerateContentWithRetry(w http.ResponseWriter, r
runtime := newGeminiStreamRuntime(w, rc, canFlush, model, finalPrompt, thinkingEnabled, searchEnabled, stripReferenceMarkersEnabled(), toolNames, toolsRaw, historySession)
completionruntime.ExecuteStreamWithRetry(r.Context(), h.DS, a, resp, payload, pow, completionruntime.StreamRetryOptions{
Surface: "gemini.generate_content",
Stream: true,
RetryEnabled: true,
MaxAttempts: 3,
UsagePrompt: finalPrompt,
Surface: "gemini.generate_content",
Stream: true,
RetryEnabled: true,
MaxAttempts: 3,
UsagePrompt: finalPrompt,
Request: stdReq,
CurrentInputFile: h.Store,
}, completionruntime.StreamRetryHooks{
ConsumeAttempt: func(currentResp *http.Response, allowDeferEmpty bool) (bool, bool) {
return h.consumeGeminiStreamAttempt(r.Context(), currentResp, runtime, thinkingEnabled, allowDeferEmpty)

View File

@@ -205,6 +205,57 @@ func TestGeminiDirectAppliesCurrentInputFile(t *testing.T) {
}
}
func TestGeminiCurrentInputFileUploadsToolsSeparately(t *testing.T) {
ds := &testGeminiDS{
resp: makeGeminiUpstreamResponse(`data: {"p":"response/content","v":"ok"}`),
}
h := &Handler{
Store: testGeminiConfig{},
Auth: testGeminiAuth{},
DS: ds,
}
reqBody := `{
"contents":[{"role":"user","parts":[{"text":"run code"}]}],
"tools":[{"functionDeclarations":[{"name":"eval_javascript","description":"eval","parameters":{"type":"object","properties":{"code":{"type":"string"}}}}]}]
}`
req := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-2.5-pro:generateContent", strings.NewReader(reqBody))
req.Header.Set("Content-Type", "application/json")
rec := httptest.NewRecorder()
r := chi.NewRouter()
RegisterRoutes(r, h)
r.ServeHTTP(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)
}
historyText := string(ds.uploadCalls[0].Data)
if strings.Contains(historyText, "Description: eval") {
t.Fatalf("history transcript should not embed tool descriptions, got %q", historyText)
}
toolsText := string(ds.uploadCalls[1].Data)
if !strings.Contains(toolsText, "# DS2API_TOOLS.txt") || !strings.Contains(toolsText, "Tool: eval_javascript") || !strings.Contains(toolsText, "Description: eval") {
t.Fatalf("expected tools transcript to include Gemini tool schema, got %q", toolsText)
}
refIDs, _ := ds.payloads[0]["ref_file_ids"].([]any)
if len(refIDs) < 2 || refIDs[0] != "file-gemini-history" || refIDs[1] != "file-gemini-tools" {
t.Fatalf("expected history and tools ref ids first, got %#v", ds.payloads[0]["ref_file_ids"])
}
prompt, _ := ds.payloads[0]["prompt"].(string)
if !strings.Contains(prompt, "DS2API_TOOLS.txt") || !strings.Contains(prompt, "TOOL CALL FORMAT") {
t.Fatalf("expected live prompt to reference tools file and retain format instructions, got %q", prompt)
}
if strings.Contains(prompt, "Description: eval") {
t.Fatalf("live prompt should not inline tool descriptions, got %q", prompt)
}
}
func TestGeminiRoutesRegistered(t *testing.T) {
h := &Handler{
Store: testGeminiConfig{},

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 == "" {

View File

@@ -99,6 +99,8 @@ func (s Service) ApplyCurrentInputFile(ctx context.Context, a *auth.RequestAuth,
stdReq.Messages = messages
stdReq.HistoryText = fileText
stdReq.CurrentInputFileApplied = true
stdReq.CurrentInputFileID = fileID
stdReq.CurrentToolsFileID = toolFileID
stdReq.RefFileIDs = prependUniqueRefFileIDs(stdReq.RefFileIDs, fileID, toolFileID)
stdReq.FinalPrompt, stdReq.ToolNames = promptcompat.BuildOpenAIPromptWithToolInstructionsOnly(messages, stdReq.ToolsRaw, "", stdReq.ToolChoice, stdReq.Thinking)
// Token accounting must reflect the actual downstream context:
@@ -112,6 +114,58 @@ func (s Service) ApplyCurrentInputFile(ctx context.Context, a *auth.RequestAuth,
return stdReq, nil
}
func (s Service) ReuploadAppliedCurrentInputFile(ctx context.Context, a *auth.RequestAuth, stdReq promptcompat.StandardRequest) (promptcompat.StandardRequest, error) {
if !stdReq.CurrentInputFileApplied || s.DS == nil || a == nil {
return stdReq, nil
}
fileText := strings.TrimSpace(stdReq.HistoryText)
if fileText == "" {
return stdReq, nil
}
modelType := "default"
if resolvedType, ok := config.GetModelType(stdReq.ResolvedModel); ok {
modelType = resolvedType
}
result, err := s.DS.UploadFile(ctx, a, dsclient.UploadFileRequest{
Filename: currentInputFilename,
ContentType: currentInputContentType,
Purpose: currentInputPurpose,
ModelType: modelType,
Data: []byte(stdReq.HistoryText),
}, 3)
if err != nil {
return stdReq, fmt.Errorf("upload current user input file: %w", err)
}
fileID := strings.TrimSpace(result.ID)
if fileID == "" {
return stdReq, errors.New("upload current user input file returned empty file id")
}
toolsText, _ := promptcompat.BuildOpenAIToolsContextTranscript(stdReq.ToolsRaw, stdReq.ToolChoice)
toolFileID := ""
if strings.TrimSpace(toolsText) != "" {
result, err := s.DS.UploadFile(ctx, a, dsclient.UploadFileRequest{
Filename: currentToolsFilename,
ContentType: currentInputContentType,
Purpose: currentInputPurpose,
ModelType: modelType,
Data: []byte(toolsText),
}, 3)
if err != nil {
return stdReq, fmt.Errorf("upload current tools file: %w", err)
}
toolFileID = strings.TrimSpace(result.ID)
if toolFileID == "" {
return stdReq, errors.New("upload current tools file returned empty file id")
}
}
stdReq.RefFileIDs = replaceGeneratedCurrentInputRefs(stdReq.RefFileIDs, stdReq.CurrentInputFileID, stdReq.CurrentToolsFileID, fileID, toolFileID)
stdReq.CurrentInputFileID = fileID
stdReq.CurrentToolsFileID = toolFileID
return stdReq, nil
}
func latestUserInputForFile(messages []any) (int, string) {
for i := len(messages) - 1; i >= 0; i-- {
msg, ok := messages[i].(map[string]any)
@@ -168,3 +222,25 @@ func prependUniqueRefFileIDs(existing []string, fileIDs ...string) []string {
}
return out
}
func replaceGeneratedCurrentInputRefs(existing []string, oldHistoryID, oldToolsID, newHistoryID, newToolsID string) []string {
filtered := make([]string, 0, len(existing))
old := map[string]struct{}{}
for _, id := range []string{oldHistoryID, oldToolsID} {
trimmed := strings.ToLower(strings.TrimSpace(id))
if trimmed != "" {
old[trimmed] = struct{}{}
}
}
for _, id := range existing {
trimmed := strings.TrimSpace(id)
if trimmed == "" {
continue
}
if _, ok := old[strings.ToLower(trimmed)]; ok {
continue
}
filtered = append(filtered, trimmed)
}
return prependUniqueRefFileIDs(filtered, newHistoryID, newToolsID)
}

View File

@@ -610,6 +610,69 @@ func TestResponsesCurrentInputFileUploadsContextAndKeepsNeutralPrompt(t *testing
}
}
func TestResponsesCurrentInputFileUploadsToolsSeparately(t *testing.T) {
ds := &inlineUploadDSStub{}
h := &openAITestSurface{
Store: mockOpenAIConfig{
currentInputEnabled: true,
},
Auth: streamStatusAuthStub{},
DS: ds,
}
r := chi.NewRouter()
registerOpenAITestRoutes(r, h)
reqBody, _ := json.Marshal(map[string]any{
"model": "deepseek-v4-flash",
"messages": historySplitTestMessages(),
"tools": []any{
map[string]any{
"type": "function",
"function": map[string]any{
"name": "search",
"description": "search docs",
"parameters": map[string]any{"type": "object"},
},
},
},
"stream": false,
})
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(string(reqBody)))
req.Header.Set("Authorization", "Bearer direct-token")
req.Header.Set("Content-Type", "application/json")
rec := httptest.NewRecorder()
r.ServeHTTP(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)
}
historyText := string(ds.uploadCalls[0].Data)
if strings.Contains(historyText, "Description: search docs") {
t.Fatalf("history transcript should not embed tool descriptions, got %q", historyText)
}
toolsText := string(ds.uploadCalls[1].Data)
if !strings.Contains(toolsText, "# DS2API_TOOLS.txt") || !strings.Contains(toolsText, "Tool: search") || !strings.Contains(toolsText, "Description: search docs") {
t.Fatalf("expected tools transcript to include schema, got %q", toolsText)
}
promptText, _ := ds.completionReq["prompt"].(string)
if !strings.Contains(promptText, "DS2API_TOOLS.txt") || !strings.Contains(promptText, "TOOL CALL FORMAT") {
t.Fatalf("expected live prompt to reference tools file and retain format instructions, got %q", promptText)
}
if strings.Contains(promptText, "Description: search docs") {
t.Fatalf("live prompt should not inline tool descriptions, got %q", promptText)
}
refIDs, _ := ds.completionReq["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", ds.completionReq["ref_file_ids"])
}
}
func TestChatCompletionsCurrentInputFileMapsManagedAuthFailureTo401(t *testing.T) {
ds := &inlineUploadDSStub{
uploadErr: &dsclient.RequestFailure{Op: "upload file", Kind: dsclient.FailureManagedUnauthorized, Message: "expired token"},

View File

@@ -26,6 +26,15 @@ func TestSanitizeLeakedOutputRemovesStandaloneMetaMarkers(t *testing.T) {
}
}
func TestSanitizeLeakedOutputRemovesFullwidthDelimitedMetaMarkers(t *testing.T) {
fw := "\uff5c"
raw := "A<" + fw + "end▁of▁sentence" + fw + ">B<" + fw + " Assistant " + fw + ">C<" + fw + "end_of_toolresults" + fw + ">D"
got := sanitizeLeakedOutput(raw)
if got != "ABCD" {
t.Fatalf("unexpected sanitize result for fullwidth-delimited meta markers: %q", got)
}
}
func TestSanitizeLeakedOutputRemovesThinkAndBosMarkers(t *testing.T) {
raw := "A<think>B</think>C<|begin▁of▁sentence|>D<| begin_of_sentence |>E<|begin_of_sentence|>F"
got := sanitizeLeakedOutput(raw)
@@ -42,6 +51,15 @@ func TestSanitizeLeakedOutputRemovesThoughtMarkers(t *testing.T) {
}
}
func TestSanitizeLeakedOutputRemovesFullwidthDelimitedBosAndThoughtMarkers(t *testing.T) {
fw := "\uff5c"
raw := "A<" + fw + "begin▁of▁sentence" + fw + ">B<" + fw + "▁of▁thought" + fw + ">C<" + fw + " begin_of_thought " + fw + ">D"
got := sanitizeLeakedOutput(raw)
if got != "ABCD" {
t.Fatalf("unexpected sanitize result for fullwidth-delimited BOS/thought markers: %q", got)
}
}
func TestSanitizeLeakedOutputRemovesDanglingThinkBlock(t *testing.T) {
raw := "Answer prefix<think>internal reasoning that never closes"
got := sanitizeLeakedOutput(raw)

View File

@@ -15,7 +15,7 @@ import (
streamengine "ds2api/internal/stream"
)
func (h *Handler) handleResponsesStreamWithRetry(w http.ResponseWriter, r *http.Request, a *auth.RequestAuth, resp *http.Response, payload map[string]any, pow, owner, responseID, model, finalPrompt string, refFileTokens int, thinkingEnabled, searchEnabled bool, toolNames []string, toolsRaw any, toolChoice promptcompat.ToolChoicePolicy, traceID string, historySession *responsehistory.Session) {
func (h *Handler) handleResponsesStreamWithRetry(w http.ResponseWriter, r *http.Request, a *auth.RequestAuth, resp *http.Response, payload map[string]any, pow, owner, responseID string, stdReq promptcompat.StandardRequest, model, finalPrompt string, refFileTokens int, thinkingEnabled, searchEnabled bool, toolNames []string, toolsRaw any, toolChoice promptcompat.ToolChoicePolicy, traceID string, historySession *responsehistory.Session) {
streamRuntime, initialType, ok := h.prepareResponsesStreamRuntime(w, resp, owner, responseID, model, finalPrompt, refFileTokens, thinkingEnabled, searchEnabled, toolNames, toolsRaw, toolChoice, traceID, historySession)
if !ok {
return
@@ -27,6 +27,8 @@ func (h *Handler) handleResponsesStreamWithRetry(w http.ResponseWriter, r *http.
RetryMaxAttempts: emptyOutputRetryMaxAttempts(),
MaxAttempts: 3,
UsagePrompt: finalPrompt,
Request: stdReq,
CurrentInputFile: h.Store,
}, completionruntime.StreamRetryHooks{
ConsumeAttempt: func(currentResp *http.Response, allowDeferEmpty bool) (bool, bool) {
return h.consumeResponsesStreamAttempt(r, currentResp, streamRuntime, initialType, thinkingEnabled, allowDeferEmpty)

View File

@@ -138,7 +138,7 @@ func (h *Handler) Responses(w http.ResponseWriter, r *http.Request) {
streamReq := start.Request
refFileTokens := streamReq.RefFileTokens
h.handleResponsesStreamWithRetry(w, r, a, start.Response, start.Payload, start.Pow, owner, responseID, streamReq.ResponseModel, streamReq.PromptTokenText, refFileTokens, streamReq.Thinking, streamReq.Search, streamReq.ToolNames, streamReq.ToolsRaw, streamReq.ToolChoice, traceID, historySession)
h.handleResponsesStreamWithRetry(w, r, a, start.Response, start.Payload, start.Pow, owner, responseID, streamReq, streamReq.ResponseModel, streamReq.PromptTokenText, refFileTokens, streamReq.Thinking, streamReq.Search, streamReq.ToolNames, streamReq.ToolsRaw, streamReq.ToolChoice, traceID, historySession)
}
func (h *Handler) handleResponsesNonStream(w http.ResponseWriter, resp *http.Response, owner, responseID, model, finalPrompt string, refFileTokens int, thinkingEnabled, searchEnabled bool, toolNames []string, toolsRaw any, toolChoice promptcompat.ToolChoicePolicy, traceID string) {

View File

@@ -13,21 +13,23 @@ var leakedToolResultBlobPattern = regexp.MustCompile(`(?is)<\s*\|\s*tool\s*\|\s*
var leakedThinkTagPattern = regexp.MustCompile(`(?is)</?\s*think\s*>`)
// leakedBOSMarkerPattern matches DeepSeek BOS markers in BOTH forms:
// leakedBOSMarkerPattern matches DeepSeek BOS markers with halfwidth or
// legacy U+FF5C fullwidth delimiters:
// - ASCII underscore: <|begin_of_sentence|>
// - U+2581 variant: <|begin▁of▁sentence|>
var leakedBOSMarkerPattern = regexp.MustCompile(`(?i)<[|\|]\s*begin[_▁]of[_▁]sentence\s*[|\|]>`)
var leakedBOSMarkerPattern = regexp.MustCompile(`(?i)<[\|\x{ff5c}]\s*begin[_▁]of[_▁]sentence\s*[\|\x{ff5c}]>`)
// leakedThoughtMarkerPattern matches leaked thought control markers in both
// explicit and compact forms:
// - ASCII underscore: <| of_thought |>, <| begin_of_thought |>
// - U+2581 variant: <|▁of▁thought|>, <|begin▁of▁thought|>
var leakedThoughtMarkerPattern = regexp.MustCompile(`(?i)<[|\|]\s*(?:begin[_▁])?[_▁]*of[_▁]thought\s*[|\|]>`)
var leakedThoughtMarkerPattern = regexp.MustCompile(`(?i)<[\|\x{ff5c}]\s*(?:begin[_▁])?[_▁]*of[_▁]thought\s*[\|\x{ff5c}]>`)
// leakedMetaMarkerPattern matches the remaining DeepSeek special tokens in BOTH forms:
// leakedMetaMarkerPattern matches the remaining DeepSeek special tokens with
// halfwidth or legacy U+FF5C fullwidth delimiters:
// - ASCII underscore: <|end_of_sentence|>, <|end_of_toolresults|>, <|end_of_instructions|>
// - U+2581 variant: <|end▁of▁sentence|>, <|end▁of▁toolresults|>, <|end▁of▁instructions|>
var leakedMetaMarkerPattern = regexp.MustCompile(`(?i)<[|\|]\s*(?:assistant|tool|end[_▁]of[_▁]sentence|end[_▁]of[_▁]thinking|end[_▁]of[_▁]thought|end[_▁]of[_▁]toolresults|end[_▁]of[_▁]instructions)\s*[|\|]>`)
var leakedMetaMarkerPattern = regexp.MustCompile(`(?i)<[\|\x{ff5c}]\s*(?:assistant|tool|end[_▁]of[_▁]sentence|end[_▁]of[_▁]thinking|end[_▁]of[_▁]thought|end[_▁]of[_▁]toolresults|end[_▁]of[_▁]instructions)\s*[\|\x{ff5c}]>`)
// leakedAgentXMLBlockPatterns catch agent-style XML blocks that leak through
// when the sieve fails to capture them. These are applied only to complete

View File

@@ -85,6 +85,33 @@ async function fetchStreamPow(req, leaseID) {
};
}
async function fetchStreamSwitch(req, leaseID) {
const url = buildInternalGoURL(req);
url.searchParams.set('__stream_switch', '1');
const upstream = await fetch(url.toString(), {
method: 'POST',
headers: buildInternalGoHeaders(req, { withInternalToken: true, withContentType: true }),
body: Buffer.from(JSON.stringify({ lease_id: leaseID })),
});
const text = await upstream.text();
let body = {};
try {
body = JSON.parse(text || '{}');
} catch (_err) {
body = {};
}
return {
ok: upstream.ok,
status: upstream.status,
contentType: upstream.headers.get('content-type') || 'application/json',
text,
body,
};
}
function relayPreparedFailure(res, prep) {
if (prep.status === 401 && looksLikeVercelAuthPage(prep.text)) {
writeOpenAIError(
@@ -223,6 +250,7 @@ module.exports = {
readRawBody,
fetchStreamPrepare,
fetchStreamPow,
fetchStreamSwitch,
relayPreparedFailure,
safeReadText,
buildInternalGoURL,

View File

@@ -7,9 +7,9 @@ const {
SKIP_EXACT_PATHS,
} = require('../shared/deepseek-constants');
const LEAKED_BOS_MARKER_PATTERN = /<[||]\s*begin[_▁]of[_▁]sentence\s*[||]>/gi;
const LEAKED_THOUGHT_MARKER_PATTERN = /<[||]\s*(?:begin[_▁])?[_▁]*of[_▁]thought\s*[||]>/gi;
const LEAKED_META_MARKER_PATTERN = /<[||]\s*(?:assistant|tool|end[_▁]of[_▁]sentence|end[_▁]of[_▁]thinking|end[_▁]of[_▁]thought|end[_▁]of[_▁]toolresults|end[_▁]of[_▁]instructions)\s*[||]>/gi;
const LEAKED_BOS_MARKER_PATTERN = /<[\|\uFF5C]\s*begin[_▁]of[_▁]sentence\s*[\|\uFF5C]>/gi;
const LEAKED_THOUGHT_MARKER_PATTERN = /<[\|\uFF5C]\s*(?:begin[_▁])?[_▁]*of[_▁]thought\s*[\|\uFF5C]>/gi;
const LEAKED_META_MARKER_PATTERN = /<[\|\uFF5C]\s*(?:assistant|tool|end[_▁]of[_▁]sentence|end[_▁]of[_▁]thinking|end[_▁]of[_▁]thought|end[_▁]of[_▁]toolresults|end[_▁]of[_▁]instructions)\s*[\|\uFF5C]>/gi;

View File

@@ -25,6 +25,7 @@ const {
isAbortError,
fetchStreamPrepare,
fetchStreamPow,
fetchStreamSwitch,
relayPreparedFailure,
createLeaseReleaser,
} = require('./http_internal');
@@ -46,11 +47,11 @@ async function handleVercelStream(req, res, rawBody, payload) {
}
const model = asString(prep.body.model) || asString(payload.model);
const sessionID = asString(prep.body.session_id) || `chatcmpl-${Date.now()}`;
const responseID = asString(prep.body.session_id) || `chatcmpl-${Date.now()}`;
const leaseID = asString(prep.body.lease_id);
const deepseekToken = asString(prep.body.deepseek_token);
let deepseekToken = asString(prep.body.deepseek_token);
const initialPowHeader = asString(prep.body.pow_header);
const completionPayload = prep.body.payload && typeof prep.body.payload === 'object' ? prep.body.payload : null;
let completionPayload = prep.body.payload && typeof prep.body.payload === 'object' ? prep.body.payload : null;
const finalPrompt = asString(prep.body.final_prompt);
const thinkingEnabled = toBool(prep.body.thinking_enabled);
const searchEnabled = toBool(prep.body.search_enabled);
@@ -133,13 +134,14 @@ async function handleVercelStream(req, res, rawBody, payload) {
}
};
const fetchCompletion = (bodyPayload) => fetchDeepSeekStream(DEEPSEEK_COMPLETION_URL, bodyPayload, currentPowHeader);
let activeDeepSeekSessionID = responseID;
const fetchContinue = async (messageID) => {
const powHeader = await refreshPowHeader('continue');
if (!powHeader) {
return null;
}
return fetchDeepSeekStream(DEEPSEEK_CONTINUE_URL, {
chat_session_id: sessionID,
chat_session_id: activeDeepSeekSessionID,
message_id: messageID,
fallback_to_resume: true,
}, powHeader);
@@ -185,7 +187,7 @@ async function handleVercelStream(req, res, rawBody, payload) {
let ended = false;
const { sendFrame, sendDeltaFrame } = createChatCompletionEmitter({
res,
sessionID,
sessionID: responseID,
created,
model,
isClosed: () => clientClosed,
@@ -242,7 +244,7 @@ async function handleVercelStream(req, res, rawBody, payload) {
}
ended = true;
sendFrame({
id: sessionID,
id: responseID,
object: 'chat.completion.chunk',
created,
model,
@@ -261,7 +263,7 @@ async function handleVercelStream(req, res, rawBody, payload) {
const processStream = async (initialResponse, allowDeferEmpty) => {
let currentResponse = initialResponse;
let continueState = createContinueState(sessionID);
let continueState = createContinueState(activeDeepSeekSessionID);
let continueRounds = 0;
// eslint-disable-next-line no-constant-condition
while (true) {
@@ -412,13 +414,39 @@ async function handleVercelStream(req, res, rawBody, payload) {
};
let retryAttempts = 0;
let accountSwitchAttempted = false;
// eslint-disable-next-line no-constant-condition
while (true) {
const processed = await processStream(completionRes, retryAttempts < EMPTY_OUTPUT_RETRY_MAX_ATTEMPTS);
const allowDeferEmpty = retryAttempts < EMPTY_OUTPUT_RETRY_MAX_ATTEMPTS || !accountSwitchAttempted;
const processed = await processStream(completionRes, allowDeferEmpty);
if (processed.terminal) {
return;
}
if (!processed.retryable || retryAttempts >= EMPTY_OUTPUT_RETRY_MAX_ATTEMPTS) {
if (!processed.retryable) {
await finish('stop');
return;
}
if (retryAttempts >= EMPTY_OUTPUT_RETRY_MAX_ATTEMPTS) {
if (!accountSwitchAttempted) {
accountSwitchAttempted = true;
const switched = await fetchStreamSwitch(req, leaseID);
if (switched.ok && switched.body && switched.body.payload && typeof switched.body.payload === 'object') {
completionPayload = switched.body.payload;
deepseekToken = asString(switched.body.deepseek_token) || deepseekToken;
currentPowHeader = asString(switched.body.pow_header) || currentPowHeader;
activeDeepSeekSessionID = asString(switched.body.session_id) || activeDeepSeekSessionID;
usagePrompt = finalPrompt;
completionRes = await fetchCompletion(completionPayload);
if (completionRes === null) {
return;
}
if (!completionRes.ok || !completionRes.body) {
await finish('stop');
return;
}
continue;
}
}
await finish('stop');
return;
}

View File

@@ -113,6 +113,9 @@ func TestBuildOpenAIPromptWithToolInstructionsOnlyOmitsSchemas(t *testing.T) {
if strings.Contains(finalPrompt, "You have access to these tools") || strings.Contains(finalPrompt, "Description: search docs") || strings.Contains(finalPrompt, "Parameters:") {
t.Fatalf("tool descriptions should be externalized, got: %q", finalPrompt)
}
if !strings.Contains(finalPrompt, "Treat DS2API_TOOLS.txt as the authoritative list of callable tools and schemas") {
t.Fatalf("expected instructions-only prompt to point model at tools file, got: %q", finalPrompt)
}
if !strings.Contains(finalPrompt, "TOOL CALL FORMAT") || !strings.Contains(finalPrompt, "Remember: The ONLY valid way to use tools") {
t.Fatalf("expected tool format instructions to remain in live prompt, got: %q", finalPrompt)
}

View File

@@ -11,6 +11,8 @@ type StandardRequest struct {
HistoryText string
PromptTokenText string
CurrentInputFileApplied bool
CurrentInputFileID string
CurrentToolsFileID string
ToolsRaw any
FinalPrompt string
ToolNames []string

View File

@@ -39,6 +39,8 @@ func injectToolPromptWithDescriptions(messages []map[string]any, tools []any, po
toolPrompt := parts.Instructions
if includeDescriptions && parts.Descriptions != "" {
toolPrompt = parts.Descriptions + "\n\n" + toolPrompt
} else if !includeDescriptions && parts.Descriptions != "" {
toolPrompt = "Available tool descriptions and parameter schemas are attached in DS2API_TOOLS.txt. Treat DS2API_TOOLS.txt as the authoritative list of callable tools and schemas; use only tools and parameters listed there.\n\n" + toolPrompt
}
for i := range messages {