feat(replay): compare the behaviour of two recorded sessions

cmd/root/replay.go

@@ -0,0 +1,102 @@
+package root
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"log/slog"
+
+	"github.com/spf13/cobra"
+
+	pathx "github.com/docker/docker-agent/pkg/path"
+	"github.com/docker/docker-agent/pkg/replay"
+	"github.com/docker/docker-agent/pkg/session/sqlitestore"
+)
+
+type replayFlags struct {
+	sessionDB  string
+	asJSON     bool
+	failOnDiff bool
+}
+
+func newReplayCmd() *cobra.Command {
+	var flags replayFlags
+
+	cmd := &cobra.Command{
+		Use:     "replay <session-a> <session-b>",
+		Short:   "Compare the behaviour of two recorded sessions",
+		GroupID: "advanced",
+		Args:    cobra.ExactArgs(2),
+		Long: `Compare two recorded sessions and report the first point where the agent
+behaved differently.
+
+Comparison is over the sequence of tool calls, not over the assistant's prose.
+Model output is nondeterministic: two runs of the same task almost always word
+things differently while doing exactly the same work, so diffing text would report
+a difference on every comparison. The tool calls are what changed the world, so
+they are what is compared.
+
+Reporting stops at the first divergence: everything after it is downstream of that
+difference and comparing it produces noise rather than information.`,
+		Example: `  docker agent replay <session-a> <session-b>
+  docker agent replay <a> <b> --json | jq '.divergence.turn_index'
+  docker agent replay <a> <b> --fail-on-divergence`,
+		RunE: flags.run,
+	}
+
+	cmd.Flags().StringVarP(&flags.sessionDB, "session-db", "s", "", "Path to the session database (default: <data-dir>/session.db)")
+	cmd.Flags().BoolVar(&flags.asJSON, "json", false, "Emit the comparison as JSON")
+	cmd.Flags().BoolVar(&flags.failOnDiff, "fail-on-divergence", false, "Exit non-zero when the two sessions diverge")
+
+	return cmd
+}
+
+func (f *replayFlags) run(cmd *cobra.Command, args []string) error {
+	ctx := cmd.Context()
+
+	dbPath, err := pathx.ExpandHomeDir(sessionDBPath(f.sessionDB))
+	if err != nil {
+		return err
+	}
+
+	store, err := sqlitestore.New(ctx, dbPath)
+	if err != nil {
+		return fmt.Errorf("opening session store: %w", err)
+	}
+	defer func() {
+		if err := store.Close(); err != nil {
+			slog.ErrorContext(ctx, "Failed to close session store", "error", err)
+		}
+	}()
+
+	sessA, err := store.GetSession(ctx, args[0])
+	if err != nil {
+		return fmt.Errorf("reading session %q: %w", args[0], err)
+	}
+	sessB, err := store.GetSession(ctx, args[1])
+	if err != nil {
+		return fmt.Errorf("reading session %q: %w", args[1], err)
+	}
+
+	result := replay.CompareSessions(sessA, sessB)
+	if err := renderReplay(cmd.OutOrStdout(), result, args[0], args[1], f.asJSON); err != nil {
+		return err
+	}
+
+	if f.failOnDiff && !result.Identical() {
+		return errors.New("sessions diverged")
+	}
+	return nil
+}
+
+// renderReplay writes the comparison as JSON or as text.
+func renderReplay(w io.Writer, result replay.Result, nameA, nameB string, asJSON bool) error {
+	if asJSON {
+		enc := json.NewEncoder(w)
+		enc.SetIndent("", "  ")
+		return enc.Encode(result)
+	}
+	replay.PrintResult(w, result, nameA, nameB)
+	return nil
+}

cmd/root/replay_test.go

@@ -0,0 +1,73 @@
+package root
+
+import (
+	"bytes"
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/chat"
+	"github.com/docker/docker-agent/pkg/replay"
+	"github.com/docker/docker-agent/pkg/session"
+	"github.com/docker/docker-agent/pkg/tools"
+)
+
+func replaySession(toolName string) *session.Session {
+	return &session.Session{Messages: []session.Item{
+		{Message: &session.Message{Message: chat.Message{
+			Role: chat.MessageRoleAssistant,
+			ToolCalls: []tools.ToolCall{
+				{Function: tools.FunctionCall{Name: toolName, Arguments: "{}"}},
+			},
+		}}},
+	}}
+}
+
+func TestRenderReplay_TextDivergence(t *testing.T) {
+	t.Parallel()
+
+	result := replay.CompareSessions(replaySession("read_file"), replaySession("shell"))
+	var buf bytes.Buffer
+	require.NoError(t, renderReplay(&buf, result, "aaa", "bbb", false))
+
+	out := buf.String()
+	assert.Contains(t, out, "First divergence at turn 0")
+	assert.Contains(t, out, "aaa")
+	assert.Contains(t, out, "bbb")
+}
+
+func TestRenderReplay_TextIdentical(t *testing.T) {
+	t.Parallel()
+
+	result := replay.CompareSessions(replaySession("read_file"), replaySession("read_file"))
+	var buf bytes.Buffer
+	require.NoError(t, renderReplay(&buf, result, "aaa", "bbb", false))
+	assert.Contains(t, buf.String(), "Identical behaviour")
+}
+
+func TestRenderReplay_JSON(t *testing.T) {
+	t.Parallel()
+
+	result := replay.CompareSessions(replaySession("read_file"), replaySession("shell"))
+	var buf bytes.Buffer
+	require.NoError(t, renderReplay(&buf, result, "aaa", "bbb", true))
+
+	var round replay.Result
+	require.NoError(t, json.Unmarshal(buf.Bytes(), &round))
+	require.NotNil(t, round.Divergence)
+	assert.Equal(t, 0, round.Divergence.TurnIndex)
+}
+
+func TestReplayCmd_FlagsAreRegistered(t *testing.T) {
+	t.Parallel()
+
+	cmd := newReplayCmd()
+	for _, name := range []string{"session-db", "json", "fail-on-divergence"} {
+		assert.NotNilf(t, cmd.Flags().Lookup(name), "flag %q must exist", name)
+	}
+	// Two session IDs, no more, no fewer.
+	require.Error(t, cmd.Args(cmd, []string{"only-one"}))
+	require.NoError(t, cmd.Args(cmd, []string{"a", "b"}))
+}

cmd/root/root.go

@@ -177,6 +177,7 @@ We collect anonymous usage data to help improve docker agent. To disable:
 		newNewCmd(),
 		newGettingStartedCmd(),
 		newEvalCmd(),
+		newReplayCmd(),
 		newShareCmd(),
 		newModelsCmd(),
 		newToolsetsCmd(),