session: export immutable fields, extract buildAgent helper
- Export ID, Uname, Ctx, Cancel, Proc, Keeper, Remote as struct fields (set-once at creation, not raced during agent loop) - Keep mutex-protected accessors only for Name, ToolsConn, IsPaused which are read/written concurrently - Extract buildAgent() shared helper used by both CreateAgent and restoreMultiAgentSession — eliminates ~100 lines of duplicated orchestration (tool server setup, prompt layers, runtime, backend) - Paused sessions skip tool server entirely during restore (minimal agent with just history) - Update fs/session_node.go delegators for exported fields - 1638 → 1523 lines (-115)
This commit is contained in:
Ollie Agent
parent
010ad83944
commit
12894027dd
|
|
@ -140,7 +140,7 @@ func requestGenerate(ctx HandlerCtx, data []byte) ([]byte, error) {
|
|||
func readAliases(_ HandlerCtx) ([]byte, error) {
|
||||
var sb strings.Builder
|
||||
for name, sess := range session.Sessions() {
|
||||
sid := sess.ID()
|
||||
sid := sess.ID
|
||||
fmt.Fprintf(&sb, "%s\tsession/%s\n", sid, name)
|
||||
for _, ag := range sess.Agents() {
|
||||
fmt.Fprintf(&sb, "%s\tsession/%s/agent/%s\n", ag.ID(), name, ag.Name())
|
||||
|
|
@ -257,8 +257,8 @@ func sessionBindings(ctx HandlerCtx) ([]Binding, error) {
|
|||
n := name // capture
|
||||
out = append(out, Binding{
|
||||
Name: sess.Name(),
|
||||
Aliases: []string{sess.ID()},
|
||||
UID: sess.Uname(),
|
||||
Aliases: []string{sess.ID},
|
||||
UID: sess.Uname,
|
||||
GID: "agent",
|
||||
Applier: func(c HandlerCtx) HandlerCtx {
|
||||
c.Context = node.Ctx()
|
||||
|
|
|
|||
|
|
@ -82,7 +82,7 @@ func NewRoot(cfg Config) *Tree {
|
|||
// Wire up the event handler so agent events flow to the AgentLog
|
||||
if ag := r.Session.FindAgent(ra.ID); ag != nil {
|
||||
ag.SetOutput(NewEventHandler(al))
|
||||
wireAgentStateEvents(r.Session.ID(), ag)
|
||||
wireAgentStateEvents(r.Session.ID, ag)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -33,18 +33,18 @@ func NewSessionNode(core *session.Session) *SessionNode {
|
|||
|
||||
// --- Delegation to Core ---
|
||||
|
||||
func (s *SessionNode) ID() string { return s.Core.ID() }
|
||||
func (s *SessionNode) ID() string { return s.Core.ID }
|
||||
func (s *SessionNode) Name() string { return s.Core.Name() }
|
||||
func (s *SessionNode) SetName(n string) { s.Core.SetName(n) }
|
||||
func (s *SessionNode) Uname() string { return s.Core.Uname() }
|
||||
func (s *SessionNode) Cancel() { s.Core.Cancel() }
|
||||
func (s *SessionNode) Uname() string { return s.Core.Uname }
|
||||
func (s *SessionNode) Cancel() { if s.Core.Cancel != nil { s.Core.Cancel() } }
|
||||
func (s *SessionNode) IsPaused() bool { return s.Core.IsPaused() }
|
||||
func (s *SessionNode) IsConnected() bool { return s.Core.IsConnected() }
|
||||
|
||||
func (s *SessionNode) Pause() error {
|
||||
err := s.Core.Pause()
|
||||
if err == nil {
|
||||
session.PublishEvent("session."+s.Core.ID()+".pause", "")
|
||||
session.PublishEvent("session."+s.Core.ID+".pause", "")
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
|
@ -52,7 +52,7 @@ func (s *SessionNode) Pause() error {
|
|||
func (s *SessionNode) Resume() error {
|
||||
err := s.Core.Resume()
|
||||
if err == nil {
|
||||
session.PublishEvent("session."+s.Core.ID()+".resume", "")
|
||||
session.PublishEvent("session."+s.Core.ID+".resume", "")
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
|
@ -62,7 +62,7 @@ func (s *SessionNode) LoadTool(name string, ag *agent.Agent) error {
|
|||
}
|
||||
|
||||
func (s *SessionNode) Ctx() context.Context {
|
||||
return s.Core.Ctx()
|
||||
return s.Core.Ctx
|
||||
}
|
||||
|
||||
func (s *SessionNode) ToolsConn() *toolsrv.Conn {
|
||||
|
|
|
|||
|
|
@ -11,12 +11,11 @@ import (
|
|||
|
||||
"ollie/agent"
|
||||
"ollie/backend"
|
||||
"ollie/toolsrv"
|
||||
)
|
||||
|
||||
// --- Persistence ---
|
||||
|
||||
// PersistedSession is the on-disk format for a saved session with multiple agents.
|
||||
// PersistedSession is the on-disk format for a saved session.
|
||||
type PersistedSession struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
|
|
@ -28,8 +27,8 @@ type PersistedSession struct {
|
|||
|
||||
// PersistedAgent is per-agent state within a session.
|
||||
type PersistedAgent struct {
|
||||
ID string `json:"id"` // uname
|
||||
Name string `json:"name"` // display name
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Profile string `json:"profile"`
|
||||
Backend string `json:"backend"`
|
||||
Model string `json:"model"`
|
||||
|
|
@ -64,21 +63,21 @@ func PersistSession(name string) error {
|
|||
}
|
||||
|
||||
ps := PersistedSession{
|
||||
ID: sess.ID(),
|
||||
ID: sess.ID,
|
||||
Name: sess.Name(),
|
||||
Remote: sess.Remote(),
|
||||
Remote: sess.Remote,
|
||||
Paused: sess.IsPaused(),
|
||||
}
|
||||
|
||||
for _, ag := range agents {
|
||||
pa := PersistedAgent{
|
||||
ID: ag.ID(),
|
||||
Name: ag.Name(),
|
||||
Profile: ag.Profile(),
|
||||
Backend: ag.BackendName(),
|
||||
Model: ag.ModelName(),
|
||||
CWD: ag.Cwd(),
|
||||
Messages: ag.Messages(),
|
||||
ID: ag.ID(),
|
||||
Name: ag.Name(),
|
||||
Profile: ag.Profile(),
|
||||
Backend: ag.BackendName(),
|
||||
Model: ag.ModelName(),
|
||||
CWD: ag.Cwd(),
|
||||
Messages: ag.Messages(),
|
||||
TaskState: ag.TaskState(),
|
||||
}
|
||||
if usage := ag.Usage(); usage != nil {
|
||||
|
|
@ -96,7 +95,7 @@ func PersistSession(name string) error {
|
|||
|
||||
dir := ActiveSessionsDir()
|
||||
os.MkdirAll(dir, 0700)
|
||||
path := filepath.Join(dir, sess.ID()+".json")
|
||||
path := filepath.Join(dir, sess.ID+".json")
|
||||
|
||||
data, err := json.MarshalIndent(ps, "", " ")
|
||||
if err != nil {
|
||||
|
|
@ -126,7 +125,6 @@ type RestoredSession struct {
|
|||
}
|
||||
|
||||
// RestoreAllSessions restores all persisted sessions.
|
||||
// Returns the restored sessions for the fs layer to wrap with SessionNodes.
|
||||
func RestoreAllSessions() []*RestoredSession {
|
||||
dir := ActiveSessionsDir()
|
||||
entries, err := os.ReadDir(dir)
|
||||
|
|
@ -140,12 +138,11 @@ func RestoreAllSessions() []*RestoredSession {
|
|||
files = append(files, filepath.Join(dir, e.Name()))
|
||||
}
|
||||
}
|
||||
|
||||
if len(files) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
var mu sync.Mutex
|
||||
var rmu sync.Mutex
|
||||
var results []*RestoredSession
|
||||
var wg sync.WaitGroup
|
||||
for _, path := range files {
|
||||
|
|
@ -154,13 +151,13 @@ func RestoreAllSessions() []*RestoredSession {
|
|||
defer wg.Done()
|
||||
rs, err := restoreSessionFromFile(path)
|
||||
if err != nil {
|
||||
pkgLog.Error("restore session %s: %v; removing from active sessions", path, err)
|
||||
pkgLog.Error("restore session %s: %v; removing", path, err)
|
||||
os.Remove(path)
|
||||
return
|
||||
}
|
||||
mu.Lock()
|
||||
rmu.Lock()
|
||||
results = append(results, rs)
|
||||
mu.Unlock()
|
||||
rmu.Unlock()
|
||||
}(path)
|
||||
}
|
||||
wg.Wait()
|
||||
|
|
@ -172,7 +169,6 @@ func restoreSessionFromFile(path string) (*RestoredSession, error) {
|
|||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var ps PersistedSession
|
||||
if err := json.Unmarshal(data, &ps); err != nil {
|
||||
return nil, err
|
||||
|
|
@ -184,20 +180,19 @@ func restoreSessionFromFile(path string) (*RestoredSession, error) {
|
|||
}
|
||||
|
||||
func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) {
|
||||
// Create empty session first
|
||||
ctx, cancel := context.WithCancel(serverCtx)
|
||||
sess := NewEmpty(ps.ID, ctx, cancel)
|
||||
sess.SetName(ps.Name)
|
||||
sess.SetRemote(ps.Remote)
|
||||
sess.Remote = ps.Remote
|
||||
if ps.Paused {
|
||||
sess.SetPaused(true)
|
||||
sess.mu.Lock()
|
||||
sess.paused = true
|
||||
sess.mu.Unlock()
|
||||
}
|
||||
|
||||
var restoredAgents []RestoredAgent
|
||||
|
||||
// Restore each agent with its own CWD
|
||||
for _, pa := range ps.Agents {
|
||||
// Get agent's CWD (fall back to session CWD for backwards compat, then pwd)
|
||||
cwd := pa.CWD
|
||||
if cwd == "" {
|
||||
cwd = ps.CWD
|
||||
|
|
@ -206,100 +201,47 @@ func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) {
|
|||
cwd, _ = os.Getwd()
|
||||
}
|
||||
|
||||
// For paused sessions, skip tool server setup
|
||||
var infra *SessionInfra
|
||||
if !ps.Paused {
|
||||
var err error
|
||||
// Reuse existing infra if available (from earlier agents)
|
||||
var reuseFrom *InfraConfig
|
||||
if sess.Proc() != nil {
|
||||
reuseFrom = &InfraConfig{
|
||||
Proc: sess.Proc(),
|
||||
Keeper: sess.Keeper(),
|
||||
ToolsConn: sess.ToolsConn(),
|
||||
}
|
||||
}
|
||||
infra, err = SetupToolServer(ToolServerConfig{
|
||||
Ctx: serverCtx,
|
||||
CWD: cwd,
|
||||
RemoteTarget: ps.Remote,
|
||||
Yolo: pkgYolo,
|
||||
ReuseFrom: reuseFrom,
|
||||
// For paused sessions, create a minimal agent without tool server
|
||||
if ps.Paused {
|
||||
history := agent.RestoreHistoryFromMessages(pa.Messages)
|
||||
ag := agent.NewAgent(agent.AgentCfg{
|
||||
History: history,
|
||||
Runtime: &agent.Runtime{},
|
||||
Profile: pa.Profile,
|
||||
AgentsDir: pkgAgentsDir,
|
||||
ID: pa.ID,
|
||||
Cwd: cwd,
|
||||
Log: pkgSink.NewLogger("agent"),
|
||||
SessionID: ps.ID,
|
||||
Save: sess.saveSession,
|
||||
Flush: sess.flushSave,
|
||||
})
|
||||
ag.SetName(pa.Name)
|
||||
sess.AddAgent(ag)
|
||||
restoredAgents = append(restoredAgents, RestoredAgent{
|
||||
ID: pa.ID,
|
||||
Remote: ps.Remote,
|
||||
Messages: pa.Messages,
|
||||
})
|
||||
if err != nil {
|
||||
pkgLog.Error("restore agent %s: setup tool server: %v", pa.ID, err)
|
||||
continue
|
||||
}
|
||||
// Store infra on session (first agent sets it up)
|
||||
if sess.Proc() == nil {
|
||||
sess.SetProc(infra.Proc)
|
||||
sess.SetKeeper(infra.Keeper)
|
||||
sess.SetToolsConn(infra.ToolsConn)
|
||||
}
|
||||
}
|
||||
|
||||
// Build prompt env
|
||||
promptEnv := agent.PromptEnv(cwd)
|
||||
if infra != nil {
|
||||
promptEnv = infra.PromptEnv
|
||||
}
|
||||
|
||||
cfg := agent.LoadConfig(pkgAgentsDir, pa.Profile, nil)
|
||||
|
||||
be, err := backend.NewWithName(pa.Backend)
|
||||
if err != nil {
|
||||
pkgLog.Error("restore agent %s: backend %s: %v", pa.ID, pa.Backend, err)
|
||||
continue
|
||||
}
|
||||
if pa.Model != "" {
|
||||
be.SetModel(pa.Model)
|
||||
}
|
||||
|
||||
// Build prompt layers
|
||||
layers := BuildPromptLayers(cfg, cwd, ps.ID, pa.ID, promptEnv, "")
|
||||
|
||||
// Load tools (only if we have a connection)
|
||||
if infra != nil && infra.ToolsConn != nil {
|
||||
LoadAutoLoadTools(cfg, infra.ToolsConn, ps.ID, pa.ID, func(f string, a ...any) {
|
||||
pkgLog.Error("session %s agent %s: "+f, append([]any{ps.ID, pa.ID}, a...)...)
|
||||
})
|
||||
}
|
||||
|
||||
env := []string{"OLLIE_SESSION_ID=" + ps.ID, "OLLIE_UNAME=" + pa.ID}
|
||||
env = append(env, promptEnv...)
|
||||
|
||||
var toolsConn *toolsrv.Conn
|
||||
var newToolServer func() *toolsrv.Conn
|
||||
if infra != nil {
|
||||
toolsConn = infra.ToolsConn
|
||||
newToolServer = infra.NewToolServer
|
||||
}
|
||||
|
||||
rt := agent.BuildRuntime(cfg, toolsConn, cwd, env, layers.SystemPrompt, layers.EnvBlock)
|
||||
rt.Backend = be
|
||||
|
||||
// Create history from messages
|
||||
// Active session — use shared buildAgent
|
||||
history := agent.RestoreHistoryFromMessages(pa.Messages)
|
||||
|
||||
ag := agent.NewAgent(agent.AgentCfg{
|
||||
History: history,
|
||||
Runtime: rt,
|
||||
Profile: pa.Profile,
|
||||
AgentsDir: pkgAgentsDir,
|
||||
ID: pa.ID,
|
||||
Cwd: cwd,
|
||||
SystemPrompt: layers.SystemPrompt,
|
||||
EnvBlock: layers.EnvBlock,
|
||||
PromptEnvExtra: promptEnv,
|
||||
NewToolServer: newToolServer,
|
||||
NewBackend: backend.NewWithName,
|
||||
Log: pkgSink.NewLogger("agent"),
|
||||
SessionID: ps.ID,
|
||||
Save: sess.saveSession,
|
||||
Flush: sess.flushSave,
|
||||
ag, err := buildAgent(sess, AgentParams{
|
||||
Profile: pa.Profile,
|
||||
ID: pa.ID,
|
||||
Name: pa.Name,
|
||||
CWD: cwd,
|
||||
Backend: pa.Backend,
|
||||
Model: pa.Model,
|
||||
Remote: ps.Remote,
|
||||
History: history,
|
||||
})
|
||||
ag.SetName(pa.Name)
|
||||
|
||||
if err != nil {
|
||||
pkgLog.Error("restore agent %s: %v", pa.ID, err)
|
||||
continue
|
||||
}
|
||||
sess.AddAgent(ag)
|
||||
restoredAgents = append(restoredAgents, RestoredAgent{
|
||||
ID: pa.ID,
|
||||
|
|
@ -313,14 +255,12 @@ func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) {
|
|||
return nil, fmt.Errorf("no agents restored")
|
||||
}
|
||||
|
||||
// Set uname to first agent's ID
|
||||
sess.SetUname(sess.Agents()[0].ID())
|
||||
|
||||
sess.Uname = sess.Agents()[0].ID()
|
||||
Register(sess.Name(), sess)
|
||||
|
||||
toolCount := 0
|
||||
if sess.ToolsConn() != nil {
|
||||
if infos, err := sess.ToolsConn().ListTools(); err == nil {
|
||||
if conn := sess.ToolsConn(); conn != nil {
|
||||
if infos, err := conn.ListTools(); err == nil {
|
||||
toolCount = len(infos)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ var (
|
|||
sessions = make(map[string]*Session)
|
||||
|
||||
// Package config, set via Init.
|
||||
serverCtx context.Context // server lifetime context; parent of all sessions
|
||||
serverCtx context.Context
|
||||
pkgLog *olog.Logger
|
||||
pkgSink *olog.Sink
|
||||
pkgAgentsDir string
|
||||
|
|
@ -40,7 +40,6 @@ type InitConfig struct {
|
|||
}
|
||||
|
||||
// Init initializes the session package with the given configuration.
|
||||
// Must be called before any other session functions.
|
||||
func Init(cfg InitConfig) {
|
||||
serverCtx = cfg.Ctx
|
||||
pkgLog = cfg.Log
|
||||
|
|
@ -69,7 +68,7 @@ func Lookup(nameOrID string) *Session {
|
|||
return s
|
||||
}
|
||||
for _, sess := range sessions {
|
||||
if sess.ID() == nameOrID {
|
||||
if sess.ID == nameOrID {
|
||||
return sess
|
||||
}
|
||||
}
|
||||
|
|
@ -94,26 +93,27 @@ func Kill(name string) error {
|
|||
delete(sessions, name)
|
||||
mu.Unlock()
|
||||
|
||||
// Remove from persistent store — kill is unconditional
|
||||
RemovePersistedSession(sess.ID())
|
||||
RemovePersistedSession(sess.ID)
|
||||
|
||||
var agentIDs []string
|
||||
for _, ag := range sess.Agents() {
|
||||
agentIDs = append(agentIDs, ag.ID())
|
||||
}
|
||||
sess.Close()
|
||||
sess.Cancel()
|
||||
if sess.Keeper() != nil {
|
||||
sess.Keeper().Close()
|
||||
} else if sess.Proc() != nil {
|
||||
sess.Proc().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+".agent."+agentID+".kill", "")
|
||||
}
|
||||
PublishEvent("session."+sess.ID()+".kill", "")
|
||||
PublishEvent("session."+sess.ID+".kill", "")
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -140,11 +140,10 @@ func Rename(oldName, newName string) error {
|
|||
delete(sessions, oldName)
|
||||
mu.Unlock()
|
||||
|
||||
// Persist immediately so the rename survives restart
|
||||
PersistSession(newName)
|
||||
|
||||
pkgLog.Info("renamed session %s -> %s", oldName, newName)
|
||||
PublishEvent("session."+sess.ID()+".rename", oldName+" "+newName)
|
||||
PublishEvent("session."+sess.ID+".rename", oldName+" "+newName)
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -166,12 +165,14 @@ func Shutdown() {
|
|||
delete(sessions, name)
|
||||
mu.Unlock()
|
||||
if sess != nil {
|
||||
sess.Cancel()
|
||||
if sess.Cancel != nil {
|
||||
sess.Cancel()
|
||||
}
|
||||
sess.Close()
|
||||
if sess.Keeper() != nil {
|
||||
sess.Keeper().Close()
|
||||
} else if sess.Proc() != nil {
|
||||
sess.Proc().Close()
|
||||
if sess.Keeper != nil {
|
||||
sess.Keeper.Close()
|
||||
} else if sess.Proc != nil {
|
||||
sess.Proc.Close()
|
||||
}
|
||||
pkgLog.Info("shutdown session %s", name)
|
||||
}
|
||||
|
|
@ -215,13 +216,11 @@ func waitIdle(poll, timeout time.Duration) {
|
|||
}
|
||||
|
||||
// BuildIndex builds the session index content.
|
||||
// Active sessions are listed first, paused sessions last.
|
||||
func BuildIndex() []byte {
|
||||
var sb strings.Builder
|
||||
mu.RLock()
|
||||
defer mu.RUnlock()
|
||||
|
||||
// Collect sessions and sort: active first (alphabetically), then paused (alphabetically)
|
||||
type entry struct {
|
||||
name string
|
||||
sess *Session
|
||||
|
|
@ -232,9 +231,8 @@ func BuildIndex() []byte {
|
|||
entries = append(entries, entry{name: name, sess: sess, paused: sess.IsPaused()})
|
||||
}
|
||||
sort.Slice(entries, func(i, j int) bool {
|
||||
// Active (not paused) comes before paused
|
||||
if entries[i].paused != entries[j].paused {
|
||||
return !entries[i].paused // active first
|
||||
return !entries[i].paused
|
||||
}
|
||||
return entries[i].name < entries[j].name
|
||||
})
|
||||
|
|
@ -256,13 +254,12 @@ func BuildIndex() []byte {
|
|||
}
|
||||
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)
|
||||
e.name, sess.ID, e.paused, connected, cwd, be, model, agents)
|
||||
}
|
||||
return []byte(sb.String())
|
||||
}
|
||||
|
||||
// BuildAgentIndex builds the agent index for a session.
|
||||
// Format: name<tab>id<tab>cwd<tab>backend<tab>model per line.
|
||||
func BuildAgentIndex(sess *Session) []byte {
|
||||
var sb strings.Builder
|
||||
for _, ag := range sess.Agents() {
|
||||
|
|
@ -272,8 +269,7 @@ func BuildAgentIndex(sess *Session) []byte {
|
|||
return []byte(sb.String())
|
||||
}
|
||||
|
||||
// CreateEmpty creates an empty session (no agent).
|
||||
// If remote is non-empty, spawns a remote tool server via SSH.
|
||||
// CreateEmpty creates an empty session (no agent) and registers it.
|
||||
func CreateEmpty(name, remote string) (*Session, error) {
|
||||
sessID := NewSessionID()
|
||||
if name == "" {
|
||||
|
|
@ -293,9 +289,8 @@ func CreateEmpty(name, remote string) (*Session, error) {
|
|||
ctx, cancel := context.WithCancel(serverCtx)
|
||||
sess := NewEmpty(sessID, ctx, cancel)
|
||||
sess.SetName(name)
|
||||
sess.SetRemote(remote)
|
||||
sess.Remote = remote
|
||||
|
||||
// Spawn tool server (local or remote)
|
||||
cwd, _ := os.Getwd()
|
||||
infra, err := SetupToolServer(ToolServerConfig{
|
||||
Ctx: ctx,
|
||||
|
|
@ -307,8 +302,8 @@ func CreateEmpty(name, remote string) (*Session, error) {
|
|||
cancel()
|
||||
return nil, fmt.Errorf("setup tool server: %w", err)
|
||||
}
|
||||
sess.SetProc(infra.Proc)
|
||||
sess.SetKeeper(infra.Keeper)
|
||||
sess.Proc = infra.Proc
|
||||
sess.Keeper = infra.Keeper
|
||||
sess.SetToolsConn(infra.ToolsConn)
|
||||
|
||||
mu.Lock()
|
||||
|
|
@ -327,76 +322,81 @@ func NextUname() string {
|
|||
return fmt.Sprintf("%06x%06x", b[:3], b[3:6])
|
||||
}
|
||||
|
||||
// CreateAgent creates a new agent within an existing session.
|
||||
// Args: [profile] [key=value...] where keys include name, backend, model, remote.
|
||||
func CreateAgent(sessName string, args []string) (*agent.Agent, error) {
|
||||
sess := Lookup(sessName)
|
||||
if sess == nil {
|
||||
return nil, fmt.Errorf("session not found: %s", sessName)
|
||||
// --- 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
|
||||
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")
|
||||
}
|
||||
|
||||
// Parse args
|
||||
profile := "default"
|
||||
var agentName, backendName, modelName, remoteTarget, cwdOverride string
|
||||
for _, arg := range args {
|
||||
if k, v, ok := strings.Cut(arg, "="); ok {
|
||||
switch k {
|
||||
case "name":
|
||||
agentName = v
|
||||
case "backend":
|
||||
backendName = v
|
||||
case "model":
|
||||
modelName = v
|
||||
case "remote":
|
||||
remoteTarget = v
|
||||
case "cwd":
|
||||
cwdOverride = v
|
||||
}
|
||||
} else if profile == "default" {
|
||||
profile = arg
|
||||
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(),
|
||||
}
|
||||
}
|
||||
|
||||
cfg := agent.LoadConfig(pkgAgentsDir, profile, nil)
|
||||
|
||||
uname := NextUname()
|
||||
sessID := sess.ID()
|
||||
if cwdOverride == "" {
|
||||
return nil, fmt.Errorf("cwd= is required when creating an agent")
|
||||
remote := p.Remote
|
||||
if remote == "" {
|
||||
remote = sess.Remote
|
||||
}
|
||||
cwd := cwdOverride
|
||||
|
||||
// Set up tool server - reuse existing session infrastructure
|
||||
infra, err := SetupToolServer(ToolServerConfig{
|
||||
Ctx: sess.Ctx(),
|
||||
CWD: cwd,
|
||||
RemoteTarget: remoteTarget,
|
||||
Ctx: sess.Ctx,
|
||||
CWD: p.CWD,
|
||||
RemoteTarget: remote,
|
||||
Yolo: pkgYolo,
|
||||
ReuseFrom: &InfraConfig{
|
||||
Proc: sess.Proc(),
|
||||
Keeper: sess.Keeper(),
|
||||
ToolsConn: sess.ToolsConn(),
|
||||
},
|
||||
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, cwd, sessID, uname, infra.PromptEnv, "")
|
||||
layers := BuildPromptLayers(cfg, p.CWD, sessID, p.ID, infra.PromptEnv, "")
|
||||
|
||||
// Load tools
|
||||
LoadAutoLoadTools(cfg, infra.ToolsConn, sessID, uname, func(f string, a ...any) {
|
||||
pkgLog.Error("session %s agent %s: "+f, append([]any{sessID, uname}, a...)...)
|
||||
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=" + uname}
|
||||
env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + p.ID}
|
||||
env = append(env, infra.PromptEnv...)
|
||||
rt := agent.BuildRuntime(cfg, infra.ToolsConn, cwd, env, layers.SystemPrompt, layers.EnvBlock)
|
||||
rt := agent.BuildRuntime(cfg, infra.ToolsConn, p.CWD, env, layers.SystemPrompt, layers.EnvBlock)
|
||||
|
||||
// Create backend - use override if provided, otherwise use config or env default
|
||||
beName := backendName
|
||||
// Create backend
|
||||
beName := p.Backend
|
||||
if beName == "" {
|
||||
beName = cfg.Backend
|
||||
}
|
||||
|
|
@ -406,21 +406,19 @@ func CreateAgent(sessName string, args []string) (*agent.Agent, error) {
|
|||
}
|
||||
rt.Backend = be
|
||||
|
||||
// Apply model override if provided
|
||||
if modelName != "" {
|
||||
rt.Backend.SetModel(modelName)
|
||||
if p.Model != "" {
|
||||
rt.Backend.SetModel(p.Model)
|
||||
} else if cfg.Model != "" {
|
||||
rt.Backend.SetModel(cfg.Model)
|
||||
}
|
||||
|
||||
// Create agent
|
||||
ag := agent.NewAgent(agent.AgentCfg{
|
||||
History: nil, // new agent starts with empty history
|
||||
History: p.History,
|
||||
Runtime: rt,
|
||||
Profile: profile,
|
||||
Profile: p.Profile,
|
||||
AgentsDir: pkgAgentsDir,
|
||||
ID: uname,
|
||||
Cwd: cwd,
|
||||
ID: p.ID,
|
||||
Cwd: p.CWD,
|
||||
SystemPrompt: layers.SystemPrompt,
|
||||
EnvBlock: layers.EnvBlock,
|
||||
PromptEnvExtra: infra.PromptEnv,
|
||||
|
|
@ -432,10 +430,51 @@ func CreateAgent(sessName string, args []string) (*agent.Agent, error) {
|
|||
Flush: sess.flushSave,
|
||||
})
|
||||
|
||||
sess.AddAgent(ag)
|
||||
if agentName != "" {
|
||||
ag.SetName(agentName)
|
||||
if p.Name != "" {
|
||||
ag.SetName(p.Name)
|
||||
}
|
||||
pkgLog.Info("created agent %s in session %s (profile=%s)", ag.Name(), sessName, profile)
|
||||
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) {
|
||||
sess := Lookup(sessName)
|
||||
if sess == nil {
|
||||
return nil, fmt.Errorf("session not found: %s", sessName)
|
||||
}
|
||||
|
||||
// Parse args
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -46,21 +46,24 @@ type Config struct {
|
|||
|
||||
// Session owns all session-level state: identity, agents, tool server lifecycle.
|
||||
type Session struct {
|
||||
// Identity
|
||||
id string // immutable UUID, set at creation
|
||||
name string // mutable friendly name, defaults to id prefix
|
||||
uname string // immutable user principal (numeric UID)
|
||||
// Identity (immutable after creation)
|
||||
ID string // UUID
|
||||
Uname string // user principal (numeric UID)
|
||||
|
||||
// Context lifecycle
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
Ctx context.Context
|
||||
Cancel context.CancelFunc
|
||||
|
||||
// Tool server lifecycle (set during setup/teardown, not raced)
|
||||
Proc *toolsrv.Process
|
||||
Keeper *toolsrv.ProcessKeeper
|
||||
Remote string
|
||||
|
||||
// Agent state
|
||||
agents []*agent.Agent
|
||||
log *olog.Logger
|
||||
sessionsDir string
|
||||
listHandlers map[string]func() []string
|
||||
remote string
|
||||
|
||||
// Environment
|
||||
envMu sync.RWMutex
|
||||
|
|
@ -68,18 +71,15 @@ type Session struct {
|
|||
plan []byte
|
||||
prevPrompt string
|
||||
|
||||
// Tool server lifecycle
|
||||
proc *toolsrv.Process // toolsrv subprocess
|
||||
keeper *toolsrv.ProcessKeeper // resilient process manager
|
||||
toolsConn *toolsrv.Conn // multiplexed RPC connection
|
||||
paused bool // true when paused
|
||||
// Mutable state protected by mu
|
||||
mu sync.RWMutex
|
||||
name string // friendly name
|
||||
toolsConn *toolsrv.Conn // multiplexed RPC connection
|
||||
paused bool
|
||||
|
||||
// Tool restrictions
|
||||
disallowTools map[string]struct{}
|
||||
|
||||
// Synchronization
|
||||
mu sync.RWMutex
|
||||
|
||||
// Autosave
|
||||
saveMu sync.Mutex
|
||||
saveDirty bool
|
||||
|
|
@ -110,6 +110,45 @@ func sweepStaleTmpDirs() {
|
|||
})
|
||||
}
|
||||
|
||||
// --- Mutable field accessors (only for fields that race) ---
|
||||
|
||||
// Name returns the mutable friendly name.
|
||||
func (s *Session) Name() string {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.name
|
||||
}
|
||||
|
||||
// SetName sets the friendly name.
|
||||
func (s *Session) SetName(name string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.name = name
|
||||
}
|
||||
|
||||
// ToolsConn returns the tool server connection.
|
||||
func (s *Session) ToolsConn() *toolsrv.Conn {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.toolsConn
|
||||
}
|
||||
|
||||
// SetToolsConn sets the tool server connection.
|
||||
func (s *Session) SetToolsConn(conn *toolsrv.Conn) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.toolsConn = conn
|
||||
}
|
||||
|
||||
// IsPaused returns true if the session is paused.
|
||||
func (s *Session) IsPaused() bool {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.paused
|
||||
}
|
||||
|
||||
// --- Constructor ---
|
||||
|
||||
// New creates a session with an owned agent from the given configuration.
|
||||
func New(cfg Config) *Session {
|
||||
sweepStaleTmpDirs()
|
||||
|
|
@ -142,13 +181,13 @@ func New(cfg Config) *Session {
|
|||
}
|
||||
|
||||
s := &Session{
|
||||
id: cfg.SessionID,
|
||||
ID: cfg.SessionID,
|
||||
name: name,
|
||||
Remote: cfg.Remote,
|
||||
env: make(map[string]string),
|
||||
agents: make([]*agent.Agent, 0, 1),
|
||||
log: log,
|
||||
sessionsDir: cfg.SessionsDir,
|
||||
remote: cfg.Remote,
|
||||
listHandlers: cfg.ListHandlers,
|
||||
}
|
||||
|
||||
|
|
@ -174,7 +213,7 @@ func New(cfg Config) *Session {
|
|||
})
|
||||
|
||||
s.agents = append(s.agents, ag)
|
||||
ag.SetSessionEnv(s.id)
|
||||
ag.SetSessionEnv(s.ID)
|
||||
return s
|
||||
}
|
||||
|
||||
|
|
@ -194,11 +233,10 @@ func (s *Session) Agents() []*agent.Agent { return s.agents }
|
|||
// AddAgent appends an agent to the session.
|
||||
func (s *Session) AddAgent(ag *agent.Agent) {
|
||||
s.agents = append(s.agents, ag)
|
||||
ag.SetSessionEnv(s.id)
|
||||
ag.SetSessionEnv(s.ID)
|
||||
}
|
||||
|
||||
// RemoveAgent removes an agent by ID and closes it.
|
||||
// Returns false if the agent was not found.
|
||||
func (s *Session) RemoveAgent(id string) bool {
|
||||
idx := -1
|
||||
for i, ag := range s.agents {
|
||||
|
|
@ -217,13 +255,13 @@ func (s *Session) RemoveAgent(id string) bool {
|
|||
|
||||
// Close releases resources for this session.
|
||||
func (s *Session) Close() {
|
||||
s.log.Debug("Close() session=%q", s.id)
|
||||
s.log.Debug("Close() session=%q", s.ID)
|
||||
s.flushSave()
|
||||
for _, ag := range s.agents {
|
||||
ag.Close()
|
||||
}
|
||||
if s.id != "" {
|
||||
os.RemoveAll(filepath.Join(ollieTmpDir(), s.id)) //nolint:errcheck
|
||||
if s.ID != "" {
|
||||
os.RemoveAll(filepath.Join(ollieTmpDir(), s.ID)) //nolint:errcheck
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -237,19 +275,9 @@ func (s *Session) SetEnv(key, value string) {
|
|||
}
|
||||
}
|
||||
|
||||
// Remote returns the remote target (e.g., SSH host) if any.
|
||||
func (s *Session) Remote() string {
|
||||
return s.remote
|
||||
}
|
||||
|
||||
// SetRemote sets the remote target for this session.
|
||||
func (s *Session) SetRemote(remote string) {
|
||||
s.remote = remote
|
||||
}
|
||||
|
||||
// SetSessionID renames the session.
|
||||
func (s *Session) SetSessionID(newID string) error {
|
||||
oldID := s.id
|
||||
oldID := s.ID
|
||||
if oldID == newID {
|
||||
return nil
|
||||
}
|
||||
|
|
@ -263,7 +291,7 @@ func (s *Session) SetSessionID(newID string) error {
|
|||
}
|
||||
}
|
||||
}
|
||||
s.id = newID
|
||||
s.ID = newID
|
||||
oldTemp := filepath.Join(ollieTmpDir(), oldID)
|
||||
newTemp := filepath.Join(ollieTmpDir(), newID)
|
||||
if _, err := os.Stat(oldTemp); err == nil {
|
||||
|
|
@ -275,6 +303,139 @@ func (s *Session) SetSessionID(newID string) error {
|
|||
return nil
|
||||
}
|
||||
|
||||
// IsConnected returns true if the tool server connection is alive.
|
||||
func (s *Session) IsConnected() bool {
|
||||
s.mu.RLock()
|
||||
paused := s.paused
|
||||
conn := s.toolsConn
|
||||
s.mu.RUnlock()
|
||||
if paused || conn == nil {
|
||||
return false
|
||||
}
|
||||
return conn.Ping() == nil
|
||||
}
|
||||
|
||||
// Pause stops the tool server to save resources.
|
||||
func (s *Session) Pause() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.paused {
|
||||
return fmt.Errorf("session already paused")
|
||||
}
|
||||
if s.Cancel != nil {
|
||||
s.Cancel()
|
||||
}
|
||||
if s.Keeper != nil {
|
||||
s.Keeper.Close()
|
||||
} else if s.Proc != nil {
|
||||
s.Proc.Close()
|
||||
}
|
||||
if s.toolsConn != nil {
|
||||
s.toolsConn.Close()
|
||||
s.toolsConn = nil
|
||||
}
|
||||
s.paused = true
|
||||
go s.saveSession()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Resume restarts the tool server after a Pause.
|
||||
func (s *Session) Resume() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if !s.paused {
|
||||
return fmt.Errorf("session not paused")
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(serverCtx)
|
||||
s.Ctx = ctx
|
||||
s.Cancel = cancel
|
||||
|
||||
if s.Keeper != nil {
|
||||
s.Keeper.SetContext(ctx)
|
||||
conn, err := s.Keeper.Dial()
|
||||
if err != nil {
|
||||
return fmt.Errorf("resume failed: %w", err)
|
||||
}
|
||||
s.toolsConn = conn
|
||||
} else {
|
||||
// Session was restored paused without infra — set up per-agent.
|
||||
for i, ag := range s.agents {
|
||||
cwd := ag.Cwd()
|
||||
if cwd == "" {
|
||||
cwd, _ = os.Getwd()
|
||||
}
|
||||
var reuseFrom *InfraConfig
|
||||
if s.Proc != nil {
|
||||
reuseFrom = &InfraConfig{
|
||||
Proc: s.Proc,
|
||||
Keeper: s.Keeper,
|
||||
ToolsConn: s.toolsConn,
|
||||
}
|
||||
}
|
||||
infra, err := SetupToolServer(ToolServerConfig{
|
||||
Ctx: ctx,
|
||||
CWD: cwd,
|
||||
RemoteTarget: s.Remote,
|
||||
Yolo: pkgYolo,
|
||||
ReuseFrom: reuseFrom,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("resume setup for agent %s failed: %w", ag.ID(), err)
|
||||
}
|
||||
if i == 0 {
|
||||
s.Proc = infra.Proc
|
||||
s.Keeper = infra.Keeper
|
||||
s.toolsConn = infra.ToolsConn
|
||||
}
|
||||
ag.SetToolServer(infra.NewToolServer, infra.ToolsConn)
|
||||
}
|
||||
}
|
||||
|
||||
s.paused = false
|
||||
go s.saveSession()
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- Tool loading ---
|
||||
|
||||
// LoadTool loads a tool into the tool server.
|
||||
func (s *Session) LoadTool(name string, ag *agent.Agent) error {
|
||||
name = strings.TrimSpace(name)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
s.mu.RLock()
|
||||
_, blocked := s.disallowTools[name]
|
||||
s.mu.RUnlock()
|
||||
if blocked {
|
||||
return fmt.Errorf("tool %q is disallowed for this session", name)
|
||||
}
|
||||
var conn *toolsrv.Conn
|
||||
if s.ToolsConn() != nil {
|
||||
conn = s.ToolsConn()
|
||||
} else if ag != nil {
|
||||
conn = ag.ToolServer()
|
||||
}
|
||||
return LoadToolOnConn(conn, name)
|
||||
}
|
||||
|
||||
// LoadToolOnConn loads a tool on the given connection.
|
||||
func LoadToolOnConn(conn *toolsrv.Conn, name string) error {
|
||||
if conn == nil {
|
||||
return nil
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
args, _ := json.Marshal(map[string]string{"name": name})
|
||||
if _, err := conn.CallTool(ctx, "tool_load", json.RawMessage(args)); err != nil {
|
||||
return fmt.Errorf("remote load: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- Autosave ---
|
||||
|
||||
func (s *Session) activeSessionPath(id, suffix string) string {
|
||||
return filepath.Join(s.sessionsDir, "active", id+suffix)
|
||||
}
|
||||
|
|
@ -297,10 +458,9 @@ func (s *Session) flushSave() {
|
|||
s.saveTimer = nil
|
||||
}
|
||||
s.saveMu.Unlock()
|
||||
if !dirty || s.id == "" || s.sessionsDir == "" {
|
||||
if !dirty || s.ID == "" || s.sessionsDir == "" {
|
||||
return
|
||||
}
|
||||
// Use PersistSession for consistent format (writes Name field correctly)
|
||||
if err := PersistSession(s.Name()); err != nil {
|
||||
s.log.Error("session save: %v", err)
|
||||
}
|
||||
|
|
@ -311,260 +471,6 @@ func (s *Session) Save() error {
|
|||
return PersistSession(s.Name())
|
||||
}
|
||||
|
||||
// --- Identity accessors ---
|
||||
|
||||
// ID returns the immutable session UUID.
|
||||
func (s *Session) ID() string { return s.id }
|
||||
|
||||
// Name returns the mutable friendly name.
|
||||
func (s *Session) Name() string {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.name
|
||||
}
|
||||
|
||||
// SetName sets the friendly name.
|
||||
func (s *Session) SetName(name string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.name = name
|
||||
}
|
||||
|
||||
// Uname returns the immutable user principal.
|
||||
func (s *Session) Uname() string { return s.uname }
|
||||
|
||||
// SetUname sets the user principal (should only be called during init).
|
||||
func (s *Session) SetUname(uname string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.uname = uname
|
||||
}
|
||||
|
||||
// --- Context lifecycle ---
|
||||
|
||||
// Ctx returns the session context.
|
||||
func (s *Session) Ctx() context.Context { return s.ctx }
|
||||
|
||||
// Cancel cancels the session context.
|
||||
func (s *Session) Cancel() {
|
||||
if s.cancel != nil {
|
||||
s.cancel()
|
||||
}
|
||||
}
|
||||
|
||||
// SetContext sets the session context and cancel func.
|
||||
func (s *Session) SetContext(ctx context.Context, cancel context.CancelFunc) {
|
||||
s.ctx = ctx
|
||||
s.cancel = cancel
|
||||
}
|
||||
|
||||
// --- Tool server lifecycle ---
|
||||
|
||||
// Proc returns the tool server process.
|
||||
func (s *Session) Proc() *toolsrv.Process {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.proc
|
||||
}
|
||||
|
||||
// SetProc sets the tool server process.
|
||||
func (s *Session) SetProc(proc *toolsrv.Process) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.proc = proc
|
||||
}
|
||||
|
||||
// Keeper returns the process keeper.
|
||||
func (s *Session) Keeper() *toolsrv.ProcessKeeper {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.keeper
|
||||
}
|
||||
|
||||
// SetKeeper sets the process keeper.
|
||||
func (s *Session) SetKeeper(keeper *toolsrv.ProcessKeeper) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.keeper = keeper
|
||||
}
|
||||
|
||||
// ToolsConn returns the tool server connection.
|
||||
func (s *Session) ToolsConn() *toolsrv.Conn {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.toolsConn
|
||||
}
|
||||
|
||||
// SetToolsConn sets the tool server connection.
|
||||
func (s *Session) SetToolsConn(conn *toolsrv.Conn) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.toolsConn = conn
|
||||
}
|
||||
|
||||
// IsPaused returns true if the session is paused.
|
||||
func (s *Session) IsPaused() bool {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.paused
|
||||
}
|
||||
|
||||
// SetPaused sets the paused state directly (used during restore).
|
||||
func (s *Session) SetPaused(paused bool) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.paused = paused
|
||||
}
|
||||
|
||||
// IsConnected returns true if the tool server connection is alive.
|
||||
// Returns false if the session is paused.
|
||||
func (s *Session) IsConnected() bool {
|
||||
s.mu.RLock()
|
||||
paused := s.paused
|
||||
conn := s.toolsConn
|
||||
s.mu.RUnlock()
|
||||
if paused || conn == nil {
|
||||
return false
|
||||
}
|
||||
return conn.Ping() == nil
|
||||
}
|
||||
|
||||
// Pause stops the tool server to save resources.
|
||||
// It cancels the session context (stopping all agent operations),
|
||||
// kills the tool server process, and marks the session as paused.
|
||||
func (s *Session) Pause() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.paused {
|
||||
return fmt.Errorf("session already paused")
|
||||
}
|
||||
// Cancel the session context to stop all agent operations.
|
||||
if s.cancel != nil {
|
||||
s.cancel()
|
||||
}
|
||||
// Close the tool server process.
|
||||
if s.keeper != nil {
|
||||
s.keeper.Close()
|
||||
} else if s.proc != nil {
|
||||
s.proc.Close()
|
||||
}
|
||||
// Close any existing connection.
|
||||
if s.toolsConn != nil {
|
||||
s.toolsConn.Close()
|
||||
s.toolsConn = nil
|
||||
}
|
||||
s.paused = true
|
||||
// Persist the paused state.
|
||||
go s.saveSession()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Resume restarts the tool server after a Pause.
|
||||
// It creates a fresh session context and respawns the tool server.
|
||||
// For sessions restored in paused state (no keeper), it sets up the tool server from scratch.
|
||||
func (s *Session) Resume() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if !s.paused {
|
||||
return fmt.Errorf("session not paused")
|
||||
}
|
||||
|
||||
// Create a fresh context since the old one was cancelled on pause.
|
||||
ctx, cancel := context.WithCancel(serverCtx)
|
||||
s.ctx = ctx
|
||||
s.cancel = cancel
|
||||
|
||||
// If we have a keeper, use it to respawn. Otherwise set up from scratch.
|
||||
if s.keeper != nil {
|
||||
s.keeper.SetContext(ctx)
|
||||
conn, err := s.keeper.Dial()
|
||||
if err != nil {
|
||||
return fmt.Errorf("resume failed: %w", err)
|
||||
}
|
||||
s.toolsConn = conn
|
||||
} else {
|
||||
// Session was restored paused without infra - set up per-agent.
|
||||
for i, ag := range s.agents {
|
||||
cwd := ag.Cwd()
|
||||
if cwd == "" {
|
||||
cwd, _ = os.Getwd()
|
||||
}
|
||||
|
||||
// First agent spawns the tool server, others reuse
|
||||
var reuseFrom *InfraConfig
|
||||
if s.proc != nil {
|
||||
reuseFrom = &InfraConfig{
|
||||
Proc: s.proc,
|
||||
Keeper: s.keeper,
|
||||
ToolsConn: s.toolsConn,
|
||||
}
|
||||
}
|
||||
|
||||
infra, err := SetupToolServer(ToolServerConfig{
|
||||
Ctx: ctx,
|
||||
CWD: cwd,
|
||||
RemoteTarget: s.remote,
|
||||
Yolo: pkgYolo,
|
||||
ReuseFrom: reuseFrom,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("resume setup for agent %s failed: %w", ag.ID(), err)
|
||||
}
|
||||
|
||||
// First agent sets session infra
|
||||
if i == 0 {
|
||||
s.proc = infra.Proc
|
||||
s.keeper = infra.Keeper
|
||||
s.toolsConn = infra.ToolsConn
|
||||
}
|
||||
|
||||
ag.SetToolServer(infra.NewToolServer, infra.ToolsConn)
|
||||
}
|
||||
}
|
||||
|
||||
s.paused = false
|
||||
// Persist the resumed state.
|
||||
go s.saveSession()
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- Tool loading ---
|
||||
|
||||
// LoadTool loads a tool into the tool server.
|
||||
func (s *Session) LoadTool(name string, ag *agent.Agent) error {
|
||||
name = strings.TrimSpace(name)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
s.mu.RLock()
|
||||
_, blocked := s.disallowTools[name]
|
||||
s.mu.RUnlock()
|
||||
if blocked {
|
||||
return fmt.Errorf("tool %q is disallowed for this session", name)
|
||||
}
|
||||
var conn *toolsrv.Conn
|
||||
if s.toolsConn != nil {
|
||||
conn = s.toolsConn
|
||||
} else if ag != nil {
|
||||
conn = ag.ToolServer()
|
||||
}
|
||||
return LoadToolOnConn(conn, name)
|
||||
}
|
||||
|
||||
// LoadToolOnConn loads a tool on the given connection.
|
||||
func LoadToolOnConn(conn *toolsrv.Conn, name string) error {
|
||||
if conn == nil {
|
||||
return nil
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
args, _ := json.Marshal(map[string]string{"name": name})
|
||||
if _, err := conn.CallTool(ctx, "tool_load", json.RawMessage(args)); err != nil {
|
||||
return fmt.Errorf("remote load: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- Constructor for empty session ---
|
||||
|
||||
// NewEmpty creates an empty session without an agent.
|
||||
|
|
@ -575,10 +481,10 @@ func NewEmpty(id string, ctx context.Context, cancel context.CancelFunc) *Sessio
|
|||
}
|
||||
log := pkgSink.NewLogger("session")
|
||||
return &Session{
|
||||
id: id,
|
||||
ID: id,
|
||||
name: name,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
Ctx: ctx,
|
||||
Cancel: cancel,
|
||||
env: make(map[string]string),
|
||||
agents: make([]*agent.Agent, 0),
|
||||
log: log,
|
||||
|
|
|
|||
Loading…
Reference in New Issue