ycode/scripts/chat-store-regression.ts

import assert from 'node:assert/strict'
import { useChatStore } from '../src/renderer/src/stores/chat'
import type { SendResult } from '../src/shared/ipc'
import type { ChatEvent, ChatMessage, SessionMeta, ToolCallRecord } from '../src/shared/types'

const session: SessionMeta = { id: 'session-1', title: 'Regression', createdAt: 1, updatedAt: 1 }

function tool(overrides: Partial<ToolCallRecord> = {}): ToolCallRecord {
  return {
    id: 'tool-1',
    name: 'read_file',
    args: { path: 'notes.txt' },
    argsRaw: '{"path":"notes.txt"}',
    status: 'running',
    ...overrides
  }
}

function sendResult(runId: string): SendResult {
  return {
    runId,
    assistantMessageId: `${runId}-assistant`,
    modelId: 'model-profile-1',
    modelName: 'Frozen Model',
    userMessage: {
      id: `${runId}-user`,
      role: 'user',
      content: 'question',
      createdAt: 10
    }
  }
}

function resetStore(): void {
  useChatStore.setState({
    sessions: [session],
    currentSessionId: session.id,
    messages: [],
    sessionMessages: {},
    runs: {},
    sendingSessions: [],
    running: null,
    sending: false,
    loading: false,
    error: null
  })
}

function installFakeWindow(send: (payload: { sessionId: string; text: string }) => Promise<SendResult>): void {
  const api = {
    session: {
      list: async () => [session],
      create: async () => session,
      messages: async () => [],
      rename: async () => undefined,
      remove: async () => undefined
    },
    chat: {
      send,
      cancel: async () => ({ cancelled: true }),
      onEvent: () => () => undefined
    }
  }
  Object.defineProperty(globalThis, 'window', {
    configurable: true,
    value: { ycode: api }
  })
}

function assistantMessage(): ChatMessage {
  const message = useChatStore.getState().messages.find((entry) => entry.role === 'assistant')
  assert.ok(message, 'send result should create an assistant placeholder')
  return message
}

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

async function testEarlyOrderedEventsAndDuplicateGuard(): Promise<void> {
  resetStore()
  const result = sendResult('run-order')
  const response = deferred<SendResult>()
  let sendCalls = 0
  const events: ChatEvent[] = [
    { runId: 'other-run', type: 'text', delta: 'must be ignored' },
    { runId: result.runId, type: 'text', delta: 'before' },
    { runId: result.runId, type: 'tool_call', call: tool() },
    { runId: result.runId, type: 'text', delta: 'after' },
    {
      runId: result.runId,
      type: 'tool_result',
      call: tool({ status: 'ok', output: 'file contents', summary: 'file contents', durationMs: 12 })
    },
    {
      runId: result.runId,
      type: 'done',
      content: 'beforeafter',
      tools: [tool({ status: 'ok', output: 'file contents', summary: 'file contents', durationMs: 12 })]
    }
  ]

  installFakeWindow(async () => {
    sendCalls++
    for (const event of events) useChatStore.getState().applyEvent(event)
    return response.promise
  })

  const firstSend = useChatStore.getState().send('question')
  assert.equal(useChatStore.getState().sending, true, 'sending should cover the pending IPC request')
  assert.equal(useChatStore.getState().messages.length, 0, 'early events wait for the send result placeholder')
  await useChatStore.getState().send('duplicate')
  assert.equal(sendCalls, 1, 'a second send must be ignored while the first send is pending')

  response.resolve(result)
  await firstSend

  const message = assistantMessage()
  assert.equal(message.turnStartedAt, undefined, 'terminal events clear the live timer timestamp')
  assert.equal(message.content, 'beforeafter')
  assert.deepEqual(message.parts, [
    { type: 'text', text: 'before' },
    { type: 'tool', toolId: 'tool-1' },
    { type: 'text', text: 'after' }
  ])
  assert.equal(message.status, 'done')
  assert.equal(message.tools?.length, 1)
  assert.equal(message.tools?.[0].output, 'file contents')
  assert.equal(useChatStore.getState().running, null, 'done should clear the running turn')
  assert.equal(useChatStore.getState().sending, false, 'the send request should leave the sending state')
  assert.equal(message.content.includes('must be ignored'), false, 'events from another run must be ignored')
  useChatStore.getState().applyEvent({ runId: result.runId, type: 'text', delta: 'late' })
  assert.equal(assistantMessage().content, 'beforeafter', 'events after a terminal event must not change the message')
}

