From c4e1950f8ca735f7699b1031df68588174a706f3 Mon Sep 17 00:00:00 2001
From: Levi Neely <141506390+lneely@users.noreply.github.com>
Date: Sun, 19 Jul 2026 21:53:01 +0200
Subject: [PATCH] remote: restore sessions with remote= reconnect to remote
 host

---
 session/session.go | 61 +++++++++++++++++++++++++++++++++++-----------
 1 file changed, 47 insertions(+), 14 deletions(-)

diff --git a/session/session.go b/session/session.go
index 7322412..f6f1b37 100644
--- a/session/session.go
+++ b/session/session.go
@@ -802,21 +802,45 @@ func (s *Manager) restoreSession(ps *agent.PersistedSession) error {
 	}
 
 	uname := s.nextUname()
-	var execOpts []execute.Option
-	if !s.cfg.NoMount {
-		execOpts = append(execOpts, WithMount9P())
+	var newDisp func() tools.Dispatcher
+	var promptEnv []string
+	remoteTarget := ps.Remote
+
+	if remoteTarget != "" {
+		rsrv, dialErr := remote.Dial(context.Background(), remote.Config{
+			SSHTarget: remoteTarget,
+			CWD:       cwd,
+		})
+		if dialErr != nil {
+			return fmt.Errorf("remote dial: %w", dialErr)
+		}
+		newDisp = tools.NewDispatcherFunc(map[string]func() tools.Server{
+			"execute": func() tools.Server { return rsrv },
+		})
+		promptEnv = []string{
+			"PRIME_CWD=" + cwd,
+			"PRIME_PLATFORM=" + rsrv.Info.Platform,
+			"PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", rsrv.Info.IsGitRepo),
+		}
+	} else {
+		var execOpts []execute.Option
+		if !s.cfg.NoMount {
+			execOpts = append(execOpts, WithMount9P())
+		}
+		if s.cfg.Strict {
+			execOpts = append(execOpts, execute.WithStrict())
+		}
+		if s.cfg.Yolo {
+			execOpts = append(execOpts, execute.WithYolo())
+		}
+		newDisp = tools.NewDispatcherFunc(map[string]func() tools.Server{
+			"execute": execute.Decl(cwd, execOpts...),
+		})
+		promptEnv = agent.PromptEnv(cwd)
 	}
-	if s.cfg.Strict {
-		execOpts = append(execOpts, execute.WithStrict())
-	}
-	if s.cfg.Yolo {
-		execOpts = append(execOpts, execute.WithYolo())
-	}
-	newDisp := tools.NewDispatcherFunc(map[string]func() tools.Server{
-		"execute": execute.Decl(cwd, execOpts...),
-	})
+
 	env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname, "OLLIE=" + MountPathFor(uname)}
-	env = append(env, agent.PromptEnv(cwd)...)
+	env = append(env, promptEnv...)
 	rt := agent.BuildRuntime(cfg, newDisp(), cwd, env)
 
 	restoredSession := agent.RestoreSession(ps.Messages)
@@ -825,7 +849,6 @@ func (s *Manager) restoreSession(ps *agent.PersistedSession) error {
 	}
 
 	var sessPtr *Session
-	promptEnv := agent.PromptEnv(cwd)
 	core := agent.NewAgentCore(agent.AgentCoreConfig{
 		Backend:        be,
 		AgentName:      agentName,
@@ -834,6 +857,7 @@ func (s *Manager) restoreSession(ps *agent.PersistedSession) error {
 		SessionID:      sessID,
 		Uname:          uname,
 		CWD:            cwd,
+		Remote:         remoteTarget,
 		Session:        restoredSession,
 		Runtime:        rt,
 		NewDispatcher:  newDisp,
@@ -855,6 +879,14 @@ func (s *Manager) restoreSession(ps *agent.PersistedSession) error {
 	sess := NewSession(sessID, core, ctx, cancel)
 	sessPtr = sess
 	sess.uname = uname
+	sess.remote = remoteTarget
+
+	// Start mount for remote sessions (not managed by local execute.Server).
+	if remoteTarget != "" && !s.cfg.NoMount {
+		ms := &mountState{envExtra: make(map[string]string)}
+		startMount(ms, uname)
+		sess.mount = ms
+	}
 
 	// Replay tail of persisted messages into the chat log so the GUI
 	// and `chat` file show recent history on restore.
@@ -1093,6 +1125,7 @@ func (s *Manager) CreateSession(args []string) (string, error) {
 			SessionID:      sessID,
 			Uname:          uname,
 			CWD:            cwd,
+			Remote:         remoteTarget,
 			Runtime:        rt,
 			NewDispatcher:  newDisp,
 			PromptEnvExtra: promptEnv,