@raidou/pi-notify 0.3.1 → 0.5.0

package/src/dashboard/session-store.ts

@@ -0,0 +1,87 @@
+import { realpathSync } from 'node:fs'
+import { basename } from 'node:path'
+
+import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
+
+import { Registrar } from '../shared/registrar.js'
+import type { StateTracker } from '../state-tracker.js'
+import type { SessionRecord, SessionState } from './state-store.js'
+import { updateState } from './state-store.js'
+
+export class SessionStore extends Registrar {
+  private readonly stateTracker: StateTracker
+  private sessionId: string | undefined = undefined
+  private meta: { pid: number; cwd: string; projectName: string } | undefined =
+    undefined
+
+  constructor(pi: ExtensionAPI, stateTracker: StateTracker) {
+    super(pi)
+    this.stateTracker = stateTracker
+  }
+
+  private async saveSession(nextState: SessionState): Promise<void> {
+    if (this.sessionId === undefined || this.meta === undefined) return
+
+    const now = Date.now()
+    const meta = this.meta
+    const sessionId = this.sessionId
+    const id = String(meta.pid)
+
+    await updateState((state) => {
+      const existing = state.sessions[id]
+      const isRunning = nextState === 'running'
+
+      const startRunningAt = isRunning ? now : undefined
+
+      const record: SessionRecord = {
+        pid: meta.pid,
+        sessionId,
+        cwd: meta.cwd,
+        projectName: meta.projectName,
+        startedAt: existing?.startedAt ?? now,
+        state: nextState,
+        startedRunningAt: startRunningAt,
+      }
+
+      return {
+        ...state,
+        sessions: { ...state.sessions, [id]: record },
+      }
+    })
+  }
+
+  protected override setup(): void {
+    this.pi.on('session_start', (_event, ctx) => {
+      let normalizedCwd = ctx.cwd
+      try {
+        normalizedCwd = realpathSync(ctx.cwd)
+      } catch {
+        // Use original path if realpathSync fails
+      }
+      const meta = {
+        pid: process.pid,
+        cwd: normalizedCwd,
+        projectName: basename(normalizedCwd),
+      }
+      const sessionId = ctx.sessionManager.getSessionId()
+      this.sessionId = sessionId
+      this.meta = meta
+      void this.saveSession('idle')
+    })
+
+    this.unsubscribes.push(
+      this.stateTracker.events.on('running', async () => {
+        await this.saveSession('running')
+      }),
+      this.stateTracker.events.on('idle', async () => {
+        await this.saveSession('idle')
+      }),
+      this.stateTracker.events.on('tool', async ({ data }) => {
+        await this.saveSession(`tool_call:${data}`)
+      }),
+      this.stateTracker.events.on('event', async ({ data }) => {
+        await this.saveSession(`event:${data}`)
+      }),
+    )
+  }
+}

package/src/dashboard/state-store.test.ts

