A lock makes one worker create the sandbox for a thread. A second run for that thread waits, then uses the same sandbox.
Only one caller at a time runs the critical section for a key.
Add withLocks to the middleware array. Put it before the middleware that uses the lock:
// api/chat.ts
import {
chat,
chatParamsFromRequest,
toServerSentEventsResponse,
} from '@tanstack/ai'
import { withLocks, InMemoryLockStore } from '@tanstack/ai/locks'
import { openaiText } from '@tanstack/ai-openai'
import { onePrepPerThread } from './one-prep-per-thread'
// Create the store once, so all requests share it.
// For more than one process, pass a distributed LockStore.
const locks = new InMemoryLockStore()
export async function POST(request: Request) {
const { messages, threadId, runId } = await chatParamsFromRequest(request)
const stream = chat({
adapter: openaiText('gpt-6.1-sol'),
messages,
threadId,
runId,
middleware: [withLocks(locks), onePrepPerThread],
})
return toServerSentEventsResponse(stream)
}Read the store with getLocks(ctx). Then wrap the critical section in withLock:
// one-prep-per-thread.ts
import { defineChatMiddleware } from '@tanstack/ai'
import { LocksCapability, getLocks } from '@tanstack/ai/locks'
export const onePrepPerThread = defineChatMiddleware({
name: 'one-prep-per-thread',
requires: [LocksCapability],
async onStart(ctx) {
const locks = getLocks(ctx)
await locks.withLock(`thread:${ctx.threadId}`, async (signal) => {
// Only one run for this thread is in here at a time.
// Pass `signal` to cancellable work when you use a lease-backed store.
console.log(`start ${ctx.threadId}`)
await new Promise((resolve) => setTimeout(resolve, 1000))
console.log(`end ${ctx.threadId}, lost lock: ${signal.aborted}`)
})
},
})When two runs for the same thread start at the same time, the server console shows start, end, start, end. The two critical sections never overlap.
When you use a sandbox, put withLocks before withSandbox:
import { withLocks, InMemoryLockStore } from '@tanstack/ai/locks'
import { withSandbox } from '@tanstack/ai-sandbox'
import type { SandboxDefinition } from '@tanstack/ai-sandbox'
declare const sandbox: SandboxDefinition
const middleware = [
withLocks(new InMemoryLockStore()),
withSandbox(sandbox), // after the providers
]Use a lock when more than one process or isolate can enter the same critical section for the same key:
You do not need a lock for:
import type { LockStore } from '@tanstack/ai/locks'
declare const locks: LockStore
// Mutual exclusion for a key. Lease-backed stores abort `signal` on loss.
await locks.withLock('thread:abc', async (signal) => {
// Critical section. Pass `signal` to cancellable work when you use leases.
void signal
})| Piece | Role |
|---|---|
| LockStore | Interface: withLock(key, fn) |
| withLocks(store) | Chat middleware that provides LocksCapability |
| InMemoryLockStore | Implementation for one process (a promise chain per key) |
| defineLock(lock) | Types your own LockStore inline |
| getLocks / provideLocks | Low-level capability accessors for custom middleware |
InMemoryLockStore is correct in one process only:
Wrap your own mutual-exclusion primitive with defineLock. It checks the object against the contract inline, so you get autocomplete and need no : LockStore annotation.
import { defineLock } from '@tanstack/ai/locks'
// Your distributed primitive. `acquire` waits until the key is free and returns
// a `release` (plus, for leases, a `signal` that fires when ownership is lost).
import { acquire } from './my-lock-backend'
export const locks = defineLock({
async withLock(key, fn) {
const { release, signal } = await acquire(key)
try {
return await fn(signal)
} finally {
release()
}
},
})Wire it as middleware with withLocks(locks). A production store must meet the rules in the next section.
A deployment with more than one instance needs a distributed store, for example a Durable Object or Redis. A good store:
A callback that ignores signal still type-checks, because () => Promise<T> is assignable. But a lease-backed store cannot protect you if the critical section keeps changing data after the abort.
To test your store:
On Redis, upstashLocks() from @upstash/agentkit-tanstack-ai is a published third-party LockStore on Upstash Redis. It renews its lease while fn runs, and it aborts signal if a renewal fails. See Upstash.
If you already own a custom middleware, call provideLocks(ctx, store) in its setup hook, and list LocksCapability in its provides. Later middleware then reads the store with getLocks(ctx).