fix: remove system role filtering, TrustyAI now supports system messages#232
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Central YAML (base), Organization UI (inherited) Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe change removes the special-case filtering of OpenAI Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Security Assessment
🚥 Pre-merge checks | ✅ 4✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
pkg/plugins/nemo/request_guard.go (1)
160-167:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winBlock non-string
contentinstead of silently turning it into"".Severity: critical. Exploit: send a chat message with OpenAI array-form content parts, e.g. a
systemorusermessage whosecontentis[{type:"text",text:"..."}]. This code forwards{"content": ""}to NeMo, while the original request still carries the real instruction to the downstream model, creating a guardrail bypass on the exact path this PR is trying to harden. CWE-20.Suggested fix
func extractOpenAIMessages(raw any) ([]map[string]string, error) { slice, ok := raw.([]any) if !ok { return nil, fmt.Errorf("messages is not an array") } - var messages []map[string]string + messages := make([]map[string]string, 0, len(slice)) - for _, m := range slice { + for i, m := range slice { msg, ok := m.(map[string]any) if !ok { - continue + return nil, fmt.Errorf("message at index %d is not an object", i) } role, _ := msg["role"].(string) - content, _ := msg["content"].(string) + if role == "" { + return nil, fmt.Errorf("message at index %d has invalid role", i) + } + content, err := extractOpenAIContent(msg["content"]) + if err != nil { + return nil, fmt.Errorf("message at index %d: %w", i, err) + } messages = append(messages, map[string]string{"role": role, "content": content}) } return messages, nil } + +func extractOpenAIContent(raw any) (string, error) { + switch v := raw.(type) { + case string: + return v, nil + case []any: + parts := make([]string, 0, len(v)) + for i, p := range v { + part, ok := p.(map[string]any) + if !ok { + return "", fmt.Errorf("content part %d is not an object", i) + } + if part["type"] != "text" { + continue + } + text, ok := part["text"].(string) + if !ok { + return "", fmt.Errorf("content part %d has invalid text", i) + } + parts = append(parts, text) + } + if len(parts) == 0 { + return "", fmt.Errorf("unsupported content format") + } + return strings.Join(parts, "\n"), nil + default: + return "", fmt.Errorf("unsupported content format") + } +}As per coding guidelines,
**: REVIEW PRIORITIES: 1. Security vulnerabilities (provide severity, exploit scenario, and remediation code).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@pkg/plugins/nemo/request_guard.go` around lines 160 - 167, The loop currently coerces non-string message content into an empty string (content, _ := msg["content"].(string)) which allows array-form or structured contents to be forwarded; instead, detect when msg["content"] is not a string and reject the request/path. Change the logic in the loop that iterates over slice (the variables m, msg, role, content, messages) to validate content using a strict type assertion (e.g., v, ok := msg["content"]; s, isString := v.(string)) and if not isString return or mark the request invalid (do not append a map with ""), logging or surfacing an error so the guard blocks any messages whose content is not a plain string. Ensure the same strict behavior for role if applicable.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@pkg/plugins/nemo/request_guard.go`:
- Around line 160-167: The loop currently coerces non-string message content
into an empty string (content, _ := msg["content"].(string)) which allows
array-form or structured contents to be forwarded; instead, detect when
msg["content"] is not a string and reject the request/path. Change the logic in
the loop that iterates over slice (the variables m, msg, role, content,
messages) to validate content using a strict type assertion (e.g., v, ok :=
msg["content"]; s, isString := v.(string)) and if not isString return or mark
the request invalid (do not append a map with ""), logging or surfacing an error
so the guard blocks any messages whose content is not a plain string. Ensure the
same strict behavior for role if applicable.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Central YAML (base), Organization UI (inherited)
Review profile: CHILL
Plan: Enterprise
Run ID: e66b650b-7947-45f5-9725-7564e910c72a
📒 Files selected for processing (2)
pkg/plugins/nemo/request_guard.gopkg/plugins/nemo/request_guard_test.go
|
@liavweiss making sure - was this change tested with latest trusty version? |
| continue | ||
| } | ||
| role, _ := msg["role"].(string) | ||
| if role == "system" { |
There was a problem hiding this comment.
would it be correct to just return “slice” without copying?
in current code we copy the whole messages array on every request.
so if we have a long multi turn chat, memory allocation for each turn is expensive.
There was a problem hiding this comment.
Thanks, I added a pre-allocation and an early return for empty slices to avoid unnecessary allocations.
There was a problem hiding this comment.
my question was actually about the messages themselves, not only about the case of empty slice.
my intention was - if we have something along these lines:
func extractOpenAIMessages(raw any) ([]map[string]string, error) {
slice, ok := raw.([]any)
if !ok {
return nil, fmt.Errorf("messages is not an array")
}
if len(slice) == 0 {
return nil, nil
}
return slicewould it work?
There was a problem hiding this comment.
Should work with some extra changes, but returning the slice directly means forwarding all original message fields to NeMo (e.g., tool_calls, function_call) instead of only role and content. NeMo ignores extra fields, but it sends more data over the wire for no reason. It also requires changing the return type from []map[string]string to []any and updating all callers and tests accordingly. If we want to go in that direction, it would be better as a follow-up PR.
There was a problem hiding this comment.
sure. I'm just thinking about the tradeoffs.
do we prefer allocating memory for every request and copying the whole chat? (that could be very large) or do we prefer sending the same request that we received over the wire to yet another http hop? (but at least we don't do both - allocate memory + sending http).
there are pros/cons for both sides.
my gut feeling says sending the original request as is would probably be "cheaper".
we can handle in a follow up, for sure.
Signed-off-by: Liav Weiss <lweiss@lweiss-thinkpadx1carbongen11.raanaii.csb>
c704557 to
32c73fc
Compare
Yes, I built the TrustyAI NeMo image locally that includes the new code #47 |
|
[APPROVALNOTIFIER] This PR is APPROVED This pull-request has been approved by: liavweiss, nirrozenbaum The full list of commands accepted by this bot can be found here. The pull request process is described here DetailsNeeds approval from an approver in each of these files:
Approvers can indicate their approval by writing |
Summary
nemo-request-guard. Previously, messages withrole: "system"were filtered out before sending to NeMo because the TrustyAI/v1/guardrail/checksendpoint rejected them with a hard error (Unsupported message role: 'system').Closes #160. #205
Changes
request_guard.go: Removed theif role == "system" { continue }filter inextractOpenAIMessages. Updated comments.request_guard_test.go: Updated tests to expect system messages in the forwarded payload.E2E verification
Verified end-to-end by building the TrustyAI NeMo image locally from the
developbranch (Dockerfile.server) which includes PR #47, and deploying it on a Kind cluster withdetect sensitive data on inputrails:system: "You are a helpful assistant"+user: "Hello"system: "card 4111-1111-1111-1111"+user: "Hello"system: "Contact John Smith at john@example.com"cc @nirrozenbaum @christinaexyou
Summary by CodeRabbit