@@ -0,0 +1,303 @@
+import { rmSync, unlinkSync, writeFileSync } from 'node:fs'
+import { dirname } from 'node:path'
+
+import {
+  afterAll,
+  afterEach,
+  beforeEach,
+  describe,
+  expect,
+  it,
+  vi,
+} from 'vitest'
+
+vi.mock('./consts.js', async () => {
+  const { mkdtempSync } = await import('node:fs')
+  const path = await import('node:path')
+  const { tmpdir } = await import('node:os')
+
+  const stateDir = mkdtempSync(path.join(tmpdir(), 'pi-notify-test-'))
+  return {
+    STATE_FILE: path.join(stateDir, 'state.json'),
+    STATE_TMP_FILE: path.join(stateDir, 'state.json.tmp'),
+  }
+})
+
+afterAll(() => {
+  rmSync(dirname(STATE_FILE), { recursive: true, force: true })
+})
+
+import { STATE_FILE } from './consts.js'
+import {
+  isProcessAlive,
+  readSessions,
+  readState,
+  updateState,
+} from './state-store.js'
+
+describe('readState', () => {
+  const testLockFile = `${STATE_FILE}.lock`
+
+  afterEach(() => {
+    try {
+      unlinkSync(testLockFile)
+      unlinkSync(STATE_FILE)
+    } catch {
+      // ignore
+    }
+  })
+
+  it('writes state with version 2', async () => {
+    await updateState((state) => ({ ...state, sessions: {} }))
+    expect(readState().version).toBe(2)
+  })
+
+  it('reads v1 records as-is with absent v2 fields', () => {
+    const record = {
+      pid: process.pid,
+      sessionId: 'v1-session',
+      cwd: process.cwd(),
+      projectName: 'v1-project',
+      startedAt: Date.now(),
+      state: 'idle',
+    }
+    writeFileSync(
+      STATE_FILE,
+      JSON.stringify({
+        version: 1,
+        sessions: { [String(process.pid)]: record },
+      }),
+      'utf8',
+    )
+
+    const state = readState()
+    expect(state.version).toBe(2)
+    expect(state.sessions[String(process.pid)]).toEqual(record)
+  })
+
+  it('drops malformed session records instead of crashing', () => {
+    writeFileSync(
+      STATE_FILE,
+      JSON.stringify({ version: 2, sessions: { x: null } }),
+      'utf8',
+    )
+    expect(readState().sessions).toEqual({ x: undefined })
+  })
+})
+
+describe('isProcessAlive', () => {
+  it('returns true for current process pid', () => {
+    expect(isProcessAlive(process.pid)).toBe(true)
+  })
+
+  it('returns false for non-existent pid', () => {
+    expect(isProcessAlive(999999)).toBe(false)
+  })
+
+  it('returns false for pid out of range', () => {
+    expect(isProcessAlive(2147483647)).toBe(false)
+  })
+
+  it('returns false for invalid pid', () => {
+    expect(isProcessAlive(0)).toBe(false)
+    expect(isProcessAlive(-1)).toBe(false)
+  })
+
+  it('handles ESRCH gracefully', () => {
+    const error = new Error('process not found') as Error & { code: string }
+    error.code = 'ESRCH'
+    vi.spyOn(process, 'kill').mockImplementation(() => {
+      throw error
+    })
+    try {
+      expect(isProcessAlive(12345)).toBe(false)
+    } finally {
+      vi.restoreAllMocks()
+    }
+  })
+
+  it('returns true for EPERM (process exists but no permission)', () => {
+    const error = new Error('permission denied') as Error & { code: string }
+    error.code = 'EPERM'
+    vi.spyOn(process, 'kill').mockImplementation(() => {
+      throw error
+    })
+    try {
+      expect(isProcessAlive(1)).toBe(true)
+    } finally {
+      vi.restoreAllMocks()
+    }
+  })
+})
+
+describe('readSessions', () => {
+  const testLockFile = `${STATE_FILE}.lock`
+
+  beforeEach(async () => {
+    try {
+      unlinkSync(testLockFile)
+    } catch {
+      // ignore
+    }
+    await updateState(() => ({ version: 2, sessions: {} }))
+  })
+
+  afterEach(() => {
+    try {
+      unlinkSync(testLockFile)
+    } catch {
+      // ignore
+    }
+  })
+
+  it('keeps alive sessions and removes dead ones', async () => {
+    const aliveSessionId = `alive-${Date.now()}`
+    const deadSessionId = `dead-${Date.now()}`
+    const deadPid = 999999
+    const aliveKey = String(process.pid)
+    const deadKey = String(deadPid)
+
+    await updateState((state) => ({
+      ...state,
+      sessions: {
+        ...state.sessions,
+        [aliveKey]: {
+          pid: process.pid,
+          sessionId: aliveSessionId,
+          cwd: process.cwd(),
+          projectName: 'alive-project',
+          startedAt: Date.now(),
+          state: 'running',
+        },
+        [deadKey]: {
+          pid: deadPid,
+          sessionId: deadSessionId,
+          cwd: process.cwd(),
+          projectName: 'dead-project',
+          startedAt: Date.now(),
+          state: 'running',
+        },
+      },
+    }))
+
+    const alive = await readSessions()
+
+    expect(alive.some((s) => s.sessionId === aliveSessionId)).toBe(true)
+    expect(alive.every((s) => s.sessionId !== deadSessionId)).toBe(true)
+
+    const state = await import('./state-store.js').then((m) => m.readState())
+    expect(aliveKey in state.sessions).toBe(true)
+    expect(deadKey in state.sessions).toBe(false)
+  })
+
+  it('returns all sessions when none are dead', async () => {
+    const sessionId = `all-alive-${Date.now()}`
+
+    await updateState((state) => ({
+      ...state,
+      sessions: {
+        ...state.sessions,
+        [sessionId]: {
+          pid: process.pid,
+          sessionId,
+          cwd: process.cwd(),
+          projectName: 'test-project',
+          startedAt: Date.now(),
+          state: 'running',
+        },
+      },
+    }))
+
+    const alive = await readSessions()
+
+    expect(alive.some((s) => s.sessionId === sessionId)).toBe(true)
+  })
+})
+
+describe('updateState', () => {
+  const testLockFile = `${STATE_FILE}.lock`
+
+  beforeEach(() => {
+    try {
+      unlinkSync(testLockFile)
+    } catch {
+      // ignore
+    }
+  })
+
+  afterEach(() => {
+    try {
+      unlinkSync(testLockFile)
+    } catch {
+      // ignore
+    }
+  })
+
+  it('writes state and preserves other fields', async () => {
+    await updateState((state) => ({
+      ...state,
+      sessions: {
+        ...state.sessions,
+        'test-session-1': {
+          pid: process.pid,
+          sessionId: 'test-session-1',
+          cwd: process.cwd(),
+          projectName: 'test-project',
+          startedAt: Date.now(),
+          state: 'running',
+        },
+      },
+    }))
+  })
+
+  it('handles concurrent updates correctly', async () => {
+    const sessionId = `test-concurrent-${Date.now()}`
+
+    await updateState((state) => ({
+      ...state,
+      sessions: {
+        ...state.sessions,
+        [sessionId]: {
+          pid: process.pid,
+          sessionId,
+          cwd: process.cwd(),
+          projectName: 'concurrent-test',
+          startedAt: Date.now(),
+          state: 'running',
+        },
+      },
+    }))
+
+    await Promise.all([
+      updateState((state) => ({
+        ...state,
+        sessions: {
+          ...state.sessions,
+          [sessionId]: state.sessions[sessionId]
+            ? {
+                ...state.sessions[sessionId],
+                startedAt: Date.now(),
+              }
+            : undefined,
+        },
+      })),
+      updateState((state) => ({
+        ...state,
+        sessions: {
+          ...state.sessions,
+          [sessionId]: state.sessions[sessionId]
+            ? {
+                ...state.sessions[sessionId],
+                startedAt: Date.now() + 1,
+              }
+            : undefined,
+        },
+      })),
+    ])
+
+    await updateState((state) => {
+      expect(state.sessions[sessionId]).toBeDefined()
+      return state
+    })
+  })
+})