async function testLiveTurnStartTimestamp(): Promise<void> {
  resetStore()
  const result = sendResult('run-live-timer')
  installFakeWindow(async () => result)

  await useChatStore.getState().send('question')
  const message = assistantMessage()
  assert.equal(message.status, 'streaming')
  assert.equal(message.modelId, 'model-profile-1', 'assistant placeholder keeps the model identity from the send result')
  assert.equal(message.modelName, 'Frozen Model', 'assistant placeholder can render its model label before reloading history')
  assert.equal(message.turnStartedAt, result.userMessage.createdAt,
    'the live timer starts from the same user-submit timestamp used by terminal duration')

  useChatStore.getState().applyEvent({ runId: result.runId, type: 'text', delta: 'answer' })
  assert.equal(assistantMessage().turnStartedAt, result.userMessage.createdAt,
    'ordinary stream updates keep the timer origin stable')
  useChatStore.getState().applyEvent({ runId: result.runId, type: 'done', content: 'answer', tools: [], turnDurationMs: 1200 })
  assert.equal(assistantMessage().turnDurationMs, 1200)
  assert.equal(assistantMessage().turnStartedAt, undefined)
}

async function testEarlyTerminalEvents(): Promise<void> {
  const cases: Array<{
    name: string
    terminal: ChatEvent
    messageStatus: 'cancelled' | 'error'
    toolStatus: 'cancelled' | 'error'
  }> = [
    {
      name: 'cancelled',
      terminal: { runId: 'run-cancel', type: 'cancelled' },
      messageStatus: 'cancelled',
      toolStatus: 'cancelled'
    },
    {
      name: 'error',
      terminal: { runId: 'run-error', type: 'error', message: 'provider failed' },
      messageStatus: 'error',
      toolStatus: 'error'
    }
  ]

  for (const scenario of cases) {
    resetStore()
    const result = sendResult(scenario.terminal.runId)
    const response = deferred<SendResult>()
    installFakeWindow(async () => {
      useChatStore.getState().applyEvent({ runId: result.runId, type: 'text', delta: 'prefix' })
      useChatStore.getState().applyEvent({ runId: result.runId, type: 'tool_call', call: tool({ output: 'partial output' }) })
      useChatStore.getState().applyEvent({ runId: result.runId, type: 'text', delta: 'suffix' })
      useChatStore.getState().applyEvent(scenario.terminal)
      return response.promise
    })

    const sending = useChatStore.getState().send('question')
    response.resolve(result)
    await sending

    const message = assistantMessage()
    assert.equal(message.content, 'prefixsuffix', `${scenario.name}: text received before the result must be retained`)
    assert.deepEqual(message.parts, [
      { type: 'text', text: 'prefix' },
      { type: 'tool', toolId: 'tool-1' },
      { type: 'text', text: 'suffix' }
    ], `${scenario.name}: ordered content should survive a terminal event`)
    assert.equal(message.status, scenario.messageStatus)
    assert.equal(message.tools?.[0].status, scenario.toolStatus)
    assert.equal(message.tools?.[0].output, 'partial output', `${scenario.name}: terminal state must preserve tool output`)
    assert.equal(useChatStore.getState().running, null)
    assert.equal(useChatStore.getState().sending, false)
    if (scenario.name === 'error') assert.equal(useChatStore.getState().error, 'provider failed')
  }
}

