/ src / services / vcr.ts
vcr.ts
  1  import { createHash, type UUID } from 'crypto'
  2  import { mkdirSync, readFileSync, writeFileSync } from 'fs'
  3  import { dirname } from 'path'
  4  import type { AssistantMessage, UserMessage } from '../query.js'
  5  import { existsSync } from 'fs'
  6  import { env } from '../utils/env.js'
  7  import { getCwd } from '../utils/state.js'
  8  import * as path from 'path'
  9  import { mapValues } from 'lodash-es'
 10  import type { ContentBlock } from '@anthropic-ai/sdk/resources/index.mjs'
 11  
 12  export async function withVCR(
 13    messages: (UserMessage | AssistantMessage)[],
 14    f: () => Promise<AssistantMessage>,
 15  ): Promise<AssistantMessage> {
 16    if (process.env.NODE_ENV !== 'test') {
 17      return await f()
 18    }
 19  
 20    const dehydratedInput = mapMessages(
 21      messages.map(_ => _.message.content),
 22      dehydrateValue,
 23    )
 24    const filename = `./fixtures/${dehydratedInput.map(_ => createHash('sha1').update(JSON.stringify(_)).digest('hex').slice(0, 6)).join('-')}.json`
 25  
 26    // Fetch cached fixture
 27    if (existsSync(filename)) {
 28      const cached = JSON.parse(readFileSync(filename, 'utf-8'))
 29      return mapAssistantMessage(cached.output, hydrateValue)
 30    }
 31  
 32    if (env.isCI) {
 33      console.warn(
 34        `Anthropic API fixture missing. Re-run npm test locally, then commit the result. ${JSON.stringify({ input: dehydratedInput }, null, 2)}`,
 35      )
 36    }
 37  
 38    // Create & write new fixture
 39    const result = await f()
 40    if (env.isCI) {
 41      return result
 42    }
 43  
 44    if (!existsSync(dirname(filename))) {
 45      mkdirSync(dirname(filename), { recursive: true })
 46    }
 47    writeFileSync(
 48      filename,
 49      JSON.stringify(
 50        {
 51          input: dehydratedInput,
 52          output: mapAssistantMessage(result, dehydrateValue),
 53        },
 54        null,
 55        2,
 56      ),
 57    )
 58    return result
 59  }
 60  
 61  function mapMessages(
 62    messages: (UserMessage | AssistantMessage)['message']['content'][],
 63    f: (s: unknown) => unknown,
 64  ): (UserMessage | AssistantMessage)['message']['content'][] {
 65    return messages.map(_ => {
 66      if (typeof _ === 'string') {
 67        return f(_)
 68      }
 69      return _.map(_ => {
 70        switch (_.type) {
 71          case 'tool_result':
 72            if (typeof _.content === 'string') {
 73              return { ..._, content: f(_.content) }
 74            }
 75            if (Array.isArray(_.content)) {
 76              return {
 77                ..._,
 78                content: _.content.map(_ => {
 79                  switch (_.type) {
 80                    case 'text':
 81                      return { ..._, text: f(_.text) }
 82                    case 'image':
 83                      return _
 84                  }
 85                }),
 86              }
 87            }
 88            return _
 89          case 'text':
 90            return { ..._, text: f(_.text) }
 91          case 'tool_use':
 92            return {
 93              ..._,
 94              input: mapValues(_.input as Record<string, unknown>, f),
 95            }
 96          case 'image':
 97            return _
 98        }
 99      })
100    }) as (UserMessage | AssistantMessage)['message']['content'][]
101  }
102  
103  function mapAssistantMessage(
104    message: AssistantMessage,
105    f: (s: unknown) => unknown,
106  ): AssistantMessage {
107    return {
108      durationMs: 'DURATION' as unknown as number,
109      costUSD: 'COST' as unknown as number,
110      uuid: 'UUID' as unknown as UUID,
111      message: {
112        ...message.message,
113        content: message.message.content
114          .map(_ => {
115            switch (_.type) {
116              case 'text':
117                return {
118                  ..._,
119                  text: f(_.text) as string,
120                  citations: _.citations || [],
121                } // Ensure citations
122              case 'tool_use':
123                return {
124                  ..._,
125                  input: mapValues(_.input as Record<string, unknown>, f),
126                }
127              default:
128                return _ // Handle other block types unchanged
129            }
130          })
131          .filter(Boolean) as ContentBlock[],
132      },
133      type: 'assistant',
134    }
135  }
136  
137  function dehydrateValue(s: unknown): unknown {
138    if (typeof s !== 'string') {
139      return s
140    }
141    const s1 = s
142      .replace(/num_files="\d+"/g, 'num_files="[NUM]"')
143      .replace(/duration_ms="\d+"/g, 'duration_ms="[DURATION]"')
144      .replace(/cost_usd="\d+"/g, 'cost_usd="[COST]"')
145      .replace(/\//g, path.sep)
146      .replaceAll(getCwd(), '[CWD]')
147    if (s1.includes('Files modified by user:')) {
148      return 'Files modified by user: [FILES]'
149    }
150    return s1
151  }
152  
153  function hydrateValue(s: unknown): unknown {
154    if (typeof s !== 'string') {
155      return s
156    }
157    return s
158      .replaceAll('[NUM]', '1')
159      .replaceAll('[DURATION]', '100')
160      .replaceAll('[CWD]', getCwd())
161  }