package/src/dashboard/state-store.ts

@@ -3,35 +3,76 @@ import {
   mkdirSync,
   readFileSync,
   renameSync,
-  rmdirSync,
   writeFileSync,
 } from 'node:fs'
-import { dirname, join } from 'node:path'
+import { dirname } from 'node:path'
 
-import { getAgentDir } from '@earendil-works/pi-coding-agent'
+import { omit, partition } from 'lodash-es'
+import lockfile from 'proper-lockfile'
+
+import { STATE_FILE, STATE_TMP_FILE } from './consts.js'
+
+const ESRCH = 'ESRCH'
+const EPERM = 'EPERM'
+
+export function isProcessAlive(pid: number): boolean {
+  if (pid <= 0) return false
+  if (pid === process.pid) return true
+
+  try {
+    process.kill(pid, 0)
+    return true
+  } catch (err: unknown) {
+    if (err instanceof Error && 'code' in err) {
+      const code = (err as { code: string }).code
+      if (code === EPERM) return true
+      if (code === ESRCH) return false
+    }
+    return false
+  }
+}
+
+export async function readSessions(): Promise<SessionRecord[]> {
+  const state = readState()
+  const sessions = Object.values(state.sessions).filter(
+    (session): session is SessionRecord => session !== undefined,
+  )
+  const [alive, dead] = partition(sessions, (session) =>
+    isProcessAlive(session.pid),
+  )
+
+  const deadIds = dead.map((session) => String(session.pid))
+
+  if (deadIds.length === 0) return alive
+
+  await updateState((s) => {
+    const sessions: typeof s.sessions = omit(s.sessions, deadIds)
+    return { ...s, sessions }
+  })
+
+  return alive
+}
+
+export type SessionState =
+  'running' | 'idle' | `tool_call:${string}` | `event:${string}`
 
 export interface SessionRecord {
   pid: number
+  sessionId: string
   cwd: string
   projectName: string
   startedAt: number
-  lastHeartbeatAt: number
-  state: 'running' | 'idle'
-  stateChangedAt: number
-  model: string | undefined
-  lastEvent: { type: string; summary: string; at: number } | undefined
+  state: SessionState
+  startedRunningAt?: number
 }
 
 export interface DashboardState {
-  version: 1
+  version: 2
   sessions: Record<string, SessionRecord | undefined>
 }
 
