diff --git a/fs/lifecycle.go b/fs/lifecycle.go index ee3d362..b50b180 100644 --- a/fs/lifecycle.go +++ b/fs/lifecycle.go @@ -155,7 +155,9 @@ func Kill(rs *rootState, id string) { sess.Core.Close() } sess.Cancel() - if sess.proc != nil { + if sess.keeper != nil { + sess.keeper.Close() + } else if sess.proc != nil { sess.proc.Close() } removePersistedAgent(rs, sess.id) @@ -221,7 +223,9 @@ func Shutdown(root *Tree) { if sess.Core != nil { sess.Core.Close() } - if sess.proc != nil { + if sess.keeper != nil { + sess.keeper.Close() + } else if sess.proc != nil { sess.proc.Close() } rs.cfg.Log.Info("shutdown session %s", id) @@ -367,6 +371,9 @@ func Create(rs *rootState, args []string) (string, error) { return "", fmt.Errorf("spawn tool server: %w", err) } sess.proc = proc + sess.keeper = toolsrv.NewProcessKeeper(sessionCtx, proc, func(ctx context.Context) (*toolsrv.Process, error) { + return toolsrv.Spawn(ctx, cwd, dialOpts...) + }) conn, err := toolsrv.Dial(sessionCtx, toolsrv.LocalAddr(proc.Socket)) if err != nil { sessionCancel() @@ -502,8 +509,15 @@ func CreateAgent(rs *rootState, sessName string, args []string) error { if err != nil { return fmt.Errorf("remote spawn: %w", err) } + keeper := toolsrv.NewProcessKeeper(rs.cfg.Ctx, proc, func(ctx context.Context) (*toolsrv.Process, error) { + return toolsrv.SpawnRemote(ctx, toolsrv.RemoteConfig{SSHTarget: remoteTarget, CWD: cwd}) + }) + sess.keeper = keeper newToolServer = func() toolsrv.Runner { - conn, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket)) + conn, err := keeper.Dial() + if err != nil { + return nil + } return conn } promptEnv = []string{ @@ -514,8 +528,19 @@ func CreateAgent(rs *rootState, sessName string, args []string) error { } else if sess.proc != nil { // Reuse the tool server spawned during Create. proc = sess.proc - newToolServer = func() toolsrv.Runner { - return sess.toolsConn + if sess.keeper != nil { + newToolServer = func() toolsrv.Runner { + conn, err := sess.keeper.Dial() + if err != nil { + rs.cfg.Log.Error("toolsrv dial: %v", err) + return nil + } + return conn + } + } else { + newToolServer = func() toolsrv.Runner { + return sess.toolsConn + } } if cwc, ok := sess.toolsConn.(interface{ SetCWD(string) }); ok { cwc.SetCWD(cwd) @@ -526,8 +551,16 @@ func CreateAgent(rs *rootState, sessName string, args []string) error { if err != nil { return fmt.Errorf("local spawn: %w", err) } + keeper := toolsrv.NewProcessKeeper(rs.cfg.Ctx, proc, func(ctx context.Context) (*toolsrv.Process, error) { + return toolsrv.Spawn(ctx, cwd, dialOpts...) + }) + sess.keeper = keeper newToolServer = func() toolsrv.Runner { - conn, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket)) + conn, err := keeper.Dial() + if err != nil { + rs.cfg.Log.Error("toolsrv dial: %v", err) + return nil + } return conn } promptEnv = agent.PromptEnv(cwd) diff --git a/fs/persist.go b/fs/persist.go index 33d8196..b8fe316 100644 --- a/fs/persist.go +++ b/fs/persist.go @@ -155,6 +155,8 @@ func restoreSession(rs *rootState, ps *agent.PersistedAgent) error { var proc *toolsrv.Process remoteTarget := ps.Remote + var keeper *toolsrv.ProcessKeeper + if remoteTarget != "" { var spawnErr error proc, spawnErr = toolsrv.SpawnRemote(rs.cfg.Ctx, toolsrv.RemoteConfig{ @@ -164,8 +166,14 @@ func restoreSession(rs *rootState, ps *agent.PersistedAgent) error { if spawnErr != nil { return fmt.Errorf("remote spawn: %w", spawnErr) } + keeper = toolsrv.NewProcessKeeper(rs.cfg.Ctx, proc, func(ctx context.Context) (*toolsrv.Process, error) { + return toolsrv.SpawnRemote(ctx, toolsrv.RemoteConfig{SSHTarget: remoteTarget, CWD: cwd}) + }) newToolServer = func() toolsrv.Runner { - conn, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket)) + conn, err := keeper.Dial() + if err != nil { + return nil + } return conn } promptEnv = []string{ @@ -183,8 +191,14 @@ func restoreSession(rs *rootState, ps *agent.PersistedAgent) error { if spawnErr != nil { return fmt.Errorf("local spawn: %w", spawnErr) } + keeper = toolsrv.NewProcessKeeper(rs.cfg.Ctx, proc, func(ctx context.Context) (*toolsrv.Process, error) { + return toolsrv.Spawn(ctx, cwd, dialOpts...) + }) newToolServer = func() toolsrv.Runner { - conn, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(proc.Socket)) + conn, err := keeper.Dial() + if err != nil { + return nil + } return conn } promptEnv = agent.PromptEnv(cwd) @@ -274,6 +288,7 @@ func restoreSession(rs *rootState, ps *agent.PersistedAgent) error { sess.uname = uname sess.SetAgentLog(uname, al) sess.proc = proc + sess.keeper = keeper // Use the multiplexed connection for both agent execution and 9P management. sess.toolsConn = toolSrv diff --git a/fs/types.go b/fs/types.go index 76de7f6..75cfc1b 100644 --- a/fs/types.go +++ b/fs/types.go @@ -23,7 +23,8 @@ type Session struct { Core *session.Session SessionCtx context.Context cancel context.CancelFunc - proc *toolsrv.Process // toolsrv subprocess (session owns lifecycle) + proc *toolsrv.Process // toolsrv subprocess (session owns lifecycle) + keeper *toolsrv.ProcessKeeper // resilient process manager (respawns on failure) // toolsConn is the multiplexed connection to the remote tool server, // shared by agent execution and 9P tool-management handlers. diff --git a/toolsrv/processkeeper.go b/toolsrv/processkeeper.go new file mode 100644 index 0000000..d44b17e --- /dev/null +++ b/toolsrv/processkeeper.go @@ -0,0 +1,104 @@ +package toolsrv + +import ( + "context" + "fmt" + "sync" + "time" +) + +// ProcessKeeper maintains a running ollie-remote process and respawns it +// on failure. Multiple agents can Dial() the same keeper to get independent +// connections. If the process dies, the next Dial() triggers a respawn. +type ProcessKeeper struct { + mu sync.Mutex + ctx context.Context + proc *Process + spawn func(context.Context) (*Process, error) +} + +// NewProcessKeeper creates a keeper with an initial process and a spawn factory. +func NewProcessKeeper(ctx context.Context, proc *Process, spawn func(context.Context) (*Process, error)) *ProcessKeeper { + return &ProcessKeeper{ + ctx: ctx, + proc: proc, + spawn: spawn, + } +} + +// Dial returns a new Conn to the managed process. If the process is dead +// (socket unreachable), it respawns before dialing. +func (pk *ProcessKeeper) Dial() (*Conn, error) { + pk.mu.Lock() + sock := pk.proc.Socket + pk.mu.Unlock() + + // Try dialing the existing socket. + conn, err := Dial(pk.ctx, LocalAddr(sock)) + if err == nil { + return conn, nil + } + + // Dial failed — process likely dead. Respawn. + return pk.respawnAndDial() +} + +// Socket returns the current process socket path. +func (pk *ProcessKeeper) Socket() string { + pk.mu.Lock() + defer pk.mu.Unlock() + return pk.proc.Socket +} + +// Process returns the current underlying Process. +func (pk *ProcessKeeper) Process() *Process { + pk.mu.Lock() + defer pk.mu.Unlock() + return pk.proc +} + +// Close shuts down the managed process. +func (pk *ProcessKeeper) Close() { + pk.mu.Lock() + defer pk.mu.Unlock() + if pk.proc != nil { + pk.proc.Close() + } +} + +func (pk *ProcessKeeper) respawnAndDial() (*Conn, error) { + pk.mu.Lock() + defer pk.mu.Unlock() + + // Another goroutine may have respawned already — try current socket first. + if conn, err := Dial(pk.ctx, LocalAddr(pk.proc.Socket)); err == nil { + return conn, nil + } + + // Clean up old process. + pk.proc.Close() + + var lastErr error + for attempt := 0; attempt < 3; attempt++ { + if pk.ctx.Err() != nil { + return nil, pk.ctx.Err() + } + proc, err := pk.spawn(pk.ctx) + if err != nil { + lastErr = err + time.Sleep(time.Duration(attempt+1) * 2 * time.Second) + continue + } + pk.proc = proc + + conn, err := Dial(pk.ctx, LocalAddr(proc.Socket)) + if err != nil { + lastErr = err + proc.Close() + time.Sleep(time.Duration(attempt+1) * 2 * time.Second) + continue + } + return conn, nil + } + return nil, fmt.Errorf("respawn failed after 3 attempts: %w", lastErr) +}