From de9d4703b996da0e165b29144040078e67541fa2 Mon Sep 17 00:00:00 2001 From: xiaoen <2768753269@qq.com> Date: Tue, 10 Mar 2026 15:02:11 +0800 Subject: [PATCH] fix(session): close session stores on shutdown and harden migration - Add Close() to AgentInstance, AgentRegistry, and AgentLoop so JSONL file handles are released during gateway shutdown and CLI exit. - Fall back to SessionManager when migration fails, preventing a split state where some sessions live in JSONL and others remain in JSON. - Add defer agentLoop.Close() in the CLI agent command path. - Document SessionStore interface methods (fire-and-forget contract). --- cmd/picoclaw/internal/agent/helpers.go | 1 + cmd/picoclaw/internal/gateway/helpers.go | 1 + pkg/agent/instance.go | 18 ++++++++++++++++-- pkg/agent/loop.go | 5 +++++ pkg/agent/registry.go | 12 ++++++++++++ pkg/session/session_store.go | 13 +++++++++++++ 6 files changed, 48 insertions(+), 2 deletions(-) diff --git a/cmd/picoclaw/internal/agent/helpers.go b/cmd/picoclaw/internal/agent/helpers.go index f754abc65..a995945d2 100644 --- a/cmd/picoclaw/internal/agent/helpers.go +++ b/cmd/picoclaw/internal/agent/helpers.go @@ -50,6 +50,7 @@ func agentCmd(message, sessionKey, model string, debug bool) error { msgBus := bus.NewMessageBus() defer msgBus.Close() agentLoop := agent.NewAgentLoop(cfg, msgBus, provider) + defer agentLoop.Close() // Print agent startup info (only for interactive mode) startupInfo := agentLoop.GetStartupInfo() diff --git a/cmd/picoclaw/internal/gateway/helpers.go b/cmd/picoclaw/internal/gateway/helpers.go index 4f93b858a..fed3d5ffb 100644 --- a/cmd/picoclaw/internal/gateway/helpers.go +++ b/cmd/picoclaw/internal/gateway/helpers.go @@ -214,6 +214,7 @@ func gatewayCmd(debug bool) error { cronService.Stop() mediaStore.Stop() agentLoop.Stop() + agentLoop.Close() fmt.Println("✓ Gateway stopped") return nil diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 82c183969..709b79203 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -277,9 +277,18 @@ func compilePatterns(patterns []string) []*regexp.Regexp { return compiled } +// Close releases resources held by the agent's session store. +func (a *AgentInstance) Close() error { + if a.Sessions != nil { + return a.Sessions.Close() + } + return nil +} + // initSessionStore creates the session persistence backend. // It uses the JSONL store by default and auto-migrates legacy JSON sessions. -// Falls back to SessionManager if the JSONL store cannot be initialized. +// Falls back to SessionManager if the JSONL store cannot be initialized or +// if migration fails (which indicates the store cannot write reliably). func initSessionStore(dir string) session.SessionStore { store, err := memory.NewJSONLStore(dir) if err != nil { @@ -288,7 +297,12 @@ func initSessionStore(dir string) session.SessionStore { } if n, merr := memory.MigrateFromJSON(context.Background(), dir, store); merr != nil { - log.Printf("memory: migration: %v", merr) + // Migration failure means the store could not write data. + // Fall back to SessionManager to avoid a split state where + // some sessions are in JSONL and others remain in JSON. + log.Printf("memory: migration failed: %v; falling back to json sessions", merr) + store.Close() + return session.NewSessionManager(dir) } else if n > 0 { log.Printf("memory: migrated %d session(s) to jsonl", n) } diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 3d13071c0..fb7806f2c 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -380,6 +380,11 @@ func (al *AgentLoop) Stop() { al.running.Store(false) } +// Close releases resources held by agent session stores. Call after Stop. +func (al *AgentLoop) Close() { + al.registry.Close() +} + func (al *AgentLoop) RegisterTool(tool tools.Tool) { for _, agentID := range al.registry.ListAgentIDs() { if agent, ok := al.registry.GetAgent(agentID); ok { diff --git a/pkg/agent/registry.go b/pkg/agent/registry.go index 0e7973dc3..58b7ce440 100644 --- a/pkg/agent/registry.go +++ b/pkg/agent/registry.go @@ -114,6 +114,18 @@ func (r *AgentRegistry) ForEachTool(name string, fn func(tools.Tool)) { } } +// Close releases resources held by all registered agents. +func (r *AgentRegistry) Close() { + r.mu.RLock() + defer r.mu.RUnlock() + for _, agent := range r.agents { + if err := agent.Close(); err != nil { + logger.WarnCF("agent", "Failed to close agent", + map[string]any{"agent_id": agent.ID, "error": err.Error()}) + } + } +} + // GetDefaultAgent returns the default agent instance. func (r *AgentRegistry) GetDefaultAgent() *AgentInstance { r.mu.RLock() diff --git a/pkg/session/session_store.go b/pkg/session/session_store.go index 177e58efe..1d1a2f967 100644 --- a/pkg/session/session_store.go +++ b/pkg/session/session_store.go @@ -6,14 +6,27 @@ import "github.com/sipeed/picoclaw/pkg/providers" // Both SessionManager (legacy JSON backend) and JSONLBackend satisfy this // interface, allowing the storage layer to be swapped without touching the // agent loop code. +// +// Write methods (Add*, Set*, Truncate*) are fire-and-forget: they do not +// return errors. Implementations should log failures internally. This +// matches the original SessionManager contract that the agent loop relies on. type SessionStore interface { + // AddMessage appends a simple role/content message to the session. AddMessage(sessionKey, role, content string) + // AddFullMessage appends a complete message including tool calls. AddFullMessage(sessionKey string, msg providers.Message) + // GetHistory returns the full message history for the session. GetHistory(key string) []providers.Message + // GetSummary returns the conversation summary, or "" if none. GetSummary(key string) string + // SetSummary replaces the conversation summary. SetSummary(key, summary string) + // SetHistory replaces the full message history. SetHistory(key string, history []providers.Message) + // TruncateHistory keeps only the last keepLast messages. TruncateHistory(key string, keepLast int) + // Save persists any pending state to durable storage. Save(key string) error + // Close releases resources held by the store. Close() error }