Skip to content

Commit bf2a46c

Browse files
committed
fix(extended-memory): scan user model state on load and fix quarantine promote lock
- Scan every string value in UserState on Load; drop fields/entries that fail ScanContent so a tampered user_model.json cannot poison prompts. - Scan the formatted Summary() output before returning it. - Use Lock instead of RLock in Quarantine.Promote because loadLocked may call saveLocked during eviction.
1 parent a001ff9 commit bf2a46c

2 files changed

Lines changed: 68 additions & 4 deletions

File tree

internal/memory/extended/quarantine.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -142,8 +142,8 @@ func (q *Quarantine) Promote(id string) (MemoryAtom, error) {
142142
if err := session.ValidateSessionID(id); err != nil {
143143
return MemoryAtom{}, fmt.Errorf("extended quarantine: invalid atom id: %w", err)
144144
}
145-
q.mu.RLock()
146-
defer q.mu.RUnlock()
145+
q.mu.Lock()
146+
defer q.mu.Unlock()
147147

148148
entries, err := q.loadLocked()
149149
if err != nil {

internal/memory/extended/usermodel.go

Lines changed: 66 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,9 @@ func (u *UserModel) Enabled() bool {
9898
}
9999

100100
// Load reads the persisted user model, if any. Missing files are not errors.
101+
// Loaded string values are scanned for injection patterns; fields that fail
102+
// the scan are dropped so a tampered user_model.json cannot poison the
103+
// system prompt.
101104
func (u *UserModel) Load() error {
102105
if u == nil || u.store == nil {
103106
return nil
@@ -106,6 +109,7 @@ func (u *UserModel) Load() error {
106109
if err != nil {
107110
return err
108111
}
112+
state = scanUserState(state)
109113
u.mu.Lock()
110114
defer u.mu.Unlock()
111115
u.state = state
@@ -384,6 +388,59 @@ func appendUnique(base, add []string) []string {
384388
return base
385389
}
386390

391+
// scanUserState scans every string value in a loaded UserState and clears
392+
// any field or slice entry that fails the content scan. This prevents a
393+
// tampered user_model.json from injecting instructions into the system prompt.
394+
func scanUserState(s UserState) UserState {
395+
if ScanContent(s.Style.Verbosity) != nil {
396+
s.Style.Verbosity = ""
397+
}
398+
if ScanContent(s.Style.Humor) != nil {
399+
s.Style.Humor = ""
400+
}
401+
if ScanContent(s.Style.Formality) != nil {
402+
s.Style.Formality = ""
403+
}
404+
if ScanContent(s.Style.ExplanationDepth) != nil {
405+
s.Style.ExplanationDepth = ""
406+
}
407+
if ScanContent(s.Style.Tone) != nil {
408+
s.Style.Tone = ""
409+
}
410+
411+
s.Technical.Languages = filterScanned(s.Technical.Languages)
412+
s.Technical.Patterns = filterScanned(s.Technical.Patterns)
413+
s.Technical.Tools = filterScanned(s.Technical.Tools)
414+
415+
if ScanContent(s.CurrentFocus.Project) != nil {
416+
s.CurrentFocus.Project = ""
417+
}
418+
if ScanContent(s.CurrentFocus.Task) != nil {
419+
s.CurrentFocus.Task = ""
420+
}
421+
if ScanContent(s.CurrentFocus.Blocker) != nil {
422+
s.CurrentFocus.Blocker = ""
423+
}
424+
425+
s.InteractionPatterns.CommonOpeners = filterScanned(s.InteractionPatterns.CommonOpeners)
426+
if ScanContent(s.InteractionPatterns.FollowupAfterRefactor) != nil {
427+
s.InteractionPatterns.FollowupAfterRefactor = ""
428+
}
429+
if ScanContent(s.InteractionPatterns.FollowupAfterBugfix) != nil {
430+
s.InteractionPatterns.FollowupAfterBugfix = ""
431+
}
432+
433+
var pending []PendingReview
434+
for _, p := range s.PendingReview {
435+
if ScanContent(p.Field) != nil || ScanContent(p.Value) != nil || ScanContent(p.Evidence) != nil {
436+
continue
437+
}
438+
pending = append(pending, p)
439+
}
440+
s.PendingReview = pending
441+
return s
442+
}
443+
387444
// ConfirmPendingReview applies a pending review to the model and persists it.
388445
func (u *UserModel) ConfirmPendingReview(id string) error {
389446
if u == nil {
@@ -451,7 +508,9 @@ func (u *UserModel) State() UserState {
451508
return u.state
452509
}
453510

454-
// Summary formats the user model for system-prompt injection.
511+
// Summary formats the user model for system-prompt injection. The formatted
512+
// output is scanned before being returned; if it fails the scan, an empty
513+
// string is returned so a poisoned value cannot reach the system prompt.
455514
func (u *UserModel) Summary() string {
456515
if u == nil {
457516
return ""
@@ -524,7 +583,12 @@ func (u *UserModel) Summary() string {
524583
}
525584
}
526585
b.WriteString("────────────────────\n")
527-
return b.String()
586+
summary := b.String()
587+
if err := ScanContent(summary); err != nil {
588+
log.Printf("extended memory: user-model summary rejected by scan: %v", err)
589+
return ""
590+
}
591+
return summary
528592
}
529593

530594
func applyPendingValue(state *UserState, p PendingReview) {

0 commit comments

Comments
 (0)