From 790ec307f9faebf664d92063adb234b0909970f9 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Sat, 20 Jun 2026 12:15:06 +0800 Subject: [PATCH] feat(config): add support for `rebuild_mid_system_message` configuration - Introduced `RebuildMidSystemMessage` field in config to move system messages into the top-level Claude system field. - Updated executor to handle mid-system message rebuilding when enabled via config or auth attributes. - Added unit tests to verify rebuilding behavior and default behavior when disabled. - Updated configuration example and API handlers to support the new field. Closes: #3792 --- config.example.yaml | 1 + .../api/handlers/management/config_lists.go | 18 ++-- internal/config/config.go | 3 + internal/runtime/executor/claude_executor.go | 94 ++++++++++++++++++ .../runtime/executor/claude_executor_test.go | 97 +++++++++++++++++++ internal/runtime/executor/claude_signing.go | 8 ++ internal/watcher/diff/config_diff.go | 3 + internal/watcher/synthesizer/config.go | 3 + internal/watcher/synthesizer/config_test.go | 12 ++- 9 files changed, 228 insertions(+), 11 deletions(-) diff --git a/config.example.yaml b/config.example.yaml index c2b422955..c480b2f53 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -256,6 +256,7 @@ nonstream-keepalive-interval: 0 # - "claude-3-*" # wildcard matching prefix (e.g. claude-3-7-sonnet-20250219) # - "*-thinking" # wildcard matching suffix (e.g. claude-opus-4-5-thinking) # - "*haiku*" # wildcard matching substring (e.g. claude-3-5-haiku-20241022) +# rebuild-mid-system-message: false # optional: default is false; when true, move messages with role "system" into the top-level Claude system field # cloak: # optional: request cloaking for non-Claude-Code clients # mode: "auto" # "auto" (default): cloak only when client is not Claude Code # # "always": always apply cloaking diff --git a/internal/api/handlers/management/config_lists.go b/internal/api/handlers/management/config_lists.go index d9050b4c8..fb4c67d21 100644 --- a/internal/api/handlers/management/config_lists.go +++ b/internal/api/handlers/management/config_lists.go @@ -307,13 +307,14 @@ func (h *Handler) PutClaudeKeys(c *gin.Context) { } func (h *Handler) PatchClaudeKey(c *gin.Context) { type claudeKeyPatch struct { - APIKey *string `json:"api-key"` - Prefix *string `json:"prefix"` - BaseURL *string `json:"base-url"` - ProxyURL *string `json:"proxy-url"` - Models *[]config.ClaudeModel `json:"models"` - Headers *map[string]string `json:"headers"` - ExcludedModels *[]string `json:"excluded-models"` + APIKey *string `json:"api-key"` + Prefix *string `json:"prefix"` + BaseURL *string `json:"base-url"` + ProxyURL *string `json:"proxy-url"` + Models *[]config.ClaudeModel `json:"models"` + Headers *map[string]string `json:"headers"` + ExcludedModels *[]string `json:"excluded-models"` + RebuildMidSystemMessage *bool `json:"rebuild-mid-system-message"` } var body struct { Index *int `json:"index"` @@ -367,6 +368,9 @@ func (h *Handler) PatchClaudeKey(c *gin.Context) { if body.Value.ExcludedModels != nil { entry.ExcludedModels = config.NormalizeExcludedModels(*body.Value.ExcludedModels) } + if body.Value.RebuildMidSystemMessage != nil { + entry.RebuildMidSystemMessage = *body.Value.RebuildMidSystemMessage + } normalizeClaudeKey(&entry) h.cfg.ClaudeKey[targetIndex] = entry h.cfg.SanitizeClaudeKeys() diff --git a/internal/config/config.go b/internal/config/config.go index 11d8f48bd..ffb67e427 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -456,6 +456,9 @@ type ClaudeKey struct { // ExcludedModels lists model IDs that should be excluded for this provider. ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"` + // RebuildMidSystemMessage moves Claude messages with role "system" into the top-level system field. + RebuildMidSystemMessage bool `yaml:"rebuild-mid-system-message,omitempty" json:"rebuild-mid-system-message,omitempty"` + // DisableCooling disables auth/model cooldown scheduling for this credential when true. DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"` diff --git a/internal/runtime/executor/claude_executor.go b/internal/runtime/executor/claude_executor.go index fec288b89..c588f315c 100644 --- a/internal/runtime/executor/claude_executor.go +++ b/internal/runtime/executor/claude_executor.go @@ -219,6 +219,9 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r if err != nil { return resp, err } + if rebuildMidSystemMessageEnabled(e.cfg, auth) { + body = rebuildMidSystemMessagesToTopLevel(body) + } // Apply cloaking (system prompt injection, fake user ID, sensitive word obfuscation) // based on client type and configuration. @@ -406,6 +409,9 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A if err != nil { return nil, err } + if rebuildMidSystemMessageEnabled(e.cfg, auth) { + body = rebuildMidSystemMessagesToTopLevel(body) + } // Apply cloaking (system prompt injection, fake user ID, sensitive word obfuscation) // based on client type and configuration. @@ -674,6 +680,9 @@ func (e *ClaudeExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut stream := from != to body := sdktranslator.TranslateRequest(from, to, baseModel, req.Payload, stream) body, _ = sjson.SetBytes(body, "model", baseModel) + if rebuildMidSystemMessageEnabled(e.cfg, auth) { + body = rebuildMidSystemMessagesToTopLevel(body) + } if !strings.HasPrefix(baseModel, "claude-3-5-haiku") { body = checkSystemInstructions(body) @@ -1141,6 +1150,91 @@ func checkSystemInstructions(payload []byte) []byte { return checkSystemInstructionsWithSigningMode(payload, false, false, false, "2.1.63", "", "") } +func rebuildMidSystemMessagesToTopLevel(payload []byte) []byte { + messages := gjson.GetBytes(payload, "messages") + if !messages.IsArray() { + return payload + } + + var movedSystemParts []string + keptMessages := make([]string, 0, int(messages.Get("#").Int())) + messages.ForEach(func(_, message gjson.Result) bool { + if strings.EqualFold(strings.TrimSpace(message.Get("role").String()), "system") { + movedSystemParts = append(movedSystemParts, claudeSystemTextParts(message.Get("content"))...) + return true + } + keptMessages = append(keptMessages, message.Raw) + return true + }) + if len(movedSystemParts) == 0 { + return payload + } + + systemParts := claudeSystemTextParts(gjson.GetBytes(payload, "system")) + systemParts = append(systemParts, movedSystemParts...) + if len(systemParts) > 0 { + if updated, errSetSystem := sjson.SetRawBytes(payload, "system", rawJSONArray(systemParts)); errSetSystem == nil { + payload = updated + } + } + if updated, errSetMessages := sjson.SetRawBytes(payload, "messages", rawJSONArray(keptMessages)); errSetMessages == nil { + payload = updated + } + return payload +} + +func claudeSystemTextParts(content gjson.Result) []string { + if !content.Exists() { + return nil + } + if content.Type == gjson.String { + text := content.String() + if strings.TrimSpace(text) == "" { + return nil + } + block := []byte(`{"type":"text","text":""}`) + block, _ = sjson.SetBytes(block, "text", text) + return []string{string(block)} + } + if !content.IsArray() { + return nil + } + + var parts []string + content.ForEach(func(_, item gjson.Result) bool { + if item.Type == gjson.String { + text := item.String() + if strings.TrimSpace(text) != "" { + block := []byte(`{"type":"text","text":""}`) + block, _ = sjson.SetBytes(block, "text", text) + parts = append(parts, string(block)) + } + return true + } + if item.IsObject() && item.Get("type").String() == "text" && strings.TrimSpace(item.Get("text").String()) != "" { + parts = append(parts, item.Raw) + } + return true + }) + return parts +} + +func rawJSONArray(items []string) []byte { + if len(items) == 0 { + return []byte("[]") + } + var builder strings.Builder + builder.WriteByte('[') + for i, item := range items { + if i > 0 { + builder.WriteByte(',') + } + builder.WriteString(item) + } + builder.WriteByte(']') + return []byte(builder.String()) +} + func isClaudeOAuthToken(apiKey string) bool { return strings.Contains(apiKey, "sk-ant-oat") } diff --git a/internal/runtime/executor/claude_executor_test.go b/internal/runtime/executor/claude_executor_test.go index be4a97190..7c0a6e763 100644 --- a/internal/runtime/executor/claude_executor_test.go +++ b/internal/runtime/executor/claude_executor_test.go @@ -2113,6 +2113,103 @@ func TestClaudeExecutor_ExperimentalCCHSigningOptInSignsFinalBody(t *testing.T) } } +func TestClaudeExecutor_RebuildMidSystemMessageDisabledByDefault(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"msg_1","type":"message","model":"claude-3-5-sonnet","role":"assistant","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":1,"output_tokens":1}}`)) + })) + defer server.Close() + + executor := NewClaudeExecutor(&config.Config{ + ClaudeKey: []config.ClaudeKey{{ + APIKey: "key-123", + BaseURL: server.URL, + }}, + }) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + payload := []byte(`{"system":[{"type":"text","text":"Top rule","cache_control":{"type":"ephemeral"}}],"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]},{"role":"system","content":"Mid rule"},{"role":"user","content":[{"type":"text","text":"continue"}]}]}`) + ctx := contextWithGinHeaders(map[string]string{"User-Agent": "claude-cli/2.1.153 (external, cli)"}) + + _, errExecute := executor.Execute(ctx, auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: payload, + }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude")}) + if errExecute != nil { + t.Fatalf("Execute() error = %v", errExecute) + } + if len(seenBody) == 0 { + t.Fatal("expected request body to be captured") + } + if got := gjson.GetBytes(seenBody, "system.0.text").String(); got != "Top rule" { + t.Fatalf("system.0.text = %q, want top-level system preserved", got) + } + if got := gjson.GetBytes(seenBody, `messages.#(role=="system").content`).String(); got != "Mid rule" { + t.Fatalf("mid system message = %q, want original message preserved", got) + } +} + +func TestClaudeExecutor_RebuildMidSystemMessageOptInMovesSystemMessages(t *testing.T) { + var seenBody []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seenBody = bytes.Clone(body) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"msg_1","type":"message","model":"claude-3-5-sonnet","role":"assistant","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":1,"output_tokens":1}}`)) + })) + defer server.Close() + + executor := NewClaudeExecutor(&config.Config{ + ClaudeKey: []config.ClaudeKey{{ + APIKey: "key-123", + BaseURL: server.URL, + RebuildMidSystemMessage: true, + }}, + }) + auth := &cliproxyauth.Auth{Attributes: map[string]string{ + "api_key": "key-123", + "base_url": server.URL, + }} + payload := []byte(`{"system":"Top rule","messages":[{"role":"user","content":[{"type":"text","text":"hi"}]},{"role":"system","content":"Mid string rule"},{"role":"assistant","content":[{"type":"text","text":"ok"}]},{"role":"system","content":[{"type":"text","text":"Mid array rule","cache_control":{"type":"ephemeral"}}]},{"role":"user","content":[{"type":"text","text":"continue"}]}]}`) + ctx := contextWithGinHeaders(map[string]string{"User-Agent": "claude-cli/2.1.153 (external, cli)"}) + + _, errExecute := executor.Execute(ctx, auth, cliproxyexecutor.Request{ + Model: "claude-3-5-sonnet-20241022", + Payload: payload, + }, cliproxyexecutor.Options{SourceFormat: sdktranslator.FromString("claude")}) + if errExecute != nil { + t.Fatalf("Execute() error = %v", errExecute) + } + if len(seenBody) == 0 { + t.Fatal("expected request body to be captured") + } + + system := gjson.GetBytes(seenBody, "system").Array() + if len(system) != 3 { + t.Fatalf("system has %d items, want 3: %s", len(system), gjson.GetBytes(seenBody, "system").Raw) + } + wantTexts := []string{"Top rule", "Mid string rule", "Mid array rule"} + for i, want := range wantTexts { + if got := system[i].Get("text").String(); got != want { + t.Fatalf("system[%d].text = %q, want %q", i, got, want) + } + } + if got := gjson.GetBytes(seenBody, "system.2.cache_control.type").String(); got != "ephemeral" { + t.Fatalf("system.2.cache_control.type = %q, want ephemeral", got) + } + if gjson.GetBytes(seenBody, `messages.#(role=="system")`).Exists() { + t.Fatalf("messages should not contain system role after rebuild: %s", gjson.GetBytes(seenBody, "messages").Raw) + } + if got := gjson.GetBytes(seenBody, "messages.#").Int(); got != 3 { + t.Fatalf("messages count = %d, want 3", got) + } +} + func TestApplyCloaking_PreservesConfiguredStrictModeAndSensitiveWordsWhenModeOmitted(t *testing.T) { cfg := &config.Config{ ClaudeKey: []config.ClaudeKey{{ diff --git a/internal/runtime/executor/claude_signing.go b/internal/runtime/executor/claude_signing.go index 060e86e84..8afd57a67 100644 --- a/internal/runtime/executor/claude_signing.go +++ b/internal/runtime/executor/claude_signing.go @@ -79,3 +79,11 @@ func experimentalCCHSigningEnabled(cfg *config.Config, auth *cliproxyauth.Auth) entry := resolveClaudeKeyConfig(cfg, auth) return entry != nil && entry.ExperimentalCCHSigning } + +func rebuildMidSystemMessageEnabled(cfg *config.Config, auth *cliproxyauth.Auth) bool { + if auth != nil && auth.Attributes != nil && strings.EqualFold(strings.TrimSpace(auth.Attributes["rebuild_mid_system_message"]), "true") { + return true + } + entry := resolveClaudeKeyConfig(cfg, auth) + return entry != nil && entry.RebuildMidSystemMessage +} diff --git a/internal/watcher/diff/config_diff.go b/internal/watcher/diff/config_diff.go index 903526951..80cc44ddc 100644 --- a/internal/watcher/diff/config_diff.go +++ b/internal/watcher/diff/config_diff.go @@ -185,6 +185,9 @@ func BuildConfigChangeDetails(oldCfg, newCfg *config.Config) []string { if oldExcluded.hash != newExcluded.hash { changes = append(changes, fmt.Sprintf("claude[%d].excluded-models: updated (%d -> %d entries)", i, oldExcluded.count, newExcluded.count)) } + if o.RebuildMidSystemMessage != n.RebuildMidSystemMessage { + changes = append(changes, fmt.Sprintf("claude[%d].rebuild-mid-system-message: %t -> %t", i, o.RebuildMidSystemMessage, n.RebuildMidSystemMessage)) + } if o.Cloak != nil && n.Cloak != nil { if strings.TrimSpace(o.Cloak.Mode) != strings.TrimSpace(n.Cloak.Mode) { changes = append(changes, fmt.Sprintf("claude[%d].cloak.mode: %s -> %s", i, o.Cloak.Mode, n.Cloak.Mode)) diff --git a/internal/watcher/synthesizer/config.go b/internal/watcher/synthesizer/config.go index 14e1cd5b2..82a75cf78 100644 --- a/internal/watcher/synthesizer/config.go +++ b/internal/watcher/synthesizer/config.go @@ -126,6 +126,9 @@ func (s *ConfigSynthesizer) synthesizeClaudeKeys(ctx *SynthesisContext) []*corea if base != "" { attrs["base_url"] = base } + if ck.RebuildMidSystemMessage { + attrs["rebuild_mid_system_message"] = "true" + } if hash := diff.ComputeClaudeModelsHash(ck.Models); hash != "" { attrs["models_hash"] = hash } diff --git a/internal/watcher/synthesizer/config_test.go b/internal/watcher/synthesizer/config_test.go index 0016844b3..5646ef871 100644 --- a/internal/watcher/synthesizer/config_test.go +++ b/internal/watcher/synthesizer/config_test.go @@ -175,10 +175,11 @@ func TestConfigSynthesizer_ClaudeKeys(t *testing.T) { Config: &config.Config{ ClaudeKey: []config.ClaudeKey{ { - APIKey: "sk-ant-api-xxx", - Prefix: "main", - BaseURL: "https://api.anthropic.com", - DisableCooling: true, + APIKey: "sk-ant-api-xxx", + Prefix: "main", + BaseURL: "https://api.anthropic.com", + DisableCooling: true, + RebuildMidSystemMessage: true, Models: []config.ClaudeModel{ {Name: "claude-3-opus"}, {Name: "claude-3-sonnet"}, @@ -213,6 +214,9 @@ func TestConfigSynthesizer_ClaudeKeys(t *testing.T) { if _, ok := auths[0].Attributes["models_hash"]; !ok { t.Error("expected models_hash in attributes") } + if got := auths[0].Attributes["rebuild_mid_system_message"]; got != "true" { + t.Errorf("expected rebuild_mid_system_message=true, got %s", got) + } if v, ok := auths[0].Metadata["disable_cooling"].(bool); !ok || !v { t.Errorf("expected disable_cooling=true, got %v", auths[0].Metadata["disable_cooling"]) }