feat: switch to undici and add ava unit tests

src/example.ts → demo.ts

 import dotenv from 'dotenv-safe'
 import { oraPromise } from 'ora'
 
-import { ChatGPTAPI } from './chatgpt-api'
+import { ChatGPTAPI } from './src'
 
 dotenv.config()
 

package.json

@@ -32,14 +32,15 @@
     "prepare": "husky install",
     "pre-commit": "lint-staged",
     "test": "run-p test:*",
+    "test:unit": "ava",
     "test:prettier": "prettier '**/*.{js,jsx,ts,tsx}' --check"
   },
   "dependencies": {
     "eventsource-parser": "^0.0.5",
     "expiry-map": "^2.0.0",
-    "node-fetch": "2",
     "remark": "^14.0.2",
     "strip-markdown": "^5.0.0",
+    "undici": "^5.13.0",
     "uuid": "^9.0.0"
   },
   "devDependencies": {
@@ -47,6 +48,7 @@
     "@types/node": "^18.11.9",
     "@types/node-fetch": "2",
     "@types/uuid": "^9.0.0",
+    "ava": "^5.1.0",
     "del-cli": "^5.0.0",
     "dotenv-safe": "^8.2.0",
     "husky": "^8.0.2",
@@ -65,6 +67,14 @@
       "prettier --write"
     ]
   },
+  "ava": {
+    "extensions": {
+      "ts": "module"
+    },
+    "nodeArguments": [
+      "--loader=tsx"
+    ]
+  },
   "keywords": [
     "openai",
     "chatgpt",

src/chatgpt-api.test.ts

+import test from 'ava'
+import dotenv from 'dotenv-safe'
+
+import { ChatGPTAPI } from './chatgpt-api'
+
+dotenv.config()
+
+const isCI = !!process.env.CI
+
+test('ChatGPTAPI invalid session token', async (t) => {
+  t.throws(() => new ChatGPTAPI({ sessionToken: null }), {
+    message: 'ChatGPT invalid session token'
+  })
+
+  await t.throwsAsync(
+    async () => {
+      const chatgpt = new ChatGPTAPI({ sessionToken: 'invalid' })
+      await chatgpt.ensureAuth()
+    },
+    {
+      message: 'ChatGPT failed to refresh auth token. Error: Unauthorized'
+    }
+  )
+})
+
+test('ChatGPTAPI expired session token', async (t) => {
+  if (isCI) {
+    return
+  }
+
+  const expiredSessionToken = process.env.TEST_EXPIRED_SESSION_TOKEN
+
+  await t.throwsAsync(
+    async () => {
+      const chatgpt = new ChatGPTAPI({ sessionToken: expiredSessionToken })
+      await chatgpt.ensureAuth()
+    },
+    {
+      message:
+        'ChatGPT failed to refresh auth token. Error: session token has expired'
+    }
+  )
+})
+
+test('ChatGPTAPI valid session token', async (t) => {
+  if (!isCI) {
+    t.timeout(2 * 60 * 1000) // 2 minutes
+  }
+
+  t.notThrows(
+    () => new ChatGPTAPI({ sessionToken: 'fake valid session token' })
+  )
+
+  await t.notThrowsAsync(
+    (async () => {
+      const api = new ChatGPTAPI({ sessionToken: process.env.SESSION_TOKEN })
+
+      // Don't make any real API calls using our session token if we're running on CI
+      if (!isCI) {
+        await api.ensureAuth()
+        const response = await api.sendMessage('test')
+        console.log('chatgpt response', response)
+
+        t.truthy(response)
+        t.is(typeof response, 'string')
+      }
+    })()
+  )
+})

src/chatgpt-api.ts

-import { createParser } from 'eventsource-parser'
 import ExpiryMap from 'expiry-map'
-import fetch from 'node-fetch'
 import { v4 as uuidv4 } from 'uuid'
 
 import * as types from './types'
+import { fetch } from './fetch'
+import { fetchSSE } from './fetch-sse'
 import { markdownToText } from './utils'
 
 const KEY_ACCESS_TOKEN = 'accessToken'
@@ -119,7 +119,7 @@ export class ChatGPTAPI {
     let response = ''
 
     return new Promise((resolve, reject) => {
-      this._fetchSSE(url, {
+      fetchSSE(url, {
         method: 'POST',
         headers: {
           Authorization: `Bearer ${accessToken}`,
@@ -179,34 +179,22 @@ export class ChatGPTAPI {
       const accessToken = res?.accessToken
 
       if (!accessToken) {
-        console.warn('no auth token', res)
         throw new Error('Unauthorized')
       }
 
+      const error = res?.error
+      if (error) {
+        if (error === 'RefreshAccessTokenError') {
+          throw new Error('session token has expired')
+        } else {
+          throw new Error(error)
+        }
+      }
+
       this._accessTokenCache.set(KEY_ACCESS_TOKEN, accessToken)
       return accessToken
     } catch (err: any) {
-      throw new Error(`ChatGPT failed to refresh auth token: ${err.toString()}`)
+      throw new Error(`ChatGPT failed to refresh auth token. ${err.toString()}`)
     }
   }
-
-  protected async _fetchSSE(
-    url: string,
-    options: Parameters<typeof fetch>[1] & { onMessage: (data: string) => void }
-  ) {
-    const { onMessage, ...fetchOptions } = options
-    const resp = await fetch(url, fetchOptions)
-    const parser = createParser((event) => {
-      if (event.type === 'event') {
-        onMessage(event.data)
-      }
-    })
-
-    resp.body.on('readable', () => {
-      let chunk: string | Buffer
-      while (null !== (chunk = resp.body.read())) {
-        parser.feed(chunk.toString())
-      }
-    })
-  }
 }

src/fetch-sse.ts

+import { createParser } from 'eventsource-parser'
+
+import { fetch } from './fetch'
+import { streamAsyncIterable } from './stream-async-iterable'
+
+export async function fetchSSE(
+  url: string,
+  options: Parameters<typeof fetch>[1] & { onMessage: (data: string) => void }
+) {
+  const { onMessage, ...fetchOptions } = options
+  const resp = await fetch(url, fetchOptions)
+  const parser = createParser((event) => {
+    if (event.type === 'event') {
+      onMessage(event.data)
+    }
+  })
+
+  for await (const chunk of streamAsyncIterable(resp.body)) {
+    const str = new TextDecoder().decode(chunk)
+    parser.feed(str)
+  }
+}

src/fetch.ts

+import { fetch } from 'undici'
+
+export { fetch }

src/stream-async-iterable.ts

+import { type ReadableStream } from 'stream/web'
+
+export async function* streamAsyncIterable(stream: ReadableStream) {
+  const reader = stream.getReader()
+  try {
+    while (true) {
+      const { done, value } = await reader.read()
+      if (done) {
+        return
+      }
+      yield value
+    }
+  } finally {
+    reader.releaseLock()
+  }
+}

src/types.ts

@@ -7,7 +7,7 @@ export type Role = 'user' | 'assistant'
  */
 export type SessionResult = {
   /**
-   * Object of the current user
+   * Authenticated user
    */
   user: User
 
@@ -20,6 +20,11 @@ export type SessionResult = {
    * The access token
    */
   accessToken: string
+
+  /**
+   * If there was an error associated with this request
+   */
+  error?: string | null
 }
 
 export type User = {