chore: distinct operations for provisioner's 'parse', 'init', 'plan', 'apply', 'graph' (#21064)

Provisioner steps broken into smaller granular actions.
Changes:
- `ExtractArchive` moved to `init` request (was in `configure`)
- Writing `tfstate` moved to `plan` (was in `configure`)
- Moved most plan/apply outputs to `GraphComplete`

provisionerd/provisionerd_test.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/hashicorp/yamux"
+	"github.com/spf13/afero"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
@@ -131,6 +132,16 @@ func TestProvisionerd(t *testing.T) {
 					}
 					return c
 				},
+				init: func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+					closerMutex.Lock()
+					defer closerMutex.Unlock()
+					err := closer.Close()
+					c := &sdkproto.InitComplete{}
+					if err != nil {
+						c.Error = err.Error()
+					}
+					return c
+				},
 			}),
 		})
 		closerMutex.Unlock()
@@ -138,47 +149,6 @@ func TestProvisionerd(t *testing.T) {
 		require.NoError(t, closer.Close())
 	})
 
-	t.Run("MaliciousTar", func(t *testing.T) {
-		// Ensures tars with "../../../etc/passwd" as the path
-		// are not allowed to run, and will fail the job.
-		t.Parallel()
-		done := make(chan struct{})
-		t.Cleanup(func() {
-			close(done)
-		})
-		var (
-			completeChan = make(chan struct{})
-			completeOnce sync.Once
-			acq          = newAcquireOne(t, &proto.AcquiredJob{
-				JobId:       "test",
-				Provisioner: "someprovisioner",
-				TemplateSourceArchive: testutil.CreateTar(t, map[string]string{
-					"../../../etc/passwd": "content",
-				}),
-				Type: &proto.AcquiredJob_TemplateImport_{
-					TemplateImport: &proto.AcquiredJob_TemplateImport{
-						Metadata: &sdkproto.Metadata{},
-					},
-				},
-			})
-		)
-
-		closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
-			return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
-				acquireJobWithCancel: acq.acquireWithCancel,
-				updateJob:            noopUpdateJob,
-				failJob: func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error) {
-					completeOnce.Do(func() { close(completeChan) })
-					return &proto.Empty{}, nil
-				},
-			}), nil
-		}, provisionerd.LocalProvisioners{
-			"someprovisioner": createProvisionerClient(t, done, provisionerTestServer{}),
-		})
-		require.Condition(t, closedWithin(completeChan, testutil.WaitMedium))
-		require.NoError(t, closer.Close())
-	})
-
 	// LargePayloads sends a 3mb tar file to the provisioner. The provisioner also
 	// returns large payload messages back. The limit should be 4mb, so all
 	// these messages should work.
@@ -227,14 +197,16 @@ func TestProvisionerd(t *testing.T) {
 						Readme: make([]byte, largeSize),
 					}
 				},
+				init: func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+					return &sdkproto.InitComplete{}
+				},
 				plan: func(
 					_ *provisionersdk.Session,
 					_ *sdkproto.PlanRequest,
 					_ <-chan struct{},
 				) *sdkproto.PlanComplete {
 					return &sdkproto.PlanComplete{
-						Resources: []*sdkproto.Resource{},
-						Plan:      make([]byte, largeSize),
+						Plan: make([]byte, largeSize),
 					}
 				},
 				apply: func(
@@ -246,6 +218,11 @@ func TestProvisionerd(t *testing.T) {
 						State: make([]byte, largeSize),
 					}
 				},
+				graph: func(s *provisionersdk.Session, r *sdkproto.GraphRequest, canceledOrComplete <-chan struct{}) *sdkproto.GraphComplete {
+					return &sdkproto.GraphComplete{
+						Resources: []*sdkproto.Resource{},
+					}
+				},
 			}),
 		})
 		require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
@@ -299,6 +276,9 @@ func TestProvisionerd(t *testing.T) {
 					<-cancelOrComplete
 					return &sdkproto.ParseComplete{}
 				},
+				init: func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+					return &sdkproto.InitComplete{}
+				},
 			}),
 		})
 		require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
@@ -349,6 +329,7 @@ func TestProvisionerd(t *testing.T) {
 			}), nil
 		}, provisionerd.LocalProvisioners{
 			"someprovisioner": createProvisionerClient(t, done, provisionerTestServer{
+				init: extractInit(t),
 				parse: func(
 					s *provisionersdk.Session,
 					_ *sdkproto.ParseRequest,
@@ -366,9 +347,7 @@ func TestProvisionerd(t *testing.T) {
 					cancelOrComplete <-chan struct{},
 				) *sdkproto.PlanComplete {
 					s.ProvisionLog(sdkproto.LogLevel_INFO, "hello")
-					return &sdkproto.PlanComplete{
-						Resources: []*sdkproto.Resource{},
-					}
+					return &sdkproto.PlanComplete{}
 				},
 				apply: func(
 					_ *provisionersdk.Session,
@@ -378,6 +357,11 @@ func TestProvisionerd(t *testing.T) {
 					t.Error("dry run should not apply")
 					return &sdkproto.ApplyComplete{}
 				},
+				graph: func(s *provisionersdk.Session, r *sdkproto.GraphRequest, canceledOrComplete <-chan struct{}) *sdkproto.GraphComplete {
+					return &sdkproto.GraphComplete{
+						Resources: []*sdkproto.Resource{},
+					}
+				},
 			}),
 		})
 
@@ -433,14 +417,15 @@ func TestProvisionerd(t *testing.T) {
 			}), nil
 		}, provisionerd.LocalProvisioners{
 			"someprovisioner": createProvisionerClient(t, done, provisionerTestServer{
+				init: func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+					return &sdkproto.InitComplete{}
+				},
 				plan: func(
 					_ *provisionersdk.Session,
 					_ *sdkproto.PlanRequest,
 					_ <-chan struct{},
 				) *sdkproto.PlanComplete {
-					return &sdkproto.PlanComplete{
-						Resources: []*sdkproto.Resource{},
-					}
+					return &sdkproto.PlanComplete{}
 				},
 				apply: func(
 					_ *provisionersdk.Session,
@@ -450,6 +435,11 @@ func TestProvisionerd(t *testing.T) {
 					t.Error("dry run should not apply")
 					return &sdkproto.ApplyComplete{}
 				},
+				graph: func(s *provisionersdk.Session, r *sdkproto.GraphRequest, canceledOrComplete <-chan struct{}) *sdkproto.GraphComplete {
+					return &sdkproto.GraphComplete{
+						Resources: []*sdkproto.Resource{},
+					}
+				},
 			}),
 		})
 
