From 4ae1584887149ddb34672558db8fcbe21bc15b04 Mon Sep 17 00:00:00 2001 From: Levi Neely Date: Tue, 4 Aug 2026 14:02:15 +0200 Subject: [PATCH] refactor: move persistence to session/, support multi-agent sessions session/persist.go: - New multi-agent format: PersistedSession with []PersistedAgent - PersistSession saves ALL agents, not just first - RestoreAllSessions returns []RestoredSession with []RestoredAgent - Backward compatible: restores legacy single-agent format fs/persist.go: - Thin wrapper that creates AgentLogs for each restored agent - Delegates to session.PersistSession/RestoreAllSessions agent/agent.go: - Added Messages() to get conversation history - Added Usage() to get token/cost stats agent/history.go: - Added RestoreHistoryFromMessages for multi-agent restore session/session.go: - Added Remote() getter --- agent/agent.go | 29 ++++ agent/history.go | 20 +++ fs/persist.go | 207 ++--------------------- session/persist.go | 399 +++++++++++++++++++++++++++++++++++++++++++++ session/session.go | 5 + 5 files changed, 469 insertions(+), 191 deletions(-) create mode 100644 session/persist.go diff --git a/agent/agent.go b/agent/agent.go index 6fa9b1a..d0acb56 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -71,6 +71,35 @@ func (ag *Agent) Name() string { return ag.name } // Profile returns the config profile name (e.g. "default"). func (ag *Agent) Profile() string { return ag.profile } +// Messages returns the conversation history messages. +func (ag *Agent) Messages() []backend.Message { + if ag.history == nil { + return nil + } + return ag.history.history() +} + +// UsageStats holds token usage and cost information. +type UsageStats struct { + TotalInputTokens int + TotalOutputTokens int + TotalRequests int + SessionCostUSD float64 +} + +// Usage returns the agent's usage statistics. +func (ag *Agent) Usage() *UsageStats { + if ag.history == nil { + return nil + } + return &UsageStats{ + TotalInputTokens: ag.history.TotalInputTokens, + TotalOutputTokens: ag.history.TotalOutputTokens, + TotalRequests: ag.history.TotalRequests, + SessionCostUSD: ag.history.SessionCostUSD, + } +} + // SetName changes the agent's display name without reloading config. func (ag *Agent) SetName(name string) { ag.stateMu.Lock() diff --git a/agent/history.go b/agent/history.go index 2b1a649..d03013b 100644 --- a/agent/history.go +++ b/agent/history.go @@ -196,6 +196,26 @@ func RestoreHistory(ps *PersistedAgent) *History { return s } +// RestoreHistoryFromMessages reconstructs a History from just the message list. +// Used for multi-agent session restore where each agent has separate messages. +func RestoreHistoryFromMessages(messages []backend.Message) *History { + s := &History{ + messages: messages, + } + for i := range s.messages { + if s.messages[i].Role == "assistant" && s.messages[i].ID == "" { + s.messages[i].ID = NewResponseID() + } + } + for _, m := range messages { + if m.Role == "user" { + s.goal = m.Content + break + } + } + return s +} + // History is an ephemeral in-memory state backend. type History struct { goal string diff --git a/fs/persist.go b/fs/persist.go index 1e0505b..94eae08 100644 --- a/fs/persist.go +++ b/fs/persist.go @@ -1,210 +1,35 @@ package fs import ( - "context" - "os" - "path/filepath" - "strings" - "sync" - - "ollie/agent" - "ollie/backend" "ollie/session" ) -// --- Persistence --- - -func activeSessionsDir(rs *RootState) string { - return filepath.Join(rs.Cfg.SessionsDir, "active") -} - +// persistSession saves a session to disk. func persistSession(rs *RootState, name string) { - sess := session.Lookup(name) - if sess == nil { - return - } - if sess.AgentCount() == 0 { - return - } - dir := activeSessionsDir(rs) - os.MkdirAll(dir, 0700) - path := filepath.Join(dir, sess.ID()+".json") - ag := sess.Agents()[0] - if err := ag.SaveFull(path, sess.ID(), sess.Name(), sess.CWD(), ""); err != nil { + if err := session.PersistSession(name); err != nil { rs.Cfg.Log.Error("persist session %s: %v", name, err) } } -func removePersistedSession(rs *RootState, id string) { - path := filepath.Join(activeSessionsDir(rs), id+".json") - os.Remove(path) -} - +// saveAllSessions persists all active sessions. func saveAllSessions(rs *RootState) { - for name := range session.Sessions() { - persistSession(rs, name) - } + session.SaveAllSessions() } +// restoreAllSessions restores persisted sessions and wraps them with SessionNodes. func restoreAllSessions(rs *RootState) { - dir := activeSessionsDir(rs) - entries, err := os.ReadDir(dir) - if err != nil { - return - } + restored := session.RestoreAllSessions() + for _, r := range restored { + // Wrap in SessionNode + node := rs.GetOrCreateNode(r.Session) - type loadedSession struct { - ps *agent.PersistedAgent - name string - } - var loaded []loadedSession - for _, e := range entries { - if !strings.HasSuffix(e.Name(), ".json") { - continue + // Create AgentLog for each agent + for _, ra := range r.Agents { + al := NewAgentLog(ra.Remote) + node.SetAgentLog(ra.ID, al) + + // Replay messages to the AgentLog + replayMessagesToLog(al, ra.Messages) } - path := filepath.Join(dir, e.Name()) - ps, err := agent.LoadPersistedAgent(path) - if err != nil { - rs.Cfg.Log.Error("restore session %s: %v", e.Name(), err) - continue - } - loaded = append(loaded, loadedSession{ps: ps, name: e.Name()}) } - - if len(loaded) == 0 { - return - } - - var wg sync.WaitGroup - for _, ls := range loaded { - wg.Add(1) - go func(ls loadedSession) { - defer wg.Done() - if err := restoreSession(rs, ls.ps); err != nil { - rs.Cfg.Log.Error("restore session %s: %v", ls.ps.ID, err) - } - }(ls) - } - wg.Wait() -} - -func restoreSession(rs *RootState, ps *agent.PersistedAgent) error { - cwd := ps.CWD - if cwd == "" { - cwd, _ = os.Getwd() - } - profile := ps.Profile - if profile == "" { - profile = "default" - } - sessID := ps.ID - remoteTarget := ps.Remote - - cfg := session.LoadAgentConfig(rs.Cfg.AgentsDir, profile, nil) - - backendName := ps.Backend - if backendName == "" && cfg != nil && cfg.Backend != "" { - backendName = cfg.Backend - } - be, err := backend.NewWithName(backendName) - if err != nil { - return 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 := ps.Uname - if uname == "" { - uname = session.NextUname() - } - - // Set up tool server infrastructure - infra, err := session.SetupToolServer(session.ToolServerConfig{ - Ctx: rs.Cfg.Ctx, - CWD: cwd, - RemoteTarget: remoteTarget, - Yolo: rs.Cfg.Yolo, - }) - if err != nil { - return err - } - - // Build prompt layers - layers := session.BuildPromptLayers(cfg, cwd, sessID, uname, infra.PromptEnv, "") - - // Load configured tools - session.LoadAutoLoadTools(cfg, infra.ToolsConn, sessID, uname, func(f string, a ...any) { - rs.Cfg.Log.Error("session %s: "+f, append([]any{sessID}, a...)...) - }) - - env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname} - env = append(env, infra.PromptEnv...) - rt := agent.BuildRuntime(cfg, infra.ToolsConn, cwd, env, layers.SystemPrompt, layers.OpModel, layers.EnvBlock) - - restoredHistory := agent.RestoreHistory(ps) - - // Create AgentLog for 9P streaming - al := NewAgentLog(remoteTarget) - - // Create session with the restored agent - sess := session.New(session.Config{ - Backend: be, - Profile: profile, - AgentsDir: rs.Cfg.AgentsDir, - SessionsDir: rs.Cfg.SessionsDir, - SessionID: sessID, - AgentID: uname, - CWD: cwd, - Remote: remoteTarget, - History: restoredHistory, - Runtime: rt, - NewToolServer: infra.NewToolServer, - PromptEnvExtra: infra.PromptEnv, - BaseLayers: layers.BaseLayers(), - Log: rs.Cfg.Sink.NewLogger("core"), - Output: NewEventHandler(al), - ReadPlanStep: func() string { - return "" - }, - }) - ctx, cancel := context.WithCancel(rs.Cfg.Ctx) - sess.SetContext(ctx, cancel) - if ps.SessionName != "" { - sess.SetName(ps.SessionName) - } - sess.SetUname(uname) - sess.SetProc(infra.Proc) - sess.SetKeeper(infra.Keeper) - sess.SetToolsConn(infra.ToolsConn) - - // Restore agent display name - if ag := sess.FindAgent(uname); ag != nil && ps.Agent != "" { - ag.SetName(ps.Agent) - } - - // Replay messages to the AgentLog - replayMessagesToLog(al, ps.Messages) - - // Register session and create SessionNode - session.Register(sess.Name(), sess) - node := rs.GetOrCreateNode(sess) - node.SetAgentLog(uname, al) - - toolCount := 0 - if infra.ToolsConn != nil { - if infos, err := infra.ToolsConn.ListTools(); err == nil { - toolCount = len(infos) - } - } - rs.Cfg.Log.Info("restored session %s: tools=%d (backend=%s model=%s)", - sessID, toolCount, backendName, modelName) - return nil } diff --git a/session/persist.go b/session/persist.go new file mode 100644 index 0000000..c181abb --- /dev/null +++ b/session/persist.go @@ -0,0 +1,399 @@ +package session + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "sync" + + "ollie/agent" + "ollie/backend" +) + +// --- Persistence --- + +// PersistedSession is the on-disk format for a saved session with multiple agents. +type PersistedSession struct { + ID string `json:"id"` + Name string `json:"name"` + CWD string `json:"cwd"` + Remote string `json:"remote,omitempty"` + Agents []PersistedAgent `json:"agents"` +} + +// PersistedAgent is per-agent state within a session. +type PersistedAgent struct { + ID string `json:"id"` // uname + Name string `json:"name"` // display name + Profile string `json:"profile"` + Backend string `json:"backend"` + Model string `json:"model"` + Messages []backend.Message `json:"messages"` + // Usage tracking + TotalInputTokens int `json:"totalInputTokens,omitempty"` + TotalOutputTokens int `json:"totalOutputTokens,omitempty"` + TotalRequests int `json:"totalRequests,omitempty"` + SessionCostUSD float64 `json:"sessionCostUSD,omitempty"` +} + +// ActiveSessionsDir returns the directory for persisted sessions. +func ActiveSessionsDir() string { + return filepath.Join(pkgSessionsDir, "active") +} + +// PersistSession saves a session (with all agents) to disk. +func PersistSession(name string) error { + sess := Lookup(name) + if sess == nil { + return nil + } + agents := sess.Agents() + if len(agents) == 0 { + return nil + } + + ps := PersistedSession{ + ID: sess.ID(), + Name: sess.Name(), + CWD: sess.CWD(), + Remote: sess.Remote(), + } + + for _, ag := range agents { + pa := PersistedAgent{ + ID: ag.ID(), + Name: ag.Name(), + Profile: ag.Profile(), + Backend: ag.BackendName(), + Model: ag.ModelName(), + Messages: ag.Messages(), + } + if usage := ag.Usage(); usage != nil { + pa.TotalInputTokens = usage.TotalInputTokens + pa.TotalOutputTokens = usage.TotalOutputTokens + pa.TotalRequests = usage.TotalRequests + pa.SessionCostUSD = usage.SessionCostUSD + } + ps.Agents = append(ps.Agents, pa) + } + + dir := ActiveSessionsDir() + os.MkdirAll(dir, 0700) + path := filepath.Join(dir, sess.ID()+".json") + + data, err := json.MarshalIndent(ps, "", " ") + if err != nil { + return err + } + return os.WriteFile(path, data, 0600) +} + +// RemovePersistedSession removes a persisted session file. +func RemovePersistedSession(id string) { + path := filepath.Join(ActiveSessionsDir(), id+".json") + os.Remove(path) +} + +// SaveAllSessions persists all active sessions. +func SaveAllSessions() { + for name := range Sessions() { + if err := PersistSession(name); err != nil { + pkgLog.Error("persist session %s: %v", name, err) + } + } +} + +// RestoredAgent holds restored agent info for fs layer. +type RestoredAgent struct { + ID string + Remote string + Messages []backend.Message +} + +// RestoredSession holds a restored session and metadata for fs layer. +type RestoredSession struct { + Session *Session + Agents []RestoredAgent + ToolCount int +} + +// RestoreAllSessions restores all persisted sessions. +// Returns the restored sessions for the fs layer to wrap with SessionNodes. +func RestoreAllSessions() []*RestoredSession { + dir := ActiveSessionsDir() + entries, err := os.ReadDir(dir) + if err != nil { + return nil + } + + var files []string + for _, e := range entries { + if strings.HasSuffix(e.Name(), ".json") { + files = append(files, filepath.Join(dir, e.Name())) + } + } + + if len(files) == 0 { + return nil + } + + var mu sync.Mutex + var results []*RestoredSession + var wg sync.WaitGroup + for _, path := range files { + wg.Add(1) + go func(path string) { + defer wg.Done() + rs, err := restoreSessionFromFile(path) + if err != nil { + pkgLog.Error("restore session %s: %v", path, err) + return + } + mu.Lock() + results = append(results, rs) + mu.Unlock() + }(path) + } + wg.Wait() + return results +} + +func restoreSessionFromFile(path string) (*RestoredSession, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + + // Try new multi-agent format first + var ps PersistedSession + if err := json.Unmarshal(data, &ps); err == nil && ps.ID != "" && len(ps.Agents) > 0 { + return restoreMultiAgentSession(&ps) + } + + // Fall back to legacy single-agent format + var legacy agent.PersistedAgent + if err := json.Unmarshal(data, &legacy); err != nil { + return nil, err + } + return restoreLegacySession(&legacy) +} + +func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) { + cwd := ps.CWD + if cwd == "" { + cwd, _ = os.Getwd() + } + + // Set up tool server infrastructure (shared by all agents) + infra, err := SetupToolServer(ToolServerConfig{ + Ctx: pkgCtx, + CWD: cwd, + RemoteTarget: ps.Remote, + Yolo: pkgYolo, + }) + if err != nil { + return nil, err + } + + // Create empty session + ctx, cancel := context.WithCancel(pkgCtx) + sess := NewEmpty(ps.ID, ctx, cancel) + sess.SetName(ps.Name) + sess.SetProc(infra.Proc) + sess.SetKeeper(infra.Keeper) + sess.SetToolsConn(infra.ToolsConn) + + var restoredAgents []RestoredAgent + + // Restore each agent + for _, pa := range ps.Agents { + cfg := LoadAgentConfig(pkgAgentsDir, pa.Profile, nil) + + be, err := backend.NewWithName(pa.Backend) + if err != nil { + pkgLog.Error("restore agent %s: backend %s: %v", pa.ID, pa.Backend, err) + continue + } + if pa.Model != "" { + be.SetModel(pa.Model) + } + + // Build prompt layers + layers := BuildPromptLayers(cfg, cwd, ps.ID, pa.ID, infra.PromptEnv, "") + + // Load tools + LoadAutoLoadTools(cfg, infra.ToolsConn, ps.ID, pa.ID, func(f string, a ...any) { + pkgLog.Error("session %s agent %s: "+f, append([]any{ps.ID, pa.ID}, a...)...) + }) + + env := []string{"OLLIE_SESSION_ID=" + ps.ID, "OLLIE_UNAME=" + pa.ID} + env = append(env, infra.PromptEnv...) + rt := agent.BuildRuntime(cfg, infra.ToolsConn, cwd, env, layers.SystemPrompt, layers.OpModel, layers.EnvBlock) + rt.Backend = be + + // Create history from messages + history := agent.RestoreHistoryFromMessages(pa.Messages) + + ag := agent.NewAgent(agent.AgentCfg{ + History: history, + Runtime: rt, + Profile: pa.Profile, + AgentsDir: pkgAgentsDir, + AgentID: pa.ID, + CWD: cwd, + BaseLayers: layers.BaseLayers(), + PromptEnvExtra: infra.PromptEnv, + NewToolServer: infra.NewToolServer, + NewBackend: backend.NewWithName, + Log: pkgSink.NewLogger("agent"), + SessionID: ps.ID, + }) + ag.SetName(pa.Name) + + sess.AddAgent(ag) + restoredAgents = append(restoredAgents, RestoredAgent{ + ID: pa.ID, + Remote: ps.Remote, + Messages: pa.Messages, + }) + } + + if len(sess.Agents()) == 0 { + cancel() + return nil, err + } + + // Set uname to first agent's ID + sess.SetUname(sess.Agents()[0].ID()) + + Register(sess.Name(), sess) + + toolCount := 0 + if infra.ToolsConn != nil { + if infos, err := infra.ToolsConn.ListTools(); err == nil { + toolCount = len(infos) + } + } + pkgLog.Info("restored session %s with %d agents, tools=%d", ps.ID, len(restoredAgents), toolCount) + + return &RestoredSession{ + Session: sess, + Agents: restoredAgents, + ToolCount: toolCount, + }, nil +} + +// restoreLegacySession handles the old single-agent format for backward compatibility. +func restoreLegacySession(ps *agent.PersistedAgent) (*RestoredSession, error) { + cwd := ps.CWD + if cwd == "" { + cwd, _ = os.Getwd() + } + profile := ps.Profile + if profile == "" { + profile = "default" + } + sessID := ps.ID + remoteTarget := ps.Remote + + cfg := LoadAgentConfig(pkgAgentsDir, profile, nil) + + backendName := ps.Backend + if backendName == "" && cfg != nil && cfg.Backend != "" { + backendName = cfg.Backend + } + be, err := backend.NewWithName(backendName) + if err != nil { + return nil, 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 := ps.Uname + if uname == "" { + uname = NextUname() + } + + infra, err := SetupToolServer(ToolServerConfig{ + Ctx: pkgCtx, + CWD: cwd, + RemoteTarget: remoteTarget, + Yolo: pkgYolo, + }) + if err != nil { + return nil, err + } + + layers := BuildPromptLayers(cfg, cwd, sessID, uname, infra.PromptEnv, "") + + LoadAutoLoadTools(cfg, infra.ToolsConn, sessID, uname, func(f string, a ...any) { + pkgLog.Error("session %s: "+f, append([]any{sessID}, a...)...) + }) + + env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname} + env = append(env, infra.PromptEnv...) + rt := agent.BuildRuntime(cfg, infra.ToolsConn, cwd, env, layers.SystemPrompt, layers.OpModel, layers.EnvBlock) + + restoredHistory := agent.RestoreHistory(ps) + + sess := New(Config{ + Backend: be, + Profile: profile, + AgentsDir: pkgAgentsDir, + SessionsDir: pkgSessionsDir, + SessionID: sessID, + AgentID: uname, + CWD: cwd, + Remote: remoteTarget, + History: restoredHistory, + Runtime: rt, + NewToolServer: infra.NewToolServer, + PromptEnvExtra: infra.PromptEnv, + BaseLayers: layers.BaseLayers(), + Log: pkgSink.NewLogger("core"), + }) + ctx, cancel := context.WithCancel(pkgCtx) + sess.SetContext(ctx, cancel) + if ps.SessionName != "" { + sess.SetName(ps.SessionName) + } + sess.SetUname(uname) + sess.SetProc(infra.Proc) + sess.SetKeeper(infra.Keeper) + sess.SetToolsConn(infra.ToolsConn) + + if ag := sess.FindAgent(uname); ag != nil && ps.Agent != "" { + ag.SetName(ps.Agent) + } + + Register(sess.Name(), sess) + + toolCount := 0 + if infra.ToolsConn != nil { + if infos, err := infra.ToolsConn.ListTools(); err == nil { + toolCount = len(infos) + } + } + pkgLog.Info("restored legacy session %s: tools=%d (backend=%s model=%s)", + sessID, toolCount, backendName, modelName) + + return &RestoredSession{ + Session: sess, + Agents: []RestoredAgent{{ + ID: uname, + Remote: remoteTarget, + Messages: ps.Messages, + }}, + ToolCount: toolCount, + }, nil +} diff --git a/session/session.go b/session/session.go index 33c5663..4f4db46 100644 --- a/session/session.go +++ b/session/session.go @@ -284,6 +284,11 @@ func (a *Session) CWD() string { return wd } +// Remote returns the remote target (e.g., SSH host) if any. +func (a *Session) Remote() string { + return a.remote +} + // SetCWD validates and sets the working directory on all agents. func (a *Session) SetCWD(dir string) error { dir = paths.ExpandHome(dir)