feat(cli): show snapshots in task logs (#21787)

cli/task_test.go

@@ -20,7 +20,11 @@ import (
 	"github.com/coder/coder/v2/agent"
 	"github.com/coder/coder/v2/agent/agenttest"
 	"github.com/coder/coder/v2/cli/clitest"
+	"github.com/coder/coder/v2/coderd"
 	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/coder/v2/coderd/database/dbfake"
 	"github.com/coder/coder/v2/coderd/util/ptr"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
@@ -271,6 +275,93 @@ func setupCLITaskTest(ctx context.Context, t *testing.T, agentAPIHandlers map[st
 	return userClient, task
 }
 
+// setupCLITaskTestWithSnapshot creates a task in the specified status with a log snapshot.
+func setupCLITaskTestWithSnapshot(ctx context.Context, t *testing.T, status codersdk.TaskStatus, messages []agentapisdk.Message) (*codersdk.Client, codersdk.Task) {
+	t.Helper()
+
+	ownerClient, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+	owner := coderdtest.CreateFirstUser(t, ownerClient)
+	userClient, user := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID)
+
+	ownerUser, err := ownerClient.User(ctx, owner.UserID.String())
+	require.NoError(t, err)
+	ownerSubject := coderdtest.AuthzUserSubject(ownerUser)
+
+	task := createTaskInStatus(t, db, owner.OrganizationID, user.ID, status)
+
+	// Create snapshot envelope with agentapi format.
+	envelope := coderd.TaskLogSnapshotEnvelope{
+		Format: "agentapi",
+		Data: agentapisdk.GetMessagesResponse{
+			Messages: messages,
+		},
+	}
+	snapshotJSON, err := json.Marshal(envelope)
+	require.NoError(t, err)
+
+	// Insert snapshot into database.
+	snapshotTime := time.Now()
+	err = db.UpsertTaskSnapshot(dbauthz.As(ctx, ownerSubject), database.UpsertTaskSnapshotParams{
+		TaskID:               task.ID,
+		LogSnapshot:          json.RawMessage(snapshotJSON),
+		LogSnapshotCreatedAt: snapshotTime,
+	})
+	require.NoError(t, err)
+
+	return userClient, task
+}
+
+// setupCLITaskTestWithoutSnapshot creates a task in the specified status without a log snapshot.
+func setupCLITaskTestWithoutSnapshot(t *testing.T, status codersdk.TaskStatus) (*codersdk.Client, codersdk.Task) {
+	t.Helper()
+
+	ownerClient, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+	owner := coderdtest.CreateFirstUser(t, ownerClient)
+	userClient, user := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID)
+
+	task := createTaskInStatus(t, db, owner.OrganizationID, user.ID, status)
+
+	return userClient, task
+}
+
+// createTaskInStatus creates a task in the specified status using dbfake.
+func createTaskInStatus(t *testing.T, db database.Store, orgID, ownerID uuid.UUID, status codersdk.TaskStatus) codersdk.Task {
+	t.Helper()
+
+	builder := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
+		OrganizationID: orgID,
+		OwnerID:        ownerID,
+	}).
+		WithTask(database.TaskTable{
+			OrganizationID: orgID,
+			OwnerID:        ownerID,
+		}, nil)
+
+	switch status {
+	case codersdk.TaskStatusPending:
+		builder = builder.Pending()
+	case codersdk.TaskStatusInitializing:
+		builder = builder.Starting()
+	case codersdk.TaskStatusPaused:
+		builder = builder.Seed(database.WorkspaceBuild{
+			Transition: database.WorkspaceTransitionStop,
+		})
+	default:
+		require.Fail(t, "unsupported task status in test helper", "status: %s", status)
+	}
+
+	resp := builder.Do()
+
+	return codersdk.Task{
+		ID:             resp.Task.ID,
+		Name:           resp.Task.Name,
+		OrganizationID: resp.Task.OrganizationID,
+		OwnerID:        resp.Task.OwnerID,
+		WorkspaceID:    resp.Task.WorkspaceID,
+		Status:         status,
+	}
+}
+
 // createAITaskTemplate creates a template configured for AI tasks with a sidebar app.
 func createAITaskTemplate(t *testing.T, client *codersdk.Client, orgID uuid.UUID, opts ...aiTemplateOpt) codersdk.Template {
 	t.Helper()