From 1411b7eb7ea9f4abed1aef7ec37a44696cee30b5 Mon Sep 17 00:00:00 2001
From: Achord Chan <achordchan@gmail.com>
Date: Sun, 20 Sep 2026 11:02:06 +0800
Subject: [PATCH] fix(gateway): recognize Baseten reasoning budget errors

---
 backend/internal/service/gateway_request.go   | 14 +++-
 .../service/gateway_thinking_budget_test.go   | 68 +++++++++++++++++++
 2 files changed, 81 insertions(+), 1 deletion(-)
 create mode 100644 backend/internal/service/gateway_thinking_budget_test.go

diff --git a/backend/internal/service/gateway_request.go b/backend/internal/service/gateway_request.go
index a099f99463..72f48ad40a 100644
--- a/backend/internal/service/gateway_request.go
+++ b/backend/internal/service/gateway_request.go
@@ -1607,12 +1607,16 @@ const (
 
 // isThinkingBudgetConstraintError detects whether an upstream error message indicates
 // a budget_tokens constraint violation (e.g. "budget_tokens >= 1024").
-// Matches three conditions (all must be true):
+// Also recognizes Baseten's final-answer reserve constraint.
+// For the original budget constraint, all three conditions must be true:
 //  1. Contains "budget_tokens" or "budget tokens"
 //  2. Contains "thinking"
 //  3. Contains ">= 1024" or "greater than or equal to 1024" or ("1024" + "input should be")
 func isThinkingBudgetConstraintError(errMsg string) bool {
 	m := strings.ToLower(errMsg)
+	if isFinalAnswerReserveError(m) {
+		return true
+	}
 
 	// Condition 1: budget_tokens or budget tokens
 	hasBudget := strings.Contains(m, "budget_tokens") || strings.Contains(m, "budget tokens")
@@ -1636,6 +1640,14 @@ func isThinkingBudgetConstraintError(errMsg string) bool {
 	return false
 }
 
+// isFinalAnswerReserveError matches the specific reserve constraint, rather than
+// treating arbitrary reasoning quota or context-length errors as repairable.
+func isFinalAnswerReserveError(errMsg string) bool {
+	m := strings.ToLower(errMsg)
+	return strings.Contains(m, "must be greater than 1024 to reserve tokens for a final answer") &&
+		strings.Contains(m, "baseten reasoning is enabled")
+}
+
 // RectifyThinkingBudget modifies the request body to fix budget_tokens constraint errors.
 // It sets thinking.budget_tokens = 32000, thinking.type = "enabled" (unless adaptive),
 // and ensures max_tokens >= 32001.
diff --git a/backend/internal/service/gateway_thinking_budget_test.go b/backend/internal/service/gateway_thinking_budget_test.go
new file mode 100644
index 0000000000..00fe39b1b2
--- /dev/null
+++ b/backend/internal/service/gateway_thinking_budget_test.go
@@ -0,0 +1,68 @@
+//go:build unit
+
+package service
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/tidwall/gjson"
+)
+
+const basetenBudgetError = "Error from provider (Console Go): Upstream request failed: [invalid_request_error] must be greater than 1024 to reserve tokens for a final answer when Baseten reasoning is enabled"
+
+func TestThinkingBudgetConstraintErrors(t *testing.T) {
+	for _, tc := range []struct {
+		message string
+		want    bool
+	}{
+		{basetenBudgetError, true},
+		{"MUST BE GREATER THAN 1024 TO RESERVE TOKENS FOR A FINAL ANSWER WHEN BASETEN REASONING IS ENABLED", true},
+		{"thinking.budget_tokens must be >= 1024", true},
+		{"thinking budget tokens must be greater than or equal to 1024", true},
+		{"thinking.budget_tokens: Input should be greater than 1024", true},
+		{"reasoning quota exceeded: 1024 tokens", false},
+		{"must be greater than 1024", false},
+		{"reserve tokens for a final answer when Baseten reasoning is enabled", false},
+		{"must be greater than 2048 to reserve tokens for a final answer when Baseten reasoning is enabled", false},
+		{"", false},
+	} {
+		t.Run(tc.message, func(t *testing.T) {
+			require.Equal(t, tc.want, isThinkingBudgetConstraintError(tc.message))
+		})
+	}
+}
+
+func TestRectifyThinkingBudgetFinalAnswerReserve(t *testing.T) {
+	for _, tc := range []struct {
+		maxTokens int
+		want      int64
+		changed   bool
+	}{
+		{1024, 64000, true},
+		{32000, 64000, true},
+		{32001, 32001, false},
+		{40000, 40000, false},
+		{64000, 64000, false},
+		{100000, 100000, false},
+	} {
+		t.Run(fmt.Sprint(tc.maxTokens), func(t *testing.T) {
+			body := []byte(fmt.Sprintf(`{"model":"example","max_tokens":%d,"thinking":{"type":"enabled","budget_tokens":32000},"messages":[{"role":"user","content":"hi"}]}`, tc.maxTokens))
+			require.True(t, isThinkingBudgetConstraintError(basetenBudgetError))
+			got, changed := RectifyThinkingBudget(body)
+			require.Equal(t, tc.changed, changed)
+			require.Equal(t, tc.want, gjson.GetBytes(got, "max_tokens").Int())
+			require.Equal(t, int64(32000), gjson.GetBytes(got, "thinking.budget_tokens").Int())
+			require.Equal(t, "hi", gjson.GetBytes(got, "messages.0.content").String())
+		})
+	}
+	adaptive := []byte(`{"max_tokens":1024,"thinking":{"type":"adaptive"}}`)
+	got, changed := RectifyThinkingBudget(adaptive)
+	require.False(t, changed)
+	require.Equal(t, adaptive, got)
+	legacy := []byte(`{"max_tokens":32001,"thinking":{"type":"enabled","budget_tokens":32000}}`)
+	got, changed = RectifyThinkingBudget(legacy)
+	require.False(t, changed)
+	require.Equal(t, legacy, got)
+}