feat(agent): integrate observation, DAG compression, and new tools
Wire the new memory subsystems into the agent loop: - Observational memory monitors interactions passively - DAG compressor manages context window budget - Skill, retrieval, and focus tools registered at startup - State manager supports KV-backed persistence via delegate - Session manager gains observation hooks and DAG integration - Cron service manages periodic observation reflection - Remove legacy pkg/agent/memory.go (superseded by delegate)
This commit is contained in:
ZanzyTHEbar
parent
6da3c30866
commit
09678e50cf
9 changed files with 597 additions and 234 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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...),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue