feat: add env_vars as input to run_sh

core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/run_python.go

@@ -203,8 +203,13 @@ func (builtin *RunPythonCapabilities) Interpret(_ string, arguments *builtin_arg
 		}
 	}
 
+	envVars, interpretationErr := extractEnvVarsIfDefined(arguments)
+	if err != nil {
+		return nil, interpretationErr
+	}
+
 	// build a service config from image and files artifacts expansion.
-	builtin.serviceConfig, err = getServiceConfig(image, filesArtifactExpansion)
+	builtin.serviceConfig, err = getServiceConfig(image, filesArtifactExpansion, envVars)
 	if err != nil {
 		return nil, startosis_errors.WrapWithInterpretationError(err, "An error occurred creating service config using image '%s'", image)
 	}

core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/run_sh.go

@@ -58,6 +58,12 @@ func NewRunShService(serviceNetwork service_network.ServiceNetwork, runtimeValue
 					IsOptional:        true,
 					ZeroValueProvider: builtin_argument.ZeroValueProvider[*starlark.List],
 				},
+				{
+					Name:              EnvVarsArgName,
+					IsOptional:        true,
+					ZeroValueProvider: builtin_argument.ZeroValueProvider[*starlark.Dict],
+					Validator:         nil,
+				},
 				{
 					Name:              WaitArgName,
 					IsOptional:        true,
@@ -90,6 +96,7 @@ func NewRunShService(serviceNetwork service_network.ServiceNetwork, runtimeValue
 			FilesArgName:      true,
 			StoreFilesArgName: true,
 			WaitArgName:       true,
+			EnvVarsArgName:    true,
 		},
 	}
 }
@@ -147,8 +154,13 @@ func (builtin *RunShCapabilities) Interpret(_ string, arguments *builtin_argumen
 		}
 	}
 
+	envVars, interpretationErr := extractEnvVarsIfDefined(arguments)
+	if err != nil {
+		return nil, interpretationErr
+	}
+
 	// build a service config from image and files artifacts expansion.
-	builtin.serviceConfig, err = getServiceConfig(image, filesArtifactExpansion)
+	builtin.serviceConfig, err = getServiceConfig(image, filesArtifactExpansion, envVars)
 	if err != nil {
 		return nil, startosis_errors.WrapWithInterpretationError(err, "An error occurred creating service config using image '%s'", image)
 	}

core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/tasks_shared.go

@@ -31,6 +31,7 @@ const (
 	StoreFilesArgName = "store"
 	WaitArgName       = "wait"
 	FilesArgName      = "files"
+	EnvVarsArgName    = "env_vars"
 
 	newlineChar = "\n"
 
@@ -248,7 +249,11 @@ func resultMapToString(resultMap map[string]starlark.Comparable, builtinNameForL
 	return fmt.Sprintf("Command returned with exit code '%v' and the following output: %v", exitCode, outputStr)
 }
 
-func getServiceConfig(image string, filesArtifactExpansion *service_directory.FilesArtifactsExpansion) (*service.ServiceConfig, error) {
+func getServiceConfig(
+	image string,
+	filesArtifactExpansion *service_directory.FilesArtifactsExpansion,
+	envVars *map[string]string,
+) (*service.ServiceConfig, error) {
 	serviceConfig, err := service.CreateServiceConfig(
 		image,
 		nil,
@@ -262,7 +267,7 @@ func getServiceConfig(image string, filesArtifactExpansion *service_directory.Fi
 		//  command is completed
 		runTailCommandToPreventContainerToStopOnCreating,
 		nil,
-		nil,
+		*envVars,
 		filesArtifactExpansion,
 		nil,
 		0,
@@ -293,3 +298,21 @@ func removeService(ctx context.Context, serviceNetwork service_network.ServiceNe
 	}
 	return nil
 }
+
+func extractEnvVarsIfDefined(arguments *builtin_argument.ArgumentValuesSet) (*map[string]string, *startosis_errors.InterpretationError) {
+	envVars := map[string]string{}
+	if arguments.IsSet(EnvVarsArgName) {
+		envVarsStarlark, err := builtin_argument.ExtractArgumentValue[*starlark.Dict](arguments, EnvVarsArgName)
+		if err != nil {
+			return nil, startosis_errors.WrapWithInterpretationError(err, "Unable to extract value for '%s' argument", EnvVarsArgName)
+		}
+		if envVarsStarlark != nil && envVarsStarlark.Len() > 0 {
+			var interpretationErr *startosis_errors.InterpretationError
+			envVars, interpretationErr = kurtosis_types.SafeCastToMapStringString(envVarsStarlark, EnvVarsArgName)
+			if interpretationErr != nil {
+				return nil, interpretationErr
+			}
+		}
+	}
+	return &envVars, nil
+}

docs/docs/api-reference/starlark-reference/plan.md

@@ -477,6 +477,13 @@ The `run_sh` instruction executes a one-time execution task. It runs the bash co
         # OPTIONAL (Default: badouralix/curl-jq)
         image = "badouralix/curl-jq",
 
+        # Defines environment variables that should be set inside the Docker container running the task.
+        # OPTIONAL (Default: {})
+        env_vars = {
+            "VAR_1": "VALUE_1",
+            "VAR_2": "VALUE_2",
+        },
+
         # A mapping of path_on_task_where_contents_will_be_mounted -> files_artifact_id_to_mount
         # For more information about file artifacts, see below.
         # CAUTION: duplicate paths to files or directories to be mounted is not supported, and it will fail

internal_testsuites/golang/scripts/test.sh

@@ -12,7 +12,7 @@ lang_root_dirpath="$(dirname "${script_dirpath}")"
 # ==================================================================================================
 PARALLELISM=2
 DOCKER_TIMEOUT="3m"   # This must be Go-parseable timeout
-KUBERNETES_TIMEOUT="8m" # K8S takes longer than docker
+KUBERNETES_TIMEOUT="510s" # K8S takes longer than docker
 
 TESTSUITE_CLUSTER_BACKEND_DOCKER="docker"
 TESTSUITE_CLUSTER_BACKEND_KUBERNETES="kubernetes"

internal_testsuites/golang/testsuite/startosis_run_sh_task_test/run_task_sh_task_test.go

@@ -55,6 +55,12 @@ def run(plan):
   result2 = plan.run_sh(run="cat /temp/tech.txt | tr -d '\n'", files={"/temp": file_artifacts[0]})
   plan.verify(result2.output, "==", "kurtosis")
 `
+
+	runShWithEnvVar = `
+def run(plan):
+  result = plan.run_sh(run="mkdir -p kurtosis && echo $EXAMPLE",image="badouralix/curl-jq",env_vars={"EXAMPLE": "value"})
+  plan.verify(result.output, "==", "value\n")
+`
 )
 
 func TestStarlark_RunshTaskSimple(t *testing.T) {
@@ -101,3 +107,9 @@ func TestStarlark_RunShWithNewLineRemovalPipe(t *testing.T) {
 	_, err := test_helpers.SetupSimpleEnclaveAndRunScript(t, ctx, runshTest, runShWithNewLineRemoval)
 	require.Nil(t, err)
 }
+
+func TestStarlark_RunShWithEnvVars(t *testing.T) {
+	ctx := context.Background()
+	_, err := test_helpers.SetupSimpleEnclaveAndRunScript(t, ctx, runshTest, runShWithEnvVar)
+	require.Nil(t, err)
+}