lkn
/
ollie
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: rename saveSession/flushSave to save/flush, remove legacy restore 

- 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
This commit is contained in:
Levi Neely 2026-08-04 14:38:09 +02:00
parent f348b968a2
commit 86f7c54805
8 changed files with 39 additions and 152 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
agent/agent.go
View File

@ -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
}

10
agent/commands.go
View File

@ -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) {

10
agent/loop.go
View File

@ -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 {

16
agent/new.go
View File

@ -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)

10
agent/turn.go
View File

@ -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 != "" {

131
session/persist.go
View File

@ -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
}

2
session/registry.go
View File

@ -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)

8
session/session.go
View File

@ -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
}