Skip to content

Commit a001ff9

Browse files
jkyberneeesjkyberneees
committed
fix(extended-memory): anaphora resolution, close race, recent-message predictor wiring, and context propagation
- Replace only the first pronoun when the top atom's semantic score is above SemanticSearchMinScore; return (resolved, bool). - Guard triggerBackgroundInference with a closed flag and lock to prevent starting goroutines after Close. - Add a 10-message recent-user buffer and pass it plus UserState to the Predictor in Recall.Query. - Thread context.Context through AnaphoraResolve, FormatExtendedContext, FormatUserStateContext, and FormatReturnAfterBreak; update loop engine callbacks and MemoryManager callers to use caller-provided contexts with timeouts. - Add TestExtendedMemoryCloseRace.
1 parent 07f706a commit a001ff9

6 files changed

Lines changed: 159 additions & 67 deletions

File tree

‎internal/loop/loop.go‎

Lines changed: 21 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -68,12 +68,12 @@ type EpisodeContextFunc func(userInput string) string
6868
// ExtendedMemoryContextFunc is an optional callback that returns formatted
6969
// Extended Memory context for the latest user input. It is injected as a
7070
// system message after the legacy memory prompt block.
71-
type ExtendedMemoryContextFunc func(userInput string) string
71+
type ExtendedMemoryContextFunc func(ctx context.Context, userInput string) string
7272

7373
// UserMessageHandler is an optional callback invoked once per new user
7474
// message. It is used by callers (e.g. odek.New) to trigger Extended Memory
7575
// atom extraction.
76-
type UserMessageHandler func(msg string)
76+
type UserMessageHandler func(ctx context.Context, msg string)
7777

7878
// ToolEventHandler is an optional callback invoked for each tool execution
7979
// during the agent loop — fires before (tool_call) and after (tool_result)
@@ -103,23 +103,23 @@ type IterationCallback func(info IterationInfo)
103103

