Takeover & Detached Runs (Advanced)

// run-setup.ts
import { chat } from '@tanstack/ai'
import { InMemoryLockStore, withLocks } from '@tanstack/ai/locks'
import { claudeCodeText } from '@tanstack/ai-claude-code'
import { durableStream } from '@tanstack/ai-durable-stream'
import { memoryPersistence, withPersistence } from '@tanstack/ai-persistence'
import {
  InMemorySandboxInstanceStore,
  createSecrets,
  defineSandbox,
  defineWorkspace,
  githubRepo,
  withSandbox,
} from '@tanstack/ai-sandbox'
import { dockerSandbox } from '@tanstack/ai-sandbox-docker'
import type { StreamChunk } from '@tanstack/ai'

// Local dev stand-ins. In production, use stores every replica reads and a
// distributed LockStore. See Requirements.
export const persistence = memoryPersistence()
export const locks = new InMemoryLockStore()
export const instances = new InMemorySandboxInstanceStore()
export const { runs } = persistence.stores

// The Durable Streams backend every replica reaches.
export const durableOptions = {
  server: 'https://streams.example.com',
  streamPrefix: 'agent-runs',
}

export const sandbox = defineSandbox({
  id: 'repo-agent',
  provider: dockerSandbox({ image: 'node:22' }),
  workspace: defineWorkspace({
    source: githubRepo({ repo: 'owner/app' }),
    setup: [
      'npm install -g @anthropic-ai/claude-code --include=optional && claude --version',
    ],
    secrets: createSecrets({
      ANTHROPIC_API_KEY: process.env.ANTHROPIC_API_KEY ?? '',
    }),
  }),
})

// The run's delivery log, from its id alone. Works with no live request,
// so the reaper can use it too.
export function durabilityFor(runId: string) {
  const url = new URL('https://runs.internal/')
  url.searchParams.set('runId', runId)
  return durableStream(new Request(url), durableOptions)
}

// The claim gives `drive` an AbortSignal that fires when this host loses the
// run. chat() takes an AbortController, so mirror one onto the other.
function controllerFor(signal: AbortSignal): AbortController {
  const controller = new AbortController()
  const abort = (): void => controller.abort(signal.reason)
  if (signal.aborted) abort()
  else signal.addEventListener('abort', abort, { once: true })
  return controller
}

// Continue an existing run. Used by the takeover route and by the reaper.
export async function* driveRun(input: {
  runId: string
  threadId: string
  signal: AbortSignal
}): AsyncIterable<StreamChunk> {
  // The client sent no history: it is reconnecting, not asking a question.
  const messages = await persistence.stores.messages.loadThread(input.threadId)
  const stream = chat({
    adapter: claudeCodeText('claude-opus-4-8'),
    messages,
    // Reuse the threadId on the run record. A new one breaks the replay.
    threadId: input.threadId,
    runId: input.runId,
    abortController: controllerFor(input.signal),
    middleware: [
      withPersistence(persistence),
      withLocks(locks),
      withSandbox(sandbox, {
        runs,
        instances,
        // attach: true reads the run's existing journal. It does not start a
        // second agent. Only a takeover sets it, never a POST handler.
        durability: { adapter: durabilityFor(input.runId), attach: true },
      }),
    ],
  })
  for await (const chunk of stream) yield chunk
}

// routes/chat.ts
import {
  chat,
  chatParamsFromRequest,
  toServerSentEventsResponse,
} from '@tanstack/ai'
import { withLocks } from '@tanstack/ai/locks'
import { claudeCodeText } from '@tanstack/ai-claude-code'
import { durableStream } from '@tanstack/ai-durable-stream'
import { withPersistence } from '@tanstack/ai-persistence'
import { withSandbox } from '@tanstack/ai-sandbox'
import {
  durableOptions,
  instances,
  locks,
  persistence,
  runs,
  sandbox,
} from './run-setup'

