diff --git a/coderd/aitasks.go b/coderd/aitasks.go index 10c3efc96131a..1f90a6afda3c9 100644 --- a/coderd/aitasks.go +++ b/coderd/aitasks.go @@ -23,6 +23,7 @@ import ( "github.com/coder/coder/v2/coderd/rbac/policy" "github.com/coder/coder/v2/coderd/searchquery" "github.com/coder/coder/v2/coderd/taskname" + "github.com/coder/coder/v2/coderd/util/slice" "github.com/coder/coder/v2/codersdk" ) @@ -440,8 +441,32 @@ func (api *API) taskGet(rw http.ResponseWriter, r *http.Request) { return } if data.builds[0].HasAITask == nil || !*data.builds[0].HasAITask { - httpapi.ResourceNotFound(rw) - return + // TODO(DanielleMaywood): + // This is a temporary workaround. When a task has just been created, but + // not yet provisioned, the workspace build will not have `HasAITask` set. + // + // When we reach this code flow, it is _either_ because the workspace is + // not a task, or it is a task that has not yet been provisioned. This + // endpoint should rarely be called with a non-task workspace so we + // should be fine with this extra database call to check if it has the + // special "AI Task" parameter. + parameters, err := api.Database.GetWorkspaceBuildParameters(ctx, data.builds[0].ID) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Internal error fetching workspace build parameters.", + Detail: err.Error(), + }) + return + } + + _, hasAITask := slice.Find(parameters, func(t database.WorkspaceBuildParameter) bool { + return t.Name == codersdk.AITaskPromptParameterName + }) + + if !hasAITask { + httpapi.ResourceNotFound(rw) + return + } } appStatus := codersdk.WorkspaceAppStatus{} diff --git a/coderd/aitasks_test.go b/coderd/aitasks_test.go index 767f52eeab6b2..1f3cd8cbbdd08 100644 --- a/coderd/aitasks_test.go +++ b/coderd/aitasks_test.go @@ -253,7 +253,6 @@ func TestTasks(t *testing.T) { {Name: codersdk.AITaskPromptParameterName, Value: wantPrompt}, } }) - coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID) // Fetch the task by ID via experimental API and verify fields. exp := codersdk.NewExperimentalClient(client)