Add a system-prompt-count option for some Claude providers to pass validation#1838
Add a system-prompt-count option for some Claude providers to pass validation#1838Kindness-Kismet wants to merge 3 commits intorouter-for-me:mainfrom
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses a validation requirement from certain Claude providers by introducing a configurable Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a system-prompt-count option for Claude providers to enforce a specific number of system prompt blocks, useful for compatibility with certain Claude-like APIs. However, a critical security vulnerability has been identified: the lack of input validation for the maximum allowed system-prompt-count creates a potential Denial of Service (DoS) vector, as a large value can lead to excessive memory allocation and server crashes. Additionally, a high-severity bug was found where non-text system prompt blocks (like cache_control) are discarded, and there's a medium-severity issue with redundant code in configuration handling.
| func enforceSystemPromptCount(payload []byte, targetCount int) []byte { | ||
| if targetCount <= 0 { | ||
| return payload | ||
| } | ||
| root, ok := parsePayloadObject(payload) | ||
| if !ok { | ||
| return payload | ||
| } | ||
|
|
||
| texts := make([]string, 0, targetCount) | ||
| appendText := func(value any) { | ||
| switch v := value.(type) { | ||
| case string: | ||
| texts = append(texts, v) | ||
| case map[string]any: | ||
| if blockType, ok := v["type"].(string); ok && blockType != "" && !strings.EqualFold(blockType, "text") { | ||
| return | ||
| } | ||
| return entry.Cloak | ||
| text, ok := v["text"].(string) | ||
| if !ok { | ||
| return | ||
| } | ||
| texts = append(texts, text) | ||
| } | ||
| } | ||
|
|
||
| return nil | ||
| if system, exists := root["system"]; exists { | ||
| if items, ok := asArray(system); ok { | ||
| for _, item := range items { | ||
| appendText(item) | ||
| } | ||
| } else { | ||
| appendText(system) | ||
| } | ||
| } | ||
|
|
||
| normalized := make([]string, 0, targetCount) | ||
| if len(texts) >= targetCount { | ||
| normalized = append(normalized, texts[:targetCount-1]...) | ||
| normalized = append(normalized, strings.Join(texts[targetCount-1:], "\n\n")) | ||
| } else { | ||
| normalized = append(normalized, texts...) | ||
| for len(normalized) < targetCount { | ||
| normalized = append(normalized, "") | ||
| } | ||
| } | ||
|
|
||
| blocks := make([]any, 0, targetCount) | ||
| for _, text := range normalized { | ||
| blocks = append(blocks, map[string]any{ | ||
| "type": "text", | ||
| "text": text, | ||
| }) | ||
| } | ||
| root["system"] = blocks | ||
| return marshalPayloadObject(payload, root) | ||
| } |
There was a problem hiding this comment.
The function enforceSystemPromptCount incorrectly handles system prompts containing non-text blocks. It extracts only the text content from the system array and then reconstructs it using only text blocks. This causes any non-text blocks, such as cache_control blocks, to be discarded from the request payload. This can lead to incorrect behavior, for example, by breaking prompt caching.
The logic should be updated to preserve non-text blocks while adjusting the number of text blocks.
| if entry.SystemPromptCount < 0 { | ||
| entry.SystemPromptCount = 0 | ||
| } |
There was a problem hiding this comment.
The system-prompt-count parameter lacks upper bound validation. An attacker with access to the management API could set this value to an extremely large number, leading to an Out-Of-Memory (OOM) crash of the server when enforceSystemPromptCount attempts to allocate large slices. It is recommended to enforce a reasonable upper limit (e.g., 100) for SystemPromptCount. Additionally, the existing if entry.SystemPromptCount < 0 { entry.SystemPromptCount = 0 } validation might be redundant here, as SanitizeClaudeKeys performs this exact sanitization; consider consolidating sanitization logic to avoid duplication.
if entry.SystemPromptCount < 0 {
entry.SystemPromptCount = 0
} else if entry.SystemPromptCount > 100 {
entry.SystemPromptCount = 100
}| if entry.SystemPromptCount < 0 { | ||
| entry.SystemPromptCount = 0 | ||
| } |
There was a problem hiding this comment.
Similar to the management API handler, the configuration sanitization logic for Claude keys lacks an upper bound check on SystemPromptCount. This could allow a maliciously crafted configuration file to cause a Denial of Service via memory exhaustion when requests are processed.
Recommendation: Enforce a reasonable upper limit (e.g., 100) for SystemPromptCount.
if entry.SystemPromptCount < 0 {
entry.SystemPromptCount = 0
} else if entry.SystemPromptCount > 100 {
entry.SystemPromptCount = 100
}
luispater
left a comment
luispater
left a comment
There was a problem hiding this comment.
Summary:
This PR adds system-prompt-count to Claude provider config, threads it through the management API and config diff output, and enforces it in the Claude executor. I ran the targeted package tests plus a server build, and they passed. I still found two blocking correctness issues in the executor path.
Key findings:
- Blocking: shrinking the
systemarray can reorder text around intervening non-text system blocks. - Blocking: the new policy runs after cloaking, so it rewrites the Claude Code billing/agent blocks when
cloakandsystem-prompt-countare enabled together.
Test plan:
go test ./internal/runtime/executor ./internal/config ./internal/watcher/diff ./internal/api/handlers/managementgo build -o test-output ./cmd/server
Follow-ups:
- Apply the count policy before cloaking, or explicitly exclude cloaking-injected blocks from the count.
- Add a regression test for a mixed sequence such as
[text, text, non-text, text]when shrinking to2.
…s non-text blocks - Merge adjacent text blocks from the end first, avoiding text content crossing non-text block boundaries. - Fall back to cross-boundary merge only when adjacent merge cannot reach the target count. - Move system-prompt-count policy before cloaking so cloaking-injected system blocks are excluded from the count. - Add regression tests for mixed text/non-text sequences and exact-count no-op.
Thanks! All issues have been addressed |
2 participants
Some Claude providers verify the number of prompts, so this option has been added.