lisone/ds2api
1
0
Fork 0
mirror of https://github.com/CJackHwang/ds2api.git synced 2026-05-13 20:57:41 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: implement managed-account rotation on 429 empty-output completion retries

Browse Source
This commit is contained in:
CJACK
2026-05-10 00:41:45 +08:00
parent 3cc7f469f3
commit ddd42e532e
15 changed files with 362 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
internal/completionruntime/nonstream.go
View File

@@ -104,6 +104,7 @@ func ExecuteNonStreamStartedWithRetry(ctx context.Context, ds DeepSeekCaller, a
pow := start.Pow
attempts := 0
accountSwitchAttempted := false
currentResp := start.Response
usagePrompt := stdReq.PromptTokenText
accumulatedThinking := ""
@@ -112,6 +113,24 @@ func ExecuteNonStreamStartedWithRetry(ctx context.Context, ds DeepSeekCaller, a
for {
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)
if switchErr != nil {
return NonStreamResult{SessionID: sessionID, Payload: payload, Attempts: attempts}, switchErr
}
if switched.Response != nil {
config.Logger.Info("[completion_runtime_account_switch_retry] retrying after 429", "surface", stdReq.Surface, "stream", false, "account", a.AccountID)
sessionID = switched.SessionID
payload = switched.Payload
pow = switched.Pow
currentResp = switched.Response
usagePrompt = stdReq.PromptTokenText
accumulatedThinking = ""
accumulatedRawThinking = ""
accumulatedToolDetectionThinking = ""
continue
}
}
return NonStreamResult{SessionID: sessionID, Payload: payload, Attempts: attempts}, outErr
}
accumulatedThinking += sse.TrimContinuationOverlap(accumulatedThinking, turn.Thinking)
@@ -134,6 +153,24 @@ func ExecuteNonStreamStartedWithRetry(ctx context.Context, ds DeepSeekCaller, a
retryMax = shared.EmptyOutputRetryMaxAttempts()
}
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)
if switchErr != nil {
return NonStreamResult{SessionID: sessionID, Payload: payload, Turn: turn, Attempts: attempts}, switchErr
}
if switched.Response != nil {
config.Logger.Info("[completion_runtime_account_switch_retry] retrying after 429", "surface", stdReq.Surface, "stream", false, "account", a.AccountID)
sessionID = switched.SessionID
payload = switched.Payload
pow = switched.Pow
currentResp = switched.Response
usagePrompt = stdReq.PromptTokenText
accumulatedThinking = ""
accumulatedRawThinking = ""
accumulatedToolDetectionThinking = ""
continue
}
}
return NonStreamResult{SessionID: sessionID, Payload: payload, Turn: turn, Attempts: attempts}, turn.Error
}
@@ -154,6 +191,37 @@ func ExecuteNonStreamStartedWithRetry(ctx context.Context, ds DeepSeekCaller, a
}
}
func canRetryOnAlternateAccount(ctx context.Context, a *auth.RequestAuth, outErr *assistantturn.OutputError, retryEnabled bool, attempted *bool) bool {
if outErr == nil || outErr.Status != http.StatusTooManyRequests {
return false
}
if !retryEnabled || attempted == nil || *attempted {
return false
}
if a == nil || !a.UseConfigToken {
return false
}
*attempted = true
return a.SwitchAccount(ctx)
}
func startStandardCompletionOnAlternateAccount(ctx context.Context, ds DeepSeekCaller, a *auth.RequestAuth, stdReq promptcompat.StandardRequest, maxAttempts int) (StartResult, *assistantturn.OutputError) {
sessionID, err := ds.CreateSession(ctx, a, maxAttempts)
if err != nil {
return StartResult{}, authOutputError(a)
}
pow, err := ds.GetPow(ctx, a, maxAttempts)
if err != nil {
return StartResult{SessionID: sessionID}, &assistantturn.OutputError{Status: http.StatusUnauthorized, Message: "Failed to get PoW (invalid token or unknown error).", Code: "error"}
}
payload := stdReq.CompletionPayload(sessionID)
resp, err := ds.CallCompletion(ctx, a, payload, pow, maxAttempts)
if err != nil {
return StartResult{SessionID: sessionID, Payload: payload, Pow: pow}, &assistantturn.OutputError{Status: http.StatusInternalServerError, Message: "Failed to get completion.", Code: "error"}
}
return StartResult{SessionID: sessionID, Payload: payload, Pow: pow, Response: resp, Request: stdReq}, 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 {