セッションローテーション実装 (Invalid token対策)

.gitignore

@@ -32,3 +32,6 @@ report.[0-9]_.[0-9]_.[0-9]_.[0-9]_.json
 
 # Finder (MacOS) folder config
 .DS_Store
+
+# Agents
+.claude

chat.ts

@@ -3,29 +3,81 @@ import {
   fetchAnonymousToken,
   generateDeviceID,
   getSignedWsUrl,
+  refreshAuthToken,
   uploadFile,
 } from './core/types'
 import type { ChatRequestMessage, ChatResponseStream } from './core/chat-types'
 
 const DEFAULT_AGENT_ID = '6812e64f9dfaf301f7000001'
 
 export class User {
-  readonly accessToken: string
+  #accessToken: string
   readonly deviceId: string
-  constructor(deviceId: string, accessToken: string) {
-    this.accessToken = accessToken
+  #refreshToken: string
+  #refreshPromise: Promise<void> | null = null
+  #refreshTimer: ReturnType<typeof setTimeout> | null = null
+  get accessToken(): string {
+    return this.#accessToken
+  }
+  constructor(deviceId: string, accessToken: string, refreshToken: string, expiresAt: number | null = null) {
+    this.#accessToken = accessToken
     this.deviceId = deviceId
+    this.#refreshToken = refreshToken
+    if (expiresAt !== null) {
+      this.#scheduleRefresh(expiresAt)
+    }
   }
   static async create(): Promise<User> {
     const deviceId = generateDeviceID()
     const i = await fetchAnonymousToken(deviceId)
-    return new User(deviceId, i.accessToken)
+    return new User(deviceId, i.accessToken, i.refreshToken, i.expiresAt)
+  }
+  #scheduleRefresh(expiresAt: number): void {
+    if (this.#refreshTimer !== null) {
+      clearTimeout(this.#refreshTimer)
+    }
+    // 有効期限の5分前にリフレッシュ (すでに5分以内なら即時)
+    const delay = Math.max(0, expiresAt - Date.now() - 5 * 60 * 1000)
+    this.#refreshTimer = setTimeout(() => {
+      this.#refreshTimer = null
+      this.refresh().catch(() => {})
+    }, delay)
+  }
+  async refresh(): Promise<void> {
+    if (this.#refreshPromise) return this.#refreshPromise
+    this.#refreshPromise = this.#doRefresh()
+    try {
+      await this.#refreshPromise
+    } finally {
+      this.#refreshPromise = null
+    }
+  }
+  async #doRefresh(): Promise<void> {
+    const tokens = await refreshAuthToken(this.deviceId, this.#refreshToken).catch(() => fetchAnonymousToken(this.deviceId))
+    this.#accessToken = tokens.accessToken
+    this.#refreshToken = tokens.refreshToken
+    if (tokens.expiresAt !== null) {
+      this.#scheduleRefresh(tokens.expiresAt)
+    }
+  }
+  async #withTokenRefresh<T>(fn: () => Promise<T>): Promise<T> {
+    try {
+      return await fn()
+    } catch (err) {
+      if (err instanceof Error && (err.message.includes('Invalid token') || err.message.includes('HTTP Error 401'))) {
+        await this.refresh()
+        return await fn()
+      }
+      throw err
+    }
   }
   async createThread(opts?: { scenarioAgentId?: string; title?: string }): Promise<Thread> {
-    const thread = await createChatThread(this.deviceId, this.accessToken, {
-      scenarioAgentId: opts?.scenarioAgentId ?? DEFAULT_AGENT_ID,
-      title: opts?.title ?? '新しいスレッド',
-    })
+    const thread = await this.#withTokenRefresh(() =>
+      createChatThread(this.deviceId, this.#accessToken, {
+        scenarioAgentId: opts?.scenarioAgentId ?? DEFAULT_AGENT_ID,
+        title: opts?.title ?? '新しいスレッド',
+      })
+    )
     const connected = await Thread.connect(thread.id, this)
     return connected
   }
