fix(session, theta, artifacts): fix P1 bugs from Codex review

- P1-1: Set currentPhase at start of execute() instead of end - P1-2: Add AbortController support to timeout execution - P1-3: Add createdAt to DynamoDB AttributeDefinitions for GSI - P1-4: Remove PendingIndex to avoid hot key problem - P1-5: Add conditional update to prevent recovery race condition
2026-01-26 07:35:38 +09:00 · 2026-01-26 07:35:38 +09:00 · 415a52055c
commit 415a52055c
parent da1c5b26e0
4 changed files with 139 additions and 37 deletions
--- a/src/artifacts/manager.ts
+++ b/src/artifacts/manager.ts
@ -59,12 +59,14 @@ export async function initializeTable(): Promise<void> {
    await ddbClient.send(
      new CreateTableCommand({
        TableName: ARTIFACTS_TABLE,
+        // P1-3修正: GSIで使用する全属性をAttributeDefinitionsに追加
        AttributeDefinitions: [
          { AttributeName: "id", AttributeType: "S" },
          { AttributeName: "sessionId", AttributeType: "S" },
          { AttributeName: "userId", AttributeType: "S" },
          { AttributeName: "type", AttributeType: "S" },
-          { AttributeName: "expiresAt", AttributeType: "N" },
+          { AttributeName: "createdAt", AttributeType: "N" }, // GSI用
+          { AttributeName: "expiresAt", AttributeType: "N" }, // GSI用
        ],
        KeySchema: [{ AttributeName: "id", KeyType: "HASH" }],
        GlobalSecondaryIndexes: [
--- a/src/session/manager.ts
+++ b/src/session/manager.ts
@ -13,6 +13,7 @@ import {
  UpdateCommand,
  QueryCommand,
  DeleteCommand,
+  ScanCommand,
 } from "@aws-sdk/lib-dynamodb";
 import type {
  PersistedSessionState,
@ -74,14 +75,8 @@ export async function initializeTable(): Promise<void> {
            ],
            Projection: { ProjectionType: "ALL" },
          },
-          {
-            IndexName: "PendingIndex",
-            KeySchema: [
-              { AttributeName: "status", KeyType: "HASH" },
-              { AttributeName: "expiresAt", KeyType: "RANGE" },
-            ],
-            Projection: { ProjectionType: "ALL" },
-          },
+          // P1-4修正: PendingIndexを削除（ホットキー問題回避）
+          // 代わりにUserIndex/GuildIndexでstatus=RUNNINGをクエリ可能
        ],
        BillingMode: "PAY_PER_REQUEST",
        // Note: TTL is configured separately via UpdateTimeToLiveCommand
@ -165,7 +160,10 @@ export async function restoreState(sessionId: string): Promise<RestoredSession |
 /**
 * 未完了セッションを取得
 *
- * @param filter - フィルタ条件
+ * P1-4修正: PendingIndex削除に伴い、userIdまたはguildId必須に変更
+ * （ホットキー問題回避のため、全件取得は非推奨）
+ *
+ * @param filter - フィルタ条件 (userIdまたはguildId必須)
 * @returns 未完了セッション一覧
 */
 export async function getPendingSessions(
@ -173,6 +171,12 @@ export async function getPendingSessions(
 ): Promise<PersistedSessionState[]> {
  const now = Math.floor(Date.now() / 1000);

+  // P1-4修正: userIdまたはguildIdが必須
+  if (!filter.userId && !filter.guildId) {
+    console.warn("[SessionManager] getPendingSessions requires userId or guildId");
+    return [];
+  }
+
  // ステータスが指定されている場合はGSIを使用
  if (filter.status) {
    let indexName: string | undefined;
@ -190,9 +194,6 @@ export async function getPendingSessions(
      indexName = "GuildIndex";
      keyCondition = "guildId = :guildId AND #status = :status";
      expressionValues[":guildId"] = filter.guildId;
-    } else {
-      indexName = "PendingIndex";
-      keyCondition = "#status = :status AND expiresAt > :now";
    }

    const response = await docClient.send(
@ -210,19 +211,31 @@ export async function getPendingSessions(
    return (response.Items as unknown as PersistedSessionState[]) || [];
  }

-  // フィルタなしで全未完了セッションを取得
+  // userIdまたはguildIdでRUNNINGセッションを取得
+  const indexName = filter.userId ? "UserIndex" : "GuildIndex";
+  const keyCondition = filter.userId
+    ? "userId = :userId AND #status = :status"
+    : "guildId = :guildId AND #status = :status";
+
+  const expressionValues: Record<string, unknown> = {
+    ":status": Status.RUNNING,
+    ":now": now,
+  };
+  if (filter.userId) {
+    expressionValues[":userId"] = filter.userId;
+  } else {
+    expressionValues[":guildId"] = filter.guildId;
+  }
+
  const response = await docClient.send(
    new QueryCommand({
      TableName: SESSIONS_TABLE,
-      IndexName: "PendingIndex",
-      KeyConditionExpression: "#status = :status AND expiresAt > :now",
+      IndexName: indexName,
+      KeyConditionExpression: keyCondition,
      ExpressionAttributeNames: {
        "#status": "status",
      },
-      ExpressionAttributeValues: {
-        ":status": Status.RUNNING,
-        ":now": now,
-      },
+      ExpressionAttributeValues: expressionValues,
    }),
  );

@ -274,19 +287,68 @@ export async function updateStatus(sessionId: string, status: SessionStatus): Pr
  );
 }

+/**
+ * セッションステータスを条件付きで更新
+ *
+ * P1-5修正: 現在のステータスが指定値の場合のみ更新（二重実行防止）
+ *
+ * @param sessionId - セッションID
+ * @param newStatus - 新しいステータス
+ * @param expectedCurrentStatus - 期待する現在のステータス
+ * @returns 更新が成功したかどうか
+ */
+export async function updateStatusIf(
+  sessionId: string,
+  newStatus: SessionStatus,
+  expectedCurrentStatus: SessionStatus,
+): Promise<boolean> {
+  try {
+    await docClient.send(
+      new UpdateCommand({
+        TableName: SESSIONS_TABLE,
+        Key: { sessionId },
+        UpdateExpression: "SET #status = :newStatus, #updated = :updated",
+        ConditionExpression: "#status = :expectedStatus",
+        ExpressionAttributeNames: {
+          "#status": "status",
+          "#updated": "lastUpdateTime",
+        },
+        ExpressionAttributeValues: {
+          ":newStatus": newStatus,
+          ":expectedStatus": expectedCurrentStatus,
+          ":updated": Date.now(),
+        },
+      }),
+    );
+    return true; // 更新成功
+  } catch (error: unknown) {
+    // ConditionalCheckFailedExceptionは正常（他のプロセスが既に更新）
+    const err = error as { name: string; code: string };
+    if (err.code === "ConditionalCheckFailedException") {
+      return false; // 条件不一致で更新スキップ
+    }
+    throw error; // その他のエラーは再スロー
+  }
+}
+
 /**
 * 期限切れセッションをクリーンアップ
 *
 * 手動クリーンアップ用。通常はTTLにより自動削除される。
+ *
+ * P1-4修正: PendingIndex削除に伴い、Scanを使用（高コストなので注意）
+ *
+ * NOTE: userId/guildIdフィルタは削除（型定義の複雑化を回避）
+ * 必要に応じてgetPendingSessionsでユーザー/ギルド単位で実行してください
 */
 export async function cleanupExpiredSessions(): Promise<number> {
  const now = Math.floor(Date.now() / 1000);

+  // Scan操作（高コストなので注意して使用）
  const response = await docClient.send(
-    new QueryCommand({
+    new ScanCommand({
      TableName: SESSIONS_TABLE,
-      IndexName: "PendingIndex",
-      KeyConditionExpression: "#status = :status AND expiresAt < :now",
+      FilterExpression: "#status = :status AND expiresAt < :now",
      ExpressionAttributeNames: {
        "#status": "status",
      },
@ -299,7 +361,10 @@ export async function cleanupExpiredSessions(): Promise<number> {

  const expired = response.Items || [];
  for (const item of expired) {
-    await deleteSession(item.sessionId as string);
+    const sessionId = (item as unknown as { sessionId: string }).sessionId;
+    if (sessionId) {
+      await deleteSession(sessionId);
+    }
  }

  return expired.length;
--- a/src/session/recovery.ts
+++ b/src/session/recovery.ts
@ -5,7 +5,7 @@
 */

 import type { PersistedSessionState, PendingSessionsFilter } from "./types.js";
-import { getPendingSessions, updateStatus, deleteSession } from "./manager.js";
+import { getPendingSessions, updateStatus, updateStatusIf, deleteSession } from "./manager.js";
 import { SessionStatus } from "./types.js";

 /**
@ -87,8 +87,21 @@ export async function recoverPendingSessions(
    }

    try {
-      // ステータスを一時停止に変更（二重実行防止）
-      await updateStatus(session.metadata.sessionId, SessionStatus.PAUSED);
+      // P1-5修正: 条件付き更新（現在のステータスがRUNNINGの場合のみ）
+      // 他のリカバリプロセスが既に処理していた場合はスキップ
+      const updated = await updateStatusIf(
+        session.metadata.sessionId,
+        SessionStatus.PAUSED,
+        SessionStatus.RUNNING,
+      );
+
+      if (!updated) {
+        console.log(
+          `[Recovery] Session ${session.metadata.sessionId} already being recovered by another process`,
+        );
+        result.skipped++;
+        continue;
+      }

      // コールバック実行
      if (onRestore) {
--- a/src/theta/execute.ts
+++ b/src/theta/execute.ts
@ -38,12 +38,15 @@ const DEFAULT_OPTIONS: ExecuteOptions = {
 export async function execute<T = unknown>(
  state: ThetaCycleState,
  decision: DecisionResult,
-  executor: (params: Record<string, unknown>) => Promise<T>,
+  executor: (params: Record<string, unknown>, signal?: AbortSignal) => Promise<T>,
  options?: Partial<ExecuteOptions>,
 ): Promise<{ state: ThetaCycleState; result: ExecutionResult }> {
  const { runId } = state;
  const opts = { ...DEFAULT_OPTIONS, ...options };

+  // P1-1修正: 実行開始時にcurrentPhaseを設定
+  state.currentPhase = ThetaPhase.EXECUTE;
+
  // フェーズ開始イベント
  emitPhaseEvent(runId, ThetaEventType.PHASE_START, {
    phase: ThetaPhase.EXECUTE,
@ -82,7 +85,6 @@ export async function execute<T = unknown>(
    };

    state.events.push(executionEvent);
-    state.currentPhase = ThetaPhase.EXECUTE;
    state.context.set("execute.result", result);

    // Agentイベントを発行
@ -140,9 +142,11 @@ export async function execute<T = unknown>(

 /**
 * リトライ付きで実行する
+ *
+ * P1-2修正: AbortControllerを使用してタイムアウト時に実行を確実にキャンセル
 */
 async function executeWithRetry<T>(
-  executor: (params: Record<string, unknown>) => Promise<T>,
+  executor: (params: Record<string, unknown>, signal?: AbortSignal) => Promise<T>,
  params: Record<string, unknown>,
  timeout: number,
  retries: number,
@ -151,14 +155,20 @@ async function executeWithRetry<T>(
  let lastError: Error | null = null;

  for (let attempt = 0; attempt <= retries; attempt++) {
+    // 各試行ごとに新しいAbortControllerを作成
+    const controller = new AbortController();
+    const signal = controller.signal;
+
    try {
      // 進度通知
      onProgress?.(attempt / (retries + 1));

-      // タイムアウト付き実行
-      const result = await withTimeout(executor(params), timeout);
+      // タイムアウト付き実行（AbortSignalを渡す）
+      const result = await withTimeout(executor(params, signal), timeout, controller);
      return result;
    } catch (error) {
+      // タイムアウトの場合はAbortControllerでキャンセル済み
+      // それ以外のエラーも記録
      lastError = error instanceof Error ? error : new Error(String(error));

      // 最後の試行で失敗した場合はエラーを投げる
@ -177,13 +187,25 @@ async function executeWithRetry<T>(

 /**
 * タイムアウト付きで実行する
+ *
+ * P1-2修正: AbortControllerで確実にキャンセル
 */
-function withTimeout<T>(promise: Promise<T>, timeoutMs: number): Promise<T> {
+function withTimeout<T>(
+  promise: Promise<T>,
+  timeoutMs: number,
+  controller: AbortController,
+): Promise<T> {
  return Promise.race([
    promise,
-    new Promise<T>((_, reject) =>
-      setTimeout(() => reject(new Error(`Timeout after ${timeoutMs}ms`)), timeoutMs),
-    ),
+    new Promise<T>((_, reject) => {
+      const timeoutId = setTimeout(() => {
+        controller.abort(); // 実行をキャンセル
+        reject(new Error(`Timeout after ${timeoutMs}ms`));
+      }, timeoutMs);
+
+      // Promiseが解決したらクリアンアップ（キャンセルされても実行）
+      promise.finally?.(() => clearTimeout(timeoutId));
+    }),
  ]);
 }

@ -231,7 +253,7 @@ function emitPhaseEvent(runId: string, type: ThetaEventType, data: Record<string
 export function executeBackground<T = unknown>(
  state: ThetaCycleState,
  decision: DecisionResult,
-  executor: (params: Record<string, unknown>) => Promise<T>,
+  executor: (params: Record<string, unknown>, signal?: AbortSignal) => Promise<T>,
  options?: Partial<ExecuteOptions>,
 ): string {
  const executionId = `bg_${Date.now()}_${Math.random().toString(36).slice(2, 11)}`;