Antigravity × Cloudflare Durable Objects で作るステートフル AI エージェント — 1セッション=1インスタンスで会話・ツール状態・コストを一元管理

// src/agent-session.ts
// 1セッション=1インスタンスで会話履歴・ツール状態・トークン消費を保持する
// Durable Object クラス。fetch() ハンドラから AgentKit を呼び出す設計です
 
import { DurableObject } from "cloudflare:workers";
import { Agent, createAgent } from "@antigravity/agentkit";
import type { Env, Message, ToolCallRecord, MetricsSnapshot } from "./types";
 
export class AgentSession extends DurableObject<Env> {
  private history: Message[] = [];
  private toolCalls: ToolCallRecord[] = [];
  private metrics: MetricsSnapshot = { inputTokens: 0, outputTokens: 0, costUsd: 0 };
  private agent: Agent | null = null;
  private hydrated = false;
 
  // 起動時にストレージから状態を復元（lazy hydration）
  private async hydrate(): Promise<void> {
    if (this.hydrated) return;
    this.history = (await this.ctx.storage.get<Message[]>("history")) ?? [];
    this.toolCalls = (await this.ctx.storage.get<ToolCallRecord[]>("toolCalls")) ?? [];
    this.metrics = (await this.ctx.storage.get<MetricsSnapshot>("metrics")) ?? this.metrics;
    this.hydrated = true;
  }
 
  async fetch(request: Request): Promise<Response> {
    await this.hydrate();
    const url = new URL(request.url);
 
    if (url.pathname === "/chat" && request.method === "POST") {
      return this.handleChat(request);
    }
    if (url.pathname === "/metrics" && request.method === "GET") {
      return Response.json(this.metrics);
    }
    if (url.pathname === "/reset" && request.method === "POST") {
      return this.handleReset();
    }
    return new Response("Not found", { status: 404 });
  }
 
  private async handleReset(): Promise<Response> {
    this.history = [];
    this.toolCalls = [];
    await this.ctx.storage.delete(["history", "toolCalls"]);
    return Response.json({ ok: true });
  }
 
  // chat / agent 実装は次のステップで
  private async handleChat(_request: Request): Promise<Response> {
    return new Response("TODO", { status: 501 });
  }
}

// src/agent-session.ts (handleChat の実装)
private async handleChat(request: Request): Promise<Response> {
  const { userMessage } = await request.json<{ userMessage: string }>();
 
  // 受け取ったユーザーメッセージを履歴に追加
  this.history.push({ role: "user", content: userMessage, timestamp: Date.now() });
 
  // Agent はインスタンスごとに1度だけ作成（再ハイドレート時を含む）
  if (!this.agent) {
    this.agent = createAgent({
      model: "gemini-2.5-pro",
      tools: this.buildTools(),                    // 後述
      systemPrompt: "あなたはユーザーの開発作業を支援するアシスタントです。",
      onToolCall: (call) => this.recordToolCall(call),
      onUsage: (usage) => this.recordUsage(usage),
    });
  }
 
  // ストリーミング実行
  const stream = await this.agent.stream({ messages: this.history });
 
  const { readable, writable } = new TransformStream();
  const writer = writable.getWriter();
  const encoder = new TextEncoder();
  let assistantText = "";
 
  // バックグラウンドで stream を読みながらフロントへ流す
  this.ctx.waitUntil((async () => {
    try {
      for await (const chunk of stream) {
        if (chunk.type === "text") {
          assistantText += chunk.delta;
          await writer.write(encoder.encode(chunk.delta));
        }
      }
      // 完了後に履歴を保存（書き込みは1回にまとめる）
      this.history.push({ role: "assistant", content: assistantText, timestamp: Date.now() });
      await this.ctx.storage.put("history", this.history);
    } catch (err) {
      await writer.write(encoder.encode("\n[stream error]"));
      console.error("agent stream failed", err);
    } finally {
      await writer.close();
    }
  })());
 
  return new Response(readable, {
    headers: { "content-type": "text/event-stream", "cache-control": "no-cache" },
  });
}
 
private recordToolCall(call: ToolCallRecord): void {
  this.toolCalls.push({ ...call, timestamp: Date.now() });
  // 書き込みは累積して保存（後述の最適化セクション参照）
  this.ctx.storage.put("toolCalls", this.toolCalls).catch(() => {});
}
 
