(fix) console output abstraction

agent/nanobot/nanobot_test.go

@@ -12,6 +12,7 @@ import (
 	"context"
 	"errors"
 	"testing"
+	"testing/fstest"
 
 	"github.com/fogfish/it/v2"
 	"github.com/kshard/chatter"
@@ -66,9 +67,15 @@ type MockChalk struct {
 	tasks  []string
 	dones  int
 	failed []error
+	subs   int
 }
 
-func (c *MockChalk) Sub(ctx context.Context) context.Context { return ctx }
+type chalkCtxKey struct{}
+
+func (c *MockChalk) Sub(ctx context.Context) context.Context {
+	c.subs++
+	return context.WithValue(ctx, chalkCtxKey{}, true)
+}
 func (c *MockChalk) Task(_ context.Context, format string, _ ...any) {
 	c.tasks = append(c.tasks, format)
 }
@@ -103,15 +110,6 @@ func TestRuntime(t *testing.T) {
 		it.Then(t).ShouldNot(it.Nil(rt2))
 		it.Then(t).ShouldNot(it.Equal(rt, rt2))
 	})
-
-	t.Run("WithStdout", func(t *testing.T) {
-		rt := nanobot.NewRuntime(nil, nil)
-		chalk := &MockChalk{}
-		rt2 := rt.WithStdout(chalk)
-
-		it.Then(t).ShouldNot(it.Nil(rt2))
-		it.Then(t).ShouldNot(it.Equal(rt, rt2))
-	})
 }
 
 // =============================================================================
@@ -348,6 +346,62 @@ func TestWhen(t *testing.T) {
 	})
 }
 
+// =============================================================================
+// TestWithTask
+// =============================================================================
+
+func TestArrWithTask(t *testing.T) {
+	t.Run("Success", func(t *testing.T) {
+		chalk := &MockChalk{}
+		//lint:ignore SA1029 We use string keys to allow zero-dep discovery
+		ctx := context.WithValue(context.Background(), "io.console.chalkboard", chalk)
+
+		arr := nanobot.Lift(func(ctx context.Context, w Work) (Work, error) {
+			if ok, _ := ctx.Value(chalkCtxKey{}).(bool); !ok {
+				t.Fatalf("expected chalk sub-context to be used")
+			}
+			w.Result = "done"
+			return w, nil
+		}).WithTask("arr-step")
+
+		result, err := arr.Prompt(ctx, Work{})
+		it.Then(t).Should(
+			it.Nil(err),
+			it.Equal(result.Result, "done"),
+			it.Equal(len(chalk.tasks), 1),
+			it.Equal(chalk.tasks[0], "arr-step"),
+			it.Equal(chalk.dones, 1),
+			it.Equal(len(chalk.failed), 0),
+			it.Equal(chalk.subs, 1),
+		)
+	})
+
+	t.Run("ErrorStillDone", func(t *testing.T) {
+		errArr := errors.New("arr error")
+		chalk := &MockChalk{}
+		//lint:ignore SA1029 We use string keys to allow zero-dep discovery
+		ctx := context.WithValue(context.Background(), "io.console.chalkboard", chalk)
+
+		arr := nanobot.Lift(func(ctx context.Context, w Work) (Work, error) {
+			if ok, _ := ctx.Value(chalkCtxKey{}).(bool); !ok {
+				t.Fatalf("expected chalk sub-context to be used")
+			}
+			return w, errArr
+		}).WithTask("arr-step")
+
+		_, err := arr.Prompt(ctx, Work{})
+		it.Then(t).ShouldNot(it.Nil(err))
+		it.Then(t).Should(
+			it.True(errors.Is(err, errArr)),
+			it.Equal(len(chalk.tasks), 1),
+			it.Equal(chalk.tasks[0], "arr-step"),
+			it.Equal(chalk.dones, 1),
+			it.Equal(len(chalk.failed), 0),
+			it.Equal(chalk.subs, 1),
+		)
+	})
+}
+
 // =============================================================================
 // TestReflect
 // =============================================================================
@@ -561,31 +615,6 @@ func TestReflect(t *testing.T) {
 		it.Then(t).ShouldNot(it.Nil(err))
 		it.Then(t).Should(it.True(errors.Is(err, errEffect)))
 	})
