diff --git a/session/session.go b/session/session.go index 7858d26..f16fa02 100644 --- a/session/session.go +++ b/session/session.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "os" + "path/filepath" "sort" "strconv" "strings" @@ -216,6 +217,7 @@ func NewManager(cfg ManagerConfig) *Manager { sessions: make(map[string]*Session), } ss.nextUID.Store(9999) + ss.restoreAllSessions() ss.tree = fs.NewTree(nil, 0, fs.WithStat(func(_ []string, name string) (os.FileInfo, error) { return ss.stat(name) }), fs.WithOpener(func(_ []string, name string) (fs.File, error) { return ss.openEntry(name) }), @@ -636,7 +638,163 @@ func (s *Manager) InterruptAll() { } // Shutdown kills all active sessions. +// --- Session Persistence --- + +func (s *Manager) activeSessionsDir() string { + return filepath.Join(paths.DataDir(), "sessions", "active-9p") +} + +func (s *Manager) persistSession(id string) { + s.mu.RLock() + sess, ok := s.sessions[id] + s.mu.RUnlock() + if !ok { + return + } + dir := s.activeSessionsDir() + os.MkdirAll(dir, 0700) + path := filepath.Join(dir, id+".json") + if err := sess.Core.SaveSession(path); err != nil { + s.cfg.Log.Error("persist session %s: %v", id, err) + } +} + +func (s *Manager) removePersistedSession(id string) { + path := filepath.Join(s.activeSessionsDir(), id+".json") + os.Remove(path) +} + +func (s *Manager) saveAllSessions() { + s.mu.RLock() + ids := make([]string, 0, len(s.sessions)) + for id := range s.sessions { + ids = append(ids, id) + } + s.mu.RUnlock() + for _, id := range ids { + s.persistSession(id) + } +} + +func (s *Manager) restoreAllSessions() { + dir := s.activeSessionsDir() + entries, err := os.ReadDir(dir) + if err != nil { + return + } + for _, e := range entries { + if !strings.HasSuffix(e.Name(), ".json") { + continue + } + path := filepath.Join(dir, e.Name()) + ps, err := agent.LoadPersistedSession(path) + if err != nil { + s.cfg.Log.Error("restore session %s: %v", e.Name(), err) + continue + } + if err := s.restoreSession(ps); err != nil { + s.cfg.Log.Error("restore session %s: %v", ps.ID, err) + } + } + // Clean up persisted files after successful restore + for _, e := range entries { + os.Remove(filepath.Join(dir, e.Name())) + } +} + +func (s *Manager) restoreSession(ps *agent.PersistedSession) error { + cwd := ps.CWD + if cwd == "" { + cwd, _ = os.Getwd() + } + agentName := ps.Agent + if agentName == "" { + agentName = "default" + } + sessID := ps.ID + + cfg := LoadAgentConfig(s.cfg.AgentsDir, agentName, nil) + + backendName := ps.Backend + if backendName == "" && cfg != nil && cfg.Backend != "" { + backendName = cfg.Backend + } + be, err := backend.NewWithName(backendName) + if err != nil { + return fmt.Errorf("backend: %w", err) + } + modelName := ps.Model + if modelName == "" && cfg != nil && cfg.Model != "" { + modelName = cfg.Model + } + if modelName == "" { + modelName = os.Getenv("OLLIE_MODEL") + } + if modelName != "" { + be.SetModel(modelName) + } + + uname := s.nextUname() + var execOpts []execute.Option + if !s.cfg.NoMount { + execOpts = append(execOpts, WithMount9P()) + } + if s.cfg.Strict { + execOpts = append(execOpts, execute.WithStrict()) + } + if s.cfg.Yolo { + execOpts = append(execOpts, execute.WithYolo()) + } + newDisp := tools.NewDispatcherFunc(map[string]func() tools.Server{ + "execute": execute.Decl(cwd, execOpts...), + }) + rt := agent.BuildRuntime(cfg, newDisp(), cwd, []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname, "OLLIE=" + MountPathFor(uname)}) + + restoredSession := agent.RestoreSession(ps.Messages) + if ps.TaskState != nil { + restoredSession.TaskState = ps.TaskState + } + + var sessPtr *Session + core := agent.NewAgentCore(agent.AgentCoreConfig{ + Backend: be, + AgentName: agentName, + AgentsDir: s.cfg.AgentsDir, + SessionsDir: s.cfg.SessionsDir, + SessionID: sessID, + Uname: uname, + CWD: cwd, + Session: restoredSession, + Runtime: rt, + NewDispatcher: newDisp, + Log: s.cfg.Sink.NewLogger("core"), + ReadPlanStep: func() string { + if sessPtr == nil { + return "" + } + sessPtr.mu.RLock() + data := make([]byte, len(sessPtr.plan)) + copy(data, sessPtr.plan) + sessPtr.mu.RUnlock() + return agent.NextUncheckedStep(data) + }, + }) + + ctx, cancel := context.WithCancel(context.Background()) + sess := NewSession(sessID, core, ctx, cancel) + sessPtr = sess + sess.uname = uname + + s.mu.Lock() + s.sessions[sessID] = sess + s.mu.Unlock() + + s.cfg.Log.Info("restored session %s (backend=%s model=%s agent=%s)", sessID, backendName, modelName, agentName) + return nil +} + func (s *Manager) Shutdown() { + s.saveAllSessions() s.mu.Lock() ids := make([]string, 0, len(s.sessions)) for id := range s.sessions { @@ -656,6 +814,7 @@ func (s *Manager) KillSession(id string) { if sess != nil { sess.Cancel() sess.Core.Close() + s.removePersistedSession(id) s.cfg.Log.Info("killed session %s", id) } }