mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-28 16:09:34 +00:00
feat(security-audit): add narrow blocking audit scope
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user