325 lines
8.3 KiB
Go
325 lines
8.3 KiB
Go
package session
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
|
|
"ollie/cmd/olliesrv/internal/agent"
|
|
"ollie/cmd/olliesrv/internal/backend"
|
|
)
|
|
|
|
// --- Persistence ---
|
|
|
|
// PersistedSession is the on-disk format for a saved session.
|
|
type PersistedSession struct {
|
|
ID string `json:"id"`
|
|
Name string `json:"name"`
|
|
CWD string `json:"cwd"`
|
|
Remote string `json:"remote,omitempty"`
|
|
Yolo bool `json:"yolo,omitempty"`
|
|
Paused bool `json:"paused,omitempty"`
|
|
Workflow string `json:"workflow,omitempty"`
|
|
Variant string `json:"variant,omitempty"`
|
|
Goal string `json:"goal,omitempty"`
|
|
GoalStatus string `json:"goalStatus,omitempty"`
|
|
Agents []PersistedAgent `json:"agents"`
|
|
}
|
|
|
|
// PersistedAgent is per-agent state within a session.
|
|
type PersistedAgent struct {
|
|
ID string `json:"id"`
|
|
Name string `json:"name"`
|
|
ParentID string `json:"parent,omitempty"`
|
|
Profile string `json:"profile"`
|
|
Backend string `json:"backend"`
|
|
Model string `json:"model"`
|
|
Peers []string `json:"peers,omitempty"`
|
|
Messages []backend.Message `json:"messages"`
|
|
// 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
|
|
}
|
|
|
|
goalText, goalStatus := sess.Goal()
|
|
ps := PersistedSession{
|
|
ID: sess.ID,
|
|
Name: sess.Name(),
|
|
CWD: sess.Cwd(),
|
|
Remote: sess.Remote,
|
|
Yolo: sess.Yolo,
|
|
Paused: sess.IsPaused(),
|
|
Workflow: sess.Workflow(),
|
|
Variant: sess.Variant(),
|
|
Goal: goalText,
|
|
GoalStatus: goalStatus,
|
|
}
|
|
|
|
for _, ag := range agents {
|
|
pa := PersistedAgent{
|
|
ID: ag.ID(),
|
|
Name: ag.Name(),
|
|
ParentID: ag.ParentID(),
|
|
Profile: ag.Profile(),
|
|
Backend: ag.BackendName(),
|
|
Model: ag.ModelName(),
|
|
Peers: ag.Peers(),
|
|
Messages: backend.SanitizeMessages(ag.Messages()),
|
|
}
|
|
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)
|
|
}
|
|
|
|
// 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.
|
|
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 rmu 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", path, err)
|
|
os.Remove(path)
|
|
return
|
|
}
|
|
rmu.Lock()
|
|
results = append(results, rs)
|
|
rmu.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) {
|
|
ctx, cancel := context.WithCancel(serverCtx)
|
|
sess := NewEmpty(ps.ID, ctx, cancel)
|
|
sess.SetName(ps.Name)
|
|
sess.Remote = ps.Remote
|
|
sess.Yolo = ps.Yolo
|
|
if ps.Paused {
|
|
sess.mu.Lock()
|
|
sess.paused = true
|
|
sess.mu.Unlock()
|
|
}
|
|
if ps.CWD != "" {
|
|
sess.SetCwd(ps.CWD)
|
|
}
|
|
if ps.Workflow != "" {
|
|
sess.SetWorkflow(ps.Workflow)
|
|
}
|
|
if ps.Variant != "" {
|
|
sess.SetVariant(ps.Variant)
|
|
}
|
|
if ps.Goal != "" {
|
|
sess.SetGoal(ps.Goal)
|
|
}
|
|
if ps.GoalStatus != "" {
|
|
sess.SetGoalStatus(ps.GoalStatus)
|
|
}
|
|
|
|
var restoredAgents []RestoredAgent
|
|
|
|
for _, pa := range ps.Agents {
|
|
// For paused sessions, create a minimal agent without tool server
|
|
if ps.Paused {
|
|
history := agent.RestoreHistoryFromMessages(pa.Messages)
|
|
ag := agent.NewAgent(agent.AgentParams{
|
|
History: history,
|
|
Runtime: &agent.Runtime{Preamble: &agent.Preamble{}},
|
|
Profile: pa.Profile,
|
|
AgentsDir: pkgAgentsDir,
|
|
ID: pa.ID,
|
|
ParentID: pa.ParentID,
|
|
GetCwd: sess.Cwd,
|
|
Log: pkgSink.NewLogger("agent"),
|
|
SessionID: ps.ID,
|
|
Save: sess.saveSession,
|
|
Flush: sess.flushSave,
|
|
})
|
|
ag.SetName(pa.Name)
|
|
ag.RequireMemoryWake()
|
|
if err := sess.AddAgent(ag); err != nil {
|
|
pkgLog.Error("restore agent %s: %v", pa.ID, err)
|
|
ag.Close()
|
|
continue
|
|
}
|
|
restoredAgents = append(restoredAgents, RestoredAgent{
|
|
ID: pa.ID,
|
|
Remote: ps.Remote,
|
|
Messages: pa.Messages,
|
|
})
|
|
continue
|
|
}
|
|
|
|
// Active session — use shared buildAgent
|
|
history := agent.RestoreHistoryFromMessages(pa.Messages)
|
|
ag, err := buildAgent(sess, AgentParams{
|
|
Profile: pa.Profile,
|
|
ID: pa.ID,
|
|
Name: pa.Name,
|
|
Backend: pa.Backend,
|
|
Model: pa.Model,
|
|
Remote: ps.Remote,
|
|
ParentID: pa.ParentID,
|
|
History: history,
|
|
})
|
|
if err != nil {
|
|
pkgLog.Error("restore agent %s: %v", pa.ID, err)
|
|
continue
|
|
}
|
|
ag.RequireMemoryWake()
|
|
if err := sess.AddAgent(ag); err != nil {
|
|
pkgLog.Error("restore agent %s: %v", pa.ID, err)
|
|
ag.Close()
|
|
continue
|
|
}
|
|
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")
|
|
}
|
|
|
|
// Restore peer links (each agent's peers were persisted independently).
|
|
for _, pa := range ps.Agents {
|
|
if len(pa.Peers) == 0 {
|
|
continue
|
|
}
|
|
ag := sess.FindAgent(pa.ID)
|
|
if ag == nil {
|
|
continue
|
|
}
|
|
for _, peerName := range pa.Peers {
|
|
ag.AddPeer(peerName)
|
|
}
|
|
}
|
|
|
|
sess.Uname = sess.Agents()[0].ID()
|
|
Register(sess.Name(), sess)
|
|
|
|
// Start bypass approval loop for restored sessions (unless paused)
|
|
if !ps.Paused {
|
|
sess.StartBypassLoop()
|
|
}
|
|
|
|
toolCount := 0
|
|
if conn := sess.ToolsConn(); conn != nil {
|
|
if infos, err := conn.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
|
|
}
|