@@ -34,11 +86,13 @@ export class User {
     threadId?: string
     isImage?: boolean
   }): Promise<UploadedFile> {
-    const res = await uploadFile(this.deviceId, this.accessToken, {
-      file: opts.file,
-      threadId: opts.threadId ?? crypto.randomUUID(),
-      isImage: opts.isImage,
-    })
+    const res = await this.#withTokenRefresh(() =>
+      uploadFile(this.deviceId, this.#accessToken, {
+        file: opts.file,
+        threadId: opts.threadId ?? crypto.randomUUID(),
+        isImage: opts.isImage,
+      })
+    )
     return {
       fileId: res.fileId,
       fileUrl: res.fileUrl,
@@ -60,17 +114,20 @@ export class Thread {
   #ws: WebSocket
   #stream: ReadableStream<ChatResponseStream>
   #streamReader: ReadableStreamDefaultReader<ChatResponseStream>
+  #disposed = false
+  #pendingReject: ((err: unknown) => void) | null = null
   constructor(id: string, user: User, ws: WebSocket) {
     this.id = id
     this.#user = user
     this.#ws = ws
 
     this.#stream = new ReadableStream<ChatResponseStream>({
       start: async (controller) => {
-        while (true) {
+        while (!this.#disposed) {
           try {
             // 現在のWSの終了を待機するPromiseを作成
             const closed = new Promise<void>((resolve, reject) => {
+              this.#pendingReject = reject
               // FIXME: 場合によっては初っ端のメッセージを取りこぼすけど知らない
               this.#ws.onmessage = (e) => {
                 try {
@@ -81,28 +138,52 @@ export class Thread {
               }
               this.#ws.onerror = reject
               this.#ws.onclose = async () => {
+                const oldWs = this.#ws
                 try {
                   this.#ws = await Thread.connectWS(this.#user)
-                  resolve()
-                } catch (err) {
-                  reject(err)
+                } catch {
+                  // WS接続失敗時にトークンリフレッシュを試行して再接続
+                  // NOTE: WS APIではHTTPステータスを取得できないため、
+                  // トークン期限切れ以外のエラーでもリフレッシュが走る
+                  try {
+                    await this.#user.refresh()
+                    this.#ws = await Thread.connectWS(this.#user)
+                  } catch (err) {
+                    reject(err)
+                    return
+                  }
                 }
+                // 旧WSのハンドラをクリーンアップ
+                oldWs.onmessage = null
+                oldWs.onerror = null
+                oldWs.onclose = null
+                resolve()
               }
             })
 
             await closed
+            this.#pendingReject = null
           } catch (err) {
+            this.#ws.onmessage = null
             this.#ws.onerror = null
             this.#ws.onclose = null
-            controller.error(err)
+            this.#pendingReject = null
+            if (!this.#disposed) {
+              controller.error(err)
+            }
             break
           }
         }
       },
       cancel: () => {
+        this.#disposed = true
+        this.#ws.onmessage = null
         this.#ws.onerror = null
         this.#ws.onclose = null
-        this.#ws.close() // 一心同体
+        this.#ws.close()
+        // awaitで停止中のPromiseを解放してwhileループを抜けさせる
+        this.#pendingReject?.(new Error('disposed'))
+        this.#pendingReject = null
       },
     })
 

core/types.ts

@@ -14,6 +14,7 @@ interface AuthTokens {
   refreshToken: string
   idToken: string
   userType: string
+  expiresAt: number | null
 }
 
 // スレッド作成リクエストの型
@@ -131,14 +132,53 @@ export async function fetchAnonymousToken(
     throw new Error(`HTTP Error: ${response.status}`)
   }
 
-  const result = (await response.json()) as ApiResponse<AuthTokens>
+  const result = (await response.json()) as ApiResponse<Omit<AuthTokens, 'expiresAt'>>
 
   // 3. 業務エラーチェック (on.handle相当)
   if (result.code !== '0') {
     throw new Error(result.message || 'Authentication Failed')
   }
 
-  return result.data
+  const expires = response.headers.get('expires')
+  const expiresAt = expires ? new Date(expires).getTime() : null
+  return { ...result.data, expiresAt }
+}
+
+/** トークンのリフレッシュ */
+export async function refreshAuthToken(
+  deviceId: string,
+  refreshToken: string,
+): Promise<AuthTokens> {
+  const endpoint = '/api/v2/auth/refresh'
+  const url = `${BASE_URL}${endpoint}`
+
+  const signedHeaders = await getSignedHeaders('POST', endpoint)
+
+  const response = await fetch(url, {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      'X-Platform': 'WEB',
+      'X-Country-Code': 'JP',
+      'Device-ID': deviceId,
+      ...signedHeaders,
+    },
+    body: JSON.stringify({ refreshToken }),
+  })
+
+  if (!response.ok) {
+    throw new Error(`HTTP Error: ${response.status}`)
+  }
+
+  const result = (await response.json()) as ApiResponse<Omit<AuthTokens, 'expiresAt'>>
+
+  if (result.code !== '0') {
+    throw new Error(result.message || 'Token Refresh Failed')
+  }
+
+  const expires = response.headers.get('expires')
+  const expiresAt = expires ? new Date(expires).getTime() : null
+  return { ...result.data, expiresAt }
 }
 
 const WS_BASE_URL = 'wss://companion.ai.rakuten.co.jp' // Br.wsBasePath