510 lines
12 KiB
Go
510 lines
12 KiB
Go
// Package session manages the lifecycle of agent sessions.
|
|
//
|
|
// A session owns one or more agents, a tool server connection, and persistence
|
|
// state. The package-level registry (initialized via Init) tracks all active
|
|
// sessions and provides lookup, creation, and teardown.
|
|
//
|
|
// Typical flow: Init → Create → Register → (use) → Kill or Shutdown.
|
|
package session
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"fmt"
|
|
"os"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"ollie/cmd/olliesrv/internal/agent"
|
|
"ollie/cmd/olliesrv/internal/backend"
|
|
olog "ollie/log"
|
|
"ollie/util"
|
|
)
|
|
|
|
// Package-level session registry with thread-safe access.
|
|
var (
|
|
mu sync.RWMutex
|
|
sessions = make(map[string]*Session)
|
|
|
|
// Package config, set via Init.
|
|
serverCtx context.Context
|
|
pkgLog *olog.Logger
|
|
pkgSink *olog.Sink
|
|
pkgAgentsDir string
|
|
pkgSessionsDir string
|
|
pkgYolo bool
|
|
pkgBypassNotify BypassNotifyFunc
|
|
)
|
|
|
|
// InitConfig configures the session package.
|
|
type InitConfig struct {
|
|
Ctx context.Context
|
|
Log *olog.Logger
|
|
Sink *olog.Sink
|
|
AgentsDir string
|
|
SessionsDir string
|
|
Yolo bool
|
|
BypassNotify BypassNotifyFunc // may be nil if bypass is disabled
|
|
}
|
|
|
|
// Init initializes the session package with the given configuration.
|
|
func Init(cfg InitConfig) {
|
|
serverCtx = cfg.Ctx
|
|
pkgLog = cfg.Log
|
|
pkgSink = cfg.Sink
|
|
pkgAgentsDir = cfg.AgentsDir
|
|
pkgSessionsDir = cfg.SessionsDir
|
|
pkgYolo = cfg.Yolo
|
|
pkgBypassNotify = cfg.BypassNotify
|
|
}
|
|
|
|
// Sessions returns a snapshot of all sessions.
|
|
func Sessions() map[string]*Session {
|
|
mu.RLock()
|
|
defer mu.RUnlock()
|
|
result := make(map[string]*Session, len(sessions))
|
|
for k, v := range sessions {
|
|
result[k] = v
|
|
}
|
|
return result
|
|
}
|
|
|
|
// Lookup returns the session by name or ID, or nil.
|
|
func Lookup(nameOrID string) *Session {
|
|
mu.RLock()
|
|
defer mu.RUnlock()
|
|
if s := sessions[nameOrID]; s != nil {
|
|
return s
|
|
}
|
|
for _, sess := range sessions {
|
|
if sess.ID == nameOrID {
|
|
return sess
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Register adds a session to the registry.
|
|
func Register(name string, sess *Session) {
|
|
mu.Lock()
|
|
sessions[name] = sess
|
|
mu.Unlock()
|
|
}
|
|
|
|
// Kill terminates a session and removes it from the registry.
|
|
func Kill(name string) error {
|
|
mu.Lock()
|
|
sess := sessions[name]
|
|
if sess == nil {
|
|
mu.Unlock()
|
|
return fmt.Errorf("session not found: %s", name)
|
|
}
|
|
delete(sessions, name)
|
|
mu.Unlock()
|
|
|
|
RemovePersistedSession(sess.ID)
|
|
|
|
var agentIDs []string
|
|
for _, ag := range sess.Agents() {
|
|
agentIDs = append(agentIDs, ag.ID())
|
|
}
|
|
sess.Close()
|
|
if sess.Cancel != nil {
|
|
sess.Cancel()
|
|
}
|
|
if sess.Keeper != nil {
|
|
sess.Keeper.Close()
|
|
} else if sess.Proc != nil {
|
|
sess.Proc.Close()
|
|
}
|
|
|
|
pkgLog.Info("killed session %s", name)
|
|
for _, agentID := range agentIDs {
|
|
PublishEvent("session."+sess.ID+".agent."+agentID+".kill", "")
|
|
}
|
|
PublishEvent("session."+sess.ID+".kill", "")
|
|
return nil
|
|
}
|
|
|
|
// Rename renames a session.
|
|
func Rename(oldName, newName string) error {
|
|
mu.Lock()
|
|
sess, ok := sessions[oldName]
|
|
if !ok {
|
|
mu.Unlock()
|
|
return fmt.Errorf("session not found: %s", oldName)
|
|
}
|
|
if _, exists := sessions[newName]; exists {
|
|
mu.Unlock()
|
|
return fmt.Errorf("session already exists: %s", newName)
|
|
}
|
|
for _, ag := range sess.Agents() {
|
|
if ag.IsRunning() {
|
|
mu.Unlock()
|
|
return fmt.Errorf("cannot rename while agent is running")
|
|
}
|
|
}
|
|
sess.SetName(newName)
|
|
sessions[newName] = sess
|
|
delete(sessions, oldName)
|
|
mu.Unlock()
|
|
|
|
PersistSession(newName)
|
|
|
|
pkgLog.Info("renamed session %s -> %s", oldName, newName)
|
|
PublishEvent("session."+sess.ID+".rename", oldName+" "+newName)
|
|
return nil
|
|
}
|
|
|
|
// Shutdown interrupts all sessions, waits for idle, and closes them.
|
|
func Shutdown() {
|
|
InterruptAll()
|
|
waitIdle(100*time.Millisecond, 5*time.Second)
|
|
|
|
mu.Lock()
|
|
names := make([]string, 0, len(sessions))
|
|
for name := range sessions {
|
|
names = append(names, name)
|
|
}
|
|
mu.Unlock()
|
|
|
|
for _, name := range names {
|
|
mu.Lock()
|
|
sess := sessions[name]
|
|
delete(sessions, name)
|
|
mu.Unlock()
|
|
if sess != nil {
|
|
if sess.Cancel != nil {
|
|
sess.Cancel()
|
|
}
|
|
sess.Close()
|
|
if sess.Keeper != nil {
|
|
sess.Keeper.Close()
|
|
} else if sess.Proc != nil {
|
|
sess.Proc.Close()
|
|
}
|
|
pkgLog.Info("shutdown session %s", name)
|
|
}
|
|
}
|
|
}
|
|
|
|
// InterruptAll interrupts every active session.
|
|
func InterruptAll() {
|
|
mu.RLock()
|
|
defer mu.RUnlock()
|
|
for _, sess := range sessions {
|
|
for _, ag := range sess.Agents() {
|
|
ag.Interrupt(agent.ErrInterrupted)
|
|
}
|
|
}
|
|
}
|
|
|
|
func waitIdle(poll, timeout time.Duration) {
|
|
deadline := time.Now().Add(timeout)
|
|
for time.Now().Before(deadline) {
|
|
allIdle := true
|
|
mu.RLock()
|
|
for _, sess := range sessions {
|
|
for _, ag := range sess.Agents() {
|
|
if ag.State() != "idle" {
|
|
allIdle = false
|
|
break
|
|
}
|
|
}
|
|
if !allIdle {
|
|
break
|
|
}
|
|
}
|
|
mu.RUnlock()
|
|
if allIdle {
|
|
return
|
|
}
|
|
time.Sleep(poll)
|
|
}
|
|
pkgLog.Warn("shutdown: timed out waiting for idle")
|
|
}
|
|
|
|
// BuildIndex builds the session index content.
|
|
func BuildIndex() []byte {
|
|
var sb strings.Builder
|
|
mu.RLock()
|
|
defer mu.RUnlock()
|
|
|
|
type entry struct {
|
|
name string
|
|
sess *Session
|
|
paused bool
|
|
}
|
|
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 {
|
|
if entries[i].paused != entries[j].paused {
|
|
return !entries[i].paused
|
|
}
|
|
return entries[i].name < entries[j].name
|
|
})
|
|
|
|
for _, e := range entries {
|
|
sess := e.sess
|
|
sess.mu.RLock()
|
|
connected := sess.toolsConn != nil && !sess.paused
|
|
sess.mu.RUnlock()
|
|
var cwd, be, model string
|
|
var agentList []string
|
|
for _, ag := range sess.Agents() {
|
|
if cwd == "" {
|
|
cwd = ag.Cwd()
|
|
}
|
|
if be == "" {
|
|
be = ag.BackendName()
|
|
model = ag.ModelName()
|
|
}
|
|
agentList = append(agentList, fmt.Sprintf("%s:%s:%s", ag.Name(), ag.ID(), ag.State()))
|
|
}
|
|
agents := strings.Join(agentList, ";")
|
|
fmt.Fprintf(&sb, "%s\t%s\t%t\t%t\t%s\t%s\t%s\t%s\n",
|
|
e.name, sess.ID, e.paused, connected, cwd, be, model, agents)
|
|
}
|
|
return []byte(sb.String())
|
|
}
|
|
|
|
// BuildAgentIndex builds the agent index for a session.
|
|
func BuildAgentIndex(sess *Session) []byte {
|
|
var sb strings.Builder
|
|
for _, ag := range sess.Agents() {
|
|
fmt.Fprintf(&sb, "%s\t%s\t%s\t%s\t%s\n",
|
|
ag.Name(), ag.ID(), ag.Cwd(), ag.BackendName(), ag.ModelName())
|
|
}
|
|
return []byte(sb.String())
|
|
}
|
|
|
|
// CreateEmpty creates an empty session (no agent) and registers it.
|
|
func CreateEmpty(name, remote string) (*Session, error) {
|
|
sessID := util.NewUUID()
|
|
if name == "" {
|
|
name = sessID
|
|
if idx := strings.IndexByte(sessID, '-'); idx > 0 {
|
|
name = sessID[:idx]
|
|
}
|
|
}
|
|
|
|
mu.RLock()
|
|
_, exists := sessions[name]
|
|
mu.RUnlock()
|
|
if exists {
|
|
return nil, fmt.Errorf("session already exists: %s", name)
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(serverCtx)
|
|
sess := NewEmpty(sessID, ctx, cancel)
|
|
sess.SetName(name)
|
|
sess.Remote = remote
|
|
|
|
cwd, _ := os.Getwd()
|
|
infra, err := SetupToolServer(ToolServerConfig{
|
|
Ctx: ctx,
|
|
CWD: cwd,
|
|
RemoteTarget: remote,
|
|
SessionID: sessID,
|
|
Yolo: pkgYolo,
|
|
})
|
|
if err != nil {
|
|
cancel()
|
|
return nil, fmt.Errorf("setup tool server: %w", err)
|
|
}
|
|
sess.Proc = infra.Proc
|
|
sess.Keeper = infra.Keeper
|
|
sess.SetToolsConn(infra.ToolsConn)
|
|
|
|
// Start bypass approval loop (if notify function is configured)
|
|
if pkgBypassNotify != nil {
|
|
sess.StartBypassLoop(pkgBypassNotify)
|
|
}
|
|
|
|
mu.Lock()
|
|
sessions[name] = sess
|
|
mu.Unlock()
|
|
|
|
pkgLog.Info("created empty session %s (id=%s)", name, sessID)
|
|
PublishEvent("session."+sessID+".new", name)
|
|
return sess, nil
|
|
}
|
|
|
|
// NextUname generates a unique agent name (12 hex chars).
|
|
func NextUname() string {
|
|
b := make([]byte, 6)
|
|
rand.Read(b)
|
|
return fmt.Sprintf("%06x%06x", b[:3], b[3:6])
|
|
}
|
|
|
|
// --- Shared agent construction ---
|
|
|
|
// AgentParams describes the parameters for building an agent within a session.
|
|
type AgentParams struct {
|
|
Profile string
|
|
ID string // uname; if empty, generate one
|
|
Name string // display name; optional
|
|
CWD string // required
|
|
Backend string // backend name; empty = config default
|
|
Model string // model override; empty = config default
|
|
Remote string // remote target override
|
|
ParentID string // immutable ID of the agent that spawned this agent
|
|
History *agent.History // nil = fresh agent
|
|
}
|
|
|
|
// buildAgent constructs a fully-wired agent within a session.
|
|
// Used by both CreateAgent and restore.
|
|
func buildAgent(sess *Session, p AgentParams) (*agent.Agent, error) {
|
|
if p.ID == "" {
|
|
p.ID = NextUname()
|
|
}
|
|
if p.CWD == "" {
|
|
return nil, fmt.Errorf("cwd is required")
|
|
}
|
|
|
|
cfg := agent.LoadConfig(pkgAgentsDir, p.Profile, nil)
|
|
sessID := sess.ID
|
|
|
|
// Set up tool server — reuse existing session infrastructure
|
|
var reuseFrom *InfraConfig
|
|
if sess.Proc != nil {
|
|
reuseFrom = &InfraConfig{
|
|
Proc: sess.Proc,
|
|
Keeper: sess.Keeper,
|
|
ToolsConn: sess.ToolsConn(),
|
|
}
|
|
}
|
|
|
|
remote := p.Remote
|
|
if remote == "" {
|
|
remote = sess.Remote
|
|
}
|
|
|
|
infra, err := SetupToolServer(ToolServerConfig{
|
|
Ctx: sess.Ctx,
|
|
CWD: p.CWD,
|
|
RemoteTarget: remote,
|
|
SessionID: sessID,
|
|
Yolo: pkgYolo,
|
|
ReuseFrom: reuseFrom,
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("setup tool server: %w", err)
|
|
}
|
|
|
|
// First agent in a session sets the shared infra
|
|
if sess.Proc == nil {
|
|
sess.Proc = infra.Proc
|
|
sess.Keeper = infra.Keeper
|
|
sess.SetToolsConn(infra.ToolsConn)
|
|
}
|
|
|
|
// Build prompt layers
|
|
layers := BuildPromptLayers(cfg, p.CWD, sessID, p.ID, infra.Platform, infra.IsGitRepo, "")
|
|
|
|
// Load tools
|
|
LoadAutoLoadTools(cfg, infra.ToolsConn, sessID, p.ID, func(f string, a ...any) {
|
|
pkgLog.Error("session %s agent %s: "+f, append([]any{sessID, p.ID}, a...)...)
|
|
})
|
|
|
|
// Build runtime
|
|
env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + p.ID}
|
|
rt := agent.BuildRuntime(cfg, infra.ToolsConn, p.CWD, env, layers.SystemPrompt, layers.EnvBlock)
|
|
|
|
// Create backend
|
|
beName := p.Backend
|
|
if beName == "" {
|
|
beName = cfg.Backend
|
|
}
|
|
be, err := backend.NewWithName(beName)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create backend: %w", err)
|
|
}
|
|
rt.Backend = be
|
|
|
|
if p.Model != "" {
|
|
rt.Backend.SetModel(p.Model)
|
|
} else if cfg.Model != "" {
|
|
rt.Backend.SetModel(cfg.Model)
|
|
}
|
|
|
|
ag := agent.NewAgent(agent.AgentParams{
|
|
History: p.History,
|
|
Runtime: rt,
|
|
Profile: p.Profile,
|
|
AgentsDir: pkgAgentsDir,
|
|
ID: p.ID,
|
|
Cwd: p.CWD,
|
|
ParentID: p.ParentID,
|
|
SystemPrompt: layers.SystemPrompt,
|
|
EnvBlock: layers.EnvBlock,
|
|
NewToolServer: infra.NewToolServer,
|
|
NewBackend: backend.NewWithName,
|
|
Log: pkgSink.NewLogger("agent"),
|
|
SessionID: sessID,
|
|
Save: sess.saveSession,
|
|
Flush: sess.flushSave,
|
|
})
|
|
|
|
if p.Name != "" {
|
|
ag.SetName(p.Name)
|
|
}
|
|
return ag, nil
|
|
}
|
|
|
|
// CreateAgent creates a new agent within an existing session.
|
|
// Args: [profile] [key=value...] where keys include name, backend, model, remote, cwd.
|
|
func CreateAgent(sessName string, args []string) (*agent.Agent, error) {
|
|
// Parse args into AgentParams.
|
|
p := AgentParams{Profile: "default"}
|
|
for _, arg := range args {
|
|
if k, v, ok := strings.Cut(arg, "="); ok {
|
|
switch k {
|
|
case "name":
|
|
p.Name = v
|
|
case "backend":
|
|
p.Backend = v
|
|
case "model":
|
|
p.Model = v
|
|
case "remote":
|
|
p.Remote = v
|
|
case "cwd":
|
|
p.CWD = v
|
|
}
|
|
} else if p.Profile == "default" {
|
|
p.Profile = arg
|
|
}
|
|
}
|
|
return CreateAgentWithParams(sessName, p)
|
|
}
|
|
|
|
// CreateAgentWithParams creates a new agent within an existing session using
|
|
// a pre-built AgentParams struct.
|
|
func CreateAgentWithParams(sessName string, p AgentParams) (*agent.Agent, error) {
|
|
sess := Lookup(sessName)
|
|
if sess == nil {
|
|
return nil, fmt.Errorf("session not found: %s", sessName)
|
|
}
|
|
|
|
if p.Profile == "" {
|
|
p.Profile = "default"
|
|
}
|
|
|
|
if p.CWD == "" {
|
|
return nil, fmt.Errorf("cwd= is required when creating an agent")
|
|
}
|
|
|
|
ag, err := buildAgent(sess, p)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
sess.AddAgent(ag)
|
|
pkgLog.Info("created agent %s in session %s (profile=%s)", ag.Name(), sessName, p.Profile)
|
|
return ag, nil
|
|
}
|