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

cmd/root/root.go

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

cmd/root/sessions.go

@@ -0,0 +1,161 @@
+package root
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"log/slog"
+	"strings"
+
+	"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"
+	"github.com/docker/docker-agent/pkg/session/sqlitestore"
+)
+
+type sessionsDiffFlags struct {
+	sessionDB  string
+	asJSON     bool
+	failOnDiff bool
+}
+
+// newSessionsCmd groups session-inspection subcommands.
+//
+// Deliberately not called "replay": pkg/recording already owns that word for
+// recording and replaying API interactions, and `--record` writes cassettes.
+// This command replays nothing — it diffs two recordings — and naming it replay
+// would also take the word from the re-run-against-another-model feature that
+// actually is a replay.
+func newSessionsCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "sessions",
+		Short:   "Inspect recorded sessions",
+		GroupID: "advanced",
+	}
+	cmd.AddCommand(newSessionsDiffCmd())
+	return cmd
+}
+
+func newSessionsDiffCmd() *cobra.Command {
+	var flags sessionsDiffFlags
+
+	cmd := &cobra.Command{
+		Use:   "diff <session-a> <session-b>",
+		Short: "Compare the behaviour of two recorded sessions",
+		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 sessions diff <session-a> <session-b>
+  docker agent sessions diff -1 -2
+  docker agent sessions diff <a> <b> --json | jq '.divergence.turn_index'
+  docker agent sessions diff <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 *sessionsDiffFlags) 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 := loadSessionRef(ctx, store, args[0])
+	if err != nil {
+		return err
+	}
+	sessB, err := loadSessionRef(ctx, store, args[1])
+	if err != nil {
+		return 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
+}
+
+// loadSessionRef resolves a user-supplied reference and loads the session.
+//
+// References go through session.ResolveSessionID like every other
+// session-consuming command, so relative forms work — "compare my last two
+// runs" is `sessions diff -1 -2`. An unambiguous ID prefix is accepted too,
+// since full UUIDs are the hardest thing for a user to produce by hand.
+func loadSessionRef(ctx context.Context, store session.Store, ref string) (*session.Session, error) {
+	id, err := session.ResolveSessionID(ctx, store, ref)
+	if err != nil {
+		return nil, err
+	}
+	if sess, err := store.GetSession(ctx, id); err == nil {
+		return sess, nil
+	}
+
+	summaries, err := store.GetSessionSummaries(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("listing sessions: %w", err)
+	}
+	var matches []string
+	for _, summary := range summaries {
+		if strings.HasPrefix(summary.ID, id) {
+			matches = append(matches, summary.ID)
+		}
+	}
+	switch len(matches) {
+	case 1:
+		sess, err := store.GetSession(ctx, matches[0])
+		if err != nil {
+			return nil, fmt.Errorf("reading session %q: %w", ref, err)
+		}
+		return sess, nil
+	case 0:
+		return nil, fmt.Errorf("no session matches %q", ref)
+	default:
+		return nil, fmt.Errorf("%q matches %d sessions; use more characters", ref, len(matches))
+	}
+}
+
+// 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/sessions_test.go

@@ -0,0 +1,185 @@
+package root
+
+import (
+	"bytes"
+	"encoding/json"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"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/session/sqlitestore"
+	"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 TestSessionsDiffCmd_FlagsAreRegistered(t *testing.T) {
+	t.Parallel()
+
+	cmd := newSessionsDiffCmd()
+	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"}))
+}
+
+// diffFixture builds a real store with two sessions whose behaviour differs.
+func diffFixture(t *testing.T) (string, *session.Session, *session.Session) {
+	t.Helper()
+
+	dbPath := filepath.Join(t.TempDir(), "s.db")
+	store, err := sqlitestore.New(t.Context(), dbPath)
+	require.NoError(t, err)
+
+	mk := func(id, toolName string, created time.Time) *session.Session {
+		s := &session.Session{ID: id, CreatedAt: created, Messages: []session.Item{
+			{Message: &session.Message{Message: chat.Message{
+				Role: chat.MessageRoleAssistant,
+				ToolCalls: []tools.ToolCall{
+					{Function: tools.FunctionCall{Name: toolName, Arguments: "{}"}},
+				},
+			}}},
+		}}
+		require.NoError(t, store.AddSession(t.Context(), s))
+		return s
+	}
+
+	now := time.Now()
+	a := mk("aaaaaaaa11111111", "read_file", now.Add(-2*time.Hour))
+	b := mk("bbbbbbbb22222222", "shell", now.Add(-time.Hour))
+	require.NoError(t, store.Close())
+
+	return dbPath, a, b
+}
+
+func runSessionsDiff(t *testing.T, dbPath string, args ...string) (string, error) {
+	t.Helper()
+
+	cmd := newSessionsDiffCmd()
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+	cmd.SetErr(&out)
+	cmd.SetContext(t.Context())
+	require.NoError(t, cmd.Flags().Set("session-db", dbPath))
+	for i := 0; i+1 < len(args); i += 2 {
+		require.NoError(t, cmd.Flags().Set(args[i], args[i+1]))
+	}
+	// Run first: operands of a return statement are evaluated left to right, so
+	// reading the buffer in the same statement would capture it before the run.
+	err := cmd.RunE(cmd, []string{"aaaaaaaa11111111", "bbbbbbbb22222222"})
+	return out.String(), err
+}
+
+// The non-zero exit is the whole contract of --fail-on-divergence for the CI
+// use case it exists for.
+func TestSessionsDiff_FailOnDivergence(t *testing.T) {
+	t.Parallel()
+
+	dbPath, _, _ := diffFixture(t)
+
+	out, err := runSessionsDiff(t, dbPath)
+	require.NoError(t, err, "without the flag a divergence is reported but does not fail")
+	assert.Contains(t, out, "First divergence")
+
+	out, err = runSessionsDiff(t, dbPath, "fail-on-divergence", "true")
+	require.Error(t, err, "with the flag a divergence must exit non-zero")
+	assert.Contains(t, err.Error(), "diverged")
+	assert.Contains(t, out, "First divergence")
+}
+
+// References go through ResolveSessionID like every other session command, so
+// "compare my last two runs" works.
+func TestSessionsDiff_ResolvesRelativeAndPrefixRefs(t *testing.T) {
+	t.Parallel()
+
+	dbPath, a, b := diffFixture(t)
+
+	store, err := sqlitestore.New(t.Context(), dbPath)
+	require.NoError(t, err)
+	defer func() { require.NoError(t, store.Close()) }()
+
+	// Relative.
+	for _, ref := range []string{"-1", "-2"} {
+		got, err := loadSessionRef(t.Context(), store, ref)
+		require.NoErrorf(t, err, "relative ref %q must resolve", ref)
+		require.NotNil(t, got)
+	}
+
+	// Prefix.
+	got, err := loadSessionRef(t.Context(), store, a.ID[:8])
+	require.NoError(t, err)
+	assert.Equal(t, a.ID, got.ID)
+
+	got, err = loadSessionRef(t.Context(), store, b.ID)
+	require.NoError(t, err)
+	assert.Equal(t, b.ID, got.ID)
+
+	_, err = loadSessionRef(t.Context(), store, "nosuchsession")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "no session matches")
+}
+
+func TestSessionsCmd_HasDiffSubcommand(t *testing.T) {
+	t.Parallel()
+
+	cmd := newSessionsCmd()
+	assert.Equal(t, "sessions", cmd.Name())
+
+	var names []string
+	for _, sub := range cmd.Commands() {
+		names = append(names, sub.Name())
+	}
+	assert.Contains(t, names, "diff")
+}