From e38b3cc175ea85f2c67ca5c25b5e927ce448287d Mon Sep 17 00:00:00 2001 From: root Date: Wed, 25 Feb 2026 12:55:55 +0800 Subject: [PATCH] Improve context compaction/memory flow and harden cron scheduling --- config/config.example.json | 30 +- pkg/agent/compaction.go | 320 ++++++++++++++ pkg/agent/context.go | 754 ++++++++++++++++++-------------- pkg/agent/context_test.go | 353 +++++++-------- pkg/agent/instance.go | 204 +++++++-- pkg/agent/loop.go | 97 +++- pkg/agent/loop_test.go | 8 +- pkg/agent/memory.go | 180 +++++++- pkg/agent/memory_test.go | 48 ++ pkg/agent/memory_tool.go | 170 +++++++ pkg/agent/memory_tool_test.go | 93 ++++ pkg/agent/memory_vector.go | 645 +++++++++++++++++++++++++++ pkg/agent/memory_vector_test.go | 107 +++++ pkg/config/config.go | 38 ++ pkg/config/config_test.go | 18 + pkg/config/defaults.go | 28 ++ pkg/session/manager.go | 83 +++- pkg/session/manager_test.go | 28 ++ pkg/tools/cron.go | 69 ++- pkg/tools/cron_test.go | 101 +++++ 20 files changed, 2760 insertions(+), 614 deletions(-) create mode 100644 pkg/agent/compaction.go create mode 100644 pkg/agent/memory_test.go create mode 100644 pkg/agent/memory_tool.go create mode 100644 pkg/agent/memory_tool_test.go create mode 100644 pkg/agent/memory_vector.go create mode 100644 pkg/agent/memory_vector_test.go create mode 100644 pkg/tools/cron_test.go diff --git a/config/config.example.json b/config/config.example.json index 52e993a97..2d10022b6 100644 --- a/config/config.example.json +++ b/config/config.example.json @@ -6,7 +6,35 @@ "model_name": "gpt4", "max_tokens": 8192, "temperature": 0.7, - "max_tool_iterations": 20 + "max_tool_iterations": 20, + "compaction": { + "mode": "safeguard", + "reserve_tokens": 2048, + "keep_recent_tokens": 2048, + "max_history_share": 0.5, + "memory_flush": { + "enabled": true, + "soft_threshold_tokens": 1500 + } + }, + "context_pruning": { + "mode": "tools_only", + "include_old_chitchat": true, + "soft_tool_result_chars": 2000, + "hard_tool_result_chars": 350, + "trigger_ratio": 0.8 + }, + "bootstrap_snapshot": { + "enabled": true + }, + "memory_vector": { + "enabled": true, + "dimensions": 256, + "top_k": 6, + "min_score": 0.15, + "max_context_chars": 1800, + "recent_daily_days": 14 + } } }, "model_list": [ diff --git a/pkg/agent/compaction.go b/pkg/agent/compaction.go new file mode 100644 index 000000000..ec9fc5671 --- /dev/null +++ b/pkg/agent/compaction.go @@ -0,0 +1,320 @@ +package agent + +import ( + "context" + "fmt" + "strings" + "time" + + "github.com/sipeed/picoclaw/pkg/logger" + "github.com/sipeed/picoclaw/pkg/providers" +) + +func (al *AgentLoop) maybeFlushMemoryBeforeCompaction( + ctx context.Context, + agent *AgentInstance, + sessionKey string, + tokenEstimate int, +) (bool, error) { + if !isCompactionModeEnabled(agent.CompactionMode) { + return false, nil + } + if !agent.MemoryFlushEnabled || sessionKey == "" { + return false, nil + } + + triggerPoint := agent.ContextWindow - agent.CompactionReserveTokens - agent.MemoryFlushSoftThreshold + if triggerPoint < agent.ContextWindow/3 { + triggerPoint = agent.ContextWindow / 3 + } + if tokenEstimate < triggerPoint { + return false, nil + } + + compactionCount, flushedAtCount, _ := agent.Sessions.GetCompactionState(sessionKey) + if flushedAtCount == compactionCount { + return false, nil + } + + if err := al.flushMemorySnapshot(ctx, agent, sessionKey); err != nil { + return false, err + } + + agent.Sessions.MarkMemoryFlush(sessionKey, compactionCount) + _ = agent.Sessions.Save(sessionKey) + return true, nil +} + +func (al *AgentLoop) flushMemorySnapshot(ctx context.Context, agent *AgentInstance, sessionKey string) error { + history := agent.Sessions.GetHistory(sessionKey) + if len(history) == 0 { + return nil + } + + recent := make([]providers.Message, 0, 12) + for i := len(history) - 1; i >= 0 && len(recent) < 12; i-- { + msg := history[i] + if (msg.Role != "user" && msg.Role != "assistant") || strings.TrimSpace(msg.Content) == "" { + continue + } + recent = append([]providers.Message{msg}, recent...) + } + if len(recent) == 0 { + return nil + } + + var prompt strings.Builder + prompt.WriteString("Extract durable memory from this chat.\n") + prompt.WriteString("Return concise markdown bullets under these headings only:\n") + prompt.WriteString("## Profile\n## Long-term Facts\n## Active Goals\n## Constraints\n## Open Threads\n## Deprecated/Resolved\n") + prompt.WriteString("\nCHAT:\n") + for _, m := range recent { + prompt.WriteString(m.Role) + prompt.WriteString(": ") + prompt.WriteString(m.Content) + prompt.WriteString("\n") + } + + resp, err := agent.Provider.Chat( + ctx, + []providers.Message{{Role: "user", Content: prompt.String()}}, + nil, + agent.Model, + map[string]any{ + "max_tokens": 700, + "temperature": 0.2, + }, + ) + if err != nil { + return err + } + if strings.TrimSpace(resp.Content) == "" { + return fmt.Errorf("empty memory flush response") + } + + memory := NewMemoryStore(agent.Workspace) + return memory.OrganizeWriteback(resp.Content) +} + +func (al *AgentLoop) compactWithSafeguard( + ctx context.Context, + agent *AgentInstance, + sessionKey string, +) (bool, error) { + switch normalizeCompactionMode(agent.CompactionMode) { + case "off": + return false, nil + case "legacy": + beforeHistory := len(agent.Sessions.GetHistory(sessionKey)) + beforeSummary := strings.TrimSpace(agent.Sessions.GetSummary(sessionKey)) + al.summarizeSession(agent, sessionKey) + afterHistory := len(agent.Sessions.GetHistory(sessionKey)) + afterSummary := strings.TrimSpace(agent.Sessions.GetSummary(sessionKey)) + if afterHistory < beforeHistory || afterSummary != beforeSummary { + agent.Sessions.IncrementCompactionCount(sessionKey) + _ = agent.Sessions.Save(sessionKey) + return true, nil + } + return false, nil + } + + history := sanitizeHistoryForProvider(agent.Sessions.GetHistory(sessionKey)) + if len(history) <= 6 { + return false, nil + } + + historyTokens := al.estimateTokens(history) + historyBudget := int(float64(agent.ContextWindow) * agent.CompactionMaxHistoryShare) + if historyBudget <= 0 { + historyBudget = agent.ContextWindow / 2 + } + if historyTokens <= historyBudget && len(history) < 24 { + return false, nil + } + + keepRecentTokens := agent.CompactionKeepRecentTokens + if keepRecentTokens <= 0 { + keepRecentTokens = maxInt(1024, agent.ContextWindow/4) + } + + keepStart := len(history) - 4 + if keepStart < 1 { + keepStart = 1 + } + acc := 0 + for i := len(history) - 1; i >= 1; i-- { + acc += al.estimateMessageTokens(history[i]) + if acc >= keepRecentTokens { + keepStart = i + break + } + } + if keepStart <= 0 || keepStart >= len(history) { + return false, nil + } + + toSummarize := history[:keepStart] + kept := history[keepStart:] + if len(toSummarize) == 0 || len(kept) == 0 { + return false, nil + } + + existingSummary := agent.Sessions.GetSummary(sessionKey) + summary, err := al.generateCompactionSummary(ctx, agent, toSummarize, existingSummary) + if err != nil { + return false, err + } + if strings.TrimSpace(summary) == "" { + return false, fmt.Errorf("compaction summary unavailable") + } + + summary = strings.TrimSpace(summary) + "\n\n[Post-compaction refresh: Re-check AGENTS.md and MEMORY.md before continuing.]" + agent.Sessions.SetSummary(sessionKey, summary) + agent.Sessions.SetHistory(sessionKey, kept) + agent.Sessions.IncrementCompactionCount(sessionKey) + _ = agent.Sessions.Save(sessionKey) + + logger.InfoCF("agent", "Compaction safeguard completed", map[string]any{ + "session_key": sessionKey, + "history_tokens_before": historyTokens, + "kept_messages": len(kept), + "summarized_messages": len(toSummarize), + }) + return true, nil +} + +func (al *AgentLoop) generateCompactionSummary( + ctx context.Context, + agent *AgentInstance, + history []providers.Message, + existingSummary string, +) (string, error) { + safeMessages := make([]providers.Message, 0, len(history)) + for _, msg := range history { + if msg.Role != "user" && msg.Role != "assistant" && msg.Role != "tool" { + continue + } + content := strings.TrimSpace(msg.Content) + if content == "" { + continue + } + if msg.Role == "tool" && len(content) > 1200 { + head := 700 + tail := 300 + content = content[:head] + "\n...\n[tool result condensed]\n...\n" + content[len(content)-tail:] + } + msg.Content = content + safeMessages = append(safeMessages, msg) + } + if len(safeMessages) == 0 { + return strings.TrimSpace(existingSummary), nil + } + + maxChunkTokens := int(float64(agent.ContextWindow)*0.35) - 1024 + if maxChunkTokens < 512 { + maxChunkTokens = 512 + } + + chunks := al.splitMessagesByTokenBudget(safeMessages, maxChunkTokens) + summary := strings.TrimSpace(existingSummary) + for _, chunk := range chunks { + next, err := al.summarizeBatchStructured(ctx, agent, chunk, summary) + if err != nil { + return "", err + } + summary = strings.TrimSpace(next) + } + return summary, nil +} + +func (al *AgentLoop) summarizeBatchStructured( + ctx context.Context, + agent *AgentInstance, + batch []providers.Message, + existingSummary string, +) (string, error) { + var sb strings.Builder + sb.WriteString("Summarize this conversation segment for future continuity.\n") + sb.WriteString("Use concise markdown with sections: Intent, Decisions, Tool Results, Pending Actions, Constraints.\n") + if existingSummary != "" { + sb.WriteString("\nExisting summary:\n") + sb.WriteString(existingSummary) + sb.WriteString("\n") + } + sb.WriteString("\nConversation:\n") + for _, m := range batch { + sb.WriteString(m.Role) + sb.WriteString(": ") + sb.WriteString(m.Content) + sb.WriteString("\n") + } + + resp, err := agent.Provider.Chat( + ctx, + []providers.Message{{Role: "user", Content: sb.String()}}, + nil, + agent.Model, + map[string]any{ + "max_tokens": 800, + "temperature": 0.2, + }, + ) + if err != nil { + return "", err + } + return strings.TrimSpace(resp.Content), nil +} + +func (al *AgentLoop) splitMessagesByTokenBudget( + messages []providers.Message, + maxTokens int, +) [][]providers.Message { + if len(messages) == 0 || maxTokens <= 0 { + return nil + } + chunks := make([][]providers.Message, 0, 4) + current := make([]providers.Message, 0, 8) + currentTokens := 0 + for _, msg := range messages { + msgTokens := al.estimateMessageTokens(msg) + if len(current) > 0 && currentTokens+msgTokens > maxTokens { + chunks = append(chunks, current) + current = make([]providers.Message, 0, 8) + currentTokens = 0 + } + current = append(current, msg) + currentTokens += msgTokens + } + if len(current) > 0 { + chunks = append(chunks, current) + } + return chunks +} + +func maxInt(a, b int) int { + if a > b { + return a + } + return b +} + +func (al *AgentLoop) safeCompactionContext() (context.Context, context.CancelFunc) { + return context.WithTimeout(context.Background(), 90*time.Second) +} + +func normalizeCompactionMode(mode string) string { + switch strings.ToLower(strings.TrimSpace(mode)) { + case "", "safeguard": + return "safeguard" + case "off", "none", "disabled": + return "off" + case "legacy": + return "legacy" + default: + return "safeguard" + } +} + +func isCompactionModeEnabled(mode string) bool { + return normalizeCompactionMode(mode) != "off" +} diff --git a/pkg/agent/context.go b/pkg/agent/context.go index d8b887839..11109eaa2 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -1,38 +1,45 @@ package agent import ( - "errors" "fmt" - "io/fs" "os" "path/filepath" "runtime" "strings" "sync" "time" + "unicode/utf8" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/skills" + "github.com/sipeed/picoclaw/pkg/tools" ) +type ContextRuntimeSettings struct { + ContextWindowTokens int + PruningMode string + IncludeOldChitChat bool + SoftToolResultChars int + HardToolResultChars int + TriggerRatio float64 + BootstrapSnapshotEnabled bool + MemoryVectorEnabled bool + MemoryVectorDimensions int + MemoryVectorTopK int + MemoryVectorMinScore float64 + MemoryVectorMaxChars int + MemoryVectorRecentDays int +} + type ContextBuilder struct { - workspace string - skillsLoader *skills.SkillsLoader - memory *MemoryStore - - // Cache for system prompt to avoid rebuilding on every call. - // This fixes issue #607: repeated reprocessing of the entire context. - // The cache auto-invalidates when workspace source files change (mtime check). - systemPromptMutex sync.RWMutex - cachedSystemPrompt string - cachedAt time.Time // max observed mtime across tracked paths at cache build time - - // existedAtCache tracks which source file paths existed the last time the - // cache was built. This lets sourceFilesChanged detect files that are newly - // created (didn't exist at cache time, now exist) or deleted (existed at - // cache time, now gone) — both of which should trigger a cache rebuild. - existedAtCache map[string]bool + workspace string + skillsLoader *skills.SkillsLoader + memory *MemoryStore + tools *tools.ToolRegistry // Direct reference to tool registry + settings ContextRuntimeSettings + bootstrapMu sync.RWMutex + bootstrapCache map[string]string } func getGlobalConfigDir() string { @@ -50,46 +57,156 @@ func NewContextBuilder(workspace string) *ContextBuilder { builtinSkillsDir := filepath.Join(wd, "skills") globalSkillsDir := filepath.Join(getGlobalConfigDir(), "skills") + defaultSettings := ContextRuntimeSettings{ + PruningMode: "tools_only", + IncludeOldChitChat: true, + SoftToolResultChars: 2000, + HardToolResultChars: 350, + TriggerRatio: 0.8, + BootstrapSnapshotEnabled: true, + MemoryVectorEnabled: true, + MemoryVectorDimensions: defaultMemoryVectorDimensions, + MemoryVectorTopK: defaultMemoryVectorTopK, + MemoryVectorMinScore: defaultMemoryVectorMinScore, + MemoryVectorMaxChars: defaultMemoryVectorMaxContextChars, + MemoryVectorRecentDays: defaultMemoryVectorRecentDailyDays, + } + + memoryStore := NewMemoryStore(workspace) + memoryStore.SetVectorSettings(MemoryVectorSettings{ + Enabled: defaultSettings.MemoryVectorEnabled, + Dimensions: defaultSettings.MemoryVectorDimensions, + TopK: defaultSettings.MemoryVectorTopK, + MinScore: defaultSettings.MemoryVectorMinScore, + MaxContextChars: defaultSettings.MemoryVectorMaxChars, + RecentDailyDays: defaultSettings.MemoryVectorRecentDays, + }) + return &ContextBuilder{ - workspace: workspace, - skillsLoader: skills.NewSkillsLoader(workspace, globalSkillsDir, builtinSkillsDir), - memory: NewMemoryStore(workspace), + workspace: workspace, + skillsLoader: skills.NewSkillsLoader(workspace, globalSkillsDir, builtinSkillsDir), + memory: memoryStore, + bootstrapCache: map[string]string{}, + settings: defaultSettings, } } +// SetToolsRegistry sets the tools registry for dynamic tool summary generation. +func (cb *ContextBuilder) SetToolsRegistry(registry *tools.ToolRegistry) { + cb.tools = registry +} + +func (cb *ContextBuilder) SetRuntimeSettings(settings ContextRuntimeSettings) { + if settings.PruningMode == "" { + settings.PruningMode = "tools_only" + } + if settings.SoftToolResultChars <= 0 { + settings.SoftToolResultChars = 2000 + } + if settings.HardToolResultChars <= 0 { + settings.HardToolResultChars = 350 + } + if settings.TriggerRatio <= 0 || settings.TriggerRatio >= 1 { + settings.TriggerRatio = 0.8 + } + if settings.MemoryVectorDimensions <= 0 { + settings.MemoryVectorDimensions = defaultMemoryVectorDimensions + } + if settings.MemoryVectorTopK <= 0 { + settings.MemoryVectorTopK = defaultMemoryVectorTopK + } + if settings.MemoryVectorMinScore < 0 || settings.MemoryVectorMinScore >= 1 { + settings.MemoryVectorMinScore = defaultMemoryVectorMinScore + } + if settings.MemoryVectorMaxChars <= 0 { + settings.MemoryVectorMaxChars = defaultMemoryVectorMaxContextChars + } + if settings.MemoryVectorRecentDays <= 0 { + settings.MemoryVectorRecentDays = defaultMemoryVectorRecentDailyDays + } + cb.settings = settings + cb.memory.SetVectorSettings(MemoryVectorSettings{ + Enabled: settings.MemoryVectorEnabled, + Dimensions: settings.MemoryVectorDimensions, + TopK: settings.MemoryVectorTopK, + MinScore: settings.MemoryVectorMinScore, + MaxContextChars: settings.MemoryVectorMaxChars, + RecentDailyDays: settings.MemoryVectorRecentDays, + }) +} + func (cb *ContextBuilder) getIdentity() string { + now := time.Now().Format("2006-01-02 15:04 (Monday)") workspacePath, _ := filepath.Abs(filepath.Join(cb.workspace)) + runtime := fmt.Sprintf("%s %s, Go %s", runtime.GOOS, runtime.GOARCH, runtime.Version()) + + // Build tools section dynamically + toolsSection := cb.buildToolsSection() return fmt.Sprintf(`# picoclaw 🦞 You are picoclaw, a helpful AI assistant. +## Current Time +%s + +## Runtime +%s + ## Workspace Your workspace is at: %s - Memory: %s/memory/MEMORY.md - Daily Notes: %s/memory/YYYYMM/YYYYMMDD.md - Skills: %s/skills/{skill-name}/SKILL.md +%s + ## Important Rules 1. **ALWAYS use tools** - When you need to perform an action (schedule reminders, send messages, execute commands, etc.), you MUST call the appropriate tool. Do NOT just say you'll do it or pretend to do it. 2. **Be helpful and accurate** - When using tools, briefly explain what you're doing. -3. **Memory** - When interacting with me if something seems memorable, update %s/memory/MEMORY.md +3. **Memory** - When interacting with me if something seems memorable, update %s/memory/MEMORY.md`, + now, runtime, workspacePath, workspacePath, workspacePath, workspacePath, toolsSection, workspacePath) +} -4. **Context summaries** - Conversation summaries provided as context are approximate references only. They may be incomplete or outdated. Always defer to explicit user instructions over summary content.`, - workspacePath, workspacePath, workspacePath, workspacePath, workspacePath) +func (cb *ContextBuilder) buildToolsSection() string { + if cb.tools == nil { + return "" + } + + summaries := cb.tools.GetSummaries() + if len(summaries) == 0 { + return "" + } + + var sb strings.Builder + sb.WriteString("## Available Tools\n\n") + sb.WriteString( + "**CRITICAL**: You MUST use tools to perform actions. Do NOT pretend to execute commands or schedule tasks.\n\n", + ) + sb.WriteString("You have access to the following tools:\n\n") + for _, s := range summaries { + sb.WriteString(s) + sb.WriteString("\n") + } + + return sb.String() } func (cb *ContextBuilder) BuildSystemPrompt() string { + return cb.BuildSystemPromptForSession("") +} + +func (cb *ContextBuilder) BuildSystemPromptForSession(sessionKey string) string { parts := []string{} // Core identity section parts = append(parts, cb.getIdentity()) // Bootstrap files - bootstrapContent := cb.LoadBootstrapFiles() + bootstrapContent := cb.LoadBootstrapFiles(sessionKey) if bootstrapContent != "" { parts = append(parts, bootstrapContent) } @@ -114,227 +231,16 @@ The following skills extend your capabilities. To use a skill, read its SKILL.md return strings.Join(parts, "\n\n---\n\n") } -// BuildSystemPromptWithCache returns the cached system prompt if available -// and source files haven't changed, otherwise builds and caches it. -// Source file changes are detected via mtime checks (cheap stat calls). -func (cb *ContextBuilder) BuildSystemPromptWithCache() string { - // Try read lock first — fast path when cache is valid - cb.systemPromptMutex.RLock() - if cb.cachedSystemPrompt != "" && !cb.sourceFilesChangedLocked() { - result := cb.cachedSystemPrompt - cb.systemPromptMutex.RUnlock() - return result - } - cb.systemPromptMutex.RUnlock() - - // Acquire write lock for building - cb.systemPromptMutex.Lock() - defer cb.systemPromptMutex.Unlock() - - // Double-check: another goroutine may have rebuilt while we waited - if cb.cachedSystemPrompt != "" && !cb.sourceFilesChangedLocked() { - return cb.cachedSystemPrompt - } - - // Snapshot the baseline (existence + max mtime) BEFORE building the prompt. - // This way cachedAt reflects the pre-build state: if a file is modified - // during BuildSystemPrompt, its new mtime will be > baseline.maxMtime, - // so the next sourceFilesChangedLocked check will correctly trigger a - // rebuild. The alternative (baseline after build) risks caching stale - // content with a too-new baseline, making the staleness invisible. - baseline := cb.buildCacheBaseline() - prompt := cb.BuildSystemPrompt() - cb.cachedSystemPrompt = prompt - cb.cachedAt = baseline.maxMtime - cb.existedAtCache = baseline.existed - - logger.DebugCF("agent", "System prompt cached", - map[string]any{ - "length": len(prompt), - }) - - return prompt -} - -// InvalidateCache clears the cached system prompt. -// Normally not needed because the cache auto-invalidates via mtime checks, -// but this is useful for tests or explicit reload commands. -func (cb *ContextBuilder) InvalidateCache() { - cb.systemPromptMutex.Lock() - defer cb.systemPromptMutex.Unlock() - - cb.cachedSystemPrompt = "" - cb.cachedAt = time.Time{} - cb.existedAtCache = nil - - logger.DebugCF("agent", "System prompt cache invalidated", nil) -} - -// sourcePaths returns the workspace source file paths tracked for cache -// invalidation (bootstrap files + memory). The skills directory is handled -// separately in sourceFilesChangedLocked because it requires both directory- -// level and recursive file-level mtime checks. -func (cb *ContextBuilder) sourcePaths() []string { - return []string{ - filepath.Join(cb.workspace, "AGENTS.md"), - filepath.Join(cb.workspace, "SOUL.md"), - filepath.Join(cb.workspace, "USER.md"), - filepath.Join(cb.workspace, "IDENTITY.md"), - filepath.Join(cb.workspace, "memory", "MEMORY.md"), - } -} - -// cacheBaseline holds the file existence snapshot and the latest observed -// mtime across all tracked paths. Used as the cache reference point. -type cacheBaseline struct { - existed map[string]bool - maxMtime time.Time -} - -// buildCacheBaseline records which tracked paths currently exist and computes -// the latest mtime across all tracked files + skills directory contents. -// Called under write lock when the cache is built. -func (cb *ContextBuilder) buildCacheBaseline() cacheBaseline { - skillsDir := filepath.Join(cb.workspace, "skills") - - // All paths whose existence we track: source files + skills dir. - allPaths := append(cb.sourcePaths(), skillsDir) - - existed := make(map[string]bool, len(allPaths)) - var maxMtime time.Time - - for _, p := range allPaths { - info, err := os.Stat(p) - existed[p] = err == nil - if err == nil && info.ModTime().After(maxMtime) { - maxMtime = info.ModTime() +func (cb *ContextBuilder) LoadBootstrapFiles(sessionKey string) string { + if cb.settings.BootstrapSnapshotEnabled && strings.TrimSpace(sessionKey) != "" { + cb.bootstrapMu.RLock() + cached, ok := cb.bootstrapCache[sessionKey] + cb.bootstrapMu.RUnlock() + if ok { + return cached } } - // Walk skills files to capture their mtimes too. - // Use os.Stat (not d.Info) to match the stat method used in - // fileChangedSince / skillFilesModifiedSince for consistency. - _ = filepath.WalkDir(skillsDir, func(path string, d fs.DirEntry, walkErr error) error { - if walkErr == nil && !d.IsDir() { - if info, err := os.Stat(path); err == nil && info.ModTime().After(maxMtime) { - maxMtime = info.ModTime() - } - } - return nil - }) - - // If no tracked files exist yet (empty workspace), maxMtime is zero. - // Use a very old non-zero time so that: - // 1. cachedAt.IsZero() won't trigger perpetual rebuilds. - // 2. Any real file created afterwards has mtime > cachedAt, so it - // will be detected by fileChangedSince (unlike time.Now() which - // could race with a file whose mtime <= Now). - if maxMtime.IsZero() { - maxMtime = time.Unix(1, 0) - } - - return cacheBaseline{existed: existed, maxMtime: maxMtime} -} - -// sourceFilesChangedLocked checks whether any workspace source file has been -// modified, created, or deleted since the cache was last built. -// -// IMPORTANT: The caller MUST hold at least a read lock on systemPromptMutex. -// Go's sync.RWMutex is not reentrant, so this function must NOT acquire the -// lock itself (it would deadlock when called from BuildSystemPromptWithCache -// which already holds RLock or Lock). -func (cb *ContextBuilder) sourceFilesChangedLocked() bool { - if cb.cachedAt.IsZero() { - return true - } - - // Check tracked source files (bootstrap + memory). - for _, p := range cb.sourcePaths() { - if cb.fileChangedSince(p) { - return true - } - } - - // --- Skills directory (handled separately from sourcePaths) --- - // - // 1. Creation/deletion: tracked via existedAtCache, same as bootstrap files. - skillsDir := filepath.Join(cb.workspace, "skills") - if cb.fileChangedSince(skillsDir) { - return true - } - - // 2. Structural changes (add/remove entries inside the dir) are reflected - // in the directory's own mtime, which fileChangedSince already checks. - // - // 3. Content-only edits to files inside skills/ do NOT update the parent - // directory mtime on most filesystems, so we recursively walk to check - // individual file mtimes at any nesting depth. - if skillFilesModifiedSince(skillsDir, cb.cachedAt) { - return true - } - - return false -} - -// fileChangedSince returns true if a tracked source file has been modified, -// newly created, or deleted since the cache was built. -// -// Four cases: -// - existed at cache time, exists now -> check mtime -// - existed at cache time, gone now -> changed (deleted) -// - absent at cache time, exists now -> changed (created) -// - absent at cache time, gone now -> no change -func (cb *ContextBuilder) fileChangedSince(path string) bool { - // Defensive: if existedAtCache was never initialized, treat as changed - // so the cache rebuilds rather than silently serving stale data. - if cb.existedAtCache == nil { - return true - } - - existedBefore := cb.existedAtCache[path] - info, err := os.Stat(path) - existsNow := err == nil - - if existedBefore != existsNow { - return true // file was created or deleted - } - if !existsNow { - return false // didn't exist before, doesn't exist now - } - return info.ModTime().After(cb.cachedAt) -} - -// errWalkStop is a sentinel error used to stop filepath.WalkDir early. -// Using a dedicated error (instead of fs.SkipAll) makes the early-exit -// intent explicit and avoids the nilerr linter warning that would fire -// if the callback returned nil when its err parameter is non-nil. -var errWalkStop = errors.New("walk stop") - -// skillFilesModifiedSince recursively walks the skills directory and checks -// whether any file was modified after t. This catches content-only edits at -// any nesting depth (e.g. skills/name/docs/extra.md) that don't update -// parent directory mtimes. -func skillFilesModifiedSince(skillsDir string, t time.Time) bool { - changed := false - err := filepath.WalkDir(skillsDir, func(path string, d fs.DirEntry, walkErr error) error { - if walkErr == nil && !d.IsDir() { - if info, statErr := os.Stat(path); statErr == nil && info.ModTime().After(t) { - changed = true - return errWalkStop // stop walking - } - } - return nil - }) - // errWalkStop is expected (early exit on first changed file). - // os.IsNotExist means the skills dir doesn't exist yet — not an error. - // Any other error is unexpected and worth logging. - if err != nil && !errors.Is(err, errWalkStop) && !os.IsNotExist(err) { - logger.DebugCF("agent", "skills walk error", map[string]any{"error": err.Error()}) - } - return changed -} - -func (cb *ContextBuilder) LoadBootstrapFiles() string { bootstrapFiles := []string{ "AGENTS.md", "SOUL.md", @@ -350,29 +256,13 @@ func (cb *ContextBuilder) LoadBootstrapFiles() string { } } - return sb.String() -} - -// buildDynamicContext returns a short dynamic context string with per-request info. -// This changes every request (time, session) so it is NOT part of the cached prompt. -// LLM-side KV cache reuse is achieved by each provider adapter's native mechanism: -// - Anthropic: per-block cache_control (ephemeral) on the static SystemParts block -// - OpenAI / Codex: prompt_cache_key for prefix-based caching -// -// See: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching -// See: https://platform.openai.com/docs/guides/prompt-caching -func (cb *ContextBuilder) buildDynamicContext(channel, chatID string) string { - now := time.Now().Format("2006-01-02 15:04 (Monday)") - rt := fmt.Sprintf("%s %s, Go %s", runtime.GOOS, runtime.GOARCH, runtime.Version()) - - var sb strings.Builder - fmt.Fprintf(&sb, "## Current Time\n%s\n\n## Runtime\n%s", now, rt) - - if channel != "" && chatID != "" { - fmt.Fprintf(&sb, "\n\n## Current Session\nChannel: %s\nChat ID: %s", channel, chatID) + content := sb.String() + if cb.settings.BootstrapSnapshotEnabled && strings.TrimSpace(sessionKey) != "" { + cb.bootstrapMu.Lock() + cb.bootstrapCache[sessionKey] = content + cb.bootstrapMu.Unlock() } - - return sb.String() + return content } func (cb *ContextBuilder) BuildMessages( @@ -381,68 +271,45 @@ func (cb *ContextBuilder) BuildMessages( currentMessage string, media []string, channel, chatID string, +) []providers.Message { + return cb.BuildMessagesForSession( + "", + history, + summary, + currentMessage, + media, + channel, + chatID, + ) +} + +func (cb *ContextBuilder) BuildMessagesForSession( + sessionKey string, + history []providers.Message, + summary string, + currentMessage string, + media []string, + channel, chatID string, ) []providers.Message { messages := []providers.Message{} - // The static part (identity, bootstrap, skills, memory) is cached locally to - // avoid repeated file I/O and string building on every call (fixes issue #607). - // Dynamic parts (time, session, summary) are appended per request. - // Everything is sent as a single system message for provider compatibility: - // - Anthropic adapter extracts messages[0] (Role=="system") and maps its content - // to the top-level "system" parameter in the Messages API request. A single - // contiguous system block makes this extraction straightforward. - // - Codex maps only the first system message to its instructions field. - // - OpenAI-compat passes messages through as-is. - staticPrompt := cb.BuildSystemPromptWithCache() + systemPrompt := cb.BuildSystemPromptForSession(sessionKey) - // Build short dynamic context (time, runtime, session) — changes per request - dynamicCtx := cb.buildDynamicContext(channel, chatID) - - // Compose a single system message: static (cached) + dynamic + optional summary. - // Keeping all system content in one message ensures every provider adapter can - // extract it correctly (Anthropic adapter -> top-level system param, - // Codex -> instructions field). - // - // SystemParts carries the same content as structured blocks so that - // cache-aware adapters (Anthropic) can set per-block cache_control. - // The static block is marked "ephemeral" — its prefix hash is stable - // across requests, enabling LLM-side KV cache reuse. - stringParts := []string{staticPrompt, dynamicCtx} - - contentBlocks := []providers.ContentBlock{ - {Type: "text", Text: staticPrompt, CacheControl: &providers.CacheControl{Type: "ephemeral"}}, - {Type: "text", Text: dynamicCtx}, + // Add Current Session info if provided + if channel != "" && chatID != "" { + systemPrompt += fmt.Sprintf("\n\n## Current Session\nChannel: %s\nChat ID: %s", channel, chatID) } - if summary != "" { - summaryText := fmt.Sprintf( - "CONTEXT_SUMMARY: The following is an approximate summary of prior conversation "+ - "for reference only. It may be incomplete or outdated — always defer to explicit instructions.\n\n%s", - summary) - stringParts = append(stringParts, summaryText) - contentBlocks = append(contentBlocks, providers.ContentBlock{Type: "text", Text: summaryText}) - } - - fullSystemPrompt := strings.Join(stringParts, "\n\n---\n\n") - - // Log system prompt summary for debugging (debug mode only). - // Read cachedSystemPrompt under lock to avoid a data race with - // concurrent InvalidateCache / BuildSystemPromptWithCache writes. - cb.systemPromptMutex.RLock() - isCached := cb.cachedSystemPrompt != "" - cb.systemPromptMutex.RUnlock() - + // Log system prompt summary for debugging (debug mode only) logger.DebugCF("agent", "System prompt built", map[string]any{ - "static_chars": len(staticPrompt), - "dynamic_chars": len(dynamicCtx), - "total_chars": len(fullSystemPrompt), - "has_summary": summary != "", - "cached": isCached, + "total_chars": len(systemPrompt), + "total_lines": strings.Count(systemPrompt, "\n") + 1, + "section_count": strings.Count(systemPrompt, "\n\n---\n\n") + 1, }) // Log preview of system prompt (avoid logging huge content) - preview := fullSystemPrompt + preview := systemPrompt if len(preview) > 500 { preview = preview[:500] + "... (truncated)" } @@ -451,21 +318,36 @@ func (cb *ContextBuilder) BuildMessages( "preview": preview, }) + if summary != "" { + systemPrompt += "\n\n## Summary of Previous Conversation\n\n" + summary + } + + if cb.settings.MemoryVectorEnabled && strings.TrimSpace(currentMessage) != "" { + hits, err := cb.memory.SearchRelevant( + currentMessage, + cb.settings.MemoryVectorTopK, + cb.settings.MemoryVectorMinScore, + ) + if err != nil { + logger.WarnCF("agent", "Semantic memory retrieval failed", map[string]any{ + "error": err.Error(), + }) + } else if section := formatRetrievedMemoryContext(hits, cb.settings.MemoryVectorMaxChars); section != "" { + systemPrompt += "\n\n---\n\n" + section + } + } + + history = sanitizeHistoryForProvider(history) + history = cb.pruneHistoryForContext(history, systemPrompt) history = sanitizeHistoryForProvider(history) - // Single system message containing all context — compatible with all providers. - // SystemParts enables cache-aware adapters to set per-block cache_control; - // Content is the concatenated fallback for adapters that don't read SystemParts. messages = append(messages, providers.Message{ - Role: "system", - Content: fullSystemPrompt, - SystemParts: contentBlocks, + Role: "system", + Content: systemPrompt, }) - // Add conversation history messages = append(messages, history...) - // Add current user message if strings.TrimSpace(currentMessage) != "" { messages = append(messages, providers.Message{ Role: "user", @@ -476,6 +358,190 @@ func (cb *ContextBuilder) BuildMessages( return messages } +func formatRetrievedMemoryContext(hits []MemoryVectorHit, maxChars int) string { + if len(hits) == 0 { + return "" + } + if maxChars <= 0 { + maxChars = defaultMemoryVectorMaxContextChars + } + + var sb strings.Builder + sb.WriteString("# Retrieved Memory\n\n") + sb.WriteString("Use these semantic hits as hints; prefer current workspace files when they conflict.\n\n") + + remaining := maxChars + for _, hit := range hits { + if remaining <= 0 { + break + } + text := compactWhitespace(hit.Text) + if text == "" { + continue + } + line := fmt.Sprintf("- (score=%.2f, source=%s) %s\n", hit.Score, hit.Source, text) + if len(line) > remaining { + if remaining > 6 { + line = line[:remaining-4] + "...\n" + } else { + break + } + } + sb.WriteString(line) + remaining -= len(line) + } + + if remaining == maxChars { + return "" + } + return strings.TrimSpace(sb.String()) +} + +func (cb *ContextBuilder) pruneHistoryForContext( + history []providers.Message, + systemPrompt string, +) []providers.Message { + if len(history) == 0 || cb.settings.PruningMode == "off" || cb.settings.ContextWindowTokens <= 0 { + return history + } + + estimateTokens := func(msg providers.Message) int { + chars := utf8.RuneCountInString(msg.Content) + for _, tc := range msg.ToolCalls { + chars += utf8.RuneCountInString(tc.Name) + if tc.Function != nil { + chars += utf8.RuneCountInString(tc.Function.Name) + chars += utf8.RuneCountInString(tc.Function.Arguments) + } + } + if chars == 0 { + return 0 + } + return chars * 2 / 5 + } + + totalTokens := utf8.RuneCountInString(systemPrompt) * 2 / 5 + for _, msg := range history { + totalTokens += estimateTokens(msg) + } + + ratio := float64(totalTokens) / float64(cb.settings.ContextWindowTokens) + if ratio < cb.settings.TriggerRatio { + return history + } + + cutoff := len(history) - 8 + if cutoff <= 0 { + return history + } + + pruned := make([]providers.Message, 0, len(history)) + for i := 0; i < cutoff; i++ { + msg := history[i] + + if cb.settings.PruningMode == "tools_only" && msg.Role == "tool" && cb.settings.SoftToolResultChars > 0 { + raw := msg.Content + if len(raw) > cb.settings.SoftToolResultChars { + head := cb.settings.SoftToolResultChars * 7 / 10 + tail := cb.settings.SoftToolResultChars * 2 / 10 + if head+tail > len(raw) { + head = len(raw) + tail = 0 + } + msg.Content = raw[:head] + + "\n...\n[tool result condensed for context stability]\n...\n" + + raw[len(raw)-tail:] + } + } + + pruned = append(pruned, msg) + } + pruned = append(pruned, history[cutoff:]...) + + if cb.settings.IncludeOldChitChat { + pruned = compactOldChitChat(pruned, cutoff) + } + + totalTokens = utf8.RuneCountInString(systemPrompt) * 2 / 5 + for _, msg := range pruned { + totalTokens += estimateTokens(msg) + } + ratio = float64(totalTokens) / float64(cb.settings.ContextWindowTokens) + if ratio < cb.settings.TriggerRatio || cb.settings.HardToolResultChars <= 0 { + return pruned + } + + scanLimit := minInt(cutoff, len(pruned)) + for i := 0; i < scanLimit; i++ { + if ratio < cb.settings.TriggerRatio { + break + } + msg := pruned[i] + if msg.Role != "tool" || len(msg.Content) <= cb.settings.HardToolResultChars { + continue + } + pruned[i].Content = "[tool result omitted for context stability; details preserved in session history]" + totalTokens = utf8.RuneCountInString(systemPrompt) * 2 / 5 + for _, m := range pruned { + totalTokens += estimateTokens(m) + } + ratio = float64(totalTokens) / float64(cb.settings.ContextWindowTokens) + } + + return pruned +} + +func compactOldChitChat(history []providers.Message, cutoff int) []providers.Message { + if len(history) == 0 || cutoff <= 0 { + return history + } + + isLowSignal := func(msg providers.Message) bool { + if msg.Role != "user" && msg.Role != "assistant" { + return false + } + if len(msg.ToolCalls) > 0 || msg.ToolCallID != "" { + return false + } + text := strings.ToLower(strings.TrimSpace(msg.Content)) + if text == "" || len(text) > 40 { + return false + } + switch text { + case "ok", "okay", "thanks", "thank you", "got it", "roger", "understood", "好的", "收到", "谢谢": + return true + } + return false + } + + result := make([]providers.Message, 0, len(history)) + i := 0 + for i < len(history) { + if i >= cutoff || !isLowSignal(history[i]) { + result = append(result, history[i]) + i++ + continue + } + + j := i + for j < cutoff && isLowSignal(history[j]) { + j++ + } + runLen := j - i + if runLen >= 2 { + result = append(result, providers.Message{ + Role: "assistant", + Content: fmt.Sprintf("[History note: %d brief acknowledgements condensed]", runLen), + }) + } else { + result = append(result, history[i]) + } + i = j + } + + return result +} + func sanitizeHistoryForProvider(history []providers.Message) []providers.Message { if len(history) == 0 { return history @@ -483,17 +549,29 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message sanitized := make([]providers.Message, 0, len(history)) var pendingToolCalls map[string]struct{} + var pendingToolCallOrder []string + flushPendingToolCalls := func() { + if len(pendingToolCalls) == 0 { + pendingToolCalls = nil + pendingToolCallOrder = nil + return + } + for _, id := range pendingToolCallOrder { + if _, ok := pendingToolCalls[id]; !ok { + continue + } + sanitized = append(sanitized, providers.Message{ + Role: "tool", + ToolCallID: id, + Content: "[tool result missing in transcript; synthesized placeholder for provider compatibility]", + }) + } + pendingToolCalls = nil + pendingToolCallOrder = nil + } for _, msg := range history { switch msg.Role { - case "system": - // Drop system messages from history. BuildMessages always - // constructs its own single system message (static + dynamic + - // summary); extra system messages would break providers that - // only accept one (Anthropic, Codex). - logger.DebugCF("agent", "Dropping system message from history", map[string]any{}) - continue - case "tool": if pendingToolCalls == nil { logger.DebugCF("agent", "Dropping orphaned tool message", map[string]any{}) @@ -509,7 +587,7 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message if _, ok := pendingToolCalls[msg.ToolCallID]; !ok { logger.DebugCF( "agent", - "Dropping orphaned tool message with unknown call id", + "Dropping duplicate/orphaned tool message with unknown call id", map[string]any{"tool_call_id": msg.ToolCallID}, ) continue @@ -519,7 +597,7 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message sanitized = append(sanitized, msg) case "assistant": - pendingToolCalls = nil + flushPendingToolCalls() if len(msg.ToolCalls) > 0 { if len(sanitized) == 0 { @@ -537,19 +615,25 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message } pendingToolCalls = make(map[string]struct{}, len(msg.ToolCalls)) + pendingToolCallOrder = make([]string, 0, len(msg.ToolCalls)) for _, tc := range msg.ToolCalls { if tc.ID != "" { + if _, exists := pendingToolCalls[tc.ID]; exists { + continue + } pendingToolCalls[tc.ID] = struct{}{} + pendingToolCallOrder = append(pendingToolCallOrder, tc.ID) } } } sanitized = append(sanitized, msg) default: - pendingToolCalls = nil + flushPendingToolCalls() sanitized = append(sanitized, msg) } } + flushPendingToolCalls() return sanitized } diff --git a/pkg/agent/context_test.go b/pkg/agent/context_test.go index 3429986dc..662f63140 100644 --- a/pkg/agent/context_test.go +++ b/pkg/agent/context_test.go @@ -1,193 +1,14 @@ package agent import ( + "os" + "path/filepath" + "strings" "testing" "github.com/sipeed/picoclaw/pkg/providers" ) -func msg(role, content string) providers.Message { - return providers.Message{Role: role, Content: content} -} - -func assistantWithTools(toolIDs ...string) providers.Message { - calls := make([]providers.ToolCall, len(toolIDs)) - for i, id := range toolIDs { - calls[i] = providers.ToolCall{ID: id, Type: "function"} - } - return providers.Message{Role: "assistant", ToolCalls: calls} -} - -func toolResult(id string) providers.Message { - return providers.Message{Role: "tool", Content: "result", ToolCallID: id} -} - -func TestSanitizeHistoryForProvider_EmptyHistory(t *testing.T) { - result := sanitizeHistoryForProvider(nil) - if len(result) != 0 { - t.Fatalf("expected empty, got %d messages", len(result)) - } - - result = sanitizeHistoryForProvider([]providers.Message{}) - if len(result) != 0 { - t.Fatalf("expected empty, got %d messages", len(result)) - } -} - -func TestSanitizeHistoryForProvider_SingleToolCall(t *testing.T) { - history := []providers.Message{ - msg("user", "hello"), - assistantWithTools("A"), - toolResult("A"), - msg("assistant", "done"), - } - - result := sanitizeHistoryForProvider(history) - if len(result) != 4 { - t.Fatalf("expected 4 messages, got %d", len(result)) - } - assertRoles(t, result, "user", "assistant", "tool", "assistant") -} - -func TestSanitizeHistoryForProvider_MultiToolCalls(t *testing.T) { - history := []providers.Message{ - msg("user", "do two things"), - assistantWithTools("A", "B"), - toolResult("A"), - toolResult("B"), - msg("assistant", "both done"), - } - - result := sanitizeHistoryForProvider(history) - if len(result) != 5 { - t.Fatalf("expected 5 messages, got %d: %+v", len(result), roles(result)) - } - assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant") -} - -func TestSanitizeHistoryForProvider_AssistantToolCallAfterPlainAssistant(t *testing.T) { - history := []providers.Message{ - msg("user", "hi"), - msg("assistant", "thinking"), - assistantWithTools("A"), - toolResult("A"), - } - - result := sanitizeHistoryForProvider(history) - if len(result) != 2 { - t.Fatalf("expected 2 messages, got %d: %+v", len(result), roles(result)) - } - assertRoles(t, result, "user", "assistant") -} - -func TestSanitizeHistoryForProvider_OrphanedLeadingTool(t *testing.T) { - history := []providers.Message{ - toolResult("A"), - msg("user", "hello"), - } - - result := sanitizeHistoryForProvider(history) - if len(result) != 1 { - t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result)) - } - assertRoles(t, result, "user") -} - -func TestSanitizeHistoryForProvider_ToolAfterUserDropped(t *testing.T) { - history := []providers.Message{ - msg("user", "hello"), - toolResult("A"), - } - - result := sanitizeHistoryForProvider(history) - if len(result) != 1 { - t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result)) - } - assertRoles(t, result, "user") -} - -func TestSanitizeHistoryForProvider_ToolAfterAssistantNoToolCalls(t *testing.T) { - history := []providers.Message{ - msg("user", "hello"), - msg("assistant", "hi"), - toolResult("A"), - } - - result := sanitizeHistoryForProvider(history) - if len(result) != 2 { - t.Fatalf("expected 2 messages, got %d: %+v", len(result), roles(result)) - } - assertRoles(t, result, "user", "assistant") -} - -func TestSanitizeHistoryForProvider_AssistantToolCallAtStart(t *testing.T) { - history := []providers.Message{ - assistantWithTools("A"), - toolResult("A"), - msg("user", "hello"), - } - - result := sanitizeHistoryForProvider(history) - if len(result) != 1 { - t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result)) - } - assertRoles(t, result, "user") -} - -func TestSanitizeHistoryForProvider_MultiToolCallsThenNewRound(t *testing.T) { - history := []providers.Message{ - msg("user", "do two things"), - assistantWithTools("A", "B"), - toolResult("A"), - toolResult("B"), - msg("assistant", "done"), - msg("user", "hi"), - assistantWithTools("C"), - toolResult("C"), - msg("assistant", "done again"), - } - - result := sanitizeHistoryForProvider(history) - if len(result) != 9 { - t.Fatalf("expected 9 messages, got %d: %+v", len(result), roles(result)) - } - assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant", "user", "assistant", "tool", "assistant") -} - -func TestSanitizeHistoryForProvider_ConsecutiveMultiToolRounds(t *testing.T) { - history := []providers.Message{ - msg("user", "start"), - assistantWithTools("A", "B"), - toolResult("A"), - toolResult("B"), - assistantWithTools("C", "D"), - toolResult("C"), - toolResult("D"), - msg("assistant", "all done"), - } - - result := sanitizeHistoryForProvider(history) - if len(result) != 8 { - t.Fatalf("expected 8 messages, got %d: %+v", len(result), roles(result)) - } - assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant", "tool", "tool", "assistant") -} - -func TestSanitizeHistoryForProvider_PlainConversation(t *testing.T) { - history := []providers.Message{ - msg("user", "hello"), - msg("assistant", "hi"), - msg("user", "how are you"), - msg("assistant", "fine"), - } - - result := sanitizeHistoryForProvider(history) - if len(result) != 4 { - t.Fatalf("expected 4 messages, got %d", len(result)) - } - assertRoles(t, result, "user", "assistant", "user", "assistant") -} - func TestSanitizeHistoryForProvider_KeepMultipleToolOutputsFromOneAssistantTurn(t *testing.T) { history := []providers.Message{ {Role: "user", Content: "check two files"}, @@ -229,27 +50,165 @@ func TestSanitizeHistoryForProvider_DropToolOutputWithUnknownCallID(t *testing.T got := sanitizeHistoryForProvider(history) - if len(got) != 2 { - t.Fatalf("len(got) = %d, want 2; got=%#v", len(got), got) + if len(got) != 3 { + t.Fatalf("len(got) = %d, want 3; got=%#v", len(got), got) + } + if got[2].Role != "tool" || got[2].ToolCallID != "call_1" { + t.Fatalf("expected synthesized placeholder for call_1, got=%#v", got[2]) } } -func roles(msgs []providers.Message) []string { - r := make([]string, len(msgs)) - for i, m := range msgs { - r[i] = m.Role +func TestPruneHistoryForContext_ToolResultCondensed(t *testing.T) { + cb := NewContextBuilder(t.TempDir()) + cb.SetRuntimeSettings(ContextRuntimeSettings{ + ContextWindowTokens: 100, + PruningMode: "tools_only", + SoftToolResultChars: 80, + HardToolResultChars: 30, + TriggerRatio: 0.1, + BootstrapSnapshotEnabled: false, + }) + + history := []providers.Message{ + {Role: "user", Content: "run command"}, + {Role: "tool", Content: strings.Repeat("x", 600)}, + {Role: "assistant", Content: "done"}, + {Role: "user", Content: "next"}, + {Role: "assistant", Content: "ok"}, + {Role: "user", Content: "continue"}, + {Role: "assistant", Content: "ready"}, + {Role: "user", Content: "go"}, + {Role: "assistant", Content: "working"}, + {Role: "user", Content: "status"}, + } + + pruned := cb.pruneHistoryForContext(history, strings.Repeat("S", 500)) + if len(pruned) != len(history) { + t.Fatalf("len(pruned) = %d, want %d", len(pruned), len(history)) + } + if !strings.Contains(pruned[1].Content, "tool result") { + t.Fatalf("expected tool result to be condensed/omitted, got: %q", pruned[1].Content) } - return r } -func assertRoles(t *testing.T, msgs []providers.Message, expected ...string) { - t.Helper() - if len(msgs) != len(expected) { - t.Fatalf("role count mismatch: got %v, want %v", roles(msgs), expected) +func TestCompactOldChitChat_CondensesRun(t *testing.T) { + history := []providers.Message{ + {Role: "user", Content: "ok"}, + {Role: "assistant", Content: "thanks"}, + {Role: "user", Content: "received"}, + {Role: "assistant", Content: "actual content"}, + {Role: "user", Content: "keep recent"}, } - for i, exp := range expected { - if msgs[i].Role != exp { - t.Errorf("message[%d]: got role %q, want %q", i, msgs[i].Role, exp) - } + + got := compactOldChitChat(history, 4) + if len(got) >= len(history) { + t.Fatalf("expected condensed history, got len=%d", len(got)) + } + if !strings.Contains(strings.ToLower(got[0].Content), "condensed") { + t.Fatalf("expected condensed marker, got: %q", got[0].Content) } } + +func TestBuildMessagesForSession_IncludesRetrievedMemory(t *testing.T) { + workspace := t.TempDir() + memoryDir := filepath.Join(workspace, "memory") + if err := os.MkdirAll(memoryDir, 0o755); err != nil { + t.Fatalf("mkdir memory dir: %v", err) + } + + memoryContent := `# MEMORY + +## Long-term Facts +- Preferred editor is Neovim +` + if err := os.WriteFile(filepath.Join(memoryDir, "MEMORY.md"), []byte(memoryContent), 0o644); err != nil { + t.Fatalf("write memory file: %v", err) + } + + cb := NewContextBuilder(workspace) + cb.SetRuntimeSettings(ContextRuntimeSettings{ + ContextWindowTokens: 4096, + PruningMode: "off", + BootstrapSnapshotEnabled: false, + MemoryVectorEnabled: true, + MemoryVectorDimensions: 128, + MemoryVectorTopK: 3, + MemoryVectorMinScore: 0.01, + MemoryVectorMaxChars: 800, + MemoryVectorRecentDays: 7, + }) + + messages := cb.BuildMessagesForSession( + "sess-1", + nil, + "", + "Which editor do I usually prefer?", + nil, + "", + "", + ) + if len(messages) == 0 { + t.Fatalf("expected at least one message") + } + + system := strings.ToLower(messages[0].Content) + if !strings.Contains(system, "retrieved memory") { + t.Fatalf("expected retrieved memory section in system prompt, got:\n%s", messages[0].Content) + } + if !strings.Contains(system, "neovim") { + t.Fatalf("expected retrieved semantic hit to mention neovim, got:\n%s", messages[0].Content) + } +} + +func TestSanitizeHistoryForProvider_SynthesizesMissingToolOutputs(t *testing.T) { + history := []providers.Message{ + {Role: "user", Content: "check file"}, + { + Role: "assistant", + ToolCalls: []providers.ToolCall{ + {ID: "call_1", Name: "read_file"}, + {ID: "call_2", Name: "list_dir"}, + }, + }, + {Role: "tool", ToolCallID: "call_1", Content: "ok"}, + {Role: "assistant", Content: "continuing"}, + } + + got := sanitizeHistoryForProvider(history) + if len(got) != 5 { + t.Fatalf("len(got) = %d, want 5; got=%#v", len(got), got) + } + if got[3].Role != "tool" || got[3].ToolCallID != "call_2" { + t.Fatalf("expected synthesized tool output for call_2 at index 3, got %#v", got[3]) + } + if !strings.Contains(strings.ToLower(got[3].Content), "synthesized") { + t.Fatalf("expected synthesized marker, got %q", got[3].Content) + } +} + +func TestPruneHistoryForContext_DoesNotPanicAfterChitChatCompaction(t *testing.T) { + cb := NewContextBuilder(t.TempDir()) + cb.SetRuntimeSettings(ContextRuntimeSettings{ + ContextWindowTokens: 100, + PruningMode: "tools_only", + IncludeOldChitChat: true, + SoftToolResultChars: 80, + HardToolResultChars: 30, + TriggerRatio: 0.2, + BootstrapSnapshotEnabled: false, + }) + + history := []providers.Message{ + {Role: "user", Content: "ok"}, + {Role: "assistant", Content: "thanks"}, + {Role: "user", Content: "ok"}, + {Role: "assistant", Content: "thanks"}, + {Role: "user", Content: "ok"}, + {Role: "assistant", Content: "thanks"}, + {Role: "tool", Content: strings.Repeat("x", 600)}, + {Role: "assistant", Content: "ready"}, + {Role: "user", Content: "go"}, + } + + _ = cb.pruneHistoryForContext(history, strings.Repeat("S", 500)) +} diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index d5172d806..d5fbdbc28 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -15,22 +15,40 @@ import ( // AgentInstance represents a fully configured agent with its own workspace, // session manager, context builder, and tool registry. type AgentInstance struct { - ID string - Name string - Model string - Fallbacks []string - Workspace string - MaxIterations int - MaxTokens int - Temperature float64 - ContextWindow int - Provider providers.LLMProvider - Sessions *session.SessionManager - ContextBuilder *ContextBuilder - Tools *tools.ToolRegistry - Subagents *config.SubagentsConfig - SkillsFilter []string - Candidates []providers.FallbackCandidate + ID string + Name string + Model string + Fallbacks []string + Workspace string + MaxIterations int + MaxTokens int + Temperature float64 + ContextWindow int + Provider providers.LLMProvider + Sessions *session.SessionManager + ContextBuilder *ContextBuilder + Tools *tools.ToolRegistry + Subagents *config.SubagentsConfig + SkillsFilter []string + Candidates []providers.FallbackCandidate + CompactionMode string + CompactionReserveTokens int + CompactionKeepRecentTokens int + CompactionMaxHistoryShare float64 + MemoryFlushEnabled bool + MemoryFlushSoftThreshold int + ContextPruningMode string + ContextPruningIncludeChitChat bool + ContextPruningSoftToolChars int + ContextPruningHardToolChars int + ContextPruningTriggerRatio float64 + BootstrapSnapshotEnabled bool + MemoryVectorEnabled bool + MemoryVectorDimensions int + MemoryVectorTopK int + MemoryVectorMinScore float64 + MemoryVectorMaxContextChars int + MemoryVectorRecentDailyDays int } // NewAgentInstance creates an agent instance from config. @@ -91,6 +109,110 @@ func NewAgentInstance( temperature = *defaults.Temperature } + compactionMode := strings.TrimSpace(defaults.Compaction.Mode) + if compactionMode == "" { + compactionMode = "safeguard" + } + + compactionReserveTokens := defaults.Compaction.ReserveTokens + if compactionReserveTokens <= 0 { + compactionReserveTokens = 2048 + } + + compactionKeepRecentTokens := defaults.Compaction.KeepRecentTokens + if compactionKeepRecentTokens <= 0 { + compactionKeepRecentTokens = 2048 + } + + compactionMaxHistoryShare := defaults.Compaction.MaxHistoryShare + if compactionMaxHistoryShare <= 0 || compactionMaxHistoryShare > 0.9 { + compactionMaxHistoryShare = 0.5 + } + + memoryFlushEnabled := defaults.Compaction.MemoryFlush.Enabled + // Preserve default behavior if omitted from config. + if !defaults.Compaction.MemoryFlush.Enabled && defaults.Compaction.MemoryFlush.SoftThresholdTokens == 0 { + memoryFlushEnabled = true + } + + memoryFlushSoftThreshold := defaults.Compaction.MemoryFlush.SoftThresholdTokens + if memoryFlushSoftThreshold <= 0 { + memoryFlushSoftThreshold = 1500 + } + + contextPruningMode := strings.TrimSpace(defaults.ContextPruning.Mode) + if contextPruningMode == "" { + contextPruningMode = "tools_only" + } + + contextPruningSoftToolChars := defaults.ContextPruning.SoftToolResultChars + if contextPruningSoftToolChars <= 0 { + contextPruningSoftToolChars = 2000 + } + + contextPruningHardToolChars := defaults.ContextPruning.HardToolResultChars + if contextPruningHardToolChars <= 0 { + contextPruningHardToolChars = 350 + } + + contextPruningTriggerRatio := defaults.ContextPruning.TriggerRatio + if contextPruningTriggerRatio <= 0 || contextPruningTriggerRatio >= 1 { + contextPruningTriggerRatio = 0.8 + } + + bootstrapSnapshotEnabled := defaults.BootstrapSnapshot.Enabled + + memoryVectorEnabled := defaults.MemoryVector.Enabled + memoryVectorDimensions := defaults.MemoryVector.Dimensions + if memoryVectorDimensions <= 0 { + memoryVectorDimensions = defaultMemoryVectorDimensions + } + + memoryVectorTopK := defaults.MemoryVector.TopK + if memoryVectorTopK <= 0 { + memoryVectorTopK = defaultMemoryVectorTopK + } + + memoryVectorMinScore := defaults.MemoryVector.MinScore + if memoryVectorMinScore < 0 || memoryVectorMinScore >= 1 { + memoryVectorMinScore = defaultMemoryVectorMinScore + } + + memoryVectorMaxContextChars := defaults.MemoryVector.MaxContextChars + if memoryVectorMaxContextChars <= 0 { + memoryVectorMaxContextChars = defaultMemoryVectorMaxContextChars + } + + memoryVectorRecentDailyDays := defaults.MemoryVector.RecentDailyDays + if memoryVectorRecentDailyDays <= 0 { + memoryVectorRecentDailyDays = defaultMemoryVectorRecentDailyDays + } + + contextBuilder.SetRuntimeSettings(ContextRuntimeSettings{ + ContextWindowTokens: maxTokens, + PruningMode: contextPruningMode, + IncludeOldChitChat: defaults.ContextPruning.IncludeOldChitChat, + SoftToolResultChars: contextPruningSoftToolChars, + HardToolResultChars: contextPruningHardToolChars, + TriggerRatio: contextPruningTriggerRatio, + BootstrapSnapshotEnabled: bootstrapSnapshotEnabled, + MemoryVectorEnabled: memoryVectorEnabled, + MemoryVectorDimensions: memoryVectorDimensions, + MemoryVectorTopK: memoryVectorTopK, + MemoryVectorMinScore: memoryVectorMinScore, + MemoryVectorMaxChars: memoryVectorMaxContextChars, + MemoryVectorRecentDays: memoryVectorRecentDailyDays, + }) + + if memoryVectorEnabled { + toolsRegistry.Register(NewMemorySearchTool( + contextBuilder.memory, + memoryVectorTopK, + memoryVectorMinScore, + )) + toolsRegistry.Register(NewMemoryGetTool(contextBuilder.memory)) + } + // Resolve fallback candidates modelCfg := providers.ModelConfig{ Primary: model, @@ -99,22 +221,40 @@ func NewAgentInstance( candidates := providers.ResolveCandidates(modelCfg, defaults.Provider) return &AgentInstance{ - ID: agentID, - Name: agentName, - Model: model, - Fallbacks: fallbacks, - Workspace: workspace, - MaxIterations: maxIter, - MaxTokens: maxTokens, - Temperature: temperature, - ContextWindow: maxTokens, - Provider: provider, - Sessions: sessionsManager, - ContextBuilder: contextBuilder, - Tools: toolsRegistry, - Subagents: subagents, - SkillsFilter: skillsFilter, - Candidates: candidates, + ID: agentID, + Name: agentName, + Model: model, + Fallbacks: fallbacks, + Workspace: workspace, + MaxIterations: maxIter, + MaxTokens: maxTokens, + Temperature: temperature, + ContextWindow: maxTokens, + Provider: provider, + Sessions: sessionsManager, + ContextBuilder: contextBuilder, + Tools: toolsRegistry, + Subagents: subagents, + SkillsFilter: skillsFilter, + Candidates: candidates, + CompactionMode: compactionMode, + CompactionReserveTokens: compactionReserveTokens, + CompactionKeepRecentTokens: compactionKeepRecentTokens, + CompactionMaxHistoryShare: compactionMaxHistoryShare, + MemoryFlushEnabled: memoryFlushEnabled, + MemoryFlushSoftThreshold: memoryFlushSoftThreshold, + ContextPruningMode: contextPruningMode, + ContextPruningIncludeChitChat: defaults.ContextPruning.IncludeOldChitChat, + ContextPruningSoftToolChars: contextPruningSoftToolChars, + ContextPruningHardToolChars: contextPruningHardToolChars, + ContextPruningTriggerRatio: contextPruningTriggerRatio, + BootstrapSnapshotEnabled: bootstrapSnapshotEnabled, + MemoryVectorEnabled: memoryVectorEnabled, + MemoryVectorDimensions: memoryVectorDimensions, + MemoryVectorTopK: memoryVectorTopK, + MemoryVectorMinScore: memoryVectorMinScore, + MemoryVectorMaxContextChars: memoryVectorMaxContextChars, + MemoryVectorRecentDailyDays: memoryVectorRecentDailyDays, } } diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 6a37f3018..19b6dd49f 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -478,7 +478,8 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt history = agent.Sessions.GetHistory(opts.SessionKey) summary = agent.Sessions.GetSummary(opts.SessionKey) } - messages := agent.ContextBuilder.BuildMessages( + messages := agent.ContextBuilder.BuildMessagesForSession( + opts.SessionKey, history, summary, opts.UserMessage, @@ -639,10 +640,41 @@ func (al *AgentLoop) runLLMIteration( }) } - al.forceCompression(agent, opts.SessionKey) + compactionCtx, cancel := al.safeCompactionContext() + currentTokens := al.estimateTokens(agent.Sessions.GetHistory(opts.SessionKey)) + if flushed, flushErr := al.maybeFlushMemoryBeforeCompaction( + compactionCtx, + agent, + opts.SessionKey, + currentTokens, + ); flushErr != nil { + logger.WarnCF("agent", "Pre-compaction memory flush failed", map[string]any{ + "error": flushErr.Error(), + }) + } else if flushed { + logger.InfoCF("agent", "Pre-compaction memory flush completed", map[string]any{ + "session_key": opts.SessionKey, + }) + } + + compacted, compactErr := al.compactWithSafeguard(compactionCtx, agent, opts.SessionKey) + cancel() + if compactErr != nil { + logger.WarnCF("agent", "Compaction safeguard cancelled", map[string]any{ + "error": compactErr.Error(), + }) + break + } + if !compacted { + logger.WarnCF("agent", "Compaction safeguard skipped; preserving history", map[string]any{ + "session_key": opts.SessionKey, + }) + continue + } newHistory := agent.Sessions.GetHistory(opts.SessionKey) newSummary := agent.Sessions.GetSummary(opts.SessionKey) - messages = agent.ContextBuilder.BuildMessages( + messages = agent.ContextBuilder.BuildMessagesForSession( + opts.SessionKey, newHistory, newSummary, "", nil, opts.Channel, opts.ChatID, ) @@ -832,13 +864,48 @@ func (al *AgentLoop) maybeSummarize(agent *AgentInstance, sessionKey, channel, c tokenEstimate := al.estimateTokens(newHistory) threshold := agent.ContextWindow * 75 / 100 - if len(newHistory) > 20 || tokenEstimate > threshold { + if len(newHistory) > 100 || tokenEstimate > threshold { summarizeKey := agent.ID + ":" + sessionKey if _, loading := al.summarizing.LoadOrStore(summarizeKey, true); !loading { go func() { defer al.summarizing.Delete(summarizeKey) logger.Debug("Memory threshold reached. Optimizing conversation history...") - al.summarizeSession(agent, sessionKey) + if !constants.IsInternalChannel(channel) { + al.bus.PublishOutbound(bus.OutboundMessage{ + Channel: channel, + ChatID: chatID, + Content: "Memory threshold reached. Optimizing conversation history...", + }) + } + ctx, cancel := al.safeCompactionContext() + defer cancel() + + if flushed, err := al.maybeFlushMemoryBeforeCompaction( + ctx, + agent, + sessionKey, + tokenEstimate, + ); err != nil { + logger.WarnCF("agent", "Background memory flush failed", map[string]any{ + "session_key": sessionKey, + "error": err.Error(), + }) + } else if flushed { + logger.InfoCF("agent", "Background memory flush completed", map[string]any{ + "session_key": sessionKey, + }) + } + + if compacted, err := al.compactWithSafeguard(ctx, agent, sessionKey); err != nil { + logger.WarnCF("agent", "Background compaction cancelled", map[string]any{ + "session_key": sessionKey, + "error": err.Error(), + }) + } else if compacted { + logger.InfoCF("agent", "Background compaction completed", map[string]any{ + "session_key": sessionKey, + }) + } }() } } @@ -1099,11 +1166,25 @@ func (al *AgentLoop) summarizeBatch( // Uses a safe heuristic of 2.5 characters per token to account for CJK and other // overheads better than the previous 3 chars/token. func (al *AgentLoop) estimateTokens(messages []providers.Message) int { - totalChars := 0 + total := 0 for _, m := range messages { - totalChars += utf8.RuneCountInString(m.Content) + total += al.estimateMessageTokens(m) + } + return total +} + +func (al *AgentLoop) estimateMessageTokens(msg providers.Message) int { + totalChars := utf8.RuneCountInString(msg.Content) + for _, tc := range msg.ToolCalls { + totalChars += utf8.RuneCountInString(tc.Name) + if tc.Function != nil { + totalChars += utf8.RuneCountInString(tc.Function.Name) + totalChars += utf8.RuneCountInString(tc.Function.Arguments) + } + } + if totalChars == 0 { + return 0 } - // 2.5 chars per token = totalChars * 2 / 5 return totalChars * 2 / 5 } diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 4414398b1..d644ac7a9 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -615,9 +615,11 @@ func TestAgentLoop_ContextExhaustionRetry(t *testing.T) { t.Errorf("Expected 'Recovered from context error', got '%s'", response) } - // We expect 2 calls: 1st failed, 2nd succeeded - if provider.currentCall != 2 { - t.Errorf("Expected 2 calls (1 fail + 1 success), got %d", provider.currentCall) + // We expect at least 2 calls: + // 1) initial failed request + // 2) retry request succeeds (with or without compaction summary call) + if provider.currentCall < 2 { + t.Errorf("Expected at least 2 calls after retry, got %d", provider.currentCall) } // Check final history length diff --git a/pkg/agent/memory.go b/pkg/agent/memory.go index dd5f4441c..b23bc98a2 100644 --- a/pkg/agent/memory.go +++ b/pkg/agent/memory.go @@ -10,6 +10,7 @@ import ( "fmt" "os" "path/filepath" + "sort" "strings" "time" ) @@ -21,6 +22,30 @@ type MemoryStore struct { workspace string memoryDir string memoryFile string + vector *memoryVectorStore +} + +var memorySectionOrder = []string{ + "Profile", + "Long-term Facts", + "Active Goals", + "Constraints", + "Open Threads", + "Deprecated/Resolved", +} + +var memorySectionAliases = map[string]string{ + "profile": "Profile", + "long-term memory": "Long-term Facts", + "long term memory": "Long-term Facts", + "long-term facts": "Long-term Facts", + "active goals": "Active Goals", + "constraints": "Constraints", + "open threads": "Open Threads", + "open tasks": "Open Threads", + "pending tasks": "Open Threads", + "deprecated/resolved": "Deprecated/Resolved", + "resolved": "Deprecated/Resolved", } // NewMemoryStore creates a new MemoryStore with the given workspace path. @@ -32,10 +57,13 @@ func NewMemoryStore(workspace string) *MemoryStore { // Ensure memory directory exists os.MkdirAll(memoryDir, 0o755) + vectorSettings := defaultMemoryVectorSettings() + return &MemoryStore{ workspace: workspace, memoryDir: memoryDir, memoryFile: memoryFile, + vector: newMemoryVectorStore(memoryDir, memoryFile, vectorSettings), } } @@ -58,7 +86,11 @@ func (ms *MemoryStore) ReadLongTerm() string { // WriteLongTerm writes content to the long-term memory file (MEMORY.md). func (ms *MemoryStore) WriteLongTerm(content string) error { - return os.WriteFile(ms.memoryFile, []byte(content), 0o644) + if err := os.WriteFile(ms.memoryFile, []byte(content), 0o644); err != nil { + return err + } + ms.refreshVectorIndex() + return nil } // ReadToday reads today's daily note. @@ -95,7 +127,11 @@ func (ms *MemoryStore) AppendToday(content string) error { newContent = existingContent + "\n" + content } - return os.WriteFile(todayFile, []byte(newContent), 0o644) + if err := os.WriteFile(todayFile, []byte(newContent), 0o644); err != nil { + return err + } + ms.refreshVectorIndex() + return nil } // GetRecentDailyNotes returns daily notes from the last N days. @@ -149,3 +185,143 @@ func (ms *MemoryStore) GetMemoryContext() string { return sb.String() } + +// OrganizeWriteback rewrites MEMORY.md using stable sections and deduplicated bullets. +// It preserves existing content while integrating newly extracted memory notes. +func (ms *MemoryStore) OrganizeWriteback(extracted string) error { + base := parseMemorySections(ms.ReadLongTerm()) + incoming := parseMemorySections(extracted) + + for section, entries := range incoming { + base[section] = append(base[section], entries...) + } + normalizeMemorySections(base) + + return ms.WriteLongTerm(renderMemorySections(base)) +} + +func (ms *MemoryStore) SetVectorSettings(settings MemoryVectorSettings) { + if ms.vector == nil { + return + } + ms.vector.SetSettings(settings) +} + +// SearchRelevant runs semantic retrieval over MEMORY.md + recent daily notes. +func (ms *MemoryStore) SearchRelevant(query string, topK int, minScore float64) ([]MemoryVectorHit, error) { + if ms.vector == nil { + return nil, nil + } + return ms.vector.Search(query, topK, minScore) +} + +func (ms *MemoryStore) GetBySource(source string) (MemoryVectorHit, bool, error) { + if ms.vector == nil { + return MemoryVectorHit{}, false, nil + } + return ms.vector.GetBySource(source) +} + +func (ms *MemoryStore) refreshVectorIndex() { + if ms.vector == nil { + return + } + _ = ms.vector.Rebuild() +} + +func parseMemorySections(content string) map[string][]string { + sections := make(map[string][]string, len(memorySectionOrder)) + if strings.TrimSpace(content) == "" { + return sections + } + + current := "Long-term Facts" + for _, raw := range strings.Split(content, "\n") { + line := strings.TrimSpace(raw) + if line == "" { + continue + } + if strings.HasPrefix(line, "#") { + heading := strings.TrimSpace(strings.TrimLeft(line, "#")) + if normalized, ok := normalizeMemorySectionName(heading); ok { + current = normalized + } + continue + } + + entry := strings.TrimSpace(strings.TrimLeft(line, "-*+")) + if entry == "" { + continue + } + sections[current] = append(sections[current], entry) + } + return sections +} + +func normalizeMemorySectionName(name string) (string, bool) { + key := strings.ToLower(strings.TrimSpace(name)) + if key == "" { + return "", false + } + if section, ok := memorySectionAliases[key]; ok { + return section, true + } + for _, section := range memorySectionOrder { + if strings.EqualFold(section, name) { + return section, true + } + } + return "", false +} + +func normalizeMemorySections(sections map[string][]string) { + for section, entries := range sections { + seen := map[string]struct{}{} + deduped := make([]string, 0, len(entries)) + for _, entry := range entries { + clean := strings.TrimSpace(entry) + if clean == "" { + continue + } + key := strings.ToLower(clean) + if _, exists := seen[key]; exists { + continue + } + seen[key] = struct{}{} + deduped = append(deduped, clean) + } + sort.Strings(deduped) + sections[section] = deduped + } +} + +func renderMemorySections(sections map[string][]string) string { + var sb strings.Builder + sb.WriteString("# MEMORY\n\n") + sb.WriteString(fmt.Sprintf("_Last organized: %s_\n\n", time.Now().Format("2006-01-02 15:04"))) + + wroteSection := false + for _, section := range memorySectionOrder { + entries := sections[section] + if len(entries) == 0 { + continue + } + wroteSection = true + sb.WriteString("## ") + sb.WriteString(section) + sb.WriteString("\n") + for _, entry := range entries { + sb.WriteString("- ") + sb.WriteString(entry) + sb.WriteString("\n") + } + sb.WriteString("\n") + } + + if !wroteSection { + sb.WriteString("## Long-term Facts\n") + sb.WriteString("- (no durable facts recorded yet)\n") + } + + return strings.TrimSpace(sb.String()) + "\n" +} diff --git a/pkg/agent/memory_test.go b/pkg/agent/memory_test.go new file mode 100644 index 000000000..7d49190b8 --- /dev/null +++ b/pkg/agent/memory_test.go @@ -0,0 +1,48 @@ +package agent + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestMemoryOrganizeWriteback_DedupesAndSections(t *testing.T) { + workspace := t.TempDir() + ms := NewMemoryStore(workspace) + + initial := `# MEMORY + +## Long-term Facts +- Likes tea +- Likes tea + +## Open Threads +- Renew passport +` + if err := os.WriteFile(filepath.Join(workspace, "memory", "MEMORY.md"), []byte(initial), 0o644); err != nil { + t.Fatalf("write initial memory: %v", err) + } + + update := `## Long-term Facts +- Likes tea +- Works remotely + +## Active Goals +- Finish migration +` + if err := ms.OrganizeWriteback(update); err != nil { + t.Fatalf("OrganizeWriteback: %v", err) + } + + got := ms.ReadLongTerm() + if strings.Count(got, "- Likes tea") != 1 { + t.Fatalf("expected deduped fact, got:\n%s", got) + } + if !strings.Contains(got, "## Active Goals") { + t.Fatalf("expected Active Goals section, got:\n%s", got) + } + if !strings.Contains(got, "- Finish migration") { + t.Fatalf("expected merged goal, got:\n%s", got) + } +} diff --git a/pkg/agent/memory_tool.go b/pkg/agent/memory_tool.go new file mode 100644 index 000000000..e934c9bdb --- /dev/null +++ b/pkg/agent/memory_tool.go @@ -0,0 +1,170 @@ +package agent + +import ( + "context" + "fmt" + "strings" + + "github.com/sipeed/picoclaw/pkg/tools" +) + +// MemorySearchTool performs semantic lookup over persisted memory files. +type MemorySearchTool struct { + memory *MemoryStore + defaultTopK int + defaultMinScore float64 +} + +func NewMemorySearchTool(memory *MemoryStore, defaultTopK int, defaultMinScore float64) *MemorySearchTool { + if defaultTopK <= 0 { + defaultTopK = defaultMemoryVectorTopK + } + if defaultMinScore < 0 || defaultMinScore >= 1 { + defaultMinScore = defaultMemoryVectorMinScore + } + return &MemorySearchTool{ + memory: memory, + defaultTopK: defaultTopK, + defaultMinScore: defaultMinScore, + } +} + +func (t *MemorySearchTool) Name() string { + return "memory_search" +} + +func (t *MemorySearchTool) Description() string { + return "Semantically search MEMORY.md and recent daily notes for relevant facts" +} + +func (t *MemorySearchTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "query": map[string]any{ + "type": "string", + "description": "Natural-language query to search semantic memory", + }, + "top_k": map[string]any{ + "type": "integer", + "description": "Maximum number of hits to return (default from agent settings)", + }, + "min_score": map[string]any{ + "type": "number", + "description": "Minimum cosine similarity in [0,1), lower means broader recall", + }, + }, + "required": []string{"query"}, + } +} + +func (t *MemorySearchTool) Execute(ctx context.Context, args map[string]any) *tools.ToolResult { + _ = ctx + + if t.memory == nil { + return tools.ErrorResult("memory store unavailable") + } + + query, ok := args["query"].(string) + if !ok || strings.TrimSpace(query) == "" { + return tools.ErrorResult("query is required") + } + + topK := t.defaultTopK + if raw, ok := args["top_k"]; ok { + switch v := raw.(type) { + case int: + if v > 0 { + topK = v + } + case int64: + if v > 0 { + topK = int(v) + } + case float64: + if int(v) > 0 { + topK = int(v) + } + } + } + + minScore := t.defaultMinScore + if raw, ok := args["min_score"]; ok { + if v, ok := raw.(float64); ok && v >= 0 && v < 1 { + minScore = v + } + } + + hits, err := t.memory.SearchRelevant(query, topK, minScore) + if err != nil { + return tools.ErrorResult(fmt.Sprintf("memory search failed: %v", err)).WithError(err) + } + if len(hits) == 0 { + return tools.SilentResult("No relevant memory hits found.") + } + + var sb strings.Builder + sb.WriteString("Memory search hits:\n") + for _, hit := range hits { + sb.WriteString(fmt.Sprintf("- (score=%.2f, source=%s) %s\n", hit.Score, hit.Source, hit.Text)) + } + + return tools.SilentResult(strings.TrimSpace(sb.String())) +} + +// MemoryGetTool returns a specific memory item by its source citation. +type MemoryGetTool struct { + memory *MemoryStore +} + +func NewMemoryGetTool(memory *MemoryStore) *MemoryGetTool { + return &MemoryGetTool{memory: memory} +} + +func (t *MemoryGetTool) Name() string { + return "memory_get" +} + +func (t *MemoryGetTool) Description() string { + return "Retrieve one memory entry by source citation returned from memory_search" +} + +func (t *MemoryGetTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "source": map[string]any{ + "type": "string", + "description": "Citation source like MEMORY.md#Long-term Facts", + }, + }, + "required": []string{"source"}, + } +} + +func (t *MemoryGetTool) Execute(ctx context.Context, args map[string]any) *tools.ToolResult { + _ = ctx + + if t.memory == nil { + return tools.ErrorResult("memory store unavailable") + } + + source, ok := args["source"].(string) + if !ok || strings.TrimSpace(source) == "" { + return tools.ErrorResult("source is required") + } + + hit, found, err := t.memory.GetBySource(source) + if err != nil { + return tools.ErrorResult(fmt.Sprintf("memory get failed: %v", err)).WithError(err) + } + if !found { + return tools.SilentResult("Memory source not found.") + } + + return tools.SilentResult(fmt.Sprintf( + "Memory entry:\n- source=%s\n- content=%s", + hit.Source, + hit.Text, + )) +} diff --git a/pkg/agent/memory_tool_test.go b/pkg/agent/memory_tool_test.go new file mode 100644 index 000000000..81e17bd3e --- /dev/null +++ b/pkg/agent/memory_tool_test.go @@ -0,0 +1,93 @@ +package agent + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestMemorySearchTool_Execute(t *testing.T) { + workspace := t.TempDir() + ms := NewMemoryStore(workspace) + ms.SetVectorSettings(MemoryVectorSettings{ + Enabled: true, + Dimensions: 128, + TopK: 5, + MinScore: 0.01, + MaxContextChars: 1200, + RecentDailyDays: 7, + }) + + if err := os.MkdirAll(filepath.Join(workspace, "memory"), 0o755); err != nil { + t.Fatalf("mkdir memory dir: %v", err) + } + content := `# MEMORY + +## Long-term Facts +- Favorite editor: neovim +` + if err := os.WriteFile(filepath.Join(workspace, "memory", "MEMORY.md"), []byte(content), 0o644); err != nil { + t.Fatalf("write memory: %v", err) + } + + tool := NewMemorySearchTool(ms, 3, 0.01) + result := tool.Execute(context.Background(), map[string]any{ + "query": "what editor do I use", + }) + + if result.IsError { + t.Fatalf("expected successful tool result, got error: %s", result.ForLLM) + } + if !result.Silent { + t.Fatalf("expected memory_search to be silent") + } + if !strings.Contains(strings.ToLower(result.ForLLM), "neovim") { + t.Fatalf("expected tool output to include retrieved memory, got: %s", result.ForLLM) + } +} + +func TestMemoryGetTool_Execute(t *testing.T) { + workspace := t.TempDir() + ms := NewMemoryStore(workspace) + ms.SetVectorSettings(MemoryVectorSettings{ + Enabled: true, + Dimensions: 128, + TopK: 5, + MinScore: 0.01, + MaxContextChars: 1200, + RecentDailyDays: 7, + }) + + if err := os.MkdirAll(filepath.Join(workspace, "memory"), 0o755); err != nil { + t.Fatalf("mkdir memory dir: %v", err) + } + content := `# MEMORY + +## Long-term Facts +- Favorite editor: neovim +` + if err := os.WriteFile(filepath.Join(workspace, "memory", "MEMORY.md"), []byte(content), 0o644); err != nil { + t.Fatalf("write memory: %v", err) + } + + search := NewMemorySearchTool(ms, 3, 0.01) + searchResult := search.Execute(context.Background(), map[string]any{ + "query": "favorite editor", + }) + if searchResult.IsError { + t.Fatalf("memory_search failed: %s", searchResult.ForLLM) + } + + get := NewMemoryGetTool(ms) + getResult := get.Execute(context.Background(), map[string]any{ + "source": "MEMORY.md#Long-term Facts", + }) + if getResult.IsError { + t.Fatalf("memory_get failed: %s", getResult.ForLLM) + } + if !strings.Contains(strings.ToLower(getResult.ForLLM), "neovim") { + t.Fatalf("expected memory_get output to include neovim, got: %s", getResult.ForLLM) + } +} diff --git a/pkg/agent/memory_vector.go b/pkg/agent/memory_vector.go new file mode 100644 index 000000000..4bcb6de5c --- /dev/null +++ b/pkg/agent/memory_vector.go @@ -0,0 +1,645 @@ +package agent + +import ( + "crypto/sha1" + "encoding/hex" + "encoding/json" + "fmt" + "hash/fnv" + "math" + "os" + "path/filepath" + "sort" + "strings" + "sync" + "time" + "unicode" +) + +const ( + defaultMemoryVectorDimensions = 256 + defaultMemoryVectorTopK = 6 + defaultMemoryVectorMinScore = 0.15 + defaultMemoryVectorMaxContextChars = 1800 + defaultMemoryVectorRecentDailyDays = 14 + memoryVectorChunkChars = 280 +) + +// MemoryVectorSettings controls semantic memory indexing and retrieval behavior. +type MemoryVectorSettings struct { + Enabled bool + Dimensions int + TopK int + MinScore float64 + MaxContextChars int + RecentDailyDays int +} + +type MemoryVectorHit struct { + Source string + Text string + Score float64 +} + +type memoryVectorDocument struct { + ID string `json:"id"` + Source string `json:"source"` + Text string `json:"text"` + Vector []float32 `json:"vector"` +} + +type memoryVectorIndex struct { + Version int `json:"version"` + BuiltAt string `json:"built_at"` + Fingerprint string `json:"fingerprint"` + Dimensions int `json:"dimensions"` + Documents []memoryVectorDocument `json:"documents"` +} + +type memoryVectorSourceFile struct { + Path string + RelPath string + Size int64 + ModUnix int64 +} + +type memoryVectorStore struct { + memoryDir string + memoryFile string + indexPath string + + mu sync.Mutex + settings MemoryVectorSettings + cache *memoryVectorIndex +} + +func defaultMemoryVectorSettings() MemoryVectorSettings { + return MemoryVectorSettings{ + Enabled: true, + Dimensions: defaultMemoryVectorDimensions, + TopK: defaultMemoryVectorTopK, + MinScore: defaultMemoryVectorMinScore, + MaxContextChars: defaultMemoryVectorMaxContextChars, + RecentDailyDays: defaultMemoryVectorRecentDailyDays, + } +} + +func normalizeMemoryVectorSettings(settings MemoryVectorSettings) MemoryVectorSettings { + if settings.Dimensions <= 0 { + settings.Dimensions = defaultMemoryVectorDimensions + } + if settings.TopK <= 0 { + settings.TopK = defaultMemoryVectorTopK + } + if settings.MinScore < 0 || settings.MinScore >= 1 { + settings.MinScore = defaultMemoryVectorMinScore + } + if settings.MaxContextChars <= 0 { + settings.MaxContextChars = defaultMemoryVectorMaxContextChars + } + if settings.RecentDailyDays <= 0 { + settings.RecentDailyDays = defaultMemoryVectorRecentDailyDays + } + return settings +} + +func newMemoryVectorStore(memoryDir, memoryFile string, settings MemoryVectorSettings) *memoryVectorStore { + settings = normalizeMemoryVectorSettings(settings) + return &memoryVectorStore{ + memoryDir: memoryDir, + memoryFile: memoryFile, + indexPath: filepath.Join(memoryDir, "vector", "index.json"), + settings: settings, + } +} + +func (vs *memoryVectorStore) SetSettings(settings MemoryVectorSettings) { + vs.mu.Lock() + defer vs.mu.Unlock() + + normalized := normalizeMemoryVectorSettings(settings) + if vs.settings != normalized { + vs.settings = normalized + vs.cache = nil + } +} + +func (vs *memoryVectorStore) Rebuild() error { + vs.mu.Lock() + defer vs.mu.Unlock() + return vs.rebuildLocked() +} + +func (vs *memoryVectorStore) Search(query string, topK int, minScore float64) ([]MemoryVectorHit, error) { + query = strings.TrimSpace(query) + if query == "" { + return nil, nil + } + + vs.mu.Lock() + defer vs.mu.Unlock() + + if !vs.settings.Enabled { + return nil, nil + } + if topK <= 0 { + topK = vs.settings.TopK + } + if minScore < 0 { + minScore = 0 + } + + if err := vs.ensureIndexLocked(); err != nil { + return nil, err + } + if vs.cache == nil || len(vs.cache.Documents) == 0 { + return nil, nil + } + + queryVec := embedHashedText(query, vs.settings.Dimensions) + queryTerms := uniqueTokenSet(tokenizeForEmbedding(query)) + if len(queryVec) == 0 { + return nil, nil + } + + hits := make([]MemoryVectorHit, 0, minInt(topK, len(vs.cache.Documents))) + for _, doc := range vs.cache.Documents { + vectorScore := cosineSimilarity(queryVec, doc.Vector) + keywordScore := lexicalSimilarity(queryTerms, tokenizeForEmbedding(doc.Text)) + // Blend semantic and lexical signals to improve recall on terse notes and identifiers. + score := 0.8*vectorScore + 0.2*keywordScore + if score < minScore { + continue + } + hits = append(hits, MemoryVectorHit{ + Source: doc.Source, + Text: doc.Text, + Score: score, + }) + } + + sort.Slice(hits, func(i, j int) bool { + if hits[i].Score == hits[j].Score { + return hits[i].Source < hits[j].Source + } + return hits[i].Score > hits[j].Score + }) + + if len(hits) > topK { + hits = hits[:topK] + } + return hits, nil +} + +func (vs *memoryVectorStore) GetBySource(source string) (MemoryVectorHit, bool, error) { + source = strings.TrimSpace(source) + if source == "" { + return MemoryVectorHit{}, false, nil + } + + vs.mu.Lock() + defer vs.mu.Unlock() + + if !vs.settings.Enabled { + return MemoryVectorHit{}, false, nil + } + + if err := vs.ensureIndexLocked(); err != nil { + return MemoryVectorHit{}, false, err + } + if vs.cache == nil || len(vs.cache.Documents) == 0 { + return MemoryVectorHit{}, false, nil + } + + for _, doc := range vs.cache.Documents { + if doc.Source != source { + continue + } + return MemoryVectorHit{ + Source: doc.Source, + Text: doc.Text, + Score: 1, + }, true, nil + } + + return MemoryVectorHit{}, false, nil +} + +func (vs *memoryVectorStore) ensureIndexLocked() error { + sources, fingerprint, err := vs.collectSourceFilesLocked(time.Now()) + if err != nil { + return err + } + + if vs.cache != nil && + vs.cache.Fingerprint == fingerprint && + vs.cache.Dimensions == vs.settings.Dimensions { + return nil + } + + if disk, loadErr := vs.loadIndexLocked(); loadErr == nil && disk != nil { + if disk.Fingerprint == fingerprint && disk.Dimensions == vs.settings.Dimensions { + vs.cache = disk + return nil + } + } + + return vs.rebuildFromSourcesLocked(sources, fingerprint) +} + +func (vs *memoryVectorStore) rebuildLocked() error { + sources, fingerprint, err := vs.collectSourceFilesLocked(time.Now()) + if err != nil { + return err + } + return vs.rebuildFromSourcesLocked(sources, fingerprint) +} + +func (vs *memoryVectorStore) rebuildFromSourcesLocked( + sources []memoryVectorSourceFile, + fingerprint string, +) error { + docs, err := vs.buildDocumentsLocked(sources) + if err != nil { + return err + } + + index := &memoryVectorIndex{ + Version: 1, + BuiltAt: time.Now().Format(time.RFC3339), + Fingerprint: fingerprint, + Dimensions: vs.settings.Dimensions, + Documents: docs, + } + if err := vs.saveIndexLocked(index); err != nil { + return err + } + vs.cache = index + return nil +} + +func (vs *memoryVectorStore) collectSourceFilesLocked(now time.Time) ([]memoryVectorSourceFile, string, error) { + sources := make([]memoryVectorSourceFile, 0, vs.settings.RecentDailyDays+1) + + if info, err := os.Stat(vs.memoryFile); err == nil && !info.IsDir() { + sources = append(sources, memoryVectorSourceFile{ + Path: vs.memoryFile, + RelPath: "MEMORY.md", + Size: info.Size(), + ModUnix: info.ModTime().Unix(), + }) + } else if err != nil && !os.IsNotExist(err) { + return nil, "", err + } + + for i := 0; i < vs.settings.RecentDailyDays; i++ { + day := now.AddDate(0, 0, -i).Format("20060102") + candidate := filepath.Join(vs.memoryDir, day[:6], day+".md") + + info, err := os.Stat(candidate) + if err != nil { + if os.IsNotExist(err) { + continue + } + return nil, "", err + } + if info.IsDir() { + continue + } + + rel, relErr := filepath.Rel(vs.memoryDir, candidate) + if relErr != nil { + rel = filepath.Base(candidate) + } + sources = append(sources, memoryVectorSourceFile{ + Path: candidate, + RelPath: filepath.ToSlash(rel), + Size: info.Size(), + ModUnix: info.ModTime().Unix(), + }) + } + + fingerprint := buildSourceFingerprint(sources, vs.settings.Dimensions) + return sources, fingerprint, nil +} + +func buildSourceFingerprint(sources []memoryVectorSourceFile, dims int) string { + h := sha1.New() + fmt.Fprintf(h, "dims=%d\n", dims) + for _, src := range sources { + fmt.Fprintf(h, "%s|%d|%d\n", src.RelPath, src.Size, src.ModUnix) + } + return hex.EncodeToString(h.Sum(nil)) +} + +func (vs *memoryVectorStore) buildDocumentsLocked(sources []memoryVectorSourceFile) ([]memoryVectorDocument, error) { + docs := make([]memoryVectorDocument, 0, len(sources)*8) + + for _, src := range sources { + data, err := os.ReadFile(src.Path) + if err != nil { + return nil, err + } + content := string(data) + if strings.TrimSpace(content) == "" { + continue + } + + if src.RelPath == "MEMORY.md" { + sections := parseMemorySections(content) + for _, section := range memorySectionOrder { + for _, entry := range sections[section] { + text := compactWhitespace(entry) + if text == "" { + continue + } + payload := section + ": " + text + docs = append(docs, memoryVectorDocument{ + ID: buildMemoryVectorID(src.RelPath, section, text), + Source: fmt.Sprintf("%s#%s", src.RelPath, section), + Text: payload, + Vector: embedHashedText(payload, vs.settings.Dimensions), + }) + } + } + continue + } + + chunks := chunkMarkdownForVectors(content, memoryVectorChunkChars) + for idx, chunk := range chunks { + if strings.TrimSpace(chunk) == "" { + continue + } + docs = append(docs, memoryVectorDocument{ + ID: buildMemoryVectorID(src.RelPath, fmt.Sprintf("%d", idx+1), chunk), + Source: fmt.Sprintf("%s#%d", src.RelPath, idx+1), + Text: chunk, + Vector: embedHashedText(chunk, vs.settings.Dimensions), + }) + } + } + + return docs, nil +} + +func chunkMarkdownForVectors(content string, maxChars int) []string { + if maxChars <= 0 { + maxChars = memoryVectorChunkChars + } + + lines := strings.Split(content, "\n") + out := make([]string, 0, 8) + var current strings.Builder + currentHeading := "" + + flush := func() { + chunk := compactWhitespace(current.String()) + if chunk != "" { + out = append(out, chunk) + } + current.Reset() + } + + for _, raw := range lines { + line := strings.TrimSpace(raw) + if line == "" || line == "---" { + flush() + continue + } + if strings.HasPrefix(line, "#") { + flush() + currentHeading = strings.TrimSpace(strings.TrimLeft(line, "#")) + continue + } + + line = strings.TrimSpace(strings.TrimLeft(line, "-*+")) + if line == "" { + continue + } + if currentHeading != "" { + line = currentHeading + ": " + line + } + + if current.Len() == 0 { + current.WriteString(line) + } else { + if current.Len()+1+len(line) > maxChars { + flush() + current.WriteString(line) + } else { + current.WriteString(" ") + current.WriteString(line) + } + } + } + flush() + + return out +} + +func (vs *memoryVectorStore) loadIndexLocked() (*memoryVectorIndex, error) { + data, err := os.ReadFile(vs.indexPath) + if err != nil { + return nil, err + } + var index memoryVectorIndex + if err := json.Unmarshal(data, &index); err != nil { + return nil, err + } + return &index, nil +} + +func (vs *memoryVectorStore) saveIndexLocked(index *memoryVectorIndex) error { + if err := os.MkdirAll(filepath.Dir(vs.indexPath), 0o755); err != nil { + return err + } + payload, err := json.MarshalIndent(index, "", " ") + if err != nil { + return err + } + return os.WriteFile(vs.indexPath, payload, 0o644) +} + +func buildMemoryVectorID(parts ...string) string { + h := sha1.New() + for _, part := range parts { + if part == "" { + continue + } + h.Write([]byte(part)) + h.Write([]byte{0}) + } + return hex.EncodeToString(h.Sum(nil)) +} + +func embedHashedText(text string, dims int) []float32 { + if dims <= 0 { + return nil + } + + tokens := tokenizeForEmbedding(text) + if len(tokens) == 0 { + return nil + } + + tf := make(map[string]int, len(tokens)) + for _, token := range tokens { + tf[token]++ + } + + vec := make([]float64, dims) + for token, count := range tf { + h := fnv.New32a() + _, _ = h.Write([]byte(token)) + sum := h.Sum32() + + index := int(sum % uint32(dims)) + sign := 1.0 + if (sum>>31)&1 == 1 { + sign = -1 + } + + weight := 1.0 + math.Log(float64(count)) + vec[index] += sign * weight + } + + norm := 0.0 + for _, v := range vec { + norm += v * v + } + if norm == 0 { + return nil + } + norm = math.Sqrt(norm) + + out := make([]float32, dims) + for i, v := range vec { + out[i] = float32(v / norm) + } + return out +} + +func tokenizeForEmbedding(text string) []string { + text = strings.ToLower(strings.TrimSpace(text)) + if text == "" { + return nil + } + + tokens := make([]string, 0, 32) + var current []rune + flush := func() { + if len(current) == 0 { + return + } + tokens = append(tokens, string(current)) + current = current[:0] + } + + for _, r := range text { + if unicode.IsLetter(r) || unicode.IsDigit(r) { + current = append(current, r) + continue + } + flush() + } + flush() + + if len(tokens) > 0 { + expanded := make([]string, 0, len(tokens)*3) + for _, token := range tokens { + expanded = append(expanded, token) + runes := []rune(token) + if len(runes) < 4 { + continue + } + for i := 0; i+3 <= len(runes); i++ { + expanded = append(expanded, string(runes[i:i+3])) + } + } + return expanded + } + + // For very short non-word strings, fall back to rune-level tokens. + for _, r := range text { + if unicode.IsSpace(r) { + continue + } + tokens = append(tokens, string(r)) + } + return tokens +} + +func cosineSimilarity(a, b []float32) float64 { + n := minInt(len(a), len(b)) + if n == 0 { + return 0 + } + + dot := 0.0 + normA := 0.0 + normB := 0.0 + for i := 0; i < n; i++ { + av := float64(a[i]) + bv := float64(b[i]) + dot += av * bv + normA += av * av + normB += bv * bv + } + + if normA == 0 || normB == 0 { + return 0 + } + return dot / math.Sqrt(normA*normB) +} + +func compactWhitespace(s string) string { + return strings.Join(strings.Fields(strings.TrimSpace(s)), " ") +} + +func uniqueTokenSet(tokens []string) map[string]struct{} { + if len(tokens) == 0 { + return nil + } + out := make(map[string]struct{}, len(tokens)) + for _, token := range tokens { + if token == "" { + continue + } + out[token] = struct{}{} + } + return out +} + +func lexicalSimilarity(queryTokens map[string]struct{}, docTokens []string) float64 { + if len(queryTokens) == 0 || len(docTokens) == 0 { + return 0 + } + + docSet := uniqueTokenSet(docTokens) + if len(docSet) == 0 { + return 0 + } + + overlap := 0 + for token := range queryTokens { + if _, ok := docSet[token]; ok { + overlap++ + } + } + if overlap == 0 { + return 0 + } + + denom := len(queryTokens) + if denom == 0 { + return 0 + } + return float64(overlap) / float64(denom) +} + +func minInt(a, b int) int { + if a < b { + return a + } + return b +} diff --git a/pkg/agent/memory_vector_test.go b/pkg/agent/memory_vector_test.go new file mode 100644 index 000000000..fa781ab06 --- /dev/null +++ b/pkg/agent/memory_vector_test.go @@ -0,0 +1,107 @@ +package agent + +import ( + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +func TestMemoryStore_SearchRelevant_BuildsVectorIndex(t *testing.T) { + workspace := t.TempDir() + ms := NewMemoryStore(workspace) + ms.SetVectorSettings(MemoryVectorSettings{ + Enabled: true, + Dimensions: 128, + TopK: 4, + MinScore: 0.01, + MaxContextChars: 1200, + RecentDailyDays: 7, + }) + + memoryDir := filepath.Join(workspace, "memory") + if err := os.MkdirAll(memoryDir, 0o755); err != nil { + t.Fatalf("mkdir memory dir: %v", err) + } + + memoryContent := `# MEMORY + +## Open Threads +- Renew passport in March +- Compare flight options for Tokyo trip +` + if err := os.WriteFile(filepath.Join(memoryDir, "MEMORY.md"), []byte(memoryContent), 0o644); err != nil { + t.Fatalf("write MEMORY.md: %v", err) + } + + day := time.Now().Format("20060102") + dayPath := filepath.Join(memoryDir, day[:6], day+".md") + if err := os.MkdirAll(filepath.Dir(dayPath), 0o755); err != nil { + t.Fatalf("mkdir daily dir: %v", err) + } + daily := "# Daily\n\n- Follow up with passport office tomorrow\n" + if err := os.WriteFile(dayPath, []byte(daily), 0o644); err != nil { + t.Fatalf("write daily note: %v", err) + } + + hits, err := ms.SearchRelevant("passport renewal", 3, 0.01) + if err != nil { + t.Fatalf("SearchRelevant failed: %v", err) + } + if len(hits) == 0 { + t.Fatalf("expected semantic hits, got none") + } + + joined := strings.ToLower(hits[0].Text) + if !strings.Contains(joined, "passport") { + t.Fatalf("expected top hit to mention passport, got %q", hits[0].Text) + } + + indexPath := filepath.Join(memoryDir, "vector", "index.json") + if _, err := os.Stat(indexPath); err != nil { + t.Fatalf("expected vector index at %s: %v", indexPath, err) + } +} + +func TestMemoryStore_GetBySource(t *testing.T) { + workspace := t.TempDir() + ms := NewMemoryStore(workspace) + ms.SetVectorSettings(MemoryVectorSettings{ + Enabled: true, + Dimensions: 128, + TopK: 4, + MinScore: 0.01, + MaxContextChars: 1200, + RecentDailyDays: 7, + }) + + memoryDir := filepath.Join(workspace, "memory") + if err := os.MkdirAll(memoryDir, 0o755); err != nil { + t.Fatalf("mkdir memory dir: %v", err) + } + + memoryContent := `# MEMORY + +## Open Threads +- Prepare tax documents by end of month +` + if err := os.WriteFile(filepath.Join(memoryDir, "MEMORY.md"), []byte(memoryContent), 0o644); err != nil { + t.Fatalf("write MEMORY.md: %v", err) + } + + if _, err := ms.SearchRelevant("tax documents", 3, 0.01); err != nil { + t.Fatalf("SearchRelevant failed: %v", err) + } + + hit, found, err := ms.GetBySource("MEMORY.md#Open Threads") + if err != nil { + t.Fatalf("GetBySource failed: %v", err) + } + if !found { + t.Fatal("expected source to be found") + } + if !strings.Contains(strings.ToLower(hit.Text), "tax") { + t.Fatalf("expected hit text to include tax info, got: %q", hit.Text) + } +} diff --git a/pkg/config/config.go b/pkg/config/config.go index 28fe8d490..1b4aec388 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -178,6 +178,44 @@ type AgentDefaults struct { MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"` Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"` MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"` + Compaction AgentCompactionConfig `json:"compaction,omitempty"` + ContextPruning AgentContextPruningConfig `json:"context_pruning,omitempty"` + BootstrapSnapshot AgentBootstrapSnapshotConfig `json:"bootstrap_snapshot,omitempty"` + MemoryVector AgentMemoryVectorConfig `json:"memory_vector,omitempty"` +} + +type AgentCompactionConfig struct { + Mode string `json:"mode,omitempty"` + ReserveTokens int `json:"reserve_tokens,omitempty"` + KeepRecentTokens int `json:"keep_recent_tokens,omitempty"` + MaxHistoryShare float64 `json:"max_history_share,omitempty"` + MemoryFlush AgentCompactionMemoryFlushConfig `json:"memory_flush,omitempty"` +} + +type AgentCompactionMemoryFlushConfig struct { + Enabled bool `json:"enabled,omitempty"` + SoftThresholdTokens int `json:"soft_threshold_tokens,omitempty"` +} + +type AgentContextPruningConfig struct { + Mode string `json:"mode,omitempty"` + IncludeOldChitChat bool `json:"include_old_chitchat,omitempty"` + SoftToolResultChars int `json:"soft_tool_result_chars,omitempty"` + HardToolResultChars int `json:"hard_tool_result_chars,omitempty"` + TriggerRatio float64 `json:"trigger_ratio,omitempty"` +} + +type AgentBootstrapSnapshotConfig struct { + Enabled bool `json:"enabled,omitempty"` +} + +type AgentMemoryVectorConfig struct { + Enabled bool `json:"enabled,omitempty"` + Dimensions int `json:"dimensions,omitempty"` + TopK int `json:"top_k,omitempty"` + MinScore float64 `json:"min_score,omitempty"` + MaxContextChars int `json:"max_context_chars,omitempty"` + RecentDailyDays int `json:"recent_daily_days,omitempty"` } // GetModelName returns the effective model name for the agent defaults. diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index bf56b7f34..f4b254dcd 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -361,6 +361,24 @@ func TestDefaultConfig_OpenAIWebSearchEnabled(t *testing.T) { } } +func TestDefaultConfig_MemoryVectorDefaults(t *testing.T) { + cfg := DefaultConfig() + mv := cfg.Agents.Defaults.MemoryVector + + if !mv.Enabled { + t.Fatal("memory vector should be enabled by default") + } + if mv.Dimensions <= 0 { + t.Fatal("memory vector dimensions should be > 0") + } + if mv.TopK <= 0 { + t.Fatal("memory vector top_k should be > 0") + } + if mv.MinScore < 0 || mv.MinScore >= 1 { + t.Fatal("memory vector min_score should be in [0,1)") + } +} + func TestLoadConfig_OpenAIWebSearchDefaultsTrueWhenUnset(t *testing.T) { dir := t.TempDir() configPath := filepath.Join(dir, "config.json") diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index a1db7ab3e..638cb85fb 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -17,6 +17,34 @@ func DefaultConfig() *Config { MaxTokens: 8192, Temperature: nil, // nil means use provider default MaxToolIterations: 20, + Compaction: AgentCompactionConfig{ + Mode: "safeguard", + ReserveTokens: 2048, + KeepRecentTokens: 2048, + MaxHistoryShare: 0.5, + MemoryFlush: AgentCompactionMemoryFlushConfig{ + Enabled: true, + SoftThresholdTokens: 1500, + }, + }, + ContextPruning: AgentContextPruningConfig{ + Mode: "tools_only", + IncludeOldChitChat: true, + SoftToolResultChars: 2000, + HardToolResultChars: 350, + TriggerRatio: 0.8, + }, + BootstrapSnapshot: AgentBootstrapSnapshotConfig{ + Enabled: true, + }, + MemoryVector: AgentMemoryVectorConfig{ + Enabled: true, + Dimensions: 256, + TopK: 6, + MinScore: 0.15, + MaxContextChars: 1800, + RecentDailyDays: 14, + }, }, }, Bindings: []AgentBinding{}, diff --git a/pkg/session/manager.go b/pkg/session/manager.go index aad54ec4c..88f27c8a1 100644 --- a/pkg/session/manager.go +++ b/pkg/session/manager.go @@ -13,11 +13,14 @@ import ( ) type Session struct { - Key string `json:"key"` - Messages []providers.Message `json:"messages"` - Summary string `json:"summary,omitempty"` - Created time.Time `json:"created"` - Updated time.Time `json:"updated"` + Key string `json:"key"` + Messages []providers.Message `json:"messages"` + Summary string `json:"summary,omitempty"` + CompactionCount int `json:"compaction_count,omitempty"` + MemoryFlushAt time.Time `json:"memory_flush_at,omitempty"` + MemoryFlushCompactionCount int `json:"memory_flush_compaction_count,omitempty"` + Created time.Time `json:"created"` + Updated time.Time `json:"updated"` } type SessionManager struct { @@ -146,6 +149,43 @@ func (sm *SessionManager) TruncateHistory(key string, keepLast int) { session.Updated = time.Now() } +func (sm *SessionManager) IncrementCompactionCount(key string) int { + sm.mu.Lock() + defer sm.mu.Unlock() + + session, ok := sm.sessions[key] + if !ok { + return 0 + } + session.CompactionCount++ + session.Updated = time.Now() + return session.CompactionCount +} + +func (sm *SessionManager) MarkMemoryFlush(key string, compactionCount int) { + sm.mu.Lock() + defer sm.mu.Unlock() + + session, ok := sm.sessions[key] + if !ok { + return + } + session.MemoryFlushAt = time.Now() + session.MemoryFlushCompactionCount = compactionCount + session.Updated = time.Now() +} + +func (sm *SessionManager) GetCompactionState(key string) (count int, flushCount int, flushAt time.Time) { + sm.mu.RLock() + defer sm.mu.RUnlock() + + session, ok := sm.sessions[key] + if !ok { + return 0, 0, time.Time{} + } + return session.CompactionCount, session.MemoryFlushCompactionCount, session.MemoryFlushAt +} + // sanitizeFilename converts a session key into a cross-platform safe filename. // Session keys use "channel:chatID" (e.g. "telegram:123456") but ':' is the // volume separator on Windows, so filepath.Base would misinterpret the key. @@ -179,10 +219,13 @@ func (sm *SessionManager) Save(key string) error { } snapshot := Session{ - Key: stored.Key, - Summary: stored.Summary, - Created: stored.Created, - Updated: stored.Updated, + Key: stored.Key, + Summary: stored.Summary, + CompactionCount: stored.CompactionCount, + MemoryFlushAt: stored.MemoryFlushAt, + MemoryFlushCompactionCount: stored.MemoryFlushCompactionCount, + Created: stored.Created, + Updated: stored.Updated, } if len(stored.Messages) > 0 { snapshot.Messages = make([]providers.Message, len(stored.Messages)) @@ -293,10 +336,13 @@ func (sm *SessionManager) GetSessionSnapshot(key string) (*Session, bool) { } snapshot := Session{ - Key: stored.Key, - Summary: stored.Summary, - Created: stored.Created, - Updated: stored.Updated, + Key: stored.Key, + Summary: stored.Summary, + CompactionCount: stored.CompactionCount, + MemoryFlushAt: stored.MemoryFlushAt, + MemoryFlushCompactionCount: stored.MemoryFlushCompactionCount, + Created: stored.Created, + Updated: stored.Updated, } if len(stored.Messages) > 0 { snapshot.Messages = make([]providers.Message, len(stored.Messages)) @@ -316,10 +362,13 @@ func (sm *SessionManager) ListSessionSnapshots() []Session { snapshots := make([]Session, 0, len(sm.sessions)) for _, stored := range sm.sessions { snapshot := Session{ - Key: stored.Key, - Summary: stored.Summary, - Created: stored.Created, - Updated: stored.Updated, + Key: stored.Key, + Summary: stored.Summary, + CompactionCount: stored.CompactionCount, + MemoryFlushAt: stored.MemoryFlushAt, + MemoryFlushCompactionCount: stored.MemoryFlushCompactionCount, + Created: stored.Created, + Updated: stored.Updated, } if len(stored.Messages) > 0 { snapshot.Messages = make([]providers.Message, len(stored.Messages)) diff --git a/pkg/session/manager_test.go b/pkg/session/manager_test.go index 0800fc0b1..e78be5f2d 100644 --- a/pkg/session/manager_test.go +++ b/pkg/session/manager_test.go @@ -116,3 +116,31 @@ func TestListSessionSnapshots_SortedByUpdatedDesc(t *testing.T) { t.Fatalf("snapshot mutation leaked into manager state, got %q", history[0].Content) } } + +func TestCompactionStateLifecycle(t *testing.T) { + sm := NewSessionManager(t.TempDir()) + key := "agent:main:main" + sm.AddMessage(key, "user", "hello") + + count, flushedCount, flushAt := sm.GetCompactionState(key) + if count != 0 || flushedCount != 0 || !flushAt.IsZero() { + t.Fatalf("initial compaction state unexpected: count=%d flushed=%d flushAt=%v", count, flushedCount, flushAt) + } + + next := sm.IncrementCompactionCount(key) + if next != 1 { + t.Fatalf("IncrementCompactionCount = %d, want 1", next) + } + + sm.MarkMemoryFlush(key, next) + count, flushedCount, flushAt = sm.GetCompactionState(key) + if count != 1 { + t.Fatalf("count = %d, want 1", count) + } + if flushedCount != 1 { + t.Fatalf("flushedCount = %d, want 1", flushedCount) + } + if flushAt.IsZero() { + t.Fatal("flushAt should be set") + } +} diff --git a/pkg/tools/cron.go b/pkg/tools/cron.go index 3aa23ea85..5cf63c4fd 100644 --- a/pkg/tools/cron.go +++ b/pkg/tools/cron.go @@ -3,6 +3,7 @@ package tools import ( "context" "fmt" + "strings" "sync" "time" @@ -149,33 +150,44 @@ func (t *CronTool) addJob(args map[string]any) *ToolResult { var schedule cron.CronSchedule - // Check for at_seconds (one-time), every_seconds (recurring), or cron_expr - atSeconds, hasAt := args["at_seconds"].(float64) - everySeconds, hasEvery := args["every_seconds"].(float64) - cronExpr, hasCron := args["cron_expr"].(string) + // Check for at_seconds (one-time), every_seconds (recurring), or cron_expr. + // Some model providers emit unused numeric args as 0. Treat 0 as "unset" to avoid + // mistakenly creating an immediate "at" schedule that fails validation. + atSeconds, hasAt, err := parsePositiveSecondsArg(args, "at_seconds") + if err != nil { + return ErrorResult(err.Error()) + } + everySeconds, hasEvery, err := parsePositiveSecondsArg(args, "every_seconds") + if err != nil { + return ErrorResult(err.Error()) + } + cronExpr, _ := args["cron_expr"].(string) + cronExpr = strings.TrimSpace(cronExpr) + hasCron := cronExpr != "" timezone, _ := args["timezone"].(string) - // Priority: at_seconds > every_seconds > cron_expr - if hasAt { - atMS := time.Now().UnixMilli() + int64(atSeconds)*1000 - schedule = cron.CronSchedule{ - Kind: "at", - AtMS: &atMS, - } - } else if hasEvery { - everyMS := int64(everySeconds) * 1000 - schedule = cron.CronSchedule{ - Kind: "every", - EveryMS: &everyMS, - } - } else if hasCron { + // Priority: cron_expr > every_seconds > at_seconds + // This is resilient when LLMs include unused defaults (e.g. at_seconds=0). + if hasCron { schedule = cron.CronSchedule{ Kind: "cron", Expr: cronExpr, TZ: timezone, } + } else if hasEvery { + everyMS := everySeconds * 1000 + schedule = cron.CronSchedule{ + Kind: "every", + EveryMS: &everyMS, + } + } else if hasAt { + atMS := time.Now().UnixMilli() + atSeconds*1000 + schedule = cron.CronSchedule{ + Kind: "at", + AtMS: &atMS, + } } else { - return ErrorResult("one of at_seconds, every_seconds, or cron_expr is required") + return ErrorResult("one of at_seconds (>0), every_seconds (>0), or cron_expr is required") } // Read deliver parameter, default to true @@ -217,6 +229,25 @@ func (t *CronTool) addJob(args map[string]any) *ToolResult { return SilentResult(fmt.Sprintf("Cron job added: %s (id: %s)", job.Name, job.ID)) } +func parsePositiveSecondsArg(args map[string]any, key string) (int64, bool, error) { + raw, exists := args[key] + if !exists || raw == nil { + return 0, false, nil + } + + n, err := toInt(raw) + if err != nil { + return 0, false, fmt.Errorf("%s must be an integer", key) + } + if n < 0 { + return 0, false, fmt.Errorf("%s must be >= 0", key) + } + if n == 0 { + return 0, false, nil + } + return int64(n), true, nil +} + func (t *CronTool) listJobs() *ToolResult { jobs := t.cronService.ListJobs(false) diff --git a/pkg/tools/cron_test.go b/pkg/tools/cron_test.go new file mode 100644 index 000000000..fd2b8acce --- /dev/null +++ b/pkg/tools/cron_test.go @@ -0,0 +1,101 @@ +package tools + +import ( + "context" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" + cronpkg "github.com/sipeed/picoclaw/pkg/cron" +) + +func newCronToolForTest(t *testing.T) *CronTool { + t.Helper() + + workspace := t.TempDir() + storePath := filepath.Join(workspace, "cron", "jobs.json") + cronService := cronpkg.NewCronService(storePath, nil) + + tool := NewCronTool( + cronService, + nil, + bus.NewMessageBus(), + workspace, + true, + 5*time.Second, + config.DefaultConfig(), + ) + tool.SetContext("cli", "direct") + return tool +} + +func TestCronToolAddJob_UsesCronExprWhenZeroNumericFieldsPresent(t *testing.T) { + tool := newCronToolForTest(t) + + result := tool.Execute(context.Background(), map[string]any{ + "action": "add", + "message": "daily check", + "at_seconds": 0, + "every_seconds": 0, + "cron_expr": "*/5 * * * *", + "timezone": "Asia/Shanghai", + }) + if result.IsError { + t.Fatalf("expected add success, got error: %s", result.ForLLM) + } + + jobs := tool.cronService.ListJobs(false) + if len(jobs) != 1 { + t.Fatalf("expected 1 job, got %d", len(jobs)) + } + if jobs[0].Schedule.Kind != "cron" { + t.Fatalf("expected cron schedule kind, got %q", jobs[0].Schedule.Kind) + } + if jobs[0].Schedule.Expr != "*/5 * * * *" { + t.Fatalf("expected cron expr to be preserved, got %q", jobs[0].Schedule.Expr) + } +} + +func TestCronToolAddJob_UsesEveryWhenAtIsZero(t *testing.T) { + tool := newCronToolForTest(t) + + result := tool.Execute(context.Background(), map[string]any{ + "action": "add", + "message": "hourly check", + "at_seconds": 0, + "every_seconds": 3600, + }) + if result.IsError { + t.Fatalf("expected add success, got error: %s", result.ForLLM) + } + + jobs := tool.cronService.ListJobs(false) + if len(jobs) != 1 { + t.Fatalf("expected 1 job, got %d", len(jobs)) + } + if jobs[0].Schedule.Kind != "every" { + t.Fatalf("expected every schedule kind, got %q", jobs[0].Schedule.Kind) + } + if jobs[0].Schedule.EveryMS == nil || *jobs[0].Schedule.EveryMS != 3600*1000 { + t.Fatalf("expected everyMs=3600000, got %+v", jobs[0].Schedule.EveryMS) + } +} + +func TestCronToolAddJob_NegativeSecondsRejected(t *testing.T) { + tool := newCronToolForTest(t) + + result := tool.Execute(context.Background(), map[string]any{ + "action": "add", + "message": "invalid", + "at_seconds": -1, + }) + if !result.IsError { + t.Fatalf("expected error for negative at_seconds, got success: %s", result.ForLLM) + } + if !strings.Contains(result.ForLLM, "at_seconds must be >= 0") { + t.Fatalf("unexpected error text: %s", result.ForLLM) + } +}