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:
parent
841e1d4275
commit
de9d4703b9
6 changed files with 48 additions and 2 deletions
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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 {
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue