AgentKit 2.0 でコンテキスト圧縮サブエージェントを設計する — 30分超のタスクで履歴肥大を止める本番パイプライン

// src/agent/compression-trigger.ts
import { encodingForModel } from "js-tiktoken";
 
const enc = encodingForModel("gpt-4o"); // 近似で十分
 
export interface CompressionDecision {
  shouldCompress: boolean;
  reason: "token-threshold" | "tool-event" | "turn-limit" | "skip";
  estimatedTokens: number;
}
 
const TOKEN_THRESHOLD_RATIO = 0.6; // モデル上限の 60%
const HEAVY_TOOLS = new Set(["search_codebase", "read_many_files"]);
 
export function evaluateCompression(
  messages: Array<{ role: string; content: string; toolName?: string }>,
  modelMaxTokens: number,
  turnsSinceLastCompression: number
): CompressionDecision {
  // 1. トークン数で測る
  const totalTokens = messages.reduce(
    (sum, m) => sum + enc.encode(m.content).length,
    0
  );
 
  if (totalTokens > modelMaxTokens * TOKEN_THRESHOLD_RATIO) {
    return { shouldCompress: true, reason: "token-threshold", estimatedTokens: totalTokens };
  }
 
  // 2. 直前が重いツール呼び出しなら、しきい値の 0.4 倍でも発火
  const last = messages[messages.length - 1];
  if (last?.toolName && HEAVY_TOOLS.has(last.toolName)) {
    if (totalTokens > modelMaxTokens * 0.4) {
      return { shouldCompress: true, reason: "tool-event", estimatedTokens: totalTokens };
    }
  }
 
  // 3. 最低保証 — 30 ターン経過したら強制圧縮
  if (turnsSinceLastCompression >= 30) {
    return { shouldCompress: true, reason: "turn-limit", estimatedTokens: totalTokens };
  }
 
  return { shouldCompress: false, reason: "skip", estimatedTokens: totalTokens };
}

// src/agent/compression-schema.ts
export interface CompressedContext {
  // ユーザーから受け取った制約・要件（消してはいけない）
  user_constraints: string[];
 
  // これまでに下した重要な決定（理由付きで残す）
  decisions: Array<{ decision: string; rationale: string }>;
 
  // 試して失敗したアプローチ（同じ間違いを避けるため）
  failed_attempts: Array<{ approach: string; reason_failed: string }>;
 
  // 実行済みのツール呼び出しサマリ（重複呼び出し防止）
  tool_calls_summary: Array<{ tool: string; target: string; outcome: "success" | "error" | "partial" }>;
 
  // 直近 5 ターンの自由記述要約（流れの維持用）
  recent_narrative: string;
 
  // タスクの現在地点（何が終わって何が残っているか）
  task_state: { completed: string[]; remaining: string[]; blockers: string[] };
}

// src/agent/compression-agent.ts
import type { Message, AgentKit } from "@google/agentkit";
import type { CompressedContext } from "./compression-schema";
 
const COMPRESSION_SYSTEM_PROMPT = `
あなたはエージェント実行履歴の構造化要約を作る専門のサブエージェントです。
入力として、メインエージェントが行ってきた一連のメッセージと
ツール呼び出し履歴を受け取ります。
 
以下の指針で要約してください:
 
1. ユーザーの制約条件（言語、ライブラリ、命名規則など）は一字一句残す
2. 「決定」と「失敗」を分けて記録する
3. ツール呼び出しは「何を / どこに / 結果は」の3点で要約する
4. 直近5ターンは流れがわかる程度に narrative として残す
5. 推測や創作はしない — 履歴に書かれていない情報は出力しない
`;
 
