feat(coderd): use task data model for send/logs (#20381)

Updates coder/internal#976

coderd/aitasks_test.go

@@ -2,7 +2,7 @@ package coderd_test
 
 import (
 	"database/sql"
-	"fmt"
+	"encoding/json"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -15,6 +15,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	agentapisdk "github.com/coder/agentapi-sdk-go"
 	"github.com/coder/coder/v2/agent"
 	"github.com/coder/coder/v2/agent/agenttest"
 	"github.com/coder/coder/v2/coderd/coderdtest"
@@ -210,8 +211,8 @@ func TestTasks(t *testing.T) {
 											Apps: []*proto.App{
 												{
 													Id:          taskAppID.String(),
-													Slug:        "task-sidebar",
-													DisplayName: "Task Sidebar",
+													Slug:        "task-app",
+													DisplayName: "Task App",
 													Url:         opt.appURL,
 												},
 											},
@@ -221,9 +222,7 @@ func TestTasks(t *testing.T) {
 							},
 							AiTasks: []*proto.AITask{
 								{
-									SidebarApp: &proto.AITaskSidebarApp{
-										Id: taskAppID.String(),
-									},
+									AppId: taskAppID.String(),
 								},
 							},
 						},
@@ -469,41 +468,40 @@ func TestTasks(t *testing.T) {
 	t.Run("Send", func(t *testing.T) {
 		t.Parallel()
 
-		t.Skip("TODO(mafredri): Remove, fixed down-stack!")
-
 		t.Run("IntegrationOK", func(t *testing.T) {
 			t.Parallel()
 
-			client, _, api := coderdtest.NewWithAPI(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
-			owner := coderdtest.CreateFirstUser(t, client)
-			userClient, _ := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
-
-			createStatusResponse := func(status string) string {
-				return `
-					{
-						"$schema": "http://localhost:3284/schemas/StatusResponseBody.json",
-						"status": "` + status + `"
-					}
-				`
-			}
-			statusResponse := createStatusResponse("stable")
+			statusResponse := agentapisdk.StatusStable
 
 			// Start a fake AgentAPI that accepts GET /status and POST /message.
 			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 				if r.Method == http.MethodGet && r.URL.Path == "/status" {
 					w.Header().Set("Content-Type", "application/json")
+					resp := agentapisdk.GetStatusResponse{
+						Status: statusResponse,
+					}
+					respBytes, err := json.Marshal(resp)
+					assert.NoError(t, err)
 					w.WriteHeader(http.StatusOK)
-					_, _ = fmt.Fprint(w, statusResponse)
+					w.Write(respBytes)
 					return
 				}
 				if r.Method == http.MethodPost && r.URL.Path == "/message" {
 					w.Header().Set("Content-Type", "application/json")
 
 					b, _ := io.ReadAll(r.Body)
-					assert.Equal(t, `{"content":"Hello, Agent!","type":"user"}`, string(b), "expected message content")
+					expectedReq := agentapisdk.PostMessageParams{
+						Content: "Hello, Agent!",
+						Type:    agentapisdk.MessageTypeUser,
+					}
+					expectedBytes, _ := json.Marshal(expectedReq)
+					assert.Equal(t, string(expectedBytes), string(b), "expected message content")
 
+					resp := agentapisdk.PostMessageResponse{Ok: true}
+					respBytes, err := json.Marshal(resp)
+					assert.NoError(t, err)
 					w.WriteHeader(http.StatusOK)
-					io.WriteString(w, `{"ok": true}`)
+					w.Write(respBytes)
 					return
 				}
 				w.WriteHeader(http.StatusInternalServerError)
@@ -511,103 +509,105 @@ func TestTasks(t *testing.T) {
 			defer srv.Close()
 
 			// Create an AI-capable template whose sidebar app points to our fake AgentAPI.
-			authToken := uuid.NewString()
-			template := createAITemplate(t, client, owner, withSidebarURL(srv.URL), withAgentToken(authToken))
+			var (
+				client, db     = coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+				ctx            = testutil.Context(t, testutil.WaitLong)
+				owner          = coderdtest.CreateFirstUser(t, client)
+				userClient, _  = coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
+				agentAuthToken = uuid.NewString()
+				template       = createAITemplate(t, client, owner, withAgentToken(agentAuthToken), withSidebarURL(srv.URL))
+				exp            = codersdk.NewExperimentalClient(userClient)
+			)
 
-			// Create a workspace (task) from the AI-capable template.
-			ws := coderdtest.CreateWorkspace(t, userClient, template.ID, func(req *codersdk.CreateWorkspaceRequest) {
-				req.RichParameterValues = []codersdk.WorkspaceBuildParameter{
-					{Name: codersdk.AITaskPromptParameterName, Value: "send a message"},
-				}
+			task, err := exp.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
+				TemplateVersionID: template.ActiveVersionID,
+				Input:             "send me food",
 			})
-			coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID)
+			require.NoError(t, err)
+			require.True(t, task.WorkspaceID.Valid)
+
+			// Get the workspace and wait for it to be ready.
+			ws, err := userClient.Workspace(ctx, task.WorkspaceID.UUID)
+			require.NoError(t, err)
+			coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, ws.LatestBuild.ID)
+
+			// Fetch the task by ID via experimental API and verify fields.
+			task, err = exp.TaskByID(ctx, task.ID)
+			require.NoError(t, err)
+			require.NotZero(t, task.WorkspaceBuildNumber)
+			require.True(t, task.WorkspaceAgentID.Valid)
+			require.True(t, task.WorkspaceAppID.Valid)
+
+			// Insert an app status for the workspace
+			_, err = db.InsertWorkspaceAppStatus(dbauthz.AsSystemRestricted(ctx), database.InsertWorkspaceAppStatusParams{
+				ID:          uuid.New(),
+				WorkspaceID: task.WorkspaceID.UUID,
+				CreatedAt:   dbtime.Now(),
+				AgentID:     task.WorkspaceAgentID.UUID,
+				AppID:       task.WorkspaceAppID.UUID,
+				State:       database.WorkspaceAppStatusStateComplete,
+				Message:     "all done",
+			})
+			require.NoError(t, err)
 
 			// Start a fake agent so the workspace agent is connected before sending the message.
-			agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(authToken))
-			_ = agenttest.New(t, client.URL, authToken, func(o *agent.Options) {
+			agentClient := agentsdk.New(userClient.URL, agentsdk.WithFixedToken(agentAuthToken))
+			_ = agenttest.New(t, userClient.URL, agentAuthToken, func(o *agent.Options) {
 				o.Client = agentClient
 			})
+			coderdtest.NewWorkspaceAgentWaiter(t, userClient, ws.ID).WaitFor(coderdtest.AgentsReady)
 
-			ctx := testutil.Context(t, testutil.WaitMedium)
-			coderdtest.NewWorkspaceAgentWaiter(t, client, ws.ID).WithContext(ctx).WaitFor(coderdtest.AgentsReady)
-
-			// Lookup the sidebar app ID.
-			w, err := client.Workspace(ctx, ws.ID)
+			// Fetch the task by ID via experimental API and verify fields.
+			task, err = exp.TaskByID(ctx, task.ID)
 			require.NoError(t, err)
-			var sidebarAppID uuid.UUID
-			for _, res := range w.LatestBuild.Resources {
-				for _, ag := range res.Agents {
-					for _, app := range ag.Apps {
-						if app.Slug == "task-sidebar" {
-							sidebarAppID = app.ID
-						}
-					}
-				}
-			}
-			require.NotEqual(t, uuid.Nil, sidebarAppID)
 
 			// Make the sidebar app unhealthy initially.
-			err = api.Database.UpdateWorkspaceAppHealthByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceAppHealthByIDParams{
-				ID:     sidebarAppID,
+			err = db.UpdateWorkspaceAppHealthByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceAppHealthByIDParams{
+				ID:     task.WorkspaceAppID.UUID,
 				Health: database.WorkspaceAppHealthUnhealthy,
 			})
 			require.NoError(t, err)
 
-			exp := codersdk.NewExperimentalClient(userClient)
-			err = exp.TaskSend(ctx, "me", ws.ID, codersdk.TaskSendRequest{
+			err = exp.TaskSend(ctx, "me", task.ID, codersdk.TaskSendRequest{
 				Input: "Hello, Agent!",
 			})
 			require.Error(t, err, "wanted error due to unhealthy sidebar app")
 
 			// Make the sidebar app healthy.
-			err = api.Database.UpdateWorkspaceAppHealthByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceAppHealthByIDParams{
-				ID:     sidebarAppID,
+			err = db.UpdateWorkspaceAppHealthByID(dbauthz.AsSystemRestricted(ctx), database.UpdateWorkspaceAppHealthByIDParams{
+				ID:     task.WorkspaceAppID.UUID,
 				Health: database.WorkspaceAppHealthHealthy,
 			})
 			require.NoError(t, err)
 
-			statusResponse = createStatusResponse("bad")
+			statusResponse = agentapisdk.AgentStatus("bad")
 
-			err = exp.TaskSend(ctx, "me", ws.ID, codersdk.TaskSendRequest{
+			err = exp.TaskSend(ctx, "me", task.ID, codersdk.TaskSendRequest{
 				Input: "Hello, Agent!",
 			})
 			require.Error(t, err, "wanted error due to bad status")
 
-			statusResponse = createStatusResponse("stable")
+			statusResponse = agentapisdk.StatusStable
 
-			// Send task input to the tasks sidebar app and expect 204.e
-			err = exp.TaskSend(ctx, "me", ws.ID, codersdk.TaskSendRequest{
-				Input: "Hello, Agent!",
-			})
-			require.NoError(t, err, "wanted no error due to healthy sidebar app and stable status")
-		})
-
-		t.Run("MissingContent", func(t *testing.T) {
-			t.Parallel()
-
-			client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
-			user := coderdtest.CreateFirstUser(t, client)
-			ctx := testutil.Context(t, testutil.WaitLong)
-
-			template := createAITemplate(t, client, user)
-
-			// Create a workspace (task).
-			ws := coderdtest.CreateWorkspace(t, client, template.ID, func(req *codersdk.CreateWorkspaceRequest) {
-				req.RichParameterValues = []codersdk.WorkspaceBuildParameter{
-					{Name: codersdk.AITaskPromptParameterName, Value: "do work"},
-				}
-			})
-			coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID)
-
-			exp := codersdk.NewExperimentalClient(client)
-			err := exp.TaskSend(ctx, "me", ws.ID, codersdk.TaskSendRequest{
-				Input: "",
+			//nolint:tparallel // Not intended to run in parallel.
+			t.Run("SendOK", func(t *testing.T) {
+				err = exp.TaskSend(ctx, "me", task.ID, codersdk.TaskSendRequest{
+					Input: "Hello, Agent!",
+				})
+				require.NoError(t, err, "wanted no error due to healthy sidebar app and stable status")
 			})
 
-			var sdkErr *codersdk.Error
-			require.Error(t, err)
-			require.ErrorAs(t, err, &sdkErr)
-			require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
+			//nolint:tparallel // Not intended to run in parallel.
+			t.Run("MissingContent", func(t *testing.T) {
+				err = exp.TaskSend(ctx, "me", task.ID, codersdk.TaskSendRequest{
+					Input: "",
+				})
+				require.Error(t, err, "wanted error due to missing content")
+
+				var sdkErr *codersdk.Error
+				require.ErrorAs(t, err, &sdkErr)
+				require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
+			})
 		})
 
 		t.Run("TaskNotFound", func(t *testing.T) {
@@ -627,108 +627,112 @@ func TestTasks(t *testing.T) {
 			require.ErrorAs(t, err, &sdkErr)
 			require.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
 		})
-
-		t.Run("NotATask", func(t *testing.T) {
-			t.Parallel()
-
-			client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
-			user := coderdtest.CreateFirstUser(t, client)
-			ctx := testutil.Context(t, testutil.WaitShort)
-
-			// Create a template without AI tasks.
-			version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
-			coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
-			template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
-
-			ws := coderdtest.CreateWorkspace(t, client, template.ID)
-			coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID)
-
-			exp := codersdk.NewExperimentalClient(client)
-			err := exp.TaskSend(ctx, "me", ws.ID, codersdk.TaskSendRequest{
-				Input: "hello",
-			})
-
-			var sdkErr *codersdk.Error
-			require.Error(t, err)
-			require.ErrorAs(t, err, &sdkErr)
-			require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
-		})
 	})
 
 	t.Run("Logs", func(t *testing.T) {
 		t.Parallel()
 
-		t.Skip("TODO(mafredri): Remove, fixed down-stack!")
-
-		t.Run("OK", func(t *testing.T) {
-			t.Parallel()
-
-			client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
-			owner := coderdtest.CreateFirstUser(t, client)
-			ctx := testutil.Context(t, testutil.WaitLong)
-
-			messageResponse := `
+		messageResponseData := agentapisdk.GetMessagesResponse{
+			Messages: []agentapisdk.Message{
 				{
-					"$schema": "http://localhost:3284/schemas/MessagesResponseBody.json",
-					"messages": [
-						{
-							"id": 0,
-							"content": "Welcome, user!",
-							"role": "agent",
-							"time": "2025-09-25T10:42:48.751774125Z"
-						},
-						{
-							"id": 1,
-							"content": "Hello, agent!",
-							"role": "user",
-							"time": "2025-09-25T10:46:42.880996296Z"
-						},
-						{
-							"id": 2,
-							"content": "What would you like to work on today?",
-							"role": "agent",
-							"time": "2025-09-25T10:46:50.747761102Z"
-						}
-					]
-				}
-			`
+					Id:      0,
+					Content: "Welcome, user!",
+					Role:    agentapisdk.RoleAgent,
+					Time:    time.Date(2025, 9, 25, 10, 42, 48, 0, time.UTC),
+				},
+				{
+					Id:      1,
+					Content: "Hello, agent!",
+					Role:    agentapisdk.RoleUser,
+					Time:    time.Date(2025, 9, 25, 10, 46, 42, 0, time.UTC),
+				},
+				{
+					Id:      2,
+					Content: "What would you like to work on today?",
+					Role:    agentapisdk.RoleAgent,
+					Time:    time.Date(2025, 9, 25, 10, 46, 50, 0, time.UTC),
+				},
+			},
+		}
+		messageResponseBytes, err := json.Marshal(messageResponseData)
+		require.NoError(t, err)
+		messageResponse := string(messageResponseBytes)
 
-			// Fake AgentAPI that returns a couple of messages.
-			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				if r.Method == http.MethodGet && r.URL.Path == "/messages" {
-					w.Header().Set("Content-Type", "application/json")
-					w.WriteHeader(http.StatusOK)
-					io.WriteString(w, messageResponse)
-					return
-				}
-				w.WriteHeader(http.StatusNotFound)
-			}))
-			t.Cleanup(srv.Close)
+		var shouldReturnError bool
 
-			// Template pointing sidebar app to our fake AgentAPI.
-			authToken := uuid.NewString()
-			template := createAITemplate(t, client, owner, withSidebarURL(srv.URL), withAgentToken(authToken))
+		// Fake AgentAPI that returns a couple of messages or an error.
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if shouldReturnError {
+				w.WriteHeader(http.StatusInternalServerError)
+				_, _ = io.WriteString(w, "boom")
+				return
+			}
+			if r.Method == http.MethodGet && r.URL.Path == "/messages" {
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(http.StatusOK)
+				io.WriteString(w, messageResponse)
+				return
+			}
+			w.WriteHeader(http.StatusNotFound)
+		}))
+		defer srv.Close()
 
-			// Create task workspace.
-			ws := coderdtest.CreateWorkspace(t, client, template.ID, func(req *codersdk.CreateWorkspaceRequest) {
-				req.RichParameterValues = []codersdk.WorkspaceBuildParameter{
-					{Name: codersdk.AITaskPromptParameterName, Value: "show logs"},
-				}
-			})
-			coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID)
+		// Create an AI-capable template whose sidebar app points to our fake AgentAPI.
+		var (
+			client, db     = coderdtest.NewWithDatabase(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
+			ctx            = testutil.Context(t, testutil.WaitLong)
+			owner          = coderdtest.CreateFirstUser(t, client)
+			agentAuthToken = uuid.NewString()
+			template       = createAITemplate(t, client, owner, withAgentToken(agentAuthToken), withSidebarURL(srv.URL))
+			exp            = codersdk.NewExperimentalClient(client)
+		)
 
-			// Start a fake agent.
-			agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(authToken))
-			_ = agenttest.New(t, client.URL, authToken, func(o *agent.Options) {
-				o.Client = agentClient
-			})
-			coderdtest.NewWorkspaceAgentWaiter(t, client, ws.ID).WithContext(ctx).WaitFor(coderdtest.AgentsReady)
+		task, err := exp.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
+			TemplateVersionID: template.ActiveVersionID,
+			Input:             "show logs",
+		})
+		require.NoError(t, err)
+		require.True(t, task.WorkspaceID.Valid)
 
-			// Omit sidebar app health as undefined is OK.
+		// Get the workspace and wait for it to be ready.
+		ws, err := client.Workspace(ctx, task.WorkspaceID.UUID)
+		require.NoError(t, err)
+		coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID)
 
+		// Fetch the task by ID via experimental API and verify fields.
+		task, err = exp.TaskByID(ctx, task.ID)
+		require.NoError(t, err)
+		require.NotZero(t, task.WorkspaceBuildNumber)
+		require.True(t, task.WorkspaceAgentID.Valid)
+		require.True(t, task.WorkspaceAppID.Valid)
+
+		// Insert an app status for the workspace
+		_, err = db.InsertWorkspaceAppStatus(dbauthz.AsSystemRestricted(ctx), database.InsertWorkspaceAppStatusParams{
+			ID:          uuid.New(),
+			WorkspaceID: task.WorkspaceID.UUID,
+			CreatedAt:   dbtime.Now(),
+			AgentID:     task.WorkspaceAgentID.UUID,
+			AppID:       task.WorkspaceAppID.UUID,
+			State:       database.WorkspaceAppStatusStateComplete,
+			Message:     "all done",
+		})
+		require.NoError(t, err)
+
+		// Start a fake agent so the workspace agent is connected before fetching logs.
+		agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(agentAuthToken))
+		_ = agenttest.New(t, client.URL, agentAuthToken, func(o *agent.Options) {
+			o.Client = agentClient
+		})
+		coderdtest.NewWorkspaceAgentWaiter(t, client, ws.ID).WaitFor(coderdtest.AgentsReady)
+
+		// Fetch the task by ID via experimental API and verify fields.
+		task, err = exp.TaskByID(ctx, task.ID)
+		require.NoError(t, err)
+
+		//nolint:tparallel // Not intended to run in parallel.
+		t.Run("OK", func(t *testing.T) {
 			// Fetch logs.
-			exp := codersdk.NewExperimentalClient(client)
-			resp, err := exp.TaskLogs(ctx, "me", ws.ID)
+			resp, err := exp.TaskLogs(ctx, "me", task.ID)
 			require.NoError(t, err)
 			require.Len(t, resp.Logs, 3)
 			assert.Equal(t, 0, resp.Logs[0].ID)
@@ -744,38 +748,11 @@ func TestTasks(t *testing.T) {
 			assert.Equal(t, "What would you like to work on today?", resp.Logs[2].Content)
 		})
 
+		//nolint:tparallel // Not intended to run in parallel.
 		t.Run("UpstreamError", func(t *testing.T) {
-			t.Parallel()
-
-			client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
-			owner := coderdtest.CreateFirstUser(t, client)
-			ctx := testutil.Context(t, testutil.WaitShort)
-
-			// Fake AgentAPI that returns 500 for messages.
-			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				w.WriteHeader(http.StatusInternalServerError)
-				_, _ = io.WriteString(w, "boom")
-			}))
-			t.Cleanup(srv.Close)
-
-			authToken := uuid.NewString()
-			template := createAITemplate(t, client, owner, withSidebarURL(srv.URL), withAgentToken(authToken))
-			ws := coderdtest.CreateWorkspace(t, client, template.ID, func(req *codersdk.CreateWorkspaceRequest) {
-				req.RichParameterValues = []codersdk.WorkspaceBuildParameter{
-					{Name: codersdk.AITaskPromptParameterName, Value: "show logs"},
-				}
-			})
-			coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID)
-
-			// Start fake agent.
-			agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(authToken))
-			_ = agenttest.New(t, client.URL, authToken, func(o *agent.Options) {
-				o.Client = agentClient
-			})
-			coderdtest.NewWorkspaceAgentWaiter(t, client, ws.ID).WithContext(ctx).WaitFor(coderdtest.AgentsReady)
-
-			exp := codersdk.NewExperimentalClient(client)
-			_, err := exp.TaskLogs(ctx, "me", ws.ID)
+			shouldReturnError = true
+			t.Cleanup(func() { shouldReturnError = false })
+			_, err := exp.TaskLogs(ctx, "me", task.ID)
 
 			var sdkErr *codersdk.Error
 			require.Error(t, err)