@@ -498,6 +488,9 @@ func TestProvisionerd(t *testing.T) {
 			}), nil
 		}, provisionerd.LocalProvisioners{
 			"someprovisioner": createProvisionerClient(t, done, provisionerTestServer{
+				init: func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+					return &sdkproto.InitComplete{}
+				},
 				plan: func(
 					s *provisionersdk.Session,
 					_ *sdkproto.PlanRequest,
@@ -513,6 +506,9 @@ func TestProvisionerd(t *testing.T) {
 				) *sdkproto.ApplyComplete {
 					return &sdkproto.ApplyComplete{}
 				},
+				graph: func(s *provisionersdk.Session, r *sdkproto.GraphRequest, canceledOrComplete <-chan struct{}) *sdkproto.GraphComplete {
+					return &sdkproto.GraphComplete{}
+				},
 			}),
 		})
 		require.Condition(t, closedWithin(acq.complete, testutil.WaitShort))
@@ -570,6 +566,9 @@ func TestProvisionerd(t *testing.T) {
 			}), nil
 		}, provisionerd.LocalProvisioners{
 			"someprovisioner": createProvisionerClient(t, done, provisionerTestServer{
+				init: func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+					return &sdkproto.InitComplete{}
+				},
 				plan: func(
 					s *provisionersdk.Session,
 					_ *sdkproto.PlanRequest,
@@ -577,14 +576,7 @@ func TestProvisionerd(t *testing.T) {
 				) *sdkproto.PlanComplete {
 					s.ProvisionLog(sdkproto.LogLevel_DEBUG, "wow")
 					return &sdkproto.PlanComplete{
-						Resources: []*sdkproto.Resource{
-							{
-								DailyCost: 10,
-							},
-							{
-								DailyCost: 15,
-							},
-						},
+						DailyCost: 25,
 					}
 				},
 				apply: func(
@@ -593,7 +585,10 @@ func TestProvisionerd(t *testing.T) {
 					_ <-chan struct{},
 				) *sdkproto.ApplyComplete {
 					t.Error("should not apply when resources exceed quota")
-					return &sdkproto.ApplyComplete{
+					return &sdkproto.ApplyComplete{}
+				},
+				graph: func(s *provisionersdk.Session, r *sdkproto.GraphRequest, canceledOrComplete <-chan struct{}) *sdkproto.GraphComplete {
+					return &sdkproto.GraphComplete{
 						Resources: []*sdkproto.Resource{
 							{
 								DailyCost: 10,
@@ -646,6 +641,12 @@ func TestProvisionerd(t *testing.T) {
 			}), nil
 		}, provisionerd.LocalProvisioners{
 			"someprovisioner": createProvisionerClient(t, done, provisionerTestServer{
+				init: func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+					return &sdkproto.InitComplete{}
+				},
+				graph: func(s *provisionersdk.Session, r *sdkproto.GraphRequest, canceledOrComplete <-chan struct{}) *sdkproto.GraphComplete {
+					return &sdkproto.GraphComplete{}
+				},
 				plan: func(
 					s *provisionersdk.Session,
 					_ *sdkproto.PlanRequest,
@@ -756,6 +757,9 @@ func TestProvisionerd(t *testing.T) {
 			}), nil
 		}, provisionerd.LocalProvisioners{
 			"someprovisioner": createProvisionerClient(t, done, provisionerTestServer{
+				init: func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+					return &sdkproto.InitComplete{}
+				},
 				plan: func(
 					s *provisionersdk.Session,
 					_ *sdkproto.PlanRequest,
@@ -844,6 +848,9 @@ func TestProvisionerd(t *testing.T) {
 			}), nil
 		}, provisionerd.LocalProvisioners{
 			"someprovisioner": createProvisionerClient(t, done, provisionerTestServer{
+				init: func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+					return &sdkproto.InitComplete{}
+				},
 				plan: func(
 					s *provisionersdk.Session,
 					_ *sdkproto.PlanRequest,
@@ -938,6 +945,9 @@ func TestProvisionerd(t *testing.T) {
 			return client, nil
 		}, provisionerd.LocalProvisioners{
 			"someprovisioner": createProvisionerClient(t, done, provisionerTestServer{
+				init: func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+					return &sdkproto.InitComplete{}
+				},
 				plan: func(
 					_ *provisionersdk.Session,
 					_ *sdkproto.PlanRequest,
@@ -1031,6 +1041,9 @@ func TestProvisionerd(t *testing.T) {
 			return client, nil
 		}, provisionerd.LocalProvisioners{
 			"someprovisioner": createProvisionerClient(t, done, provisionerTestServer{
+				init: func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+					return &sdkproto.InitComplete{}
+				},
 				plan: func(
 					_ *provisionersdk.Session,
 					_ *sdkproto.PlanRequest,
@@ -1045,6 +1058,9 @@ func TestProvisionerd(t *testing.T) {
 				) *sdkproto.ApplyComplete {
 					return &sdkproto.ApplyComplete{}
 				},
+				graph: func(s *provisionersdk.Session, r *sdkproto.GraphRequest, canceledOrComplete <-chan struct{}) *sdkproto.GraphComplete {
+					return &sdkproto.GraphComplete{}
+				},
 			}),
 		})
 		require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
@@ -1125,6 +1141,12 @@ func TestProvisionerd(t *testing.T) {
 			}), nil
 		}, provisionerd.LocalProvisioners{
 			"someprovisioner": createProvisionerClient(t, done, provisionerTestServer{
+				init: func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+					return &sdkproto.InitComplete{}
+				},
+				graph: func(s *provisionersdk.Session, r *sdkproto.GraphRequest, canceledOrComplete <-chan struct{}) *sdkproto.GraphComplete {
+					return &sdkproto.GraphComplete{}
+				},
 				plan: func(
 					s *provisionersdk.Session,
 					_ *sdkproto.PlanRequest,
@@ -1253,9 +1275,15 @@ func createProvisionerClient(t *testing.T, done <-chan struct{}, server provisio
 }
 
 type provisionerTestServer struct {
+	init  func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete
 	parse func(s *provisionersdk.Session, r *sdkproto.ParseRequest, canceledOrComplete <-chan struct{}) *sdkproto.ParseComplete
 	plan  func(s *provisionersdk.Session, r *sdkproto.PlanRequest, canceledOrComplete <-chan struct{}) *sdkproto.PlanComplete
 	apply func(s *provisionersdk.Session, r *sdkproto.ApplyRequest, canceledOrComplete <-chan struct{}) *sdkproto.ApplyComplete
+	graph func(s *provisionersdk.Session, r *sdkproto.GraphRequest, canceledOrComplete <-chan struct{}) *sdkproto.GraphComplete
+}
+
+func (p *provisionerTestServer) Init(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+	return p.init(s, r, canceledOrComplete)
 }
 
 func (p *provisionerTestServer) Parse(s *provisionersdk.Session, r *sdkproto.ParseRequest, canceledOrComplete <-chan struct{}) *sdkproto.ParseComplete {
@@ -1270,6 +1298,10 @@ func (p *provisionerTestServer) Apply(s *provisionersdk.Session, r *sdkproto.App
 	return p.apply(s, r, canceledOrComplete)
 }
 
+func (p *provisionerTestServer) Graph(s *provisionersdk.Session, r *sdkproto.GraphRequest, canceledOrComplete <-chan struct{}) *sdkproto.GraphComplete {
+	return p.graph(s, r, canceledOrComplete)
+}
+
 func (p *provisionerDaemonTestServer) UploadFile(stream proto.DRPCProvisionerDaemon_UploadFileStream) error {
 	return p.uploadFile(stream)
 }
@@ -1359,3 +1391,16 @@ func (a *acquireOne) acquireWithCancel(stream proto.DRPCProvisionerDaemon_Acquir
 	}
 	return nil
 }
+
+func extractInit(t *testing.T) func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+	logger := slogtest.Make(t, nil)
+	return func(s *provisionersdk.Session, r *sdkproto.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete {
+		err := s.Files.ExtractArchive(s.Context(), logger, afero.NewOsFs(), r.TemplateSourceArchive)
+		if err != nil {
+			return &sdkproto.InitComplete{
+				Error: fmt.Sprintf("failed to extract template source archive: %v", err),
+			}
+		}
+		return &sdkproto.InitComplete{}
+	}
+}

provisionerd/runner/apply.go

@@ -0,0 +1,64 @@
+package runner
+
+import (
+	"context"
+	"time"
+
+	"cdr.dev/slog"
+	"github.com/coder/coder/v2/provisionerd/proto"
+	sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
+)
+
+func (r *Runner) apply(ctx context.Context, stage string, req *sdkproto.ApplyRequest) (
+	*sdkproto.ApplyComplete, *proto.FailedJob,
+) {
+	// use the notStopped so that if we attempt to gracefully cancel, the stream
+	// will still be available for us to send the cancel to the provisioner
+	err := r.session.Send(&sdkproto.Request{Type: &sdkproto.Request_Apply{Apply: req}})
+	if err != nil {
+		return nil, r.failedWorkspaceBuildf("start provision: %s", err)
+	}
+	nevermind := make(chan struct{})
+	defer close(nevermind)
+	go func() {
+		select {
+		case <-nevermind:
+			return
+		case <-r.notStopped.Done():
+			return
+		case <-r.notCanceled.Done():
+			_ = r.session.Send(&sdkproto.Request{
+				Type: &sdkproto.Request_Cancel{
+					Cancel: &sdkproto.CancelRequest{},
+				},
+			})
+		}
+	}()
+
+	for {
+		msg, err := r.session.Recv()
+		if err != nil {
+			return nil, r.failedWorkspaceBuildf("recv workspace provision: %s", err)
+		}
+		switch msgType := msg.Type.(type) {
+		case *sdkproto.Response_Log:
+			r.logProvisionerJobLog(context.Background(), msgType.Log.Level, "workspace provisioner job logged",
+				slog.F("level", msgType.Log.Level),
+				slog.F("output", msgType.Log.Output),
+				slog.F("workspace_build_id", r.job.GetWorkspaceBuild().WorkspaceBuildId),
+			)
+
+			r.queueLog(ctx, &proto.Log{
+				Source:    proto.LogSource_PROVISIONER,
+				Level:     msgType.Log.Level,
+				CreatedAt: time.Now().UnixMilli(),
+				Output:    msgType.Log.Output,
+				Stage:     stage,
+			})
+		case *sdkproto.Response_Apply:
+			return msgType.Apply, nil
+		default:
+			return nil, r.failedJobf("unexpected plan response type %T", msg.Type)
+		}
+	}
+}

provisionerd/runner/graph.go

@@ -0,0 +1,64 @@
+package runner
+
+import (
+	"context"
+	"time"
+
+	"cdr.dev/slog"
+	"github.com/coder/coder/v2/coderd/tracing"
+	"github.com/coder/coder/v2/provisionerd/proto"
+	sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
+)
+
+func (r *Runner) graph(ctx context.Context, req *sdkproto.GraphRequest) (*sdkproto.GraphComplete, *proto.FailedJob) {
+	ctx, span := r.startTrace(ctx, tracing.FuncName())
+	defer span.End()
+
+	err := r.session.Send(&sdkproto.Request{Type: &sdkproto.Request_Graph{Graph: req}})
+	if err != nil {
+		return nil, r.failedJobf("send graph request: %v", err)
+	}
+
+	nevermind := make(chan struct{})
+	defer close(nevermind)
+	go func() {
+		select {
+		case <-nevermind:
+			return
+		case <-r.notStopped.Done():
+			return
+		case <-r.notCanceled.Done():
+			_ = r.session.Send(&sdkproto.Request{
+				Type: &sdkproto.Request_Cancel{
+					Cancel: &sdkproto.CancelRequest{},
+				},
+			})
+		}
+	}()
+
+	for {
+		msg, err := r.session.Recv()
+		if err != nil {
+			return nil, r.failedJobf("receive graph response: %v", err)
+		}
+		switch msgType := msg.Type.(type) {
+		case *sdkproto.Response_Log:
+			r.logProvisionerJobLog(context.Background(), msgType.Log.Level, "terraform graphing",
+				slog.F("level", msgType.Log.Level),
+				slog.F("output", msgType.Log.Output),
+			)
+
+			r.queueLog(ctx, &proto.Log{
+				Source:    proto.LogSource_PROVISIONER,
+				Level:     msgType.Log.Level,
+				CreatedAt: time.Now().UnixMilli(),
+				Output:    msgType.Log.Output,
+				Stage:     "Graphing Infrastructure",
+			})
+		case *sdkproto.Response_Graph:
+			return msgType.Graph, nil
+		default:
+			return nil, r.failedJobf("unexpected graph response type %T", msg.Type)
+		}
+	}
+}

provisionerd/runner/init.go

@@ -0,0 +1,113 @@
+package runner
+
+import (
+	"bytes"
+	"context"
+	"time"
+
+	"cdr.dev/slog"
+	"github.com/coder/coder/v2/coderd/tracing"
+	"github.com/coder/coder/v2/provisionerd/proto"
+	sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
+)
+
+//nolint:revive
+func (r *Runner) init(ctx context.Context, omitModules bool, templateArchive []byte) (*sdkproto.InitComplete, *proto.FailedJob) {
+	ctx, span := r.startTrace(ctx, tracing.FuncName())
+	defer span.End()
+
+	err := r.session.Send(&sdkproto.Request{Type: &sdkproto.Request_Init{Init: &sdkproto.InitRequest{
+		TemplateSourceArchive: templateArchive,
+		OmitModuleFiles:       omitModules,
+	}}})
+	if err != nil {
+		return nil, r.failedJobf("send init request: %v", err)
+	}
+
+	nevermind := make(chan struct{})
+	defer close(nevermind)
+	go func() {
+		select {
+		case <-nevermind:
+			return
+		case <-r.notStopped.Done():
+			return
+		case <-r.notCanceled.Done():
+			_ = r.session.Send(&sdkproto.Request{
+				Type: &sdkproto.Request_Cancel{
+					Cancel: &sdkproto.CancelRequest{},
+				},
+			})
+		}
+	}()
+
+	var moduleFilesUpload *sdkproto.DataBuilder
+	for {
+		msg, err := r.session.Recv()
+		if err != nil {
+			return nil, r.failedJobf("receive init response: %v", err)
+		}
+		switch msgType := msg.Type.(type) {
+		case *sdkproto.Response_Log:
+			r.logProvisionerJobLog(context.Background(), msgType.Log.Level, "terraform initialization",
+				slog.F("level", msgType.Log.Level),
+				slog.F("output", msgType.Log.Output),
+			)
+
+			r.queueLog(ctx, &proto.Log{
+				Source:    proto.LogSource_PROVISIONER,
+				Level:     msgType.Log.Level,
+				CreatedAt: time.Now().UnixMilli(),
+				Output:    msgType.Log.Output,
+				Stage:     "Initializing Terraform Directory",
+			})
+		case *sdkproto.Response_DataUpload:
+			if omitModules {
+				return nil, r.failedJobf("received unexpected module files data upload when omitModules is true")
+			}
+			c := msgType.DataUpload
+			if c.UploadType != sdkproto.DataUploadType_UPLOAD_TYPE_MODULE_FILES {
+				return nil, r.failedJobf("invalid data upload type: %q", c.UploadType)
+			}
+
+			if moduleFilesUpload != nil {
+				return nil, r.failedJobf("multiple module data uploads received, only expect 1")
+			}
+
+			moduleFilesUpload, err = sdkproto.NewDataBuilder(c)
+			if err != nil {
+				return nil, r.failedJobf("create data builder: %s", err.Error())
+			}
+		case *sdkproto.Response_ChunkPiece:
+			if omitModules {
+				return nil, r.failedJobf("received unexpected module files data upload when omitModules is true")
+			}
+			c := msgType.ChunkPiece
+			if moduleFilesUpload == nil {
+				return nil, r.failedJobf("received chunk piece before module files data upload")
+			}
+
+			_, err := moduleFilesUpload.Add(c)
+			if err != nil {
+				return nil, r.failedJobf("module files, add chunk piece: %s", err.Error())
+			}
+		case *sdkproto.Response_Init:
+			if moduleFilesUpload != nil {
+				// If files were uploaded in multiple chunks, put them back together.
+				moduleFilesData, err := moduleFilesUpload.Complete()
+				if err != nil {
+					return nil, r.failedJobf("complete module files data upload: %s", err.Error())
+				}
+
+				if !bytes.Equal(msgType.Init.ModuleFilesHash, moduleFilesUpload.Hash) {
+					return nil, r.failedJobf("module files hash mismatch, uploaded: %x, expected: %x", moduleFilesUpload.Hash, msgType.Init.ModuleFilesHash)
+				}
+				msgType.Init.ModuleFiles = moduleFilesData
+			}
+
+			return msgType.Init, nil
+		default:
+			return nil, r.failedJobf("unexpected init response type %T", msg.Type)
+		}
+	}
+}

provisionerd/runner/plan.go

@@ -0,0 +1,64 @@
+package runner
+
+import (
+	"context"
+	"time"
+
+	"cdr.dev/slog"
+	"github.com/coder/coder/v2/coderd/tracing"
+	"github.com/coder/coder/v2/provisionerd/proto"
+	sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
+)
+
+func (r *Runner) plan(ctx context.Context, stage string, req *sdkproto.PlanRequest) (*sdkproto.PlanComplete, *proto.FailedJob) {
+	ctx, span := r.startTrace(ctx, tracing.FuncName())
+	defer span.End()
+
+	err := r.session.Send(&sdkproto.Request{Type: &sdkproto.Request_Plan{Plan: req}})
+	if err != nil {
+		return nil, r.failedJobf("send plan request: %v", err)
+	}
+
+	nevermind := make(chan struct{})
+	defer close(nevermind)
+	go func() {
+		select {
+		case <-nevermind:
+			return
+		case <-r.notStopped.Done():
+			return
+		case <-r.notCanceled.Done():
+			_ = r.session.Send(&sdkproto.Request{
+				Type: &sdkproto.Request_Cancel{
+					Cancel: &sdkproto.CancelRequest{},
+				},
+			})
+		}
+	}()
+
+	for {
+		msg, err := r.session.Recv()
+		if err != nil {
+			return nil, r.failedJobf("receive plan response: %v", err)
+		}
+		switch msgType := msg.Type.(type) {
+		case *sdkproto.Response_Log:
+			r.logProvisionerJobLog(context.Background(), msgType.Log.Level, "terraform planning",
+				slog.F("level", msgType.Log.Level),
+				slog.F("output", msgType.Log.Output),
+			)
+
+			r.queueLog(ctx, &proto.Log{
+				Source:    proto.LogSource_PROVISIONER,
+				Level:     msgType.Log.Level,
+				CreatedAt: time.Now().UnixMilli(),
+				Output:    msgType.Log.Output,
+				Stage:     stage,
+			})
+		case *sdkproto.Response_Plan:
+			return msgType.Plan, nil
+		default:
+			return nil, r.failedJobf("unexpected plan response type %T", msg.Type)
+		}
+	}
+}

provisionerd/runner/quota.go

@@ -1,11 +0,0 @@
-package runner
-
-import "github.com/coder/coder/v2/provisionersdk/proto"
-
-func sumDailyCost(resources []*proto.Resource) int {
-	var sum int
-	for _, r := range resources {
-		sum += int(r.DailyCost)
-	}
-	return sum
-}

provisionerd/runner/runner.go

@@ -1,7 +1,6 @@
 package runner
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
@@ -515,7 +514,6 @@ func (r *Runner) runTemplateImport(ctx context.Context) (*proto.CompletedJob, *p
 	defer span.End()
 
 	failedJob := r.configure(&sdkproto.Config{
-		TemplateSourceArchive:      r.job.GetTemplateSourceArchive(),
 		TemplateId:                 strings2.EmptyToNil(r.job.GetTemplateImport().Metadata.TemplateId),
 		TemplateVersionId:          strings2.EmptyToNil(r.job.GetTemplateImport().Metadata.TemplateVersionId),
 		ExpReuseTerraformWorkspace: ptr.Ref(false),
@@ -524,6 +522,18 @@ func (r *Runner) runTemplateImport(ctx context.Context) (*proto.CompletedJob, *p
 		return nil, failedJob
 	}
 
+	// Initialize the Terraform working directory
+	initResp, failedInit := r.init(ctx, false, r.job.GetTemplateSourceArchive())
+	if failedInit != nil {
+		return nil, failedInit
+	}
+	if initResp == nil {
+		return nil, r.failedJobf("template import init returned nil response")
+	}
+	if initResp.Error != "" {
+		return nil, r.failedJobf("template import init error: %s", initResp.Error)
+	}
+
 	// Parse parameters and update the job with the parameter specs
 	r.queueLog(ctx, &proto.Log{
 		Source:    proto.LogSource_PROVISIONER_DAEMON,
@@ -560,7 +570,7 @@ func (r *Runner) runTemplateImport(ctx context.Context) (*proto.CompletedJob, *p
 		CoderUrl:             r.job.GetTemplateImport().Metadata.CoderUrl,
 		WorkspaceOwnerGroups: r.job.GetTemplateImport().Metadata.WorkspaceOwnerGroups,
 		WorkspaceTransition:  sdkproto.WorkspaceTransition_START,
-	}, false)
+	})
 	if err != nil {
 		return nil, r.failedJobf("template import provision for start: %s", err)
 	}
@@ -576,8 +586,7 @@ func (r *Runner) runTemplateImport(ctx context.Context) (*proto.CompletedJob, *p
 		CoderUrl:             r.job.GetTemplateImport().Metadata.CoderUrl,
 		WorkspaceOwnerGroups: r.job.GetTemplateImport().Metadata.WorkspaceOwnerGroups,
 		WorkspaceTransition:  sdkproto.WorkspaceTransition_STOP,
-	}, true, // Modules downloaded on the start provision
-	)
+	})
 	if err != nil {
 		return nil, r.failedJobf("template import provision for stop: %s", err)
 	}
@@ -597,12 +606,13 @@ func (r *Runner) runTemplateImport(ctx context.Context) (*proto.CompletedJob, *p
 				RichParameters:             startProvision.Parameters,
 				ExternalAuthProvidersNames: externalAuthProviderNames,
 				ExternalAuthProviders:      startProvision.ExternalAuthProviders,
-				StartModules:               startProvision.Modules,
-				StopModules:                stopProvision.Modules,
-				Presets:                    startProvision.Presets,
-				Plan:                       startProvision.Plan,
-				// ModuleFiles are not on the stopProvision. So grab from the startProvision.
-				ModuleFiles: startProvision.ModuleFiles,
+				// TODO: These are defined as different, but can they be?
+				//   Terraform downloads modules regardless of `count`, so this should be the same
+				StartModules: initResp.Modules,
+				StopModules:  initResp.Modules,
+				Presets:      startProvision.Presets,
+				Plan:         startProvision.Plan,
+				ModuleFiles:  initResp.ModuleFiles,
 				// ModuleFileHash will be populated if the file is uploaded async
 				ModuleFilesHash:   []byte{},
 				HasAiTasks:        startProvision.HasAITasks,
@@ -666,10 +676,8 @@ type templateImportProvision struct {
 	Resources             []*sdkproto.Resource
 	Parameters            []*sdkproto.RichParameter
 	ExternalAuthProviders []*sdkproto.ExternalAuthProviderResource
-	Modules               []*sdkproto.Module
 	Presets               []*sdkproto.Preset
 	Plan                  json.RawMessage
-	ModuleFiles           []byte
 	HasAITasks            bool
 	HasExternalAgents     bool
 }
@@ -677,8 +685,8 @@ type templateImportProvision struct {
 // Performs a dry-run provision when importing a template.
 // This is used to detect resources that would be provisioned for a workspace in various states.
 // It doesn't define values for rich parameters as they're unknown during template import.
-func (r *Runner) runTemplateImportProvision(ctx context.Context, variableValues []*sdkproto.VariableValue, metadata *sdkproto.Metadata, omitModules bool) (*templateImportProvision, error) {
-	return r.runTemplateImportProvisionWithRichParameters(ctx, variableValues, nil, metadata, omitModules)
+func (r *Runner) runTemplateImportProvision(ctx context.Context, variableValues []*sdkproto.VariableValue, metadata *sdkproto.Metadata) (*templateImportProvision, error) {
+	return r.runTemplateImportProvisionWithRichParameters(ctx, variableValues, nil, metadata)
 }
 
 // Performs a dry-run provision with provided rich parameters.
@@ -688,7 +696,6 @@ func (r *Runner) runTemplateImportProvisionWithRichParameters(
 	variableValues []*sdkproto.VariableValue,
 	richParameterValues []*sdkproto.RichParameterValue,
 	metadata *sdkproto.Metadata,
-	omitModules bool,
 ) (*templateImportProvision, error) {
 	ctx, span := r.startTrace(ctx, tracing.FuncName())
 	defer span.End()
@@ -700,126 +707,48 @@ func (r *Runner) runTemplateImportProvisionWithRichParameters(
 	case sdkproto.WorkspaceTransition_STOP:
 		stage = "Detecting ephemeral resources"
 	}
-	// use the notStopped so that if we attempt to gracefully cancel, the stream will still be available for us
-	// to send the cancel to the provisioner
-	err := r.session.Send(&sdkproto.Request{Type: &sdkproto.Request_Plan{Plan: &sdkproto.PlanRequest{
-		Metadata:            metadata,
-		RichParameterValues: richParameterValues,
-		// Template import has no previous values
-		PreviousParameterValues: make([]*sdkproto.RichParameterValue, 0),
+
+	planComplete, failed := r.plan(ctx, stage, &sdkproto.PlanRequest{
+		Metadata:                metadata,
+		RichParameterValues:     richParameterValues,
 		VariableValues:          variableValues,
-		OmitModuleFiles:         omitModules,
-	}}})
-	if err != nil {
-		return nil, xerrors.Errorf("start provision: %w", err)
+		ExternalAuthProviders:   nil,
+		PreviousParameterValues: nil,
+		State:                   nil,
+	})
+	if failed != nil {
+		return nil, xerrors.Errorf("plan during template import provision: %w", failed)
 	}
-	nevermind := make(chan struct{})
-	defer close(nevermind)
-	go func() {
-		select {
-		case <-nevermind:
-			return
-		case <-r.notStopped.Done():
-			return
-		case <-r.notCanceled.Done():
-			_ = r.session.Send(&sdkproto.Request{
-				Type: &sdkproto.Request_Cancel{
-					Cancel: &sdkproto.CancelRequest{},
-				},
-			})
-		}
-	}()
-
-	var moduleFilesUpload *sdkproto.DataBuilder
-	for {
-		msg, err := r.session.Recv()
-		if err != nil {
-			return nil, xerrors.Errorf("recv import provision: %w", err)
-		}
-
-		switch msgType := msg.Type.(type) {
-		case *sdkproto.Response_Log:
-			r.logProvisionerJobLog(context.Background(), msgType.Log.Level, "template import provision job logged",
-				slog.F("level", msgType.Log.Level),
-				slog.F("output", msgType.Log.Output),
-			)
-			r.queueLog(ctx, &proto.Log{
-				Source:    proto.LogSource_PROVISIONER,
-				Level:     msgType.Log.Level,
-				CreatedAt: time.Now().UnixMilli(),
-				Output:    msgType.Log.Output,
-				Stage:     stage,
-			})
-		case *sdkproto.Response_DataUpload:
-			c := msgType.DataUpload
-			if c.UploadType != sdkproto.DataUploadType_UPLOAD_TYPE_MODULE_FILES {
-				return nil, xerrors.Errorf("invalid data upload type: %q", c.UploadType)
-			}
-
-			if moduleFilesUpload != nil {
-				return nil, xerrors.New("multiple module data uploads received, only expect 1")
-			}
-
-			moduleFilesUpload, err = sdkproto.NewDataBuilder(c)
-			if err != nil {
-				return nil, xerrors.Errorf("create data builder: %w", err)
-			}
-		case *sdkproto.Response_ChunkPiece:
-			c := msgType.ChunkPiece
-			if moduleFilesUpload == nil {
-				return nil, xerrors.New("received chunk piece before module files data upload")
-			}
-
-			_, err := moduleFilesUpload.Add(c)
-			if err != nil {
-				return nil, xerrors.Errorf("module files, add chunk piece: %w", err)
-			}
-		case *sdkproto.Response_Plan:
-			c := msgType.Plan
-			if c.Error != "" {
-				r.logger.Info(context.Background(), "dry-run provision failure",
-					slog.F("error", c.Error),
-				)
-
-				return nil, xerrors.New(c.Error)
-			}
-
-			if moduleFilesUpload != nil && len(c.ModuleFiles) > 0 {
-				return nil, xerrors.New("module files were uploaded and module files were returned in the plan response. Only one of these should be set")
-			}
-
-			r.logger.Info(context.Background(), "parse dry-run provision successful",
-				slog.F("resource_count", len(c.Resources)),
-				slog.F("resources", resourceNames(c.Resources)),
-			)
-
-			moduleFilesData := c.ModuleFiles
-			if moduleFilesUpload != nil {
-				uploadData, err := moduleFilesUpload.Complete()
-				if err != nil {
-					return nil, xerrors.Errorf("module files, complete upload: %w", err)
-				}
-				moduleFilesData = uploadData
-				if !bytes.Equal(c.ModuleFilesHash, moduleFilesUpload.Hash) {
-					return nil, xerrors.Errorf("module files hash mismatch, uploaded: %x, expected: %x", moduleFilesUpload.Hash, c.ModuleFilesHash)
-				}
-			}
-			return &templateImportProvision{
-				Resources:             c.Resources,
-				Parameters:            c.Parameters,
-				ExternalAuthProviders: c.ExternalAuthProviders,
-				Modules:               c.Modules,
-				Presets:               c.Presets,
-				Plan:                  c.Plan,
-				ModuleFiles:           moduleFilesData,
-				HasAITasks:            c.HasAiTasks,
-				HasExternalAgents:     c.HasExternalAgents,
-			}, nil
-		default:
-			return nil, xerrors.Errorf("invalid message type %q received from provisioner",
-				reflect.TypeOf(msg.Type).String())
-		}
+	if planComplete == nil {
+		return nil, xerrors.New("plan during template import provision returned nil response")
 	}
+	if planComplete.Error != "" {
+		return nil, xerrors.Errorf("plan during template import provision error: %s", planComplete.Error)
+	}
+
+	graphComplete, failed := r.graph(ctx, &sdkproto.GraphRequest{
+		Metadata: metadata,
+		Source:   sdkproto.GraphSource_SOURCE_PLAN,
+	})
+	if failed != nil {
+		return nil, xerrors.Errorf("graph during template import provision: %w", failed)
+	}
+	if graphComplete == nil {
+		return nil, xerrors.New("graph during template import provision returned nil response")
+	}
+	if graphComplete.Error != "" {
+		return nil, xerrors.Errorf("graph during template import provision error: %s", graphComplete.Error)
+	}
+
+	return &templateImportProvision{
+		Resources:             graphComplete.Resources,
+		Parameters:            graphComplete.Parameters,
+		ExternalAuthProviders: graphComplete.ExternalAuthProviders,
+		Presets:               graphComplete.Presets,
+		Plan:                  planComplete.Plan,
+		HasAITasks:            graphComplete.HasAiTasks,
+		HasExternalAgents:     graphComplete.HasExternalAgents,
+	}, nil
 }
 
 func (r *Runner) runTemplateDryRun(ctx context.Context) (*proto.CompletedJob, *proto.FailedJob) {
@@ -854,19 +783,28 @@ func (r *Runner) runTemplateDryRun(ctx context.Context) (*proto.CompletedJob, *p
 		metadata.WorkspaceOwnerId = id.String()
 	}
 
-	failedJob := r.configure(&sdkproto.Config{
-		TemplateSourceArchive: r.job.GetTemplateSourceArchive(),
-	})
+	failedJob := r.configure(&sdkproto.Config{})
 	if failedJob != nil {
 		return nil, failedJob
 	}
 
+	// Initialize the Terraform working directory
+	initResp, failedJob := r.init(ctx, false, r.job.GetTemplateSourceArchive())
+	if failedJob != nil {
+		return nil, failedJob
+	}
+	if initResp == nil {
+		return nil, r.failedJobf("template dry-run init returned nil response")
+	}
+	if initResp.Error != "" {
+		return nil, r.failedJobf("template dry-run init error: %s", initResp.Error)
+	}
+
 	// Run the template import provision task since it's already a dry run.
 	provision, err := r.runTemplateImportProvisionWithRichParameters(ctx,
 		r.job.GetTemplateDryRun().GetVariableValues(),
 		r.job.GetTemplateDryRun().GetRichParameterValues(),
 		metadata,
-		false,
 	)
 	if err != nil {
 		return nil, r.failedJobf("run dry-run provision job: %s", err)
@@ -877,73 +815,14 @@ func (r *Runner) runTemplateDryRun(ctx context.Context) (*proto.CompletedJob, *p
 		Type: &proto.CompletedJob_TemplateDryRun_{
 			TemplateDryRun: &proto.CompletedJob_TemplateDryRun{
 				Resources: provision.Resources,
-				Modules:   provision.Modules,
+				Modules:   initResp.Modules,
 			},
 		},
 	}, nil
 }
 
-func (r *Runner) buildWorkspace(ctx context.Context, stage string, req *sdkproto.Request) (
-	*sdkproto.Response, *proto.FailedJob,
-) {
-	// use the notStopped so that if we attempt to gracefully cancel, the stream
-	// will still be available for us to send the cancel to the provisioner
-	err := r.session.Send(req)
-	if err != nil {
-		return nil, r.failedWorkspaceBuildf("start provision: %s", err)
-	}
-	nevermind := make(chan struct{})
-	defer close(nevermind)
-	go func() {
-		select {
-		case <-nevermind:
-			return
-		case <-r.notStopped.Done():
-			return
-		case <-r.notCanceled.Done():
-			_ = r.session.Send(&sdkproto.Request{
-				Type: &sdkproto.Request_Cancel{
-					Cancel: &sdkproto.CancelRequest{},
-				},
-			})
-		}
-	}()
-
-	for {
-		msg, err := r.session.Recv()
-		if err != nil {
-			return nil, r.failedWorkspaceBuildf("recv workspace provision: %s", err)
-		}
-		switch msgType := msg.Type.(type) {
-		case *sdkproto.Response_Log:
-			r.logProvisionerJobLog(context.Background(), msgType.Log.Level, "workspace provisioner job logged",
-				slog.F("level", msgType.Log.Level),
-				slog.F("output", msgType.Log.Output),
-				slog.F("workspace_build_id", r.job.GetWorkspaceBuild().WorkspaceBuildId),
-			)
-
-			r.queueLog(ctx, &proto.Log{
-				Source:    proto.LogSource_PROVISIONER,
-				Level:     msgType.Log.Level,
-				CreatedAt: time.Now().UnixMilli(),
-				Output:    msgType.Log.Output,
-				Stage:     stage,
-			})
-		case *sdkproto.Response_DataUpload:
-			continue // Only for template imports
-		case *sdkproto.Response_ChunkPiece:
-			continue // Only for template imports
-		default:
-			// Stop looping!
-			return msg, nil
-		}
-	}
-}
-
-func (r *Runner) commitQuota(ctx context.Context, resources []*sdkproto.Resource) *proto.FailedJob {
-	cost := sumDailyCost(resources)
+func (r *Runner) commitQuota(ctx context.Context, cost int32) *proto.FailedJob {
 	r.logger.Debug(ctx, "committing quota",
-		slog.F("resources", resourceNames(resources)),
 		slog.F("cost", cost),
 	)
 	if cost == 0 {
@@ -953,9 +832,8 @@ func (r *Runner) commitQuota(ctx context.Context, resources []*sdkproto.Resource
 	const stage = "Commit quota"
 
 	resp, err := r.quotaCommitter.CommitQuota(ctx, &proto.CommitQuotaRequest{
-		JobId: r.job.JobId,
-		// #nosec G115 - Safe conversion as cost is expected to be within int32 range for provisioning costs
-		DailyCost: int32(cost),
+		JobId:     r.job.JobId,
+		DailyCost: cost,
 	})
 	if err != nil {
 		r.queueLog(ctx, &proto.Log{
@@ -1014,8 +892,6 @@ func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *p
 	}
 
 	failedJob := r.configure(&sdkproto.Config{
-		TemplateSourceArchive:      r.job.GetTemplateSourceArchive(),
-		State:                      r.job.GetWorkspaceBuild().State,
 		ProvisionerLogLevel:        r.job.GetWorkspaceBuild().LogLevel,
 		TemplateId:                 strings2.EmptyToNil(r.job.GetWorkspaceBuild().Metadata.TemplateId),
 		TemplateVersionId:          strings2.EmptyToNil(r.job.GetWorkspaceBuild().Metadata.TemplateVersionId),
@@ -1025,25 +901,53 @@ func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *p
 		return nil, failedJob
 	}
 
-	resp, failed := r.buildWorkspace(ctx, "Planning infrastructure", &sdkproto.Request{
-		Type: &sdkproto.Request_Plan{
-			Plan: &sdkproto.PlanRequest{
-				OmitModuleFiles:         true, // Only useful for template imports
-				Metadata:                r.job.GetWorkspaceBuild().Metadata,
-				RichParameterValues:     r.job.GetWorkspaceBuild().RichParameterValues,
-				PreviousParameterValues: r.job.GetWorkspaceBuild().PreviousParameterValues,
-				VariableValues:          r.job.GetWorkspaceBuild().VariableValues,
-				ExternalAuthProviders:   r.job.GetWorkspaceBuild().ExternalAuthProviders,
+	// timings collects all timings from each phase of the build
+	timings := make([]*sdkproto.Timing, 0)
+
+	// Initialize the Terraform working directory
+	initComplete, failedJob := r.init(ctx, true, r.job.GetTemplateSourceArchive())
+	if failedJob != nil {
+		return nil, failedJob
+	}
+	if initComplete == nil {
+		return nil, r.failedWorkspaceBuildf("invalid message type received from provisioner during init")
+	}
+	// Collect init timings
+	timings = append(timings, initComplete.Timings...)
+	if initComplete.Error != "" {
+		r.logger.Warn(context.Background(), "init request failed",
+			slog.F("error", initComplete.Error),
+		)
+
+		return nil, &proto.FailedJob{
+			JobId: r.job.JobId,
+			Error: initComplete.Error,
+			Type: &proto.FailedJob_WorkspaceBuild_{
+				WorkspaceBuild: &proto.FailedJob_WorkspaceBuild{
+					State:   r.job.GetWorkspaceBuild().State,
+					Timings: timings,
+				},
 			},
-		},
+		}
+	}
+
+	// Run `terraform plan`
+	planComplete, failed := r.plan(ctx, "Planning Infrastructure", &sdkproto.PlanRequest{
+		Metadata:                r.job.GetWorkspaceBuild().Metadata,
+		RichParameterValues:     r.job.GetWorkspaceBuild().RichParameterValues,
+		VariableValues:          r.job.GetWorkspaceBuild().VariableValues,
+		ExternalAuthProviders:   r.job.GetWorkspaceBuild().ExternalAuthProviders,
+		PreviousParameterValues: r.job.GetWorkspaceBuild().PreviousParameterValues,
+		State:                   r.job.GetWorkspaceBuild().State,
 	})
 	if failed != nil {
 		return nil, failed
 	}
-	planComplete := resp.GetPlan()
 	if planComplete == nil {
-		return nil, r.failedWorkspaceBuildf("invalid message type %T received from provisioner", resp.Type)
+		return nil, r.failedWorkspaceBuildf("invalid message type received from provisioner during plan")
 	}
+	// Collect plan timings
+	timings = append(timings, planComplete.Timings...)
 	if planComplete.Error != "" {
 		r.logger.Warn(context.Background(), "plan request failed",
 			slog.F("error", planComplete.Error),
@@ -1053,27 +957,28 @@ func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *p
 			JobId: r.job.JobId,
 			Error: planComplete.Error,
 			Type: &proto.FailedJob_WorkspaceBuild_{
-				WorkspaceBuild: &proto.FailedJob_WorkspaceBuild{},
+				WorkspaceBuild: &proto.FailedJob_WorkspaceBuild{
+					Timings: timings,
+				},
 			},
 		}
 	}
-	if len(planComplete.AiTasks) > 1 {
-		return nil, r.failedWorkspaceBuildf("only one 'coder_ai_task' resource can be provisioned per template")
+
+	if planComplete.AiTaskCount > 1 {
+		return nil, r.failedWorkspaceBuildf("only one 'coder_ai_task' resource can be provisioned per template, found %d", planComplete.AiTaskCount)
 	}
 
-	r.logger.Info(context.Background(), "plan request successful",
-		slog.F("resource_count", len(planComplete.Resources)),
-		slog.F("resources", resourceNames(planComplete.Resources)),
-	)
+	r.logger.Info(context.Background(), "plan request successful")
 	r.flushQueuedLogs(ctx)
 	if commitQuota {
-		failed = r.commitQuota(ctx, planComplete.Resources)
+		failed = r.commitQuota(ctx, planComplete.GetDailyCost())
 		r.flushQueuedLogs(ctx)
 		if failed != nil {
 			return nil, failed
 		}
 	}
 
+	// Run Terraform Apply
 	r.queueLog(ctx, &proto.Log{
 		Source:    proto.LogSource_PROVISIONER_DAEMON,
 		Level:     sdkproto.LogLevel_INFO,
@@ -1081,24 +986,17 @@ func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *p
 		CreatedAt: time.Now().UnixMilli(),
 	})
 
-	resp, failed = r.buildWorkspace(ctx, applyStage, &sdkproto.Request{
-		Type: &sdkproto.Request_Apply{
-			Apply: &sdkproto.ApplyRequest{
-				Metadata: r.job.GetWorkspaceBuild().Metadata,
-			},
-		},
+	applyComplete, failed := r.apply(ctx, applyStage, &sdkproto.ApplyRequest{
+		Metadata: r.job.GetWorkspaceBuild().Metadata,
 	})
 	if failed != nil {
 		return nil, failed
 	}
-	applyComplete := resp.GetApply()
 	if applyComplete == nil {
-		return nil, r.failedWorkspaceBuildf("invalid message type %T received from provisioner", resp.Type)
+		return nil, r.failedWorkspaceBuildf("invalid message type received from provisioner during apply")
 	}
-
-	// Prepend the plan timings (since they occurred first).
-	applyComplete.Timings = append(planComplete.Timings, applyComplete.Timings...)
-
+	// Collect apply timings
+	timings = append(timings, applyComplete.Timings...)
 	if applyComplete.Error != "" {
 		r.logger.Warn(context.Background(), "apply failed; updating state",
 			slog.F("error", applyComplete.Error),
@@ -1111,15 +1009,46 @@ func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *p
 			Type: &proto.FailedJob_WorkspaceBuild_{
 				WorkspaceBuild: &proto.FailedJob_WorkspaceBuild{
 					State:   applyComplete.State,
-					Timings: applyComplete.Timings,
+					Timings: timings,
+				},
+			},
+		}
+	}
+
+	// Run Terraform Graph
+	graphComplete, failed := r.graph(ctx, &sdkproto.GraphRequest{
+		Metadata: r.job.GetWorkspaceBuild().Metadata,
+		Source:   sdkproto.GraphSource_SOURCE_STATE,
+	})
+	if failed != nil {
+		return nil, failed
+	}
+	if graphComplete == nil {
+		return nil, r.failedWorkspaceBuildf("invalid message type received from provisioner during graph")
+	}
+	// Collect graph timings
+	timings = append(timings, graphComplete.Timings...)
+	if graphComplete.Error != "" {
+		r.logger.Warn(context.Background(), "graph request failed",
+			slog.F("error", planComplete.Error),
+		)
+
+		return nil, &proto.FailedJob{
+			JobId: r.job.JobId,
+			Error: graphComplete.Error,
+			Type: &proto.FailedJob_WorkspaceBuild_{
+				WorkspaceBuild: &proto.FailedJob_WorkspaceBuild{
+					// Graph does not change the state, so return the state returned from apply.
+					State:   applyComplete.State,
+					Timings: timings,
 				},
 			},
 		}
 	}
 
 	r.logger.Info(context.Background(), "apply successful",
-		slog.F("resource_count", len(applyComplete.Resources)),
-		slog.F("resources", resourceNames(applyComplete.Resources)),
+		slog.F("resource_count", len(graphComplete.Resources)),
+		slog.F("resources", resourceNames(graphComplete.Resources)),
 		slog.F("state_len", len(applyComplete.State)),
 	)
 	r.flushQueuedLogs(ctx)
@@ -1129,15 +1058,14 @@ func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *p
 		Type: &proto.CompletedJob_WorkspaceBuild_{
 			WorkspaceBuild: &proto.CompletedJob_WorkspaceBuild{
 				State:     applyComplete.State,
-				Resources: applyComplete.Resources,
-				Timings:   applyComplete.Timings,
-				// Modules are created on disk by `terraform init`, and that is only
-				// called by `plan`. `apply` does not modify them, so we can use the
-				// modules from the plan response.
-				Modules: planComplete.Modules,
+				Resources: graphComplete.Resources,
+				Timings:   timings,
+				// Modules files are omitted for workspace builds, but the modules.json metadata
+				// is available from init to return.
+				Modules: initComplete.Modules,
 				// Resource replacements are discovered at plan time, only.
 				ResourceReplacements: planComplete.ResourceReplacements,
-				AiTasks:              applyComplete.AiTasks,
+				AiTasks:              graphComplete.AiTasks,
 			},
 		},
 	}, nil