export async function compressContext(
  agentkit: AgentKit,
  messages: Message[]
): Promise<CompressedContext> {
  // 直近 5 ターンは「圧縮しない」ので分離
  const compressTarget = messages.slice(0, -5);
  const recent = messages.slice(-5);
 
  if (compressTarget.length === 0) {
    throw new Error("compressContext: 圧縮対象が 0 件です");
  }
 
  // structured output で JSON を強制
  const result = await agentkit.subAgents.run({
    name: "context-compressor",
    model: "gemma-4-flash",
    systemPrompt: COMPRESSION_SYSTEM_PROMPT,
    messages: [
      {
        role: "user",
        content: `次のメッセージ列を CompressedContext スキーマに従って要約してください:\n\n${JSON.stringify(compressTarget)}`,
      },
    ],
    responseSchema: COMPRESSED_CONTEXT_JSON_SCHEMA,
    temperature: 0.0, // 圧縮は決定論的にする
    maxOutputTokens: 2048,
  });
 
  // JSON パース失敗は明示的にエラー化
  let parsed: CompressedContext;
  try {
    parsed = JSON.parse(result.text);
  } catch (e) {
    throw new Error(`圧縮結果のJSONパースに失敗: ${result.text.slice(0, 200)}`);
  }
 
  // 最低限の妥当性検証
  if (!Array.isArray(parsed.user_constraints) || !parsed.task_state) {
    throw new Error("圧縮結果に必須フィールドが欠落しています");
  }
 
  // 直近 5 ターンの narrative にマージ
  const recentText = recent.map((m) => `[${m.role}] ${m.content.slice(0, 200)}`).join("\n");
  parsed.recent_narrative = `${parsed.recent_narrative}\n\n[直近5ターン保全]\n${recentText}`;
 
  return parsed;
}

{
  "user_constraints": [
    "TypeScript で書く",
    "テストは Vitest を使う",
    "認証ロジックは src/auth/ 配下に集約する"
  ],
  "decisions": [
    {
      "decision": "JWT のリフレッシュトークンは httpOnly Cookie に保存",
      "rationale": "XSS リスクを下げるため、localStorage は使わない"
    }
  ],
  "failed_attempts": [
    {
      "approach": "useAuth フックを App.tsx に直接配置",
      "reason_failed": "Context Provider が未設定でランタイムエラー"
    }
  ],
  "task_state": {
    "completed": ["ログインフォーム実装", "JWT 発行 API"],
    "remaining": ["リフレッシュトークン更新フロー", "ログアウト処理"],
    "blockers": []
  }
}

// src/agent/main-agent.ts
import { evaluateCompression } from "./compression-trigger";
import { compressContext } from "./compression-agent";
 
const mainAgent = agentkit.agents.create({
  name: "main",
  model: "gemini-2.5-pro",
  systemPrompt: MAIN_PROMPT,
 
  hooks: {
    onBeforeStep: async (ctx) => {
      const decision = evaluateCompression(
        ctx.messages,
        128_000, // gemini-2.5-pro の上限
        ctx.metadata.turnsSinceLastCompression ?? 0
      );
 
      ctx.logger.info({ decision }, "compression check");
 
      if (!decision.shouldCompress) {
        ctx.metadata.turnsSinceLastCompression =
          (ctx.metadata.turnsSinceLastCompression ?? 0) + 1;
        return;
      }
 
      // 直近 5 ターンを残し、それ以外を圧縮
      const compressed = await compressContext(ctx.agentkit, ctx.messages);
      const recentMessages = ctx.messages.slice(-5);
 
      ctx.replaceMessages([
        {
          role: "system",
          content: `[CONTEXT_SUMMARY]\n${JSON.stringify(compressed, null, 2)}`,
        },
        ...recentMessages,
      ]);
 
      ctx.metadata.turnsSinceLastCompression = 0;
      ctx.metrics.recordCompression({
        reason: decision.reason,
        beforeTokens: decision.estimatedTokens,
      });
    },
  },
});

// src/agent/eval/compression-ab-test.ts
interface EvalResult {
  taskId: string;
  completed: boolean;
  contradictions: number;
  totalInputTokens: number;
  totalOutputTokens: number;
  durationMs: number;
}
 
export async function runABTest(
  tasks: Array<{ id: string; prompt: string }>,
  options: { withCompression: boolean }
): Promise<EvalResult[]> {
  const results: EvalResult[] = [];
 
  for (const task of tasks) {
    const agent = createMainAgent({ enableCompression: options.withCompression });
    const start = Date.now();
 
    try {
      const out = await agent.run({ prompt: task.prompt, maxSteps: 100 });
      results.push({
        taskId: task.id,
        completed: out.status === "success",
        contradictions: countContradictions(out.history),
        totalInputTokens: out.usage.input,
        totalOutputTokens: out.usage.output,
        durationMs: Date.now() - start,
      });
    } catch (e) {
      results.push({
        taskId: task.id,
        completed: false,
        contradictions: 0,
        totalInputTokens: 0,
        totalOutputTokens: 0,
        durationMs: Date.now() - start,
      });
    }
  }
 
  return results;
}
 
// 矛盾検出は別 LLM に判定させる（ここでは割愛）
function countContradictions(history: Message[]): number {
  // 例: 「TypeScript で書く」と決めたあとに .js を生成していないか
  // をルールベースまたは judge LLM で数える
  return 0;
}