async function testFailedSendCanRetry(): Promise<void> {
  resetStore()
  let sendCalls = 0
  const retryResult = sendResult('run-retry')
  installFakeWindow(async () => {
    sendCalls++
    if (sendCalls === 1) {
      useChatStore.getState().applyEvent({ runId: 'run-failed', type: 'text', delta: 'discarded failed run' })
      throw new Error('send failed')
    }
    useChatStore.getState().applyEvent({
      runId: retryResult.runId,
      type: 'done',
      content: 'retry complete',
      tools: []
    })
    return retryResult
  })

  await assert.rejects(useChatStore.getState().send('question'), /send failed/)
  assert.equal(useChatStore.getState().sending, false, 'a rejected send must reset sending')
  assert.equal(useChatStore.getState().error, 'send failed')
  assert.equal(useChatStore.getState().messages.length, 0, 'a failed send must not install a placeholder')

  await useChatStore.getState().send('retry')
  assert.equal(sendCalls, 2, 'a new send should be accepted after a failed request')
  const message = assistantMessage()
  assert.equal(message.content, 'retry complete', 'failed-run buffered events must not leak into a retry')
  assert.equal(message.status, 'done')
  assert.equal(useChatStore.getState().sending, false)
  assert.equal(useChatStore.getState().running, null)
}

async function testCancelFlushesPendingAndIgnoresLateEvents(): Promise<void> {
  resetStore()
  const result = sendResult('run-user-cancel')
  installFakeWindow(async () => result)
  let cancelCalls = 0
  window.ycode.chat.cancel = async (runId) => {
    cancelCalls++
    assert.equal(runId, result.runId)
    return { cancelled: true }
  }

  await useChatStore.getState().send('question')
  useChatStore.getState().applyEvent({ runId: result.runId, type: 'text', delta: 'already rendered' })
  useChatStore.getState().applyEvent({ runId: result.runId, type: 'text', delta: ' pending text' })
  await useChatStore.getState().cancel()
  assert.equal(cancelCalls, 1)
  assert.equal(assistantMessage().content, 'already rendered pending text', 'cancel must flush pending text before requesting cancellation')

  useChatStore.getState().applyEvent({ runId: result.runId, type: 'text', delta: 'late before terminal' })
  assert.equal(assistantMessage().content, 'already rendered pending text', 'text after a cancel request must be ignored')
  useChatStore.getState().applyEvent({ runId: result.runId, type: 'cancelled' })
  useChatStore.getState().applyEvent({ runId: result.runId, type: 'text', delta: 'late after terminal' })
  assert.equal(assistantMessage().content, 'already rendered pending text')
  assert.equal(assistantMessage().status, 'cancelled')
}

async function testSessionSwitchFlushesPendingIntoItsOwner(): Promise<void> {
  resetStore()
  const result = sendResult('run-session-switch')
  installFakeWindow(async () => result)
  await useChatStore.getState().send('question')
  useChatStore.getState().applyEvent({ runId: result.runId, type: 'text', delta: 'session one' })
  useChatStore.getState().applyEvent({ runId: result.runId, type: 'text', delta: ' pending' })

  await useChatStore.getState().selectSession('session-2')
  assert.equal(useChatStore.getState().messages.length, 0, 'switching sessions must not show the outgoing assistant message')
  assert.equal(useChatStore.getState().sessionMessages[session.id]?.find((message) => message.id === result.assistantMessageId)?.content,
    'session one pending', 'switching must flush buffered text into the session that owns the run')

  await useChatStore.getState().selectSession(session.id)
  assert.equal(assistantMessage().content, 'session one pending', 'restoring the session must preserve the flushed stream order')
  useChatStore.getState().applyEvent({ runId: result.runId, type: 'done', content: 'session one pending', tools: [] })
}

