feat: store coder_ai_task resource state

Signed-off-by: Danny Kopping <dannykopping@gmail.com>
coder · dannykopping · Jun 24, 2025 · Jun 18, 2025 · Jun 18, 2025 · Jun 18, 2025
commit 438335039ab6d62488d90dd4ff06ca44b7507de0
diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go
@@ -28,6 +28,7 @@ import (
 	protobuf "google.golang.org/protobuf/proto"
 
 	"cdr.dev/slog"
+
 	"github.com/coder/coder/v2/coderd/util/slice"
 
 	"github.com/coder/coder/v2/codersdk/drpcsdk"
@@ -1654,6 +1655,17 @@ func (s *server) completeTemplateImportJob(ctx context.Context, job database.Pro
 		if err != nil {
 			return xerrors.Errorf("update template version external auth providers: %w", err)
 		}
+		err = db.UpdateTemplateVersionAITaskByJobID(ctx, database.UpdateTemplateVersionAITaskByJobIDParams{
+			JobID: jobID,
+			HasAITask: sql.NullBool{
+				Bool:  jobType.TemplateImport.HasAiTasks,
+				Valid: true,
+			},
+			UpdatedAt: now,
+		})
+		if err != nil {
+			return xerrors.Errorf("update template version external auth providers: %w", err)
+		}
 
 		// Process terraform values
 		plan := jobType.TemplateImport.Plan
@@ -1866,6 +1878,34 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
 			}
 		}
 
+	
8000
	var sidebarAppID uuid.NullUUID
+		if len(jobType.WorkspaceBuild.AiTasks) == 1 {
+			task := jobType.WorkspaceBuild.AiTasks[0]
+			if task.SidebarApp == nil {
+				return xerrors.Errorf("update ai task: sidebar app is nil")
+			}
+
+			id, err := uuid.Parse(task.SidebarApp.Id)
+			if err != nil {
+				return xerrors.Errorf("parse sidebar app id: %w", err)
+			}
+
+			sidebarAppID = uuid.NullUUID{UUID: id, Valid: true}
+		}
+
+		err = db.UpdateWorkspaceBuildAITaskByID(ctx, database.UpdateWorkspaceBuildAITaskByIDParams{
+			ID: workspaceBuild.ID,
+			HasAITask: sql.NullBool{
+				Bool:  len(jobType.WorkspaceBuild.AiTasks) > 0,
+				Valid: true,
+			},
+			SidebarAppID: sidebarAppID,
+			UpdatedAt:    now,
+		})
+		if err != nil {
+			return xerrors.Errorf("update workspace build ai tasks flag: %w", err)
+		}
+
 		// Insert timings inside the transaction now
 		// nolint:exhaustruct // The other fields are set further down.
 		params := database.InsertProvisionerJobTimingsParams{

diff --git a/provisioner/terraform/executor.go b/provisioner/terraform/executor.go
@@ -361,6 +361,8 @@ func (e *executor) plan(ctx, killCtx context.Context, env, vars []string, logr l
 		Plan:                  planJSON,
 		ResourceReplacements:  resReps,
 		ModuleFiles:           moduleFiles,
+		HasAiTasks:            state.HasAITasks,
+		AiTasks:               state.AITasks,
 	}
 
 	return msg, nil
@@ -577,6 +579,7 @@ func (e *executor) apply(
 		ExternalAuthProviders: state.ExternalAuthProviders,
 		State:                 stateContent,
 		Timings:               e.timings.aggregate(),
+		AiTasks:               state.AITasks,
 	}, nil
 }
 

diff --git a/provisioner/terraform/resources.go b/provisioner/terraform/resources.go
@@ -162,10 +162,29 @@ type State struct {
 	Parameters            []*proto.RichParameter
 	Presets               []*proto.Preset
 	ExternalAuthProviders []*proto.ExternalAuthProviderResource
+	AITasks               []*proto.AITask
+	HasAITasks            bool
 }
 
 var ErrInvalidTerraformAddr = xerrors.New("invalid terraform address")
 
