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