361 lines
8.4 KiB
Go
361 lines
8.4 KiB
Go
package session
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"ollie/agent"
|
|
"ollie/backend"
|
|
coresession "ollie/session"
|
|
"ollie/toolsrv"
|
|
"ollie/tools"
|
|
"olliesrv/fs"
|
|
"olliesrv/prompts"
|
|
)
|
|
|
|
// Kill terminates a session and removes it from the tree.
|
|
func Kill(rs *rootState, id string) {
|
|
rs.mu.Lock()
|
|
sess := rs.sessions[id]
|
|
delete(rs.sessions, id)
|
|
rs.mu.Unlock()
|
|
if sess != nil {
|
|
sess.Cancel()
|
|
sess.Core.Close()
|
|
removePersistedAgent(rs, id)
|
|
rs.cfg.Log.Info("killed session %s", id)
|
|
if rs.cfg.OnSessionKilled != nil {
|
|
rs.cfg.OnSessionKilled(id)
|
|
}
|
|
}
|
|
}
|
|
|
|
// KillFromRoot is the public API for callers with *fs.Tree.
|
|
func KillFromRoot(root *fs.Tree, id string) {
|
|
Kill(state(root), id)
|
|
}
|
|
|
|
// Rename renames a session.
|
|
func Rename(rs *rootState, old, new string) error {
|
|
rs.mu.Lock()
|
|
sess, ok := rs.sessions[old]
|
|
if !ok {
|
|
rs.mu.Unlock()
|
|
return fmt.Errorf("session not found: %s", old)
|
|
}
|
|
if _, exists := rs.sessions[new]; exists {
|
|
rs.mu.Unlock()
|
|
return fmt.Errorf("session already exists: %s", new)
|
|
}
|
|
if sess.Core.Agent().IsRunning() {
|
|
rs.mu.Unlock()
|
|
return fmt.Errorf("cannot rename while agent is running")
|
|
}
|
|
|
|
if err := sess.Core.SetSessionID(new); err != nil {
|
|
rs.mu.Unlock()
|
|
return err
|
|
}
|
|
|
|
sess.id = new
|
|
rs.sessions[new] = sess
|
|
delete(rs.sessions, old)
|
|
rs.mu.Unlock()
|
|
|
|
sess.AppendLog([]byte(fmt.Sprintf(":: session renamed: %s -> %s\n", old, new)))
|
|
rs.cfg.Log.Info("renamed session %s -> %s", old, new)
|
|
if rs.cfg.OnSessionRenamed != nil {
|
|
rs.cfg.OnSessionRenamed(old, new)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// RenameFromRoot is the public API for callers with *fs.Tree.
|
|
func RenameFromRoot(root *fs.Tree, old, new string) error {
|
|
return Rename(state(root), old, new)
|
|
}
|
|
|
|
// Shutdown interrupts all sessions, waits for idle, persists, and closes.
|
|
func Shutdown(root *fs.Tree) {
|
|
rs := state(root)
|
|
InterruptAll(root)
|
|
waitIdle(rs, 100*time.Millisecond, 5*time.Second)
|
|
saveAllSessions(rs)
|
|
rs.mu.Lock()
|
|
ids := make([]string, 0, len(rs.sessions))
|
|
for id := range rs.sessions {
|
|
ids = append(ids, id)
|
|
}
|
|
rs.mu.Unlock()
|
|
for _, id := range ids {
|
|
rs.mu.Lock()
|
|
sess := rs.sessions[id]
|
|
delete(rs.sessions, id)
|
|
rs.mu.Unlock()
|
|
if sess != nil {
|
|
sess.Cancel()
|
|
sess.Core.Close()
|
|
rs.cfg.Log.Info("shutdown session %s", id)
|
|
}
|
|
}
|
|
}
|
|
|
|
func waitIdle(rs *rootState, poll, timeout time.Duration) {
|
|
deadline := time.Now().Add(timeout)
|
|
for time.Now().Before(deadline) {
|
|
allIdle := true
|
|
rs.mu.RLock()
|
|
for _, sess := range rs.sessions {
|
|
if sess.Core.Agent().State() != "idle" {
|
|
allIdle = false
|
|
break
|
|
}
|
|
}
|
|
rs.mu.RUnlock()
|
|
if allIdle {
|
|
return
|
|
}
|
|
time.Sleep(poll)
|
|
}
|
|
rs.cfg.Log.Warn("shutdown: timed out waiting for sessions to become idle, saving anyway")
|
|
}
|
|
|
|
// --- 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
|
|
}
|
|
dir := activeSessionsDir(rs)
|
|
os.MkdirAll(dir, 0700)
|
|
path := filepath.Join(dir, id+".json")
|
|
if err := sess.Core.SaveSession(path); 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()
|
|
}
|
|
agentName := ps.Agent
|
|
if agentName == "" {
|
|
agentName = "default"
|
|
}
|
|
sessID := ps.ID
|
|
|
|
cfg := LoadAgentConfig(rs.cfg.AgentsDir, agentName, 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 := rs.nextUname()
|
|
var newToolServer func() toolsrv.Runner
|
|
var promptEnv []string
|
|
remoteTarget := ps.Remote
|
|
|
|
if remoteTarget != "" {
|
|
rsrv, dialErr := toolsrv.RemoteDial(rs.cfg.Ctx, toolsrv.RemoteConfig{
|
|
SSHTarget: remoteTarget,
|
|
CWD: cwd,
|
|
})
|
|
if dialErr != nil {
|
|
return fmt.Errorf("remote dial: %w", dialErr)
|
|
}
|
|
newToolServer = func() toolsrv.Runner { 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 []toolsrv.Option
|
|
if rs.cfg.Strict {
|
|
execOpts = append(execOpts, toolsrv.WithStrict())
|
|
}
|
|
if rs.cfg.Yolo {
|
|
execOpts = append(execOpts, toolsrv.WithYolo())
|
|
}
|
|
if rs.cfg.ToolRegistry != nil {
|
|
execOpts = append(execOpts, toolsrv.WithToolRegistry(rs.cfg.ToolRegistry, sessID))
|
|
}
|
|
if rs.cfg.SkillsRegistry != nil {
|
|
execOpts = append(execOpts, toolsrv.WithSkillsRegistry(rs.cfg.SkillsRegistry))
|
|
execOpts = append(execOpts, toolsrv.WithBuiltins(tools.Builtins()))
|
|
}
|
|
newToolServer = func() toolsrv.Runner {
|
|
srv, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(cwd), execOpts...)
|
|
return srv
|
|
}
|
|
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(rs.cfg.Enable9P, rs.cfg.EnableDBus, 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()
|
|
rt := agent.BuildRuntime(cfg, toolSrv, cwd, env, sysPrompt, opModel, envBlock)
|
|
|
|
restoredSession := agent.RestoreHistory(ps)
|
|
|
|
var sessPtr *Session
|
|
core := coresession.New(coresession.Config{
|
|
Backend: be,
|
|
AgentName: agentName,
|
|
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"),
|
|
ReadPlanStep: func() string {
|
|
if sessPtr == nil {
|
|
return ""
|
|
}
|
|
sessPtr.mu.RLock()
|
|
data := make([]byte, len(sessPtr.plan))
|
|
copy(data, sessPtr.plan)
|
|
sessPtr.mu.RUnlock()
|
|
return coresession.NextUncheckedStep(data)
|
|
},
|
|
})
|
|
|
|
sessionCtx, sessionCancel := context.WithCancel(rs.cfg.Ctx)
|
|
sess := NewSession(sessID, core, sessionCtx, sessionCancel)
|
|
sessPtr = sess
|
|
sess.uname = uname
|
|
sess.remote = remoteTarget
|
|
|
|
replayMessagesToLog(sess, ps.Messages)
|
|
|
|
rs.mu.Lock()
|
|
rs.sessions[sessID] = sess
|
|
rs.mu.Unlock()
|
|
|
|
rs.cfg.Log.Info("restored session %s (backend=%s model=%s agent=%s)", sessID, backendName, modelName, agentName)
|
|
if rs.cfg.OnSessionCreated != nil {
|
|
rs.cfg.OnSessionCreated(sessID, sess)
|
|
}
|
|
return nil
|
|
}
|