345 lines
8.8 KiB
Go
345 lines
8.8 KiB
Go
package session
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
|
|
"ollie/agent"
|
|
"ollie/backend"
|
|
"ollie/toolsrv"
|
|
)
|
|
|
|
// --- Persistence ---
|
|
|
|
// PersistedSession is the on-disk format for a saved session with multiple agents.
|
|
type PersistedSession struct {
|
|
ID string `json:"id"`
|
|
Name string `json:"name"`
|
|
CWD string `json:"cwd"`
|
|
Remote string `json:"remote,omitempty"`
|
|
Paused bool `json:"paused,omitempty"`
|
|
Agents []PersistedAgent `json:"agents"`
|
|
}
|
|
|
|
// PersistedAgent is per-agent state within a session.
|
|
type PersistedAgent struct {
|
|
ID string `json:"id"` // uname
|
|
Name string `json:"name"` // display name
|
|
Profile string `json:"profile"`
|
|
Backend string `json:"backend"`
|
|
Model string `json:"model"`
|
|
CWD string `json:"cwd"`
|
|
Messages []backend.Message `json:"messages"`
|
|
TaskState *agent.TaskState `json:"taskState,omitempty"`
|
|
// Usage tracking
|
|
TotalInputTokens int `json:"totalInputTokens,omitempty"`
|
|
TotalCachedInputTokens int `json:"totalCachedInputTokens,omitempty"`
|
|
TotalCacheCreationTokens int `json:"totalCacheCreationTokens,omitempty"`
|
|
TotalOutputTokens int `json:"totalOutputTokens,omitempty"`
|
|
TotalRequests int `json:"totalRequests,omitempty"`
|
|
Estimated bool `json:"estimated,omitempty"`
|
|
LastTurnCostUSD float64 `json:"lastTurnCostUSD,omitempty"`
|
|
SessionCostUSD float64 `json:"sessionCostUSD,omitempty"`
|
|
}
|
|
|
|
// ActiveSessionsDir returns the directory for persisted sessions.
|
|
func ActiveSessionsDir() string {
|
|
return filepath.Join(pkgSessionsDir, "active")
|
|
}
|
|
|
|
// PersistSession saves a session (with all agents) to disk.
|
|
func PersistSession(name string) error {
|
|
sess := Lookup(name)
|
|
if sess == nil {
|
|
return nil
|
|
}
|
|
agents := sess.Agents()
|
|
if len(agents) == 0 {
|
|
return nil
|
|
}
|
|
|
|
ps := PersistedSession{
|
|
ID: sess.ID(),
|
|
Name: sess.Name(),
|
|
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(),
|
|
TaskState: ag.TaskState(),
|
|
}
|
|
if usage := ag.Usage(); usage != nil {
|
|
pa.TotalInputTokens = usage.TotalInputTokens
|
|
pa.TotalCachedInputTokens = usage.TotalCachedInputTokens
|
|
pa.TotalCacheCreationTokens = usage.TotalCacheCreationTokens
|
|
pa.TotalOutputTokens = usage.TotalOutputTokens
|
|
pa.TotalRequests = usage.TotalRequests
|
|
pa.Estimated = usage.Estimated
|
|
pa.LastTurnCostUSD = usage.LastTurnCostUSD
|
|
pa.SessionCostUSD = usage.SessionCostUSD
|
|
}
|
|
ps.Agents = append(ps.Agents, pa)
|
|
}
|
|
|
|
dir := ActiveSessionsDir()
|
|
os.MkdirAll(dir, 0700)
|
|
path := filepath.Join(dir, sess.ID()+".json")
|
|
|
|
data, err := json.MarshalIndent(ps, "", " ")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return os.WriteFile(path, data, 0600)
|
|
}
|
|
|
|
// RemovePersistedSession removes a persisted session file.
|
|
func RemovePersistedSession(id string) {
|
|
path := filepath.Join(ActiveSessionsDir(), id+".json")
|
|
os.Remove(path)
|
|
}
|
|
|
|
// SaveAllSessions persists all active sessions.
|
|
func SaveAllSessions() {
|
|
for name := range Sessions() {
|
|
if err := PersistSession(name); err != nil {
|
|
pkgLog.Error("persist session %s: %v", name, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
// RestoredAgent holds restored agent info for fs layer.
|
|
type RestoredAgent struct {
|
|
ID string
|
|
Remote string
|
|
Messages []backend.Message
|
|
}
|
|
|
|
// RestoredSession holds a restored session and metadata for fs layer.
|
|
type RestoredSession struct {
|
|
Session *Session
|
|
Agents []RestoredAgent
|
|
ToolCount int
|
|
}
|
|
|
|
// 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)
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
|
|
var files []string
|
|
for _, e := range entries {
|
|
if strings.HasSuffix(e.Name(), ".json") {
|
|
files = append(files, filepath.Join(dir, e.Name()))
|
|
}
|
|
}
|
|
|
|
if len(files) == 0 {
|
|
return nil
|
|
}
|
|
|
|
var mu sync.Mutex
|
|
var results []*RestoredSession
|
|
var wg sync.WaitGroup
|
|
for _, path := range files {
|
|
wg.Add(1)
|
|
go func(path string) {
|
|
defer wg.Done()
|
|
rs, err := restoreSessionFromFile(path)
|
|
if err != nil {
|
|
pkgLog.Error("restore session %s: %v; removing from active sessions", path, err)
|
|
os.Remove(path)
|
|
return
|
|
}
|
|
mu.Lock()
|
|
results = append(results, rs)
|
|
mu.Unlock()
|
|
}(path)
|
|
}
|
|
wg.Wait()
|
|
return results
|
|
}
|
|
|
|
func restoreSessionFromFile(path string) (*RestoredSession, error) {
|
|
data, err := os.ReadFile(path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var ps PersistedSession
|
|
if err := json.Unmarshal(data, &ps); err != nil {
|
|
return nil, err
|
|
}
|
|
if ps.ID == "" || len(ps.Agents) == 0 {
|
|
return nil, fmt.Errorf("invalid session file: missing ID or agents")
|
|
}
|
|
return restoreMultiAgentSession(&ps)
|
|
}
|
|
|
|
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)
|
|
if ps.Paused {
|
|
sess.SetPaused(true)
|
|
}
|
|
|
|
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
|
|
}
|
|
if cwd == "" {
|
|
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,
|
|
})
|
|
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 := LoadAgentConfig(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.Runner
|
|
var newToolServer func() toolsrv.Runner
|
|
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
|
|
history := agent.RestoreHistoryFromMessages(pa.Messages)
|
|
|
|
ag := agent.NewAgent(agent.AgentCfg{
|
|
History: history,
|
|
Runtime: rt,
|
|
Profile: pa.Profile,
|
|
AgentsDir: pkgAgentsDir,
|
|
AgentID: pa.ID,
|
|
CWD: cwd,
|
|
BaseLayers: layers.BaseLayers(),
|
|
PromptEnvExtra: promptEnv,
|
|
NewToolServer: newToolServer,
|
|
NewBackend: backend.NewWithName,
|
|
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 len(sess.Agents()) == 0 {
|
|
cancel()
|
|
return nil, fmt.Errorf("no agents restored")
|
|
}
|
|
|
|
// Set uname to first agent's ID
|
|
sess.SetUname(sess.Agents()[0].ID())
|
|
|
|
Register(sess.Name(), sess)
|
|
|
|
toolCount := 0
|
|
if sess.ToolsConn() != nil {
|
|
if infos, err := sess.ToolsConn().ListTools(); err == nil {
|
|
toolCount = len(infos)
|
|
}
|
|
}
|
|
pkgLog.Info("restored session %s with %d agents, tools=%d, paused=%v", ps.ID, len(restoredAgents), toolCount, ps.Paused)
|
|
|
|
return &RestoredSession{
|
|
Session: sess,
|
|
Agents: restoredAgents,
|
|
ToolCount: toolCount,
|
|
}, nil
|
|
}
|
|
|
|
|