From 415a52055caf15a7d148b4b5c9e42835ca847810 Mon Sep 17 00:00:00 2001 From: Shunsuke Hayashi Date: Mon, 26 Jan 2026 07:35:38 +0900 Subject: [PATCH] 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 --- src/artifacts/manager.ts | 4 +- src/session/manager.ts | 111 +++++++++++++++++++++++++++++++-------- src/session/recovery.ts | 19 +++++-- src/theta/execute.ts | 42 +++++++++++---- 4 files changed, 139 insertions(+), 37 deletions(-) diff --git a/src/artifacts/manager.ts b/src/artifacts/manager.ts index 695c01ead..e0267eca9 100644 --- a/src/artifacts/manager.ts +++ b/src/artifacts/manager.ts @@ -59,12 +59,14 @@ export async function initializeTable(): Promise { 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: [ diff --git a/src/session/manager.ts b/src/session/manager.ts index da3352fdd..7fee632d7 100644 --- 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 { ], 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 { 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 = { + ":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 { + 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 { 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 { 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; diff --git a/src/session/recovery.ts b/src/session/recovery.ts index c7a353842..d6d3a2a4e 100644 --- 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) { diff --git a/src/theta/execute.ts b/src/theta/execute.ts index 862de758d..ea0338837 100644 --- a/src/theta/execute.ts +++ b/src/theta/execute.ts @@ -38,12 +38,15 @@ const DEFAULT_OPTIONS: ExecuteOptions = { export async function execute( state: ThetaCycleState, decision: DecisionResult, - executor: (params: Record) => Promise, + executor: (params: Record, signal?: AbortSignal) => Promise, options?: Partial, ): 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( }; state.events.push(executionEvent); - state.currentPhase = ThetaPhase.EXECUTE; state.context.set("execute.result", result); // Agentイベントを発行 @@ -140,9 +142,11 @@ export async function execute( /** * リトライ付きで実行する + * + * P1-2修正: AbortControllerを使用してタイムアウト時に実行を確実にキャンセル */ async function executeWithRetry( - executor: (params: Record) => Promise, + executor: (params: Record, signal?: AbortSignal) => Promise, params: Record, timeout: number, retries: number, @@ -151,14 +155,20 @@ async function executeWithRetry( 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( /** * タイムアウト付きで実行する + * + * P1-2修正: AbortControllerで確実にキャンセル */ -function withTimeout(promise: Promise, timeoutMs: number): Promise { +function withTimeout( + promise: Promise, + timeoutMs: number, + controller: AbortController, +): Promise { return Promise.race([ promise, - new Promise((_, reject) => - setTimeout(() => reject(new Error(`Timeout after ${timeoutMs}ms`)), timeoutMs), - ), + new Promise((_, 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( state: ThetaCycleState, decision: DecisionResult, - executor: (params: Record) => Promise, + executor: (params: Record, signal?: AbortSignal) => Promise, options?: Partial, ): string { const executionId = `bg_${Date.now()}_${Math.random().toString(36).slice(2, 11)}`;