332 lines
8.1 KiB
Go
332 lines
8.1 KiB
Go
package fs
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
|
|
"ollie/agent"
|
|
"ollie/backend"
|
|
"ollie/prompts"
|
|
coresession "ollie/session"
|
|
"ollie/toolsrv"
|
|
)
|
|
|
|
// KillFromRoot is the public API for callers with *Tree.
|
|
func KillFromRoot(root *Tree, id string) {
|
|
Kill(state(root), id)
|
|
}
|
|
|
|
// RenameFromRoot is the public API for callers with *Tree.
|
|
func RenameFromRoot(root *Tree, old, new string) error {
|
|
return Rename(state(root), old, new)
|
|
}
|
|
|
|
// --- Persistence ---
|
|
|
|
func activeSessionsDir(rs *rootState) string {
|
|
return filepath.Join(rs.cfg.SessionsDir, "active")
|
|
}
|
|
|
|
func persistSession(rs *rootState, id string) {
|
|
rs.mu.RLock()
|
|
sess, ok := rs.sessions[id]
|
|
rs.mu.RUnlock()
|
|
if !ok {
|
|
return
|
|
}
|
|
if sess.Core == nil || sess.Core.AgentCount() == 0 {
|
|
// Empty session — no agent to persist. Silently skip.
|
|
return
|
|
}
|
|
dir := activeSessionsDir(rs)
|
|
os.MkdirAll(dir, 0700)
|
|
path := filepath.Join(dir, sess.id+".json")
|
|
ag := sess.Core.Agents()[0]
|
|
if err := ag.SaveFull(path, sess.id, sess.Name(), sess.Core.CWD(), ""); err != nil {
|
|
rs.cfg.Log.Error("persist session %s: %v", id, err)
|
|
}
|
|
}
|
|
|
|
func removePersistedAgent(rs *rootState, id string) {
|
|
path := filepath.Join(activeSessionsDir(rs), id+".json")
|
|
os.Remove(path)
|
|
}
|
|
|
|
func saveAllSessions(rs *rootState) {
|
|
rs.mu.RLock()
|
|
ids := make([]string, 0, len(rs.sessions))
|
|
for id := range rs.sessions {
|
|
ids = append(ids, id)
|
|
}
|
|
rs.mu.RUnlock()
|
|
for _, id := range ids {
|
|
persistSession(rs, id)
|
|
}
|
|
}
|
|
|
|
func restoreAllSessions(rs *rootState) {
|
|
dir := activeSessionsDir(rs)
|
|
entries, err := os.ReadDir(dir)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
type loadedSession struct {
|
|
ps *agent.PersistedAgent
|
|
name string
|
|
}
|
|
var loaded []loadedSession
|
|
for _, e := range entries {
|
|
if !strings.HasSuffix(e.Name(), ".json") {
|
|
continue
|
|
}
|
|
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 and agent display name are independent persisted fields. Never
|
|
// use the agent name (or its immutable ID) as a profile fallback: doing so
|
|
// makes restored sessions report an ID such as "eda74c4e8fb4" as the
|
|
// active profile.
|
|
profile := ps.Profile
|
|
if profile == "" {
|
|
profile = "default"
|
|
}
|
|
sessID := ps.ID
|
|
|
|
cfg := 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 fmt.Errorf("backend: %w", 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(rs)
|
|
}
|
|
var newToolServer func() toolsrv.Runner
|
|
var promptEnv []string
|
|
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{
|
|
SSHTarget: remoteTarget,
|
|
CWD: cwd,
|
|
})
|
|
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, err := keeper.Dial()
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
return conn
|
|
}
|
|
promptEnv = []string{
|
|
"PRIME_CWD=" + cwd,
|
|
"PRIME_PLATFORM=" + proc.Info.Platform,
|
|
"PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", proc.Info.IsGitRepo),
|
|
}
|
|
} else {
|
|
var dialOpts []toolsrv.Option
|
|
if rs.cfg.Yolo {
|
|
dialOpts = append(dialOpts, toolsrv.WithYolo())
|
|
}
|
|
var spawnErr error
|
|
proc, spawnErr = toolsrv.Spawn(rs.cfg.Ctx, cwd, dialOpts...)
|
|
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, err := keeper.Dial()
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
return conn
|
|
}
|
|
promptEnv = agent.PromptEnv(cwd)
|
|
}
|
|
|
|
env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname}
|
|
env = append(env, promptEnv...)
|
|
|
|
var spOverride string
|
|
if cfg != nil {
|
|
spOverride = cfg.SystemPrompt
|
|
}
|
|
sysPrompt := prompts.ResolveSystemPrompt(spOverride)
|
|
envMap := make(map[string]string)
|
|
for _, e := range env {
|
|
if k, v, ok := strings.Cut(e, "="); ok {
|
|
envMap[k] = v
|
|
}
|
|
}
|
|
opModel := prompts.OperationalModel(true, false, envMap)
|
|
platform := "linux"
|
|
isGitRepo := false
|
|
for _, e := range promptEnv {
|
|
if k, v, ok := strings.Cut(e, "="); ok {
|
|
switch k {
|
|
case "PRIME_PLATFORM":
|
|
platform = v
|
|
case "PRIME_IS_GIT_REPO":
|
|
isGitRepo = v == "true"
|
|
}
|
|
}
|
|
}
|
|
envBlock := prompts.Environment(cwd, platform, isGitRepo, "")
|
|
|
|
toolSrv := newToolServer()
|
|
// Load configured tools before building the runtime. BuildRuntime snapshots
|
|
// the loaded-tool list to construct the restored agent's tool definitions.
|
|
if cfg != nil {
|
|
for _, tl := range cfg.AutoLoad {
|
|
if err := loadToolOnRunner(toolSrv, tl); err != nil {
|
|
rs.cfg.Log.Error("session %s: autoLoad tool %q: %v", sessID, tl, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
rt := agent.BuildRuntime(cfg, toolSrv, cwd, env, sysPrompt, opModel, envBlock)
|
|
|
|
restoredSession := agent.RestoreHistory(ps)
|
|
|
|
var sessPtr *Session
|
|
al := NewAgentLog(remoteTarget)
|
|
core := coresession.New(coresession.Config{
|
|
Backend: be,
|
|
Profile: profile,
|
|
AgentsDir: rs.cfg.AgentsDir,
|
|
SessionsDir: rs.cfg.SessionsDir,
|
|
SessionID: sessID,
|
|
AgentID: uname,
|
|
CWD: cwd,
|
|
Remote: remoteTarget,
|
|
History: restoredSession,
|
|
Runtime: rt,
|
|
NewToolServer: newToolServer,
|
|
PromptEnvExtra: promptEnv,
|
|
BaseLayers: []string{sysPrompt, opModel, envBlock},
|
|
Log: rs.cfg.Sink.NewLogger("core"),
|
|
Output: NewEventHandler(al),
|
|
ReadPlanStep: func() string {
|
|
if sessPtr == nil || sessPtr.AgentLog() == nil {
|
|
return ""
|
|
}
|
|
al := sessPtr.AgentLog()
|
|
al.mu.RLock()
|
|
data := make([]byte, len(al.plan))
|
|
copy(data, al.plan)
|
|
al.mu.RUnlock()
|
|
return coresession.NextUncheckedStep(data)
|
|
},
|
|
})
|
|
|
|
sessionCtx, sessionCancel := context.WithCancel(rs.cfg.Ctx)
|
|
sess := NewSession(sessID, core, sessionCtx, sessionCancel)
|
|
if ps.SessionName != "" {
|
|
sess.SetName(ps.SessionName)
|
|
}
|
|
sessPtr = sess
|
|
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
|
|
|
|
// Restore agent display name from persisted session.
|
|
ag := sess.Core.FindAgent(uname)
|
|
if ps.Agent != "" {
|
|
ag.SetName(ps.Agent)
|
|
}
|
|
|
|
replayMessagesToLog(sess.AgentLog(), ps.Messages)
|
|
|
|
// Pre-load autoLoad tools into the remote.
|
|
if cfg != nil {
|
|
for _, tl := range cfg.AutoLoad {
|
|
if err := sess.loadTool(tl, ag); err != nil {
|
|
rs.cfg.Log.Error("session %s: autoLoad tool %q: %v", sessID, tl, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
name := sess.Name()
|
|
rs.mu.Lock()
|
|
rs.sessions[name] = sess
|
|
rs.mu.Unlock()
|
|
|
|
toolCount := 0
|
|
if toolSrv != nil {
|
|
if infos, err := toolSrv.ListTools(); err == nil {
|
|
toolCount = len(infos)
|
|
}
|
|
}
|
|
rs.cfg.Log.Info("restored session %s: cfg=%v autoload=%d tools=%d (backend=%s model=%s agent=%s)",
|
|
sessID, cfg != nil, func() int {
|
|
if cfg != nil {
|
|
return len(cfg.AutoLoad)
|
|
}
|
|
return -1
|
|
}(), toolCount, backendName, modelName, profile)
|
|
return nil
|
|
}
|