From 12bce129524b1ab7a87a41f97a44d67e7e83f0de Mon Sep 17 00:00:00 2001
From: Danielle Maywood <danielle@themaywoods.com>
Date: Tue, 2 Sep 2025 13:25:32 +0100
Subject: [PATCH] fix(coderd): ensure a newly created task can be fetched
 (#19670)

Due to how we currently label a workspace as a task, there is a delay
between when a task workspace is created and when it is labelled as a
task.

This PR introduces fallback check for when a workspace does _not_ have
`HasAITask` set. This fallback check tests to see if the special "AI
Prompt" parameter is present in the workspace's build parameters.
---
 coderd/aitasks.go      | 29 +++++++++++++++++++++++++++--
 coderd/aitasks_test.go |  1 -
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/coderd/aitasks.go b/coderd/aitasks.go
index 10c3efc961..1f90a6afda 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 767f52eeab..1f3cd8cbbd 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)