STAFFA SOFTWARE SOLUTIONS
[DE]
Eine Fotomontage aus einer Baumkrone und einer Schraubzwinge

How to save time, money, and tokens with ast-grep

Sebastian Staffa

LLMs are great tools for automating the annoying tasks of software development, like refactoring large projects. For a lot of things there were already sensible tools in the past, but some things still had to be done by hand. One example: Renaming a function and all of its calls has been a core feature of many IDEs for decades. But tasks that required context about the location to be refactored were mostly manual work before LLMs.

For example, I recently moved information from a passed-through context object into an AsyncLocalStorage in a TypeScript application I maintain. In the course of this, all HTTP handlers had to be refactored so that their entire content was moved into a withContext(() => ...) callback. Hundreds of locations across the entire application were affected. To make these changes, I sent out an agent with the corresponding task.

The problem: Such a task is quite token- and time-intensive. The agent in question failed several times to find all the locations, and had to be sent out multiple times, even though tests were available which, when executed, output all affected locations. I wasn’t thrilled.

The necessary changes weren’t actually that complicated at first glance. The affected files were known through the tests. The necessary refactoring looked like this:

// Ausgangssituation
class HandlerName implements Handler {
  async handle(...) {
     // content to be wrapped
  }
}

// Ziel
class HandlerName implements Handler {
  async handle(...) {
    return await this.contextService.withContext(async () => {
       // content to be wrapped
    })
  }
}

There has to be something better, I thought to myself, and talked to my long-time friend Christian Rades about the topic, who promptly delivered the brilliant idea: “Hmmm, you’d need something like a CSS selector over the AST”.

A quick Perplexity search later, it turns out – exactly such a thing already exists: ast-grep.

ast-grep provides a more sophisticated way to find your code: Rules are like CSS selectors that can compose together to filter AST nodes based on certain criteria.

— 

“Rule Essentials”, ast-grep Documentation

ast-grep already supports dozens of programming languages and yes - of course there is already a skill for agents, but more on that later. Because before handing everything over to the LLM, I wanted to understand what the tool is capable of first.

ast-grep

To find out, I wanted to reimplement the solution for some refactoring problems from the recent past, where I was annoyed by the way LLMs work when implementing them. One of these problems was converting a function header with many, partly optional parameters into a function that only accepts an options object in order to make handling the optional parameters simple:

// Ausgangssituation
async myMethodWithOptionalParams(
    myParam1: string,
    myParam2: string,
    myOptionalParam1?: string,
    myOptionalParam2?: string
): Promise<void> {
 // ...
}

// Ziel
async myMethodWithOptionalParams(
    opts: {
        myParam1: string,
        myParam2: string,
        myOptionalParam1?: string,
        myOptionalParam2?: string
    }
): Promise<void> {
 // ...
}

ast-grep already supports rewriting code using so-called rewrite rules - however, before I could get started with this, I first had to define a matcher that finds the function calls I wanted to rewrite.

My naive first solution worked purely with the pattern construct of ast-grep:

ast-grep --pattern 'myMethodWithOptionalParams($$$)' ./src

And with this pattern I couldn’t find anything at all. Why? myMethodWithOptionalParams is not a standalone function, but a member of a service. The calls follow the form:

await this.myAwesomeService.myMethodWithOptionalParams( ... );

But ast-grep patterns always work on exactly one AST node. A look at the very useful ast-grep playground shows: myMethodWithOptionalParams in the AST is a property access (AST node property_identifier) on the service. The parameters are housed in another node: A call_expression AST node. The call_expression includes the variable name of the service, the method name, as well as the parameters. At the same time, the name of the variable in which the service is stored is not predefined in my codebase and can be different depending on the caller. The pattern therefore must be adapted so that the entire call_expression is matched, regardless of the local service name. The solution: Include the variable name as a metavariable in the pattern:

ast-grep --pattern '$SERVICE.myMethodWithOptionalParams($P1,$P2,$P3,$P4)' ./src

