Fix context propagation in execute operations

cmd/common/compile.go

@@ -1,6 +1,7 @@
 package common
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"os"
@@ -33,15 +34,15 @@ type WorkflowCompileOptions struct {
 }
 
 // getBuildCmd returns a single step that builds the workflow and returns the WASM bytes.
-func getBuildCmd(workflowRootFolder, mainFile, language string, opts WorkflowCompileOptions) (func() ([]byte, error), error) {
+func getBuildCmd(ctx context.Context, workflowRootFolder, mainFile, language string, opts WorkflowCompileOptions) (func() ([]byte, error), error) {
 	tmpPath := filepath.Join(workflowRootFolder, ".cre_build_tmp.wasm")
 	switch language {
 	case constants.WorkflowLanguageTypeScript:
 		args := []string{"cre-compile", mainFile, tmpPath}
 		if opts.SkipTypeChecks {
 			args = append(args, SkipTypeChecksFlag)
 		}
-		cmd := exec.Command("bun", args...)
+		cmd := exec.CommandContext(ctx, "bun", args...)
 		cmd.Dir = workflowRootFolder
 		return func() ([]byte, error) {
 			out, err := cmd.CombinedOutput()
@@ -67,7 +68,7 @@ func getBuildCmd(workflowRootFolder, mainFile, language string, opts WorkflowCom
 		if opts.StripSymbols {
 			ldflags = "-buildid= -w -s"
 		}
-		cmd := exec.Command(
+		cmd := exec.CommandContext(ctx,
 			"go", "build",
 			"-o", tmpPath,
 			"-trimpath",
@@ -92,7 +93,7 @@ func getBuildCmd(workflowRootFolder, mainFile, language string, opts WorkflowCom
 		if err != nil {
 			return nil, err
 		}
-		makeCmd := exec.Command("make", "build")
+		makeCmd := exec.CommandContext(ctx, "make", "build")
 		makeCmd.Dir = makeRoot
 		builtPath := filepath.Join(makeRoot, defaultWasmOutput)
 		return func() ([]byte, error) {
@@ -108,7 +109,7 @@ func getBuildCmd(workflowRootFolder, mainFile, language string, opts WorkflowCom
 		if opts.StripSymbols {
 			ldflags = "-buildid= -w -s"
 		}
-		cmd := exec.Command(
+		cmd := exec.CommandContext(ctx,
 			"go", "build",
 			"-o", tmpPath,
 			"-trimpath",
@@ -135,7 +136,7 @@ func getBuildCmd(workflowRootFolder, mainFile, language string, opts WorkflowCom
 // opts.StripSymbols: for Go builds, true strips debug symbols (deploy); false keeps them (simulate).
 // opts.SkipTypeChecks: for TypeScript, passes SkipTypeChecksFlag to cre-compile.
 // For custom Makefile WASM builds, StripSymbols and SkipTypeChecks have no effect.
-func CompileWorkflowToWasm(workflowPath string, opts WorkflowCompileOptions) ([]byte, error) {
+func CompileWorkflowToWasm(ctx context.Context, workflowPath string, opts WorkflowCompileOptions) ([]byte, error) {
 	workflowRootFolder, workflowMainFile, err := WorkflowPathRootAndMain(workflowPath)
 	if err != nil {
 		return nil, fmt.Errorf("workflow path: %w", err)
@@ -167,7 +168,7 @@ func CompileWorkflowToWasm(workflowPath string, opts WorkflowCompileOptions) ([]
 		return nil, fmt.Errorf("unsupported workflow language for file %s", workflowMainFile)
 	}
 
-	buildStep, err := getBuildCmd(workflowRootFolder, workflowMainFile, language, opts)
+	buildStep, err := getBuildCmd(ctx, workflowRootFolder, workflowMainFile, language, opts)
 	if err != nil {
 		return nil, err
 	}

cmd/common/compile_test.go

@@ -2,6 +2,7 @@ package common
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"os"
 	"os/exec"
@@ -47,29 +48,29 @@ func TestFindMakefileRoot(t *testing.T) {
 func TestCompileWorkflowToWasm_Go_Success(t *testing.T) {
 	t.Run("basic_workflow", func(t *testing.T) {
 		path := deployTestdataPath("basic_workflow", "main.go")
-		wasm, err := CompileWorkflowToWasm(path, WorkflowCompileOptions{StripSymbols: true})
+		wasm, err := CompileWorkflowToWasm(context.Background(), path, WorkflowCompileOptions{StripSymbols: true})
 		require.NoError(t, err)
 		assert.NotEmpty(t, wasm)
 	})
 
 	t.Run("configless_workflow", func(t *testing.T) {
 		path := deployTestdataPath("configless_workflow", "main.go")
-		wasm, err := CompileWorkflowToWasm(path, WorkflowCompileOptions{StripSymbols: true})
+		wasm, err := CompileWorkflowToWasm(context.Background(), path, WorkflowCompileOptions{StripSymbols: true})
 		require.NoError(t, err)
 		assert.NotEmpty(t, wasm)
 	})
 
 	t.Run("missing_go_mod", func(t *testing.T) {
 		path := deployTestdataPath("missing_go_mod", "main.go")
-		wasm, err := CompileWorkflowToWasm(path, WorkflowCompileOptions{StripSymbols: true})
+		wasm, err := CompileWorkflowToWasm(context.Background(), path, WorkflowCompileOptions{StripSymbols: true})
 		require.NoError(t, err)
 		assert.NotEmpty(t, wasm)
 	})
 }
 
 func TestCompileWorkflowToWasm_Go_Malformed_Fails(t *testing.T) {
 	path := deployTestdataPath("malformed_workflow", "main.go")
-	_, err := CompileWorkflowToWasm(path, WorkflowCompileOptions{StripSymbols: true})
+	_, err := CompileWorkflowToWasm(context.Background(), path, WorkflowCompileOptions{StripSymbols: true})
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "failed to compile workflow")
 	assert.Contains(t, err.Error(), "undefined: sdk.RemovedFunctionThatFailsCompilation")
@@ -80,7 +81,7 @@ func TestCompileWorkflowToWasm_Wasm_Success(t *testing.T) {
 	_ = os.Remove(wasmPath)
 	t.Cleanup(func() { _ = os.Remove(wasmPath) })
 
-	wasm, err := CompileWorkflowToWasm(wasmPath, WorkflowCompileOptions{StripSymbols: true})
+	wasm, err := CompileWorkflowToWasm(context.Background(), wasmPath, WorkflowCompileOptions{StripSymbols: true})
 	require.NoError(t, err)
 	assert.NotEmpty(t, wasm)
 
@@ -96,14 +97,14 @@ func TestCompileWorkflowToWasm_Wasm_Fails(t *testing.T) {
 		wasmPath := filepath.Join(wasmDir, "workflow.wasm")
 		require.NoError(t, os.WriteFile(wasmPath, []byte("not really wasm"), 0600))
 
-		_, err := CompileWorkflowToWasm(wasmPath, WorkflowCompileOptions{StripSymbols: true})
+		_, err := CompileWorkflowToWasm(context.Background(), wasmPath, WorkflowCompileOptions{StripSymbols: true})
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "no Makefile found")
 	})
 
 	t.Run("make_build_fails", func(t *testing.T) {
 		path := deployTestdataPath("wasm_make_fails", "wasm", "workflow.wasm")
-		_, err := CompileWorkflowToWasm(path, WorkflowCompileOptions{StripSymbols: true})
+		_, err := CompileWorkflowToWasm(context.Background(), path, WorkflowCompileOptions{StripSymbols: true})
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "failed to compile workflow")
 		assert.Contains(t, err.Error(), "build output:")
@@ -138,7 +139,7 @@ func TestCompileWorkflowToWasm_TS_Success(t *testing.T) {
   "include": ["main.ts"]
 }
 `), 0600))
-	wasm, err := CompileWorkflowToWasm(mainPath, WorkflowCompileOptions{StripSymbols: true})
+	wasm, err := CompileWorkflowToWasm(context.Background(), mainPath, WorkflowCompileOptions{StripSymbols: true})
 	if err != nil {
 		t.Skipf("TS compile failed (published cre-sdk may lack full layout): %v", err)
 	}

cmd/common/fetch.go

@@ -1,6 +1,7 @@
 package common
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -28,8 +29,13 @@ func IsURL(s string) bool {
 }
 
 // FetchURL performs an HTTP GET and returns the response body bytes.
-func FetchURL(url string) ([]byte, error) {
-	resp, err := http.Get(url) //nolint:gosec,noctx
+func FetchURL(ctx context.Context, url string) ([]byte, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return nil, fmt.Errorf("HTTP GET %s: %w", url, err)
+	}
+
+	resp, err := http.DefaultClient.Do(req) //nolint:gosec
 	if err != nil {
 		return nil, fmt.Errorf("HTTP GET %s: %w", url, err)
 	}

cmd/common/fetch_test.go

@@ -1,6 +1,7 @@
 package common
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -42,7 +43,7 @@ func TestFetchURL(t *testing.T) {
 		}))
 		defer srv.Close()
 
-		data, err := FetchURL(srv.URL)
+		data, err := FetchURL(context.Background(), srv.URL)
 		require.NoError(t, err)
 		assert.Equal(t, body, data)
 	})
@@ -53,13 +54,13 @@ func TestFetchURL(t *testing.T) {
 		}))
 		defer srv.Close()
 
-		_, err := FetchURL(srv.URL)
+		_, err := FetchURL(context.Background(), srv.URL)
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "returned status 404")
 	})
 
 	t.Run("unreachable host", func(t *testing.T) {
-		_, err := FetchURL("http://127.0.0.1:1")
+		_, err := FetchURL(context.Background(), "http://127.0.0.1:1")
 		require.Error(t, err)
 	})
 }

cmd/workflow/build/build.go

@@ -1,6 +1,7 @@
 package build
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -26,15 +27,15 @@ func New(runtimeContext *runtime.Context) *cobra.Command {
 		RunE: func(cmd *cobra.Command, args []string) error {
 			outputPath, _ := cmd.Flags().GetString("output")
 			skipTypeChecks, _ := cmd.Flags().GetBool(cmdcommon.SkipTypeChecksCLIFlag)
-			return execute(args[0], outputPath, skipTypeChecks)
+			return execute(cmd.Context(), args[0], outputPath, skipTypeChecks)
 		},
 	}
 	buildCmd.Flags().StringP("output", "o", "", "Output file path for the compiled WASM binary (default: <workflow-folder>/binary.wasm)")
 	buildCmd.Flags().Bool(cmdcommon.SkipTypeChecksCLIFlag, false, "Skip TypeScript project typecheck during compilation (passes "+cmdcommon.SkipTypeChecksFlag+" to cre-compile)")
 	return buildCmd
 }
 
-func execute(workflowFolder, outputPath string, skipTypeChecks bool) error {
+func execute(ctx context.Context, workflowFolder, outputPath string, skipTypeChecks bool) error {
 	workflowDir, err := filepath.Abs(workflowFolder)
 	if err != nil {
 		return fmt.Errorf("resolve workflow folder: %w", err)
@@ -60,7 +61,7 @@ func execute(workflowFolder, outputPath string, skipTypeChecks bool) error {
 	outputPath = cmdcommon.EnsureWasmExtension(outputPath)
 
 	ui.Dim("Compiling workflow...")
-	wasmBytes, err := cmdcommon.CompileWorkflowToWasm(resolvedPath, cmdcommon.WorkflowCompileOptions{
+	wasmBytes, err := cmdcommon.CompileWorkflowToWasm(ctx, resolvedPath, cmdcommon.WorkflowCompileOptions{
 		StripSymbols:   true,
 		SkipTypeChecks: skipTypeChecks,
 	})

cmd/workflow/deploy/artifacts.go

@@ -9,6 +9,10 @@ import (
 )
 
 func (h *handler) uploadArtifacts() error {
+	if err := h.executionContext().Err(); err != nil {
+		return err
+	}
+
 	if h.workflowArtifact == nil {
 		return fmt.Errorf("workflowArtifact is nil")
 	}
@@ -46,7 +50,7 @@ func (h *handler) uploadArtifacts() error {
 	if !binaryFromURL {
 		ui.Success(fmt.Sprintf("Loaded binary from: %s", h.inputs.OutputPath))
 		binaryResp, err := storageClient.UploadArtifactWithRetriesAndGetURL(
-			workflowID, storageclient.ArtifactTypeBinary, binaryData, "application/octet-stream")
+			h.executionContext(), workflowID, storageclient.ArtifactTypeBinary, binaryData, "application/octet-stream")
 		if err != nil {
 			return fmt.Errorf("uploading binary artifact: %w", err)
 		}
@@ -59,7 +63,7 @@ func (h *handler) uploadArtifacts() error {
 		ui.Success(fmt.Sprintf("Loaded config from: %s", h.inputs.ConfigPath))
 		var err error
 		configURL, err = storageClient.UploadArtifactWithRetriesAndGetURL(
-			workflowID, storageclient.ArtifactTypeConfig, configData, "text/plain")
+			h.executionContext(), workflowID, storageclient.ArtifactTypeConfig, configData, "text/plain")
 		if err != nil {
 			return fmt.Errorf("uploading config artifact: %w", err)
 		}

cmd/workflow/deploy/artifacts_test.go

@@ -71,7 +71,7 @@ func TestUpload_SuccessAndErrorCases(t *testing.T) {
 
 	simulatedEnvironment := chainsim.NewSimulatedEnvironment(t)
 	ctx, buf := simulatedEnvironment.NewRuntimeContextWithBufferedOutput()
-	h := newHandler(ctx, buf)
+	h := newTestHandler(ctx, buf)
 	h.inputs.WorkflowOwner = chainsim.TestAddress
 	h.inputs.WorkflowName = "test_workflow"
 	h.inputs.DonFamily = "test_label"
@@ -147,7 +147,7 @@ func TestUploadArtifactToStorageService_OriginError(t *testing.T) {
 
 	simulatedEnvironment := chainsim.NewSimulatedEnvironment(t)
 	runtimeContext, buf := simulatedEnvironment.NewRuntimeContextWithBufferedOutput()
-	h := newHandler(runtimeContext, buf)
+	h := newTestHandler(runtimeContext, buf)
 	h.inputs.WorkflowOwner = chainsim.TestAddress
 	h.inputs.WorkflowName = "test_workflow"
 	h.inputs.DonFamily = "test_label"
@@ -187,7 +187,7 @@ func TestUploadArtifactToStorageService_AlreadyExistsError(t *testing.T) {
 
 	simulatedEnvironment := chainsim.NewSimulatedEnvironment(t)
 	runtimeContext, buf := simulatedEnvironment.NewRuntimeContextWithBufferedOutput()
-	h := newHandler(runtimeContext, buf)
+	h := newTestHandler(runtimeContext, buf)
 	h.inputs.WorkflowOwner = chainsim.TestAddress
 	h.inputs.WorkflowName = "test_workflow"
 	h.inputs.DonFamily = "test_label"
@@ -255,7 +255,7 @@ func TestUpload_UsesResolvedWorkflowOwnerForPresignedUrls(t *testing.T) {
 	simulatedEnvironment := chainsim.NewSimulatedEnvironment(t)
 	t.Cleanup(simulatedEnvironment.Close)
 	ctx, buf := simulatedEnvironment.NewRuntimeContextWithBufferedOutput()
-	h := newHandler(ctx, buf)
+	h := newTestHandler(ctx, buf)
 	h.inputs.WorkflowOwner = "0x2222222222222222222222222222222222222222"
 	h.inputs.WorkflowName = "test_workflow"
 	h.inputs.DonFamily = "test_label"

cmd/workflow/deploy/auto_link.go

@@ -1,7 +1,6 @@
 package deploy
 
 import (
-	"context"
 	"fmt"
 	"strings"
 	"time"
@@ -25,7 +24,7 @@ const (
 func (h *handler) ensureOwnerLinkedOrFail(onChain *settings.OnChainRegistry) error {
 	ownerAddr := common.HexToAddress(h.inputs.WorkflowOwner)
 
-	linked, err := h.wrc.IsOwnerLinked(h.execCtx, ownerAddr)
+	linked, err := h.wrc.IsOwnerLinked(h.executionContext(), ownerAddr)
 	if err != nil {
 		return fmt.Errorf("failed to check owner link status: %w", err)
 	}
@@ -66,7 +65,7 @@ func (h *handler) ensureOwnerLinkedOrFail(onChain *settings.OnChainRegistry) err
 func (h *handler) autoLinkMSIGAndExit(onChain *settings.OnChainRegistry) (halt bool, err error) {
 	ownerAddr := common.HexToAddress(h.inputs.WorkflowOwner)
 
-	linked, err := h.wrc.IsOwnerLinked(h.execCtx, ownerAddr)
+	linked, err := h.wrc.IsOwnerLinked(h.executionContext(), ownerAddr)
 	if err != nil {
 		return false, fmt.Errorf("failed to check owner link status: %w", err)
 	}
@@ -107,7 +106,7 @@ func (h *handler) tryAutoLink(onChain *settings.OnChainRegistry) error {
 		EnvironmentSet: h.environmentSet,
 	}
 
-	return linkkey.Exec(h.execCtx, rtx, linkkey.Inputs{
+	return linkkey.Exec(h.executionContext(), rtx, linkkey.Inputs{
 		WorkflowOwner:                   h.inputs.WorkflowOwner,
 		WorkflowRegistryContractAddress: onChain.Address(),
 		WorkflowOwnerLabel:              h.inputs.OwnerLabel,
@@ -137,7 +136,7 @@ func (h *handler) checkLinkStatusViaGraphQL(ownerAddr common.Address) (bool, err
 	}
 
 	gql := graphqlclient.New(h.credentials, h.environmentSet, h.log)
-	if err := gql.Execute(context.Background(), req, &resp); err != nil {
+	if err := gql.Execute(h.executionContext(), req, &resp); err != nil {
 		return false, fmt.Errorf("GraphQL query failed: %w", err)
 	}
 
@@ -181,7 +180,12 @@ func (h *handler) waitForBackendLinkProcessing(ownerAddr common.Address) error {
 	ui.Line()
 
 	// Wait for 3 block confirmations before polling
-	time.Sleep(initialBlockWait)
+	ctx := h.executionContext()
+	select {
+	case <-time.After(initialBlockWait):
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 
 	err := retry.Do(
 		func() error {
@@ -199,6 +203,7 @@ func (h *handler) waitForBackendLinkProcessing(ownerAddr common.Address) error {
 		retry.Delay(retryDelay),
 		retry.DelayType(retry.FixedDelay), // Use fixed 3s delay between retries
 		retry.LastErrorOnly(true),
+		retry.Context(ctx),
 		retry.OnRetry(func(n uint, err error) {
 			h.log.Debug().Uint("attempt", n+1).Uint("maxAttempts", maxAttempts).Err(err).Msg("Retrying link status check")
 			ui.Dim(fmt.Sprintf("  Waiting for verification... (attempt %d/%d)", n+1, maxAttempts))

cmd/workflow/deploy/auto_link_test.go

@@ -158,7 +158,7 @@ func TestCheckLinkStatusViaGraphQL(t *testing.T) {
 				AuthType:    credentials.AuthTypeApiKey,
 				IsValidated: true,
 			}
-			h := newHandler(ctx, nil)
+			h := newTestHandler(ctx, nil)
 			h.inputs.WorkflowOwner = tt.ownerAddress
 			h.environmentSet.GraphQLURL = server.URL + "/graphql"
 
@@ -329,7 +329,7 @@ func TestWaitForBackendLinkProcessing(t *testing.T) {
 				AuthType:    credentials.AuthTypeApiKey,
 				IsValidated: true,
 			}
-			h := newHandler(ctx, nil)
+			h := newTestHandler(ctx, nil)
 			h.inputs.WorkflowOwner = tt.ownerAddress
 			h.environmentSet.GraphQLURL = server.URL + "/graphql"
 

cmd/workflow/deploy/compile.go

@@ -67,7 +67,7 @@ func (h *handler) Compile() error {
 		h.runtimeContext.Workflow.Language = cmdcommon.GetWorkflowLanguage(workflowMainFile)
 	}
 
-	wasmFile, err = cmdcommon.CompileWorkflowToWasm(resolvedWorkflowPath, cmdcommon.WorkflowCompileOptions{
+	wasmFile, err = cmdcommon.CompileWorkflowToWasm(h.executionContext(), resolvedWorkflowPath, cmdcommon.WorkflowCompileOptions{
 		StripSymbols:   true,
 		SkipTypeChecks: h.inputs.SkipTypeChecks,
 	})