restore active sessions from data directory

main.go

@@ -135,12 +135,6 @@ func cmdMount() {
 	fmt.Printf("mounted %s at %s (pid %d)\n", addr, mnt, cmd.Process.Pid)
 }
 
-
-
-
-
-
-
 func runServer(sockPath string) {
 	env.EnsureDefaults()
 
@@ -162,7 +156,7 @@ func runServer(sockPath string) {
 	sink := olog.NewSink(os.Stdout, os.Stderr, olog.ParseLevel(os.Getenv("OLLIE_LOG"), olog.LevelWarn))
 
 	agentsDirs := agent.AgentsDirs()
-	sessionsDir := paths.CfgDir() + "/sessions"
+	sessionsDir := paths.DataDir() + "/sessions"
 
 	// D-Bus adapter (initialized after manager so callbacks can reference it).
 	var dbusAdapter *DBusAdapter
@@ -273,9 +267,6 @@ func runServer(sockPath string) {
 	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
 	<-sigChan
 
-	if srv != nil {
-		srv.InterruptAll()
-	}
 	fmt.Println("shutting down")
 	if dbusAdapter != nil {
 		dbusAdapter.Close()
@@ -295,7 +286,6 @@ func runServer(sockPath string) {
 	sink.Flush()
 }
 
-
 // ToolIndex generates a tool index from a tree's file listing.
 // NewRootStore returns a read-only FileTree for synthetic root-level files.
 func NewRootStore() *fs.Tree {

session/session.go

@@ -212,8 +212,6 @@ func (sess *Session) SetPlan(p []byte) { sess.plan = p }
 // PrevPrompt returns the last submitted prompt. Caller must hold Mu().RLock().
 func (sess *Session) PrevPrompt() []byte { return sess.prevPrompt }
 
-
-
 var sessionStoreOrder = []string{"new", "idx", "ls", "kill", "sh", "b", "bfg", "bbg", "cleanup"}
 
 // FileMode returns the mode for a fixed session file,
@@ -266,6 +264,7 @@ type Manager struct {
 
 // Tree returns the Tree view of the session namespace.
 func (s *Manager) Tree() *fs.Tree { return s.tree }
+
 // nextUname generates the next uname atomically.
 func (s *Manager) nextUname() string {
 	return fmt.Sprintf("%d", s.nextUID.Add(1))
@@ -569,7 +568,9 @@ func (s *Manager) openEntry(name string) (fs.File, error) {
 	switch name {
 	case "new":
 		return &fs.FileConfig{
-			StatFn: func() (os.FileInfo, error) { return &fs.SyntheticFileInfo{Name_: "new", Mode_: fs.Perms[fs.PathSessions].Files["new"]}, nil },
+			StatFn: func() (os.FileInfo, error) {
+				return &fs.SyntheticFileInfo{Name_: "new", Mode_: fs.Perms[fs.PathSessions].Files["new"]}, nil
+			},
 			ReadFn: func() ([]byte, error) {
 				return []byte("name=\ncwd=\nremote=\nbackend=\nmodel=\nagent=\nmaxTokens=\nmaxCompletionTokens=\ntemperature=\ntopP=\ntopK=\nminP=\ntopA=\nfrequencyPenalty=\npresencePenalty=\nrepetitionPenalty=\nreasoning=\nreasoningEffort=\nincludeReasoning=\nresponseFormat=\nstop=\nverbosity=\n"), nil
 			},
@@ -581,7 +582,9 @@ func (s *Manager) openEntry(name string) (fs.File, error) {
 		}, nil
 	case "idx":
 		return &fs.FileConfig{
-			StatFn:         func() (os.FileInfo, error) { return &fs.SyntheticFileInfo{Name_: "idx", Mode_: fs.Perms[fs.PathSessions].Files["idx"]}, nil },
+			StatFn: func() (os.FileInfo, error) {
+				return &fs.SyntheticFileInfo{Name_: "idx", Mode_: fs.Perms[fs.PathSessions].Files["idx"]}, nil
+			},
 			ReadFn:         func() ([]byte, error) { return s.index(), nil },
 			WriteFn:        func([]byte) error { return fmt.Errorf("idx: read-only") },
 			BlockingReadFn: notBlocking,
@@ -589,7 +592,9 @@ func (s *Manager) openEntry(name string) (fs.File, error) {
 	}
 	if _, ok := fs.Perms[fs.PathSessions].Files[name]; ok {
 		return &fs.FileConfig{
-			StatFn: func() (os.FileInfo, error) { return &fs.SyntheticFileInfo{Name_: name, Mode_: fs.Perms[fs.PathSessions].Files[name]}, nil },
+			StatFn: func() (os.FileInfo, error) {
+				return &fs.SyntheticFileInfo{Name_: name, Mode_: fs.Perms[fs.PathSessions].Files[name]}, nil
+			},
 			ReadFn: func() ([]byte, error) {
 				return s.cfg.ReadFile(paths.CfgDir() + "/scripts/s/" + name)
 			},
@@ -713,7 +718,7 @@ func (s *Manager) InterruptAll() {
 // --- Session Persistence ---
 
 func (s *Manager) activeSessionsDir() string {
-	return filepath.Join(paths.DataDir(), "sessions", "active")
+	return filepath.Join(s.cfg.SessionsDir, "active")
 }
 
 func (s *Manager) persistSession(id string) {
@@ -766,11 +771,10 @@ func (s *Manager) restoreAllSessions() {
 		}
 		if err := s.restoreSession(ps); err != nil {
 			s.cfg.Log.Error("restore session %s: %v", ps.ID, err)
+			continue
 		}
-	}
-	// Clean up persisted files after successful restore
-	for _, e := range entries {
-		os.Remove(filepath.Join(dir, e.Name()))
+		// Keep active snapshots until the session is explicitly killed.
+		// Shutdown refreshes these files with the latest state.
 	}
 }
 
@@ -938,6 +942,10 @@ func (s *Manager) restoreSession(ps *agent.PersistedSession) error {
 }
 
 func (s *Manager) Shutdown() {
+	// Interrupt all in-progress turns and wait for them to finish
+	// before persisting state, so we capture the latest messages.
+	s.InterruptAll()
+	s.waitIdle(100*time.Millisecond, 5*time.Second)
 	s.saveAllSessions()
 	s.mu.Lock()
 	ids := make([]string, 0, len(s.sessions))
@@ -961,6 +969,27 @@ func (s *Manager) Shutdown() {
 	}
 }
 
+// waitIdle polls all sessions until they are idle or the timeout expires.
+func (s *Manager) waitIdle(poll, timeout time.Duration) {
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		allIdle := true
+		s.mu.RLock()
+		for _, sess := range s.sessions {
+			if sess.Core.State() != "idle" {
+				allIdle = false
+				break
+			}
+		}
+		s.mu.RUnlock()
+		if allIdle {
+			return
+		}
+		time.Sleep(poll)
+	}
+	s.cfg.Log.Warn("shutdown: timed out waiting for sessions to become idle, saving anyway")
+}
+
 func (s *Manager) KillSession(id string) {
 	s.mu.Lock()
 	sess := s.sessions[id]
@@ -1424,7 +1453,6 @@ func (s *Manager) Complete(cwd, filePath, prefix, suffix, extraContext string) (
 	return result, nil
 }
 
-
 func stripCompletionNoise(s string) string {
 	var lines []string
 	for _, line := range strings.Split(s, "\n") {