From f764761ac57e4b7e3da122bce1d41cc918cf10c9 Mon Sep 17 00:00:00 2001 From: Dmitry Kolesnikov Date: Tue, 28 Apr 2026 07:48:08 +0300 Subject: [PATCH] (fix) optional formatting for done i/o at Arr[S] --- agent/nanobot/runtime.go | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/agent/nanobot/runtime.go b/agent/nanobot/runtime.go index 5d57460..160982a 100644 --- a/agent/nanobot/runtime.go +++ b/agent/nanobot/runtime.go @@ -192,15 +192,21 @@ func (f Arr[S]) WithTask(name string) Arr[S] { return f.WithTaskf(func(S) string // WithTaskf wraps the arrow with a progress report. // The task name is generated by applying fn to the current state S at the time of execution. // The task is automatically marked done when the arrow returns, even if it returns an error. -func (f Arr[S]) WithTaskf(fn func(S) string) Arr[S] { +func (f Arr[S]) WithTaskf(taskf func(S) string, donef ...func(S) string) Arr[S] { return func(ctx context.Context, s S, opt ...chatter.Opt) (S, error) { c, ok := ctx.Value(chalkboard).(Chalk) - if !ok || c == nil || fn == nil { + if !ok || c == nil || taskf == nil { return f(ctx, s, opt...) } - c.Task(ctx, fn(s)) - defer c.Done() + c.Task(ctx, taskf(s)) + defer func() { + if len(donef) > 0 && donef[0] != nil { + c.Done(donef[0](s)) + } else { + c.Done() + } + }() return f(c.Sub(ctx), s, opt...) } }