This pattern led to the desired results and was “good enough” for my use case. If my codebase contained a second service with a function of the same name and four parameters, this pattern would be too naive and would have to be refined further. However, this allowed me to start with the rewrite rule. For the sake of clarity, I moved the rule into a reusable YAML file:

id: transform-to-opts
language: TypeScript
rule:
  pattern: $SERVICE.myMethodWithOptionalParams($P1,$P2,$P3,$P4)
fix:
  "$SERVICE.myMethodWithOptionalParams({ myParam1: $P1,$ myParam2: $P2,customerId: $P3,recipientId: $P4})"

In my newly added fix rule, I sort the four parameters of the method as needed; I keep the variable name of the service. I don’t care about the formatting of the output, the already existing linter tooling will take care of that later.

More Complex Refactoring Problems

Even though I was very satisfied with this solution, it seemed to me like it was just scratching the surface of the possibilities ast-grep offers. To dive deeper, I tried another refactoring problem from the past few weeks:

In the unit tests of some HTTP request handlers, stub objects were used in which not all necessary fields of an HTTP request, such as headers or the route, were set, which is why these objects were explicitly cast to any directly during the call to appease the compiler:

it("works", () => {
  // ...
  const result = myHandlerFunction({ requestBody: { attrib: 42 } } as any);
  // ...
});

To make the mocking cleaner, I had introduced a buildStubRequest function that fills the missing attributes of such a stub request with sensible default values. In the future, the call of a handler in a test should accordingly look like this:

it("works", () => {
  // ...
  const result = myHandlerFunction(
    buildStubRequest({ requestBody: { attrib: 42 } }),
  );
  // ...
});

The task in the refactoring was to keep the existing stub object, pass it as a parameter in buildStubRequest, and remove the any cast. The challenge here: Not just match any object construction. The following solution was sufficient for my codebase as long as it was only applied to the tests:

id: transform-mock-obj
language: TypeScript
rule:
  pattern: "$OBJ as $TYPE"
  kind: as_expression
  has:
    kind: object
    precedes:
      stopBy: end
      kind: predefined_type
      regex: "any"
    has:
      stopBy: end
      kind: property_identifier
      regex: "requestBody"
fix: "buildStubRequest($OBJ)"

The above rule matches all casts (all as_expression AST nodes), but only those that define an object. This object definition must stand directly before a hardcoded any cast and have at least one property key requestBody. The pattern breaks the found AST node down into two parts: Object definition and the type from the cast. Only the former is used in the fix rule and passed as a parameter to buildStubRequest.

Solving the Initial Problem

Armed with the experience from these simpler examples, it’s time to return to the initial problem: I want to move the contents of the handle method into a withContextHandler wrapper, as long as they are located in a class that implements the Handler interface. My solution is as follows:

id: wrap-handler-context
language: TypeScript
rule:
  pattern: "$BODY"
  kind: statement_block
  inside:
    kind: method_definition
    has:
      kind: property_identifier
      regex: "^handle$"
    inside:
      stopBy: end
      kind: class_declaration
      has:
        stopBy: end
        kind: implements_clause
        has:
          stopBy: end
          kind: type_identifier
          regex: "^Handler$"
fix: "{return await this.contextService.withContext(async () => $BODY)}"

This rule matches all statement blocks of a method definition, as long as the method is named exactly handle. Furthermore, this method definition must be located within a class declaration that defines the Handler interface. The task for the reader is to define a second rule that adds the contextService as a constructor parameter in the affected files. Bonus points if the constructors that already accept a contextService as a parameter are not changed ;)

Agent Skill

Now that I had gained an understanding of the possibilities that ast-grep offers, I want to try out the already available Claude Skill.

Instead of Claude, as described in the README of the skill, I’m using Google’s Gemini models. When it comes to coding tasks, however, on paper these are not far from the currently available Anthropic models (Comparison of models [Archive link] ). I keep the prompt deliberately vague:

please use the `ast-grep` skill to write a ast-grep rule that matches all
`handle` method inside a `Handler` and provide a `fix` to wrap the body inside a
`contextService.withContext` call

In the first attempt, I only use Gemini 3.5 Flash, a cost-effective model. The generated following output

id: ai2-wrap-handle-with-context
language: typescript
rule:
  kind: method_definition
  all:
    - has:
        field: name
        regex: ^handle$
    - inside:
        kind: class_declaration
        has:
          kind: implements_clause
          has:
            kind: type_identifier
            regex: Handler
        stopBy: end
    - not:
        has:
          pattern: this.contextService.withContext($$$)
          stopBy: end
    # Capture the method components for the fix
    - any:
        - pattern: "async handle($$$ARGS)$RET { $$$BODY }"
        - pattern: "public async handle($$$ARGS)$RET { $$$BODY }"
fix: |
  async handle($$$ARGS)$RET {
    return await this.contextService.withContext(async () => {
      $$$BODY
    });
  }

… is not a valid ast-grep rule, however:

✖ Caused by
╰▻ Fail to parse yaml as RuleConfig
╰▻ mapping values are not allowed in this context at line 11 column 26

ASTs are complex constructs. Instead of feeding the error back into the small model, I decide to switch to the larger Gemini 3.1 Pro and repeat the prompt:

id: ai3-wrap-handle-with-context
language: typescript
rule:
  pattern: |
    class $A implements $H {
      $$$BEFORE
      async handle($CNAME: $CTYPE, $ENAME: $ETYPE): $RET {
        $$$BODY
      }
      $$$AFTER
    }
  constraints:
    H:
      regex: Handler
fix: |
  class $A implements $H {
    $$$BEFORE
    async handle($CNAME: $CTYPE, $ENAME: $ETYPE): $RET {
      return await this.contextService.withContext($CNAME, $ENAME, async () => {
        $$$BODY
      });
    }
    $$$AFTER
  }

This output looks much more compact and elegant right away, but it also doesn’t work:

✖ Caused by
╰▻ Fail to parse yaml as RuleConfig
╰▻ unknown field `constraints`, expected one of `pattern`, `kind`, `regex`, `nthChild`, `range`, `inside`, `has`, `precedes`, `follows`, `all`, `any`, `not`, `matches`

The error would be easy to fix: A constraint may only exist at the top level. If this statement is moved accordingly, the rule works flawlessly. But since our goal is not to “think”, but to “prompt”, we feed the error message back into the LLM and get a simplified rule that works:

id: wrap-handle-with-context
language: typescript
rule:
  pattern: |
    class $A implements Handler {
      $$$BEFORE
      async handle($CNAME: $CTYPE, $ENAME: $ETYPE): $RET {
        $$$BODY
      }
      $$$AFTER
    }
fix: |
  class $A implements Handler {
    $$$BEFORE
    async handle($CNAME: $CTYPE, $ENAME: $ETYPE): $RET {
      return await this.contextService.withContext($CNAME, $ENAME, async () => {
        $$$BODY
      });
    }
    $$$AFTER
  }

The cost for the two prompts with the Pro model on my 300k LOC codebase according to OpenCode: $0.45. Significantly cheaper than reading every affected handler and paying for the entire content as output tokens again.

Summary

ast-grep is one of those tools that I wish I had found years ago. In my opinion, it deserves much more than the 14,500 stars it currently has on GitHub. Whether operated by hand or instrumented with a suitably equipped LLM: When working in large software projects, it can save a lot of time and/or money (for tokens).

And ast-grep can do much more that didn’t fit into this blog post: For example, the same rule engine can also be used to create linting rules for your own codebase, there is a json mode and APIs to process the output of the tool in your own programs, and when it comes to code rewrites, we haven’t even scratched the surface of what else would be possible with additional transforms.