-
-	t.Run("WithChalk", func(t *testing.T) {
-		chalk := &MockChalk{}
-		rt2 := nanobot.NewRuntime(nil, nil).WithStdout(chalk)
-
-		rawJudge := &MockBot[Work, string]{
-			fn: func(_ context.Context, w Work, _ ...chatter.Opt) (string, error) {
-				return w.Result, nil
-			},
-		}
-		react := &MockBot[Work, string]{
-			fn: func(_ context.Context, _ Work, _ ...chatter.Opt) (string, error) {
-				return "fix", nil
-			},
-		}
-
-		// no eff: always-accept on first attempt
-		judge := nanobot.Judge[Work, string](rawJudge)
-
-		bot, err := nanobot.NewReflect(rt2, judge, nanobot.Arrow[Work, string](react))
-		it.Then(t).Should(it.Nil(err))
-		_, err = bot.Prompt(context.Background(), Work{Result: "state"})
-		it.Then(t).Should(it.Nil(err))
-		it.Then(t).Should(it.True(chalk.dones >= 2))
-	})
 }
 
 // =============================================================================
@@ -851,3 +880,62 @@ func TestJsonify(t *testing.T) {
 		it.Then(t).ShouldNot(it.Nil(err))
 	})
 }
+
+// =============================================================================
+// TestReActWithTask
+// =============================================================================
+
+func TestReActWithTask(t *testing.T) {
+	newBot := func(t *testing.T, llm chatter.Chatter) *nanobot.BotReAct[Work, string] {
+		t.Helper()
+
+		fs := fstest.MapFS{
+			"react.prompt": &fstest.MapFile{
+				Data: []byte("Return {{.Result}}"),
+			},
+		}
+
+		bot, err := nanobot.NewReAct[Work, string](nanobot.NewRuntime(fs, &MockLLMs{
+			models: map[string]chatter.Chatter{"base": llm},
+		}), "react.prompt")
+		it.Then(t).Should(it.Nil(err))
+
+		return bot.WithTask("react-step")
+	}
+
+	t.Run("Success", func(t *testing.T) {
+		chalk := &MockChalk{}
+		//lint:ignore SA1029 We use string keys to allow zero-dep discovery
+		ctx := context.WithValue(context.Background(), "io.console.chalkboard", chalk)
+		bot := newBot(t, &MockChatter{response: "final answer"})
+
+		result, err := bot.Prompt(ctx, Work{Result: "input"})
+		it.Then(t).Should(
+			it.Nil(err),
+			it.Equal(result, "final answer"),
+			it.Equal(len(chalk.tasks), 1),
+			it.Equal(chalk.tasks[0], "react-step"),
+			it.Equal(chalk.dones, 1),
+			it.Equal(len(chalk.failed), 0),
+		)
+	})
+
+	t.Run("FailureCallsFail", func(t *testing.T) {
+		errLLM := errors.New("llm failure")
+		chalk := &MockChalk{}
+		//lint:ignore SA1029 We use string keys to allow zero-dep discovery
+		ctx := context.WithValue(context.Background(), "io.console.chalkboard", chalk)
+		bot := newBot(t, &MockChatter{err: errLLM})
+
+		_, err := bot.Prompt(ctx, Work{Result: "input"})
+		it.Then(t).ShouldNot(it.Nil(err))
+		it.Then(t).Should(
+			it.True(errors.Is(err, errLLM)),
+			it.Equal(len(chalk.tasks), 1),
+			it.Equal(chalk.tasks[0], "react-step"),
+			it.Equal(chalk.dones, 0),
+			it.Equal(len(chalk.failed), 1),
+			it.True(errors.Is(chalk.failed[0], errLLM)),
+		)
+	})
+}

agent/nanobot/react.go

@@ -41,7 +41,8 @@ type BotReAct[A, B any] struct {
 	registry *command.SeqRegistry
 	prompt   *prompt.Prompt
 	t        *template.Template
-	chalk    Chalk
+	taskf    func(A) string
+	donef    func(B) string
 }
 
 // ReAct is like NewReAct but panics on error.
@@ -99,10 +100,7 @@ func NewReAct[A, B any](rt *Runtime, file string) (*BotReAct[A, B], error) {
 		runner = aio.NewJsonLogger(os.Stderr, runner)
 	}
 
-	bot := &BotReAct[A, B]{prompt: prompt, t: t, chalk: rt.Chalk}
-	if len(bot.prompt.Name) == 0 {
-		bot.chalk = devnull{}
-	}
+	bot := &BotReAct[A, B]{prompt: prompt, t: t}
 
 	bot.registry = command.NewSeqRegistry()
 	bot.registry.Bind(registry)