104104
// Engine runs the agent loop: observe → think → act → repeat.
105105
type Engine struct {
106-
client *llm.Client
107-
registry *tool.Registry
108-
renderer *render.Renderer // optional: colored terminal output
109-
maxIter int
110-
system string
111-
baseSystem string // original system message without memory/skills
112-
maxContext int // max context tokens (0 = no limit)
113-
skillLoader SkillLoader // optional: loads matching skills
114-
lastSkillMsg string // last user message that triggered skill loading (dedup)
115-
lastEpiMsg string // last user message that triggered episode search (dedup)
116-
lastExtMsg string // last user message that triggered extended memory search (dedup)
117-
lastUserMsg string // last user message passed to userMsgHandler (dedup)
118-
skillVerbose bool // show full skill banners (default: condensed)
119-
episodeCtx EpisodeContextFunc // optional: per-turn episode search
120-
extendedCtx ExtendedMemoryContextFunc // optional: per-turn extended memory search
121-
userMsgHandler UserMessageHandler // optional: called once per new user message
122-
wrapUntrusted func(source, content string) string // optional: wraps skill/episode content
106+
client *llm.Client
107+
registry *tool.Registry
108+
renderer *render.Renderer // optional: colored terminal output
109+
maxIter int
110+
system string
111+
baseSystem string // original system message without memory/skills
112+
maxContext int // max context tokens (0 = no limit)
113+
skillLoader SkillLoader // optional: loads matching skills
114+
lastSkillMsg string // last user message that triggered skill loading (dedup)
115+
lastEpiMsg string // last user message that triggered episode search (dedup)
116+
lastExtMsg string // last user message that triggered extended memory search (dedup)
117+
lastUserMsg string // last user message passed to userMsgHandler (dedup)
118+
skillVerbose bool // show full skill banners (default: condensed)
119+
episodeCtx EpisodeContextFunc // optional: per-turn episode search
120+
extendedCtx ExtendedMemoryContextFunc // optional: per-turn extended memory search
121+
userMsgHandler UserMessageHandler // optional: called once per new user message
122+
wrapUntrusted func(source, content string) string // optional: wraps skill/episode content
123123

124124
toolEventHandler ToolEventHandler // optional: fires during tool execution
125125
signalHandler SignalHandler // optional: fires on internal loop signals
@@ -654,7 +654,7 @@ func (e *Engine) runLoop(ctx context.Context, messages []llm.Message) (string, [
654654
if e.userMsgHandler != nil {
655655
if userMsg := lastUserMessage(messages); userMsg != "" && userMsg != e.lastUserMsg {
656656
e.lastUserMsg = userMsg
657-
e.userMsgHandler(userMsg)
657+
e.userMsgHandler(ctx, userMsg)
658658
}
659659
}
660660

@@ -767,7 +767,7 @@ func (e *Engine) runLoop(ctx context.Context, messages []llm.Message) (string, [
767767
// user messages.
768768
if e.extendedCtx != nil {
769769
if userMsg := lastUserMessage(messages); userMsg != "" && userMsg != e.lastExtMsg {
770-
if extContext := e.extendedCtx(userMsg); extContext != "" {
770+
if extContext := e.extendedCtx(ctx, userMsg); extContext != "" {
771771
wrapped := extContext
772772
if e.wrapUntrusted != nil {
773773
wrapped = e.wrapUntrusted("extended_memory", extContext)

‎internal/memory/extended/extended_memory.go‎

Lines changed: 72 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,16 @@ type ExtendedMemory struct {
4141
closeOnce sync.Once
4242
pendingWg sync.WaitGroup
4343
userStateTurns int
44+
45+
recentUserMessages []string
46+
recentMu sync.Mutex
47+
48+
inferenceMu sync.Mutex
49+
closed bool
4450
}
4551

52+
const recentUserMessageLimit = 10
53+
4654
// New creates an ExtendedMemory instance rooted at dir.
4755
func New(dir string, llm LLMClient, cfg Config) *ExtendedMemory {
4856
cfg = Resolve(cfg)
@@ -268,8 +276,15 @@ func (em *ExtendedMemory) triggerBackgroundInference() {
268276
if em == nil || !em.Enabled() || em.userModel == nil || !em.userModel.Enabled() {
269277
return
270278
}
271-
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
279+
em.inferenceMu.Lock()
280+
if em.closed {
281+
em.inferenceMu.Unlock()
282+
return
283+
}
272284
em.pendingWg.Add(1)
285+
em.inferenceMu.Unlock()
286+
287+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
273288
go func() {
274289
defer func() {
275290
cancel()
@@ -313,7 +328,7 @@ func (em *ExtendedMemory) ListPendingReview() ([]PendingReview, error) {
313328
}
314329

315330
// FormatUserStateContext returns formatted user-model context.
316-
func (em *ExtendedMemory) FormatUserStateContext() string {
331+
func (em *ExtendedMemory) FormatUserStateContext(ctx context.Context) string {
317332
if em == nil || !em.Enabled() || em.userModel == nil {
318333
return ""
319334
}
@@ -369,23 +384,44 @@ User model: %s`, recentJSON, stateJSON)
369384

370385
var pronounRE = regexp.MustCompile(`(?i)\b(it|that|this|them|those)\b`)
371386

372-
// AnaphoraResolve replaces pronouns in a user message with the most likely
373-
// antecedent from recent trusted atoms when the semantic score is high enough.
374-
func (em *ExtendedMemory) AnaphoraResolve(msg string) string {
387+
// AnaphoraResolve replaces the first pronoun in a user message with the
388+
// most likely antecedent from recent trusted atoms when the semantic score
389+
// is high enough. It returns the resolved message and true when a replacement
390+
// happened, otherwise the original message and false.
391+
func (em *ExtendedMemory) AnaphoraResolve(ctx context.Context, msg string) (string, bool) {
375392
if em == nil || !em.Enabled() || em.recall == nil || em.cfg.AnaphoraResolutionEnabled == nil || !*em.cfg.AnaphoraResolutionEnabled {
376-
return msg
393+
return msg, false
377394
}
378395
if !pronounRE.MatchString(msg) {
379-
return msg
396+
return msg, false
380397
}
381-
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
398+
ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
382399
defer cancel()
400+
383401
atoms, err := em.recall.queryAtoms(ctx, msg)
384402
if err != nil || len(atoms) == 0 {
385-
return msg
403+
return msg, false
386404
}
387-
resolved := pronounRE.ReplaceAllString(msg, atoms[0].Text)
388-
return resolved
405+
406+
// Match the top returned atom to its vector similarity score.
407+
candidates := em.index.search(msg, em.cfg.SemanticSearchTopK*em.cfg.SemanticSearchOverfetch)
408+
var topScore float32
409+
for _, c := range candidates {
410+
if c.ID == atoms[0].ID {
411+
topScore = c.Score
412+
break
413+
}
414+
}
415+
if topScore < em.cfg.SemanticSearchMinScore {
416+
return msg, false
417+
}
418+
419+
loc := pronounRE.FindStringIndex(msg)
420+
if loc == nil {
421+
return msg, false
422+
}
423+
resolved := msg[:loc[0]] + atoms[0].Text + msg[loc[1]:]
424+
return resolved, true
389425
}
390426

391427
// SearchAtoms performs an explicit semantic search and returns ranked atoms.
@@ -448,13 +484,19 @@ func (em *ExtendedMemory) PinAtom(id string) error {
448484

449485
// FormatExtendedContext returns formatted Extended Memory context for the
450486
// query, or empty string if nothing matches or Extended Memory is disabled.
451-
func (em *ExtendedMemory) FormatExtendedContext(query string) string {
487+
func (em *ExtendedMemory) FormatExtendedContext(ctx context.Context, query string) string {
452488
if em == nil || !em.Enabled() {
453489
return ""
454490
}
455-
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
491+
ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
456492
defer cancel()
457-
context, err := em.recall.Query(ctx, query)
493+
494+
em.recentMu.Lock()
495+
recent := make([]string, len(em.recentUserMessages))
496+
copy(recent, em.recentUserMessages)
497+
em.recentMu.Unlock()
498+
499+
context, err := em.recall.Query(ctx, query, recent, em.userModel.State())
458500
if err != nil {
459501
log.Printf("extended memory: format context failed: %v", err)
460502
return ""
@@ -464,17 +506,15 @@ func (em *ExtendedMemory) FormatExtendedContext(query string) string {
464506

465507
// FormatContext is an alias for FormatExtendedContext.
466508
func (em *ExtendedMemory) FormatContext(ctx context.Context, query string) string {
467-
return em.FormatExtendedContext(query)
509+
return em.FormatExtendedContext(ctx, query)
468510
}
469511

470512
// OnUserMessage extracts atoms from a user message and stores them.
471513
func (em *ExtendedMemory) OnUserMessage(ctx AtomContext, msg string) {
472514
if em == nil || !em.Enabled() {
473515
return
474516
}
475-
if em.cfg.AutoExtractPerTurn == nil || !*em.cfg.AutoExtractPerTurn {
476-
return
477-
}
517+
478518
em.mu.Lock()
479519
em.lastUser = msg
480520
if ctx.SessionID != "" {
@@ -490,6 +530,17 @@ func (em *ExtendedMemory) OnUserMessage(ctx AtomContext, msg string) {
490530
}
491531
em.mu.Unlock()
492532

533+
em.recentMu.Lock()
534+
em.recentUserMessages = append(em.recentUserMessages, msg)
535+
if len(em.recentUserMessages) > recentUserMessageLimit {
536+
em.recentUserMessages = em.recentUserMessages[len(em.recentUserMessages)-recentUserMessageLimit:]
537+
}
538+
em.recentMu.Unlock()
539+
540+
if em.cfg.AutoExtractPerTurn == nil || !*em.cfg.AutoExtractPerTurn {
541+
return
542+
}
543+
493544
c, cancel := context.WithTimeout(context.Background(), 30*time.Second)
494545
defer cancel()
495546
atoms, err := em.extractor.Extract(c, msg)
@@ -613,6 +664,9 @@ func (em *ExtendedMemory) Close() error {
613664
return nil
614665
}
615666
em.closeOnce.Do(func() {
667+
em.inferenceMu.Lock()
668+
em.closed = true
669+
em.inferenceMu.Unlock()
616670
em.index.Wait()
617671
em.pendingWg.Wait()
618672
})

‎internal/memory/extended/extended_memory_test.go‎

Lines changed: 36 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -974,7 +974,7 @@ func TestExtendedMemoryFormatUserStateContext(t *testing.T) {
974974
em.OnUserMessage(AtomContext{SessionID: "s1", Turn: 1}, "Use dry tone, I code in Go")
975975
em.Close()
976976

977-
ctx := em.FormatUserStateContext()
977+
ctx := em.FormatUserStateContext(context.Background())
978978
if ctx == "" {
979979
t.Error("expected non-empty user state context")
980980
}
@@ -1036,7 +1036,10 @@ func TestExtendedMemoryAnaphoraResolve(t *testing.T) {
10361036
_ = em.AddAtom(context.Background(), MemoryAtom{Text: "Postgres database", SourceClass: SourceUserSaid, Type: TypeFact})
10371037
em.index.Compact()
10381038

1039-
resolved := em.AnaphoraResolve("How do I configure it?")
1039+
resolved, ok := em.AnaphoraResolve(context.Background(), "How do I configure it?")
1040+
if !ok {
1041+
t.Fatal("expected anaphora resolution to replace pronoun")
1042+
}
10401043
if !strings.Contains(resolved, "Postgres database") {
10411044
t.Errorf("expected anaphora resolution to replace pronoun, got %q", resolved)
10421045
}
@@ -1050,8 +1053,8 @@ func TestExtendedMemoryAnaphoraResolveDisabled(t *testing.T) {
10501053
em := New(dir, newMockLLM(), cfg)
10511054
defer em.Close()
10521055
msg := "How do I configure it?"
1053-
if got := em.AnaphoraResolve(msg); got != msg {
1054-
t.Errorf("expected unchanged message when disabled, got %q", got)
1056+
if got, ok := em.AnaphoraResolve(context.Background(), msg); got != msg || ok {
1057+
t.Errorf("expected unchanged message when disabled, got %q (ok=%v)", got, ok)
10551058
}
10561059
}
10571060

@@ -1060,10 +1063,10 @@ func TestExtendedMemoryNilSafeMethods(t *testing.T) {
10601063
if em.UserStateStyle() != nil {
10611064
t.Error("expected nil UserStateStyle on nil em")
10621065
}
1063-
if em.FormatUserStateContext() != "" {
1066+
if em.FormatUserStateContext(context.Background()) != "" {
10641067
t.Error("expected empty FormatUserStateContext on nil em")
10651068
}
1066-
if em.AnaphoraResolve("x") != "x" {
1069+
if got, ok := em.AnaphoraResolve(context.Background(), "x"); got != "x" || ok {
10671070
t.Error("expected AnaphoraResolve passthrough on nil em")
10681071
}
10691072
if em.ReturnAfterBreak(context.Background()) != "" {
@@ -1073,3 +1076,30 @@ func TestExtendedMemoryNilSafeMethods(t *testing.T) {
10731076
t.Error("expected ListPendingReview no error on nil em")
10741077
}
10751078
}
1079+
1080+
func TestExtendedMemoryCloseRace(t *testing.T) {
1081+
dir := t.TempDir()
1082+
cfg := DefaultConfig()
1083+
cfg.Enabled = boolPtr(true)
1084+
cfg.UserStateTurnInterval = 1
1085+
em := New(dir, newMockLLM(), cfg)
1086+
defer em.Close()
1087+
em.SetSessionContext("s1", "p1")
1088+
1089+
var wg sync.WaitGroup
1090+
for i := 0; i < 20; i++ {
1091+
wg.Add(1)
1092+
go func() {
1093+
defer wg.Done()
1094+
for j := 0; j < 20; j++ {
1095+
em.OnUserMessage(AtomContext{SessionID: "s1", Turn: j}, "race test message")
1096+
}
1097+
}()
1098+
}
1099+
go func() {
1100+
for i := 0; i < 50; i++ {
1101+
_ = em.Close()
1102+
}
1103+
}()
1104+
wg.Wait()
1105+
}

‎internal/memory/extended/recall.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,11 @@ type QueryResult struct {
3737

3838
// Query searches for relevant atoms. It returns a formatted context string
3939
// bounded by MemoryBudgetChars, or empty string if nothing matches.
40-
func (r *Recall) Query(ctx context.Context, query string) (string, error) {
40+
func (r *Recall) Query(ctx context.Context, query string, recent []string, state UserState) (string, error) {
4141
if r.store == nil || r.index == nil {
4242
return "", nil
4343
}
44-
res, err := r.queryAtomsWithPrediction(ctx, query)
44+
res, err := r.queryAtomsWithPrediction(ctx, query, recent, state)
4545
if err != nil {
4646
log.Printf("extended memory: recall query failed: %v", err)
4747
return "", err
@@ -54,7 +54,7 @@ func (r *Recall) Query(ctx context.Context, query string) (string, error) {
5454

5555
// queryAtomsWithPrediction unions literal-query results with predicted-intent
5656
// results when prediction is enabled.
57-
func (r *Recall) queryAtomsWithPrediction(ctx context.Context, query string) ([]MemoryAtom, error) {
57+
func (r *Recall) queryAtomsWithPrediction(ctx context.Context, query string, recent []string, state UserState) ([]MemoryAtom, error) {
5858
all := make(map[string]MemoryAtom)
5959
literal, err := r.queryAtoms(ctx, query)
6060
if err != nil {
@@ -66,7 +66,7 @@ func (r *Recall) queryAtomsWithPrediction(ctx context.Context, query string) ([]
6666

6767
if r.predictor != nil && r.cfg.PredictiveIntents > 0 &&
6868
r.cfg.FollowUpAnticipationEnabled != nil && *r.cfg.FollowUpAnticipationEnabled {
69-
intents, err := r.predictor.Predict(ctx, query, nil, UserState{})
69+
intents, err := r.predictor.Predict(ctx, query, recent, state)
7070
if err != nil {
7171
log.Printf("extended memory: predicted-intent generation failed: %v", err)
7272
}

0 commit comments

Comments
 (0)