ycode/scripts/claude-append-protocol-probe.ts

import assert from 'node:assert/strict'
import { randomUUID } from 'node:crypto'
import { existsSync } from 'node:fs'
import { mkdir, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises'
import { createServer, type ServerResponse } from 'node:http'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { query, type SDKMessage, type SDKUserMessage } from '@anthropic-ai/claude-agent-sdk'

interface CapturedRequest {
  index: number
  url: string
  headers: Record<string, string | string[] | undefined>
  body: Record<string, unknown>
  userTexts: string[]
  latestText: string
  completed: boolean
  aborted: boolean
  responseText?: string
}

interface Deferred<T> {
  promise: Promise<T>
  resolve(value: T): void
}

class AsyncQueue<T> implements AsyncIterable<T> {
  private readonly values: T[] = []
  private readonly waiters: Array<(result: IteratorResult<T>) => void> = []
  private closed = false

  push(value: T): void {
    if (this.closed) throw new Error('input queue is closed')
    const waiter = this.waiters.shift()
    if (waiter) waiter({ value, done: false })
    else this.values.push(value)
  }

  close(): void {
    if (this.closed) return
    this.closed = true
    for (const waiter of this.waiters.splice(0)) waiter({ value: undefined, done: true })
  }

  [Symbol.asyncIterator](): AsyncIterator<T> {
    return {
      next: () => {
        const value = this.values.shift()
        if (value !== undefined) return Promise.resolve({ value, done: false })
        if (this.closed) return Promise.resolve({ value: undefined, done: true })
        return new Promise<IteratorResult<T>>((resolve) => this.waiters.push(resolve))
      }
    }
  }
}

function deferred<T>(): Deferred<T> {
  let resolve!: (value: T) => void
  const promise = new Promise<T>((resolvePromise) => { resolve = resolvePromise })
  return { promise, resolve }
}

function sse(res: ServerResponse, event: string, data: unknown): void {
  res.write(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`)
}

function messageStart(res: ServerResponse, id: string): void {
  sse(res, 'message_start', {
    type: 'message_start',
    message: {
      id,
      type: 'message',
      role: 'assistant',
      model: 'claude-append-probe',
      content: [],
      stop_reason: null,
      stop_sequence: null,
      usage: { input_tokens: 8, output_tokens: 1 }
    }
  })
}

function textBlock(res: ServerResponse, text: string): void {
  sse(res, 'content_block_start', {
    type: 'content_block_start',
    index: 0,
    content_block: { type: 'text', text: '' }
  })
  sse(res, 'content_block_delta', {
    type: 'content_block_delta',
    index: 0,
    delta: { type: 'text_delta', text }
  })
  sse(res, 'content_block_stop', { type: 'content_block_stop', index: 0 })
}

function messageEnd(res: ServerResponse, stopReason = 'end_turn'): void {
  sse(res, 'message_delta', {
    type: 'message_delta',
    delta: { stop_reason: stopReason, stop_sequence: null },
    usage: { output_tokens: 4 }
  })
  sse(res, 'message_stop', { type: 'message_stop' })
  res.end()
}

function plainText(res: ServerResponse, id: string, text: string): void {
  messageStart(res, id)
  textBlock(res, text)
  messageEnd(res)
}

function contentTexts(body: Record<string, unknown>): string[] {
  const messages = Array.isArray(body.messages) ? body.messages : []
  const texts: string[] = []
  for (const rawMessage of messages) {
    if (!rawMessage || typeof rawMessage !== 'object') continue
    const message = rawMessage as Record<string, unknown>
    if (message.role !== 'user') continue
    const content = message.content
    if (typeof content === 'string') {
      texts.push(content)
      continue
    }
    if (!Array.isArray(content)) continue
    for (const block of content) {
      if (!block || typeof block !== 'object') continue
      const item = block as Record<string, unknown>
      if (item.type === 'text' && typeof item.text === 'string') texts.push(item.text)
    }
  }
  return texts
}

function requestTrace(body: Record<string, unknown>): Array<Record<string, unknown>> {
  const messages = Array.isArray(body.messages) ? body.messages : []
  return messages.flatMap((rawMessage) => {
    if (!rawMessage || typeof rawMessage !== 'object') return []
    const message = rawMessage as Record<string, unknown>
    const content = message.content
    const parts = typeof content === 'string'
      ? [{ type: 'text', text: content }]
      : Array.isArray(content)
        ? content.flatMap((block) => {
          if (!block || typeof block !== 'object') return []
          const item = block as Record<string, unknown>
          return item.type === 'text' && typeof item.text === 'string'
            ? [{ type: 'text', text: item.text }]
            : [{ type: String(item.type) }]
        })
        : []
    const text = parts.flatMap((part) => 'text' in part ? [part.text] : []).join('')
    if (!text.includes('PROBE_') && !text.includes('answer-')) return []
    return [{ role: message.role, content: parts }]
  })
}

function compactFrame(message: SDKMessage): Record<string, unknown> {
  const value = message as unknown as Record<string, unknown>
  const inner = value.message && typeof value.message === 'object'
    ? value.message as Record<string, unknown>
    : undefined
  return {
    type: value.type,
    subtype: value.subtype,
    uuid: value.uuid,
    prompt_id: value.prompt_id,
    user_message_uuid: value.user_message_uuid,
    user_message_uuids: value.user_message_uuids,
    result_index: value.result_index,
    queued_turn_count: value.queued_turn_count,
    is_error: value.is_error,
    result: value.result,
    parent_tool_use_id: value.parent_tool_use_id,
    message_id: inner?.id,
    message_role: inner?.role,
    message_stop_reason: inner?.stop_reason,
    content_types: Array.isArray(inner?.content)
      ? inner.content.map((block) => block && typeof block === 'object'
        ? (block as Record<string, unknown>).type
        : typeof block)
      : undefined
  }
}

function inputMessage(sessionId: string, uuid: string, text: string): SDKUserMessage {
  return {
    type: 'user',
    message: { role: 'user', content: [{ type: 'text', text }] },
    parent_tool_use_id: null,
    uuid: uuid as SDKUserMessage['uuid'],
    session_id: sessionId
  }
}

function executablePath(): string {
  const executable = process.platform === 'win32' ? 'claude.exe' : 'claude'
  const packageName = `@anthropic-ai/claude-agent-sdk-${process.platform}-${process.arch}`
  const candidates = [
    join(process.cwd(), 'node_modules', ...packageName.split('/'), executable),
    join(process.cwd(), 'node_modules', '@anthropic-ai', 'claude-agent-sdk', 'vendor', executable)
  ]
  const found = candidates.find((candidate) => existsSync(candidate))
  if (!found) throw new Error(`Claude executable not found: ${candidates.join(', ')}`)
  return found
}

function isolatedEnv(tempRoot: string, baseUrl: string): NodeJS.ProcessEnv {
  const env: NodeJS.ProcessEnv = {
    ...process.env,
    HOME: join(tempRoot, 'home'),
    USERPROFILE: join(tempRoot, 'home'),
    CLAUDE_CONFIG_DIR: join(tempRoot, 'claude-config'),
    CLAUDE_SECURESTORAGE_CONFIG_DIR: join(tempRoot, 'claude-config'),
    ANTHROPIC_BASE_URL: baseUrl,
    ANTHROPIC_API_KEY: 'ycode-loopback-probe-key',
    ANTHROPIC_MODEL: 'claude-append-probe',
    ANTHROPIC_DEFAULT_OPUS_MODEL: 'claude-append-probe',
    ANTHROPIC_DEFAULT_SONNET_MODEL: 'claude-append-probe',
    ANTHROPIC_DEFAULT_HAIKU_MODEL: 'claude-append-probe',
    ANTHROPIC_SMALL_FAST_MODEL: 'claude-append-probe',
    CLAUDE_CODE_MAX_CONTEXT_TOKENS: '200000',
    CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC: '1',
    DISABLE_TELEMETRY: '1'
  }
  for (const name of [
    'ANTHROPIC_AUTH_TOKEN',
    'CLAUDE_CODE_OAUTH_TOKEN',
    'CLAUDE_CODE_USE_BEDROCK',
    'CLAUDE_CODE_USE_VERTEX',
    'CLAUDE_CODE_USE_FOUNDRY',
    'CLAUDE_CODE_USE_ANTHROPIC_AWS',
    'CLAUDE_CODE_USE_ANTHROPIC_GOOGLE_CLOUD',
    'CLAUDE_CODE_USE_MANTLE',
    'CLAUDE_CODE_USE_GATEWAY'
  ]) delete env[name]
  return env
}

async function waitFor<T>(description: string, predicate: () => T | undefined | false, timeoutMs = 20_000): Promise<T> {
  const deadline = Date.now() + timeoutMs
  while (Date.now() < deadline) {
    const value = predicate()
    if (value) return value
    await new Promise((resolve) => setTimeout(resolve, 20))
  }
  throw new Error(`timed out waiting for ${description}`)
}

const sdkPackage = JSON.parse(await readFile(
  join(process.cwd(), 'node_modules', '@anthropic-ai', 'claude-agent-sdk', 'package.json'),
  'utf8'
)) as { version?: string; claudeCodeVersion?: string }
const tempRoot = await mkdtemp(join(tmpdir(), 'ycode-claude-append-probe-'))
const workspace = join(tempRoot, 'workspace')
const outputPath = join(process.cwd(), '.tmp', 'claude-append-protocol-probe.json')
await mkdir(workspace, { recursive: true })
await mkdir(join(tempRoot, 'home'), { recursive: true })
await mkdir(join(tempRoot, 'claude-config'), { recursive: true })
await mkdir(join(process.cwd(), '.tmp'), { recursive: true })
await writeFile(join(workspace, 'README.md'), 'append protocol probe\n', 'utf8')

const requests: CapturedRequest[] = []
const runtimeMessages: SDKMessage[] = []
const stderr: string[] = []
const firstResponseGate = deferred<void>()
const processAbort = new AbortController()
const server = createServer((req, res) => {
  let raw = ''
  req.on('data', (chunk) => { raw += chunk })
  req.on('end', async () => {
    let body: Record<string, unknown> = {}
    try { body = JSON.parse(raw || '{}') as Record<string, unknown> } catch { body = { raw } }
    if (!req.url?.includes('/messages')) {
      res.writeHead(200, { 'content-type': 'application/json' })
      res.end('{}')
      return
    }

    const userTexts = contentTexts(body)
    const captured: CapturedRequest = {
      index: requests.length + 1,
      url: req.url ?? '',
      headers: req.headers,
      body,
      userTexts,
      latestText: userTexts.at(-1) ?? '',
      completed: false,
      aborted: false
    }
    requests.push(captured)
    res.on('close', () => {
      if (!captured.completed) captured.aborted = true
    })
    res.writeHead(200, {
      'content-type': 'text/event-stream',
      'cache-control': 'no-cache',
      connection: 'keep-alive'
    })
    const responseText = `answer-${captured.index}:${captured.userTexts.join('|')}`
    captured.responseText = responseText

    if (captured.index === 1) {
      await firstResponseGate.promise
    }
    if (!res.destroyed) {
      captured.completed = true
      plainText(res, `msg_append_${captured.index}`, responseText)
    }
  })
})

await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve))
const address = server.address()
if (!address || typeof address === 'string') throw new Error('loopback model did not bind')
const baseUrl = `http://127.0.0.1:${address.port}`

const sessionId = randomUUID()
const uuids = {
  initial: randomUUID(),
  appendA: randomUUID(),
  appendB: randomUUID()
}
const input = new AsyncQueue<SDKUserMessage>()
let instance: ReturnType<typeof query> | null = null
let pumpError: unknown
let pump = Promise.resolve()
let appendsQueuedBeforeFirstResponse = false

try {
  instance = query({
    prompt: input,
    options: {
      cwd: workspace,
      model: 'claude-append-probe',
      pathToClaudeCodeExecutable: executablePath(),
      env: isolatedEnv(tempRoot, baseUrl),
      includePartialMessages: true,
      abortController: processAbort,
      permissionMode: 'default',
      settingSources: [],
      strictMcpConfig: true,
      tools: [],
      allowedTools: [],
      plugins: [],
      agents: {},
      skills: [],
      systemPrompt: { type: 'preset', preset: 'claude_code' },
      sessionId,
      maxTurns: 4,
      stderr: (data) => {
        stderr.push(data)
        if (process.env.DEBUG_CLAUDE_APPEND_PROBE === '1') process.stderr.write(`[claude-probe] ${data}`)
      }
    }
  })

  const currentInstance = instance
  pump = (async () => {
    try {
      for await (const message of currentInstance) runtimeMessages.push(message)
    } catch (error) {
      pumpError = error
    }
  })()

  input.push(inputMessage(sessionId, uuids.initial, 'PROBE_INITIAL'))
  await waitFor('first model request', () => requests[0])
  // Hold request 1 so both appends are submitted while the turn is active.

  input.push(inputMessage(sessionId, uuids.appendA, 'PROBE_APPEND_A'))
  input.push(inputMessage(sessionId, uuids.appendB, 'PROBE_APPEND_B'))
  appendsQueuedBeforeFirstResponse = true
  await new Promise((resolve) => setTimeout(resolve, 250))
  firstResponseGate.resolve()

  await waitFor('all queued turns to settle', () => {
    const results = runtimeMessages.filter((message) => message.type === 'result')
    const last = results.at(-1) as unknown as Record<string, unknown> | undefined
    return results.length >= 2 && last?.queued_turn_count === 0
  }, 30_000)
  await new Promise((resolve) => setTimeout(resolve, 300))

  const results = runtimeMessages
    .filter((message) => message.type === 'result')
    .map((message) => message as unknown as Record<string, unknown>)
  const assistants = runtimeMessages
    .filter((message) => message.type === 'assistant')
    .map((message) => message as unknown as Record<string, unknown>)
  const userFrames = runtimeMessages
    .filter((message) => message.type === 'user')
    .map((message) => message as unknown as Record<string, unknown>)

  const firstResult = results[0]
  const finalResult = results.at(-1)
  const firstResultUuids = Array.isArray(firstResult?.user_message_uuids)
    ? firstResult.user_message_uuids as string[]
    : []
  const finalResultUuids = Array.isArray(finalResult?.user_message_uuids)
    ? finalResult.user_message_uuids as string[]
    : []

  assert.equal(requests[0]?.aborted, false, 'the first model request must complete without cancellation')
  assert.ok(requests.length >= 2, 'queued appends must produce another model request')
  assert.ok(firstResult, 'the initial turn must emit a result')
  assert.ok(finalResult, 'the queued append turn must emit a result')
  assert.ok(firstResultUuids.includes(uuids.initial), 'the first result must be attributed to the initial UUID')
  assert.ok(!firstResultUuids.includes(uuids.appendA), 'the first result must not claim append A')
  assert.ok(!firstResultUuids.includes(uuids.appendB), 'the first result must not claim append B')
  assert.ok(
    finalResultUuids.includes(uuids.appendA) && finalResultUuids.includes(uuids.appendB),
    'a later result must cover both queued appends'
  )
  assert.ok(results.length >= 2, 'one result must not stand for every queued send')
  assert.equal(userFrames.length, 0, 'text sends must not be replayed as SDK user frames')
  assert.equal(assistants.length, 2, 'each logical turn must emit one complete assistant frame')
  const assistantUuidSets = assistants.map((assistant) => Array.isArray(assistant.user_message_uuids)
    ? assistant.user_message_uuids as string[]
    : [])
  assert.ok(!assistantUuidSets[0]?.includes(uuids.appendA), 'the first assistant must not claim append A')
  assert.ok(
    assistantUuidSets[1]?.includes(uuids.appendA) && assistantUuidSets[1]?.includes(uuids.appendB),
    'the second assistant must cover both coalesced appends'
  )
  for (const [index, frame] of [...assistants, ...results].entries()) {
    const frameUuids = Array.isArray(frame.user_message_uuids) ? frame.user_message_uuids as string[] : []
    assert.equal(frameUuids.at(-1), frame.user_message_uuid, `frame ${index} must use the last member as user_message_uuid`)
  }
  assert.equal(pumpError, undefined, 'the SDK iterator must not fail during the probe')

  const report = {
    result: 'PASS',
    sdkVersion: sdkPackage.version,
    claudeCodeVersion: sdkPackage.claudeCodeVersion,
    loopback: baseUrl,
    activeAppend: {
      appendsQueuedBeforeFirstResponse,
      interruptCalled: false,
      inputClosedBeforeResults: false
    },
    uuids,
    requests: requests.map((request) => ({
      index: request.index,
      url: request.url,
      completed: request.completed,
      aborted: request.aborted,
      latestText: request.latestText,
      userTexts: request.userTexts
        .filter((text) => text.includes('PROBE_') || text.includes('answer-'))
        .slice(-8),
      messages: requestTrace(request.body)
    })),
    results: results.map((result) => ({
      subtype: result.subtype,
      is_error: result.is_error,
      result_index: result.result_index,
      queued_turn_count: result.queued_turn_count,
      user_message_uuid: result.user_message_uuid,
      user_message_uuids: result.user_message_uuids,
      prompt_id: result.prompt_id,
      result: result.result
    })),
    assistantKeys: assistants.map((assistant) => Object.keys(assistant).sort()),
    resultKeys: results.map((result) => Object.keys(result).sort()),
    assistants: assistants.map((assistant) => ({
      uuid: assistant.uuid,
      user_message_uuid: assistant.user_message_uuid,
      user_message_uuids: assistant.user_message_uuids,
      prompt_id: assistant.prompt_id,
      parent_tool_use_id: assistant.parent_tool_use_id,
      message_id: (assistant.message as Record<string, unknown> | undefined)?.id
    })),
    userFrameCount: userFrames.length,
    userFrames: userFrames.map((frame) => ({
      uuid: frame.uuid,
      user_message_uuid: frame.user_message_uuid,
      user_message_uuids: frame.user_message_uuids,
      prompt_id: frame.prompt_id,
      message: frame.message
    })),
    promptIdObserved: runtimeMessages.some((message) => {
      const value = message as unknown as Record<string, unknown>
      return typeof value.prompt_id === 'string'
    }),
    assistantPromptIdObserved: assistants.some((assistant) => typeof assistant.prompt_id === 'string'),
    frames: runtimeMessages.map(compactFrame),
    stderr: stderr.slice(-20)
  }
  await writeFile(outputPath, `${JSON.stringify(report, null, 2)}\n`, 'utf8')
  console.log(JSON.stringify({
    result: report.result,
    outputPath,
    uuids,
    activeAppend: report.activeAppend,
    requests: report.requests,
    results: report.results,
    assistants: report.assistants,
    userFrameCount: report.userFrameCount,
    userFrames: report.userFrames,
    assistantKeys: report.assistantKeys,
    resultKeys: report.resultKeys,
    promptIdObserved: report.promptIdObserved,
    assistantPromptIdObserved: report.assistantPromptIdObserved
  }, null, 2))
} catch (error) {
  const diagnostic = {
    result: 'FAIL',
    error: error instanceof Error ? `${error.name}: ${error.message}` : String(error),
    pumpError: pumpError instanceof Error ? `${pumpError.name}: ${pumpError.message}` : pumpError,
    uuids,
    requests: requests.map((request) => ({
      index: request.index,
      completed: request.completed,
      aborted: request.aborted,
      userTexts: request.userTexts
    })),
    frames: runtimeMessages.map(compactFrame),
    stderr: stderr.slice(-40)
  }
  await writeFile(outputPath, `${JSON.stringify(diagnostic, null, 2)}\n`, 'utf8')
  console.error(JSON.stringify(diagnostic, null, 2))
  throw error
} finally {
  input.close()
  processAbort.abort()
  instance?.close()
  await Promise.race([
    pump,
    new Promise((resolve) => setTimeout(resolve, 2000))
  ])
  await new Promise<void>((resolve) => server.close(() => resolve()))
  await rm(tempRoot, { recursive: true, force: true, maxRetries: 5, retryDelay: 200 }).catch(() => undefined)
}