feat(security-audit): add narrow blocking audit scope

This commit is contained in:
gaoren002
2026-07-31 01:47:53 +00:00
parent 5a6143097d
commit d74e669a23
16 changed files with 317 additions and 105 deletions
+84 -77
View File
@@ -64,20 +64,21 @@ type StorageEndpoint struct {
}
type storageConfig struct {
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
StorePassEvents bool `json:"store_pass_events"`
Strategy string `json:"strategy"`
WorkerCount int `json:"worker_count"`
QueueCapacity int `json:"queue_capacity"`
Scanners []string `json:"scanners"`
AllGroups bool `json:"all_groups"`
GroupIDs []int64 `json:"group_ids"`
Endpoints []StorageEndpoint `json:"endpoints"`
ConfigVersion int64 `json:"config_version"`
UpdatedAt time.Time `json:"updated_at"`
UpdatedBy int64 `json:"updated_by"`
ChangeSummary string `json:"change_summary"`
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
BlockingLatestTurnOnly bool `json:"blocking_latest_turn_only"`
StorePassEvents bool `json:"store_pass_events"`
Strategy string `json:"strategy"`
WorkerCount int `json:"worker_count"`
QueueCapacity int `json:"queue_capacity"`
Scanners []string `json:"scanners"`
AllGroups bool `json:"all_groups"`
GroupIDs []int64 `json:"group_ids"`
Endpoints []StorageEndpoint `json:"endpoints"`
ConfigVersion int64 `json:"config_version"`
UpdatedAt time.Time `json:"updated_at"`
UpdatedBy int64 `json:"updated_by"`
ChangeSummary string `json:"change_summary"`
}
type ActiveEndpoint struct {
@@ -98,21 +99,22 @@ type ActiveEndpoint struct {
}
type ActiveConfig struct {
RiskControlEnabled bool
Enabled bool
BlockingEnabled bool
StorePassEvents bool
Strategy string
WorkerCount int
QueueCapacity int
Scanners []string
AllGroups bool
GroupIDs []int64
Endpoints []ActiveEndpoint
ConfigVersion int64
UpdatedAt time.Time
UpdatedBy int64
ChangeSummary string
RiskControlEnabled bool
Enabled bool
BlockingEnabled bool
BlockingLatestTurnOnly bool
StorePassEvents bool
Strategy string
WorkerCount int
QueueCapacity int
Scanners []string
AllGroups bool
GroupIDs []int64
Endpoints []ActiveEndpoint
ConfigVersion int64
UpdatedAt time.Time
UpdatedBy int64
ChangeSummary string
}
type PublicEndpoint struct {
@@ -129,21 +131,22 @@ type PublicEndpoint struct {
}
type PublicConfig struct {
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
StorePassEvents bool `json:"store_pass_events"`
EffectiveMode Mode `json:"effective_mode"`
Strategy string `json:"strategy"`
WorkerCount int `json:"worker_count"`
QueueCapacity int `json:"queue_capacity"`
Scanners []string `json:"scanners"`
AllGroups bool `json:"all_groups"`
GroupIDs []int64 `json:"group_ids"`
Endpoints []PublicEndpoint `json:"endpoints"`
ConfigVersion int64 `json:"config_version"`
UpdatedAt time.Time `json:"updated_at"`
UpdatedBy int64 `json:"updated_by"`
ChangeSummary string `json:"change_summary"`
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
BlockingLatestTurnOnly bool `json:"blocking_latest_turn_only"`
StorePassEvents bool `json:"store_pass_events"`
EffectiveMode Mode `json:"effective_mode"`
Strategy string `json:"strategy"`
WorkerCount int `json:"worker_count"`
QueueCapacity int `json:"queue_capacity"`
Scanners []string `json:"scanners"`
AllGroups bool `json:"all_groups"`
GroupIDs []int64 `json:"group_ids"`
Endpoints []PublicEndpoint `json:"endpoints"`
ConfigVersion int64 `json:"config_version"`
UpdatedAt time.Time `json:"updated_at"`
UpdatedBy int64 `json:"updated_by"`
ChangeSummary string `json:"change_summary"`
}
type UpdateEndpoint struct {
@@ -160,32 +163,34 @@ type UpdateEndpoint struct {
}
type UpdateConfigRequest struct {
ExpectedConfigVersion int64 `json:"expected_config_version" binding:"required"`
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
StorePassEvents bool `json:"store_pass_events"`
Strategy string `json:"strategy"`
WorkerCount int `json:"worker_count"`
QueueCapacity int `json:"queue_capacity"`
Scanners []string `json:"scanners"`
AllGroups bool `json:"all_groups"`
GroupIDs []int64 `json:"group_ids"`
Endpoints []UpdateEndpoint `json:"endpoints"`
ExpectedConfigVersion int64 `json:"expected_config_version" binding:"required"`
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
BlockingLatestTurnOnly bool `json:"blocking_latest_turn_only"`
StorePassEvents bool `json:"store_pass_events"`
Strategy string `json:"strategy"`
WorkerCount int `json:"worker_count"`
QueueCapacity int `json:"queue_capacity"`
Scanners []string `json:"scanners"`
AllGroups bool `json:"all_groups"`
GroupIDs []int64 `json:"group_ids"`
Endpoints []UpdateEndpoint `json:"endpoints"`
}
func DefaultStorageConfig() storageConfig {
return storageConfig{
Enabled: false,
BlockingEnabled: false,
StorePassEvents: false,
Strategy: "priority",
WorkerCount: DefaultWorkerCount,
QueueCapacity: DefaultQueueCapacity,
Scanners: append([]string(nil), AllScannerIDs...),
AllGroups: true,
GroupIDs: []int64{},
Endpoints: []StorageEndpoint{},
ConfigVersion: 1,
Enabled: false,
BlockingEnabled: false,
BlockingLatestTurnOnly: false,
StorePassEvents: false,
Strategy: "priority",
WorkerCount: DefaultWorkerCount,
QueueCapacity: DefaultQueueCapacity,
Scanners: append([]string(nil), AllScannerIDs...),
AllGroups: true,
GroupIDs: []int64{},
Endpoints: []StorageEndpoint{},
ConfigVersion: 1,
}
}
@@ -407,7 +412,7 @@ func PublicFromStorage(cfg storageConfig, riskControlEnabled bool, invalidTokenE
}
active := ActiveConfig{RiskControlEnabled: riskControlEnabled, Enabled: cfg.Enabled, BlockingEnabled: cfg.BlockingEnabled}
return PublicConfig{
Enabled: cfg.Enabled, BlockingEnabled: cfg.BlockingEnabled, StorePassEvents: cfg.StorePassEvents,
Enabled: cfg.Enabled, BlockingEnabled: cfg.BlockingEnabled, BlockingLatestTurnOnly: cfg.BlockingLatestTurnOnly, StorePassEvents: cfg.StorePassEvents,
EffectiveMode: active.EffectiveMode(), Strategy: cfg.Strategy, WorkerCount: cfg.WorkerCount,
QueueCapacity: cfg.QueueCapacity, Scanners: scanners, AllGroups: cfg.AllGroups,
GroupIDs: groupIDs, Endpoints: endpoints, ConfigVersion: cfg.ConfigVersion,
@@ -418,7 +423,8 @@ func PublicFromStorage(cfg storageConfig, riskControlEnabled bool, invalidTokenE
func ActiveFromStorage(cfg storageConfig, riskControlEnabled bool, encryptor SecretEncryptor) (ActiveConfig, error) {
active := ActiveConfig{
RiskControlEnabled: riskControlEnabled, Enabled: cfg.Enabled, BlockingEnabled: cfg.BlockingEnabled,
StorePassEvents: cfg.StorePassEvents, Strategy: cfg.Strategy, WorkerCount: cfg.WorkerCount,
BlockingLatestTurnOnly: cfg.BlockingLatestTurnOnly,
StorePassEvents: cfg.StorePassEvents, Strategy: cfg.Strategy, WorkerCount: cfg.WorkerCount,
QueueCapacity: cfg.QueueCapacity, Scanners: append([]string(nil), cfg.Scanners...), AllGroups: cfg.AllGroups,
GroupIDs: append([]int64(nil), cfg.GroupIDs...), ConfigVersion: cfg.ConfigVersion,
UpdatedAt: cfg.UpdatedAt, UpdatedBy: cfg.UpdatedBy, ChangeSummary: cfg.ChangeSummary,
@@ -454,15 +460,16 @@ func ActiveFromStorage(cfg storageConfig, riskControlEnabled bool, encryptor Sec
func changeSummary(cfg storageConfig) string {
summary := struct {
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
StorePassEvents bool `json:"store_pass_events"`
EndpointCount int `json:"endpoint_count"`
ScannerCount int `json:"scanner_count"`
AllGroups bool `json:"all_groups"`
GroupCount int `json:"group_count"`
GroupHash string `json:"group_hash"`
}{cfg.Enabled, cfg.BlockingEnabled, cfg.StorePassEvents, len(cfg.Endpoints), len(cfg.Scanners), cfg.AllGroups, len(cfg.GroupIDs), ""}
Enabled bool `json:"enabled"`
BlockingEnabled bool `json:"blocking_enabled"`
BlockingLatestTurnOnly bool `json:"blocking_latest_turn_only"`
StorePassEvents bool `json:"store_pass_events"`
EndpointCount int `json:"endpoint_count"`
ScannerCount int `json:"scanner_count"`
AllGroups bool `json:"all_groups"`
GroupCount int `json:"group_count"`
GroupHash string `json:"group_hash"`
}{cfg.Enabled, cfg.BlockingEnabled, cfg.BlockingLatestTurnOnly, cfg.StorePassEvents, len(cfg.Endpoints), len(cfg.Scanners), cfg.AllGroups, len(cfg.GroupIDs), ""}
rawGroups, _ := json.Marshal(cfg.GroupIDs)
digest := sha256.Sum256(rawGroups)
summary.GroupHash = hex.EncodeToString(digest[:])
@@ -335,7 +335,7 @@ func (m *ConfigManager) buildNextStorage(current storageConfig, req UpdateConfig
currentByID[endpoint.ID] = endpoint
}
next := storageConfig{
Enabled: req.Enabled, BlockingEnabled: req.BlockingEnabled, StorePassEvents: req.StorePassEvents,
Enabled: req.Enabled, BlockingEnabled: req.BlockingEnabled, BlockingLatestTurnOnly: req.BlockingLatestTurnOnly, StorePassEvents: req.StorePassEvents,
Strategy: strings.TrimSpace(req.Strategy), WorkerCount: req.WorkerCount,
QueueCapacity: req.QueueCapacity, Scanners: append([]string(nil), req.Scanners...),
AllGroups: req.AllGroups, GroupIDs: append([]int64(nil), req.GroupIDs...),
@@ -32,6 +32,7 @@ func TestDefaultConfigIsOff(t *testing.T) {
storage, err := ParseStorageConfig("")
require.NoError(t, err)
require.False(t, storage.Enabled)
require.False(t, storage.BlockingLatestTurnOnly)
active, err := ActiveFromStorage(storage, true, prefixEncryptor{})
require.NoError(t, err)
require.Equal(t, ModeOff, active.EffectiveMode())
@@ -42,6 +43,28 @@ func TestDefaultConfigIsOff(t *testing.T) {
require.Contains(t, string(publicJSON), `"endpoints":[]`)
}
func TestBlockingLatestTurnOnlyConfigRoundTrip(t *testing.T) {
manager := &ConfigManager{encryptor: prefixEncryptor{}, encryptionKeyConfigured: true}
request := UpdateConfigRequest{
ExpectedConfigVersion: 1, Enabled: true, BlockingEnabled: true, BlockingLatestTurnOnly: true,
Strategy: "priority", WorkerCount: 1, QueueCapacity: 10, Scanners: []string{"pii"}, AllGroups: true,
Endpoints: []UpdateEndpoint{{
ID: "guard-1", Name: "Guard", Protocol: "openai_compatible", BaseURL: "http://127.0.0.1:8080",
Model: DefaultGuardModel, TimeoutMS: 1000, InputLimit: 1000, Enabled: true,
}},
}
next, err := manager.buildNextStorage(DefaultStorageConfig(), request, 9)
require.NoError(t, err)
require.True(t, next.BlockingLatestTurnOnly)
require.Contains(t, changeSummary(next), `"blocking_latest_turn_only":true`)
active, err := ActiveFromStorage(next, true, prefixEncryptor{})
require.NoError(t, err)
require.True(t, active.BlockingLatestTurnOnly)
public := PublicFromStorage(next, true, nil)
require.True(t, public.BlockingLatestTurnOnly)
}
func TestConfigRejectsBlockingWithoutAudit(t *testing.T) {
storage := DefaultStorageConfig()
storage.BlockingEnabled = true
@@ -228,7 +228,8 @@ func configAuditFields(request UpdateConfigRequest, saved *PublicConfig) map[str
}
return map[string]any{
"enabled": request.Enabled, "blocking_enabled": request.BlockingEnabled,
"config_version": version, "endpoint_count": len(request.Endpoints),
"blocking_latest_turn_only": request.BlockingLatestTurnOnly,
"config_version": version, "endpoint_count": len(request.Endpoints),
"scanner_count": len(request.Scanners), "all_groups": request.AllGroups,
"group_count": len(request.GroupIDs),
}
@@ -156,7 +156,7 @@ func (s *PromptService) Evaluate(ctx context.Context, req Request) (*PromptDecis
if cfg.EffectiveMode() != ModeBlocking || !cfg.IncludesGroup(req.GroupID) {
return &PromptDecision{Kind: DecisionAllow, AllowNextStage: true}, nil
}
snapshot, err := ExtractPromptSnapshot(req)
snapshot, err := ExtractBlockingPromptSnapshot(req, cfg.BlockingLatestTurnOnly)
if errors.Is(err, ErrNoPromptText) {
return &PromptDecision{Kind: DecisionAllow, AllowNextStage: true}, nil
}
@@ -68,6 +68,25 @@ func TestPromptServiceStartReportsDependencyFailureWithoutPanic(t *testing.T) {
require.NoError(t, service.Shutdown(ctx))
}
func TestPromptServiceBlockingLatestTurnOnlyUsesNarrowSnapshot(t *testing.T) {
seen := make([]string, 0, 2)
evaluator := newGuardEvaluator(PromptScannerFunc(func(_ context.Context, _ ActiveEndpoint, chunk string, _ []string) (*NormalizedResult, error) {
seen = append(seen, chunk)
return &NormalizedResult{Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow, ScannerScores: map[string]float64{}, ScannerEvidence: map[string]string{}}, nil
}), nil, NewAtomicMetrics(), 2, 2)
service := &PromptService{
config: &fakeConfigStore{active: true, cfg: ActiveConfig{
RiskControlEnabled: true, Enabled: true, BlockingEnabled: true, BlockingLatestTurnOnly: true, AllGroups: true,
Scanners: AllScannerIDs, Endpoints: []ActiveEndpoint{{ID: "guard-1", Enabled: true, TimeoutMS: 1000, InputLimit: 4096}},
}},
evaluator: evaluator,
}
decision, err := service.Evaluate(context.Background(), Request{Protocol: "openai_chat_completions", Body: []byte(`{"messages":[{"role":"system","content":"system instruction"},{"role":"user","content":"older user input"},{"role":"assistant","content":"previous output"},{"role":"user","content":"latest user input"}]}`)})
require.NoError(t, err)
require.Equal(t, DecisionAllow, decision.Kind)
require.Equal(t, []string{"latest user input", "previous output"}, seen)
}
func TestPromptServiceRejectsInvalidDeleteConfirmationClaims(t *testing.T) {
now := time.Date(2026, 7, 16, 10, 0, 0, 0, time.UTC)
start, end := now.Add(-time.Hour), now.Add(time.Hour)
+109 -20
View File
@@ -26,15 +26,30 @@ const promptAuditPrioritySeparator = "\x00SUB2API_PROMPT_AUDIT_PRIORITY_END\x00"
type promptSegment struct {
text string
user bool
role string
}
func ExtractPromptSnapshot(req Request) (PromptSnapshot, error) {
return extractPromptSnapshot(req, false)
}
// ExtractBlockingPromptSnapshot builds the narrow, low-latency blocking input
// when configured. Asynchronous auditing always uses ExtractPromptSnapshot so
// the complete client-controlled transcript is retained for review.
func ExtractBlockingPromptSnapshot(req Request, latestTurnOnly bool) (PromptSnapshot, error) {
return extractPromptSnapshot(req, latestTurnOnly)
}
func extractPromptSnapshot(req Request, latestTurnOnly bool) (PromptSnapshot, error) {
var document any
if err := json.Unmarshal(req.Body, &document); err != nil {
return PromptSnapshot{}, errors.New("prompt audit request JSON is invalid")
}
extracted := extractProtocolSegments(req.Protocol, document)
segments := normalizeSegmentsLatestUserFirst(extracted)
if latestTurnOnly {
segments = blockingSegmentsLatestUserAndPreviousOutput(extracted)
}
if len(segments) == 0 {
return PromptSnapshot{}, ErrNoPromptText
}
@@ -138,7 +153,7 @@ func extractMessages(value any, wantedRoles ...string) []promptSegment {
}
texts := contentTexts(message["content"])
for _, text := range texts {
result = append(result, promptSegment{text: text, user: role == "user"})
result = append(result, promptSegment{text: text, user: role == "user", role: role})
}
}
return result
@@ -148,7 +163,7 @@ func extractInstructions(value any) []promptSegment {
switch typed := value.(type) {
case string:
if text := strings.TrimSpace(typed); text != "" {
return []promptSegment{{text: text}}
return []promptSegment{{text: text, role: "system"}}
}
case []any:
return systemPromptSegments(contentTexts(typed))
@@ -162,7 +177,7 @@ func extractAnthropicSystem(value any) []promptSegment {
switch typed := value.(type) {
case string:
if text := strings.TrimSpace(typed); text != "" {
return []promptSegment{{text: text}}
return []promptSegment{{text: text, role: "system"}}
}
case []any:
return systemPromptSegments(contentTexts(typed))
@@ -175,13 +190,13 @@ func extractAnthropicSystem(value any) []promptSegment {
func extractResponses(value any) []promptSegment {
switch typed := value.(type) {
case string:
return []promptSegment{{text: typed, user: true}}
return []promptSegment{{text: typed, user: true, role: "user"}}
case []any:
result := make([]promptSegment, 0, len(typed))
for _, item := range typed {
switch entry := item.(type) {
case string:
result = append(result, promptSegment{text: entry, user: true})
result = append(result, promptSegment{text: entry, user: true, role: "user"})
case map[string]any:
role := strings.ToLower(stringValue(entry["role"]))
if role != "" && !isClientInstructionRole(role) {
@@ -189,10 +204,10 @@ func extractResponses(value any) []promptSegment {
}
if content, exists := entry["content"]; exists {
for _, text := range contentTexts(content) {
result = append(result, promptSegment{text: text, user: role == "" || role == "user"})
result = append(result, promptSegment{text: text, user: role == "" || role == "user", role: role})
}
} else if text := stringValue(entry["text"]); text != "" {
result = append(result, promptSegment{text: text, user: role == "" || role == "user"})
result = append(result, promptSegment{text: text, user: role == "" || role == "user", role: role})
}
}
}
@@ -241,7 +256,7 @@ func extractGemini(value any) []promptSegment {
for _, part := range parts {
if object, ok := part.(map[string]any); ok {
if text := stringValue(object["text"]); text != "" {
result = append(result, promptSegment{text: text, user: role == "" || role == "user"})
result = append(result, promptSegment{text: text, user: role == "" || role == "user", role: role})
}
}
}
@@ -278,7 +293,7 @@ func extractGeminiSystemInstruction(value any) []promptSegment {
switch typed := value.(type) {
case string:
if text := strings.TrimSpace(typed); text != "" {
return []promptSegment{{text: text}}
return []promptSegment{{text: text, role: "system"}}
}
case map[string]any:
if parts, ok := typed["parts"].([]any); ok {
@@ -286,7 +301,7 @@ func extractGeminiSystemInstruction(value any) []promptSegment {
for _, part := range parts {
if object, ok := part.(map[string]any); ok {
if text := stringValue(object["text"]); text != "" {
result = append(result, promptSegment{text: text})
result = append(result, promptSegment{text: text, role: "system"})
}
}
}
@@ -297,6 +312,7 @@ func extractGeminiSystemInstruction(value any) []promptSegment {
segments := extractGemini(typed)
for index := range segments {
segments[index].user = false
segments[index].role = "system"
}
return segments
}
@@ -312,7 +328,7 @@ func extractGeminiInstances(value any) []promptSegment {
for _, item := range instances {
if instance, ok := item.(map[string]any); ok {
if prompt := stringValue(instance["prompt"]); prompt != "" {
result = append(result, promptSegment{text: prompt, user: true})
result = append(result, promptSegment{text: prompt, user: true, role: "user"})
}
}
}
@@ -420,19 +436,13 @@ func contentTexts(value any) []string {
}
func normalizeSegmentsLatestUserFirst(values []promptSegment) []string {
normalized := make([]promptSegment, 0, len(values))
for _, value := range values {
value.text = strings.TrimSpace(value.text)
if value.text != "" {
normalized = append(normalized, value)
}
}
normalized := normalizedPromptSegments(values)
if len(normalized) == 0 {
return nil
}
priorityIndex := len(normalized) - 1
for index := len(normalized) - 1; index >= 0; index-- {
if normalized[index].user {
if isUserSegment(normalized[index]) {
priorityIndex = index
break
}
@@ -447,6 +457,85 @@ func normalizeSegmentsLatestUserFirst(values []promptSegment) []string {
return result
}
// blockingSegmentsLatestUserAndPreviousOutput limits synchronous guard input to
// the current user turn and the nearest preceding assistant/model turn. It is
// deliberately opt-in because full transcript scanning remains stronger at
// finding client-controlled content placed in older or non-user messages.
func blockingSegmentsLatestUserAndPreviousOutput(values []promptSegment) []string {
normalized := normalizedPromptSegments(values)
latestUserStart := latestUserSegmentStart(normalized)
if latestUserStart < 0 {
// A request without user content cannot be narrowed safely. Preserve the
// established full-snapshot behavior for unusual protocol payloads.
return normalizeSegmentsLatestUserFirst(values)
}
latestUserEnd := latestUserStart
for latestUserEnd < len(normalized) && isUserSegment(normalized[latestUserEnd]) {
latestUserEnd++
}
currentUserText := make([]string, 0, latestUserEnd-latestUserStart)
for _, segment := range normalized[latestUserStart:latestUserEnd] {
currentUserText = append(currentUserText, segment.text)
}
// A single client turn may have several text content parts. Keep it in one
// priority segment so every part of the latest input is scanned before the
// prior output begins.
selected := []promptSegment{{text: strings.Join(currentUserText, "\n\n"), user: true, role: "user"}}
for index := latestUserStart - 1; index >= 0; index-- {
if !isAssistantOutputSegment(normalized[index]) {
continue
}
start := index
for start > 0 && isAssistantOutputSegment(normalized[start-1]) {
start--
}
selected = append(selected, normalized[start:index+1]...)
break
}
return promptSegmentTexts(selected)
}
func normalizedPromptSegments(values []promptSegment) []promptSegment {
normalized := make([]promptSegment, 0, len(values))
for _, value := range values {
value.text = strings.TrimSpace(value.text)
if value.text != "" {
normalized = append(normalized, value)
}
}
return normalized
}
func latestUserSegmentStart(values []promptSegment) int {
latest := -1
for index := len(values) - 1; index >= 0; index-- {
if isUserSegment(values[index]) {
latest = index
break
}
}
for latest > 0 && isUserSegment(values[latest-1]) {
latest--
}
return latest
}
func isUserSegment(segment promptSegment) bool {
return segment.user || segment.role == "user"
}
func isAssistantOutputSegment(segment promptSegment) bool {
return segment.role == "assistant" || segment.role == "model"
}
func promptSegmentTexts(values []promptSegment) []string {
result := make([]string, 0, len(values))
for _, value := range values {
result = append(result, value.text)
}
return result
}
func buildPrioritizedScanText(segments []string) (scanText string, metadataText string) {
metadataText = strings.Join(segments, "\n\n")
if len(segments) <= 1 {
@@ -458,7 +547,7 @@ func buildPrioritizedScanText(segments []string) (scanText string, metadataText
func promptSegmentsForRole(texts []string, role string) []promptSegment {
result := make([]promptSegment, 0, len(texts))
for _, text := range texts {
result = append(result, promptSegment{text: text, user: role == "" || role == "user"})
result = append(result, promptSegment{text: text, user: role == "" || role == "user", role: role})
}
return result
}
@@ -293,6 +293,65 @@ func TestPromptSnapshotIncludesClientControlledInstructions(t *testing.T) {
}
}
func TestBlockingPromptSnapshotLimitsInputToLatestUserAndPreviousOutput(t *testing.T) {
tests := []struct {
name, protocol, body, want string
omitted []string
}{
{
name: "chat keeps multipart latest user and prior assistant",
protocol: "openai_chat_completions",
body: `{"messages":[
{"role":"system","content":"system instruction"},
{"role":"user","content":"older user input"},
{"role":"assistant","content":"older assistant output"},
{"role":"tool","content":"tool payload"},
{"role":"assistant","content":"previous assistant output"},
{"role":"user","content":[{"type":"text","text":"latest user first part"},{"type":"text","text":"latest user second part"}]}
]}`,
want: "latest user first part\n\nlatest user second part" + promptAuditPrioritySeparator + "previous assistant output",
omitted: []string{"system instruction", "older user input", "older assistant output", "tool payload"},
},
{
name: "gemini keeps prior model output",
protocol: "gemini",
body: `{"systemInstruction":{"parts":[{"text":"system instruction"}]},"contents":[
{"role":"user","parts":[{"text":"older user input"}]},
{"role":"model","parts":[{"text":"previous model output"}]},
{"role":"user","parts":[{"text":"latest user input"}]}
]}`,
want: "latest user input" + promptAuditPrioritySeparator + "previous model output",
omitted: []string{"system instruction", "older user input"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
snapshot, err := ExtractBlockingPromptSnapshot(Request{Protocol: tt.protocol, Body: []byte(tt.body)}, true)
require.NoError(t, err)
require.Equal(t, tt.want, snapshot.ScanText)
for _, omitted := range tt.omitted {
require.NotContains(t, snapshot.ScanText, omitted)
}
})
}
}
func TestBlockingPromptSnapshotPreservesFullScopeByDefaultAndWithoutUserInput(t *testing.T) {
req := Request{Protocol: "openai_chat_completions", Body: []byte(`{"messages":[{"role":"system","content":"system instruction"},{"role":"user","content":"older user input"},{"role":"assistant","content":"previous output"},{"role":"user","content":"latest user input"}]}`)}
full, err := ExtractPromptSnapshot(req)
require.NoError(t, err)
defaultBlocking, err := ExtractBlockingPromptSnapshot(req, false)
require.NoError(t, err)
require.Equal(t, full, defaultBlocking)
noUser := Request{Protocol: "openai_chat_completions", Body: []byte(`{"messages":[{"role":"system","content":"system instruction"},{"role":"assistant","content":"assistant output"}]}`)}
fullWithoutUser, err := ExtractPromptSnapshot(noUser)
require.NoError(t, err)
narrowWithoutUser, err := ExtractBlockingPromptSnapshot(noUser, true)
require.NoError(t, err)
require.Equal(t, fullWithoutUser, narrowWithoutUser)
}
func TestBuildPromptPreviewWithholdsMajorityOfOrdinaryText(t *testing.T) {
prompt := strings.Repeat("机密业务提示词内容", 40)
preview := BuildPromptPreview(prompt, DefaultPromptPreviewMaxRunes)