@pikku/inspector 0.12.14 → 0.12.17

package/src/utils/workflow/dsl/extract-dsl-workflow.ts

@@ -6,7 +6,6 @@ import type {
   BranchStepMeta,
   ParallelGroupStepMeta,
   FanoutStepMeta,
-  ReturnStepMeta,
   CancelStepMeta,
   SetStepMeta,
   SwitchStepMeta,
@@ -436,21 +435,34 @@ function extractExpressionStatement(
       outputVar = expr.left.text
 
       // Check if this is an assignment to a context variable (set step)
+      // But if the RHS is a workflow.do() call, fall through to RPC extraction —
+      // reassigning a pre-declared variable with a workflow step is valid and common.
       if (context.contextVars.has(outputVar)) {
-        const literalValue = extractLiteralValue(expr.right)
-        if (literalValue !== undefined) {
+        const rhs = expr.right
+        const rhsCall =
+          ts.isAwaitExpression(rhs) && ts.isCallExpression(rhs.expression)
+            ? rhs.expression
+            : null
+        const isWorkflowCall = rhsCall
+          ? isWorkflowDoCall(rhsCall, context.checker)
+          : false
+
+        if (!isWorkflowCall) {
+          const literalValue = extractLiteralValue(expr.right)
+          if (literalValue !== undefined) {
+            return {
+              type: 'set',
+              variable: outputVar,
+              value: literalValue,
+            } as SetStepMeta
+          }
+          // Non-literal assignment to context var - use expression as string
           return {
             type: 'set',
             variable: outputVar,
-            value: literalValue,
+            value: getSourceText(expr.right),
           } as SetStepMeta
         }
-        // Non-literal assignment to context var - use expression as string
-        return {
-          type: 'set',
-          variable: outputVar,
-          value: getSourceText(expr.right),
-        } as SetStepMeta
       }
     }
     // Use right side as the expression to extract from
@@ -1327,11 +1339,38 @@ function extractOutputBinding(
 function extractReturn(
   statement: ts.ReturnStatement,
   context: ExtractionContext
-): ReturnStepMeta | null {
+): WorkflowStepMeta | null {
   if (!statement.expression) {
     return null
   }
 
+  if (
+    ts.isAwaitExpression(statement.expression) &&
+    ts.isCallExpression(statement.expression.expression)
+  ) {
+    const call = statement.expression.expression
+    if (isWorkflowDoCall(call, context.checker)) {
+      return isInlineDoCall(call)
+        ? extractInlineStep(call, context)
+        : extractRpcStep(call, context)
+    }
+    if (isWorkflowSleepCall(call, context.checker)) {
+      return extractSleepStep(call, context)
+    }
+  }
+
+  if (ts.isCallExpression(statement.expression)) {
+    const call = statement.expression
+    if (isWorkflowDoCall(call, context.checker)) {
+      return isInlineDoCall(call)
+        ? extractInlineStep(call, context)
+        : extractRpcStep(call, context)
+    }
+    if (isWorkflowSleepCall(call, context.checker)) {
+      return extractSleepStep(call, context)
+    }
+  }
+
   if (!ts.isObjectLiteralExpression(statement.expression)) {
     return null
   }