From ae0860a775353f1806da2dba63850d556e8ffea8 Mon Sep 17 00:00:00 2001
From: Adam Chalmers <adam.chalmers@zoo.dev>
Date: Tue, 14 Jan 2025 11:36:15 -0600
Subject: [PATCH] Fixed another queryAst test

There were 2 problems:
 - Test was looking for the old style of `line` call to choose an offset
   for pathToNode
 - Test assumed that the `tag` param was always the third one, but in
   a kwarg call, you have to look it up by label
---
 src/lang/queryAst.test.ts |  2 +-
 src/lang/queryAst.ts      | 26 ++++++++++++++------------
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/src/lang/queryAst.test.ts b/src/lang/queryAst.test.ts
index cb3404d1b..866bb9faa 100644
--- a/src/lang/queryAst.test.ts
+++ b/src/lang/queryAst.test.ts
@@ -443,7 +443,7 @@ describe('Testing findUsesOfTagInPipe', () => {
   it('finds the current segment', async () => {
     const ast = assertParse(exampleCode)
 
-    const lineOfInterest = `198.85], %, $seg01`
+    const lineOfInterest = `198.85], tag = $seg01`
     const characterIndex =
       exampleCode.indexOf(lineOfInterest) + lineOfInterest.length
     const pathToNode = getNodePathFromSourceRange(
diff --git a/src/lang/queryAst.ts b/src/lang/queryAst.ts
index ac8835096..f52ac1185 100644
--- a/src/lang/queryAst.ts
+++ b/src/lang/queryAst.ts
@@ -5,6 +5,7 @@ import {
   ArtifactGraph,
   BinaryExpression,
   CallExpression,
+  CallExpressionKw,
   Expr,
   ExpressionStatement,
   ObjectExpression,
@@ -25,7 +26,7 @@ import {
 import { createIdentifier, splitPathAtLastIndex } from './modifyAst'
 import { getSketchSegmentFromSourceRange } from './std/sketchConstraints'
 import { getAngle } from '../lib/utils'
-import { getFirstArg } from './std/sketch'
+import { ARG_TAG, getFirstArg } from './std/sketch'
 import {
   getConstraintLevelFromSourceRange,
   getConstraintType,
@@ -33,7 +34,8 @@ import {
 import { err, Reason } from 'lib/trap'
 import { ImportStatement } from 'wasm-lib/kcl/bindings/ImportStatement'
 import { Node } from 'wasm-lib/kcl/bindings/Node'
-import { codeRefFromRange } from './std/artifactGraph'
+import { ArtifactGraph, codeRefFromRange } from './std/artifactGraph'
+import { findKwArg } from './util'
 
 /**
  * Retrieves a node from a given path within a Program node structure, optionally stopping at a specified node type.
@@ -885,27 +887,27 @@ export function findUsesOfTagInPipe(
     'segEndY',
     'segLen',
   ]
-  const nodeMeta = getNodeFromPath<CallExpression>(
+  const nodeMeta = getNodeFromPath<CallExpression | CallExpressionKw>(
     ast,
     pathToNode,
-    'CallExpression'
+    ['CallExpression', 'CallExpressionKw']
   )
   if (err(nodeMeta)) {
     console.error(nodeMeta)
     return []
   }
   const node = nodeMeta.node
-  if (node.type !== 'CallExpression') return []
+  if (node.type !== 'CallExpressionKw' && node.type !== 'CallExpression')
+    return []
   const tagIndex = node.callee.name === 'close' ? 1 : 2
-  const thirdParam = node.arguments[tagIndex]
-  if (
-    !(thirdParam?.type === 'TagDeclarator' || thirdParam?.type === 'Identifier')
-  )
+  const tagParam =
+    node.type === 'CallExpression'
+      ? node.arguments[tagIndex]
+      : findKwArg(ARG_TAG, node)
+  if (!(tagParam?.type === 'TagDeclarator' || tagParam?.type === 'Identifier'))
     return []
   const tag =
-    thirdParam?.type === 'TagDeclarator'
-      ? String(thirdParam.value)
-      : thirdParam.name
+    tagParam?.type === 'TagDeclarator' ? String(tagParam.value) : tagParam.name
 
   const varDec = getNodeFromPath<Node<VariableDeclaration>>(
     ast,