diff --git a/cmd/picoclaw/main.go b/cmd/picoclaw/main.go index 7d535e6dd..198ff92c1 100644 --- a/cmd/picoclaw/main.go +++ b/cmd/picoclaw/main.go @@ -580,7 +580,11 @@ func gatewayCmd() { // Setup cron tool and service execTimeout := time.Duration(cfg.Tools.Cron.ExecTimeoutMinutes) * time.Minute - cronService := setupCronTool(agentLoop, msgBus, cfg.WorkspacePath(), cfg.Agents.Defaults.RestrictToWorkspace, execTimeout) + var cronOpts []cron.CronOption + if del := agentLoop.MemoryDelegate(); del != nil { + cronOpts = append(cronOpts, cron.WithCronDelegate(del, "picoclaw")) + } + cronService := setupCronTool(agentLoop, msgBus, cfg.WorkspacePath(), cfg.Agents.Defaults.RestrictToWorkspace, execTimeout, cronOpts...) heartbeatService := heartbeat.NewHeartbeatService( cfg.WorkspacePath(), @@ -1141,11 +1145,10 @@ func getConfigPath() string { return filepath.Join(home, ".picoclaw", "config.json") } -func setupCronTool(agentLoop *agent.AgentLoop, msgBus *bus.MessageBus, workspace string, restrict bool, execTimeout time.Duration) *cron.CronService { +func setupCronTool(agentLoop *agent.AgentLoop, msgBus *bus.MessageBus, workspace string, restrict bool, execTimeout time.Duration, cronOpts ...cron.CronOption) *cron.CronService { cronStorePath := filepath.Join(workspace, "cron", "jobs.json") - // Create cron service - cronService := cron.NewCronService(cronStorePath, nil) + cronService := cron.NewCronService(cronStorePath, nil, cronOpts...) // Create and register CronTool cronTool := tools.NewCronTool(cronService, agentLoop, msgBus, workspace, restrict, execTimeout) diff --git a/pkg/agent/context.go b/pkg/agent/context.go index bcbda68e5..9c8009451 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -17,11 +17,14 @@ import ( ) type ContextBuilder struct { - workspace string - skillsLoader *skills.SkillsLoader - memory *MemoryStore // Legacy file-based memory - memoryStore memory.Memory // New 3-tier MemGPT memory (may be nil) - tools *tools.ToolRegistry // Direct reference to tool registry + workspace string + skillsLoader *skills.SkillsLoader + memoryStore memory.Memory // 3-tier MemGPT memory (may be nil) + delegate memory.MemoryDelegate // Direct delegate for document loading (may be nil) + tools *tools.ToolRegistry // Direct reference to tool registry + observationBlock string // Pre-rendered observation block for prompt injection + knowledgeBlock string // Pre-rendered knowledge block from Focus completions + dagBlock string // Pre-rendered DAG compressed history } func getGlobalConfigDir() string { @@ -42,7 +45,6 @@ func NewContextBuilder(workspace string) *ContextBuilder { return &ContextBuilder{ workspace: workspace, skillsLoader: skills.NewSkillsLoader(workspace, globalSkillsDir, builtinSkillsDir), - memory: NewMemoryStore(workspace), } } @@ -56,6 +58,29 @@ func (cb *ContextBuilder) SetMemoryStore(ms memory.Memory) { cb.memoryStore = ms } +// SetDelegate sets the memory delegate for document loading. +func (cb *ContextBuilder) SetDelegate(del memory.MemoryDelegate) { + cb.delegate = del +} + +func (cb *ContextBuilder) SkillsLoader() *skills.SkillsLoader { + return cb.skillsLoader +} + +func (cb *ContextBuilder) SetObservationBlock(block string) { + cb.observationBlock = block +} + +// SetKnowledgeBlock sets the pre-rendered knowledge block from completed Focus sessions. +func (cb *ContextBuilder) SetKnowledgeBlock(block string) { + cb.knowledgeBlock = block +} + +// SetDAGBlock sets the pre-rendered DAG compressed history for prompt injection. +func (cb *ContextBuilder) SetDAGBlock(block string) { + cb.dagBlock = block +} + func (cb *ContextBuilder) getIdentity() string { now := time.Now().Format("2006-01-02 15:04 (Monday)") workspacePath, _ := filepath.Abs(filepath.Join(cb.workspace)) @@ -76,8 +101,6 @@ You are picoclaw, a helpful AI assistant. ## 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 @@ -88,8 +111,10 @@ Your workspace is at: %s 2. **Be helpful and accurate** - When using tools, briefly explain what you're doing. -3. **Memory** - When remembering something, write to %s/memory/MEMORY.md`, - now, runtime, workspacePath, workspacePath, workspacePath, workspacePath, toolsSection, workspacePath) +3. **Memory** - Use the memory tool to store important facts, preferences, and decisions. + +4. **Context Management** - You MUST consolidate your context to stay effective during long tasks. Use start_focus at the beginning of any investigation or multi-step task. After 10-15 tool calls, call complete_focus with a summary of what you learned and accomplished. This compresses your working context and persists knowledge for future reference. Failing to consolidate will degrade your performance as context grows.`, + now, runtime, workspacePath, workspacePath, toolsSection) } func (cb *ContextBuilder) buildToolsSection() string { @@ -136,10 +161,17 @@ The following skills extend your capabilities. To use a skill, read its SKILL.md %s`, skillsSummary)) } - // Legacy file-based memory context - memoryContext := cb.memory.GetMemoryContext() - if memoryContext != "" { - parts = append(parts, "# Memory\n\n"+memoryContext) + // Observation block (stable prefix for prompt cache alignment) + if cb.observationBlock != "" { + parts = append(parts, "# Observations\n\n"+cb.observationBlock) + } + + if cb.knowledgeBlock != "" { + parts = append(parts, cb.knowledgeBlock) + } + + if cb.dagBlock != "" { + parts = append(parts, "# Conversation History (Compressed)\n\n"+cb.dagBlock) } // 3-tier MemGPT working context injection @@ -155,6 +187,20 @@ The following skills extend your capabilities. To use a skill, read its SKILL.md } func (cb *ContextBuilder) LoadBootstrapFiles() string { + if cb.delegate != nil { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + docs, err := cb.delegate.ListDocumentsByCategory(ctx, "picoclaw", "bootstrap") + if err == nil && len(docs) > 0 { + var result string + for _, doc := range docs { + result += fmt.Sprintf("## %s\n\n%s\n\n", doc.Name, doc.Content) + } + return result + } + } + bootstrapFiles := []string{ "AGENTS.md", "SOUL.md", diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index c512f492e..0fe8fb055 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -23,9 +23,12 @@ import ( "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/constants" picofantasy "github.com/sipeed/picoclaw/pkg/fantasy" + "github.com/sipeed/picoclaw/pkg/ids" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/memory" + "github.com/sipeed/picoclaw/pkg/memory/dag" "github.com/sipeed/picoclaw/pkg/memory/delegate" + "github.com/sipeed/picoclaw/pkg/memory/observation" memstore "github.com/sipeed/picoclaw/pkg/memory/store" "github.com/sipeed/picoclaw/pkg/messages" "github.com/sipeed/picoclaw/pkg/session" @@ -45,7 +48,10 @@ type AgentLoop struct { state *state.Manager contextBuilder *ContextBuilder tools *tools.ToolRegistry - memoryStore *memstore.MemoryStore // 3-tier MemGPT memory (nil if init failed) + memoryStore *memstore.MemoryStore // 3-tier MemGPT memory (nil if init failed) + memDelegate memory.MemoryDelegate // DB delegate (nil if memory disabled) + obsManager *observation.Manager // Observational memory (nil if memory disabled) + activeSessionKey atomic.Value // Current session key for tool access running atomic.Bool summarizing sync.Map // Tracks which sessions are currently being summarized summarizeFailures sync.Map // Tracks consecutive summarization failures per session (string -> int) @@ -138,17 +144,19 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, model fantasy.Lang subagentTool := tools.NewSubagentTool(subagentManager) toolsRegistry.Register(subagentTool) - sessionsManager := session.NewSessionManager(filepath.Join(workspace, "sessions")) - - // Create state manager for atomic state persistence - stateManager := state.NewManager(workspace) - // Create context builder and set tools registry contextBuilder := NewContextBuilder(workspace) contextBuilder.SetToolsRegistry(toolsRegistry) + // Progressive skill disclosure tools (skill_search → skill_read → skill_traverse) + sl := contextBuilder.SkillsLoader() + toolsRegistry.Register(tools.NewSkillSearchTool(sl)) + toolsRegistry.Register(tools.NewSkillReadTool(sl)) + toolsRegistry.Register(tools.NewSkillTraverseTool(sl)) + // Initialize 3-tier MemGPT memory system var ms *memstore.MemoryStore + var memDelegate memory.MemoryDelegate if cfg.Memory.Enabled { memDBPath := filepath.Join(workspace, "memory", "picoclaw.db") os.MkdirAll(filepath.Dir(memDBPath), 0755) @@ -163,13 +171,13 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, model fantasy.Lang map[string]interface{}{"error": err.Error()}) del.Close() } else { + memDelegate = del offloadThreshold := cfg.Memory.OffloadThresholdTokens if offloadThreshold <= 0 { offloadThreshold = 4000 } chunker := memstore.NewMarkdownChunker(memstore.DefaultMarkdownChunkerConfig()) - // Create embedding provider from config (nil = FTS5-only search) embedder, embErr := memstore.NewEmbedderFromConfig(cfg.Memory.Embedding, cfg.Providers) if embErr != nil { logger.WarnCF("agent", "Failed to create embedding provider, archival search will use FTS5 only", @@ -185,10 +193,29 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, model fantasy.Lang memTool := NewMemGPTTool(ms, "picoclaw", "default") toolsRegistry.Register(memTool) - // One-time migration of file-based sessions into recall memory - sessionsDir := filepath.Join(workspace, "sessions") - if _, migErr := memory.MigrateFileSessions(context.Background(), del, "picoclaw", sessionsDir); migErr != nil { - logger.WarnCF("agent", "Session migration failed (non-fatal)", + // Agentic retrieval tools (keyword_search → semantic_search → chunk_read) + toolsRegistry.Register(tools.NewKeywordSearchTool(ms, "picoclaw")) + toolsRegistry.Register(tools.NewSemanticSearchTool(ms, "picoclaw")) + toolsRegistry.Register(tools.NewChunkReadTool(ms, "picoclaw")) + + contextBuilder.SetDelegate(del) + + // One-time migrations + mctx := context.Background() + if migErr := memory.MigrateState(mctx, workspace, del, "picoclaw"); migErr != nil { + logger.WarnCF("agent", "State KV migration failed (non-fatal)", + map[string]interface{}{"error": migErr.Error()}) + } + if migErr := memory.MigrateDocuments(mctx, workspace, del, "picoclaw"); migErr != nil { + logger.WarnCF("agent", "Document migration failed (non-fatal)", + map[string]interface{}{"error": migErr.Error()}) + } + if migErr := memory.MigrateLongTermMemory(mctx, workspace, del, "picoclaw"); migErr != nil { + logger.WarnCF("agent", "Long-term memory migration failed (non-fatal)", + map[string]interface{}{"error": migErr.Error()}) + } + if migErr := memory.MigrateDailyNotes(mctx, workspace, del, "picoclaw"); migErr != nil { + logger.WarnCF("agent", "Daily notes migration failed (non-fatal)", map[string]interface{}{"error": migErr.Error()}) } } @@ -197,6 +224,21 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, model fantasy.Lang logger.InfoCF("agent", "Memory system disabled by config", nil) } + // Create state manager -- use delegate-backed KV when memory system is active + var stateOpts []state.Option + if memDelegate != nil { + stateOpts = append(stateOpts, state.WithDelegate(memDelegate)) + } + stateManager := state.NewManager(workspace, stateOpts...) + + // Create session manager -- use delegate for DB persistence when available + sessionsDir := filepath.Join(workspace, "sessions") + var sessionOpts []session.SessionOption + if memDelegate != nil { + sessionOpts = append(sessionOpts, session.WithSessionDelegate(memDelegate, "picoclaw")) + } + sessionsManager := session.NewSessionManager(sessionsDir, sessionOpts...) + // Register meta-tools for progressive disclosure (tool_search + tool_call) toolsRegistry.RegisterMetaTools() @@ -212,7 +254,28 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, model fantasy.Lang map[string]interface{}{"gateway_tools": toolsRegistry.ListVisible()}) } - return &AgentLoop{ + // Initialize observation manager if memory is enabled + var obsManager *observation.Manager + if memDelegate != nil { + callModelFn := func(ctx context.Context, prompt string) (string, error) { + temp := 0.3 + maxTokens := int64(1024) + resp, err := model.Generate(ctx, fantasy.Call{ + Prompt: fantasy.Prompt{ + fantasy.NewUserMessage(prompt), + }, + Temperature: &temp, + MaxOutputTokens: &maxTokens, + }) + if err != nil { + return "", err + } + return resp.Content.Text(), nil + } + obsManager = observation.NewManager(memDelegate, "picoclaw", callModelFn, observation.DefaultManagerConfig()) + } + + al := &AgentLoop{ bus: msgBus, languageModel: model, workspace: workspace, @@ -224,9 +287,26 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, model fantasy.Lang contextBuilder: contextBuilder, tools: toolsRegistry, memoryStore: ms, + memDelegate: memDelegate, + obsManager: obsManager, summarizing: sync.Map{}, cfg: cfg, } + + // Register focus tools (start_focus / complete_focus) when memory is available. + // The sessionKeyFn closure reads the activeSessionKey set at the start of each agent turn. + if memDelegate != nil { + sessionKeyFn := func() string { + if v := al.activeSessionKey.Load(); v != nil { + return v.(string) + } + return "" + } + toolsRegistry.Register(tools.NewStartFocusTool(memDelegate, sessionsManager, sessionKeyFn)) + toolsRegistry.Register(tools.NewCompleteFocusTool(memDelegate, sessionsManager, sessionKeyFn)) + } + + return al } func (al *AgentLoop) Run(ctx context.Context) error { @@ -450,6 +530,8 @@ func (al *AgentLoop) processSystemMessage(ctx context.Context, msg bus.InboundMe // It handles context building, Fantasy agent creation, tool execution, and response handling. // When opts.Streaming is true, delegates to runAgentLoopStreaming for real-time token delivery. func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (string, error) { + al.activeSessionKey.Store(opts.SessionKey) + if opts.Streaming { return al.runAgentLoopStreaming(ctx, opts) } @@ -466,13 +548,29 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str // 1. Update tool contexts al.updateToolContexts(opts.Channel, opts.ChatID) - // 2. Build messages (skip history for heartbeat) + // 2. Load observation block for system prompt injection + if al.obsManager != nil { + block := al.obsManager.LoadBlock(ctx, opts.SessionKey) + al.contextBuilder.SetObservationBlock(block) + } + + // 2b. Load knowledge block from completed Focus sessions + if al.memDelegate != nil { + kb := tools.LoadKnowledgeBlock(ctx, al.memDelegate, opts.SessionKey) + al.contextBuilder.SetKnowledgeBlock(kb) + } + + // 3. Build messages with DAG compression (skip history for heartbeat) var history []messages.Message var summary string if !opts.NoHistory { history = al.sessions.GetHistory(opts.SessionKey) summary = al.sessions.GetSummary(opts.SessionKey) } + + // 3a. DAG compression: compress old history, keep raw tail + history = al.applyDAGCompression(history) + builtMsgs := al.contextBuilder.BuildMessages( history, summary, @@ -482,7 +580,7 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str opts.ChatID, ) - // 3. Save user message to session + // 3b. Save user message to session al.sessions.AddMessage(opts.SessionKey, "user", opts.UserMessage) // 4. Split built messages into system prompt, conversation history, and current user prompt. @@ -542,13 +640,14 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str return "", fmt.Errorf("agent Generate failed: %w", err) } - // 9. Save all step messages to session + // 9. Save all step messages to session and audit tool calls stepCount := len(result.Steps) for _, step := range result.Steps { stepMsgs := picofantasy.StepToMessages(step) for _, m := range stepMsgs { al.sessions.AddFullMessage(opts.SessionKey, m) } + al.auditStep(ctx, step, opts.SessionKey) } // 10. Extract final text @@ -567,6 +666,12 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str al.maybeSummarize(opts.SessionKey, opts.Channel, opts.ChatID) } + // 13b. Trigger async observation if tail exceeds token threshold + if al.obsManager != nil { + tail := al.sessionsToMessagePairs(opts.SessionKey) + al.obsManager.MaybeObserveAsync(ctx, opts.SessionKey, tail) + } + // 14. Optional: send response via bus if opts.SendResponse { al.bus.PublishOutbound(bus.OutboundMessage{ @@ -605,19 +710,35 @@ func (al *AgentLoop) runAgentLoopStreaming(ctx context.Context, opts processOpti // 1. Update tool contexts al.updateToolContexts(opts.Channel, opts.ChatID) - // 2. Build messages + // 2. Load observation block for system prompt injection + if al.obsManager != nil { + block := al.obsManager.LoadBlock(ctx, opts.SessionKey) + al.contextBuilder.SetObservationBlock(block) + } + + // 2b. Load knowledge block from completed Focus sessions + if al.memDelegate != nil { + kb := tools.LoadKnowledgeBlock(ctx, al.memDelegate, opts.SessionKey) + al.contextBuilder.SetKnowledgeBlock(kb) + } + + // 3. Build messages with DAG compression var history []messages.Message var summary string if !opts.NoHistory { history = al.sessions.GetHistory(opts.SessionKey) summary = al.sessions.GetSummary(opts.SessionKey) } + + // 3a. DAG compression + history = al.applyDAGCompression(history) + builtMsgs := al.contextBuilder.BuildMessages(history, summary, opts.UserMessage, nil, opts.Channel, opts.ChatID) - // 3. Save user message + // 4. Save user message al.sessions.AddMessage(opts.SessionKey, "user", opts.UserMessage) - // 4. Split into system/history/user + // 5. Split into system/history/user systemPrompt := "" var historyMsgs []messages.Message userPrompt := opts.UserMessage @@ -677,12 +798,12 @@ func (al *AgentLoop) runAgentLoopStreaming(ctx context.Context, opts processOpti return nil }, - // Save each step's messages to session as they complete OnStepFinish: func(step fantasy.StepResult) error { stepMsgs := picofantasy.StepToMessages(step) for _, m := range stepMsgs { al.sessions.AddFullMessage(opts.SessionKey, m) } + al.auditStep(ctx, step, opts.SessionKey) return nil }, @@ -719,6 +840,12 @@ func (al *AgentLoop) runAgentLoopStreaming(ctx context.Context, opts processOpti al.maybeSummarize(opts.SessionKey, opts.Channel, opts.ChatID) } + // 12b. Trigger async observation if tail exceeds token threshold + if al.obsManager != nil { + tail := al.sessionsToMessagePairs(opts.SessionKey) + al.obsManager.MaybeObserveAsync(ctx, opts.SessionKey, tail) + } + // 13. Log response stepCount := len(result.Steps) responsePreview := utils.Truncate(finalContent, 120) @@ -735,6 +862,35 @@ func (al *AgentLoop) runAgentLoopStreaming(ctx context.Context, opts processOpti // runLLMIteration — DELETED. Replaced by Fantasy's internal agent loop. +// auditStep logs tool calls from a Fantasy step result to the audit log. +func (al *AgentLoop) auditStep(ctx context.Context, step fantasy.StepResult, sessionKey string) { + if al.memDelegate == nil { + return + } + + toolCalls := step.Content.ToolCalls() + if len(toolCalls) == 0 { + return + } + + for _, tc := range toolCalls { + entry := &memory.AuditEntry{ + ID: ids.New(), + AgentID: "picoclaw", + SessionKey: sessionKey, + Action: "tool_call", + Target: tc.ToolName, + Input: tc.Input, + } + aCtx, cancel := context.WithTimeout(ctx, time.Second) + if err := al.memDelegate.InsertAuditEntry(aCtx, entry); err != nil { + logger.WarnCF("agent", "Failed to log audit entry", + map[string]interface{}{"tool": tc.ToolName, "error": err.Error()}) + } + cancel() + } +} + // updateToolContexts updates the context for tools that need channel/chatID info. func (al *AgentLoop) updateToolContexts(channel, chatID string) { // Use ContextualTool interface instead of type assertions @@ -839,6 +995,11 @@ func (al *AgentLoop) forceCompression(sessionKey string) { }) } +// MemoryDelegate returns the active memory delegate (nil if memory system is disabled). +func (al *AgentLoop) MemoryDelegate() memory.MemoryDelegate { + return al.memDelegate +} + // GetStartupInfo returns information about loaded tools and skills for logging. func (al *AgentLoop) GetStartupInfo() map[string]interface{} { info := make(map[string]interface{}) @@ -1013,6 +1174,72 @@ func (al *AgentLoop) callModel(ctx context.Context, prompt string) (string, erro return resp.Content.Text(), nil } +// sessionsToMessagePairs converts the session history to observation.MessagePair +// for token estimation by the observation manager. +func (al *AgentLoop) sessionsToMessagePairs(sessionKey string) []observation.MessagePair { + history := al.sessions.GetHistory(sessionKey) + pairs := make([]observation.MessagePair, len(history)) + for i, m := range history { + pairs[i] = observation.MessagePair{Role: m.Role, Content: m.Content} + } + return pairs +} + +// applyDAGCompression compresses old history into a DAG summary block and +// returns only the tail messages that should be passed as raw conversation. +// The compressed portion is injected into the system prompt via contextBuilder. +func (al *AgentLoop) applyDAGCompression(history []messages.Message) []messages.Message { + const minHistoryForDAG = 16 + + if len(history) < minHistoryForDAG { + al.contextBuilder.SetDAGBlock("") + return history + } + + budget := dag.ComputeBudget(al.contextWindow, dag.DefaultBudgetConfig()) + tailCount := dag.TailMessageCount(budget.RawTail) + if tailCount >= len(history) { + al.contextBuilder.SetDAGBlock("") + return history + } + + // Split: compress old, keep tail raw + compressible := history[:len(history)-tailCount] + tail := history[len(history)-tailCount:] + + // Tool-call-aware: don't split on a "tool" message + for len(tail) > 0 && tail[0].Role == "tool" && len(compressible) > 0 { + tail = append([]messages.Message{compressible[len(compressible)-1]}, tail...) + compressible = compressible[:len(compressible)-1] + } + + if len(compressible) == 0 { + al.contextBuilder.SetDAGBlock("") + return history + } + + dagMsgs := make([]dag.Message, len(compressible)) + for i, m := range compressible { + dagMsgs[i] = dag.Message{Role: m.Role, Content: m.Content} + } + + compressor := dag.NewCompressor(dag.DefaultCompressorConfig()) + d := compressor.Compress(dagMsgs) + + rendered := dag.RenderDAGForBudget(d, budget.DAGSummaries) + al.contextBuilder.SetDAGBlock(rendered) + + logger.DebugCF("agent", "DAG compression applied", + map[string]interface{}{ + "total_msgs": len(history), + "compressed_msgs": len(compressible), + "tail_msgs": len(tail), + "dag_nodes": len(d.Nodes), + }) + + return tail +} + // estimateTokens estimates the number of tokens in a message list. func (al *AgentLoop) estimateTokens(msgs []messages.Message) int { totalChars := 0 diff --git a/pkg/agent/memory.go b/pkg/agent/memory.go deleted file mode 100644 index 3f6896f91..000000000 --- a/pkg/agent/memory.go +++ /dev/null @@ -1,161 +0,0 @@ -// PicoClaw - Ultra-lightweight personal AI agent -// Inspired by and based on nanobot: https://github.com/HKUDS/nanobot -// License: MIT -// -// Copyright (c) 2026 PicoClaw contributors - -package agent - -import ( - "fmt" - "os" - "path/filepath" - "time" -) - -// MemoryStore manages persistent memory for the agent. -// - Long-term memory: memory/MEMORY.md -// - Daily notes: memory/YYYYMM/YYYYMMDD.md -type MemoryStore struct { - workspace string - memoryDir string - memoryFile string -} - -// NewMemoryStore creates a new MemoryStore with the given workspace path. -// It ensures the memory directory exists. -func NewMemoryStore(workspace string) *MemoryStore { - memoryDir := filepath.Join(workspace, "memory") - memoryFile := filepath.Join(memoryDir, "MEMORY.md") - - // Ensure memory directory exists - os.MkdirAll(memoryDir, 0755) - - return &MemoryStore{ - workspace: workspace, - memoryDir: memoryDir, - memoryFile: memoryFile, - } -} - -// getTodayFile returns the path to today's daily note file (memory/YYYYMM/YYYYMMDD.md). -func (ms *MemoryStore) getTodayFile() string { - today := time.Now().Format("20060102") // YYYYMMDD - monthDir := today[:6] // YYYYMM - filePath := filepath.Join(ms.memoryDir, monthDir, today+".md") - return filePath -} - -// ReadLongTerm reads the long-term memory (MEMORY.md). -// Returns empty string if the file doesn't exist. -func (ms *MemoryStore) ReadLongTerm() string { - if data, err := os.ReadFile(ms.memoryFile); err == nil { - return string(data) - } - return "" -} - -// 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), 0644) -} - -// ReadToday reads today's daily note. -// Returns empty string if the file doesn't exist. -func (ms *MemoryStore) ReadToday() string { - todayFile := ms.getTodayFile() - if data, err := os.ReadFile(todayFile); err == nil { - return string(data) - } - return "" -} - -// AppendToday appends content to today's daily note. -// If the file doesn't exist, it creates a new file with a date header. -func (ms *MemoryStore) AppendToday(content string) error { - todayFile := ms.getTodayFile() - - // Ensure month directory exists - monthDir := filepath.Dir(todayFile) - os.MkdirAll(monthDir, 0755) - - var existingContent string - if data, err := os.ReadFile(todayFile); err == nil { - existingContent = string(data) - } - - var newContent string - if existingContent == "" { - // Add header for new day - header := fmt.Sprintf("# %s\n\n", time.Now().Format("2006-01-02")) - newContent = header + content - } else { - // Append to existing content - newContent = existingContent + "\n" + content - } - - return os.WriteFile(todayFile, []byte(newContent), 0644) -} - -// GetRecentDailyNotes returns daily notes from the last N days. -// Contents are joined with "---" separator. -func (ms *MemoryStore) GetRecentDailyNotes(days int) string { - var notes []string - - for i := 0; i < days; i++ { - date := time.Now().AddDate(0, 0, -i) - dateStr := date.Format("20060102") // YYYYMMDD - monthDir := dateStr[:6] // YYYYMM - filePath := filepath.Join(ms.memoryDir, monthDir, dateStr+".md") - - if data, err := os.ReadFile(filePath); err == nil { - notes = append(notes, string(data)) - } - } - - if len(notes) == 0 { - return "" - } - - // Join with separator - var result string - for i, note := range notes { - if i > 0 { - result += "\n\n---\n\n" - } - result += note - } - return result -} - -// GetMemoryContext returns formatted memory context for the agent prompt. -// Includes long-term memory and recent daily notes. -func (ms *MemoryStore) GetMemoryContext() string { - var parts []string - - // Long-term memory - longTerm := ms.ReadLongTerm() - if longTerm != "" { - parts = append(parts, "## Long-term Memory\n\n"+longTerm) - } - - // Recent daily notes (last 3 days) - recentNotes := ms.GetRecentDailyNotes(3) - if recentNotes != "" { - parts = append(parts, "## Recent Daily Notes\n\n"+recentNotes) - } - - if len(parts) == 0 { - return "" - } - - // Join parts with separator - var result string - for i, part := range parts { - if i > 0 { - result += "\n\n---\n\n" - } - result += part - } - return fmt.Sprintf("# Memory\n\n%s", result) -} diff --git a/pkg/cron/service.go b/pkg/cron/service.go index 9f62c743b..0e425de2a 100644 --- a/pkg/cron/service.go +++ b/pkg/cron/service.go @@ -1,6 +1,7 @@ package cron import ( + "context" "crypto/rand" "encoding/hex" "encoding/json" @@ -12,6 +13,7 @@ import ( "time" "github.com/adhocore/gronx" + "github.com/sipeed/picoclaw/pkg/memory" ) type CronSchedule struct { @@ -57,6 +59,19 @@ type CronStore struct { type JobHandler func(job *CronJob) (string, error) +// CronOption configures a CronService. +type CronOption func(*CronService) + +// WithCronDelegate injects a memory delegate for KV-backed cron store persistence. +func WithCronDelegate(del memory.MemoryDelegate, agentID string) CronOption { + return func(cs *CronService) { + cs.delegate = del + cs.agentID = agentID + } +} + +const cronStoreKVKey = "cron:store" + type CronService struct { storePath string store *CronStore @@ -65,15 +80,19 @@ type CronService struct { running bool stopChan chan struct{} gronx *gronx.Gronx + delegate memory.MemoryDelegate + agentID string } -func NewCronService(storePath string, onJob JobHandler) *CronService { +func NewCronService(storePath string, onJob JobHandler, opts ...CronOption) *CronService { cs := &CronService{ storePath: storePath, onJob: onJob, gronx: gronx.New(), } - // Initialize and load store on creation + for _, opt := range opts { + opt(cs) + } cs.loadStore() return cs } @@ -318,6 +337,10 @@ func (cs *CronService) loadStore() error { Jobs: []CronJob{}, } + if cs.delegate != nil { + return cs.loadStoreFromDelegate() + } + data, err := os.ReadFile(cs.storePath) if err != nil { if os.IsNotExist(err) { @@ -329,7 +352,22 @@ func (cs *CronService) loadStore() error { return json.Unmarshal(data, cs.store) } +func (cs *CronService) loadStoreFromDelegate() error { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + val, err := cs.delegate.GetKV(ctx, cs.agentID, cronStoreKVKey) + if err != nil || val == "" { + return nil + } + return json.Unmarshal([]byte(val), cs.store) +} + func (cs *CronService) saveStoreUnsafe() error { + if cs.delegate != nil { + return cs.saveStoreToDelegate() + } + dir := filepath.Dir(cs.storePath) if err := os.MkdirAll(dir, 0755); err != nil { return err @@ -343,6 +381,16 @@ func (cs *CronService) saveStoreUnsafe() error { return os.WriteFile(cs.storePath, data, 0600) } +func (cs *CronService) saveStoreToDelegate() error { + data, err := json.Marshal(cs.store) + if err != nil { + return err + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + return cs.delegate.UpsertKV(ctx, cs.agentID, cronStoreKVKey, string(data)) +} + func (cs *CronService) AddJob(name string, schedule CronSchedule, message string, deliver bool, channel, to string) (*CronJob, error) { cs.mu.Lock() defer cs.mu.Unlock() diff --git a/pkg/heartbeat/service.go b/pkg/heartbeat/service.go index 45d970045..b5e8a132f 100644 --- a/pkg/heartbeat/service.go +++ b/pkg/heartbeat/service.go @@ -43,9 +43,9 @@ type HeartbeatService struct { stopChan chan struct{} } -// NewHeartbeatService creates a new heartbeat service -func NewHeartbeatService(workspace string, intervalMinutes int, enabled bool) *HeartbeatService { - // Apply minimum interval +// NewHeartbeatService creates a new heartbeat service. +// stateOpts are forwarded to the internal state.Manager. +func NewHeartbeatService(workspace string, intervalMinutes int, enabled bool, stateOpts ...state.Option) *HeartbeatService { if intervalMinutes < minIntervalMinutes && intervalMinutes != 0 { intervalMinutes = minIntervalMinutes } @@ -58,7 +58,7 @@ func NewHeartbeatService(workspace string, intervalMinutes int, enabled bool) *H workspace: workspace, interval: time.Duration(intervalMinutes) * time.Minute, enabled: enabled, - state: state.NewManager(workspace), + state: state.NewManager(workspace, stateOpts...), } } diff --git a/pkg/session/manager.go b/pkg/session/manager.go index b0bf19318..65f621dfa 100644 --- a/pkg/session/manager.go +++ b/pkg/session/manager.go @@ -1,6 +1,7 @@ package session import ( + "context" "encoding/json" "os" "path/filepath" @@ -9,7 +10,9 @@ import ( "time" "github.com/sipeed/picoclaw/pkg/cache" + "github.com/sipeed/picoclaw/pkg/ids" "github.com/sipeed/picoclaw/pkg/logger" + "github.com/sipeed/picoclaw/pkg/memory" "github.com/sipeed/picoclaw/pkg/messages" ) @@ -33,26 +36,44 @@ type SessionManagerConfig struct { SessionTTL time.Duration } +// SessionOption configures a SessionManager. +type SessionOption func(*SessionManager) + +// WithSessionDelegate injects a memory delegate for DB-backed session persistence. +// When set, sessions persist through recall_items instead of JSON files. +func WithSessionDelegate(del memory.MemoryDelegate, agentID string) SessionOption { + return func(sm *SessionManager) { + sm.delegate = del + sm.agentID = agentID + } +} + type SessionManager struct { sessions map[string]*Session // primary store (always authoritative) lru *cache.LRU[string, bool] // tracks access order; value is just a presence flag mu sync.RWMutex storage string cfg SessionManagerConfig + delegate memory.MemoryDelegate + agentID string } -func NewSessionManager(storage string) *SessionManager { - return NewSessionManagerWithConfig(storage, SessionManagerConfig{}) +func NewSessionManager(storage string, opts ...SessionOption) *SessionManager { + return NewSessionManagerWithConfig(storage, SessionManagerConfig{}, opts...) } // NewSessionManagerWithConfig creates a SessionManager with LRU cache settings. -func NewSessionManagerWithConfig(storage string, cfg SessionManagerConfig) *SessionManager { +func NewSessionManagerWithConfig(storage string, cfg SessionManagerConfig, opts ...SessionOption) *SessionManager { sm := &SessionManager{ sessions: make(map[string]*Session), storage: storage, cfg: cfg, } + for _, opt := range opts { + opt(sm) + } + if cfg.MaxCachedSessions > 0 { sm.lru = cache.New(cache.Options[string, bool]{ MaxSize: cfg.MaxCachedSessions, @@ -63,7 +84,9 @@ func NewSessionManagerWithConfig(storage string, cfg SessionManagerConfig) *Sess }) } - if storage != "" { + if sm.delegate != nil { + sm.loadSessionsFromDelegate() + } else if storage != "" { os.MkdirAll(storage, 0755) sm.loadSessions() } @@ -71,6 +94,41 @@ func NewSessionManagerWithConfig(storage string, cfg SessionManagerConfig) *Sess return sm } +func (sm *SessionManager) loadSessionsFromDelegate() { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + items, err := sm.delegate.ListRecallItems(ctx, sm.agentID, "", 500, 0) + if err != nil { + logger.WarnCF("session", "Failed to load sessions from delegate", + map[string]interface{}{"error": err.Error()}) + return + } + + for _, item := range items { + if !strings.Contains(item.Tags, "session-message") { + continue + } + session := sm.sessions[item.SessionKey] + if session == nil { + session = &Session{ + Key: item.SessionKey, + Messages: []messages.Message{}, + Created: item.CreatedAt, + Updated: item.CreatedAt, + } + sm.sessions[item.SessionKey] = session + } + session.Messages = append(session.Messages, messages.Message{ + Role: item.Role, + Content: item.Content, + }) + if item.CreatedAt.After(session.Updated) { + session.Updated = item.CreatedAt + } + } +} + // touchLRU records an access in the LRU tracker, which may evict cold sessions. func (sm *SessionManager) touchLRU(key string) { if sm.lru != nil { @@ -181,6 +239,10 @@ func (sm *SessionManager) AddFullMessage(sessionKey string, msg messages.Message session.Updated = time.Now() sm.touchLRU(sessionKey) + if sm.delegate != nil { + sm.persistMessageToDelegate(sessionKey, msg) + } + // Hard cap: prevent unbounded growth if summarization keeps failing. // Keep last 50 messages when we exceed 200. const hardCap = 200 @@ -196,6 +258,31 @@ func (sm *SessionManager) AddFullMessage(sessionKey string, msg messages.Message } } +func (sm *SessionManager) persistMessageToDelegate(sessionKey string, msg messages.Message) { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + now := time.Now() + item := &memory.RecallItem{ + ID: ids.New(), + AgentID: sm.agentID, + SessionKey: sessionKey, + Role: msg.Role, + Sector: memory.SectorEpisodic, + Importance: 0.5, + Salience: 0.5, + DecayRate: 0.01, + Content: msg.Content, + Tags: "session-message", + CreatedAt: now, + UpdatedAt: now, + } + if err := sm.delegate.InsertRecallItem(ctx, item); err != nil { + logger.WarnCF("session", "Failed to persist message to delegate", + map[string]interface{}{"session": sessionKey, "error": err.Error()}) + } +} + func (sm *SessionManager) GetHistory(key string) []messages.Message { sm.mu.RLock() session, ok := sm.sessions[key] @@ -340,6 +427,9 @@ func sanitizeFilename(key string) string { return strings.ReplaceAll(key, ":", "_") } func (sm *SessionManager) Save(key string) error { + if sm.delegate != nil { + return nil + } if sm.storage == "" { return nil } diff --git a/pkg/session/manager_test.go b/pkg/session/manager_test.go index 5338d0e4b..d5c3e5c09 100644 --- a/pkg/session/manager_test.go +++ b/pkg/session/manager_test.go @@ -1,11 +1,13 @@ package session import ( + "context" "os" "path/filepath" "testing" "time" + "github.com/sipeed/picoclaw/pkg/memory/delegate" "github.com/sipeed/picoclaw/pkg/messages" ) @@ -189,3 +191,64 @@ func TestCleanupStale(t *testing.T) { t.Errorf("expected active session to remain") } } + +func TestSessionManager_DelegatePersistence(t *testing.T) { + del, err := delegate.NewLibSQLInMemory() + if err != nil { + t.Fatalf("NewLibSQLInMemory: %v", err) + } + if err := del.Init(context.Background()); err != nil { + t.Fatalf("Init: %v", err) + } + defer del.Close() + + sm := NewSessionManager("", WithSessionDelegate(del, "test-agent")) + + key := "delegate-session" + sm.AddMessage(key, "user", "hello from delegate") + sm.AddMessage(key, "assistant", "hi back") + + history := sm.GetHistory(key) + if len(history) != 2 { + t.Fatalf("expected 2 messages in-memory, got %d", len(history)) + } + + items, err := del.ListRecallItems(context.Background(), "test-agent", key, 100, 0) + if err != nil { + t.Fatalf("ListRecallItems: %v", err) + } + if len(items) != 2 { + t.Fatalf("expected 2 recall items in DB, got %d", len(items)) + } + if items[0].Content != "hello from delegate" { + t.Errorf("expected first item content 'hello from delegate', got %q", items[0].Content) + } +} + +func TestSessionManager_DelegateSaveIsNoop(t *testing.T) { + del, err := delegate.NewLibSQLInMemory() + if err != nil { + t.Fatalf("NewLibSQLInMemory: %v", err) + } + if err := del.Init(context.Background()); err != nil { + t.Fatalf("Init: %v", err) + } + defer del.Close() + + tmpDir := t.TempDir() + sm := NewSessionManager(tmpDir, WithSessionDelegate(del, "test-agent")) + + key := "telegram:999" + sm.AddMessage(key, "user", "test") + + if err := sm.Save(key); err != nil { + t.Fatalf("Save: %v", err) + } + + entries, _ := os.ReadDir(tmpDir) + for _, e := range entries { + if filepath.Ext(e.Name()) == ".json" { + t.Errorf("delegate mode should not write JSON files, found %s", e.Name()) + } + } +} diff --git a/pkg/state/state.go b/pkg/state/state.go index 0bb9cd497..061516208 100644 --- a/pkg/state/state.go +++ b/pkg/state/state.go @@ -1,6 +1,7 @@ package state import ( + "context" "encoding/json" "fmt" "log" @@ -8,8 +9,12 @@ import ( "path/filepath" "sync" "time" + + "github.com/sipeed/picoclaw/pkg/memory" ) +const kvAgentID = "picoclaw" + // State represents the persistent state for a workspace. // It includes information about the last active channel/chat. type State struct { @@ -23,64 +28,87 @@ type State struct { Timestamp time.Time `json:"timestamp"` } +// Option configures a Manager. +type Option func(*Manager) + +// WithDelegate injects a memory delegate for KV-backed persistence. +// When set, state is stored in the agent_kv table instead of on disk. +func WithDelegate(del memory.MemoryDelegate) Option { + return func(m *Manager) { m.delegate = del } +} + // Manager manages persistent state with atomic saves. +// When a delegate is present, state persists through agent_kv. +// Otherwise, it falls back to file-based atomic JSON writes. type Manager struct { workspace string state *State mu sync.RWMutex stateFile string + delegate memory.MemoryDelegate } // NewManager creates a new state manager for the given workspace. -func NewManager(workspace string) *Manager { +func NewManager(workspace string, opts ...Option) *Manager { + sm := &Manager{ + workspace: workspace, + state: &State{}, + } + + for _, opt := range opts { + opt(sm) + } + + if sm.delegate != nil { + sm.loadFromDelegate() + return sm + } + stateDir := filepath.Join(workspace, "state") stateFile := filepath.Join(stateDir, "state.json") oldStateFile := filepath.Join(workspace, "state.json") - // Create state directory if it doesn't exist os.MkdirAll(stateDir, 0755) + sm.stateFile = stateFile - sm := &Manager{ - workspace: workspace, - stateFile: stateFile, - state: &State{}, - } - - // Try to load from new location first if _, err := os.Stat(stateFile); os.IsNotExist(err) { - // New file doesn't exist, try migrating from old location if data, err := os.ReadFile(oldStateFile); err == nil { if err := json.Unmarshal(data, sm.state); err == nil { - // Migrate to new location sm.saveAtomic() log.Printf("[INFO] state: migrated state from %s to %s", oldStateFile, stateFile) } } } else { - // Load from new location sm.load() } return sm } +func (sm *Manager) loadFromDelegate() { + ctx := context.Background() + if v, err := sm.delegate.GetKV(ctx, kvAgentID, "state:last_channel"); err == nil && v != "" { + sm.state.LastChannel = v + } + if v, err := sm.delegate.GetKV(ctx, kvAgentID, "state:last_chat_id"); err == nil && v != "" { + sm.state.LastChatID = v + } + if v, err := sm.delegate.GetKV(ctx, kvAgentID, "state:timestamp"); err == nil && v != "" { + if t, err := time.Parse(time.RFC3339Nano, v); err == nil { + sm.state.Timestamp = t + } + } +} + // SetLastChannel atomically updates the last channel and saves the state. -// This method uses a temp file + rename pattern for atomic writes, -// ensuring that the state file is never corrupted even if the process crashes. func (sm *Manager) SetLastChannel(channel string) error { sm.mu.Lock() defer sm.mu.Unlock() - // Update state sm.state.LastChannel = channel sm.state.Timestamp = time.Now() - // Atomic save using temp file + rename - if err := sm.saveAtomic(); err != nil { - return fmt.Errorf("failed to save state atomically: %w", err) - } - - return nil + return sm.persist() } // SetLastChatID atomically updates the last chat ID and saves the state. @@ -88,15 +116,34 @@ func (sm *Manager) SetLastChatID(chatID string) error { sm.mu.Lock() defer sm.mu.Unlock() - // Update state sm.state.LastChatID = chatID sm.state.Timestamp = time.Now() - // Atomic save using temp file + rename - if err := sm.saveAtomic(); err != nil { - return fmt.Errorf("failed to save state atomically: %w", err) - } + return sm.persist() +} +// persist writes the current state to the delegate (KV) or file. +// Must be called with the lock held. +func (sm *Manager) persist() error { + if sm.delegate != nil { + return sm.persistToDelegate() + } + return sm.saveAtomic() +} + +func (sm *Manager) persistToDelegate() error { + ctx := context.Background() + ts := sm.state.Timestamp.Format(time.RFC3339Nano) + + if err := sm.delegate.UpsertKV(ctx, kvAgentID, "state:last_channel", sm.state.LastChannel); err != nil { + return fmt.Errorf("upsert last_channel: %w", err) + } + if err := sm.delegate.UpsertKV(ctx, kvAgentID, "state:last_chat_id", sm.state.LastChatID); err != nil { + return fmt.Errorf("upsert last_chat_id: %w", err) + } + if err := sm.delegate.UpsertKV(ctx, kvAgentID, "state:timestamp", ts); err != nil { + return fmt.Errorf("upsert timestamp: %w", err) + } return nil }