@@ -131,6 +129,18 @@ func (bot *BotReAct[A, B]) WithRegistry(r *command.Registry) *BotReAct[A, B] {
 	return bot
 }
 
+func (bot *BotReAct[A, B]) WithTask(name string) *BotReAct[A, B] {
+	return bot.WithTaskf(func(A) string { return name })
+}
+
+func (bot *BotReAct[A, B]) WithTaskf(taskf func(A) string, donef ...func(B) string) *BotReAct[A, B] {
+	bot.taskf = taskf
+	if len(donef) > 0 {
+		bot.donef = donef[0]
+	}
+	return bot
+}
+
 // Prompt encodes the input using the prompt template, runs the Manifold
 // ReAct loop until the model returns a final answer, and decodes the result
 // into B. Progress is reported via the Chalk sink when the prompt file
@@ -140,17 +150,22 @@ func (bot *BotReAct[A, B]) Prompt(ctx context.Context, input A, opt ...chatter.O
 		bot.memory.Reset()
 	}
 
-	bot.chalk.Task(ctx, bot.prompt.Name)
+	chalk, ok := ctx.Value(chalkboard).(Chalk)
+	if !ok || chalk == nil || bot.taskf == nil {
+		return bot.manifold.Prompt(ctx, input, opt...)
+	}
+
+	chalk.Task(ctx, bot.taskf(input))
 	val, err := bot.manifold.Prompt(ctx, input, opt...)
 	if err != nil {
-		if len(bot.prompt.Name) > 0 {
-			bot.chalk.Fail(err)
-		}
+		chalk.Fail(err)
 		return val, err
 	}
 
-	if len(bot.prompt.Name) > 0 {
-		bot.chalk.Done()
+	if bot.donef != nil {
+		chalk.Done(bot.donef(val))
+	} else {
+		chalk.Done()
 	}
 
 	return val, nil

agent/nanobot/reflect.go

@@ -114,7 +114,6 @@ type BotReflect[S any] struct {
 	judge    Bot[S, Vote[S]]
 	correct  Arr[S]
 	attempts int
-	chalk    Chalk
 }
 
 // Reflect creates a BotReflect. Panics on error.
@@ -133,7 +132,6 @@ func NewReflect[S any](rt *Runtime, judge Bot[S, Vote[S]], correct Arr[S]) (*Bot
 		judge:    judge,
 		correct:  correct,
 		attempts: 1,
-		chalk:    rt.Chalk,
 	}, nil
 }
 
@@ -149,43 +147,29 @@ func (bot *BotReflect[S]) WithAttempts(attempts int) *BotReflect[S] {
 // correct arrow and the loop retries. An error is returned on a hard reject
 // or when the retry budget is exhausted.
 func (bot *BotReflect[S]) Prompt(ctx context.Context, input S, opt ...chatter.Opt) (S, error) {
-	bot.chalk.Task(ctx, "Reflect (%d attempts)", bot.attempts)
-
 	s := input
-	for i := range bot.attempts {
-		bot.chalk.Task(bot.chalk.Sub(ctx), "attempt %d of %d", i+1, bot.attempts)
-
-		v, err := bot.judge.Prompt(bot.chalk.Sub(ctx), s, opt...)
+	for range bot.attempts {
+		v, err := bot.judge.Prompt(ctx, s, opt...)
 		if err != nil {
-			bot.chalk.Fail(err)
 			return *new(S), err
 		}
 
 		switch {
 		case v.accepted > 0:
-			bot.chalk.Done()
-			bot.chalk.Done()
 			return v.State, nil
 
 		case v.accepted < 0:
-			bot.chalk.Done()
 			err := fmt.Errorf("rejected")
-			bot.chalk.Fail(err)
 			return v.State, err
 		}
 
 		// neutral: v.State carries the critique — forward to corrector
-		s, err = bot.correct(bot.chalk.Sub(ctx), v.State, opt...)
+		s, err = bot.correct(ctx, v.State, opt...)
 		if err != nil {
-			bot.chalk.Done()
-			bot.chalk.Fail(err)
 			return s, err
 		}
-
-		bot.chalk.Done()
 	}
 
 	err := fmt.Errorf("rejected")
-	bot.chalk.Fail(err)
 	return s, err
 }