From 3d31e12795b71c74d4a3b11ca658548b5e75ae5d Mon Sep 17 00:00:00 2001 From: blak0p Date: Tue, 25 Aug 2026 10:45:24 +0200 Subject: [PATCH 1/2] feat(store): add atomic find and replace updates --- DOCS.md | 21 ++++- internal/mcp/mcp.go | 19 +++- internal/mcp/mcp_test.go | 73 ++++++++++++++++ internal/server/server.go | 12 ++- internal/server/server_test.go | 87 +++++++++++++++++++ internal/store/store.go | 31 +++++-- internal/store/store_test.go | 154 +++++++++++++++++++++++++++++++++ 7 files changed, 385 insertions(+), 12 deletions(-) diff --git a/DOCS.md b/DOCS.md index d9b90e110..fed8dfd97 100644 --- a/DOCS.md +++ b/DOCS.md @@ -126,7 +126,12 @@ Engram is local-first: local SQLite is authoritative; cloud features are optiona - `POST /observations` — Add observation. Body: `{session_id, type, title, content, tool_name?, project?, scope?, topic_key?}` - `GET /observations/recent` — Recent observations. Query: `?project=X&scope=project|personal&limit=N` - `GET /observations/{id}` — Get single observation by ID -- `PATCH /observations/{id}` — Update fields. Body: `{title?, content?, type?, project?, scope?, topic_key?}` +- `PATCH /observations/{id}` — Update fields. Body: `{title?, content?, find?, replace?, type?, project?, scope?, topic_key?}` + - `find` and `replace` provide an atomic literal find/replace against the observation's current content. A complete pair replaces every literal occurrence; regular expressions, case-insensitive matching, and replacement limits are not supported. + - `find` and `replace` must be supplied together. Supplying only one (including an empty string) returns `400 {"error":"find and replace must be used together"}`. + - `find/replace` is mutually exclusive with `content`. Supplying all three returns `400 {"error":"find/replace is mutually exclusive with content"}`. + - An empty or non-matching `find` leaves content unchanged, but still performs the normal update: revision and sync semantics remain intact, so the revision advances and an observation-upsert sync mutation is recorded. + - Example: `PATCH /observations/7` with `{"find":"alpha","replace":"beta"}` changes `alpha alpha` to `beta beta` in one transaction. - `DELETE /observations/{id}` — Delete observation (`?hard=true` for hard delete, soft delete by default) - `200` when deleted - `404` when observation does not exist @@ -419,7 +424,19 @@ When `topic_key` is provided, `mem_save` upserts the latest observation in the s ### mem_update -Update an observation by ID. Supports partial updates for `title`, `content`, `type`, `project`, `scope`, and `topic_key`. +Update an observation by ID. Supports partial updates for `title`, `content`, `find`, `replace`, `type`, `project`, `scope`, and `topic_key`. + +`find` and `replace` atomically replace every literal occurrence in the stored content. Both fields are required together; a lone field returns exactly `find and replace must be used together`. A complete pair cannot be sent with `content`; that combination returns exactly `find/replace is mutually exclusive with content`. An empty or non-matching `find` is a content no-op, while preserving the normal revision and sync side effects. + +Example: + +```json +{ + "id": 7, + "find": "alpha", + "replace": "beta" +} +``` ### mem_suggest_topic_key diff --git a/internal/mcp/mcp.go b/internal/mcp/mcp.go index 7b372154a..da00646bb 100644 --- a/internal/mcp/mcp.go +++ b/internal/mcp/mcp.go @@ -334,7 +334,7 @@ Examples: if shouldRegister("mem_update", allowlist) { srv.AddTool( mcp.NewTool("mem_update", - mcp.WithDescription("Update an existing observation by ID. Only provided fields are changed."), + mcp.WithDescription("Update an existing observation by ID. Only provided fields are changed. find and replace must be supplied together and cannot be combined with content."), mcp.WithDeferLoading(true), mcp.WithTitleAnnotation("Update Memory"), mcp.WithReadOnlyHintAnnotation(false), @@ -351,6 +351,12 @@ Examples: mcp.WithString("content", mcp.Description("New content"), ), + mcp.WithString("find", + mcp.Description("Literal text to replace everywhere in the current content; requires replace"), + ), + mcp.WithString("replace", + mcp.Description("Replacement text for every literal find occurrence; requires find"), + ), mcp.WithString("type", mcp.Description("New type/category"), ), @@ -1066,6 +1072,12 @@ func handleUpdate(s *store.Store) server.ToolHandlerFunc { if v, ok := req.GetArguments()["content"].(string); ok { update.Content = &v } + if v, ok := req.GetArguments()["find"].(string); ok { + update.Find = &v + } + if v, ok := req.GetArguments()["replace"].(string); ok { + update.Replace = &v + } if v, ok := req.GetArguments()["type"].(string); ok { update.Type = &v } @@ -1080,7 +1092,7 @@ func handleUpdate(s *store.Store) server.ToolHandlerFunc { update.TopicKey = &v } - if update.Title == nil && update.Content == nil && update.Type == nil && update.Project == nil && update.Scope == nil && update.TopicKey == nil { + if update.Title == nil && update.Content == nil && update.Find == nil && update.Replace == nil && update.Type == nil && update.Project == nil && update.Scope == nil && update.TopicKey == nil { return mcp.NewToolResultError("provide at least one field to update"), nil } @@ -1091,6 +1103,9 @@ func handleUpdate(s *store.Store) server.ToolHandlerFunc { obs, err := s.UpdateObservation(id, update) if err != nil { + if errors.Is(err, store.ErrFindReplacePairRequired) || errors.Is(err, store.ErrFindReplaceWithContent) { + return mcp.NewToolResultError(err.Error()), nil + } return mcp.NewToolResultError("Failed to update memory: " + err.Error()), nil } diff --git a/internal/mcp/mcp_test.go b/internal/mcp/mcp_test.go index 0c42e9423..2ad0e066a 100644 --- a/internal/mcp/mcp_test.go +++ b/internal/mcp/mcp_test.go @@ -823,6 +823,79 @@ func TestHandleUpdateAcceptsAllOptionalFields(t *testing.T) { } } +func TestHandleUpdateFindReplaceContract(t *testing.T) { + s := newMCPTestStore(t) + if err := s.CreateSession("mcp-find-replace", "engram", "/tmp/engram"); err != nil { + t.Fatalf("create session: %v", err) + } + id, err := s.AddObservation(store.AddObservationParams{ + SessionID: "mcp-find-replace", + Type: "note", + Title: "replacement", + Content: "alpha alpha", + Project: "engram", + }) + if err != nil { + t.Fatalf("add observation: %v", err) + } + + h := handleUpdate(s) + updated, err := h(context.Background(), mcppkg.CallToolRequest{Params: mcppkg.CallToolParams{Arguments: map[string]any{ + "id": float64(id), + "find": "alpha", + "replace": "beta", + }}}) + if err != nil { + t.Fatalf("replace handler error: %v", err) + } + if updated.IsError { + t.Fatalf("expected replacement success, got %q", callResultText(t, updated)) + } + observation, err := s.GetObservation(id) + if err != nil { + t.Fatalf("get replaced observation: %v", err) + } + if observation.Content != "beta beta" { + t.Fatalf("expected all occurrences replaced, got %q", observation.Content) + } + + for _, args := range []map[string]any{ + {"id": float64(id), "find": ""}, + {"id": float64(id), "replace": "beta"}, + {"id": float64(id), "find": "beta", "replace": "gamma", "content": "replacement"}, + } { + res, err := h(context.Background(), mcppkg.CallToolRequest{Params: mcppkg.CallToolParams{Arguments: args}}) + if err != nil { + t.Fatalf("invalid handler error: %v", err) + } + if !res.IsError { + t.Fatalf("expected tool error for %#v", args) + } + got := callResultText(t, res) + want := "find and replace must be used together" + if args["content"] != nil { + want = "find/replace is mutually exclusive with content" + } + if got != want { + t.Fatalf("expected exact error %q, got %q", want, got) + } + } +} + +func TestMemUpdateSchemaIncludesFindAndReplace(t *testing.T) { + srv := NewServer(newMCPTestStore(t)) + tool := srv.GetTool("mem_update") + if tool == nil { + t.Fatal("mem_update tool not registered") + } + + for _, property := range []string{"find", "replace"} { + if _, ok := tool.Tool.InputSchema.Properties[property]; !ok { + t.Errorf("mem_update schema missing %q property", property) + } + } +} + func TestHandleContextWithSessionOnlyUsesNoneProjects(t *testing.T) { s := newMCPTestStore(t) if err := s.CreateSession("s-context-none", "engram", "/tmp/engram"); err != nil { diff --git a/internal/server/server.go b/internal/server/server.go index 458b1f1e1..3a06e5d4d 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -5,6 +5,7 @@ package server import ( + "database/sql" "encoding/json" "errors" "fmt" @@ -320,14 +321,21 @@ func (s *Server) handleUpdateObservation(w http.ResponseWriter, r *http.Request) return } - if body.Type == nil && body.Title == nil && body.Content == nil && body.Project == nil && body.Scope == nil && body.TopicKey == nil { + if body.Type == nil && body.Title == nil && body.Content == nil && body.Find == nil && body.Replace == nil && body.Project == nil && body.Scope == nil && body.TopicKey == nil { jsonError(w, http.StatusBadRequest, "at least one field is required") return } obs, err := s.store.UpdateObservation(id, body) if err != nil { - jsonError(w, http.StatusNotFound, err.Error()) + switch { + case errors.Is(err, store.ErrFindReplacePairRequired), errors.Is(err, store.ErrFindReplaceWithContent): + jsonError(w, http.StatusBadRequest, err.Error()) + case errors.Is(err, sql.ErrNoRows), errors.Is(err, store.ErrObservationNotFound): + jsonError(w, http.StatusNotFound, err.Error()) + default: + jsonError(w, http.StatusInternalServerError, err.Error()) + } return } diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 02007852a..dc4711435 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -8,6 +8,7 @@ import ( "net" "net/http" "net/http/httptest" + "os" "strings" "sync/atomic" "testing" @@ -234,6 +235,92 @@ func TestExportRejectsExplicitBlankProjectQuery(t *testing.T) { } } +func TestHandleUpdateObservationFindReplaceContract(t *testing.T) { + st := newServerTestStore(t) + if err := st.CreateSession("http-find-replace", "engram", "/tmp/engram"); err != nil { + t.Fatalf("create session: %v", err) + } + id, err := st.AddObservation(store.AddObservationParams{ + SessionID: "http-find-replace", + Type: "note", + Title: "replacement", + Content: "alpha alpha", + Project: "engram", + }) + if err != nil { + t.Fatalf("add observation: %v", err) + } + + srv := New(st, 0) + h := srv.Handler() + + replaceReq := httptest.NewRequest(http.MethodPatch, fmt.Sprintf("/observations/%d", id), strings.NewReader(`{"find":"alpha","replace":"beta"}`)) + replaceRec := httptest.NewRecorder() + h.ServeHTTP(replaceRec, replaceReq) + if replaceRec.Code != http.StatusOK { + t.Fatalf("expected replacement status 200, got %d: %s", replaceRec.Code, replaceRec.Body.String()) + } + var updated store.Observation + if err := json.NewDecoder(replaceRec.Body).Decode(&updated); err != nil { + t.Fatalf("decode replacement response: %v", err) + } + if updated.Content != "beta beta" { + t.Fatalf("expected all occurrences replaced, got %q", updated.Content) + } + + tests := []struct { + name string + body string + want string + }{ + {name: "find only", body: `{"find":""}`, want: "find and replace must be used together"}, + {name: "replace only", body: `{"replace":"gamma"}`, want: "find and replace must be used together"}, + {name: "pair with content", body: `{"find":"beta","replace":"gamma","content":"replacement"}`, want: "find/replace is mutually exclusive with content"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req := httptest.NewRequest(http.MethodPatch, fmt.Sprintf("/observations/%d", id), strings.NewReader(tt.body)) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusBadRequest { + t.Fatalf("expected status 400, got %d: %s", rec.Code, rec.Body.String()) + } + var response map[string]string + if err := json.NewDecoder(rec.Body).Decode(&response); err != nil { + t.Fatalf("decode error response: %v", err) + } + if response["error"] != tt.want { + t.Fatalf("expected exact error %q, got %q", tt.want, response["error"]) + } + }) + } + + missingReq := httptest.NewRequest(http.MethodPatch, "/observations/999999", strings.NewReader(`{"find":"beta","replace":"gamma"}`)) + missingRec := httptest.NewRecorder() + h.ServeHTTP(missingRec, missingReq) + if missingRec.Code != http.StatusNotFound { + t.Fatalf("expected missing observation status 404, got %d: %s", missingRec.Code, missingRec.Body.String()) + } +} + +func TestDocsDescribeFindReplaceUpdateContract(t *testing.T) { + docs, err := os.ReadFile("../../DOCS.md") + if err != nil { + t.Fatalf("read DOCS.md: %v", err) + } + for _, want := range []string{ + "find and replace must be used together", + "find/replace is mutually exclusive with content", + "every literal occurrence", + "revision and sync", + `{"find":"alpha","replace":"beta"}`, + } { + if !strings.Contains(string(docs), want) { + t.Fatalf("DOCS.md must describe %q", want) + } + } +} + // ─── Sync Status Tests ─────────────────────────────────────────────────────── // stubSyncStatusProvider is a fake SyncStatusProvider for tests. diff --git a/internal/store/store.go b/internal/store/store.go index cf29e8ba5..7689954de 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -43,11 +43,13 @@ var sqliteWriteRetryBackoffs = []time.Duration{ // Sentinel errors returned by delete operations so callers can use errors.Is. var ( - ErrSessionNotFound = errors.New("session not found") - ErrSessionHasObservations = errors.New("session still has observations") - ErrSessionDeleteBlocked = errors.New("session deletion is blocked while cloud sync enrollment is active") - ErrObservationNotFound = errors.New("observation not found") - ErrPromptNotFound = errors.New("prompt not found") + ErrSessionNotFound = errors.New("session not found") + ErrSessionHasObservations = errors.New("session still has observations") + ErrSessionDeleteBlocked = errors.New("session deletion is blocked while cloud sync enrollment is active") + ErrObservationNotFound = errors.New("observation not found") + ErrPromptNotFound = errors.New("prompt not found") + ErrFindReplacePairRequired = errors.New("find and replace must be used together") + ErrFindReplaceWithContent = errors.New("find/replace is mutually exclusive with content") ) // ─── Types ─────────────────────────────────────────────────────────────────── @@ -150,6 +152,8 @@ type UpdateObservationParams struct { Type *string `json:"type,omitempty"` Title *string `json:"title,omitempty"` Content *string `json:"content,omitempty"` + Find *string `json:"find,omitempty"` + Replace *string `json:"replace,omitempty"` Project *string `json:"project,omitempty"` Scope *string `json:"scope,omitempty"` TopicKey *string `json:"topic_key,omitempty"` @@ -2347,6 +2351,13 @@ func (s *Store) GetObservation(id int64) (*Observation, error) { } func (s *Store) UpdateObservation(id int64, p UpdateObservationParams) (*Observation, error) { + if (p.Find == nil) != (p.Replace == nil) { + return nil, ErrFindReplacePairRequired + } + if p.Find != nil && p.Content != nil { + return nil, ErrFindReplaceWithContent + } + var updated *Observation err := s.withTx(func(tx *sql.Tx) error { obs, err := s.getObservationTx(tx, id) @@ -2368,7 +2379,15 @@ func (s *Store) UpdateObservation(id int64, p UpdateObservationParams) (*Observa title = stripPrivateTags(*p.Title) } if p.Content != nil { - content = stripPrivateTags(*p.Content) + content = *p.Content + } + if p.Find != nil { + if *p.Find != "" { + content = strings.ReplaceAll(content, *p.Find, *p.Replace) + } + } + if p.Content != nil || p.Find != nil { + content = stripPrivateTags(content) if len(content) > s.cfg.MaxObservationLength { content = content[:s.cfg.MaxObservationLength] + "... [truncated]" } diff --git a/internal/store/store_test.go b/internal/store/store_test.go index 73a15817c..664ee7611 100644 --- a/internal/store/store_test.go +++ b/internal/store/store_test.go @@ -262,6 +262,160 @@ func TestUpdateAndSoftDeleteExcludedFromSearchAndTimeline(t *testing.T) { } } +func TestUpdateObservationFindReplacePreservesPersistenceSemantics(t *testing.T) { + s := newTestStore(t) + if err := s.EnrollProject("engram"); err != nil { + t.Fatalf("enroll project: %v", err) + } + if err := s.CreateSession("find-replace", "engram", "/tmp/engram"); err != nil { + t.Fatalf("create session: %v", err) + } + + id, err := s.AddObservation(AddObservationParams{ + SessionID: "find-replace", + Type: "note", + Title: "replacement", + Content: "alpha alpha", + Project: "engram", + Scope: "project", + }) + if err != nil { + t.Fatalf("add observation: %v", err) + } + + before, err := s.GetObservation(id) + if err != nil { + t.Fatalf("get observation before update: %v", err) + } + beforeMutations, err := s.ListPendingSyncMutations(DefaultSyncTargetKey, 20) + if err != nil { + t.Fatalf("list mutations before update: %v", err) + } + + find, replace := "alpha", "secretbeta" + updated, err := s.UpdateObservation(id, UpdateObservationParams{Find: &find, Replace: &replace}) + if err != nil { + t.Fatalf("replace observation content: %v", err) + } + if updated.Content != "[REDACTED]beta [REDACTED]beta" { + t.Fatalf("expected replacement before sanitization, got %q", updated.Content) + } + if updated.RevisionCount != before.RevisionCount+1 { + t.Fatalf("expected revision increment from %d, got %d", before.RevisionCount, updated.RevisionCount) + } + + var persistedHash string + if err := s.db.QueryRow(`SELECT normalized_hash FROM observations WHERE id = ?`, id).Scan(&persistedHash); err != nil { + t.Fatalf("read normalized hash: %v", err) + } + if persistedHash != hashNormalized(updated.Content) { + t.Fatalf("expected normalized hash for replacement content, got %q", persistedHash) + } + + afterMutations, err := s.ListPendingSyncMutations(DefaultSyncTargetKey, 20) + if err != nil { + t.Fatalf("list mutations after update: %v", err) + } + if len(afterMutations) != len(beforeMutations)+1 { + t.Fatalf("expected one additional sync mutation, before=%d after=%d", len(beforeMutations), len(afterMutations)) + } + lastMutation := afterMutations[len(afterMutations)-1] + if lastMutation.Entity != SyncEntityObservation || lastMutation.EntityKey != updated.SyncID || lastMutation.Op != SyncOpUpsert { + t.Fatalf("expected observation upsert for update, got %+v", lastMutation) + } + + emptyFind, ignoredReplacement := "", "ignored" + emptyUpdated, err := s.UpdateObservation(id, UpdateObservationParams{Find: &emptyFind, Replace: &ignoredReplacement}) + if err != nil { + t.Fatalf("empty find update: %v", err) + } + if emptyUpdated.Content != updated.Content { + t.Fatalf("expected empty find to preserve content %q, got %q", updated.Content, emptyUpdated.Content) + } + if emptyUpdated.RevisionCount != updated.RevisionCount+1 { + t.Fatalf("expected empty find to increment revision, got %d", emptyUpdated.RevisionCount) + } + + nonMatchingFind, replacement := "missing", "changed" + nonMatchingUpdated, err := s.UpdateObservation(id, UpdateObservationParams{Find: &nonMatchingFind, Replace: &replacement}) + if err != nil { + t.Fatalf("non-matching find update: %v", err) + } + if nonMatchingUpdated.Content != updated.Content { + t.Fatalf("expected non-matching find to preserve content %q, got %q", updated.Content, nonMatchingUpdated.Content) + } + if nonMatchingUpdated.RevisionCount != emptyUpdated.RevisionCount+1 { + t.Fatalf("expected non-matching find to increment revision, got %d", nonMatchingUpdated.RevisionCount) + } +} + +func TestUpdateObservationFindReplaceValidatesFieldCombinations(t *testing.T) { + s := newTestStore(t) + if err := s.CreateSession("find-replace-validation", "engram", "/tmp/engram"); err != nil { + t.Fatalf("create session: %v", err) + } + id, err := s.AddObservation(AddObservationParams{SessionID: "find-replace-validation", Type: "note", Title: "validation", Content: "alpha", Project: "engram"}) + if err != nil { + t.Fatalf("add observation: %v", err) + } + + find, replace, content := "", "beta", "replacement content" + tests := []struct { + name string + params UpdateObservationParams + want error + }{ + {name: "find only", params: UpdateObservationParams{Find: &find}, want: ErrFindReplacePairRequired}, + {name: "replace only", params: UpdateObservationParams{Replace: &replace}, want: ErrFindReplacePairRequired}, + {name: "pair with content", params: UpdateObservationParams{Find: &find, Replace: &replace, Content: &content}, want: ErrFindReplaceWithContent}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := s.UpdateObservation(id, tt.params) + if !errors.Is(err, tt.want) { + t.Fatalf("expected %q, got %v", tt.want, err) + } + }) + } +} + +func TestUpdateObservationFindReplaceTruncatesReplacementOutput(t *testing.T) { + s := newTestStore(t) + if err := s.CreateSession("find-replace-truncation", "engram", "/tmp/engram"); err != nil { + t.Fatalf("create session: %v", err) + } + id, err := s.AddObservation(AddObservationParams{ + SessionID: "find-replace-truncation", + Type: "note", + Title: "truncation", + Content: "replace-me", + Project: "engram", + }) + if err != nil { + t.Fatalf("add observation: %v", err) + } + + find := "replace-me" + replace := strings.Repeat("x", s.MaxObservationLength()+1) + updated, err := s.UpdateObservation(id, UpdateObservationParams{Find: &find, Replace: &replace}) + if err != nil { + t.Fatalf("replace observation content: %v", err) + } + + want := strings.Repeat("x", s.MaxObservationLength()) + "... [truncated]" + if updated.Content != want { + t.Fatalf("expected replacement output to follow truncation behavior, got length=%d want length=%d", len(updated.Content), len(want)) + } + + persisted, err := s.GetObservation(id) + if err != nil { + t.Fatalf("get truncated observation: %v", err) + } + if persisted.Content != want { + t.Fatalf("expected persisted replacement output to be truncated, got length=%d want length=%d", len(persisted.Content), len(want)) + } +} + func TestTopicKeyUpsertUpdatesSameTopicWithoutCreatingNewRow(t *testing.T) { s := newTestStore(t) From 877ce784cab440521ec26181bc82b9e7651c6728 Mon Sep 17 00:00:00 2001 From: blak0p Date: Tue, 25 Aug 2026 11:30:06 +0200 Subject: [PATCH 2/2] fix(store): bound replacement expansion --- internal/store/store.go | 261 ++++++++++++++++++++++++++++++++++- internal/store/store_test.go | 135 ++++++++++++++++++ 2 files changed, 390 insertions(+), 6 deletions(-) diff --git a/internal/store/store.go b/internal/store/store.go index 25a74e81e..d74f35c7e 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -20,6 +20,8 @@ import ( "strconv" "strings" "time" + "unicode" + "unicode/utf8" "github.com/Gentleman-Programming/engram/internal/timeutil" sqlite "modernc.org/sqlite" @@ -2896,12 +2898,9 @@ func (s *Store) UpdateObservation(id int64, p UpdateObservationParams) (*Observa if p.Content != nil { content = *p.Content } - if p.Find != nil { - if *p.Find != "" { - content = strings.ReplaceAll(content, *p.Find, *p.Replace) - } - } - if p.Content != nil || p.Find != nil { + if p.Find != nil && *p.Find != "" { + content = replaceAndNormalizeObservationContent(content, *p.Find, *p.Replace, s.cfg.MaxObservationLength) + } else if p.Content != nil || p.Find != nil { content = stripPrivateTags(content) if len(content) > s.cfg.MaxObservationLength { content = content[:s.cfg.MaxObservationLength] + "... [truncated]" @@ -6589,6 +6588,11 @@ func normalizeExistingSyncID(existing, prefix string) string { // Supports multiline and nested content. Case-insensitive. var privateTagRegex = regexp.MustCompile(`(?is).*?`) +const ( + privateOpenTag = "" + privateCloseTag = "" +) + // stripPrivateTags removes all ... content from a string. // This ensures sensitive information (API keys, passwords, personal data) // is never persisted to the memory database. @@ -6599,6 +6603,251 @@ func stripPrivateTags(s string) string { return result } +// replaceAndNormalizeObservationContent streams the logical ReplaceAll output +// through private-tag sanitization, retaining only the bytes needed for the +// existing truncation behavior. This avoids materializing a large expanded +// string when a replacement occurs many times. +func replaceAndNormalizeObservationContent(content, find, replace string, max int) string { + sanitizer := privateTagSanitizer{ + output: boundedTrimmedContent{max: max}, + fallback: boundedTrimmedContent{max: max}, + } + + for sourcePos := 0; sourcePos < len(content); { + matchOffset := strings.Index(content[sourcePos:], find) + if matchOffset < 0 { + sanitizer.writeString(content[sourcePos:]) + break + } + sanitizer.writeString(content[sourcePos : sourcePos+matchOffset]) + sanitizer.writeString(replace) + sourcePos += matchOffset + len(find) + } + + return sanitizer.finish() +} + +type privateTagSanitizer struct { + output boundedTrimmedContent + fallback boundedTrimmedContent + openCandidate []byte + closeCandidate []byte + inPrivate bool +} + +func (s *privateTagSanitizer) writeByte(b byte) { + if s.inPrivate { + s.writePrivateByte(b) + return + } + s.writeNormalByte(b) +} + +func (s *privateTagSanitizer) writeString(value string) { + for i := 0; i < len(value); i++ { + s.writeByte(value[i]) + } +} + +func (s *privateTagSanitizer) writeNormalByte(b byte) { + if len(s.openCandidate) == 0 { + if b == privateOpenTag[0] { + s.openCandidate = append(s.openCandidate[:0], b) + return + } + s.output.writeByte(b) + return + } + + if len(s.openCandidate) < len(privateOpenTag) && asciiEqualFold(b, privateOpenTag[len(s.openCandidate)]) { + s.openCandidate = append(s.openCandidate, b) + if len(s.openCandidate) == len(privateOpenTag) { + s.inPrivate = true + s.fallback.reset() + s.fallback.writeBytes(s.openCandidate) + s.openCandidate = s.openCandidate[:0] + } + return + } + + s.output.writeBytes(s.openCandidate) + s.openCandidate = s.openCandidate[:0] + if b == privateOpenTag[0] { + s.openCandidate = append(s.openCandidate, b) + } else { + s.output.writeByte(b) + } +} + +func (s *privateTagSanitizer) writePrivateByte(b byte) { + s.fallback.writeByte(b) + + if len(s.closeCandidate) == 0 { + if b == privateCloseTag[0] { + s.closeCandidate = append(s.closeCandidate[:0], b) + } + return + } + + if len(s.closeCandidate) < len(privateCloseTag) && asciiEqualFold(b, privateCloseTag[len(s.closeCandidate)]) { + s.closeCandidate = append(s.closeCandidate, b) + if len(s.closeCandidate) == len(privateCloseTag) { + s.inPrivate = false + s.closeCandidate = s.closeCandidate[:0] + s.fallback.reset() + s.output.writeString("[REDACTED]") + } + return + } + + if b == privateCloseTag[0] { + s.closeCandidate = append(s.closeCandidate[:0], b) + } else { + s.closeCandidate = s.closeCandidate[:0] + } +} + +func (s *privateTagSanitizer) finish() string { + if s.inPrivate { + s.fallback.appendTo(&s.output) + } else { + for _, b := range s.openCandidate { + s.output.writeByte(b) + } + } + return s.output.finish() +} + +func asciiEqualFold(a, b byte) bool { + if a >= 'A' && a <= 'Z' { + a += 'a' - 'A' + } + if b >= 'A' && b <= 'Z' { + b += 'a' - 'A' + } + return a == b +} + +type boundedTrimmedContent struct { + max int + content strings.Builder + pendingSpace []byte + pendingLong bool + hasContent bool + truncated bool + runeBuffer []byte +} + +func (s *boundedTrimmedContent) reset() { + s.content.Reset() + s.pendingSpace = s.pendingSpace[:0] + s.pendingLong = false + s.hasContent = false + s.truncated = false + s.runeBuffer = s.runeBuffer[:0] +} + +func (s *boundedTrimmedContent) writeByte(b byte) { + s.runeBuffer = append(s.runeBuffer, b) + s.flushRunes(false) +} + +func (s *boundedTrimmedContent) writeString(value string) { + for i := 0; i < len(value); i++ { + s.writeByte(value[i]) + } +} + +func (s *boundedTrimmedContent) writeBytes(value []byte) { + for _, b := range value { + s.writeByte(b) + } +} + +func (s *boundedTrimmedContent) flushRunes(force bool) { + for len(s.runeBuffer) > 0 { + if !force && !utf8.FullRune(s.runeBuffer) { + return + } + r, size := utf8.DecodeRune(s.runeBuffer) + if size == 0 { + return + } + s.writeRune(s.runeBuffer[:size], unicode.IsSpace(r)) + s.runeBuffer = s.runeBuffer[size:] + } +} + +func (s *boundedTrimmedContent) writeRune(value []byte, space bool) { + if space { + if !s.hasContent { + return + } + s.appendPending(value) + return + } + + if !s.hasContent { + s.hasContent = true + } + s.commit(s.pendingSpace) + if s.pendingLong { + s.truncated = true + } + s.pendingSpace = s.pendingSpace[:0] + s.pendingLong = false + s.commit(value) +} + +func (s *boundedTrimmedContent) appendPending(value []byte) { + if s.pendingLong { + return + } + remaining := s.max - len(s.pendingSpace) + if len(value) > remaining { + s.pendingLong = true + } + if remaining > len(value) { + remaining = len(value) + } + s.pendingSpace = append(s.pendingSpace, value[:remaining]...) +} + +func (s *boundedTrimmedContent) commit(value []byte) { + if len(value) == 0 { + return + } + remaining := s.max - s.content.Len() + if len(value) > remaining { + s.truncated = true + value = value[:remaining] + } + _, _ = s.content.Write(value) +} + +func (s *boundedTrimmedContent) appendTo(destination *boundedTrimmedContent) { + s.flushRunes(true) + if !s.hasContent { + return + } + destination.writeString(s.content.String()) + if s.truncated { + destination.truncated = true + } +} + +func (s *boundedTrimmedContent) finish() string { + s.flushRunes(true) + if !s.hasContent { + return "" + } + result := s.content.String() + if s.truncated { + return result + "... [truncated]" + } + return result +} + // sanitizeFTS wraps each word in quotes so FTS5 doesn't choke on special chars. // "fix auth bug" → `"fix" "auth" "bug"` func sanitizeFTS(query string) string { diff --git a/internal/store/store_test.go b/internal/store/store_test.go index bb4c3483d..7df13d087 100644 --- a/internal/store/store_test.go +++ b/internal/store/store_test.go @@ -417,6 +417,141 @@ func TestUpdateObservationFindReplaceTruncatesReplacementOutput(t *testing.T) { } } +func TestUpdateObservationFindReplaceBoundsExpandedOutput(t *testing.T) { + cfg := mustDefaultConfig(t) + cfg.DataDir = t.TempDir() + cfg.MaxObservationLength = 32 + s, err := New(cfg) + if err != nil { + t.Fatalf("new store: %v", err) + } + t.Cleanup(func() { _ = s.Close() }) + + if err := s.CreateSession("find-replace-bounded", "engram", "/tmp/engram"); err != nil { + t.Fatalf("create session: %v", err) + } + id, err := s.AddObservation(AddObservationParams{ + SessionID: "find-replace-bounded", + Type: "note", + Title: "bounded replacement", + Content: strings.Repeat("a", cfg.MaxObservationLength), + Project: "engram", + }) + if err != nil { + t.Fatalf("add observation: %v", err) + } + + find := "a" + replace := strings.Repeat("b", 256*1024) + updated, err := s.UpdateObservation(id, UpdateObservationParams{Find: &find, Replace: &replace}) + if err != nil { + t.Fatalf("replace observation content: %v", err) + } + + want := strings.Repeat("b", cfg.MaxObservationLength) + "... [truncated]" + if updated.Content != want { + t.Fatalf("expected bounded replacement output to be truncated, got length=%d want length=%d", len(updated.Content), len(want)) + } +} + +func TestReplaceAndNormalizeObservationContentPreservesMixedCaseUnmatchedTag(t *testing.T) { + const max = 100 + got := replaceAndNormalizeObservationContent("\tsecret\n", "missing", "replacement", max) + if want := "secret"; got != want { + t.Fatalf("unmatched private tag changed: got %q, want %q", got, want) + } +} + +func TestReplaceAndNormalizeObservationContentMatchesLegacyPipeline(t *testing.T) { + tests := []struct { + name string + content string + find string + replace string + max int + }{ + { + name: "repeated replacements", + content: "alpha alpha alpha", + find: "alpha", + replace: "beta", + max: 100, + }, + { + name: "opening tag spans source and replacement", + content: "startend", + find: "X", + replace: "vate>", + max: 100, + }, + { + name: "closing tag spans source and replacement", + content: "startsecretXend", + find: "X", + replace: "", + max: 100, + }, + { + name: "multiple mixed-case tags", + content: "one middle two", + find: "middle", + replace: "between", + max: 100, + }, + { + name: "unmatched mixed-case tag", + content: " \u2003secret\n", + find: "missing", + replace: "replacement", + max: 100, + }, + { + name: "malformed private tag", + content: "prefixsecretsecret0123456789", + max: 12, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := replaceAndNormalizeObservationContent(tt.content, tt.find, tt.replace, tt.max) + want := legacyReplaceAndNormalizeObservationContent(tt.content, tt.find, tt.replace, tt.max) + if got != want { + t.Fatalf("streaming pipeline differs: got %q, want %q", got, want) + } + }) + } +} + +// legacyReplaceAndNormalizeObservationContent is the pre-remediation logical +// pipeline used as a differential-test oracle. It intentionally materializes +// output because the production path must not. +func legacyReplaceAndNormalizeObservationContent(content, find, replace string, max int) string { + content = strings.ReplaceAll(content, find, replace) + content = privateTagRegex.ReplaceAllString(content, "[REDACTED]") + content = strings.TrimSpace(content) + if len(content) > max { + content = content[:max] + "... [truncated]" + } + return content +} + func TestPinnedObservationsAndFormatContextPriority(t *testing.T) { cfg := mustDefaultConfig(t) cfg.DataDir = t.TempDir()