async function testParallelSessionsAndStaleLoads(): Promise<void> {
  resetStore()
  const a = deferred<SendResult>()
  const b = deferred<SendResult>()
  installFakeWindow(async (payload) => payload.sessionId === session.id ? a.promise : b.promise)
  const first = useChatStore.getState().send('A')
  await useChatStore.getState().selectSession('session-2')
  assert.equal(useChatStore.getState().sending, false)
  const second = useChatStore.getState().send('B')
  useChatStore.getState().applyEvent({ sessionId: session.id, runId: 'a', type: 'text', delta: 'early A' })
  useChatStore.getState().applyEvent({ sessionId: 'session-2', runId: 'b', type: 'text', delta: 'early B' })
  b.resolve(sendResult('b'))
  await second
  a.resolve(sendResult('a'))
  await first
  assert.equal(assistantMessage().content, 'early B', 'a background send cannot replace the foreground')
  assert.equal(useChatStore.getState().running?.runId, 'b')
  useChatStore.getState().applyEvent({ sessionId: session.id, runId: 'a', type: 'done', content: 'finished A', tools: [] })
  assert.equal(useChatStore.getState().running?.runId, 'b', 'background terminal cannot stop foreground run')
  await useChatStore.getState().selectSession(session.id)
  assert.equal(assistantMessage().content, 'finished A', 'switch reads current in-memory data')
  assert.equal(useChatStore.getState().running, null)
  await useChatStore.getState().selectSession('session-2')
  assert.equal(useChatStore.getState().running?.runId, 'b')
  useChatStore.getState().applyEvent({ runId: 'b', type: 'cancelled' })

  const slow = deferred<ChatMessage[]>()
  window.ycode.session.messages = async (id) => id === 'slow' ? slow.promise : [{ id: 'fast-user', role: 'user', content: 'fast', createdAt: 1 }]
  const oldSelection = useChatStore.getState().selectSession('slow')
  await useChatStore.getState().selectSession('fast')
  slow.resolve([{ id: 'old', role: 'user', content: 'slow', createdAt: 1 }])
  await oldSelection
  assert.equal(useChatStore.getState().currentSessionId, 'fast')
  assert.equal(useChatStore.getState().messages[0].content, 'fast', 'stale session fetch cannot replace a newer selection')
}

async function main(): Promise<void> {
  await testEarlyOrderedEventsAndDuplicateGuard()
  await testLiveTurnStartTimestamp()
  await testEarlyTerminalEvents()
  await testFailedSendCanRetry()
  await testCancelFlushesPendingAndIgnoresLateEvents()
  await testSessionSwitchFlushesPendingIntoItsOwner()
  await testParallelSessionsAndStaleLoads()
  resetStore()
  installFakeWindow(async () => sendResult('unused'))
  const creation = deferred<SessionMeta>()
  window.ycode.session.create = async () => creation.promise
  const creating = useChatStore.getState().newSession()
  await useChatStore.getState().selectSession('chosen-later')
  creation.resolve({ ...session, id: 'slow-created' })
  await creating
  assert.equal(useChatStore.getState().currentSessionId, 'chosen-later', 'late session creation cannot override a later selection')

  resetStore()
  useChatStore.setState({ currentSessionId: null })
  let sentTo = ''
  installFakeWindow(async (payload) => { sentTo = payload.sessionId; return sendResult('owned-create') })
  const ownedCreation = deferred<SessionMeta>()
  window.ycode.session.create = async () => ownedCreation.promise
  const ownedSend = useChatStore.getState().send('original draft')
  await useChatStore.getState().selectSession('other-session')
  ownedCreation.resolve({ ...session, id: 'owned-session' })
  await ownedSend
  assert.equal(sentTo, 'owned-session', 'pending creation keeps the send bound to its own session')
  assert.equal(useChatStore.getState().currentSessionId, 'other-session')
  resetStore()
  installFakeWindow(async () => sendResult('unused'))
  window.ycode.session.list = async () => { throw new Error('corrupt session store') }
  await useChatStore.getState().init()
  assert.equal(useChatStore.getState().error, 'corrupt session store', 'load failure is visible instead of silently creating an empty session')
  console.log('chat-store-regression: all checks passed')
}

void main()