+// hasAITaskResources is used to determine if a template has *any* `coder_ai_task` resources defined. During template
+// import, it's possible that none of these have `count=1` since count may be dependent on the value of a `coder_parameter`
+// or something else.
+// We need to know at template import if these resources exist to inform the frontend of their existence.
+func hasAITaskResources(graph *gographviz.Graph) bool {
+	for _, node := range graph.Nodes.Lookup {
+		// Check if this node is a coder_ai_task resource
+		if label, exists := node.Attrs["label"]; exists {
+			labelValue := strings.Trim(label, `"`)
+			if strings.Contains(labelValue, "coder_ai_task.") {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // ConvertState consumes Terraform state and a GraphViz representation
 // produced by `terraform graph` to produce resources consumable by Coder.
 // nolint:gocognit // This function makes more sense being large for now, until refactored.
@@ -189,6 +208,7 @@ func ConvertState(ctx context.Context, modules []*tfjson.StateModule, rawGraph s
 	// Extra array to preserve the order of rich parameters.
 	tfResourcesRichParameters := make([]*tfjson.StateResource, 0)
 	tfResourcesPresets := make([]*tfjson.StateResource, 0)
+	tfResourcesAITasks := make([]*tfjson.StateResource, 0)
 	var findTerraformResources func(mod *tfjson.StateModule)
 	findTerraformResources = func(mod *tfjson.StateModule) {
 		for _, module := range mod.ChildModules {
@@ -201,6 +221,9 @@ func ConvertState(ctx context.Context, modules []*tfjson.StateModule, rawGraph s
 			if resource.Type == "coder_workspace_preset" {
 				tfResourcesPresets = append(tfResourcesPresets, resource)
 			}
+			if resource.Type == "coder_ai_task" {
+				tfResourcesAITasks = append(tfResourcesAITasks, resource)
+			}
 
 			label := convertAddressToLabel(resource.Address)
 			if tfResourcesByLabel[label] == nil {
@@ -957,6 +980,27 @@ func ConvertState(ctx context.Context, modules []*tfjson.StateModule, rawGraph s
 		)
 	}
 
+	// This will only pick up resources which will actually be created.
+	aiTasks := make([]*proto.AITask, 0, len(tfResourcesAITasks))
+	for _, resource := range tfResourcesAITasks {
+		var task provider.AITask
+		err = mapstructure.Decode(resource.AttributeValues, &task)
+		if err != nil {
+			return nil, xerrors.Errorf("decode coder_ai_task attributes: %w", err)
+		}
+
+		if len(task.SidebarApp) < 1 {
+			return nil, xerrors.Errorf("coder_ai_task has no sidebar_app defined")
+		}
+
+		aiTasks = append(aiTasks, &proto.AITask{
+			Id: task.ID,
+			SidebarApp: &proto.AITaskSidebarApp{
+				Id: task.SidebarApp[0].ID,
+			},
+		})
+	}
+
 	// A map is used to ensure we don't have duplicates!
 	externalAuthProvidersMap := map[string]*proto.ExternalAuthProviderResource{}
 	for _, tfResources := range tfResourcesByLabel {
@@ -992,6 +1036,8 @@ func ConvertState(ctx context.Context, modules []*tfjson.StateModule, rawGraph s
 		Parameters:            parameters,
 		Presets:               presets,
 		ExternalAuthProviders: externalAuthProviders,
+		HasAITasks:            hasAITaskResources(graph),
+		AITasks:               aiTasks,
 	}, nil
 }
 

diff --git a/provisionerd/runner/runner.go b/provisionerd/runner/runner.go
@@ -7,11 +7,13 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"slices"
 	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
 
+	"github.com/coder/terraform-provider-coder/v2/provider"
 	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
 	"go.opentelemetry.io/otel/attribute"
@@ -584,8 +586,6 @@ func (r *Runner) runTemplateImport(ctx context.Context) (*proto.CompletedJob, *p
 		externalAuthProviderNames = append(externalAuthProviderNames, it.Id)
 	}
 
-	// fmt.Println("completed job: template import: graph:", startProvision.Graph)
-
 	return &proto.CompletedJob{
 		JobId: r.job.JobId,
 		Type: &proto.CompletedJob_TemplateImport_{
@@ -603,6 +603,7 @@ func (r *Runner) runTemplateImport(ctx context.Context) (*proto.CompletedJob, *p
 				ModuleFiles: startProvision.ModuleFiles,
 				// ModuleFileHash will be populated if the file is uploaded async
 				ModuleFilesHash: []byte{},
+				HasAiTasks:      startProvision.HasAITasks,
 			},
 		},
 	}, nil
@@ -666,6 +667,7 @@ type templateImportProvision struct {
 	Presets               []*sdkproto.Preset
 	Plan                  json.RawMessage
 	ModuleFiles           []byte
+	HasAITasks            bool
 }
 
 // Performs a dry-run provision when importing a template.
@@ -799,6 +801,15 @@ func (r *Runner) runTemplateImportProvisionWithRichParameters(
 				}
 			}
 
+			if c.HasAiTasks {
+				hasPromptParam := slices.ContainsFunc(c.Parameters, func(param *sdkproto.RichParameter) bool {
+					return param.Name == provider.TaskPromptParameterName
+				})
+				if !hasPromptParam {
+					return nil, xerrors.Errorf("coder_parameter named '%s' is required when 'coder_ai_task' resource is defined", provider.TaskPromptParameterName)
+				}
+			}
+
 			return &templateImportProvision{
 				Resources:             c.Resources,
 				Parameters:            c.Parameters,
@@ -807,6 +818,7 @@ func (r *Runner) runTemplateImportProvisionWithRichParameters(
 				Presets:               c.Presets,
 				Plan:                  c.Plan,
 				ModuleFiles:           moduleFilesData,
+				HasAITasks:            c.HasAiTasks,
 			}, nil
 		default:
 			return nil, xerrors.Errorf("invalid message type %q received from provisioner",
@@ -1047,6 +1059,9 @@ func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *p
 			},
 		}
 	}
+	if len(planComplete.AiTasks) > 1 {
+		return nil, r.failedWorkspaceBuildf("only one 'coder_ai_task' resource can be provisioned per template")
+	}
 
 	r.logger.Info(context.Background(), "plan request successful",
 		slog.F("resource_count", len(planComplete.Resources)),
@@ -1124,6 +1139,7 @@ func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *p
 				Modules: planComplete.Modules,
 				// Resource replacements are discovered at plan time, only.
 				ResourceReplacements: planComplete.ResourceReplacements,
+				AiTasks:              applyComplete.AiTasks,
 			},
 		},
 	}, nil