lkn
/
ollie
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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
This commit is contained in:
Levi Neely 2026-08-04 16:05:22 +02:00
parent d200fbad05
commit fbb2a767f6
7 changed files with 174 additions and 58 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
agent/agent.go
View File

@ -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)
}

5
agent/build_runtime.go
View File

@ -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

2
kde

@ -1 +1 @@
Subproject commit 265a1fdda9c47244d0238b855c2d3ac0ae1b6efb
Subproject commit 22302ce9f139ed8f4aec3e63f7da4306754424c9

80
session/persist.go
View File

@ -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,

50
session/registry.go
View File

@ -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
}

67
session/session.go
View File

@ -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
}

8
toolsrv/processkeeper.go
View File

@ -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()