From 86f7c54805271384a58f8bc4063a62ad716567db Mon Sep 17 00:00:00 2001 From: Levi Neely Date: Tue, 4 Aug 2026 14:38:09 +0200 Subject: [PATCH] refactor: rename saveSession/flushSave to save/flush, remove legacy restore MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Rename AgentCfg.SaveSession/FlushSave to Save/Flush for brevity - Remove restoreLegacySession — all sessions now use multi-agent format - Pass Save/Flush callbacks to agents in CreateAgent and restore paths - Add AuditLog fallback warning in NewAgent --- agent/agent.go | 4 +- agent/commands.go | 10 +--- agent/loop.go | 10 +--- agent/new.go | 16 ++++-- agent/turn.go | 10 ++-- session/persist.go | 131 ++++---------------------------------------- session/registry.go | 2 + session/session.go | 8 +-- 8 files changed, 39 insertions(+), 152 deletions(-) diff --git a/agent/agent.go b/agent/agent.go index 1580dd6..5279547 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -54,8 +54,8 @@ type Agent struct { sessionID string // the owning session's ID startupMessages []string readPlanStep func() string - saveSession func() // trigger debounced persistence - flushSave func() // immediately flush persistence + save func() // trigger debounced persistence + flush func() // immediately flush persistence turnError func(ctx context.Context, errType, errMsg string) HookResult // overridable for tests onStateChange func(agentID, state string) // optional callback for state changes } diff --git a/agent/commands.go b/agent/commands.go index 01a9e38..07a2727 100644 --- a/agent/commands.go +++ b/agent/commands.go @@ -206,12 +206,8 @@ func (ag *Agent) cmdAgent(_ context.Context, args []string) { ag.wireToolsChanged() // Profile switches clear history, so mark the session dirty and flush its // metadata now rather than waiting for a future turn (which may never occur). - if ag.saveSession != nil { - ag.saveSession() - } - if ag.flushSave != nil { - ag.flushSave() - } + ag.save() + ag.flush() ag.notifyChange() for _, msg := range rt.Messages { ag.emit(infoEvent(msg)) @@ -240,7 +236,7 @@ func (ag *Agent) cmdCompact(ctx context.Context, _ []string) { return } ag.emit(infoEvent(fmt.Sprintf("compacted %d messages", n))) - ag.saveSession() + ag.save() } func (ag *Agent) cmdContext(ctx context.Context, _ []string) { diff --git a/agent/loop.go b/agent/loop.go index a64a137..e089c3c 100644 --- a/agent/loop.go +++ b/agent/loop.go @@ -62,7 +62,7 @@ type agentConfig struct { GenerationParams backend.GenerationParams PopInject func() string // returns and clears pending inject, or "" AutoCompact func(ctx context.Context) // called after each tool round; may compact in-place - SaveSession func() // called after each state.update(); persists mid-turn progress + Save func() // called after each state.update(); persists mid-turn progress PreTool func(ctx context.Context, name string, args json.RawMessage) HookResult // called before each tool; exit 2 blocks execution PostTool func(ctx context.Context, name string, args json.RawMessage, result string) HookResult // called after each tool; exit 0 appends, exit 2 replaces result TurnError func(ctx context.Context, errType, errMsg string) HookResult // called on first backend error; if ran, skips retries @@ -228,9 +228,7 @@ func run(ctx context.Context, cfg agentConfig, state state) error { }) } state.update(msg, results) - if cfg.SaveSession != nil { - cfg.SaveSession() - } + cfg.Save() return ctx.Err() } @@ -517,9 +515,7 @@ func run(ctx context.Context, cfg agentConfig, state state) error { } state.update(msg, results) - if cfg.SaveSession != nil { - cfg.SaveSession() - } + cfg.Save() // Auto-update TaskState from what actually happened this round. if ts := state.taskState(); ts != nil { diff --git a/agent/new.go b/agent/new.go index 3c40741..1096fc7 100644 --- a/agent/new.go +++ b/agent/new.go @@ -26,12 +26,18 @@ type AgentCfg struct { SessionID string StartupMsgs []string ReadPlanStep func() string - SaveSession func() - FlushSave func() + Save func() + Flush func() } // NewAgent constructs an Agent from the given configuration. func NewAgent(cfg AgentCfg) *Agent { + auditLog := cfg.AuditLog + if auditLog == nil && cfg.Log != nil { + cfg.Log.Warn("AuditLog not provided for agent %s, deriving from Log", cfg.AgentID) + auditLog = cfg.Log.Sub("audit") + } + ag := &Agent{ history: cfg.History, runtime: cfg.Runtime, @@ -46,12 +52,12 @@ func NewAgent(cfg AgentCfg) *Agent { newBackend: cfg.NewBackend, output: cfg.Output, log: cfg.Log, - auditLog: cfg.AuditLog, + auditLog: auditLog, sessionID: cfg.SessionID, startupMessages: cfg.StartupMsgs, readPlanStep: cfg.ReadPlanStep, - saveSession: cfg.SaveSession, - flushSave: cfg.FlushSave, + save: cfg.Save, + flush: cfg.Flush, state: "idle", } ag.changeCond = sync.NewCond(&ag.changeMu) diff --git a/agent/turn.go b/agent/turn.go index b51268d..da027fb 100644 --- a/agent/turn.go +++ b/agent/turn.go @@ -223,7 +223,7 @@ func (ag *Agent) executeTurn(ctx context.Context, input string) string { ag.cfg.IncrToolCallCount = func() int64 { return ag.toolCallCount.Add(1) } - ag.cfg.SaveSession = func() { ag.saveSession() } + ag.cfg.Save = ag.save ag.cfg.ResultCache = &ag.resultCache ag.cfg.AutoCompact = func(ctx context.Context) { if ctx.Err() != nil || ag.history == nil { @@ -324,9 +324,7 @@ func (ag *Agent) executeTurn(ctx context.Context, input string) string { ag.SetReply(replyBuf.String()) replyBuf.Reset() ag.SetState("idle") - if ag.flushSave != nil { - ag.flushSave() - } + ag.flush() if err != nil { // Keep completed work — only remove cancelled tool results. @@ -335,7 +333,7 @@ func (ag *Agent) executeTurn(ctx context.Context, input string) string { } if errors.Is(err, context.Canceled) || errors.Is(err, ErrInterrupted) { ag.auditLog.Debug("turn: interrupted session=%s", ag.sessionID) - ag.saveSession() + ag.save() return "" } ag.emit(Event{Role: "error", Content: err.Error()}) @@ -370,7 +368,7 @@ func (ag *Agent) executeTurn(ctx context.Context, input string) string { auditTruncate(ag.Reply()), ag.history.LastTurnCostUSD, ag.history.SessionCostUSD, ag.sessionID) ag.notifyChange() } - ag.saveSession() + ag.save() // Post-turn hook said "continue" — its context becomes the next prompt. if stopResult.Blocked && stopResult.Context != "" { diff --git a/session/persist.go b/session/persist.go index 1b6991c..e1a261f 100644 --- a/session/persist.go +++ b/session/persist.go @@ -3,6 +3,7 @@ package session import ( "context" "encoding/json" + "fmt" "os" "path/filepath" "strings" @@ -158,7 +159,8 @@ func RestoreAllSessions() []*RestoredSession { defer wg.Done() rs, err := restoreSessionFromFile(path) if err != nil { - pkgLog.Error("restore session %s: %v", path, err) + pkgLog.Error("restore session %s: %v; removing from active sessions", path, err) + os.Remove(path) return } mu.Lock() @@ -176,18 +178,14 @@ func restoreSessionFromFile(path string) (*RestoredSession, error) { return nil, err } - // Try new multi-agent format first var ps PersistedSession - if err := json.Unmarshal(data, &ps); err == nil && ps.ID != "" && len(ps.Agents) > 0 { - return restoreMultiAgentSession(&ps) - } - - // Fall back to legacy single-agent format - var legacy agent.PersistedAgent - if err := json.Unmarshal(data, &legacy); err != nil { + if err := json.Unmarshal(data, &ps); err != nil { return nil, err } - return restoreLegacySession(&legacy) + 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) { @@ -259,6 +257,8 @@ func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) { NewBackend: backend.NewWithName, Log: pkgSink.NewLogger("agent"), SessionID: ps.ID, + Save: sess.saveSession, + Flush: sess.flushSave, }) ag.SetName(pa.Name) @@ -295,115 +295,4 @@ func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) { }, nil } -// restoreLegacySession handles the old single-agent format for backward compatibility. -func restoreLegacySession(ps *agent.PersistedAgent) (*RestoredSession, error) { - cwd := ps.CWD - if cwd == "" { - cwd, _ = os.Getwd() - } - profile := ps.Profile - if profile == "" { - profile = "default" - } - sessID := ps.ID - remoteTarget := ps.Remote - cfg := LoadAgentConfig(pkgAgentsDir, profile, nil) - - backendName := ps.Backend - if backendName == "" && cfg != nil && cfg.Backend != "" { - backendName = cfg.Backend - } - be, err := backend.NewWithName(backendName) - if err != nil { - return nil, err - } - modelName := ps.Model - if modelName == "" && cfg != nil && cfg.Model != "" { - modelName = cfg.Model - } - if modelName == "" { - modelName = os.Getenv("OLLIE_MODEL") - } - if modelName != "" { - be.SetModel(modelName) - } - - uname := ps.Uname - if uname == "" { - uname = NextUname() - } - - infra, err := SetupToolServer(ToolServerConfig{ - Ctx: pkgCtx, - CWD: cwd, - RemoteTarget: remoteTarget, - Yolo: pkgYolo, - }) - if err != nil { - return nil, err - } - - layers := BuildPromptLayers(cfg, cwd, sessID, uname, infra.PromptEnv, "") - - LoadAutoLoadTools(cfg, infra.ToolsConn, sessID, uname, func(f string, a ...any) { - pkgLog.Error("session %s: "+f, append([]any{sessID}, a...)...) - }) - - 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) - - restoredHistory := agent.RestoreHistory(ps) - - sess := New(Config{ - Backend: be, - Profile: profile, - AgentsDir: pkgAgentsDir, - SessionsDir: pkgSessionsDir, - SessionID: sessID, - AgentID: uname, - CWD: cwd, - Remote: remoteTarget, - History: restoredHistory, - Runtime: rt, - NewToolServer: infra.NewToolServer, - PromptEnvExtra: infra.PromptEnv, - BaseLayers: layers.BaseLayers(), - Log: pkgSink.NewLogger("core"), - }) - ctx, cancel := context.WithCancel(pkgCtx) - sess.SetContext(ctx, cancel) - if ps.SessionName != "" { - sess.SetName(ps.SessionName) - } - sess.SetUname(uname) - sess.SetProc(infra.Proc) - sess.SetKeeper(infra.Keeper) - sess.SetToolsConn(infra.ToolsConn) - - if ag := sess.FindAgent(uname); ag != nil && ps.Agent != "" { - ag.SetName(ps.Agent) - } - - Register(sess.Name(), sess) - - toolCount := 0 - if infra.ToolsConn != nil { - if infos, err := infra.ToolsConn.ListTools(); err == nil { - toolCount = len(infos) - } - } - pkgLog.Info("restored legacy session %s: tools=%d (backend=%s model=%s)", - sessID, toolCount, backendName, modelName) - - return &RestoredSession{ - Session: sess, - Agents: []RestoredAgent{{ - ID: uname, - Remote: remoteTarget, - Messages: ps.Messages, - }}, - ToolCount: toolCount, - }, nil -} diff --git a/session/registry.go b/session/registry.go index 7ad63c9..9d295ca 100644 --- a/session/registry.go +++ b/session/registry.go @@ -387,6 +387,8 @@ func CreateAgent(sessName string, args []string) (*agent.Agent, error) { NewBackend: backend.NewWithName, Log: pkgSink.NewLogger("agent"), SessionID: sessID, + Save: sess.saveSession, + Flush: sess.flushSave, }) sess.AddAgent(ag) diff --git a/session/session.go b/session/session.go index c27248c..0766bad 100644 --- a/session/session.go +++ b/session/session.go @@ -186,8 +186,8 @@ func New(cfg Config) *Session { SessionID: cfg.SessionID, StartupMsgs: rt.Messages, ReadPlanStep: cfg.ReadPlanStep, - SaveSession: a.saveSession, - FlushSave: a.flushSave, + Save: a.saveSession, + Flush: a.flushSave, }) a.agents = append(a.agents, ag) @@ -220,8 +220,8 @@ func (a *Session) AgentAt(idx int) *agent.Agent { func (a *Session) AgentCount() int { return len(a.agents) } // SaveFuncs returns save and flush callbacks suitable for use as -// agent.AgentCfg.SaveSession and FlushSave when adding agents -// to this session. +// SaveFuncs returns the save and flush callbacks for use in +// agent.AgentCfg.Save and Flush when adding agents to this session. func (a *Session) SaveFuncs() (save func(), flush func()) { return a.saveSession, a.flushSave }