From 080536eb18f79d76964682a76b559ad37100e0e1 Mon Sep 17 00:00:00 2001
From: Shahed Nasser <shahednasser@gmail.com>
Date: Mon, 12 Aug 2024 11:45:02 +0300
Subject: [PATCH] docs-util: associate an API route with the workflow it uses

---
 .../docs-generator/src/classes/kinds/oas.ts   | 80 +++++++++++++++++++
 .../docs-generator/src/types/index.d.ts       |  1 +
 2 files changed, 81 insertions(+)

diff --git a/www/utils/packages/docs-generator/src/classes/kinds/oas.ts b/www/utils/packages/docs-generator/src/classes/kinds/oas.ts
index 529fd7c3fdab8..060e92d7141c3 100644
--- a/www/utils/packages/docs-generator/src/classes/kinds/oas.ts
+++ b/www/utils/packages/docs-generator/src/classes/kinds/oas.ts
@@ -346,6 +346,9 @@ class OasKindGenerator extends FunctionKindGenerator {
       areaTags?.add(tagName)
     }
 
+    // get associated workflow
+    oas["x-workflow"] = this.getAssociatedWorkflow(node)
+
     return formatOas(oas, oasPrefix)
   }
 
@@ -610,6 +613,9 @@ class OasKindGenerator extends FunctionKindGenerator {
       areaTags?.add(tagName)
     }
 
+    // get associated workflow
+    oas["x-workflow"] = this.getAssociatedWorkflow(node)
+
     return formatOas(oas, oasPrefix)
   }
 
@@ -1808,6 +1814,80 @@ class OasKindGenerator extends FunctionKindGenerator {
     return oldSchemaObj
   }
 
+  /**
+   * This method retrieves the workflow used in an API route.
+   *
+   * @param node - The node to search through
+   * @returns The workflow's name.
+   */
+  getAssociatedWorkflow(node: FunctionNode): string | undefined {
+    let workflow: string | undefined
+
+    if (
+      (!ts.isFunctionDeclaration(node) && !ts.isArrowFunction(node)) ||
+      !node.body ||
+      !ts.isBlock(node.body)
+    ) {
+      return
+    }
+
+    const sourceFile = node.getSourceFile()
+    const fileLocalSymbols: Map<string, ts.Symbol> =
+      "locals" in sourceFile
+        ? (sourceFile.locals as Map<string, ts.Symbol>)
+        : new Map()
+
+    if (!fileLocalSymbols.size) {
+      return
+    }
+
+    node.body.statements.some((statement) => {
+      let awaitExpression: ts.AwaitExpression | undefined
+      if (
+        ts.isVariableStatement(statement) &&
+        statement.declarationList.declarations[0].initializer &&
+        ts.isAwaitExpression(
+          statement.declarationList.declarations[0].initializer
+        )
+      ) {
+        awaitExpression = statement.declarationList.declarations[0].initializer
+      } else if (ts.isAwaitExpression(statement)) {
+        awaitExpression = statement
+      }
+
+      if (
+        !awaitExpression ||
+        !ts.isCallExpression(awaitExpression.expression) ||
+        !ts.isPropertyAccessExpression(awaitExpression.expression.expression) ||
+        !awaitExpression.expression.expression.name.getText().startsWith("run")
+      ) {
+        return false
+      }
+
+      const fullExpressionText = awaitExpression.expression.getText()
+      const parenIndex = fullExpressionText.indexOf("(")
+      const dotIndex = fullExpressionText.indexOf(".")
+      const cutOffIndex =
+        parenIndex !== -1 && dotIndex === -1
+          ? parenIndex
+          : parenIndex === -1 && dotIndex !== -1
+            ? dotIndex
+            : parenIndex === -1 && dotIndex === -1
+              ? fullExpressionText.length
+              : Math.min(parenIndex, dotIndex)
+      const expressionName = fullExpressionText.substring(0, cutOffIndex)
+
+      if (!fileLocalSymbols.has(expressionName)) {
+        return false
+      }
+
+      workflow = expressionName
+      return true
+    })
+
+    return workflow
+  }
+
   /**
    * Retrieve the file name that's used to write the OAS operation of a node.
    *
diff --git a/www/utils/packages/docs-generator/src/types/index.d.ts b/www/utils/packages/docs-generator/src/types/index.d.ts
index 7c2e047a8d146..c827dc02c23af 100644
--- a/www/utils/packages/docs-generator/src/types/index.d.ts
+++ b/www/utils/packages/docs-generator/src/types/index.d.ts
@@ -9,6 +9,7 @@ declare type CodeSample = {
 export declare type OpenApiOperation = Partial<OpenAPIV3.OperationObject> & {
   "x-authenticated"?: boolean
   "x-codeSamples"?: CodeSample[]
+  "x-workflow"?: string
 }
 
 export declare type CommonCliOptions = {