diff --git a/cli/exp_task_logs_test.go b/cli/exp_task_logs_test.go index 5dc3a99581..69905aa434 100644 --- a/cli/exp_task_logs_test.go +++ b/cli/exp_task_logs_test.go @@ -1,11 +1,8 @@ package cli_test import ( - "context" "encoding/json" - "fmt" "net/http" - "net/http/httptest" "strings" "testing" "time" @@ -14,7 +11,10 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + agentapisdk "github.com/coder/agentapi-sdk-go" + "github.com/coder/coder/v2/cli/clitest" + "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/httpapi" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" @@ -23,178 +23,165 @@ import ( func Test_TaskLogs(t *testing.T) { t.Parallel() - var ( - clock = time.Date(2025, 8, 26, 12, 34, 56, 0, time.UTC) - - taskID = uuid.MustParse("11111111-1111-1111-1111-111111111111") - taskName = "task-workspace" - - taskLogs = []codersdk.TaskLogEntry{ - { - ID: 0, - Content: "What is 1 + 1?", - Type: codersdk.TaskLogTypeInput, - Time: clock, - }, - { - ID: 1, - Content: "2", - Type: codersdk.TaskLogTypeOutput, - Time: clock.Add(1 * time.Second), - }, - } - ) - - tests := []struct { - args []string - expectTable string - expectLogs []codersdk.TaskLogEntry - expectError string - handler func(t *testing.T, ctx context.Context) http.HandlerFunc - }{ + testMessages := []agentapisdk.Message{ { - args: []string{taskName, "--output", "json"}, - expectLogs: taskLogs, - handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case fmt.Sprintf("/api/v2/users/me/workspace/%s", taskName): - httpapi.Write(ctx, w, http.StatusOK, codersdk.Workspace{ - ID: taskID, - }) - case fmt.Sprintf("/api/experimental/tasks/me/%s/logs", taskID.String()): - httpapi.Write(ctx, w, http.StatusOK, codersdk.TaskLogsResponse{ - Logs: taskLogs, - }) - default: - t.Errorf("unexpected path: %s", r.URL.Path) - } - } - }, + Id: 0, + Role: agentapisdk.RoleUser, + Content: "What is 1 + 1?", + Time: time.Now().Add(-2 * time.Minute), }, { - args: []string{taskID.String(), "--output", "json"}, - expectLogs: taskLogs, - handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case fmt.Sprintf("/api/experimental/tasks/me/%s/logs", taskID.String()): - httpapi.Write(ctx, w, http.StatusOK, codersdk.TaskLogsResponse{ - Logs: taskLogs, - }) - default: - t.Errorf("unexpected path: %s", r.URL.Path) - } - } - }, - }, - { - args: []string{taskID.String()}, - expectTable: ` -TYPE CONTENT -input What is 1 + 1? -output 2`, - handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case fmt.Sprintf("/api/experimental/tasks/me/%s/logs", taskID.String()): - httpapi.Write(ctx, w, http.StatusOK, codersdk.TaskLogsResponse{ - Logs: taskLogs, - }) - default: - t.Errorf("unexpected path: %s", r.URL.Path) - } - } - }, - }, - { - args: []string{"doesnotexist"}, - expectError: httpapi.ResourceNotFoundResponse.Message, - handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case "/api/v2/users/me/workspace/doesnotexist": - httpapi.ResourceNotFound(w) - default: - t.Errorf("unexpected path: %s", r.URL.Path) - } - } - }, - }, - { - args: []string{uuid.Nil.String()}, // uuid does not exist - expectError: httpapi.ResourceNotFoundResponse.Message, - handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case fmt.Sprintf("/api/experimental/tasks/me/%s/logs", uuid.Nil.String()): - httpapi.ResourceNotFound(w) - default: - t.Errorf("unexpected path: %s", r.URL.Path) - } - } - }, - }, - { - args: []string{"err-fetching-logs"}, - expectError: assert.AnError.Error(), - handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case "/api/v2/users/me/workspace/err-fetching-logs": - httpapi.Write(ctx, w, http.StatusOK, codersdk.Workspace{ - ID: taskID, - }) - case fmt.Sprintf("/api/experimental/tasks/me/%s/logs", taskID.String()): - httpapi.InternalServerError(w, assert.AnError) - default: - t.Errorf("unexpected path: %s", r.URL.Path) - } - } - }, + Id: 1, + Role: agentapisdk.RoleAgent, + Content: "2", + Time: time.Now().Add(-1 * time.Minute), }, } - for _, tt := range tests { - t.Run(strings.Join(tt.args, ","), func(t *testing.T) { - t.Parallel() + t.Run("ByWorkspaceName_JSON", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) - var ( - ctx = testutil.Context(t, testutil.WaitShort) - srv = httptest.NewServer(tt.handler(t, ctx)) - client = codersdk.New(testutil.MustURL(t, srv.URL)) - args = []string{"exp", "task", "logs"} - stdout strings.Builder - err error - ) + client, workspace := setupCLITaskTest(ctx, t, fakeAgentAPITaskLogsOK(testMessages)) + userClient := client // user already has access to their own workspace - t.Cleanup(srv.Close) + var stdout strings.Builder + inv, root := clitest.New(t, "exp", "task", "logs", workspace.Name, "--output", "json") + inv.Stdout = &stdout + clitest.SetupConfig(t, userClient, root) - inv, root := clitest.New(t, append(args, tt.args...)...) - inv.Stdout = &stdout - inv.Stderr = &stdout - clitest.SetupConfig(t, client, root) + err := inv.WithContext(ctx).Run() + require.NoError(t, err) - err = inv.WithContext(ctx).Run() - if tt.expectError == "" { - assert.NoError(t, err) - } else { - assert.ErrorContains(t, err, tt.expectError) - } + var logs []codersdk.TaskLogEntry + err = json.NewDecoder(strings.NewReader(stdout.String())).Decode(&logs) + require.NoError(t, err) - if tt.expectTable != "" { - if diff := tableDiff(tt.expectTable, stdout.String()); diff != "" { - t.Errorf("unexpected output diff (-want +got):\n%s", diff) - } - } + require.Len(t, logs, 2) + require.Equal(t, "What is 1 + 1?", logs[0].Content) + require.Equal(t, codersdk.TaskLogTypeInput, logs[0].Type) + require.Equal(t, "2", logs[1].Content) + require.Equal(t, codersdk.TaskLogTypeOutput, logs[1].Type) + }) - if tt.expectLogs != nil { - var logs []codersdk.TaskLogEntry - err = json.NewDecoder(strings.NewReader(stdout.String())).Decode(&logs) - require.NoError(t, err) + t.Run("ByWorkspaceID_JSON", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) - assert.Equal(t, tt.expectLogs, logs) - } - }) + client, workspace := setupCLITaskTest(ctx, t, fakeAgentAPITaskLogsOK(testMessages)) + userClient := client + + var stdout strings.Builder + inv, root := clitest.New(t, "exp", "task", "logs", workspace.ID.String(), "--output", "json") + inv.Stdout = &stdout + clitest.SetupConfig(t, userClient, root) + + err := inv.WithContext(ctx).Run() + require.NoError(t, err) + + var logs []codersdk.TaskLogEntry + err = json.NewDecoder(strings.NewReader(stdout.String())).Decode(&logs) + require.NoError(t, err) + + require.Len(t, logs, 2) + require.Equal(t, "What is 1 + 1?", logs[0].Content) + require.Equal(t, codersdk.TaskLogTypeInput, logs[0].Type) + require.Equal(t, "2", logs[1].Content) + require.Equal(t, codersdk.TaskLogTypeOutput, logs[1].Type) + }) + + t.Run("ByWorkspaceID_Table", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + client, workspace := setupCLITaskTest(ctx, t, fakeAgentAPITaskLogsOK(testMessages)) + userClient := client + + var stdout strings.Builder + inv, root := clitest.New(t, "exp", "task", "logs", workspace.ID.String()) + inv.Stdout = &stdout + clitest.SetupConfig(t, userClient, root) + + err := inv.WithContext(ctx).Run() + require.NoError(t, err) + + output := stdout.String() + require.Contains(t, output, "What is 1 + 1?") + require.Contains(t, output, "2") + require.Contains(t, output, "input") + require.Contains(t, output, "output") + }) + + t.Run("WorkspaceNotFound_ByName", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true}) + owner := coderdtest.CreateFirstUser(t, client) + userClient, _ := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID) + + var stdout strings.Builder + inv, root := clitest.New(t, "exp", "task", "logs", "doesnotexist") + inv.Stdout = &stdout + clitest.SetupConfig(t, userClient, root) + + err := inv.WithContext(ctx).Run() + require.Error(t, err) + require.ErrorContains(t, err, httpapi.ResourceNotFoundResponse.Message) + }) + + t.Run("WorkspaceNotFound_ByID", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true}) + owner := coderdtest.CreateFirstUser(t, client) + userClient, _ := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID) + + var stdout strings.Builder + inv, root := clitest.New(t, "exp", "task", "logs", uuid.Nil.String()) + inv.Stdout = &stdout + clitest.SetupConfig(t, userClient, root) + + err := inv.WithContext(ctx).Run() + require.Error(t, err) + require.ErrorContains(t, err, httpapi.ResourceNotFoundResponse.Message) + }) + + t.Run("ErrorFetchingLogs", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + client, workspace := setupCLITaskTest(ctx, t, fakeAgentAPITaskLogsErr(assert.AnError)) + userClient := client + + inv, root := clitest.New(t, "exp", "task", "logs", workspace.ID.String()) + clitest.SetupConfig(t, userClient, root) + + err := inv.WithContext(ctx).Run() + require.ErrorContains(t, err, assert.AnError.Error()) + }) +} + +func fakeAgentAPITaskLogsOK(messages []agentapisdk.Message) map[string]http.HandlerFunc { + return map[string]http.HandlerFunc{ + "/messages": func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "messages": messages, + }) + }, + } +} + +func fakeAgentAPITaskLogsErr(err error) map[string]http.HandlerFunc { + return map[string]http.HandlerFunc{ + "/messages": func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]interface{}{ + "error": err.Error(), + }) + }, } } diff --git a/cli/exp_task_send_test.go b/cli/exp_task_send_test.go index 1d7b863f5f..cb8ee74d06 100644 --- a/cli/exp_task_send_test.go +++ b/cli/exp_task_send_test.go @@ -1,173 +1,171 @@ package cli_test import ( - "context" - "fmt" + "encoding/json" "net/http" - "net/http/httptest" "strings" "testing" + "time" "github.com/google/uuid" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + agentapisdk "github.com/coder/agentapi-sdk-go" "github.com/coder/coder/v2/cli/clitest" + "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/httpapi" - "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" ) func Test_TaskSend(t *testing.T) { t.Parallel() - var ( - taskName = "task-workspace" - taskID = uuid.MustParse("11111111-1111-1111-1111-111111111111") - ) + t.Run("ByWorkspaceName_WithArgument", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) - tests := []struct { - args []string - stdin string - expectError string - handler func(t *testing.T, ctx context.Context) http.HandlerFunc - }{ - { - args: []string{taskName, "carry on with the task"}, - handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case fmt.Sprintf("/api/v2/users/me/workspace/%s", taskName): - httpapi.Write(ctx, w, http.StatusOK, codersdk.Workspace{ - ID: taskID, - }) - case fmt.Sprintf("/api/experimental/tasks/me/%s/send", taskID.String()): - var req codersdk.TaskSendRequest - if !httpapi.Read(ctx, w, r, &req) { - return - } + client, workspace := setupCLITaskTest(ctx, t, fakeAgentAPITaskSendOK(t, "carry on with the task", "you got it")) + userClient := client - assert.Equal(t, "carry on with the task", req.Input) + var stdout strings.Builder + inv, root := clitest.New(t, "exp", "task", "send", workspace.Name, "carry on with the task") + inv.Stdout = &stdout + clitest.SetupConfig(t, userClient, root) - httpapi.Write(ctx, w, http.StatusNoContent, nil) - default: - t.Errorf("unexpected path: %s", r.URL.Path) - } - } - }, + err := inv.WithContext(ctx).Run() + require.NoError(t, err) + }) + + t.Run("ByWorkspaceID_WithArgument", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + client, workspace := setupCLITaskTest(ctx, t, fakeAgentAPITaskSendOK(t, "carry on with the task", "you got it")) + userClient := client + + var stdout strings.Builder + inv, root := clitest.New(t, "exp", "task", "send", workspace.ID.String(), "carry on with the task") + inv.Stdout = &stdout + clitest.SetupConfig(t, userClient, root) + + err := inv.WithContext(ctx).Run() + require.NoError(t, err) + }) + + t.Run("ByWorkspaceName_WithStdin", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + client, workspace := setupCLITaskTest(ctx, t, fakeAgentAPITaskSendOK(t, "carry on with the task", "you got it")) + userClient := client + + var stdout strings.Builder + inv, root := clitest.New(t, "exp", "task", "send", workspace.Name, "--stdin") + inv.Stdout = &stdout + inv.Stdin = strings.NewReader("carry on with the task") + clitest.SetupConfig(t, userClient, root) + + err := inv.WithContext(ctx).Run() + require.NoError(t, err) + }) + + t.Run("WorkspaceNotFound_ByName", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true}) + owner := coderdtest.CreateFirstUser(t, client) + userClient, _ := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID) + + var stdout strings.Builder + inv, root := clitest.New(t, "exp", "task", "send", "doesnotexist", "some task input") + inv.Stdout = &stdout + clitest.SetupConfig(t, userClient, root) + + err := inv.WithContext(ctx).Run() + require.Error(t, err) + require.ErrorContains(t, err, httpapi.ResourceNotFoundResponse.Message) + }) + + t.Run("WorkspaceNotFound_ByID", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true}) + owner := coderdtest.CreateFirstUser(t, client) + userClient, _ := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID) + + var stdout strings.Builder + inv, root := clitest.New(t, "exp", "task", "send", uuid.Nil.String(), "some task input") + inv.Stdout = &stdout + clitest.SetupConfig(t, userClient, root) + + err := inv.WithContext(ctx).Run() + require.Error(t, err) + require.ErrorContains(t, err, httpapi.ResourceNotFoundResponse.Message) + }) + + t.Run("SendError", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + userClient, workspace := setupCLITaskTest(ctx, t, fakeAgentAPITaskSendErr(t, assert.AnError)) + + var stdout strings.Builder + inv, root := clitest.New(t, "exp", "task", "send", workspace.Name, "some task input") + inv.Stdout = &stdout + clitest.SetupConfig(t, userClient, root) + + err := inv.WithContext(ctx).Run() + require.ErrorContains(t, err, assert.AnError.Error()) + }) +} + +func fakeAgentAPITaskSendOK(t *testing.T, expectMessage, returnMessage string) map[string]http.HandlerFunc { + return map[string]http.HandlerFunc{ + "/status": func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]string{ + "status": "stable", + }) }, - { - args: []string{taskID.String(), "carry on with the task"}, - handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case fmt.Sprintf("/api/experimental/tasks/me/%s/send", taskID.String()): - var req codersdk.TaskSendRequest - if !httpapi.Read(ctx, w, r, &req) { - return - } - - assert.Equal(t, "carry on with the task", req.Input) - - httpapi.Write(ctx, w, http.StatusNoContent, nil) - default: - t.Errorf("unexpected path: %s", r.URL.Path) - } - } - }, - }, - { - args: []string{taskName, "--stdin"}, - stdin: "carry on with the task", - handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case fmt.Sprintf("/api/v2/users/me/workspace/%s", taskName): - httpapi.Write(ctx, w, http.StatusOK, codersdk.Workspace{ - ID: taskID, - }) - case fmt.Sprintf("/api/experimental/tasks/me/%s/send", taskID.String()): - var req codersdk.TaskSendRequest - if !httpapi.Read(ctx, w, r, &req) { - return - } - - assert.Equal(t, "carry on with the task", req.Input) - - httpapi.Write(ctx, w, http.StatusNoContent, nil) - default: - t.Errorf("unexpected path: %s", r.URL.Path) - } - } - }, - }, - { - args: []string{"doesnotexist", "some task input"}, - expectError: httpapi.ResourceNotFoundResponse.Message, - handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case "/api/v2/users/me/workspace/doesnotexist": - httpapi.ResourceNotFound(w) - default: - t.Errorf("unexpected path: %s", r.URL.Path) - } - } - }, - }, - { - args: []string{uuid.Nil.String(), "some task input"}, - expectError: httpapi.ResourceNotFoundResponse.Message, - handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case fmt.Sprintf("/api/experimental/tasks/me/%s/send", uuid.Nil.String()): - httpapi.ResourceNotFound(w) - default: - t.Errorf("unexpected path: %s", r.URL.Path) - } - } - }, - }, - { - args: []string{uuid.Nil.String(), "some task input"}, - expectError: assert.AnError.Error(), - handler: func(t *testing.T, ctx context.Context) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case fmt.Sprintf("/api/experimental/tasks/me/%s/send", uuid.Nil.String()): - httpapi.InternalServerError(w, assert.AnError) - default: - t.Errorf("unexpected path: %s", r.URL.Path) - } - } - }, - }, - } - - for _, tt := range tests { - t.Run(strings.Join(tt.args, ","), func(t *testing.T) { - t.Parallel() - - var ( - ctx = testutil.Context(t, testutil.WaitShort) - srv = httptest.NewServer(tt.handler(t, ctx)) - client = codersdk.New(testutil.MustURL(t, srv.URL)) - args = []string{"exp", "task", "send"} - err error - ) - - t.Cleanup(srv.Close) - - inv, root := clitest.New(t, append(args, tt.args...)...) - inv.Stdin = strings.NewReader(tt.stdin) - clitest.SetupConfig(t, client, root) - - err = inv.WithContext(ctx).Run() - if tt.expectError == "" { - assert.NoError(t, err) - } else { - assert.ErrorContains(t, err, tt.expectError) + "/message": func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + var msg agentapisdk.PostMessageParams + if err := json.NewDecoder(r.Body).Decode(&msg); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return } - }) + assert.Equal(t, expectMessage, msg.Content) + message := agentapisdk.Message{ + Id: 999, + Role: agentapisdk.RoleAgent, + Content: returnMessage, + Time: time.Now(), + } + _ = json.NewEncoder(w).Encode(message) + }, + } +} + +func fakeAgentAPITaskSendErr(t *testing.T, returnErr error) map[string]http.HandlerFunc { + return map[string]http.HandlerFunc{ + "/status": func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]string{ + "status": "stable", + }) + }, + "/message": func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(returnErr.Error())) + }, } } diff --git a/cli/exp_task_test.go b/cli/exp_task_test.go new file mode 100644 index 0000000000..4cf99e1dd4 --- /dev/null +++ b/cli/exp_task_test.go @@ -0,0 +1,202 @@ +package cli_test + +import ( + "context" + "net/http" + "net/http/httptest" + "sync" + "testing" + + "github.com/google/uuid" + + "github.com/coder/coder/v2/agent" + "github.com/coder/coder/v2/agent/agenttest" + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agentsdk" + "github.com/coder/coder/v2/provisioner/echo" + "github.com/coder/coder/v2/provisionersdk/proto" +) + +// setupCLITaskTest creates a test workspace with an AI task template and agent, +// with a fake agent API configured with the provided set of handlers. +// Returns the user client and workspace. +func setupCLITaskTest(ctx context.Context, t *testing.T, agentAPIHandlers map[string]http.HandlerFunc) (*codersdk.Client, codersdk.Workspace) { + t.Helper() + + client := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true}) + owner := coderdtest.CreateFirstUser(t, client) + userClient, _ := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID) + + fakeAPI := startFakeAgentAPI(t, agentAPIHandlers) + + authToken := uuid.NewString() + template := createAITaskTemplate(t, client, owner.OrganizationID, withSidebarURL(fakeAPI.URL()), withAgentToken(authToken)) + + wantPrompt := "test prompt" + workspace := coderdtest.CreateWorkspace(t, userClient, template.ID, func(req *codersdk.CreateWorkspaceRequest) { + req.RichParameterValues = []codersdk.WorkspaceBuildParameter{ + {Name: codersdk.AITaskPromptParameterName, Value: wantPrompt}, + } + }) + coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID) + + 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, workspace.ID). + WaitFor(coderdtest.AgentsReady) + + return userClient, workspace +} + +// 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() + + opt := aiTemplateOpts{ + authToken: uuid.NewString(), + } + for _, o := range opts { + o(&opt) + } + + taskAppID := uuid.New() + version := coderdtest.CreateTemplateVersion(t, client, orgID, &echo.Responses{ + Parse: echo.ParseComplete, + ProvisionPlan: []*proto.Response{ + { + Type: &proto.Response_Plan{ + Plan: &proto.PlanComplete{ + Parameters: []*proto.RichParameter{{Name: codersdk.AITaskPromptParameterName, Type: "string"}}, + HasAiTasks: true, + }, + }, + }, + }, + ProvisionApply: []*proto.Response{ + { + Type: &proto.Response_Apply{ + Apply: &proto.ApplyComplete{ + Resources: []*proto.Resource{ + { + Name: "example", + Type: "aws_instance", + Agents: []*proto.Agent{ + { + Id: uuid.NewString(), + Name: "example", + Auth: &proto.Agent_Token{ + Token: opt.authToken, + }, + Apps: []*proto.App{ + { + Id: taskAppID.String(), + Slug: "task-sidebar", + DisplayName: "Task Sidebar", + Url: opt.appURL, + }, + }, + }, + }, + }, + }, + AiTasks: []*proto.AITask{ + { + SidebarApp: &proto.AITaskSidebarApp{ + Id: taskAppID.String(), + }, + }, + }, + }, + }, + }, + }, + }) + coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) + template := coderdtest.CreateTemplate(t, client, orgID, version.ID) + + return template +} + +// fakeAgentAPI implements a fake AgentAPI HTTP server for testing. +type fakeAgentAPI struct { + t *testing.T + server *httptest.Server + handlers map[string]http.HandlerFunc + called map[string]bool + mu sync.Mutex +} + +// startFakeAgentAPI starts an HTTP server that implements the AgentAPI endpoints. +// handlers is a map of path -> handler function. +func startFakeAgentAPI(t *testing.T, handlers map[string]http.HandlerFunc) *fakeAgentAPI { + t.Helper() + + fake := &fakeAgentAPI{ + t: t, + handlers: handlers, + called: make(map[string]bool), + } + + mux := http.NewServeMux() + + // Register all provided handlers with call tracking + for path, handler := range handlers { + mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) { + fake.mu.Lock() + fake.called[path] = true + fake.mu.Unlock() + handler(w, r) + }) + } + + knownEndpoints := []string{"/status", "/messages", "/message"} + for _, endpoint := range knownEndpoints { + if handlers[endpoint] == nil { + endpoint := endpoint // capture loop variable + mux.HandleFunc(endpoint, func(w http.ResponseWriter, r *http.Request) { + t.Fatalf("unexpected call to %s %s - no handler defined", r.Method, endpoint) + }) + } + } + // Default handler for unknown endpoints should cause the test to fail. + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + t.Fatalf("unexpected call to %s %s - no handler defined", r.Method, r.URL.Path) + }) + + fake.server = httptest.NewServer(mux) + + // Register cleanup to check that all defined handlers were called + t.Cleanup(func() { + fake.server.Close() + fake.mu.Lock() + for path := range handlers { + if !fake.called[path] { + t.Errorf("handler for %s was defined but never called", path) + } + } + }) + return fake +} + +func (f *fakeAgentAPI) URL() string { + return f.server.URL +} + +type aiTemplateOpts struct { + appURL string + authToken string +} + +type aiTemplateOpt func(*aiTemplateOpts) + +func withSidebarURL(url string) aiTemplateOpt { + return func(o *aiTemplateOpts) { o.appURL = url } +} + +func withAgentToken(token string) aiTemplateOpt { + return func(o *aiTemplateOpts) { o.authToken = token } +}