export async function POST(request: Request) {
  const { messages, threadId, runId } = await chatParamsFromRequest(request)
  // @tanstack/ai-client sends X-Run-Id on every POST, so this addresses
  // agent-runs/<runId>. Use ONE instance for the middleware and the response.
  const adapter = durableStream(request, durableOptions)

  const stream = chat({
    adapter: claudeCodeText('claude-opus-4-8'),
    messages,
    threadId,
    runId,
    // No abortController that mirrors request.signal here. See below.
    middleware: [
      withPersistence(persistence),
      withLocks(locks),
      withSandbox(sandbox, { runs, instances, durability: { adapter } }),
    ],
  })

  return toServerSentEventsResponse(stream, { durability: { adapter } })
}

// routes/chat.ts: add this next to POST, and merge the imports.
import { resumeServerSentEventsResponse } from '@tanstack/ai'
import { durableStream } from '@tanstack/ai-durable-stream'
import { sandboxRunDriver } from '@tanstack/ai-sandbox'
import {
  durabilityFor,
  durableOptions,
  driveRun,
  locks,
  runs,
} from './run-setup'

export function GET(request: Request) {
  return resumeServerSentEventsResponse({
    // Replays the log from the offset on this request (?offset=-1 from a
    // join, or Last-Event-ID from a reconnect).
    adapter: durableStream(request, durableOptions),
    // Claims the run and keeps driving it next to the replay.
    driver: sandboxRunDriver({
      request,
      runs,
      locks,
      durability: durabilityFor,
      drive: driveRun,
    }),
  })
}

// routes/chat-active.ts
import { runs } from './run-setup'

export async function GET(request: Request) {
  const threadId = new URL(request.url).searchParams.get('threadId')
  if (threadId === null) {
    return new Response('threadId is required', { status: 400 })
  }
  const active = await runs.findActiveRun(threadId)
  return Response.json({ runId: active?.runId ?? null })
}

import { fetchServerSentEvents } from '@tanstack/ai-client'
import { useEffect, useState } from 'react'
import type { StreamChunk } from '@tanstack/ai'

export function ResumeInFlight({ threadId }: { threadId: string }) {
  const [chunks, setChunks] = useState<Array<StreamChunk>>([])

  useEffect(() => {
    const controller = new AbortController()
    const connection = fetchServerSentEvents('/api/chat')

    async function rejoin(): Promise<void> {
      const response = await fetch(
        `/api/chat/active?threadId=${encodeURIComponent(threadId)}`,
        { signal: controller.signal },
      )
      const body: unknown = await response.json()
      if (typeof body !== 'object' || body === null || !('runId' in body)) return
      const runId = body.runId
      if (typeof runId !== 'string') return
      for await (const chunk of connection.joinRun(runId, controller.signal)) {
        setChunks((previous) => [...previous, chunk])
      }
    }

    void rejoin().catch(() => {
      // The run finished, or there was none. Nothing to rejoin.
    })
    return () => controller.abort()
  }, [threadId])

  return <p>{chunks.length} events replayed</p>
}

// routes/chat-cancel.ts
import { RUN_CANCEL_REASON, requestRunCancel } from '@tanstack/ai'
import { runs } from './run-setup'

// Runs this process drives now. Add an entry where you create a run's
// AbortController, and remove it when the run ends.
export const driving = new Map<string, AbortController>()

export async function POST(request: Request) {
  const body: unknown = await request.json()
  if (typeof body !== 'object' || body === null || !('threadId' in body)) {
    return new Response('threadId is required', { status: 400 })
  }
  const threadId = body.threadId
  if (typeof threadId !== 'string') {
    return new Response('threadId must be a string', { status: 400 })
  }

  const active = await runs.findActiveRun(threadId)
  if (!active) return new Response(null, { status: 204 })

  // Durable: a driver on another host sees it.
  await requestRunCancel(runs, active.runId)
  // In-process: a driver on this host stops at once.
  driving.get(active.runId)?.abort(RUN_CANCEL_REASON)

  return new Response(null, { status: 204 })
}

import { fetchServerSentEvents, useChat } from '@tanstack/ai-react'

export function StoppableChat({ threadId }: { threadId: string }) {
  const chat = useChat({
    threadId,
    connection: fetchServerSentEvents('/api/chat'),
  })

  async function stopRun(): Promise<void> {
    // Local: stop rendering the stream now.
    chat.stop()
    // Remote: tell the server that this is a cancel, not a lost connection.
    await fetch('/api/chat/cancel', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ threadId }),
    })
  }

  return <button onClick={() => void stopRun()}>Stop</button>
}