AST を活用した Kiro の高精度なコード編集

本記事は「Surgical precision with AST-based code editing in Kiro」を翻訳したものです。

TL;DR: ここ数週間、私たちは新しい AST ベースのコードナビゲーション・編集エンジンをテストしてきました。このエンジンは、Kiro で最も頻繁なクエリタイプである機能リクエストの例を含むベンチマーク SWE-PolyBench において、トークン使用量を 20% 削減します。また、正確で耐障害性の高い、プロダクショングレードのコード変換を可能にします。脆弱な正規表現やファイル全体のダンプはもう不要です！

エージェントが 1 つの関数を見つけるために何千行も読み込み、わずかなフォーマットの違いのせいで更新に失敗することは AI コーディングアシスタントを使っているすべての開発者が経験したことがあるでしょう。現在のアプローチはファイル全体を読み込み、完全一致の文字列マッチングを行いますが、トークンを大量に消費し、簡単に壊れてしまいます。私たちはより良いものを構築しました。

本日、Kiro のコードベースでの作業に外科的精度を与える AST ベースのコードナビゲーション・編集システムを紹介します。任意の行範囲を探索したり脆弱な文字列マッチングを行う代わりに、このツールはコードを構造で捉えます。関数、クラス、インポートなどを対象とし、意味的な理解に基づいた型付き操作を適用します。

＊AST（抽象構文木） : コードから言語の意味に関係ある情報のみを取り出し、木構造で表現したものです。

テキストベースのコードツールの問題点

これまで、Kiro IDE はコードの確認に readFile、編集に strReplace という 2 つのテキストベースのツールに依存していました。シンプルではありますが、このアプローチには 2 つの根本的な問題があります。

1. 高いトークンコストとレイテンシ

特定の関数を探す際、エージェントはファイルの大部分、多くの場合ファイル全体を読み込む必要があります。これは、タスクに関係のないコンテキストに何千ものトークンが費やされることを意味します。

2. 脆弱な文字列マッチング

編集には完全一致の文字列マッチングが必要です。エージェントの期待と実際のコードの間にある空白、フォーマット、コメントのわずかな違いが、編集の失敗や意図しない複数マッチを引き起こします。その後、エージェントはミスを修正するために追加のイテレーションが必要になります。

これらの問題は複合的に作用します。エージェントがミスをし、診断のためにファイルを再度読み込み、別の編集を試み、このサイクルが繰り返されます。各イテレーションがより多くのトークンを消費し、レイテンシを増加させます。

AST ベースエンジンの仕組み

このエンジンは、テキストベースの操作を AST パースに置き換えます。コードを文字列として扱う代わりに、関数、クラス、メソッド、インポート、およびそれらの関係性といった構造化されたエンティティとしてコードを理解します。

コード読み取り: ターゲットを絞った情報抽出

ファイルの内容全体をダンプする代わりに、エンジンのコード読み取りツールは必要なものだけを返します:

• シグネチャ: 実装の詳細を含まない関数やクラスの定義
• 構造: 高レベルのコード構成と関係性
• 検索結果: 条件に一致する特定の定義

Java クラスを調べる 2 つのアプローチを比較してみましょう:

従来のアプローチ（1,309 トークン）:

Read entire AsyncRpcResult.java file including:
- License header
- All imports
- Complete implementation of 25+ methods
- Full method bodies with all logic

AST ベースのアプローチ（545 トークン）:

Class signatures in AsyncRpcResult.java:

class: AsyncRpcResult (Line 41)
  method: AsyncRpcResult(Invocation invocation) (Line 53)
  method: getValue() (Line 70)
  method: getException() (Line 82)
  method: thenApplyWithContext(Function<Result, Result> fn) (Line 139)
  [additional signatures...]

この例では、エンジンはナビゲーションと意思決定に必要な本質的な構造情報を提供しながら、トークンを 58% 削減しています。

コード書き込み: セマンティック編集

エンジンのコード書き込みツールは、セレクタを使用してコード要素を正確にターゲットします:

• ClassName.methodName で特定のメソッドをターゲット
• function:functionName で関数をターゲット
• field:fieldName でクラスフィールドをターゲット
• end でファイルの末尾に追加

4 つの型付き操作をサポートしています:

• insert_node: 特定の場所に新しいコードを追加
• replace_node: 関数やクラス全体を置換
• delete_node: コード要素をクリーンに削除
• replace_in_node: コードブロック内で外科的な編集を実行

ベンチマークからの実例として、TypeScript ファイルに新しい関数を追加するケースを紹介します:

従来のアプローチ（361 トークン）:

{
  "oldStr": "/**\n * Get a cached parsed configuration...\n */\nfunction getCachedParsedConfiguration(tsConfigPath: string): ParsedConfiguration {\n  if (tsConfigCache.has(tsConfigPath)) {\n    return tsConfigCache.get(tsConfigPath)!;\n  }\n  \n  const parsedConfig = readConfiguration(tsConfigPath);\n  tsConfigCache.set(tsConfigPath, parsedConfig);\n  return parsedConfig;\n}",
  "newStr": "/**\n * Get a cached parsed configuration...\n */\nfunction getCachedParsedConfiguration(tsConfigPath: string): ParsedConfiguration {\n  if (tsConfigCache.has(tsConfigPath)) {\n    return tsConfigCache.get(tsConfigPath)!;\n  }\n  \n  const parsedConfig = readConfiguration(tsConfigPath);\n  tsConfigCache.set(tsConfigPath, parsedConfig);\n  return parsedConfig;\n}\n\n/**\n * Clear the tsconfig cache...\n */\nexport function clearTsConfigCache(): void {\n  tsConfigCache.clear();\n}"
}

AST ベースのアプローチ（96 トークン）:

{
  "path": "packages/compiler-cli/ngcc/src/ngcc_options.ts",
  "selector": "end",
  "replacement": "/**\n * Clear the tsconfig cache...\n */\nexport function clearTsConfigCache(): void {\n  tsConfigCache.clear();\n}",
  "operation": "insert_node"
}

この例では、完全一致の文字列マッチングに必要な周辺コンテキストの記述を省略することで、エンジンはトークンを 73% 削減しています。

実際の効果

私たちは AST ベースエンジンを従来のツールと比較して、2 つのベンチマークで評価しました。

PolyBench50 の結果

PolyBench50（SWE-PolyBench のサブセット）で AST の読み取り・書き込み操作を実行したところ、一貫した改善が見られました:

機能リクエストのデモンストレーション

現実的な機能リクエストで両方のアプローチをテストしました:「AWS Resource Explorer にサードパーティ統合を追加」（Slack/Teams 通知、Jira チケット、SIEM エクスポート、AWS Config 統合）。

AST がプロダクション開発に重要な理由

AST ベースのコード操作の利点は、トークン効率にとどまりません:

耐障害性: フォーマットの変更が編集を壊しません。スペース 2 つでも 4 つでも、タブでもスペースでも、構造的な編集は成功します。

精度: ファイル内の他の類似コードを気にすることなく、変更したい箇所を正確にターゲットできます。

保守性: 今日機能する操作は、周囲のコードが進化しても明日も機能します。

理解力: AST パースはエージェントにコード構造の真の理解を提供し、どこでどのように変更を行うかについてよりスマートな判断を可能にします。

これは特に機能リクエストにおいて重要です。機能リクエストは Kiro で最も頻繁なクエリタイプであり、Vibe モードのトラフィックの 45%、Spec モードのトラフィックの 67.6% を占めています。これらのリクエストはコードベース全体にわたる複数のファイル変更を伴うことが多く、脆弱な文字列マッチングの複合的な影響が大きな摩擦を生み出します。これらの初期結果に私たちは興奮しており、現在プロダクションで利用可能です。