From 60c6bcfe5bb9182dbc9f9b0ecfbe53bbe04285a5 Mon Sep 17 00:00:00 2001 From: dtapps Date: Sat, 28 Mar 2026 23:37:21 +0800 Subject: [PATCH] feat(config): support model-specific max_tokens and fix config key corruption --- pkg/agent/instance.go | 3 +++ pkg/config/config.go | 6 ++++++ pkg/config/config_old.go | 2 ++ pkg/config/migration.go | 1 + pkg/config/multikey_test.go | 4 ++++ pkg/gateway/gateway.go | 3 +++ web/backend/api/models.go | 2 ++ web/frontend/src/api/models.ts | 1 + .../src/components/models/add-model-sheet.tsx | 16 ++++++++++++++++ .../src/components/models/edit-model-sheet.tsx | 17 +++++++++++++++++ web/frontend/src/i18n/locales/en.json | 4 +++- web/frontend/src/i18n/locales/zh.json | 4 +++- 12 files changed, 61 insertions(+), 2 deletions(-) diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 880725660..9af38c1d0 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -131,6 +131,9 @@ func NewAgentInstance( } maxTokens := defaults.MaxTokens + if mc, err := cfg.GetModelConfig(defaults.OriginalModelName); err == nil { + maxTokens = mc.MaxTokens + } if maxTokens == 0 { maxTokens = 8192 } diff --git a/pkg/config/config.go b/pkg/config/config.go index 533f45a44..d37e1a973 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -327,6 +327,9 @@ type AgentDefaults struct { SubTurn SubTurnConfig `json:"subturn" envPrefix:"PICOCLAW_AGENTS_DEFAULTS_SUBTURN_"` ToolFeedback ToolFeedbackConfig `json:"tool_feedback,omitempty"` SplitOnMarker bool `json:"split_on_marker" env:"PICOCLAW_AGENTS_DEFAULTS_SPLIT_ON_MARKER"` // split messages on <|[SPLIT]|> marker + // OriginalModelName holds the original model name as specified in the config. + // This is used to resolve the model name in the model_list configuration. + OriginalModelName string `json:"-"` } const DefaultMaxMediaSize = 20 * 1024 * 1024 // 20 MB @@ -668,6 +671,7 @@ type ModelConfig struct { // Optional optimizations RPM int `json:"rpm,omitempty"` // Requests per minute limit + MaxTokens int `json:"max_tokens,omitempty"` // Maximum number of tokens per request MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens") RequestTimeout int `json:"request_timeout,omitempty"` ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive @@ -1287,6 +1291,7 @@ func expandMultiKeyModels(models []*ModelConfig) []*ModelConfig { ConnectMode: m.ConnectMode, Workspace: m.Workspace, RPM: m.RPM, + MaxTokens: m.MaxTokens, MaxTokensField: m.MaxTokensField, RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, @@ -1307,6 +1312,7 @@ func expandMultiKeyModels(models []*ModelConfig) []*ModelConfig { ConnectMode: m.ConnectMode, Workspace: m.Workspace, RPM: m.RPM, + MaxTokens: m.MaxTokens, MaxTokensField: m.MaxTokensField, RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, diff --git a/pkg/config/config_old.go b/pkg/config/config_old.go index fd54c9e08..b2ad686ba 100644 --- a/pkg/config/config_old.go +++ b/pkg/config/config_old.go @@ -660,6 +660,7 @@ type modelConfigV0 struct { // Optional optimizations RPM int `json:"rpm,omitempty"` // Requests per minute limit + MaxTokens int `json:"max_tokens,omitempty"` // Maximum number of tokens per request MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens") RequestTimeout int `json:"request_timeout,omitempty"` ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive @@ -744,6 +745,7 @@ func (c *configV0) Migrate() (*Config, error) { ConnectMode: m.ConnectMode, Workspace: m.Workspace, RPM: m.RPM, + MaxTokens: m.MaxTokens, MaxTokensField: m.MaxTokensField, RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, diff --git a/pkg/config/migration.go b/pkg/config/migration.go index fee800a76..075fa408e 100644 --- a/pkg/config/migration.go +++ b/pkg/config/migration.go @@ -499,6 +499,7 @@ func loadConfigV0(data []byte) (migratable, error) { ConnectMode: m.ConnectMode, Workspace: m.Workspace, RPM: m.RPM, + MaxTokens: m.MaxTokens, MaxTokensField: m.MaxTokensField, RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, diff --git a/pkg/config/multikey_test.go b/pkg/config/multikey_test.go index e58c6dc9e..0a4e56fa3 100644 --- a/pkg/config/multikey_test.go +++ b/pkg/config/multikey_test.go @@ -192,6 +192,7 @@ func TestExpandMultiKeyModels_PreservesOtherFields(t *testing.T) { APIBase: "https://api.example.com", Proxy: "http://proxy:8080", RPM: 60, + MaxTokens: 1024, MaxTokensField: "max_completion_tokens", RequestTimeout: 30, ThinkingLevel: "high", @@ -212,6 +213,9 @@ func TestExpandMultiKeyModels_PreservesOtherFields(t *testing.T) { if primary.RPM != 60 { t.Errorf("expected rpm preserved, got %d", primary.RPM) } + if primary.MaxTokens != 1024 { + t.Errorf("expected max_tokens preserved, got %d", primary.MaxTokens) + } if primary.MaxTokensField != "max_completion_tokens" { t.Errorf("expected max_tokens_field preserved, got %q", primary.MaxTokensField) } diff --git a/pkg/gateway/gateway.go b/pkg/gateway/gateway.go index c35b3e744..b003d0733 100644 --- a/pkg/gateway/gateway.go +++ b/pkg/gateway/gateway.go @@ -113,6 +113,9 @@ func Run(debug bool, homePath, configPath string, allowEmptyStartup bool) error } if modelID != "" { + if cfg.Agents.Defaults.OriginalModelName == "" { + cfg.Agents.Defaults.OriginalModelName = cfg.Agents.Defaults.ModelName + } cfg.Agents.Defaults.ModelName = modelID } diff --git a/web/backend/api/models.go b/web/backend/api/models.go index 38a55948b..9c5d57082 100644 --- a/web/backend/api/models.go +++ b/web/backend/api/models.go @@ -35,6 +35,7 @@ type modelResponse struct { ConnectMode string `json:"connect_mode,omitempty"` Workspace string `json:"workspace,omitempty"` RPM int `json:"rpm,omitempty"` + MaxTokens int `json:"max_tokens,omitempty"` MaxTokensField string `json:"max_tokens_field,omitempty"` RequestTimeout int `json:"request_timeout,omitempty"` ThinkingLevel string `json:"thinking_level,omitempty"` @@ -81,6 +82,7 @@ func (h *Handler) handleListModels(w http.ResponseWriter, r *http.Request) { ConnectMode: m.ConnectMode, Workspace: m.Workspace, RPM: m.RPM, + MaxTokens: m.MaxTokens, MaxTokensField: m.MaxTokensField, RequestTimeout: m.RequestTimeout, ThinkingLevel: m.ThinkingLevel, diff --git a/web/frontend/src/api/models.ts b/web/frontend/src/api/models.ts index aa66a7389..ca7f1d54e 100644 --- a/web/frontend/src/api/models.ts +++ b/web/frontend/src/api/models.ts @@ -14,6 +14,7 @@ export interface ModelInfo { connect_mode?: string workspace?: string rpm?: number + max_tokens?: number max_tokens_field?: string request_timeout?: number thinking_level?: string diff --git a/web/frontend/src/components/models/add-model-sheet.tsx b/web/frontend/src/components/models/add-model-sheet.tsx index de9481391..0e338d835 100644 --- a/web/frontend/src/components/models/add-model-sheet.tsx +++ b/web/frontend/src/components/models/add-model-sheet.tsx @@ -32,6 +32,7 @@ interface AddForm { connectMode: string workspace: string rpm: string + maxTokens: string maxTokensField: string requestTimeout: string thinkingLevel: string @@ -48,6 +49,7 @@ const EMPTY_ADD_FORM: AddForm = { connectMode: "", workspace: "", rpm: "", + maxTokens: "", maxTokensField: "", requestTimeout: "", thinkingLevel: "", @@ -128,6 +130,7 @@ export function AddModelSheet({ connect_mode: form.connectMode.trim() || undefined, workspace: form.workspace.trim() || undefined, rpm: form.rpm ? Number(form.rpm) : undefined, + max_tokens: form.maxTokens ? Number(form.maxTokens) : undefined, max_tokens_field: form.maxTokensField.trim() || undefined, request_timeout: form.requestTimeout ? Number(form.requestTimeout) @@ -302,6 +305,19 @@ export function AddModelSheet({ /> + + + + + + + +