fix(workflow): target the -main run in Wait and AbortRun (EN-1220)

internal/workflow/manager.go

@@ -117,13 +117,18 @@ func (m *WorkflowManager) RunWorkflow(ctx context.Context, id string, variables
 }
 
 func (m *WorkflowManager) Wait(ctx context.Context, instanceID string) error {
+	// The actual work runs in the detached child workflow "<instanceID>-main";
+	// the Initiate workflow (id == instanceID) completes as soon as that child
+	// has started. Waiting on instanceID would therefore return immediately,
+	// before the run is terminated, so we wait on the running child.
 	if err := m.temporalClient.
-		GetWorkflow(ctx, instanceID, "").
+		GetWorkflow(ctx, instanceID+"-main", "").
 		Get(ctx, nil); err != nil {
-		if errors.Is(err, &serviceerror.NotFound{}) {
+		var notFound *serviceerror.NotFound
+		if errors.As(err, &notFound) {
 			return ErrInstanceNotFound
 		}
-		return errors.Unwrap(err)
+		return err
 	}
 	return nil
 }
@@ -176,7 +181,10 @@ func (m *WorkflowManager) AbortRun(ctx context.Context, instanceID string) error
 		return errors.Wrap(err, "retrieving workflow execution")
 	}
 
-	return m.temporalClient.CancelWorkflow(ctx, instanceID, "")
+	// Cancel the detached child workflow that carries the actual run; the
+	// Initiate workflow (id == instanceID) has already completed, so cancelling
+	// it would be a no-op and never reach the running stages.
+	return m.temporalClient.CancelWorkflow(ctx, instanceID+"-main", "")
 }
 
 func (m *WorkflowManager) ListInstances(ctx context.Context, pagination ListInstancesQuery) (*bunpaginate.Cursor[Instance], error) {

internal/workflow/manager_test.go

@@ -76,3 +76,58 @@ func TestConfig(t *testing.T) {
 		return len(updatedInstance.Statuses) == 1
 	}, 2*time.Second, 100*time.Millisecond)
 }
+
+func TestWait(t *testing.T) {
+	t.Parallel()
+
+	database := srv.NewDatabase(t)
+	db, err := bunconnect.OpenSQLDB(logging.TestingContext(), bunconnect.ConnectionOptions{
+		DatabaseSourceName: database.ConnString(),
+	})
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		_ = db.Close()
+	})
+	require.NoError(t, storage.Migrate(logging.TestingContext(), db))
+
+	taskQueue := uuid.NewString()
+	w := temporalworker.New(logging.Testing(), devServer.Client(), taskQueue,
+		[]temporalworker.DefinitionSet{
+			NewWorkflows("test", false).DefinitionSet(),
+			temporalworker.NewDefinitionSet().Append(temporalworker.Definition{
+				Name: "NoOp",
+				Func: (&stages.NoOp{}).GetWorkflow(),
+			}),
+		},
+		[]temporalworker.DefinitionSet{
+			NewActivities(publish.NoOpPublisher, db).DefinitionSet(),
+		},
+		worker.Options{},
+	)
+	require.NoError(t, w.Start())
+	t.Cleanup(w.Stop)
+
+	manager := NewManager(db, devServer.Client(), "test", taskQueue, false)
+
+	t.Run("waits for the -main run to terminate", func(t *testing.T) {
+		config := Config{Stages: []RawStage{{"noop": map[string]any{}}}}
+		wf, err := manager.Create(logging.TestingContext(), config)
+		require.NoError(t, err)
+		i, err := manager.RunWorkflow(logging.TestingContext(), wf.ID, map[string]string{})
+		require.NoError(t, err)
+
+		// Wait must block on the detached "-main" child, not the Initiate
+		// workflow (which returns immediately). Once it returns, the instance
+		// must already be terminated.
+		require.NoError(t, manager.Wait(logging.TestingContext(), i.ID))
+
+		updated, err := manager.GetInstance(logging.TestingContext(), i.ID)
+		require.NoError(t, err)
+		require.True(t, updated.Terminated)
+	})
+
+	t.Run("unknown instance returns ErrInstanceNotFound", func(t *testing.T) {
+		err := manager.Wait(logging.TestingContext(), "does-not-exist")
+		require.ErrorIs(t, err, ErrInstanceNotFound)
+	})
+}