refactor: Improve error handling and token management

Author: ericc-ch

README.md

@@ -40,18 +40,17 @@ npx copilot-api --port 8080
 
 ## Command Line Options
 
-The command accepts several command line options:
+The following command line options are available:
 
-| Option        | Description                          | Default |
-| ------------- | ------------------------------------ | ------- |
-| --port, -p    | Port to listen on                    | 4141    |
-| --verbose, -v | Enable verbose logging               | false   |
-| --log-file    | File to log request/response details | -       |
+| Option        | Description              | Default |
+| ------------- | ------------------------ | ------- |
+| --port, -p    | Port to listen on        | 4141    |
+| --verbose, -v | Enable verbose logging   | false   |
 
-Example with options:
+Example usage:
 
 ```sh
-npx copilot-api@latest --port 8080 --verbose --log-file copilot-api.txt
+npx copilot-api@latest --port 8080 --verbose
 ```
 
 ## Running from Source
@@ -70,15 +69,18 @@ bun run dev
 bun run start
 ```
 
-## Tips to not hit the rate limit
+## Usage Tips
 
-- Use a free model from free provider like Gemini/Mistral/Openrouter for the weak model
-- Rarely use architect mode
-- Do not enable automatic yes in aider config
-- Claude 3.7 thinking mode uses more tokens. Use it sparingly
+To avoid rate limiting and optimize your experience:
+
+- Consider using free models (e.g., Gemini, Mistral, Openrouter) as the `weak-model`
+- Use architect mode sparingly
+- Disable `yes-always` in your aider configuration
+- Be mindful that Claude 3.7 thinking mode consume more tokens
 
 ## Roadmap
 
-- Manual approval for every request
-- Rate limiting (only allow request every X seconds)
-- Token counting
+- [ ] Manual request approval system
+- [ ] Rate limiting implementation
+- [ ] Token usage tracking and monitoring
+- [ ] Enhanced error handling and recovery

src/lib/api-config.ts

@@ -1,9 +1,14 @@
 import type { State } from "./state"
 
+export const standardHeaders = () => ({
+  "content-type": "application/json",
+  accept: "application/json",
+})
+
 export const COPILOT_API_BASE_URL = "https://api.individual.githubcopilot.com"
 export const copilotHeaders = (state: State) => ({
   Authorization: `Bearer ${state.copilotToken}`,
-  "content-type": "application/json",
+  "content-type": standardHeaders()["content-type"],
   "copilot-integration-id": "vscode-chat",
   "editor-version": `vscode/${state.vsCodeVersion}`,
   "editor-plugin-version": "copilot-chat/0.24.1",
@@ -15,6 +20,7 @@ export const copilotHeaders = (state: State) => ({
 
 export const GITHUB_API_BASE_URL = "https://api.github.com"
 export const githubHeaders = (state: State) => ({
+  ...standardHeaders(),
   authorization: `token ${state.githubToken}`,
   "editor-version": `vscode/${state.vsCodeVersion}`,
   "editor-plugin-version": "copilot-chat/0.24.1",

src/lib/token.ts

@@ -7,6 +7,7 @@ import { getDeviceCode } from "~/services/github/get-device-code"
 import { getGitHubUser } from "~/services/github/get-user"
 import { pollAccessToken } from "~/services/github/poll-access-token"
 
+import { HTTPError } from "./http-error"
 import { state } from "./state"
 
 const readGithubToken = () => fs.readFile(PATHS.GITHUB_TOKEN_PATH, "utf8")
@@ -33,27 +34,38 @@ export const setupCopilotToken = async () => {
 }
 
 export async function setupGitHubToken(): Promise<void> {
-  const githubToken = await readGithubToken()
+  try {
+    const githubToken = await readGithubToken()
 
-  if (githubToken) {
-    state.githubToken = githubToken
-    await logUser()
+    if (githubToken) {
+      state.githubToken = githubToken
+      await logUser()
 
-    return
-  }
+      return
+    }
+
+    consola.info("Not logged in, getting new access token")
+    const response = await getDeviceCode()
+    consola.debug("Device code response:", response)
 
-  consola.info("Not logged in, getting new access token")
-  const response = await getDeviceCode()
+    consola.info(
+      `Please enter the code "${response.user_code}" in ${response.verification_uri}`,
+    )
 
-  consola.info(
-    `Please enter the code "${response.user_code}" in ${response.verification_uri}`,
-  )
+    const token = await pollAccessToken(response)
+    await writeGithubToken(token)
+    state.githubToken = token
 
-  const token = await pollAccessToken(response)
-  await writeGithubToken(token)
-  state.githubToken = token
+    await logUser()
+  } catch (error) {
+    if (error instanceof HTTPError) {
+      consola.error("Failed to get GitHub token:", await error.response.json())
+      throw error
+    }
 
-  await logUser()
+    consola.error("Failed to get GitHub token:", error)
+    throw error
+  }
 }
 
 async function logUser() {

src/services/copilot/get-models.ts

@@ -1,12 +1,10 @@
-import { COPILOT_API_BASE_URL } from "~/lib/api-config"
+import { COPILOT_API_BASE_URL, copilotHeaders } from "~/lib/api-config"
 import { HTTPError } from "~/lib/http-error"
 import { state } from "~/lib/state"
 
 export const getModels = async () => {
   const response = await fetch(`${COPILOT_API_BASE_URL}/models`, {
-    headers: {
-      authorization: `Bearer ${state.copilotToken}`,
-    },
+    headers: copilotHeaders(state),
   })
 
   if (!response.ok) throw new HTTPError("Failed to get models", response)

src/services/github/get-device-code.ts

@@ -2,12 +2,14 @@ import {
   GITHUB_APP_SCOPES,
   GITHUB_BASE_URL,
   GITHUB_CLIENT_ID,
+  standardHeaders,
 } from "~/lib/api-config"
 import { HTTPError } from "~/lib/http-error"
 
 export async function getDeviceCode(): Promise<DeviceCodeResponse> {
   const response = await fetch(`${GITHUB_BASE_URL}/login/device/code`, {
     method: "POST",
+    headers: standardHeaders(),
     body: JSON.stringify({
       client_id: GITHUB_CLIENT_ID,
       scope: GITHUB_APP_SCOPES,

src/services/github/get-user.ts

@@ -1,11 +1,12 @@
-import { GITHUB_API_BASE_URL } from "~/lib/api-config"
+import { GITHUB_API_BASE_URL, standardHeaders } from "~/lib/api-config"
 import { HTTPError } from "~/lib/http-error"
 import { state } from "~/lib/state"
 
 export async function getGitHubUser() {
   const response = await fetch(`${GITHUB_API_BASE_URL}/user`, {
     headers: {
       authorization: `token ${state.githubToken}`,
+      ...standardHeaders(),
     },
   })
 

src/services/github/poll-access-token.ts

@@ -1,16 +1,28 @@
-import { GITHUB_BASE_URL, GITHUB_CLIENT_ID } from "~/lib/api-config"
+import consola from "consola"
+
+import {
+  GITHUB_BASE_URL,
+  GITHUB_CLIENT_ID,
+  standardHeaders,
+} from "~/lib/api-config"
 import { sleep } from "~/lib/sleep"
 
 import type { DeviceCodeResponse } from "./get-device-code"
 
 export async function pollAccessToken(
   deviceCode: DeviceCodeResponse,
 ): Promise<string> {
+  // Interval is in seconds, we need to multiply by 1000 to get milliseconds
+  // I'm also adding another second, just to be safe
+  const sleepDuration = (deviceCode.interval + 1) * 1000
+  consola.debug(`Polling access token with interval of ${sleepDuration}ms`)
+
   while (true) {
     const response = await fetch(
       `${GITHUB_BASE_URL}/login/oauth/access_token`,
       {
         method: "POST",
+        headers: standardHeaders(),
         body: JSON.stringify({
           client_id: GITHUB_CLIENT_ID,
           device_code: deviceCode.device_code,
@@ -19,16 +31,17 @@ export async function pollAccessToken(
       },
     )
 
-    // Interval is in seconds, we need to multiply by 1000 to get milliseconds
-    // I'm also adding another second, just to be safe
-    const sleepDuration = (deviceCode.interval + 1) * 1000
-
     if (!response.ok) {
       await sleep(sleepDuration)
+      consola.error("Failed to poll access token:", await response.text())
+
       continue
     }
 
-    const { access_token } = (await response.json()) as AccessTokenResponse
+    const json = await response.json()
+    consola.debug("Polling access token response:", json)
+
+    const { access_token } = json as AccessTokenResponse
 
     if (access_token) {
       return access_token