private recordUsage(usage: { inputTokens: number; outputTokens: number }): void {
  this.metrics.inputTokens += usage.inputTokens;
  this.metrics.outputTokens += usage.outputTokens;
  // gemini-2.5-pro の概算単価で USD を計算（参考値）
  this.metrics.costUsd =
    (this.metrics.inputTokens * 0.00000125) + (this.metrics.outputTokens * 0.000005);
  this.ctx.storage.put("metrics", this.metrics).catch(() => {});
}
 
private buildTools() {
  // ここでは省略（後で詳述）
  return [];
}

// 長時間ジョブを alarm() で分割実行する例
async runLongJob(jobId: string, payload: unknown): Promise<void> {
  await this.ctx.storage.put(`job:${jobId}:state`, { status: "queued", payload });
  // 1秒後に自分自身を起こす
  await this.ctx.storage.setAlarm(Date.now() + 1000);
}
 
async alarm(): Promise<void> {
  // alarm() は DO 自身が指定時刻に呼ばれるコールバック
  const jobs = await this.ctx.storage.list({ prefix: "job:" });
  for (const [key, value] of jobs) {
    const state = value as { status: string; payload: unknown };
    if (state.status !== "queued") continue;
 
    try {
      // 1ジョブあたり最大20秒で区切る（次の alarm に持ち越し）
      const result = await this.executeJobWithTimeout(state.payload, 20_000);
      await this.ctx.storage.put(key, { status: "done", result });
    } catch (err) {
      const errMessage = err instanceof Error ? err.message : "unknown";
      if (errMessage === "TIMEOUT") {
        // 次のスロットで続きを実行
        await this.ctx.storage.setAlarm(Date.now() + 1000);
        return;
      }
      await this.ctx.storage.put(key, { status: "failed", error: errMessage });
    }
  }
}
 
private async executeJobWithTimeout(payload: unknown, ms: number): Promise<unknown> {
  return new Promise(async (resolve, reject) => {
    const timer = setTimeout(() => reject(new Error("TIMEOUT")), ms);
    try {
      // 実際のツール処理（agent.run など）
      const result = await this.agent!.run({ messages: [{ role: "user", content: String(payload) }] });
      clearTimeout(timer);
      resolve(result);
    } catch (e) {
      clearTimeout(timer);
      reject(e);
    }
  });
}

// 同一 DO インスタンス内でツールを定義する例
private buildTools() {
  return [
    {
      name: "get_user_preference",
      description: "ユーザーが設定した好みやテーマカラーなどを返す",
      parameters: {
        type: "object",
        properties: { key: { type: "string" } },
        required: ["key"],
      },
      handler: async ({ key }: { key: string }) => {
        // ストレージは hydrate 済みなので追加 I/O はゼロ
        const prefs = (await this.ctx.storage.get<Record<string, string>>("prefs")) ?? {};
        return { value: prefs[key] ?? null };
      },
    },
    {
      name: "set_user_preference",
      description: "ユーザーの好みを更新する。書き込みはアトミック",
      parameters: {
        type: "object",
        properties: { key: { type: "string" }, value: { type: "string" } },
        required: ["key", "value"],
      },
      handler: async ({ key, value }: { key: string; value: string }) => {
        await this.ctx.storage.transaction(async (txn) => {
          const prefs = (await txn.get<Record<string, string>>("prefs")) ?? {};
          prefs[key] = value;
          await txn.put("prefs", prefs);
        });
        return { ok: true };
      },
    },
  ];
}

// セッションが「実行中」のとき、新しいリクエストはキューに入れる
private inFlight: Promise<void> | null = null;
 
private async handleChatGuarded(request: Request): Promise<Response> {
  // 直前のストリームが終わるまで待つ
  if (this.inFlight) {
    await this.inFlight.catch(() => {});
  }
  let resolve!: () => void;
  this.inFlight = new Promise((r) => { resolve = r; });
  try {
    return await this.handleChat(request);
  } finally {
    resolve();
    this.inFlight = null;
  }
}

// メトリクスとイベントを Analytics Engine に書き出す
private writeAnalytics(env: Env, event: { type: string; data: Record<string, unknown> }) {
  env.ANALYTICS.writeDataPoint({
    blobs: [event.type, this.ctx.id.toString()],
    doubles: [Date.now()],
    indexes: [event.type],
  });
}
 
// chat の最後に呼ぶ
this.writeAnalytics(this.env, {
  type: "chat_complete",
  data: { sessionId: this.ctx.id.toString(), tokens: this.metrics.outputTokens },
});