fix(engine): bound Remember/Add prompt budget + per-stage timings (#41)

engine/engine.go

@@ -33,18 +33,32 @@ type EngineConfig struct {
 	ReEvalSimilarityMin float64 // min similarity for re-eval candidates (default: 0.5)
 	ReEvalSimilarityMax float64 // max similarity for re-eval candidates (default: 0.85)
 	ReEvalImportanceMin float64 // min importance for re-eval candidates (default: 0.5)
+
+	// Extraction prompt budgets — cap inputs into ExtractMemories so repeated
+	// same-entity writes don't balloon the prompt and push local LLMs over
+	// their usable latency window. See issue #41.
+	ExtractionExistingMaxCount   int     // max existing-memory neighbors considered (default: 10)
+	ExtractionExistingMinScore   float64 // min similarity for existing-memory neighbors (default: 0.5)
+	ExtractionExistingMaxBytes   int     // byte cap on formatted existingMemories block (default: 4096; 0 = unlimited)
+	ExtractionContextMaxBytes    int     // byte cap on formatted conversationCtx block (default: 4096; 0 = unlimited)
+	ExtractionMemoryContentTrim  int     // per-memory content trim for the existing-memories prompt (default: 240; 0 = no trim)
 }
 
 // DefaultEngineConfig returns a default engine configuration.
 func DefaultEngineConfig() EngineConfig {
 	return EngineConfig{
-		ContextTurns:          5,
-		DefaultMinScore:       0.3,
-		FTSBaselineSimilarity: 0.4,
-		DiversityThreshold:    0.9,
-		ReEvalSimilarityMin:   0.5,
-		ReEvalSimilarityMax:   0.85,
-		ReEvalImportanceMin:   0.5,
+		ContextTurns:                5,
+		DefaultMinScore:             0.3,
+		FTSBaselineSimilarity:       0.4,
+		DiversityThreshold:          0.9,
+		ReEvalSimilarityMin:         0.5,
+		ReEvalSimilarityMax:         0.85,
+		ReEvalImportanceMin:         0.5,
+		ExtractionExistingMaxCount:  10,
+		ExtractionExistingMinScore:  0.5,
+		ExtractionExistingMaxBytes:  4096,
+		ExtractionContextMaxBytes:   4096,
+		ExtractionMemoryContentTrim: 240,
 	}
 }
 
@@ -100,6 +114,27 @@ func NewEngine(
 	if config.ReEvalImportanceMin <= 0 {
 		config.ReEvalImportanceMin = 0.5
 	}
+	if config.ExtractionExistingMaxCount <= 0 {
+		config.ExtractionExistingMaxCount = 10
+	}
+	if config.ExtractionExistingMinScore <= 0 {
+		config.ExtractionExistingMinScore = 0.5
+	}
+	if config.ExtractionExistingMaxBytes < 0 {
+		config.ExtractionExistingMaxBytes = 0
+	} else if config.ExtractionExistingMaxBytes == 0 {
+		config.ExtractionExistingMaxBytes = 4096
+	}
+	if config.ExtractionContextMaxBytes < 0 {
+		config.ExtractionContextMaxBytes = 0
+	} else if config.ExtractionContextMaxBytes == 0 {
+		config.ExtractionContextMaxBytes = 4096
+	}
+	if config.ExtractionMemoryContentTrim < 0 {
+		config.ExtractionMemoryContentTrim = 0
+	} else if config.ExtractionMemoryContentTrim == 0 {
+		config.ExtractionMemoryContentTrim = 240
+	}
 
 	graphEngine := NewGraphEngine(store, DefaultGraphConfig())
 

engine/engine_add.go

@@ -6,6 +6,7 @@ package engine
 import (
 	"context"
 	"fmt"
+	"log"
 	"time"
 
 	"github.com/keyoku-ai/keyoku-engine/llm"
@@ -80,21 +81,22 @@ func (e *Engine) Add(ctx context.Context, entityID string, req AddRequest) (*Add
 	}
 
 	// Get conversation context
+	t0 := time.Now()
 	contextMsgs, err := e.store.GetRecentSessionMessages(ctx, entityID, e.config.ContextTurns)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get session context: %w", err)
 	}
+	sessionMs := time.Since(t0).Milliseconds()
 
-	var conversationCtx []string
-	for _, msg := range contextMsgs {
-		conversationCtx = append(conversationCtx, fmt.Sprintf("[%s]: %s", msg.Role, msg.Content))
-	}
+	conversationCtx, ctxBytes := buildConversationContext(contextMsgs, e.config.ExtractionContextMaxBytes)
 
 	// Embed content to find similar existing memories
+	t0 = time.Now()
 	queryEmbedding, err := e.embedder.Embed(ctx, req.Content)
 	if err != nil {
 		return nil, fmt.Errorf("failed to embed query: %w", err)
 	}
+	embedMs := time.Since(t0).Milliseconds()
 
 	// Build visibility context for similarity search so dedup/conflict only considers visible memories
 	var visibilityFor *storage.VisibilityContext
@@ -106,26 +108,46 @@ func (e *Engine) Add(ctx context.Context, entityID string, req AddRequest) (*Add
 		visibilityFor = &storage.VisibilityContext{AgentID: agentID, TeamID: req.TeamID}
 	}
 
-	similarMemories, err := e.store.FindSimilarWithOptions(ctx, queryEmbedding, entityID, 10, 0.5, storage.SimilarityOptions{
-		AgentID:       req.AgentID,
-		VisibilityFor: visibilityFor,
-	})
+	t0 = time.Now()
+	similarMemories, err := e.store.FindSimilarWithOptions(
+		ctx, queryEmbedding, entityID,
+		e.config.ExtractionExistingMaxCount,
+		e.config.ExtractionExistingMinScore,
+		storage.SimilarityOptions{
+			AgentID:       req.AgentID,
+			VisibilityFor: visibilityFor,
+		},
+	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to find similar memories: %w", err)
 	}
+	similarMs := time.Since(t0).Milliseconds()
 
-	var existingMemories []string
-	for _, sm := range similarMemories {
-		existingMemories = append(existingMemories, fmt.Sprintf("[%s, %s, importance:%.1f] %s",
-			sm.Memory.Type, sm.Memory.State, sm.Memory.Importance, sm.Memory.Content))
-	}
+	existingMemories, existingBytes := buildExistingMemoriesBlock(
+		similarMemories,
+		e.config.ExtractionExistingMaxBytes,
+		e.config.ExtractionMemoryContentTrim,
+	)
 
 	// Extract memories using LLM
+	t0 = time.Now()
 	extractResp, err := e.provider.ExtractMemories(ctx, llm.ExtractionRequest{
 		Content:          req.Content,
 		ConversationCtx:  conversationCtx,
 		ExistingMemories: existingMemories,
 	})
+	extractMs := time.Since(t0).Milliseconds()
+
+	// Per-stage timing + payload-size trace. See harness/40-invariants/observability.md
+	// (Add() stage timings) and issue #41. Structured key=value so ops/tooling can grep.
+	log.Printf(
+		"INFO [remember/add] stage=extract entity_id=%q session_ms=%d embed_ms=%d similar_ms=%d extract_ms=%d existing_count=%d existing_bytes=%d ctx_msgs=%d ctx_bytes=%d err=%v",
+		entityID, sessionMs, embedMs, similarMs, extractMs,
+		len(existingMemories), existingBytes,
+		len(conversationCtx), ctxBytes,
+		err,
+	)
+
 	if err != nil {
 		return nil, fmt.Errorf("extraction failed: %w", err)
 	}

engine/engine_add_prompt_budget.go

@@ -0,0 +1,78 @@
+// SPDX-License-Identifier: BSL-1.1
+// Copyright (c) 2026 Keyoku. All rights reserved.
+
+package engine
+
+import (
+	"fmt"
+
+	"github.com/keyoku-ai/keyoku-engine/storage"
+)
+
+// Prompt-budget helpers for the Add() / ExtractMemories path.
+//
+// Rationale (issue #41): repeated same-entity writes returned steadily more
+// 0.5+ similarity neighbors, which inflated the formatted existingMemories
+// block inside the extraction prompt and pushed local LLMs (e.g. qwen2.5:7b)
+// past their usable latency window (25–55s, then 60s timeouts). The v2 harness
+// flagged k=10/minScore=0.5 as hardcoded and observability.md did not require
+// per-stage timings. These helpers cap the prompt inputs by byte budget
+// (most-similar first, since similarMemories arrives ordered by score) and are
+// pure functions so they're cheap to test without provider stubs.
+//
+// Byte-budget semantics:
+//   - budget > 0  → hard cap; stop once the formatted block would exceed it.
+//   - budget == 0 → treated as "unlimited" (no trim). Defaults are applied in
+//     NewEngine, so a zero-valued config struct still gets sensible caps.
+
+// buildConversationContext formats recent session messages for the extraction
+// prompt and trims by total byte budget (includes newline separators in the
+// running count, approximated by adding 1 per entry).
+func buildConversationContext(msgs []*storage.SessionMessage, maxBytes int) ([]string, int) {
+	if len(msgs) == 0 {
+		return nil, 0
+	}
+	out := make([]string, 0, len(msgs))
+	total := 0
+	for _, m := range msgs {
+		if m == nil {
+			continue
+		}
+		line := fmt.Sprintf("[%s]: %s", m.Role, m.Content)
+		if maxBytes > 0 && total+len(line)+1 > maxBytes {
+			break
+		}
+		out = append(out, line)
+		total += len(line) + 1
+	}
+	return out, total
+}
+
+// buildExistingMemoriesBlock formats the similar-memory neighbors for the
+// extraction prompt and trims by total byte budget. Each memory's content is
+// trimmed per-entry via perMemContentTrim (0 = no trim). The format is kept
+// compact: `[Type, State, importance:%.1f] Content`.
+func buildExistingMemoriesBlock(sims []*storage.SimilarityResult, maxBytes int, perMemContentTrim int) ([]string, int) {
+	if len(sims) == 0 {
+		return nil, 0
+	}
+	out := make([]string, 0, len(sims))
+	total := 0
+	for _, sm := range sims {
+		if sm == nil || sm.Memory == nil {
+			continue
+		}
+		content := sm.Memory.Content
+		if perMemContentTrim > 0 && len(content) > perMemContentTrim {
+			content = content[:perMemContentTrim] + "..."
+		}
+		line := fmt.Sprintf("[%s, %s, importance:%.1f] %s",
+			sm.Memory.Type, sm.Memory.State, sm.Memory.Importance, content)
+		if maxBytes > 0 && total+len(line)+1 > maxBytes {
+			break
+		}
+		out = append(out, line)
+		total += len(line) + 1
+	}
+	return out, total
+}

engine/engine_add_prompt_budget_test.go

@@ -0,0 +1,109 @@
+// SPDX-License-Identifier: BSL-1.1
+// Copyright (c) 2026 Keyoku. All rights reserved.
+
+package engine
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/keyoku-ai/keyoku-engine/storage"
+)
+
+func TestBuildConversationContext_Unlimited(t *testing.T) {
+	msgs := []*storage.SessionMessage{
+		{Role: "user", Content: "hello"},
+		{Role: "assistant", Content: "hi there"},
+	}
+	out, total := buildConversationContext(msgs, 0)
+	if len(out) != 2 {
+		t.Fatalf("len(out) = %d, want 2", len(out))
+	}
+	if total == 0 {
+		t.Fatalf("total bytes = 0, want > 0")
+	}
+	if out[0] != "[user]: hello" {
+		t.Fatalf("out[0] = %q", out[0])
+	}
+}
+
+func TestBuildConversationContext_TrimsByBudget(t *testing.T) {
+	// Each formatted line is "[user]: msgN" ~= 13 bytes + 1 newline accounting
+	msgs := []*storage.SessionMessage{
+		{Role: "user", Content: "msg1"},
+		{Role: "user", Content: "msg2"},
+		{Role: "user", Content: "msg3"},
+		{Role: "user", Content: "msg4"},
+	}
+	out, total := buildConversationContext(msgs, 30)
+	// 30 bytes budget should admit roughly 2 entries (~14 bytes each with newline)
+	if len(out) >= len(msgs) {
+		t.Fatalf("expected trimming, got all %d entries", len(out))
+	}
+	if total > 30 {
+		t.Fatalf("total %d exceeds budget 30", total)
+	}
+}
+
+func TestBuildExistingMemoriesBlock_TrimsPerMemoryContent(t *testing.T) {
+	longContent := strings.Repeat("x", 500)
+	sims := []*storage.SimilarityResult{
+		{Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.8, Content: longContent}},
+	}
+	out, _ := buildExistingMemoriesBlock(sims, 0, 100)
+	if len(out) != 1 {
+		t.Fatalf("len(out) = %d, want 1", len(out))
+	}
+	// 100 trim + "..." = 103-char content. Line prefix adds more, but content portion must be bounded.
+	if !strings.Contains(out[0], strings.Repeat("x", 100)+"...") {
+		t.Fatalf("expected per-memory trim to 100 chars with ellipsis, got: %q", out[0])
+	}
+	if strings.Count(out[0], "x") > 101 {
+		t.Fatalf("trim leaked: too many x in output: %d", strings.Count(out[0], "x"))
+	}
+}
+
+func TestBuildExistingMemoriesBlock_TrimsByByteBudget(t *testing.T) {
+	sims := []*storage.SimilarityResult{
+		{Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.8, Content: "one"}},
+		{Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.8, Content: "two"}},
+		{Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.8, Content: "three"}},
+	}
+	// Each formatted line is ~"[event, active, importance:0.8] xxx" ~= 35 bytes + 1.
+	// Budget 40 should allow one, not more.
+	out, total := buildExistingMemoriesBlock(sims, 40, 0)
+	if len(out) != 1 {
+		t.Fatalf("len(out) = %d, want 1 under tight budget", len(out))
+	}
+	if total > 40 {
+		t.Fatalf("total %d exceeds budget 40", total)
+	}
+}
+
+func TestBuildExistingMemoriesBlock_NilSafe(t *testing.T) {
+	sims := []*storage.SimilarityResult{
+		nil,
+		{Memory: nil},
+		{Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.5, Content: "ok"}},
+	}
+	out, _ := buildExistingMemoriesBlock(sims, 0, 0)
+	if len(out) != 1 {
+		t.Fatalf("len(out) = %d, want 1 (nil entries skipped)", len(out))
+	}
+}
+
+func TestBuildExistingMemoriesBlock_PreservesMostSimilarFirst(t *testing.T) {
+	// Input order mirrors store ordering (highest similarity first). We must
+	// keep that order so tight budgets retain the most relevant neighbors.
+	sims := []*storage.SimilarityResult{
+		{Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.9, Content: "most_similar"}},
+		{Memory: &storage.Memory{Type: storage.TypeEvent, State: "active", Importance: 0.5, Content: "least_similar"}},
+	}
+	out, _ := buildExistingMemoriesBlock(sims, 50, 0)
+	if len(out) < 1 {
+		t.Fatalf("expected at least 1 entry")
+	}
+	if !strings.Contains(out[0], "most_similar") {
+		t.Fatalf("first entry should be most_similar, got: %q", out[0])
+	}
+}