nanakura/moeka-project
1
0
Fork 0
Code Issues Pull Requests 1 Actions 1 Packages Projects Releases Wiki Activity

refactor(inference): replace custom AsyncMutex with async-mutex package (#1660)

Browse Source 
## Summary

Address maintainer feedback from #1622 ([@nekomeowww's
comment](https://github.com/moeru-ai/airi/pull/1622#discussion_r2342632282)).

The project already uses `async-mutex` in other packages
(`stage-tamagotchi`, `electron-screen-capture`), and it's in the
workspace catalog. This replaces the custom `AsyncMutex` implementation
with the shared dependency.

- Replace `AsyncMutex.run()` → `Mutex.runExclusive()`
- Replace `AsyncMutex.reset()` → `Mutex.cancel()`
- Remove `async-mutex.ts` and its unit tests
- Remove `AsyncMutex` from inference barrel exports
- Add `async-mutex` as direct dependency of `@proj-airi/stage-ui`

**Regarding `@moeru/eventa` suggestion**
([comment](https://github.com/moeru-ai/airi/pull/1622#discussion_r2342631692)):
The `waitForMessage` pattern is a thin request-response abstraction over
`postMessage`, where the worker is driven by `@huggingface/transformers`
internally. Eventa's transport-agnostic RPC is better suited for
bidirectional channels (Electron IPC, WebSocket) than for this one-way
"post and wait" pattern. No change for now.

## Test plan

- [x] `pnpm exec vitest run packages/stage-ui/src/libs/inference/` — 11
tests pass (4 removed with custom impl)
- [x] `pnpm -F @proj-airi/stage-ui typecheck` — no errors
- [x] `pnpm lint:fix` — no new errors in changed files

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
NJX
2026-04-14 11:47:33 +08:00
committed by GitHub
co-authored by Claude Opus 4.6 autofix-ci[bot]
parent 803442f5d8
commit 0c3655ee8d
9 changed files with 40 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/stage-ui/package.json
+1
View File
@@ -103,6 +103,7 @@
"@xsai/tool": "catalog:",
"@xsai/utils-chat": "catalog:",
"animejs": "^4.3.6",
"async-mutex": "catalog:",
"better-auth": "catalog:",
"culori": "^4.0.2",
"d3": "catalog:",
packages/stage-ui/src/libs/inference/adapters/background-removal.ts
+7 -7
View File
@@ -10,9 +10,9 @@ import type { AllocationToken } from '../gpu-resource-coordinator'
import type { ProgressPayload } from '../protocol'
import { defaultPerfTracer } from '@proj-airi/stage-shared'
import { Mutex } from 'async-mutex'
import { removeInferenceStatus, updateInferenceStatus } from '../../../composables/use-inference-status'
import { AsyncMutex } from '../async-mutex'
import { MODEL_IDS, MODEL_NAMES, TIMEOUTS } from '../constants'
import { getGPUCoordinator, getLoadQueue, MODEL_VRAM_ESTIMATES } from '../coordinator'
import { LOAD_PRIORITY } from '../load-queue'
@@ -58,7 +58,7 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter {
let state: BackgroundRemovalAdapter['state'] = 'idle'
let allocationToken: AllocationToken | null = null
const operationMutex = new AsyncMutex()
const operationMutex = new Mutex()
function ensureWorker(): Worker {
if (!worker) {
@@ -66,9 +66,9 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter {
new URL('../../../workers/background-removal/worker.ts', import.meta.url),
{ type: 'module' },
)
worker.addEventListener('error', (event) => {
worker.addEventListener('error', (_event) => {
state = 'error'
operationMutex.reset(new Error(event.message ?? 'Worker error'))
operationMutex.cancel()
})
}
return worker
@@ -119,7 +119,7 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter {
}
async function load(onProgress?: (p: ProgressPayload) => void): Promise<void> {
return operationMutex.run(async () => {
return operationMutex.runExclusive(async () => {
state = 'loading'
updateInferenceStatus(MODEL_NAMES.BG_REMOVAL, { state: 'downloading', device: 'webgpu' })
@@ -172,7 +172,7 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter {
}
async function processImage(imageData: ImageData): Promise<ImageData> {
return defaultPerfTracer.withMeasure('inference', 'bg-removal-process', () => operationMutex.run(async () => {
return defaultPerfTracer.withMeasure('inference', 'bg-removal-process', () => operationMutex.runExclusive(async () => {
if (!worker || (state !== 'ready' && state !== 'processing'))
throw new Error('Model not loaded. Call load() first.')
@@ -222,7 +222,7 @@ export function createBackgroundRemovalAdapter(): BackgroundRemovalAdapter {
}
function terminateAdapter(): void {
operationMutex.reset(new Error('Adapter terminated'))
operationMutex.cancel()
if (worker) {
worker.terminate()
worker = null
packages/stage-ui/src/libs/inference/adapters/kokoro.ts
+11 -15
View File
@@ -10,9 +10,9 @@ import type { AllocationToken } from '../gpu-resource-coordinator'
import type { ProgressPayload } from '../protocol'
import { defaultPerfTracer } from '@proj-airi/stage-shared'
import { Mutex } from 'async-mutex'
import { removeInferenceStatus, updateInferenceStatus } from '../../../composables/use-inference-status'
import { AsyncMutex } from '../async-mutex'
import { MAX_RESTARTS, MODEL_NAMES, RESTART_DELAY_MS, TIMEOUTS } from '../constants'
import { getGPUCoordinator, getLoadQueue, MODEL_VRAM_ESTIMATES } from '../coordinator'
import { LOAD_PRIORITY } from '../load-queue'
@@ -161,8 +161,8 @@ export function createKokoroAdapter(): KokoroAdapter {
let allocationToken: AllocationToken | null = null
let currentModelStatusId: string | null = null
const operationMutex = new AsyncMutex()
const lifecycleMutex = new AsyncMutex()
const operationMutex = new Mutex()
const lifecycleMutex = new Mutex()
function initializeWorker(): void {
worker = new Worker(
@@ -172,13 +172,9 @@ export function createKokoroAdapter(): KokoroAdapter {
worker.addEventListener('error', handleWorkerError)
}
function handleWorkerError(event: ErrorEvent | Error): void {
const message = event instanceof Error
? event.message
: (event as ErrorEvent).message ?? 'Unknown worker error'
function handleWorkerError(_event: ErrorEvent | Error): void {
state = 'error'
operationMutex.reset(new Error(message))
operationMutex.cancel()
destroyWorker()
scheduleRestart()
}
@@ -218,7 +214,7 @@ export function createKokoroAdapter(): KokoroAdapter {
}
async function ensureStarted(): Promise<void> {
await lifecycleMutex.run(async () => {
await lifecycleMutex.runExclusive(async () => {
if (!worker) {
initializeWorker()
state = 'idle'
@@ -235,7 +231,7 @@ export function createKokoroAdapter(): KokoroAdapter {
): Promise<Voices> {
await ensureStarted()
return defaultPerfTracer.withMeasure('inference', 'kokoro-load-model', () => operationMutex.run(async () => {
return defaultPerfTracer.withMeasure('inference', 'kokoro-load-model', () => operationMutex.runExclusive(async () => {
state = 'loading'
const modelStatusId = `kokoro-${quantization}`
@@ -298,7 +294,7 @@ export function createKokoroAdapter(): KokoroAdapter {
}
async function generate(text: string, voice: VoiceKey): Promise<ArrayBuffer> {
return defaultPerfTracer.withMeasure('inference', 'kokoro-generate', () => operationMutex.run(async () => {
return defaultPerfTracer.withMeasure('inference', 'kokoro-generate', () => operationMutex.runExclusive(async () => {
if (!worker)
throw new Error('Worker not initialized. Call loadModel() first.')
@@ -341,7 +337,7 @@ export function createKokoroAdapter(): KokoroAdapter {
}
function terminateAdapter(): void {
operationMutex.reset(new Error('Adapter terminated'))
operationMutex.cancel()
destroyWorker()
if (allocationToken) {
removeInferenceStatus(allocationToken.modelId)
@@ -366,14 +362,14 @@ export function createKokoroAdapter(): KokoroAdapter {
// ---------------------------------------------------------------------------
let globalAdapter: KokoroAdapter | null = null
const singletonMutex = new AsyncMutex()
const singletonMutex = new Mutex()
/**
* Get the global Kokoro adapter instance.
* Creates and starts the worker on first call.
*/
export async function getKokoroAdapter(): Promise<KokoroAdapter> {
return singletonMutex.run(async () => {
return singletonMutex.runExclusive(async () => {
if (!globalAdapter)
globalAdapter = createKokoroAdapter()
return globalAdapter
packages/stage-ui/src/libs/inference/adapters/whisper.ts
+7 -11
View File
@@ -10,9 +10,9 @@ import type { AllocationToken } from '../gpu-resource-coordinator'
import type { ProgressPayload } from '../protocol'
import { defaultPerfTracer } from '@proj-airi/stage-shared'
import { Mutex } from 'async-mutex'
import { removeInferenceStatus, updateInferenceStatus } from '../../../composables/use-inference-status'
import { AsyncMutex } from '../async-mutex'
import { MAX_RESTARTS, MODEL_NAMES, RESTART_DELAY_MS, TIMEOUTS } from '../constants'
import { getGPUCoordinator, getLoadQueue, MODEL_VRAM_ESTIMATES } from '../coordinator'
import { LOAD_PRIORITY } from '../load-queue'
@@ -84,15 +84,11 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter {
let restartAttempts = 0
const messageHandlers = new Set<(event: WhisperEvent) => void>()
const operationMutex = new AsyncMutex()
function handleWorkerError(event: ErrorEvent | Error): void {
const message = event instanceof Error
? event.message
: (event as ErrorEvent).message ?? 'Unknown worker error'
const operationMutex = new Mutex()
function handleWorkerError(_event: ErrorEvent | Error): void {
state = 'error'
operationMutex.reset(new Error(message))
operationMutex.cancel()
destroyWorker()
scheduleRestart()
}
@@ -199,7 +195,7 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter {
async function load(
onProgress?: (p: ProgressPayload) => void,
): Promise<void> {
return operationMutex.run(async () => {
return operationMutex.runExclusive(async () => {
state = 'loading'
updateInferenceStatus(MODEL_NAMES.WHISPER, { state: 'downloading', device: 'webgpu' })
@@ -253,7 +249,7 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter {
}
async function transcribe(input: WhisperTranscribeInput): Promise<string> {
return defaultPerfTracer.withMeasure('inference', 'whisper-transcribe', () => operationMutex.run(async () => {
return defaultPerfTracer.withMeasure('inference', 'whisper-transcribe', () => operationMutex.runExclusive(async () => {
if (!worker || state !== 'ready')
throw new Error('Model not loaded. Call load() first.')
@@ -286,7 +282,7 @@ export function createWhisperAdapter(workerUrl: string | URL): WhisperAdapter {
}
function terminateAdapter(): void {
operationMutex.reset(new Error('Adapter terminated'))
operationMutex.cancel()
if (worker) {
worker.terminate()
worker = null
packages/stage-ui/src/libs/inference/async-mutex.test.ts
-68
View File
@@ -1,68 +0,0 @@
import { describe, expect, it } from 'vitest'
import { AsyncMutex } from './async-mutex'
describe('asyncMutex', () => {
it('should serialize concurrent operations', async () => {
const mutex = new AsyncMutex()
const order: number[] = []
const p1 = mutex.run(async () => {
await new Promise(r => setTimeout(r, 50))
order.push(1)
return 'a'
})
const p2 = mutex.run(async () => {
order.push(2)
return 'b'
})
const [r1, r2] = await Promise.all([p1, p2])
expect(r1).toBe('a')
expect(r2).toBe('b')
expect(order).toEqual([1, 2])
})
it('should handle errors without blocking subsequent operations', async () => {
const mutex = new AsyncMutex()
await expect(mutex.run(async () => {
throw new Error('fail')
})).rejects.toThrow('fail')
// Should still work after error
const result = await mutex.run(async () => 'ok')
expect(result).toBe('ok')
})
it('should reset and reject waiting operations', async () => {
const mutex = new AsyncMutex()
// Queue a waiting operation that will never resolve on its own
let resolveHold!: () => void
const holdPromise = new Promise<void>(r => resolveHold = r)
const hold = mutex.run(() => holdPromise).catch(() => 'hold-rejected')
const waiting = mutex.run(async () => 'waited').catch(() => 'waiting-rejected')
// Reset — should reject waiting operations
await new Promise(r => setTimeout(r, 10))
mutex.reset(new Error('reset'))
resolveHold()
const [holdResult, waitingResult] = await Promise.all([hold, waiting])
// Waiting should be rejected
expect(waitingResult).toBe('waiting-rejected')
// Hold may or may not be rejected (it was active)
expect(['hold-rejected', undefined]).toContain(holdResult)
})
it('should support synchronous callbacks', async () => {
const mutex = new AsyncMutex()
const result = await mutex.run(() => 42)
expect(result).toBe(42)
})
})
packages/stage-ui/src/libs/inference/async-mutex.ts
-62
View File
@@ -1,62 +0,0 @@
/**
* An async mutex that ensures only one callback runs at a time.
* Waiters queue up and are processed in FIFO order.
*
* Extracted from KokoroWorkerManager for reuse across inference workers.
*/
export class AsyncMutex {
private locked = false
private waiters: { resolve: () => void, reject: (error: Error) => void }[] = []
// Incremented on reset() to invalidate stale lock holders
private generation = 0
/**
* Executes the callback with exclusive access to the mutex.
* If the mutex is locked, waits in queue until it's our turn.
*/
async run<T>(callback: () => Promise<T> | T): Promise<T> {
const myGeneration = this.generation
if (this.locked) {
await new Promise<void>((resolve, reject) => {
this.waiters.push({ resolve, reject })
})
if (myGeneration !== this.generation) {
throw new Error('Mutex was reset')
}
}
this.locked = true
try {
return await callback()
}
finally {
if (myGeneration === this.generation) {
const next = this.waiters.shift()
if (next) {
next.resolve()
}
else {
this.locked = false
}
}
}
}
/**
* Cancels all waiting tasks and releases the lock.
* Atomic thanks to JavaScript's Run-To-Completion semantics.
*/
reset(error: Error = new Error('Mutex reset')): void {
this.generation++
this.locked = false
const waitersToReject = this.waiters
this.waiters = []
for (const waiter of waitersToReject) {
waiter.reject(error)
}
}
}
packages/stage-ui/src/libs/inference/index.ts
-2
View File
@@ -1,5 +1,3 @@
// Core
export { AsyncMutex } from './async-mutex'
// Cache utilities
export {
clearModelCache,
packages/stage-ui/src/libs/inference/worker-manager.ts
+10 -10
View File
@@ -2,7 +2,7 @@
* Generic inference worker manager.
*
* Provides lifecycle management (start / restart / terminate), request
* serialisation via AsyncMutex, timeout handling, and a unified
* serialisation via async-mutex, timeout handling, and a unified
* message protocol for any inference worker.
*
* NOTICE: Currently not consumed by any adapter. Adapters implement their
@@ -20,8 +20,8 @@ import type {
} from './protocol'
import { errorMessageFrom } from '@moeru/std'
import { Mutex } from 'async-mutex'
import { AsyncMutex } from './async-mutex'
import { createRequestId } from './protocol'
// ---------------------------------------------------------------------------
@@ -158,8 +158,8 @@ export function createInferenceWorkerManager(
let lastError: ErrorPayload | null = null
let restartAttempts = 0
const operationMutex = new AsyncMutex()
const lifecycleMutex = new AsyncMutex()
const operationMutex = new Mutex()
const lifecycleMutex = new Mutex()
// -- Worker lifecycle -----------------------------------------------------
@@ -181,7 +181,7 @@ export function createInferenceWorkerManager(
state = 'error'
// Reject all pending operations
operationMutex.reset(new Error(message))
operationMutex.cancel()
// Clean up and try to restart
destroyWorker()
@@ -223,7 +223,7 @@ export function createInferenceWorkerManager(
}
async function ensureStarted(): Promise<void> {
await lifecycleMutex.run(async () => {
await lifecycleMutex.runExclusive(async () => {
if (!worker) {
initializeWorker()
state = 'idle'
@@ -239,7 +239,7 @@ export function createInferenceWorkerManager(
): Promise<ModelReadyResponse> {
await ensureStarted()
return operationMutex.run(async () => {
return operationMutex.runExclusive(async () => {
state = 'loading'
const requestId = createRequestId()
@@ -278,7 +278,7 @@ export function createInferenceWorkerManager(
input: TInput,
onProgress?: (p: ProgressPayload) => void,
): Promise<TOutput> {
return operationMutex.run(async () => {
return operationMutex.runExclusive(async () => {
if (!worker)
throw new Error('Worker not initialized. Call loadModel() first.')
@@ -319,7 +319,7 @@ export function createInferenceWorkerManager(
}
async function unloadModel(): Promise<void> {
return operationMutex.run(async () => {
return operationMutex.runExclusive(async () => {
if (!worker)
return
@@ -341,7 +341,7 @@ export function createInferenceWorkerManager(
}
function terminateManager(): void {
operationMutex.reset(new Error('Manager terminated'))
operationMutex.cancel()
destroyWorker()
state = 'terminated'
}
pnpm-lock.yaml
Generated
+4
View File
@@ -3065,6 +3065,9 @@ importers:
animejs:
specifier: ^4.3.6
version: 4.3.6
async-mutex:
specifier: 'catalog:'
version: 0.5.0
better-auth:
specifier: 'catalog:'
version: 1.6.2(@opentelemetry/api@1.9.1)(@prisma/client@5.22.0)(better-sqlite3@12.5.0)(drizzle-kit@0.31.10)(drizzle-orm@0.45.2(@electric-sql/pglite@0.4.4)(@opentelemetry/api@1.9.1)(@prisma/client@5.22.0)(@types/pg@8.20.0)(better-sqlite3@12.5.0)(kysely@0.28.14)(pg@8.20.0)(postgres@3.4.9))(pg@8.20.0)(react@19.2.3)(vitest@4.1.4)(vue@3.5.32(typescript@5.9.3))
@@ -8733,6 +8736,7 @@ packages:
'@simple-git/argv-parser@1.1.0':
resolution: {integrity: sha512-sUKOu2lb5vGIWADNNLpscyj07DAeQZU3KLbnE2Tj53tW6BbDQKMly2CCfnR4oYzqtRELCPWfwaPg+Q0T8qfKBg==}
deprecated: Contains a breaking change that should be a major version bump
'@sindresorhus/base62@1.0.0':
resolution: {integrity: sha512-TeheYy0ILzBEI/CO55CP6zJCSdSWeRtGnHy8U8dWSUH4I68iqTsy7HkMktR4xakThc9jotkPQUXT4ITdbV7cHA==}