simplify function

Author: SasSwart

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/docker/docker v26.1.5+incompatible
 	github.com/google/uuid v1.6.0
 	github.com/hashicorp/go-cty v1.4.1-0.20200414143053-d3edf31b6320
+	github.com/hashicorp/terraform-plugin-log v0.9.0
 	github.com/hashicorp/terraform-plugin-sdk/v2 v2.36.1
 	github.com/masterminds/semver v1.5.0
 	github.com/mitchellh/mapstructure v1.5.0
@@ -50,7 +51,6 @@ require (
 	github.com/hashicorp/terraform-exec v0.22.0 // indirect
 	github.com/hashicorp/terraform-json v0.24.0 // indirect
 	github.com/hashicorp/terraform-plugin-go v0.26.0 // indirect
-	github.com/hashicorp/terraform-plugin-log v0.9.0 // indirect
 	github.com/hashicorp/terraform-registry-address v0.2.4 // indirect
 	github.com/hashicorp/terraform-svchost v0.1.1 // indirect
 	github.com/hashicorp/yamux v0.1.1 // indirect

provider/agent.go

@@ -478,27 +478,22 @@ func updateInitScript(resourceData *schema.ResourceData, i interface{}) diag.Dia
 }
 
 func agentAuthToken(ctx context.Context, agentID string) string {
-	// Most of the time, we will generate a new token for the agent.
-	// In the case of a prebuilt workspace being claimed, we will override with
-	// an existing token provided below.
-	token := uuid.NewString()
-
 	existingToken := helpers.OptionalEnv(RunningAgentTokenEnvironmentVariable(agentID))
-	logFields := map[string]interface{}{
-		"agent_id":       agentID,
-		"token_provided": existingToken != "",
-	}
-	if existingToken != "" {
-		// An existing token was provided for this agent. That means that this
-		// is a prebuilt workspace in the process of being claimed.
-		// We should reuse the token.
-		tflog.Info(ctx, "using provided agent token for prebuild", logFields)
-		token = existingToken
-	} else {
-		tflog.Info(ctx, "using a new agent token", logFields)
+	if existingToken == "" {
+		// Most of the time, we will generate a new token for the agent.
+		// In the case of a prebuilt workspace being claimed, we will override with
+		// an existing token provided below.
+		token := uuid.NewString()
+		return token
 	}
 
-	return token
+	// An existing token was provided for this agent. That means that this
+	// is a prebuilt workspace in the process of being claimed.
+	// We should reuse the token.
+	tflog.Info(ctx, "using provided agent token for prebuild", map[string]interface{}{
+		"agent_id": agentID,
+	})
+	return existingToken
 }
 
 // RunningAgentTokenEnvironmentVariable returns the name of an environment variable

provider/workspace.go

@@ -27,13 +27,13 @@ func workspaceDataSource() *schema.Resource {
 			}
 			_ = rd.Set("start_count", count)
 
-			isPrebuild := isPrebuiltWorkspace()
-			prebuildCount := 0
-			if isPrebuild {
-				prebuildCount = 1
+			if isPrebuiltWorkspace() {
+				_ = rd.Set("prebuild_count", 1)
 				_ = rd.Set("is_prebuild", true)
+			} else {
+				_ = rd.Set("prebuild_count", 0)
+				_ = rd.Set("is_prebuild", false)
 			}
-			_ = rd.Set("prebuild_count", prebuildCount)
 
 			name := helpers.OptionalEnvOrDefault("CODER_WORKSPACE_NAME", "default")
 			rd.Set("name", name)