-const STATE_FILE = join(getAgentDir(), 'pi-notify', 'state.json')
-const LOCK_FILE = `${STATE_FILE}.lock`
-const LOCK_RETRY_DELAY_MS = 50
+const LOCK_RETRY_INTERVAL_MS = 50
 const LOCK_MAX_RETRIES = 20
-const STALE_THRESHOLD_MS = 30000
 
 function ensureStateDir(): void {
   const dir = dirname(STATE_FILE)
@@ -40,117 +81,101 @@ function ensureStateDir(): void {
   }
 }
 
-async function acquireLock(retries = LOCK_MAX_RETRIES): Promise<boolean> {
-  for (let i = 0; i < retries; i++) {
-    try {
-      mkdirSync(LOCK_FILE, { recursive: false })
-      return true
-    } catch (e) {
-      const code =
-        e instanceof Error && 'code' in e
-          ? (e as NodeJS.ErrnoException).code
-          : undefined
-      if (code === 'EEXIST') {
-        await new Promise((resolve) => setTimeout(resolve, LOCK_RETRY_DELAY_MS))
-        continue
-      }
-      if (code === 'ENOENT') {
-        ensureStateDir()
-        await new Promise((resolve) => setTimeout(resolve, LOCK_RETRY_DELAY_MS))
-        continue
-      }
-      throw e
-    }
+ensureStateDir()
+
+function parseSessionRecord(value: unknown): SessionRecord | undefined {
+  if (typeof value !== 'object' || value === null) return undefined
+  const record = value as Record<string, unknown>
+  if (
+    typeof record.pid !== 'number' ||
+    typeof record.sessionId !== 'string' ||
+    typeof record.cwd !== 'string' ||
+    typeof record.projectName !== 'string' ||
+    typeof record.startedAt !== 'number' ||
+    !isSessionState(record.state)
+  ) {
+    return undefined
+  }
+  return {
+    pid: record.pid,
+    sessionId: record.sessionId,
+    cwd: record.cwd,
+    projectName: record.projectName,
+    startedAt: record.startedAt,
+    state: record.state,
+    startedRunningAt:
+      typeof record.startedRunningAt === 'number'
+        ? record.startedRunningAt
+        : undefined,
   }
-  return false
 }
 
-function releaseLock(): void {
-  try {
-    rmdirSync(LOCK_FILE)
-  } catch {
-    // Lock file may not exist
-  }
+function isActivityState(
+  value: string,
+): value is `tool_call:${string}` | `event:${string}` {
+  return value.startsWith('tool_call:') || value.startsWith('event:')
 }
 
-function isPidAlive(pid: number): boolean {
+function isSessionState(value: unknown): value is SessionState {
+  return (
+    value === 'running' ||
+    value === 'idle' ||
+    (typeof value === 'string' && isActivityState(value))
+  )
+}
+
+function parseState(data: string): DashboardState | undefined {
+  let parsed: unknown
   try {
-    process.kill(pid, 0)
-    return true
+    parsed = JSON.parse(data)
   } catch {
-    return false
+    return undefined
   }
+  if (typeof parsed !== 'object' || parsed === null) return undefined
+  const raw = parsed as Record<string, unknown>
+  if (typeof raw.sessions !== 'object' || raw.sessions === null)
+    return undefined
+
+  const sessions: Record<string, SessionRecord | undefined> = {}
+  for (const [key, value] of Object.entries(raw.sessions)) {
+    sessions[key] = parseSessionRecord(value)
+  }
+  // Records from earlier versions without optional fields parse as-is.
+  return { version: 2, sessions }
 }
 
-export function readState(options?: { filterStale?: boolean }): DashboardState {
-  ensureStateDir()
+export function readState(): DashboardState {
   if (!existsSync(STATE_FILE)) {
-    return { version: 1, sessions: {} }
+    return { version: 2, sessions: {} }
   }
 
-  let state: DashboardState
   try {
-    const data = readFileSync(STATE_FILE, 'utf8')
-    state = JSON.parse(data) as DashboardState
+    const state = parseState(readFileSync(STATE_FILE, 'utf8'))
+    return state ?? { version: 2, sessions: {} }
   } catch {
-    return { version: 1, sessions: {} }
-  }
-
-  if (options?.filterStale === false) {
-    return state
+    return { version: 2, sessions: {} }
   }
-
-  return cleanupStale(state).cleaned
-}
-
-export function readSessions(): SessionRecord[] {
-  const state = readState()
-  return Object.values(state.sessions).filter(
-    (s): s is SessionRecord => s !== undefined,
-  )
 }
 
 export async function updateState(
   mutator: (state: DashboardState) => DashboardState,
 ): Promise<void> {
-  ensureStateDir()
-
-  const locked = await acquireLock()
-  if (!locked) {
-    throw new Error('Failed to acquire lock for state update')
-  }
+  const release = await lockfile.lock(STATE_FILE, {
+    realpath: false,
+    stale: 30000,
+    retries: {
+      minTimeout: LOCK_RETRY_INTERVAL_MS,
+      maxTimeout: LOCK_RETRY_INTERVAL_MS,
+      retries: LOCK_MAX_RETRIES,
+    },
+  })
 
   try {
     const state = readState()
     const newState = mutator(state)
-    const tmpFile = `${STATE_FILE}.tmp`
-    writeFileSync(tmpFile, JSON.stringify(newState, null, 2), 'utf8')
-    renameSync(tmpFile, STATE_FILE)
+    writeFileSync(STATE_TMP_FILE, JSON.stringify(newState, null, 2), 'utf8')
+    renameSync(STATE_TMP_FILE, STATE_FILE)
   } finally {
-    releaseLock()
+    await release()
   }
 }
-
-function cleanupStale(
-  state: DashboardState,
-  staleThresholdMs = STALE_THRESHOLD_MS,
-): { cleaned: DashboardState; removedIds: string[] } {
-  const now = Date.now()
-  const removedIds: string[] = []
-  const sessions: Record<string, SessionRecord | undefined> = {}
-
-  for (const [id, record] of Object.entries(state.sessions)) {
-    if (!record) continue
-    const heartbeatAge = now - record.lastHeartbeatAt
-    const isStale = heartbeatAge > staleThresholdMs
-    const isPidAliveValue = isPidAlive(record.pid)
-
-    if (!isStale && isPidAliveValue) {
-      sessions[id] = record
-    } else {
-      removedIds.push(id)
-    }
-  }
-
-  return { cleaned: { version: 1, sessions }, removedIds }
-}