From 2aa7db84458df5cefd06d7230bfec71525e621a3 Mon Sep 17 00:00:00 2001 From: Levi Neely Date: Tue, 4 Aug 2026 13:52:28 +0200 Subject: [PATCH] refactor: split fs/ and session/ packages session/ now owns all session business logic: - Manager: session lifecycle (create, kill, rename, lookup) - CreateAgent: agent creation within sessions - SetupToolServer, BuildPromptLayers, LoadAutoLoadTools - PublishEvent for SSE notifications - NextUname generation fs/ is now a thin 9P presentation layer: - SessionNode wraps session.Session, adds AgentLogs and models cache - RootState caches SessionNodes (GetOrCreateNode pattern) - Handlers delegate to session.Manager - persist.go uses session package for restore logic --- fs/handlers.go | 107 ++++--- fs/lifecycle.go | 653 ++++---------------------------------- fs/newroot.go | 34 +- fs/persist.go | 116 +++---- fs/spec.go | 2 +- fs/types.go | 444 ++++++++++---------------- session/event.go | 74 +++++ {fs => session}/format.go | 41 +-- session/manager.go | 383 ++++++++++++++++++++++ session/session.go | 266 +++++++++++++++- {fs => session}/setup.go | 4 +- 11 files changed, 1059 insertions(+), 1065 deletions(-) create mode 100644 session/event.go rename {fs => session}/format.go (57%) create mode 100644 session/manager.go rename {fs => session}/setup.go (98%) diff --git a/fs/handlers.go b/fs/handlers.go index feb0488..9432bc1 100644 --- a/fs/handlers.go +++ b/fs/handlers.go @@ -13,6 +13,7 @@ import ( "ollie/backend" "ollie/elevate" "ollie/fsedsl" + "ollie/session" "ollie/toolsrv" ) @@ -105,7 +106,7 @@ func readEventwait(_ HandlerCtx) ([]byte, error) { } func blockEventwait(_ HandlerCtx, cctx context.Context, base string) ([]byte, string, error) { - return WaitEvent(cctx, base) + return session.WaitEvent(cctx, base) } func requestComplete(_ HandlerCtx, data []byte) ([]byte, error) { @@ -152,20 +153,12 @@ func requestRoute(_ HandlerCtx, data []byte) ([]byte, error) { func readAliases(_ HandlerCtx) ([]byte, error) { rs := globalRootState - rs.Mu.RLock() - defer rs.Mu.RUnlock() - var sb strings.Builder - for _, sess := range rs.Sessions { - sess.mu.RLock() - name := sess.name - sid := sess.id - sess.mu.RUnlock() + for name, sess := range rs.Manager.Sessions() { + sid := sess.ID() fmt.Fprintf(&sb, "%s\tsession/%s\n", sid, name) - if sess.Core != nil { - for _, ag := range sess.Core.Agents() { - fmt.Fprintf(&sb, "%s\tsession/%s/agent/%s\n", ag.ID(), name, ag.Name()) - } + for _, ag := range sess.Agents() { + fmt.Fprintf(&sb, "%s\tsession/%s/agent/%s\n", ag.ID(), name, ag.Name()) } } return []byte(sb.String()), nil @@ -238,41 +231,66 @@ func writeElevateRequest(ctx HandlerCtx, data []byte) error { } func readSessionIdx(_ HandlerCtx) ([]byte, error) { - return BuildIndex(globalRootState), nil + return globalRootState.Manager.BuildIndex(), nil } func requestSessionNew(_ HandlerCtx, data []byte) ([]byte, error) { args := strings.Fields(string(data)) - name, err := Create(globalRootState, args) + name := "" + for _, arg := range args { + if k, v, ok := strings.Cut(arg, "="); ok && k == "name" { + name = v + } + } + sess, err := globalRootState.Manager.CreateEmpty(name) if err != nil { return nil, err } - return []byte(name + "\n"), nil + return []byte(sess.Name() + "\n"), nil } // sessionBindings returns bindings for all sessions. // The Blueprint pattern clones the structure and binds ctx.Session. func sessionBindings(_ HandlerCtx) ([]Binding, error) { rs := globalRootState - rs.Mu.RLock() - defer rs.Mu.RUnlock() + sessions := rs.Manager.Sessions() - out := make([]Binding, 0, len(rs.Sessions)) - for _, sess := range rs.Sessions { - s := sess // capture + out := make([]Binding, 0, len(sessions)) + for name, sess := range sessions { + // Get or create cached SessionNode wrapper + node := rs.GetOrCreateNode(sess) + n := name // capture out = append(out, Binding{ - Name: s.Name(), - Aliases: []string{s.ID()}, - UID: s.Uname(), + Name: sess.Name(), + Aliases: []string{sess.ID()}, + UID: sess.Uname(), GID: "agent", Applier: func(c HandlerCtx) HandlerCtx { - c.Session = s - c.Remove = func() { Kill(rs, s.Name()) } - c.Rename = func(newName string) error { return Rename(rs, s.Name(), newName) } + c.Session = node + c.Remove = func() { + rs.Manager.Kill(n) + rs.RemoveNode(n) + } + c.Rename = func(newName string) error { + if err := rs.Manager.Rename(n, newName); err != nil { + return err + } + rs.RenameNode(n, newName) + return nil + } return c }, - Remove: func() { Kill(rs, s.Name()) }, - Rename: func(newName string) error { return Rename(rs, s.Name(), newName) }, + Remove: func() { + rs.Manager.Kill(n) + rs.RemoveNode(n) + }, + Rename: func(newName string) error { + if err := rs.Manager.Rename(n, newName); err != nil { + return err + } + rs.RenameNode(n, newName) + return nil + }, }) } return out, nil @@ -354,7 +372,7 @@ func readSessionElevate(ctx HandlerCtx) ([]byte, error) { if ctx.Elevate == nil { return nil, fmt.Errorf("elevate broker not available") } - p := ctx.Elevate.SessionPolicy(ctx.Session.id) + p := ctx.Elevate.SessionPolicy(ctx.Session.ID()) return p.Marshal() } @@ -366,17 +384,20 @@ func writeSessionElevate(ctx HandlerCtx, data []byte) error { if err := elevate.ParsePolicy(data, &p); err != nil { return err } - ctx.Elevate.SetSessionPolicy(ctx.Session.id, p) + ctx.Elevate.SetSessionPolicy(ctx.Session.ID(), p) return nil } func requestAgentNew(ctx HandlerCtx, data []byte) ([]byte, error) { args := strings.Fields(string(data)) - if err := CreateAgent(globalRootState, ctx.Session.Name(), args); err != nil { + ag, err := globalRootState.Manager.CreateAgent(ctx.Session.Name(), args) + if err != nil { return nil, fmt.Errorf("create agent: %w", err) } - agents := ctx.Session.Core.Agents() - return []byte(agents[len(agents)-1].ID() + "\n"), nil + // Create AgentLog for this agent + al := NewAgentLog("") + ctx.Session.SetAgentLog(ag.ID(), al) + return []byte(ag.ID() + "\n"), nil } // agentBindings returns bindings for all agents in a session. @@ -427,7 +448,7 @@ func writeAgentPrompt(ctx HandlerCtx, data []byte) error { ctx.AgentLog.prevPrompt = []byte(input) ctx.AgentLog.mu.Unlock() go func() { - ctx.Agent.Submit(ctx.Session.SessionCtx, input) + ctx.Agent.Submit(ctx.Session.Ctx(), input) ctx.AgentLog.EnsureTrailingNewline() }() return nil @@ -478,10 +499,10 @@ func streamAgentChat(ctx HandlerCtx, cctx context.Context, base string) ([]byte, } func readAgentConnection(ctx HandlerCtx) ([]byte, error) { - if ctx.Session.toolsConn == nil { + if ctx.Session.ToolsConn() == nil { return nil, fmt.Errorf("agent tool server unavailable") } - if _, err := ctx.Session.toolsConn.ListTools(); err != nil { + if _, err := ctx.Session.ToolsConn().ListTools(); err != nil { return nil, err } return []byte("connected\n"), nil @@ -494,7 +515,7 @@ func readAgentState(ctx HandlerCtx) ([]byte, error) { func blockAgentStateWait(ctx HandlerCtx, cctx context.Context, base string) ([]byte, string, error) { waitCtx, cancel := context.WithCancel(cctx) defer cancel() - context.AfterFunc(ctx.Session.SessionCtx, cancel) + context.AfterFunc(ctx.Session.Ctx(), cancel) if base == "" { base = ctx.Agent.State() } @@ -545,7 +566,7 @@ func readAgentCfg(ctx HandlerCtx) ([]byte, error) { defer ctx.AgentLog.mu.RUnlock() p := ctx.Agent.GenParams() var sb strings.Builder - fmt.Fprintf(&sb, "name=%s\n", ctx.Session.id) + fmt.Fprintf(&sb, "name=%s\n", ctx.Session.ID()) fmt.Fprintf(&sb, "backend=%s\n", ctx.Agent.BackendName()) fmt.Fprintf(&sb, "model=%s\n", ctx.Agent.ModelName()) fmt.Fprintf(&sb, "id=%s\n", ctx.Agent.ID()) @@ -596,7 +617,7 @@ func writeAgentCtl(ctx HandlerCtx, data []byte) error { case "kill": if ctx.Session.Core != nil { ctx.Session.Core.RemoveAgent(ctx.Agent.ID()) - PublishEvent("session."+ctx.Session.ID()+".agent."+ctx.Agent.ID()+".kill", "") + session.PublishEvent("session."+ctx.Session.ID()+".agent."+ctx.Agent.ID()+".kill", "") } case "stop": ctx.Agent.Interrupt(agent.ErrInterrupted) @@ -675,10 +696,10 @@ func readAgentTail(_ HandlerCtx) ([]byte, error) { } func readAgentTools(ctx HandlerCtx) ([]byte, error) { - if ctx.Session.toolsConn == nil { + if ctx.Session.ToolsConn() == nil { return []byte("(no tool server)\n"), nil } - loaded, err := ctx.Session.toolsConn.ListTools() + loaded, err := ctx.Session.ToolsConn().ListTools() if err != nil { return []byte(fmt.Sprintf("(error: %v)\n", err)), nil } @@ -694,7 +715,7 @@ func readAgentTools(ctx HandlerCtx) ([]byte, error) { } func writeAgentTools(ctx HandlerCtx, data []byte) error { - return ctx.Session.loadTool(string(data), ctx.Agent) + return ctx.Session.LoadTool(string(data), ctx.Agent) } func readAgentName(ctx HandlerCtx) ([]byte, error) { diff --git a/fs/lifecycle.go b/fs/lifecycle.go index 176e1b2..f324826 100644 --- a/fs/lifecycle.go +++ b/fs/lifecycle.go @@ -1,626 +1,81 @@ package fs import ( - "context" - "crypto/rand" - "fmt" - "os" - "sort" - "strings" "sync" - "time" - "github.com/simonfxr/pubsub" - - "ollie/agent" - "ollie/backend" - "ollie/paths" - coresession "ollie/session" - "ollie/toolsrv" + "ollie/session" ) // globalRootState is set by NewRoot so handler functions can access -// the session map without going through ctx.Root (which is nil at -// build time when closures are created). +// the session manager without going through ctx.Root. var globalRootState *RootState -// RootState is the package-level state stored in root.Data. -// Exported for use by fshandlers package. +// RootState holds the root-level state for the 9P filesystem. type RootState struct { - Cfg Config - Mu sync.RWMutex - Sessions map[string]*Session + Cfg Config + Manager *session.Manager + + // SessionNodes caches the fs-layer wrappers for sessions. + // Keyed by session name (mutable), must be kept in sync with Manager. + nodesMu sync.RWMutex + SessionNodes map[string]*SessionNode +} + +// GetOrCreateNode returns the SessionNode for the given session, creating one if needed. +func (rs *RootState) GetOrCreateNode(sess *session.Session) *SessionNode { + rs.nodesMu.RLock() + node, ok := rs.SessionNodes[sess.Name()] + rs.nodesMu.RUnlock() + if ok { + return node + } + + rs.nodesMu.Lock() + defer rs.nodesMu.Unlock() + // Double-check after acquiring write lock + if node, ok := rs.SessionNodes[sess.Name()]; ok { + return node + } + node = NewSessionNode(sess) + rs.SessionNodes[sess.Name()] = node + return node +} + +// RemoveNode removes the SessionNode for a session. +func (rs *RootState) RemoveNode(name string) { + rs.nodesMu.Lock() + delete(rs.SessionNodes, name) + rs.nodesMu.Unlock() +} + +// RenameNode updates the SessionNodes map when a session is renamed. +func (rs *RootState) RenameNode(oldName, newName string) { + rs.nodesMu.Lock() + defer rs.nodesMu.Unlock() + if node, ok := rs.SessionNodes[oldName]; ok { + delete(rs.SessionNodes, oldName) + rs.SessionNodes[newName] = node + } } func state(root *Tree) *RootState { return root.Data.(*RootState) } -var eventBus = pubsub.NewBus() - -// Event represents a published event with topic and payload. -type Event struct { - Topic string - Payload string -} - -// PublishEvent publishes an event to a specific topic. -// Topic hierarchy: -// session.{sid}.new/kill/pause/resume/rename -// session.{sid}.agent.{aid}.new/kill/state -// -// Wildcard routing publishes to all ancestor wildcards: -// session.{sid}.agent.{aid}.state -> also publishes to: -// session.{sid}.agent.{aid}.* -// session.{sid}.agent.* -// session.{sid}.* -// session.* -// * -func PublishEvent(topic, payload string) { - ev := Event{Topic: topic, Payload: payload} - eventBus.Publish(topic, ev) - - // Publish to hierarchical wildcards - parts := strings.Split(topic, ".") - for i := len(parts) - 1; i >= 1; i-- { - wildcard := strings.Join(parts[:i], ".") + ".*" - eventBus.Publish(wildcard, ev) - } - eventBus.Publish("*", ev) -} - -// SubscribeEvents returns a channel that receives all events matching the topic. -// Use "*" for all events, "session.*" for all session events, etc. -func SubscribeEvents(ctx context.Context, topic string) <-chan Event { - ch := make(chan Event, 64) - sub := eventBus.SubscribeChan(topic, ch, pubsub.KeepOpen) - go func() { - <-ctx.Done() - eventBus.Unsubscribe(sub) - close(ch) - }() - return ch -} - -// WaitEvent blocks until an event is available. Returns the event as "topic payload\n". -// For backwards compatibility with the old eventwait interface. -func WaitEvent(ctx context.Context, _ string) ([]byte, string, error) { - ch := SubscribeEvents(ctx, "*") - select { - case ev, ok := <-ch: - if !ok { - return nil, "", nil - } - line := ev.Topic - if ev.Payload != "" { - line += " " + ev.Payload - } - return []byte(line + "\n"), "", nil - case <-ctx.Done(): - return nil, "", nil - } -} - -// BuildIndex builds the session/idx content. -// -// Format: one line per session (not per agent) -// Fields (tab-separated): -// name, id, paused, connected, cwd, backend, model, agents -// Where agents is a semicolon-separated list of: agentName:agentId:state -// -// Example: -// myproject\tf60604bf-...\tfalse\ttrue\t/home/user/proj\tkiro\tclaude-opus-4.5\tdev:7c36c9e4:idle;test:abc123:thinking - -func BuildIndex(rs *RootState) []byte { - var sb strings.Builder - rs.Mu.RLock() - ids := make([]string, 0, len(rs.Sessions)) - for id := range rs.Sessions { - ids = append(ids, id) - } - sort.Strings(ids) - for _, id := range ids { - sess := rs.Sessions[id] - sess.mu.RLock() - name := sess.name - sid := sess.id - paused := sess.paused - - // Check connected status - connected := false - if conn := sess.toolsConn; conn != nil { - if pinger, ok := conn.(interface{ Ping() error }); ok { - connected = pinger.Ping() == nil - } else { - connected = true - } - } - - var cwd, be, model string - var agentList []string - if sess.Core != nil { - cwd = sess.Core.CWD() - for _, ag := range sess.Core.Agents() { - if be == "" { - be = ag.BackendName() - model = ag.ModelName() - } - // Format: name:id:state - 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", name, sid, paused, connected, cwd, be, model, agents) - sess.mu.RUnlock() - } - rs.Mu.RUnlock() - return []byte(sb.String()) -} - -// Kill terminates a session and removes it from the tree. -func Kill(rs *RootState, id string) { - rs.Mu.Lock() - sess := rs.Sessions[id] - if sess == nil { - rs.Mu.Unlock() - rs.Cfg.Log.Error("Kill: session %q not found in sessions map", id) - return - } - delete(rs.Sessions, id) - rs.Mu.Unlock() - - var agentNames []string - if sess.Core != nil { - for _, ag := range sess.Core.Agents() { - agentNames = append(agentNames, ag.ID()) - } - sess.Core.Close() - } - sess.Cancel() - if sess.keeper != nil { - sess.keeper.Close() - } else if sess.proc != nil { - sess.proc.Close() - } - removePersistedAgent(rs, sess.id) - rs.Cfg.Log.Info("killed session %s", id) - for _, agentID := range agentNames { - PublishEvent("session."+sess.id+".agent."+agentID+".kill", "") - } - PublishEvent("session."+sess.id+".kill", "") -} - -// Rename renames a session by updating its Name and the sessions map. -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 != nil { - for _, ag := range sess.Core.Agents() { - if ag.IsRunning() { - rs.Mu.Unlock() - return fmt.Errorf("cannot rename while agent is running") - } - } - } - sess.name = new - rs.Sessions[new] = sess - delete(rs.Sessions, old) - rs.Mu.Unlock() - // Persist the updated session and agent display names immediately. The - // agent's own autosave does not know the fs-level session display name. - persistSession(rs, new) - rs.Cfg.Log.Info("renamed session %s -> %s", old, new) - PublishEvent("session."+sess.id+".rename", old+" "+new) - return nil -} - // Shutdown interrupts all sessions, waits for idle, persists, and closes. func Shutdown(root *Tree) { - rs := root.Data.(*RootState) - 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() - if sess.Core != nil { - sess.Core.Close() - } - if sess.keeper != nil { - sess.keeper.Close() - } else if sess.proc != nil { - sess.proc.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 != nil { - for _, ag := range sess.Core.Agents() { - if ag.State() != "idle" { - allIdle = false - break - } - } - } - if !allIdle { - 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") + rs := state(root) + rs.Manager.Shutdown() } // InterruptAll interrupts every active session. func InterruptAll(root *Tree) { - rs := root.Data.(*RootState) - rs.Mu.RLock() - defer rs.Mu.RUnlock() - for _, sess := range rs.Sessions { - if sess.Core != nil { - for _, ag := range sess.Core.Agents() { - ag.Interrupt(agent.ErrInterrupted) - } - } - } + rs := state(root) + rs.Manager.InterruptAll() } // Lookup returns the session for the given ID, or nil. -func Lookup(root *Tree, id string) *Session { - rs := root.Data.(*RootState) - rs.Mu.RLock() - defer rs.Mu.RUnlock() - // Try by display name first (map key). - if s := rs.Sessions[id]; s != nil { - return s - } - // Fall back to immutable session ID. - for _, sess := range rs.Sessions { - if sess.id == id { - return sess - } - } - return nil -} - -func nextUname(rs *RootState) string { - b := make([]byte, 6) - rand.Read(b) - return fmt.Sprintf("%06x%06x", b[:3], b[3:6]) -} - -// Create creates a new empty session from key=value args and registers it. -func Create(rs *RootState, args []string) (string, error) { - name := "" - for _, arg := range args { - k, v, ok := strings.Cut(arg, "=") - if !ok { - return "", fmt.Errorf("invalid option %q (expected key=value)", arg) - } - if v == "" { - continue - } - switch k { - case "name": - name = v - default: - return "", fmt.Errorf("unknown option %q (valid: name)", k) - } - } - - sessID := coresession.NewSessionID() - if name == "" { - name = sessID - if idx := strings.IndexByte(sessID, '-'); idx > 0 { - name = sessID[:idx] - } - } - - rs.Mu.RLock() - _, exists := rs.Sessions[name] - rs.Mu.RUnlock() - if exists { - return "", fmt.Errorf("session already exists: %s", name) - } - - sessionCtx, sessionCancel := context.WithCancel(rs.Cfg.Ctx) - sess := NewSession(sessID, nil, sessionCtx, sessionCancel) - sess.SetName(name) - - // Spawn the tool server so the session is immediately operational. - cwd, _ := os.Getwd() - var dialOpts []toolsrv.Option - if rs.Cfg.Yolo { - dialOpts = append(dialOpts, toolsrv.WithYolo()) - } - proc, err := toolsrv.Spawn(sessionCtx, cwd, dialOpts...) - if err != nil { - sessionCancel() - return "", fmt.Errorf("spawn tool server: %w", err) - } - sess.proc = proc - sess.keeper = toolsrv.NewProcessKeeper(sessionCtx, proc, func(ctx context.Context) (*toolsrv.Process, error) { - return toolsrv.Spawn(ctx, cwd, dialOpts...) - }) - conn, err := toolsrv.Dial(sessionCtx, toolsrv.LocalAddr(proc.Socket)) - if err != nil { - sessionCancel() - proc.Close() - return "", fmt.Errorf("dial tool server: %w", err) - } - sess.toolsConn = conn - - rs.Mu.Lock() - rs.Sessions[name] = sess - rs.Mu.Unlock() - - rs.Cfg.Log.Info("created empty session %s (id=%s)", name, sessID) - PublishEvent("session."+sessID+".new", name) - return name, nil -} - -// CreateFromRoot is the public API for callers that have a *Tree. -func CreateFromRoot(root *Tree, args []string) (string, error) { - return Create(root.Data.(*RootState), args) -} - -// CreateAgent creates and configures an agent within a session. -func CreateAgent(rs *RootState, sessName string, args []string) error { - backendOverride := "" - modelOverride := "" - profile := "" - name := "" - cwd := "" - remoteTarget := "" - systemPromptOverride := "" - for _, arg := range args { - k, v, ok := strings.Cut(arg, "=") - if !ok { - return fmt.Errorf("invalid option %q (expected key=value)", arg) - } - if v == "" { - continue - } - switch k { - case "name": - name = v - case "backend": - backendOverride = v - case "model": - modelOverride = v - case "agent": - profile = v - case "cwd": - cwd = v - case "remote": - remoteTarget = v - case "system_prompt": - systemPromptOverride = v - default: - return fmt.Errorf("unknown option %q (valid: name, backend, model, agent, cwd, remote, system_prompt)", k) - } - } - - rs.Mu.RLock() - sess, ok := rs.Sessions[sessName] - rs.Mu.RUnlock() - if !ok { - sessID := coresession.NewSessionID() - sessionCtx, sessionCancel := context.WithCancel(rs.Cfg.Ctx) - sess = NewSession(sessID, nil, sessionCtx, sessionCancel) - sess.SetName(sessName) - } - - if profile == "" { - if v := os.Getenv("OLLIE_DEFAULT_AGENT"); v != "" { - profile = v - } else { - profile = "default" - } - } - - cwd = paths.ExpandHome(os.ExpandEnv(cwd)) - if cwd == "" { - return fmt.Errorf("cwd is required (set cwd= when creating an agent)") - } - if remoteTarget == "" { - if info, err := os.Stat(cwd); err != nil { - return fmt.Errorf("cwd %q does not exist", cwd) - } else if !info.IsDir() { - return fmt.Errorf("cwd %q is not a directory", cwd) - } - } - - uname := nextUname(rs) - sessID := sess.id - - cfg := LoadAgentConfig(rs.Cfg.AgentsDir, profile, nil) - if cfg != nil { - if backendOverride == "" && cfg.Backend != "" { - backendOverride = cfg.Backend - } - if modelOverride == "" && cfg.Model != "" { - modelOverride = cfg.Model - } - } - - be, err := backend.NewWithName(backendOverride) - if err != nil { - return fmt.Errorf("backend: %w", err) - } - - if modelOverride == "" { - modelOverride = os.Getenv("OLLIE_MODEL") - } - if modelOverride != "" { - be.SetModel(modelOverride) - } - - if err := rs.Cfg.MkdirAll(rs.Cfg.SessionsDir, PermMkdirPrivate); err != nil { - return fmt.Errorf("sessions dir: %w", err) - } - - // Set up tool server infrastructure. - // If session already has a proc from Create(), reuse it. - var reuseFrom *InfraConfig - if sess.proc != nil && remoteTarget == "" { - reuseFrom = &InfraConfig{ - Proc: sess.proc, - Keeper: sess.keeper, - ToolsConn: sess.toolsConn, - } - } - - infra, err := SetupToolServer(ToolServerConfig{ - Ctx: rs.Cfg.Ctx, - CWD: cwd, - RemoteTarget: remoteTarget, - Yolo: rs.Cfg.Yolo, - ReuseFrom: reuseFrom, - }) - if err != nil { - return err - } - - // Update session's keeper if we spawned new infrastructure. - if reuseFrom == nil { - sess.keeper = infra.Keeper - } - - // Build prompt layers. - layers := BuildPromptLayers(cfg, cwd, sessID, uname, infra.PromptEnv, systemPromptOverride) - - toolSrv := infra.ToolsConn - if toolSrv == nil { - // Tool server is vital infrastructure — retry once with a fresh spawn. - rs.Cfg.Log.Error("session %s: initial tool server dial failed, retrying...", sessName) - if sess.keeper != nil { - if conn, err := sess.keeper.Dial(); err == nil { - toolSrv = conn - } - } - if toolSrv == nil { - return fmt.Errorf("tool server unavailable after retry (session infrastructure broken — try again)") - } - } - - // Load configured tools before building the runtime. - LoadAutoLoadTools(cfg, toolSrv, sessID, uname, func(f string, a ...any) { - rs.Cfg.Log.Error("session %s: "+f, append([]any{sessName}, a...)...) - }) - - env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname} - env = append(env, infra.PromptEnv...) - rt := agent.BuildRuntime(cfg, toolSrv, cwd, env, layers.SystemPrompt, layers.OpModel, layers.EnvBlock) - sess.toolsConn = toolSrv - - al := NewAgentLog(remoteTarget) - - if sess.Core == nil { - 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, - Runtime: rt, - NewToolServer: infra.NewToolServer, - PromptEnvExtra: infra.PromptEnv, - BaseLayers: layers.BaseLayers(), - Log: rs.Cfg.Sink.NewLogger("core"), - Output: NewEventHandler(al), - ReadPlanStep: func() string { - return "" - }, - }) - sess.mu.Lock() - sess.Core = core - sess.uname = uname - sess.setAgentLogLocked(uname, al) - sess.proc = infra.Proc - sess.mu.Unlock() - } else { - rt.Backend = be - saveFn, flushFn := sess.Core.SaveFuncs() - ag := agent.NewAgent(agent.AgentCfg{ - History: nil, - Runtime: rt, - Profile: profile, - AgentsDir: rs.Cfg.AgentsDir, - AgentID: uname, - CWD: cwd, - BaseLayers: layers.BaseLayers(), - PromptEnvExtra: infra.PromptEnv, - NewToolServer: infra.NewToolServer, - NewBackend: backend.NewWithName, - Output: NewEventHandler(al), - Log: rs.Cfg.Sink.NewLogger("core"), - AuditLog: rs.Cfg.Sink.NewLogger("core").Sub("audit"), - SessionID: sessID, - ReadPlanStep: func() string { - return "" - }, - SaveSession: saveFn, - FlushSave: flushFn, - }) - sess.Core.AddAgent(ag) - // Wire up state change events - capture sessID for hierarchical topic - sid := sessID - ag.SetOnStateChange(func(agentID, state string) { - PublishEvent("session."+sid+".agent."+agentID+".state", state) - }) - sess.mu.Lock() - sess.setAgentLogLocked(uname, al) - sess.mu.Unlock() - } - - rs.Mu.Lock() - if _, exists := rs.Sessions[sessName]; !exists { - rs.Sessions[sessName] = sess - } - rs.Mu.Unlock() - - ag := sess.Core.FindAgent(uname) - ag.Submit(sess.SessionCtx, "/agent "+profile) - if name != "" { - ag.SetName(name) - } - - rs.Cfg.Log.Info("created agent in session %s (backend=%s model=%s agent=%s)", - sessName, ag.BackendName(), ag.ModelName(), ag.Name()) - PublishEvent("session."+sessID+".agent."+ag.ID()+".new", ag.Name()) - return nil +func Lookup(root *Tree, id string) *session.Session { + rs := state(root) + return rs.Manager.Lookup(id) } diff --git a/fs/newroot.go b/fs/newroot.go index 0daaa89..11b4306 100644 --- a/fs/newroot.go +++ b/fs/newroot.go @@ -8,8 +8,8 @@ import ( "ollie/elevate" "ollie/fsedsl" olog "ollie/log" + "ollie/session" "ollie/skills" - coresession "ollie/session" ) // Config carries all runtime dependencies for the filesystem tree. @@ -21,7 +21,6 @@ type Config struct { AgentsDir string SessionsDir string MkdirAll func(string, os.FileMode) error - NewCore func(sessionID, profile, cwd string) (*coresession.Session, error) Yolo bool SkillsRegistry *skills.Registry ModelCache *ModelCache @@ -44,26 +43,37 @@ func NewRoot(cfg Config) *Tree { treeSpec.UID = treeSpecUID // root node: owner is system user treeSpec.GID = treeSpecGID // root node: group is "agent" + // Create session manager + mgr := session.NewManager(session.ManagerConfig{ + Ctx: cfg.Ctx, + Log: cfg.Log, + Sink: cfg.Sink, + AgentsDir: cfg.AgentsDir, + SessionsDir: cfg.SessionsDir, + MkdirAll: cfg.MkdirAll, + Yolo: cfg.Yolo, + }) + // Build the root context with all available deps. ctx := HandlerCtx{ - Root: nil, // set after tree is built - Log: cfg.Log, - Skills: cfg.SkillsRegistry, - Models: cfg.ModelCache, - Elevate: cfg.ElevateBroker, - Shutdown: cfg.Shutdown, + Root: nil, // set after tree is built + Log: cfg.Log, + Skills: cfg.SkillsRegistry, + Models: cfg.ModelCache, + Elevate: cfg.ElevateBroker, + Shutdown: cfg.Shutdown, Invalidate: cfg.Invalidate, } // Build the tree from the spec. - // Validation happens inside BuildTree. tree := fsedsl.BuildTree(treeSpec, ctx) - ctx.Root = tree // now handlers can access root tree via ctx.Root + ctx.Root = tree // Initialize root state. rs := &RootState{ - Cfg: cfg, - Sessions: make(map[string]*Session), + Cfg: cfg, + Manager: mgr, + SessionNodes: make(map[string]*SessionNode), } tree.Data = rs globalRootState = rs diff --git a/fs/persist.go b/fs/persist.go index c5b5d23..8a5c6f2 100644 --- a/fs/persist.go +++ b/fs/persist.go @@ -2,7 +2,6 @@ package fs import ( "context" - "fmt" "os" "path/filepath" "strings" @@ -10,7 +9,7 @@ import ( "ollie/agent" "ollie/backend" - coresession "ollie/session" + "ollie/session" ) // --- Persistence --- @@ -19,40 +18,31 @@ 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 { +func persistSession(rs *RootState, name string) { + sess := rs.Manager.Lookup(name) + if sess == nil { return } - if sess.Core == nil || sess.Core.AgentCount() == 0 { - // Empty session — no agent to persist. Silently skip. + if sess.AgentCount() == 0 { 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) + path := filepath.Join(dir, sess.ID()+".json") + ag := sess.Agents()[0] + if err := ag.SaveFull(path, sess.ID(), sess.Name(), sess.CWD(), ""); err != nil { + rs.Cfg.Log.Error("persist session %s: %v", name, err) } } -func removePersistedAgent(rs *RootState, id string) { +func removePersistedSession(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) + for name := range rs.Manager.Sessions() { + persistSession(rs, name) } } @@ -103,10 +93,6 @@ func restoreSession(rs *RootState, ps *agent.PersistedAgent) error { 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" @@ -114,7 +100,7 @@ func restoreSession(rs *RootState, ps *agent.PersistedAgent) error { sessID := ps.ID remoteTarget := ps.Remote - cfg := LoadAgentConfig(rs.Cfg.AgentsDir, profile, nil) + cfg := session.LoadAgentConfig(rs.Cfg.AgentsDir, profile, nil) backendName := ps.Backend if backendName == "" && cfg != nil && cfg.Backend != "" { @@ -122,7 +108,7 @@ func restoreSession(rs *RootState, ps *agent.PersistedAgent) error { } be, err := backend.NewWithName(backendName) if err != nil { - return fmt.Errorf("backend: %w", err) + return err } modelName := ps.Model if modelName == "" && cfg != nil && cfg.Model != "" { @@ -137,11 +123,11 @@ func restoreSession(rs *RootState, ps *agent.PersistedAgent) error { uname := ps.Uname if uname == "" { - uname = nextUname(rs) + uname = session.NextUname() } - // Set up tool server infrastructure. - infra, err := SetupToolServer(ToolServerConfig{ + // Set up tool server infrastructure + infra, err := session.SetupToolServer(session.ToolServerConfig{ Ctx: rs.Cfg.Ctx, CWD: cwd, RemoteTarget: remoteTarget, @@ -151,11 +137,11 @@ func restoreSession(rs *RootState, ps *agent.PersistedAgent) error { return err } - // Build prompt layers. - layers := BuildPromptLayers(cfg, cwd, sessID, uname, infra.PromptEnv, "") + // Build prompt layers + layers := session.BuildPromptLayers(cfg, cwd, sessID, uname, infra.PromptEnv, "") - // Load configured tools before building the runtime. - LoadAutoLoadTools(cfg, infra.ToolsConn, sessID, uname, func(f string, a ...any) { + // Load configured tools + session.LoadAutoLoadTools(cfg, infra.ToolsConn, sessID, uname, func(f string, a ...any) { rs.Cfg.Log.Error("session %s: "+f, append([]any{sessID}, a...)...) }) @@ -163,11 +149,13 @@ func restoreSession(rs *RootState, ps *agent.PersistedAgent) error { env = append(env, infra.PromptEnv...) rt := agent.BuildRuntime(cfg, infra.ToolsConn, cwd, env, layers.SystemPrompt, layers.OpModel, layers.EnvBlock) - restoredSession := agent.RestoreHistory(ps) + restoredHistory := agent.RestoreHistory(ps) - var sessPtr *Session + // Create AgentLog for 9P streaming al := NewAgentLog(remoteTarget) - core := coresession.New(coresession.Config{ + + // Create session with the restored agent + sess := session.New(session.Config{ Backend: be, Profile: profile, AgentsDir: rs.Cfg.AgentsDir, @@ -176,7 +164,7 @@ func restoreSession(rs *RootState, ps *agent.PersistedAgent) error { AgentID: uname, CWD: cwd, Remote: remoteTarget, - History: restoredSession, + History: restoredHistory, Runtime: rt, NewToolServer: infra.NewToolServer, PromptEnvExtra: infra.PromptEnv, @@ -184,42 +172,31 @@ func restoreSession(rs *RootState, ps *agent.PersistedAgent) error { 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) + return "" }, }) - - sessionCtx, sessionCancel := context.WithCancel(rs.Cfg.Ctx) - sess := NewSession(sessID, core, sessionCtx, sessionCancel) + ctx, cancel := context.WithCancel(rs.Cfg.Ctx) + sess.SetContext(ctx, cancel) if ps.SessionName != "" { sess.SetName(ps.SessionName) } - sessPtr = sess - sess.uname = uname - sess.SetAgentLog(uname, al) - sess.proc = infra.Proc - sess.keeper = infra.Keeper - sess.toolsConn = infra.ToolsConn + sess.SetUname(uname) + sess.SetProc(infra.Proc) + sess.SetKeeper(infra.Keeper) + sess.SetToolsConn(infra.ToolsConn) - // Restore agent display name from persisted session. - ag := sess.Core.FindAgent(uname) - if ps.Agent != "" { + // Restore agent display name + if ag := sess.FindAgent(uname); ag != nil && ps.Agent != "" { ag.SetName(ps.Agent) } - replayMessagesToLog(sess.AgentLog(), ps.Messages) + // Replay messages to the AgentLog + replayMessagesToLog(al, ps.Messages) - name := sess.Name() - rs.Mu.Lock() - rs.Sessions[name] = sess - rs.Mu.Unlock() + // Register session and create SessionNode + rs.Manager.Register(sess.Name(), sess) + node := rs.GetOrCreateNode(sess) + node.SetAgentLog(uname, al) toolCount := 0 if infra.ToolsConn != nil { @@ -227,12 +204,7 @@ func restoreSession(rs *RootState, ps *agent.PersistedAgent) error { 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) + rs.Cfg.Log.Info("restored session %s: tools=%d (backend=%s model=%s)", + sessID, toolCount, backendName, modelName) return nil } diff --git a/fs/spec.go b/fs/spec.go index 72a7634..68f11d3 100644 --- a/fs/spec.go +++ b/fs/spec.go @@ -28,7 +28,7 @@ type HandlerCtx struct { Shutdown func() // called by root ctl "kill" // Session scope. - Session *Session + Session *SessionNode Remove func() Rename func(string) error SaveFn func([]byte) error diff --git a/fs/types.go b/fs/types.go index 299c41f..cbb5fdd 100644 --- a/fs/types.go +++ b/fs/types.go @@ -2,62 +2,143 @@ package fs import ( "context" - "encoding/json" - "fmt" "strings" "sync" "time" "ollie/agent" + "ollie/backend" "ollie/session" "ollie/toolsrv" ) -// Session holds session-level state exposed via 9P. -// Agent-scoped state (log, plan, etc.) lives in AgentLog. -type Session struct { - mu sync.RWMutex - id string // immutable UUID, set at creation, never changes - name string // mutable friendly name, defaults to id - uname string // immutable user principal (numeric UID), set at creation - Core *session.Session - SessionCtx context.Context - cancel context.CancelFunc - proc *toolsrv.Process // toolsrv subprocess (session owns lifecycle) - keeper *toolsrv.ProcessKeeper // resilient process manager (respawns on failure) - paused bool // true when session is paused (process stopped) +// SessionNode wraps session.Session for 9P handlers. +// It adds AgentLogs (9P chat streaming state) and models cache. +type SessionNode struct { + Core *session.Session - // toolsConn is the multiplexed connection to the remote tool server, - // shared by agent execution and 9P tool-management handlers. - toolsConn toolsrv.Runner - - // AgentLogs for the session's agents, keyed by agent ID. - // Backward compat: AgentLog() returns the first agent's log (by map iteration order). + // AgentLogs for chat streaming, keyed by agent ID. + mu sync.RWMutex AgentLogs map[string]*AgentLog // Cached model list (expensive API call; refreshed every 24h). modelsMu sync.Mutex modelsCache string modelsCacheAt time.Time - - // disallowTools is the set of tool names that cannot be loaded - // for agents in this session. Set from agent config at creation. - disallowTools map[string]struct{} } -// AgentLog holds all per-agent state exposed via 9P. +// NewSessionNode creates a SessionNode wrapping a session. +func NewSessionNode(core *session.Session) *SessionNode { + return &SessionNode{Core: core, AgentLogs: make(map[string]*AgentLog)} +} + +// --- Delegation to Core --- + +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) 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", "") + } + return err +} + +func (s *SessionNode) Resume() error { + err := s.Core.Resume() + if err == nil { + session.PublishEvent("session."+s.Core.ID()+".resume", "") + } + return err +} + +func (s *SessionNode) LoadTool(name string, ag *agent.Agent) error { + return s.Core.LoadTool(name, ag) +} + +func (s *SessionNode) Ctx() context.Context { + return s.Core.Ctx() +} + +func (s *SessionNode) ToolsConn() toolsrv.Runner { + return s.Core.ToolsConn() +} + +// --- AgentLog management --- + +func (s *SessionNode) AgentLog() *AgentLog { + s.mu.RLock() + defer s.mu.RUnlock() + for _, al := range s.AgentLogs { + return al + } + return nil +} + +func (s *SessionNode) AgentLogFor(agentID string) *AgentLog { + s.mu.RLock() + defer s.mu.RUnlock() + return s.AgentLogs[agentID] +} + +func (s *SessionNode) SetAgentLog(agentID string, al *AgentLog) { + s.mu.Lock() + defer s.mu.Unlock() + if s.AgentLogs == nil { + s.AgentLogs = make(map[string]*AgentLog) + } + s.AgentLogs[agentID] = al +} + +// --- Models cache --- + +const modelsCacheTTL = 24 * time.Hour + +func (s *SessionNode) CachedListModels() string { + s.modelsMu.Lock() + if s.modelsCache != "" && time.Since(s.modelsCacheAt) < modelsCacheTTL { + result := s.modelsCache + s.modelsMu.Unlock() + return result + } + s.modelsMu.Unlock() + + agents := s.Core.Agents() + if len(agents) == 0 { + return "" + } + models := agents[0].ListModels() + result := strings.Join(models, "\n") + s.modelsMu.Lock() + s.modelsCache = result + s.modelsCacheAt = time.Now() + s.modelsMu.Unlock() + return result +} + +func (s *SessionNode) InvalidateModelsCache() { + s.modelsMu.Lock() + s.modelsCache = "" + s.modelsCacheAt = time.Time{} + s.modelsMu.Unlock() +} + +// AgentLog holds per-agent 9P streaming state. type AgentLog struct { mu sync.RWMutex log []byte logVers uint32 ChatOffset int plan []byte - prevPrompt []byte // last submitted prompt; overwritten on each new submission - remote string // SSH target for remote execution (empty = local) - - // Chat stream: condition variable broadcast on every log append. - // Readers of the `chat` file block on this until new data arrives. - chatCond *sync.Cond + prevPrompt []byte + remote string + chatCond *sync.Cond } func NewAgentLog(remote string) *AgentLog { @@ -66,222 +147,6 @@ func NewAgentLog(remote string) *AgentLog { return al } -func NewSession(id string, core *session.Session, ctx context.Context, cancel context.CancelFunc) *Session { - name := id - if idx := strings.IndexByte(id, '-'); idx > 0 { - name = id[:idx] - } - sess := &Session{id: id, name: name, Core: core, SessionCtx: ctx, cancel: cancel} - return sess -} - -// AgentLog returns the first agent's log for backward compatibility. -// The iteration order over the map is non-deterministic; when multi-agent -// support is fully wired, callers should use AgentLogFor or AgentLogs directly. -func (sess *Session) AgentLog() *AgentLog { - sess.mu.RLock() - defer sess.mu.RUnlock() - for _, al := range sess.AgentLogs { - return al - } - return nil -} - -// AgentLogFor returns the AgentLog for the given agent ID, or nil. -func (sess *Session) AgentLogFor(agentID string) *AgentLog { - sess.mu.RLock() - defer sess.mu.RUnlock() - return sess.AgentLogs[agentID] -} - -// SetAgentLog stores the agent log for the given agent ID. -func (sess *Session) SetAgentLog(agentID string, al *AgentLog) { - sess.mu.Lock() - defer sess.mu.Unlock() - sess.setAgentLogLocked(agentID, al) -} - -// setAgentLogLocked is like SetAgentLog but assumes sess.mu is already held. -func (sess *Session) setAgentLogLocked(agentID string, al *AgentLog) { - if sess.AgentLogs == nil { - sess.AgentLogs = make(map[string]*AgentLog) - } - sess.AgentLogs[agentID] = al -} - -const modelsCacheTTL = 24 * time.Hour - -// CachedListModels returns the model list, using a 24h cache to avoid -// repeated expensive API calls. -// NOTE: Must NOT be called while holding sess.mu (caller content() holds RLock). -// Uses its own modelsMu to avoid deadlock. -func (sess *Session) CachedListModels() string { - sess.modelsMu.Lock() - if sess.modelsCache != "" && time.Since(sess.modelsCacheAt) < modelsCacheTTL { - result := sess.modelsCache - sess.modelsMu.Unlock() - return result - } - sess.modelsMu.Unlock() - - // Cache miss — fetch and store. - agents := sess.Core.Agents() - if len(agents) == 0 { - return "" - } - models := agents[0].ListModels() - result := strings.Join(models, "\n") - sess.modelsMu.Lock() - sess.modelsCache = result - sess.modelsCacheAt = time.Now() - sess.modelsMu.Unlock() - return result -} - -// InvalidateModelsCache clears the cached model list. -func (sess *Session) InvalidateModelsCache() { - sess.modelsMu.Lock() - sess.modelsCache = "" - sess.modelsCacheAt = time.Time{} - sess.modelsMu.Unlock() -} - -// loadTool loads a tool into the remote ollie-remote process via the -// session's management toolsConn. Returns an error if the tool is in -// the session's disallow list or if the remote rejects it. -// -// This is the single convergent path for tool loading — both the 9P -// agent/{id}/tools write handler and the autoLoad initialization call this. -func (sess *Session) loadTool(name string, ag *agent.Agent) error { - name = strings.TrimSpace(name) - if name == "" { - return nil - } - - // Check disallow list. - sess.mu.RLock() - _, blocked := sess.disallowTools[name] - sess.mu.RUnlock() - if blocked { - return fmt.Errorf("tool %q is disallowed for this session", name) - } - - // Load on the remote ollie-remote process (for agent execution). - // Prefer the session's management toolsConn (avoids deadlocking with - // the agent loop), fall back to the agent's ToolServer. - var runner toolsrv.Runner - if sess.toolsConn != nil { - runner = sess.toolsConn - } else if ag != nil { - runner = ag.ToolServer() - } - return loadToolOnRunner(runner, name) -} - -func loadToolOnRunner(runner toolsrv.Runner, name string) error { - if runner == nil { - return nil - } - ctxTO, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - args, _ := json.Marshal(map[string]string{"name": name}) - if _, err := runner.CallTool(ctxTO, "tool_load", json.RawMessage(args)); err != nil { - return fmt.Errorf("remote load: %w", err) - } - return nil -} - -func (sess *Session) ID() string { return sess.id } -func (sess *Session) Uname() string { return sess.uname } -func (sess *Session) Name() string { - sess.mu.RLock() - defer sess.mu.RUnlock() - return sess.name -} -func (sess *Session) SetName(name string) { - sess.mu.Lock() - defer sess.mu.Unlock() - sess.name = name -} - -func (sess *Session) Cancel() { - sess.cancel() -} - -// IsPaused returns true if the session is paused. -func (sess *Session) IsPaused() bool { - sess.mu.RLock() - defer sess.mu.RUnlock() - return sess.paused -} - -// IsConnected returns true if the RPC connection to the tool server is alive. -func (sess *Session) IsConnected() bool { - sess.mu.RLock() - conn := sess.toolsConn - sess.mu.RUnlock() - if conn == nil { - return false - } - // Check if conn has Ping method (toolsrv.Conn does) - if pinger, ok := conn.(interface{ Ping() error }); ok { - return pinger.Ping() == nil - } - // No ping available, assume connected if non-nil - return true -} - -// Pause stops the tool server process to save resources. Tool calls will -// return errors until Resume is called. Returns an error if already paused. -func (sess *Session) Pause() error { - sess.mu.Lock() - defer sess.mu.Unlock() - if sess.paused { - return fmt.Errorf("session already paused") - } - // Close the keeper which kills the underlying process - if sess.keeper != nil { - sess.keeper.Close() - } else if sess.proc != nil { - sess.proc.Close() - } - // Close the connection (it will fail on next use anyway) - if sess.toolsConn != nil { - if c, ok := sess.toolsConn.(interface{ Close() }); ok { - c.Close() - } - sess.toolsConn = nil - } - sess.paused = true - PublishEvent("session."+sess.id+".pause", "") - return nil -} - -// Resume restarts the tool server process after a Pause. Returns an error -// if not paused or if respawn fails. -func (sess *Session) Resume() error { - sess.mu.Lock() - defer sess.mu.Unlock() - if !sess.paused { - return fmt.Errorf("session not paused") - } - if sess.keeper == nil { - return fmt.Errorf("no process keeper available") - } - // Dial triggers respawn if needed - conn, err := sess.keeper.Dial() - if err != nil { - return fmt.Errorf("resume failed: %w", err) - } - sess.toolsConn = conn - sess.paused = false - PublishEvent("session."+sess.id+".resume", "") - return nil -} - -// --- AgentLog methods --- - -// AppendLog appends data to the agent's log and bumps the version. func (as *AgentLog) AppendLog(data []byte) { if len(data) == 0 { return @@ -293,7 +158,6 @@ func (as *AgentLog) AppendLog(data []byte) { as.chatCond.Broadcast() } -// EnsureTrailingNewline appends a newline if the log doesn't already end with one. func (as *AgentLog) EnsureTrailingNewline() { as.mu.Lock() if len(as.log) > 0 && as.log[len(as.log)-1] != '\n' { @@ -303,33 +167,20 @@ func (as *AgentLog) EnsureTrailingNewline() { as.mu.Unlock() } -// LogInfo returns the current log length and version atomically. func (as *AgentLog) LogInfo() (length int, vers uint32) { as.mu.RLock() defer as.mu.RUnlock() return len(as.log), as.logVers } -// Log returns the raw chat log bytes. Caller must hold Mu().RLock(). -func (as *AgentLog) Log() []byte { return as.log } - -// Plan returns the plan bytes. Caller must hold Mu().RLock(). -func (as *AgentLog) Plan() []byte { return as.plan } - -// SetPlan sets the plan. Caller must hold Mu().Lock(). -func (as *AgentLog) SetPlan(p []byte) { as.plan = p } - -// PrevPrompt returns the last submitted prompt. Caller must hold Mu().RLock(). +func (as *AgentLog) Log() []byte { return as.log } +func (as *AgentLog) Plan() []byte { return as.plan } +func (as *AgentLog) SetPlan(p []byte) { as.plan = p } func (as *AgentLog) PrevPrompt() []byte { return as.prevPrompt } - -// Mu returns the agent state's mutex. func (as *AgentLog) Mu() *sync.RWMutex { return &as.mu } +func (as *AgentLog) Remote() string { return as.remote } -// Remote returns the SSH target for remote execution. -func (as *AgentLog) Remote() string { return as.remote } - -// NewEventHandler returns an agent.EventHandler that writes formatted events -// to this AgentLog. Use as the Output callback when creating an agent. +// NewEventHandler returns an event handler that writes to an AgentLog. func NewEventHandler(as *AgentLog) agent.EventHandler { streamingRole := "" streamingResponseID := "" @@ -359,24 +210,12 @@ func NewEventHandler(as *AgentLog) agent.EventHandler { } streamingRole = ev.Role } - as.AppendLog(FormatEvent(ev)) + as.AppendLog(session.FormatEvent(ev)) case "reasoning": - // Reasoning is included in the LLM context but not appended - // to the chat log, so the log parser doesn't misinterpret - // header-like text that reasoning might produce. - // - // Don't set streamingRole here either — since we never write - // a [[[reasoning]]] header to the log, the role-transition - // logic would emit a spurious [[[end]]] when the next role - // (assistant) fires. Just absorb the event silently. + // Absorbed silently case "tool": - // reasoning_think writes reasoning content — suppress from log - // same as the reasoning role, so the block parser never sees - // two [[[end]]] markers per turn. Don't set streamingRole - // either — reasoning was never opened with a header. Just - // return silently. if ev.Name == "reasoning_think" { return } @@ -406,7 +245,7 @@ func NewEventHandler(as *AgentLog) agent.EventHandler { streamingFence = "```" streamingRole = "tool" } else { - as.AppendLog(FormatEvent(ev)) + as.AppendLog(session.FormatEvent(ev)) } case "retry": @@ -420,21 +259,56 @@ func NewEventHandler(as *AgentLog) agent.EventHandler { } as.AppendLog([]byte("[[[retry]]]\n")) streamingRole = "retry" - as.AppendLog(FormatEvent(ev)) + as.AppendLog(session.FormatEvent(ev)) case "call": if streamingRole != "" { as.AppendLog([]byte("\n[[[end]]]\n")) streamingRole = "" } - as.AppendLog(FormatEvent(ev)) + as.AppendLog(session.FormatEvent(ev)) default: if streamingRole != "" { as.AppendLog([]byte("\n[[[end]]]\n")) streamingRole = "" } - as.AppendLog(FormatEvent(ev)) + as.AppendLog(session.FormatEvent(ev)) + } + } +} + +// replayMessagesToLog renders persisted messages into an AgentLog. +func replayMessagesToLog(as *AgentLog, messages []backend.Message) { + const maxReplay = 20 + start := len(messages) - maxReplay + if start < 0 { + start = 0 + } + for start < len(messages) && messages[start].Role == "system" { + start++ + } + for _, m := range messages[start:] { + switch m.Role { + case "system": + case "user": + as.AppendLog([]byte("[[[user]]]\n" + m.Content + "\n[[[end]]]\n")) + case "assistant": + header := "[[[assistant]]]" + if m.ID != "" { + header = "[[[assistant:" + m.ID + "]]]" + } + as.AppendLog([]byte(header + "\n")) + if m.Content != "" { + as.AppendLog([]byte(m.Content + "\n")) + } + for _, tc := range m.ToolCalls { + args := strings.Join(strings.Fields(string(tc.Arguments)), " ") + as.AppendLog([]byte("[[[call:" + tc.Name + "]]]\n" + args + "\n[[[end]]]\n")) + } + as.AppendLog([]byte("[[[end]]]\n")) + case "tool": + as.AppendLog([]byte("[[[tool]]]\n```\n" + strings.TrimRight(m.Content, "\n") + "\n```\n[[[end]]]\n")) } } } diff --git a/session/event.go b/session/event.go new file mode 100644 index 0000000..3af9d24 --- /dev/null +++ b/session/event.go @@ -0,0 +1,74 @@ +package session + +import ( + "context" + "strings" + + "github.com/simonfxr/pubsub" +) + +var eventBus = pubsub.NewBus() + +// Event represents a published event with topic and payload. +type Event struct { + Topic string + Payload string +} + +// PublishEvent publishes an event to a specific topic. +// Topic hierarchy: +// +// session.{sid}.new/kill/pause/resume/rename +// session.{sid}.agent.{aid}.new/kill/state +// +// Wildcard routing publishes to all ancestor wildcards: +// +// session.{sid}.agent.{aid}.state -> also publishes to: +// session.{sid}.agent.{aid}.* +// session.{sid}.agent.* +// session.{sid}.* +// session.* +// * +func PublishEvent(topic, payload string) { + ev := Event{Topic: topic, Payload: payload} + eventBus.Publish(topic, ev) + + // Publish to hierarchical wildcards + parts := strings.Split(topic, ".") + for i := len(parts) - 1; i >= 1; i-- { + wildcard := strings.Join(parts[:i], ".") + ".*" + eventBus.Publish(wildcard, ev) + } + eventBus.Publish("*", ev) +} + +// SubscribeEvents returns a channel that receives all events matching the topic. +// Use "*" for all events, "session.*" for all session events, etc. +func SubscribeEvents(ctx context.Context, topic string) <-chan Event { + ch := make(chan Event, 64) + sub := eventBus.SubscribeChan(topic, ch, pubsub.KeepOpen) + go func() { + <-ctx.Done() + eventBus.Unsubscribe(sub) + close(ch) + }() + return ch +} + +// WaitEvent blocks until an event is available. Returns the event as "topic payload\n". +func WaitEvent(ctx context.Context, _ string) ([]byte, string, error) { + ch := SubscribeEvents(ctx, "*") + select { + case ev, ok := <-ch: + if !ok { + return nil, "", nil + } + line := ev.Topic + if ev.Payload != "" { + line += " " + ev.Payload + } + return []byte(line + "\n"), "", nil + case <-ctx.Done(): + return nil, "", nil + } +} diff --git a/fs/format.go b/session/format.go similarity index 57% rename from fs/format.go rename to session/format.go index b76e091..e0ba382 100644 --- a/fs/format.go +++ b/session/format.go @@ -1,8 +1,7 @@ -package fs +package session import ( "ollie/agent" - "ollie/backend" "os" "strings" ) @@ -60,41 +59,3 @@ func FormatEvent(ev agent.Event) []byte { func squashWhitespace(s string) string { return strings.Join(strings.Fields(s), " ") } - -// replayMessagesToLog renders the tail of a persisted message list into the -// session's chat log so that restored sessions show recent history. -func replayMessagesToLog(as *AgentLog, messages []backend.Message) { - const maxReplay = 20 - start := len(messages) - maxReplay - if start < 0 { - start = 0 - } - // Skip system messages at the start of the replay window. - for start < len(messages) && messages[start].Role == "system" { - start++ - } - for _, m := range messages[start:] { - switch m.Role { - case "system": - // skip - case "user": - as.AppendLog([]byte("[[[user]]]\n" + m.Content + "\n[[[end]]]\n")) - case "assistant": - header := "[[[assistant]]]" - if m.ID != "" { - header = "[[[assistant:" + m.ID + "]]]" - } - as.AppendLog([]byte(header + "\n")) - if m.Content != "" { - as.AppendLog([]byte(m.Content + "\n")) - } - for _, tc := range m.ToolCalls { - args := squashWhitespace(string(tc.Arguments)) - as.AppendLog([]byte("[[[call:" + tc.Name + "]]]\n" + args + "\n[[[end]]]\n")) - } - as.AppendLog([]byte("[[[end]]]\n")) - case "tool": - as.AppendLog([]byte("[[[tool]]]\n```\n" + strings.TrimRight(m.Content, "\n") + "\n```\n[[[end]]]\n")) - } - } -} diff --git a/session/manager.go b/session/manager.go new file mode 100644 index 0000000..f5c6284 --- /dev/null +++ b/session/manager.go @@ -0,0 +1,383 @@ +package session + +import ( + "context" + "crypto/rand" + "fmt" + "os" + "sort" + "strings" + "sync" + "time" + + "ollie/agent" + "ollie/backend" + olog "ollie/log" + "ollie/toolsrv" +) + +// ManagerConfig configures the session manager. +type ManagerConfig struct { + Ctx context.Context + Log *olog.Logger + Sink *olog.Sink + AgentsDir string + SessionsDir string + MkdirAll func(string, os.FileMode) error + Yolo bool +} + +// Manager manages the lifecycle of all sessions. +type Manager struct { + cfg ManagerConfig + mu sync.RWMutex + sessions map[string]*Session +} + +// NewManager creates a session manager. +func NewManager(cfg ManagerConfig) *Manager { + return &Manager{ + cfg: cfg, + sessions: make(map[string]*Session), + } +} + +// Sessions returns a snapshot of all sessions. +func (m *Manager) Sessions() map[string]*Session { + m.mu.RLock() + defer m.mu.RUnlock() + result := make(map[string]*Session, len(m.sessions)) + for k, v := range m.sessions { + result[k] = v + } + return result +} + +// Lookup returns the session by name or ID, or nil. +func (m *Manager) Lookup(nameOrID string) *Session { + m.mu.RLock() + defer m.mu.RUnlock() + if s := m.sessions[nameOrID]; s != nil { + return s + } + for _, sess := range m.sessions { + if sess.ID() == nameOrID { + return sess + } + } + return nil +} + +// Register adds a session to the manager. +func (m *Manager) Register(name string, sess *Session) { + m.mu.Lock() + defer m.mu.Unlock() + m.sessions[name] = sess +} + +// Kill terminates a session and removes it from the manager. +func (m *Manager) Kill(name string) error { + m.mu.Lock() + sess := m.sessions[name] + if sess == nil { + m.mu.Unlock() + return fmt.Errorf("session not found: %s", name) + } + delete(m.sessions, name) + m.mu.Unlock() + + 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() + } + + m.cfg.Log.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 (m *Manager) Rename(oldName, newName string) error { + m.mu.Lock() + sess, ok := m.sessions[oldName] + if !ok { + m.mu.Unlock() + return fmt.Errorf("session not found: %s", oldName) + } + if _, exists := m.sessions[newName]; exists { + m.mu.Unlock() + return fmt.Errorf("session already exists: %s", newName) + } + for _, ag := range sess.Agents() { + if ag.IsRunning() { + m.mu.Unlock() + return fmt.Errorf("cannot rename while agent is running") + } + } + sess.SetName(newName) + m.sessions[newName] = sess + delete(m.sessions, oldName) + m.mu.Unlock() + + m.cfg.Log.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 (m *Manager) Shutdown() { + m.InterruptAll() + m.waitIdle(100*time.Millisecond, 5*time.Second) + + m.mu.Lock() + names := make([]string, 0, len(m.sessions)) + for name := range m.sessions { + names = append(names, name) + } + m.mu.Unlock() + + for _, name := range names { + m.mu.Lock() + sess := m.sessions[name] + delete(m.sessions, name) + m.mu.Unlock() + if sess != nil { + sess.Cancel() + sess.Close() + if sess.Keeper() != nil { + sess.Keeper().Close() + } else if sess.Proc() != nil { + sess.Proc().Close() + } + m.cfg.Log.Info("shutdown session %s", name) + } + } +} + +// InterruptAll interrupts every active session. +func (m *Manager) InterruptAll() { + m.mu.RLock() + defer m.mu.RUnlock() + for _, sess := range m.sessions { + for _, ag := range sess.Agents() { + ag.Interrupt(agent.ErrInterrupted) + } + } +} + +func (m *Manager) waitIdle(poll, timeout time.Duration) { + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + allIdle := true + m.mu.RLock() + for _, sess := range m.sessions { + for _, ag := range sess.Agents() { + if ag.State() != "idle" { + allIdle = false + break + } + } + if !allIdle { + break + } + } + m.mu.RUnlock() + if allIdle { + return + } + time.Sleep(poll) + } + m.cfg.Log.Warn("shutdown: timed out waiting for idle") +} + +// BuildIndex builds the session index content. +func (m *Manager) BuildIndex() []byte { + var sb strings.Builder + m.mu.RLock() + names := make([]string, 0, len(m.sessions)) + for name := range m.sessions { + names = append(names, name) + } + sort.Strings(names) + for _, name := range names { + sess := m.sessions[name] + paused := sess.IsPaused() + connected := sess.IsConnected() + cwd := sess.CWD() + var be, model string + var agentList []string + for _, ag := range sess.Agents() { + 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", + name, sess.ID(), paused, connected, cwd, be, model, agents) + } + m.mu.RUnlock() + return []byte(sb.String()) +} + +// CreateEmpty creates an empty session (no agent). +func (m *Manager) CreateEmpty(name string) (*Session, error) { + sessID := NewSessionID() + if name == "" { + name = sessID + if idx := strings.IndexByte(sessID, '-'); idx > 0 { + name = sessID[:idx] + } + } + + m.mu.RLock() + _, exists := m.sessions[name] + m.mu.RUnlock() + if exists { + return nil, fmt.Errorf("session already exists: %s", name) + } + + ctx, cancel := context.WithCancel(m.cfg.Ctx) + sess := NewEmpty(sessID, ctx, cancel) + + // Spawn tool server + cwd, _ := os.Getwd() + var dialOpts []toolsrv.Option + if m.cfg.Yolo { + dialOpts = append(dialOpts, toolsrv.WithYolo()) + } + proc, err := toolsrv.Spawn(ctx, cwd, dialOpts...) + if err != nil { + cancel() + return nil, fmt.Errorf("spawn tool server: %w", err) + } + sess.SetProc(proc) + sess.SetKeeper(toolsrv.NewProcessKeeper(ctx, proc, func(ctx context.Context) (*toolsrv.Process, error) { + return toolsrv.Spawn(ctx, cwd, dialOpts...) + })) + conn, err := toolsrv.Dial(ctx, toolsrv.LocalAddr(proc.Socket)) + if err != nil { + cancel() + proc.Close() + return nil, fmt.Errorf("dial tool server: %w", err) + } + sess.SetToolsConn(conn) + + m.mu.Lock() + m.sessions[name] = sess + m.mu.Unlock() + + m.cfg.Log.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]) +} + +// CreateAgent creates a new agent within an existing session. +// Args: [profile] [key=value...] where keys include backend, model, remote. +func (m *Manager) CreateAgent(sessName string, args []string) (*agent.Agent, error) { + sess := m.Lookup(sessName) + if sess == nil { + return nil, fmt.Errorf("session not found: %s", sessName) + } + + // Parse args + profile := "default" + var backendName, modelName, remoteTarget string + for _, arg := range args { + if k, v, ok := strings.Cut(arg, "="); ok { + switch k { + case "backend": + backendName = v + case "model": + modelName = v + case "remote": + remoteTarget = v + } + } else if profile == "default" { + profile = arg + } + } + + cfg := LoadAgentConfig(m.cfg.AgentsDir, profile, nil) + + uname := NextUname() + sessID := sess.ID() + cwd := sess.CWD() + + // Set up tool server - reuse existing session infrastructure + infra, err := SetupToolServer(ToolServerConfig{ + Ctx: sess.Ctx(), + CWD: cwd, + RemoteTarget: remoteTarget, + Yolo: m.cfg.Yolo, + ReuseFrom: &InfraConfig{ + Proc: sess.Proc(), + Keeper: sess.Keeper(), + ToolsConn: sess.ToolsConn(), + }, + }) + if err != nil { + return nil, fmt.Errorf("setup tool server: %w", err) + } + + // Build prompt layers + layers := BuildPromptLayers(cfg, cwd, sessID, uname, infra.PromptEnv, "") + + // Load tools + LoadAutoLoadTools(cfg, infra.ToolsConn, sessID, uname, func(f string, a ...any) { + m.cfg.Log.Error("session %s agent %s: "+f, append([]any{sessID, uname}, a...)...) + }) + + // Build runtime + env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname} + env = append(env, infra.PromptEnv...) + rt := agent.BuildRuntime(cfg, infra.ToolsConn, cwd, env, layers.SystemPrompt, layers.OpModel, layers.EnvBlock) + + // Apply backend/model overrides to runtime + if backendName != "" { + if be, err := backend.NewWithName(backendName); err == nil { + rt.Backend = be + } + } + if modelName != "" { + rt.Backend.SetModel(modelName) + } + + // Create agent + ag := agent.NewAgent(agent.AgentCfg{ + History: nil, // new agent starts with empty history + Runtime: rt, + Profile: profile, + AgentsDir: m.cfg.AgentsDir, + AgentID: uname, + CWD: cwd, + BaseLayers: layers.BaseLayers(), + PromptEnvExtra: infra.PromptEnv, + NewToolServer: infra.NewToolServer, + NewBackend: backend.NewWithName, + Log: m.cfg.Sink.NewLogger("agent"), + SessionID: sessID, + }) + + sess.AddAgent(ag) + m.cfg.Log.Info("created agent %s in session %s (profile=%s)", uname, sessName, profile) + return ag, nil +} diff --git a/session/session.go b/session/session.go index d682a50..33c5663 100644 --- a/session/session.go +++ b/session/session.go @@ -1,7 +1,9 @@ package session import ( + "context" "crypto/rand" + "encoding/json" "fmt" "io" "os" @@ -17,7 +19,7 @@ import ( "ollie/toolsrv" ) -// Config is the configuration for creating a session. +// Config is the configuration for creating a session with an agent. type Config struct { Backend backend.Backend ModelName string @@ -41,23 +43,44 @@ type Config struct { Output agent.EventHandler } -// Session is the concrete session type. It owns session-level state and -// delegates agent operations to its owned agents. +// Session owns all session-level state: identity, agents, tool server lifecycle. type Session struct { - id string + // Identity + id string // immutable UUID, set at creation + name string // mutable friendly name, defaults to id prefix + uname string // immutable user principal (numeric UID) + + // Context lifecycle + ctx context.Context + cancel context.CancelFunc + + // Agent state + agents []*agent.Agent + log *olog.Logger + auditLog *olog.Logger + sessionsDir string + listHandlers map[string]func() []string + remote string + + // Environment envMu sync.RWMutex env map[string]string plan []byte prevPrompt string - agents []*agent.Agent - log *olog.Logger - sessionsDir string - listHandlers map[string]func() []string - remote string - mu sync.RWMutex - auditLog *olog.Logger + // Tool server lifecycle + proc *toolsrv.Process // toolsrv subprocess + keeper *toolsrv.ProcessKeeper // resilient process manager + toolsConn toolsrv.Runner // multiplexed RPC connection + paused bool // true when paused + // Tool restrictions + disallowTools map[string]struct{} + + // Synchronization + mu sync.RWMutex + + // Autosave saveMu sync.Mutex saveDirty bool saveTimer *time.Timer @@ -129,8 +152,14 @@ func New(cfg Config) *Session { auditLog := log.Sub("audit") + name := cfg.SessionID + if idx := strings.IndexByte(cfg.SessionID, '-'); idx > 0 { + name = cfg.SessionID[:idx] + } + a := &Session{ id: cfg.SessionID, + name: name, env: make(map[string]string), agents: make([]*agent.Agent, 0, 1), log: log, @@ -348,3 +377,218 @@ func (a *Session) SaveSession(path string) error { } return nil } + +// --- Identity accessors --- + +// ID returns the immutable session UUID. +func (a *Session) ID() string { return a.id } + +// Name returns the mutable friendly name. +func (a *Session) Name() string { + a.mu.RLock() + defer a.mu.RUnlock() + return a.name +} + +// SetName sets the friendly name. +func (a *Session) SetName(name string) { + a.mu.Lock() + defer a.mu.Unlock() + a.name = name +} + +// Uname returns the immutable user principal. +func (a *Session) Uname() string { return a.uname } + +// SetUname sets the user principal (should only be called during init). +func (a *Session) SetUname(uname string) { + a.mu.Lock() + defer a.mu.Unlock() + a.uname = uname +} + +// --- Context lifecycle --- + +// Ctx returns the session context. +func (a *Session) Ctx() context.Context { return a.ctx } + +// Cancel cancels the session context. +func (a *Session) Cancel() { + if a.cancel != nil { + a.cancel() + } +} + +// SetContext sets the session context and cancel func. +func (a *Session) SetContext(ctx context.Context, cancel context.CancelFunc) { + a.ctx = ctx + a.cancel = cancel +} + +// --- Tool server lifecycle --- + +// Proc returns the tool server process. +func (a *Session) Proc() *toolsrv.Process { + a.mu.RLock() + defer a.mu.RUnlock() + return a.proc +} + +// SetProc sets the tool server process. +func (a *Session) SetProc(proc *toolsrv.Process) { + a.mu.Lock() + defer a.mu.Unlock() + a.proc = proc +} + +// Keeper returns the process keeper. +func (a *Session) Keeper() *toolsrv.ProcessKeeper { + a.mu.RLock() + defer a.mu.RUnlock() + return a.keeper +} + +// SetKeeper sets the process keeper. +func (a *Session) SetKeeper(keeper *toolsrv.ProcessKeeper) { + a.mu.Lock() + defer a.mu.Unlock() + a.keeper = keeper +} + +// ToolsConn returns the tool server connection. +func (a *Session) ToolsConn() toolsrv.Runner { + a.mu.RLock() + defer a.mu.RUnlock() + return a.toolsConn +} + +// SetToolsConn sets the tool server connection. +func (a *Session) SetToolsConn(conn toolsrv.Runner) { + a.mu.Lock() + defer a.mu.Unlock() + a.toolsConn = conn +} + +// IsPaused returns true if the session is paused. +func (a *Session) IsPaused() bool { + a.mu.RLock() + defer a.mu.RUnlock() + return a.paused +} + +// IsConnected returns true if the tool server connection is alive. +func (a *Session) IsConnected() bool { + a.mu.RLock() + conn := a.toolsConn + a.mu.RUnlock() + if conn == nil { + return false + } + if pinger, ok := conn.(interface{ Ping() error }); ok { + return pinger.Ping() == nil + } + return true +} + +// Pause stops the tool server to save resources. +func (a *Session) Pause() error { + a.mu.Lock() + defer a.mu.Unlock() + if a.paused { + return fmt.Errorf("session already paused") + } + if a.keeper != nil { + a.keeper.Close() + } else if a.proc != nil { + a.proc.Close() + } + if a.toolsConn != nil { + if c, ok := a.toolsConn.(interface{ Close() }); ok { + c.Close() + } + a.toolsConn = nil + } + a.paused = true + return nil +} + +// Resume restarts the tool server after a Pause. +func (a *Session) Resume() error { + a.mu.Lock() + defer a.mu.Unlock() + if !a.paused { + return fmt.Errorf("session not paused") + } + if a.keeper == nil { + return fmt.Errorf("no process keeper available") + } + conn, err := a.keeper.Dial() + if err != nil { + return fmt.Errorf("resume failed: %w", err) + } + a.toolsConn = conn + a.paused = false + return nil +} + +// --- Tool loading --- + +// SetDisallowTools sets the tool disallow list. +func (a *Session) SetDisallowTools(disallow map[string]struct{}) { + a.mu.Lock() + defer a.mu.Unlock() + a.disallowTools = disallow +} + +// LoadTool loads a tool into the tool server. +func (a *Session) LoadTool(name string, ag *agent.Agent) error { + name = strings.TrimSpace(name) + if name == "" { + return nil + } + a.mu.RLock() + _, blocked := a.disallowTools[name] + a.mu.RUnlock() + if blocked { + return fmt.Errorf("tool %q is disallowed for this session", name) + } + var runner toolsrv.Runner + if a.toolsConn != nil { + runner = a.toolsConn + } else if ag != nil { + runner = ag.ToolServer() + } + return LoadToolOnRunner(runner, name) +} + +// LoadToolOnRunner loads a tool on the given runner. +func LoadToolOnRunner(runner toolsrv.Runner, name string) error { + if runner == nil { + return nil + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + args, _ := json.Marshal(map[string]string{"name": name}) + if _, err := runner.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. +func NewEmpty(id string, ctx context.Context, cancel context.CancelFunc) *Session { + name := id + if idx := strings.IndexByte(id, '-'); idx > 0 { + name = id[:idx] + } + return &Session{ + id: id, + name: name, + ctx: ctx, + cancel: cancel, + env: make(map[string]string), + agents: make([]*agent.Agent, 0), + } +} diff --git a/fs/setup.go b/session/setup.go similarity index 98% rename from fs/setup.go rename to session/setup.go index 522787b..e7da5b6 100644 --- a/fs/setup.go +++ b/session/setup.go @@ -1,4 +1,4 @@ -package fs +package session import ( "context" @@ -210,7 +210,7 @@ func LoadAutoLoadTools(cfg *agent.AgentConfig, runner toolsrv.Runner, sessID, un } for _, tl := range cfg.AutoLoad { - if err := loadToolOnRunner(runner, tl); err != nil { + if err := LoadToolOnRunner(runner, tl); err != nil { if logError != nil { logError("autoLoad tool %q: %v", tl, err) }