From fbb2a767f648534643bb42f2cf741d29cd65ffcc Mon Sep 17 00:00:00 2001 From: Levi Neely Date: Tue, 4 Aug 2026 16:05:22 +0200 Subject: [PATCH] session: pause/resume persistence, restore paused sessions without spawning processes - Skip SetupToolServer for paused sessions during restore - Pause() and Resume() now persist state to disk - Resume() sets up tool server from scratch for sessions restored paused - IsConnected() returns false when paused - Add SetPaused(), SetToolServer() methods - ProcessKeeper.SetContext() for resume after pause - Session index sorts active-first, paused-last - CreateAgent accepts name= parameter - Remove duplicate SetCwd method --- agent/agent.go | 20 ++++++---- agent/build_runtime.go | 5 ++- kde | 2 +- session/persist.go | 80 +++++++++++++++++++++++++++------------- session/registry.go | 50 +++++++++++++++++-------- session/session.go | 67 +++++++++++++++++++++++++++++---- toolsrv/processkeeper.go | 8 ++++ 7 files changed, 174 insertions(+), 58 deletions(-) diff --git a/agent/agent.go b/agent/agent.go index 5279547..a5c0731 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -243,6 +243,15 @@ func (ag *Agent) SetOutput(h EventHandler) { ag.output = h } +// SetToolServer updates the tool server connection and factory. +// Used when resuming a paused session that was restored without infra. +func (ag *Agent) SetToolServer(newToolServer func() toolsrv.Runner, conn toolsrv.Runner) { + ag.newToolServer = newToolServer + if ag.runtime != nil { + ag.runtime.ToolServer = conn + } +} + // CWD returns the agent's working directory. func (ag *Agent) Cwd() string { ag.stateMu.RLock() @@ -251,13 +260,6 @@ func (ag *Agent) Cwd() string { return c } -// SetCWD sets the agent's working directory (no validation — caller must validate). -func (ag *Agent) SetCwd(dir string) { - ag.stateMu.Lock() - ag.cwd = dir - ag.stateMu.Unlock() -} - // effectiveCwd returns the agent's cwd, falling back to os.Getwd() if empty. func (ag *Agent) effectiveCwd() string { if c := ag.Cwd(); c != "" { @@ -373,7 +375,9 @@ func (ag *Agent) Close() { // SetCWD updates the agent's working directory, preamble references, and dispatcher. func (ag *Agent) SetCWD(dir string) { oldCwd := ag.Cwd() - ag.SetCwd(dir) + ag.stateMu.Lock() + ag.cwd = dir + ag.stateMu.Unlock() if oldCwd != "" && dir != "" && oldCwd != dir { ag.runtime.Preamble = strings.ReplaceAll(ag.runtime.Preamble, oldCwd, dir) } diff --git a/agent/build_runtime.go b/agent/build_runtime.go index 0830b40..23c87ff 100644 --- a/agent/build_runtime.go +++ b/agent/build_runtime.go @@ -23,7 +23,7 @@ func BuildRuntime(cfg *AgentConfig, srv toolsrv.Runner, cwd string, env []string var allToolInfos []toolsrv.ToolInfo var allTools []backend.Tool - if cfg == nil || cfg.ToolsEnabled() { + if srv != nil && (cfg == nil || cfg.ToolsEnabled()) { var listErr error allToolInfos, listErr = srv.ListTools() if listErr != nil { @@ -73,6 +73,9 @@ func BuildRuntime(cfg *AgentConfig, srv toolsrv.Runner, cwd string, env []string } exec := func(ctx context.Context, name string, args json.RawMessage) (string, []backend.ContentBlock, error) { + if srv == nil { + return "", nil, fmt.Errorf("no tool server available") + } raw, err := srv.CallTool(ctx, name, args) if err != nil { return "", nil, err diff --git a/kde b/kde index 265a1fd..22302ce 160000 --- a/kde +++ b/kde @@ -1 +1 @@ -Subproject commit 265a1fdda9c47244d0238b855c2d3ac0ae1b6efb +Subproject commit 22302ce9f139ed8f4aec3e63f7da4306754424c9 diff --git a/session/persist.go b/session/persist.go index e1a261f..0dbf30e 100644 --- a/session/persist.go +++ b/session/persist.go @@ -11,6 +11,7 @@ import ( "ollie/agent" "ollie/backend" + "ollie/toolsrv" ) // --- Persistence --- @@ -21,6 +22,7 @@ type PersistedSession struct { Name string `json:"name"` CWD string `json:"cwd"` Remote string `json:"remote,omitempty"` + Paused bool `json:"paused,omitempty"` Agents []PersistedAgent `json:"agents"` } @@ -65,6 +67,7 @@ func PersistSession(name string) error { Name: sess.Name(), CWD: sess.CWD(), Remote: sess.Remote(), + Paused: sess.IsPaused(), } for _, ag := range agents { @@ -194,27 +197,44 @@ func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) { 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 + // For paused sessions, skip tool server setup entirely. + // It will be set up on Resume(). + var infra *SessionInfra + if !ps.Paused { + var err error + infra, err = SetupToolServer(ToolServerConfig{ + Ctx: serverCtx, + CWD: cwd, + RemoteTarget: ps.Remote, + Yolo: pkgYolo, + }) + if err != nil { + return nil, err + } } // Create empty session - ctx, cancel := context.WithCancel(pkgCtx) + ctx, cancel := context.WithCancel(serverCtx) sess := NewEmpty(ps.ID, ctx, cancel) sess.SetName(ps.Name) - sess.SetProc(infra.Proc) - sess.SetKeeper(infra.Keeper) - sess.SetToolsConn(infra.ToolsConn) + sess.SetRemote(ps.Remote) + if infra != nil { + sess.SetProc(infra.Proc) + sess.SetKeeper(infra.Keeper) + sess.SetToolsConn(infra.ToolsConn) + } + if ps.Paused { + sess.SetPaused(true) + } var restoredAgents []RestoredAgent + // Build prompt env (works without live infra) + promptEnv := agent.PromptEnv(cwd) + if infra != nil { + promptEnv = infra.PromptEnv + } + // Restore each agent for _, pa := range ps.Agents { cfg := LoadAgentConfig(pkgAgentsDir, pa.Profile, nil) @@ -229,16 +249,26 @@ func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) { } // Build prompt layers - layers := BuildPromptLayers(cfg, cwd, ps.ID, pa.ID, infra.PromptEnv, "") + layers := BuildPromptLayers(cfg, cwd, ps.ID, pa.ID, 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...)...) - }) + // Load tools (only if we have a connection) + if infra != nil && infra.ToolsConn != nil { + 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) + env = append(env, promptEnv...) + + var toolsConn toolsrv.Runner + var newToolServer func() toolsrv.Runner + if infra != nil { + toolsConn = infra.ToolsConn + newToolServer = infra.NewToolServer + } + + rt := agent.BuildRuntime(cfg, toolsConn, cwd, env, layers.SystemPrompt, layers.OpModel, layers.EnvBlock) rt.Backend = be // Create history from messages @@ -252,8 +282,8 @@ func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) { AgentID: pa.ID, CWD: cwd, BaseLayers: layers.BaseLayers(), - PromptEnvExtra: infra.PromptEnv, - NewToolServer: infra.NewToolServer, + PromptEnvExtra: promptEnv, + NewToolServer: newToolServer, NewBackend: backend.NewWithName, Log: pkgSink.NewLogger("agent"), SessionID: ps.ID, @@ -272,7 +302,7 @@ func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) { if len(sess.Agents()) == 0 { cancel() - return nil, err + return nil, fmt.Errorf("no agents restored") } // Set uname to first agent's ID @@ -281,12 +311,12 @@ func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) { Register(sess.Name(), sess) toolCount := 0 - if infra.ToolsConn != nil { + if infra != nil && 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) + pkgLog.Info("restored session %s with %d agents, tools=%d, paused=%v", ps.ID, len(restoredAgents), toolCount, ps.Paused) return &RestoredSession{ Session: sess, diff --git a/session/registry.go b/session/registry.go index da5d43f..5be1f9c 100644 --- a/session/registry.go +++ b/session/registry.go @@ -21,7 +21,7 @@ var ( sessions = make(map[string]*Session) // Package config, set via Init. - pkgCtx context.Context + serverCtx context.Context // server lifetime context; parent of all sessions pkgLog *olog.Logger pkgSink *olog.Sink pkgAgentsDir string @@ -42,7 +42,7 @@ type InitConfig struct { // Init initializes the session package with the given configuration. // Must be called before any other session functions. func Init(cfg InitConfig) { - pkgCtx = cfg.Ctx + serverCtx = cfg.Ctx pkgLog = cfg.Log pkgSink = cfg.Sink pkgAgentsDir = cfg.AgentsDir @@ -212,17 +212,32 @@ func waitIdle(poll, timeout time.Duration) { } // BuildIndex builds the session index content. +// Active sessions are listed first, paused sessions last. func BuildIndex() []byte { var sb strings.Builder mu.RLock() - names := make([]string, 0, len(sessions)) - for name := range sessions { - names = append(names, name) + defer mu.RUnlock() + + // Collect sessions and sort: active first (alphabetically), then paused (alphabetically) + type entry struct { + name string + sess *Session + paused bool } - sort.Strings(names) - for _, name := range names { - sess := sessions[name] - paused := sess.IsPaused() + entries := make([]entry, 0, len(sessions)) + for name, sess := range sessions { + entries = append(entries, entry{name: name, sess: sess, paused: sess.IsPaused()}) + } + sort.Slice(entries, func(i, j int) bool { + // Active (not paused) comes before paused + if entries[i].paused != entries[j].paused { + return !entries[i].paused // active first + } + return entries[i].name < entries[j].name + }) + + for _, e := range entries { + sess := e.sess connected := sess.IsConnected() cwd := sess.CWD() var be, model string @@ -236,9 +251,8 @@ func BuildIndex() []byte { } agents := strings.Join(agentList, ";") fmt.Fprintf(&sb, "%s\t%s\t%t\t%t\t%s\t%s\t%s\t%s\n", - name, sess.ID(), paused, connected, cwd, be, model, agents) + e.name, sess.ID(), e.paused, connected, cwd, be, model, agents) } - mu.RUnlock() return []byte(sb.String()) } @@ -260,8 +274,9 @@ func CreateEmpty(name, remote string) (*Session, error) { return nil, fmt.Errorf("session already exists: %s", name) } - ctx, cancel := context.WithCancel(pkgCtx) + ctx, cancel := context.WithCancel(serverCtx) sess := NewEmpty(sessID, ctx, cancel) + sess.SetName(name) sess.SetRemote(remote) // Spawn tool server (local or remote) @@ -297,7 +312,7 @@ func NextUname() string { } // CreateAgent creates a new agent within an existing session. -// Args: [profile] [key=value...] where keys include backend, model, remote. +// Args: [profile] [key=value...] where keys include name, backend, model, remote. func CreateAgent(sessName string, args []string) (*agent.Agent, error) { sess := Lookup(sessName) if sess == nil { @@ -306,10 +321,12 @@ func CreateAgent(sessName string, args []string) (*agent.Agent, error) { // Parse args profile := "default" - var backendName, modelName, remoteTarget string + var agentName, backendName, modelName, remoteTarget string for _, arg := range args { if k, v, ok := strings.Cut(arg, "="); ok { switch k { + case "name": + agentName = v case "backend": backendName = v case "model": @@ -386,6 +403,9 @@ func CreateAgent(sessName string, args []string) (*agent.Agent, error) { }) sess.AddAgent(ag) - pkgLog.Info("created agent %s in session %s (profile=%s)", uname, sessName, profile) + if agentName != "" { + ag.SetName(agentName) + } + pkgLog.Info("created agent %s in session %s (profile=%s)", ag.Name(), sessName, profile) return ag, nil } diff --git a/session/session.go b/session/session.go index 65b9702..e45b791 100644 --- a/session/session.go +++ b/session/session.go @@ -476,12 +476,21 @@ func (a *Session) IsPaused() bool { return a.paused } +// SetPaused sets the paused state directly (used during restore). +func (a *Session) SetPaused(paused bool) { + a.mu.Lock() + defer a.mu.Unlock() + a.paused = paused +} + // IsConnected returns true if the tool server connection is alive. +// Returns false if the session is paused. func (a *Session) IsConnected() bool { a.mu.RLock() + paused := a.paused conn := a.toolsConn a.mu.RUnlock() - if conn == nil { + if paused || conn == nil { return false } if pinger, ok := conn.(interface{ Ping() error }); ok { @@ -491,17 +500,25 @@ func (a *Session) IsConnected() bool { } // Pause stops the tool server to save resources. +// It cancels the session context (stopping all agent operations), +// kills the tool server process, and marks the session as paused. func (a *Session) Pause() error { a.mu.Lock() defer a.mu.Unlock() if a.paused { return fmt.Errorf("session already paused") } + // Cancel the session context to stop all agent operations. + if a.cancel != nil { + a.cancel() + } + // Close the tool server process. if a.keeper != nil { a.keeper.Close() } else if a.proc != nil { a.proc.Close() } + // Close any existing connection. if a.toolsConn != nil { if c, ok := a.toolsConn.(interface{ Close() }); ok { c.Close() @@ -509,25 +526,59 @@ func (a *Session) Pause() error { a.toolsConn = nil } a.paused = true + // Persist the paused state. + go a.saveSession() return nil } // Resume restarts the tool server after a Pause. +// It creates a fresh session context and respawns the tool server. +// For sessions restored in paused state (no keeper), it sets up the tool server from scratch. func (a *Session) Resume() error { a.mu.Lock() defer a.mu.Unlock() if !a.paused { return fmt.Errorf("session not paused") } - if a.keeper == nil { - return fmt.Errorf("no process keeper available") + + // Create a fresh context since the old one was cancelled on pause. + ctx, cancel := context.WithCancel(serverCtx) + a.ctx = ctx + a.cancel = cancel + + // If we have a keeper, use it to respawn. Otherwise set up from scratch. + if a.keeper != nil { + a.keeper.SetContext(ctx) + conn, err := a.keeper.Dial() + if err != nil { + return fmt.Errorf("resume failed: %w", err) + } + a.toolsConn = conn + } else { + // Session was restored paused without infra - set up now. + cwd := a.CWD() + infra, err := SetupToolServer(ToolServerConfig{ + Ctx: ctx, + CWD: cwd, + RemoteTarget: a.remote, + Yolo: pkgYolo, + }) + if err != nil { + return fmt.Errorf("resume setup failed: %w", err) + } + a.proc = infra.Proc + a.keeper = infra.Keeper + a.toolsConn = infra.ToolsConn + + // Update agents with the new tool server connection. + for _, ag := range a.agents { + ag.SetToolServer(infra.NewToolServer, infra.ToolsConn) + } } - conn, err := a.keeper.Dial() - if err != nil { - return fmt.Errorf("resume failed: %w", err) - } - a.toolsConn = conn + a.paused = false + // Persist the resumed state. + go a.saveSession() return nil } diff --git a/toolsrv/processkeeper.go b/toolsrv/processkeeper.go index d44b17e..485cd9a 100644 --- a/toolsrv/processkeeper.go +++ b/toolsrv/processkeeper.go @@ -66,6 +66,14 @@ func (pk *ProcessKeeper) Close() { } } +// SetContext updates the keeper's context. Used on session resume to +// provide a fresh context after the previous one was cancelled on pause. +func (pk *ProcessKeeper) SetContext(ctx context.Context) { + pk.mu.Lock() + defer pk.mu.Unlock() + pk.ctx = ctx +} + func (pk *ProcessKeeper) respawnAndDial() (*Conn, error) { pk.mu.Lock() defer pk.mu.Unlock()