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).
This commit is contained in:
xiaoen 2026-03-10 15:02:11 +08:00
parent 841e1d4275
commit de9d4703b9
6 changed files with 48 additions and 2 deletions

View file

@ -50,6 +50,7 @@ func agentCmd(message, sessionKey, model string, debug bool) error {
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
defer msgBus.Close() defer msgBus.Close()
agentLoop := agent.NewAgentLoop(cfg, msgBus, provider) agentLoop := agent.NewAgentLoop(cfg, msgBus, provider)
defer agentLoop.Close()
// Print agent startup info (only for interactive mode) // Print agent startup info (only for interactive mode)
startupInfo := agentLoop.GetStartupInfo() startupInfo := agentLoop.GetStartupInfo()

View file

@ -214,6 +214,7 @@ func gatewayCmd(debug bool) error {
cronService.Stop() cronService.Stop()
mediaStore.Stop() mediaStore.Stop()
agentLoop.Stop() agentLoop.Stop()
agentLoop.Close()
fmt.Println("✓ Gateway stopped") fmt.Println("✓ Gateway stopped")
return nil return nil

View file

@ -277,9 +277,18 @@ func compilePatterns(patterns []string) []*regexp.Regexp {
return compiled 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. // initSessionStore creates the session persistence backend.
// It uses the JSONL store by default and auto-migrates legacy JSON sessions. // 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 { func initSessionStore(dir string) session.SessionStore {
store, err := memory.NewJSONLStore(dir) store, err := memory.NewJSONLStore(dir)
if err != nil { if err != nil {
@ -288,7 +297,12 @@ func initSessionStore(dir string) session.SessionStore {
} }
if n, merr := memory.MigrateFromJSON(context.Background(), dir, store); merr != nil { 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 { } else if n > 0 {
log.Printf("memory: migrated %d session(s) to jsonl", n) log.Printf("memory: migrated %d session(s) to jsonl", n)
} }

View file

@ -380,6 +380,11 @@ func (al *AgentLoop) Stop() {
al.running.Store(false) 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) { func (al *AgentLoop) RegisterTool(tool tools.Tool) {
for _, agentID := range al.registry.ListAgentIDs() { for _, agentID := range al.registry.ListAgentIDs() {
if agent, ok := al.registry.GetAgent(agentID); ok { if agent, ok := al.registry.GetAgent(agentID); ok {

View file

@ -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. // GetDefaultAgent returns the default agent instance.
func (r *AgentRegistry) GetDefaultAgent() *AgentInstance { func (r *AgentRegistry) GetDefaultAgent() *AgentInstance {
r.mu.RLock() r.mu.RLock()

View file

@ -6,14 +6,27 @@ import "github.com/sipeed/picoclaw/pkg/providers"
// Both SessionManager (legacy JSON backend) and JSONLBackend satisfy this // Both SessionManager (legacy JSON backend) and JSONLBackend satisfy this
// interface, allowing the storage layer to be swapped without touching the // interface, allowing the storage layer to be swapped without touching the
// agent loop code. // 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 { type SessionStore interface {
// AddMessage appends a simple role/content message to the session.
AddMessage(sessionKey, role, content string) AddMessage(sessionKey, role, content string)
// AddFullMessage appends a complete message including tool calls.
AddFullMessage(sessionKey string, msg providers.Message) AddFullMessage(sessionKey string, msg providers.Message)
// GetHistory returns the full message history for the session.
GetHistory(key string) []providers.Message GetHistory(key string) []providers.Message
// GetSummary returns the conversation summary, or "" if none.
GetSummary(key string) string GetSummary(key string) string
// SetSummary replaces the conversation summary.
SetSummary(key, summary string) SetSummary(key, summary string)
// SetHistory replaces the full message history.
SetHistory(key string, history []providers.Message) SetHistory(key string, history []providers.Message)
// TruncateHistory keeps only the last keepLast messages.
TruncateHistory(key string, keepLast int) TruncateHistory(key string, keepLast int)
// Save persists any pending state to durable storage.
Save(key string) error Save(key string) error
// Close releases resources held by the store.
Close() error Close() error
} }