From fa56334c1f2106ddd32468de06c0260f055cb5aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aur=C3=A9lien=20Sibiril?= <81782+aureliensibiril@users.noreply.github.com> Date: Sun, 26 Apr 2026 19:09:11 +0200 Subject: [PATCH] Add OnRunSnapshot hook for checkpoint persistence MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Run hooks already exposed OnRunRestore for the read side of a suspend/restore cycle. The write side -- every coreLoop or restore call that persists a checkpoint to the Checkpointer -- had no corresponding hook, so callers wanting to record metrics, audit events, or trigger external state transitions on every snapshot had no insertion point. Add OnRunSnapshot to RunHooks and emit it after each successful Checkpointer.Save: the suspend, awaiting-approval, nested-approval, post-tool-turn, and restore-progress sites. The hook fires only on durable saves; save failures still log and skip the hook so observers never see a checkpoint that did not land. Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com> --- pkg/agent/hooks.go | 2 ++ pkg/agent/restore.go | 1 + pkg/agent/run.go | 10 ++++++++++ 3 files changed, 13 insertions(+) diff --git a/pkg/agent/hooks.go b/pkg/agent/hooks.go index 67b02fd83..276a7d495 100644 --- a/pkg/agent/hooks.go +++ b/pkg/agent/hooks.go @@ -25,6 +25,7 @@ type RunHooks interface { OnRunStart(ctx context.Context, agent *Agent, messages []llm.Message) OnRunEnd(ctx context.Context, agent *Agent, result *Result, err error) OnRunRestore(ctx context.Context, agent *Agent, checkpoint *Checkpoint) + OnRunSnapshot(ctx context.Context, agent *Agent, checkpoint *Checkpoint) OnLLMStart(ctx context.Context, agent *Agent, messages []llm.Message) OnLLMEnd(ctx context.Context, agent *Agent, response *llm.ChatCompletionResponse, err error) OnToolStart(ctx context.Context, agent *Agent, tool Tool, arguments string) @@ -41,6 +42,7 @@ var _ RunHooks = NoOpHooks{} func (NoOpHooks) OnRunStart(context.Context, *Agent, []llm.Message) {} func (NoOpHooks) OnRunEnd(context.Context, *Agent, *Result, error) {} func (NoOpHooks) OnRunRestore(context.Context, *Agent, *Checkpoint) {} +func (NoOpHooks) OnRunSnapshot(context.Context, *Agent, *Checkpoint) {} func (NoOpHooks) OnLLMStart(context.Context, *Agent, []llm.Message) {} func (NoOpHooks) OnLLMEnd(context.Context, *Agent, *llm.ChatCompletionResponse, error) {} func (NoOpHooks) OnToolStart(context.Context, *Agent, Tool, string) {} diff --git a/pkg/agent/restore.go b/pkg/agent/restore.go index bd7912318..50b6f6a59 100644 --- a/pkg/agent/restore.go +++ b/pkg/agent/restore.go @@ -250,6 +250,7 @@ func restoreNestedSuspended( if err := store.Save(ctx, runID, &next); err != nil { return nil, fmt.Errorf("cannot save nested restore progress: %w", err) } + emitHook(agent, func(h RunHooks) { h.OnRunSnapshot(ctx, agent, &next) }) } return &next, nil } diff --git a/pkg/agent/run.go b/pkg/agent/run.go index ccac27aa0..21f2533dd 100644 --- a/pkg/agent/run.go +++ b/pkg/agent/run.go @@ -386,6 +386,8 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil { s.logger.ErrorCtx(ctx, "cannot save suspension checkpoint", log.Error(saveErr)) se.Checkpoint = cp + } else { + emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, cp) }) } } else { se.Checkpoint = cp @@ -559,6 +561,8 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag if s.opts.checkpointer != nil && s.opts.runID != "" { if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, outerCP); saveErr != nil { s.logger.ErrorCtx(ctx, "cannot save checkpoint", log.Error(saveErr)) + } else { + emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, outerCP) }) } } return s.finishRun(ctx, nil, &SuspendedError{RunID: s.opts.runID, Checkpoint: outerCP}) @@ -580,6 +584,8 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag cp.PendingApprovals = nae.pendingApprovals if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil { s.logger.ErrorCtx(ctx, "cannot save approval checkpoint", log.Error(saveErr)) + } else { + emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, cp) }) } } @@ -627,6 +633,8 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag } if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil { s.logger.ErrorCtx(ctx, "cannot save nested approval checkpoint", log.Error(saveErr)) + } else { + emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, cp) }) } } @@ -695,6 +703,8 @@ func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Messag cp := s.buildCheckpoint(AgentStatusSuspended) if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil { s.logger.ErrorCtx(ctx, "cannot save checkpoint", log.Error(saveErr)) + } else { + emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, cp) }) } }