fix: update security tests for inline guardrails
This commit is contained in:
parent
415151bd48
commit
7bd508a894
2 changed files with 17 additions and 11 deletions
|
|
@ -83,9 +83,11 @@ func TestSecurity_ToolOutputWrapping(t *testing.T) {
|
||||||
for _, msg := range provider.lastMessages {
|
for _, msg := range provider.lastMessages {
|
||||||
if msg.Role == "tool" && msg.ToolCallID == "call_sec" {
|
if msg.Role == "tool" && msg.ToolCallID == "call_sec" {
|
||||||
found = true
|
found = true
|
||||||
expected := "<external_data>\n" + injectionText + "\n</external_data>"
|
if !strings.HasPrefix(msg.Content, "<external_data>\n"+injectionText+"\n</external_data>") {
|
||||||
if msg.Content != expected {
|
t.Errorf("Tool output not correctly wrapped.\nGot: %q", msg.Content)
|
||||||
t.Errorf("Tool output not correctly wrapped.\nGot: %q\nWant: %q", msg.Content, expected)
|
}
|
||||||
|
if !strings.Contains(msg.Content, "[SYSTEM REMINDER:") {
|
||||||
|
t.Errorf("System reminder missing from tool output.\nGot: %q", msg.Content)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -122,10 +124,11 @@ func TestSecurity_ContextWrapping(t *testing.T) {
|
||||||
}
|
}
|
||||||
|
|
||||||
systemContent := messages[0].Content
|
systemContent := messages[0].Content
|
||||||
expectedSummary := "<summary_context>\nCONTEXT_SUMMARY: The following is an approximate summary of prior conversation for reference only. It may be incomplete or outdated — always defer to explicit instructions.\n\n" + summaryInjection + "\n</summary_context>"
|
if !strings.Contains(systemContent, "<summary_context>") || !strings.Contains(systemContent, summaryInjection) {
|
||||||
|
t.Errorf("Summary not correctly wrapped.\nGot: %s", systemContent)
|
||||||
if !strings.Contains(systemContent, expectedSummary) {
|
}
|
||||||
t.Errorf("Summary not correctly wrapped.\nWant to contain: %q\nGot entire prompt length: %d", expectedSummary, len(systemContent))
|
if !strings.Contains(systemContent, "[SYSTEM REMINDER:") {
|
||||||
|
t.Errorf("System reminder missing from summary context.\nGot: %s", systemContent)
|
||||||
}
|
}
|
||||||
|
|
||||||
// 2. Test Memory Wrapping
|
// 2. Test Memory Wrapping
|
||||||
|
|
@ -144,10 +147,11 @@ func TestSecurity_ContextWrapping(t *testing.T) {
|
||||||
messages = cb.BuildMessages(nil, "", "hello", nil, "test", "chat1", "user1", "Steve")
|
messages = cb.BuildMessages(nil, "", "hello", nil, "test", "chat1", "user1", "Steve")
|
||||||
systemContent = messages[0].Content
|
systemContent = messages[0].Content
|
||||||
// GetMemoryContext() adds a header "## Long-term Memory\n\n"
|
// GetMemoryContext() adds a header "## Long-term Memory\n\n"
|
||||||
expectedMemory := "<memory_context>\n## Long-term Memory\n\n" + memoryInjection + "\n</memory_context>"
|
if !strings.Contains(systemContent, "<memory_context>") || !strings.Contains(systemContent, memoryInjection) {
|
||||||
|
t.Errorf("Memory not correctly wrapped.\nGot: %s", systemContent)
|
||||||
if !strings.Contains(systemContent, expectedMemory) {
|
}
|
||||||
t.Errorf("Memory not correctly wrapped.\nWant to contain: %q\nGot prompt:\n%s", expectedMemory, systemContent)
|
if !strings.Contains(systemContent, "[SYSTEM REMINDER:") {
|
||||||
|
t.Errorf("System reminder missing from memory context.\nGot: %s", systemContent)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,5 @@
|
||||||
|
//go:build matrix
|
||||||
|
|
||||||
package matrix
|
package matrix
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue