refactor: Refactor initialization and token management

Author: ericc-ch

src/lib/cli.ts

@@ -7,23 +7,12 @@ const args = {
     default: false,
     description: "Show this help message",
   },
-  "emulate-streaming": {
-    type: "boolean",
-    default: false,
-    description: "Emulate streaming response for chat completions",
-  },
   port: {
     alias: "p",
     type: "string",
     default: "4141",
     description: "Port to listen on",
   },
-  logs: {
-    type: "boolean",
-    default: false,
-    description:
-      "Write logs to the app directory. Only works with emulate-streaming",
-  },
 } satisfies ArgsDef
 
 export async function getOptions() {

src/lib/initialization.ts

@@ -4,127 +4,72 @@ import { FetchError } from "ofetch"
 
 import { PATHS } from "~/lib/paths"
 import { getGitHubUser } from "~/services/github/get-user/service"
-
-import type { getOptions } from "./cli"
+import { tokenService } from "~/services/token-service"
 
 import { getModels } from "../services/copilot/get-models/service"
-import { getCopilotToken } from "../services/copilot/get-token/copilot-token"
 import { getGitHubToken } from "../services/github/get-token/service"
-import { initializeLogger } from "./logger"
-import { getGithubToken, saveGithubToken, TOKENS } from "./tokens"
 
-interface InitStep {
-  name: string
-  run: () => Promise<void> | void
+// Extract to individual functions for each initialization step
+async function initializeAppDirectory(): Promise<void> {
+  await fs.mkdir(PATHS.APP_DIR, { recursive: true })
+  await fs.writeFile(PATHS.GITHUB_TOKEN_PATH, "", { flag: "a" })
 }
 
-const initSteps: Array<InitStep> = [
-  {
-    name: "App directory",
-    run: async () => {
-      await fs.mkdir(PATHS.APP_DIR, { recursive: true })
-      await fs.writeFile(PATHS.GITHUB_TOKEN_PATH, "", { flag: "a" })
-    },
-  },
-  {
-    name: "GitHub authentication",
-    run: async () => {
-      TOKENS.GITHUB_TOKEN = await getGithubToken()
+async function initializeGithubAuthentication(): Promise<void> {
+  const githubToken = await tokenService.getGithubToken()
 
-      try {
-        await logUser()
-      } catch (error) {
-        if (!(error instanceof FetchError)) throw error
-        if (error.statusCode !== 401) {
-          consola.error("Authentication error:", {
-            error,
-            request: error.request,
-            options: error.options,
-            response: error.response,
-            data: error.response?._data as Record<string, unknown>,
-          })
-          throw error
-        }
+  try {
+    if (githubToken) {
+      // Set token in the service so github fetcher can use it
+      await tokenService.setGithubToken(githubToken)
+      await logUser()
+    } else {
+      throw new Error("No GitHub token available")
+    }
+  } catch (error) {
+    if (error instanceof FetchError && error.statusCode !== 401) {
+      consola.error("Authentication error:", {
+        error,
+        request: error.request,
+        options: error.options,
+        response: error.response,
+        data: error.response?._data as Record<string, unknown>,
+      })
+      throw error
+    }
 
-        consola.info("Not logged in, getting new access token")
-        TOKENS.GITHUB_TOKEN = await initializeGithubToken()
-        await logUser()
-      }
-    },
-  },
-  {
-    name: "Copilot token",
-    run: async () => {
-      const { token, refresh_in } = await getCopilotToken()
-      TOKENS.COPILOT_TOKEN = token
+    consola.info("Not logged in, getting new access token")
+    const newToken = await initializeGithubToken()
+    await tokenService.setGithubToken(newToken)
+    await logUser()
+  }
+}
 
-      const refreshInterval = (refresh_in - 60) * 1000
-      setInterval(async () => {
-        consola.start("Refreshing copilot token")
-        const { token: newToken } = await getCopilotToken()
-        TOKENS.COPILOT_TOKEN = newToken
-      }, refreshInterval)
-    },
-  },
-  {
-    name: "Model information",
-    run: async () => {
-      const models = await getModels()
-      consola.info(
-        `Available models: \n${models.data.map((model) => `- ${model.id}`).join("\n")}`,
-      )
-    },
-  },
-]
+async function initializeCopilotToken(): Promise<void> {
+  await tokenService.initCopilotToken()
+}
+
+async function initializeModelInformation(): Promise<void> {
+  const models = await getModels()
+  consola.info(
+    `Available models: \n${models.data.map((model) => `- ${model.id}`).join("\n")}`,
+  )
+}
 
 async function initializeGithubToken() {
   consola.start("Getting GitHub device code")
-  const token = await getGitHubToken()
-  await saveGithubToken(token.access_token)
-  return token.access_token
+  const tokenResponse = await getGitHubToken()
+  return tokenResponse.access_token
 }
 
 async function logUser() {
   const user = await getGitHubUser()
   consola.info(`Logged in as ${JSON.stringify(user.login)}`)
 }
 
-import { configManager } from "./config"
-import { initializePort } from "./port"
-
-export async function initializeApp(
-  options: Awaited<ReturnType<typeof getOptions>>,
-) {
-  configManager.setConfig({
-    EMULATE_STREAMING: options["emulate-streaming"],
-    LOGGING_ENABLED: options.logs,
-  })
-
-  // Get available port, trying the CLI option first
-  const port = await initializePort()
-
-  // Initialize logger if enabled
-  await initializeLogger()
-
-  await initialize()
-
-  const serverUrl = `http://localhost:${port}`
-  consola.success(`Server started at ${serverUrl}`)
-
-  return {
-    port,
-  }
-}
-
-async function initialize() {
-  for (const step of initSteps) {
-    try {
-      consola.start(`Initializing ${step.name}...`)
-      await step.run()
-      consola.success(`${step.name} initialized`)
-    } catch (error) {
-      consola.error(`Failed to initialize ${step.name}:`, error)
-      throw error
-    }
-  }
+export async function initializeApp() {
+  await initializeAppDirectory()
+  await initializeGithubAuthentication()
+  await initializeCopilotToken()
+  await initializeModelInformation()
 }

src/lib/logger.ts

@@ -3,13 +3,12 @@ import { getPort } from "get-port-please"
 
 import { configManager } from "./config"
 
-export async function initializePort(): Promise<number> {
+export async function initializePort(requestedPort?: number): Promise<number> {
   const config = configManager.getConfig()
-  const requestedPort = config.PORT
 
   const port = await getPort({
     name: "copilot-api",
-    port: config.PORT,
+    port: requestedPort ?? config.PORT,
     portRange: config.PORT_RANGE,
     random: false,
   })

src/lib/port.ts

@@ -1,16 +1,37 @@
 #!/usr/bin/env node
 
+import { defineCommand, runMain } from "citty"
+import consola from "consola"
 import { serve, type ServerHandler } from "srvx"
 
-import { getOptions } from "./lib/cli"
 import { initializeApp } from "./lib/initialization"
+import { initializePort } from "./lib/port"
 import { server } from "./server"
 
-const options = await getOptions()
+const main = defineCommand({
+  args: {
+    port: {
+      alias: "p",
+      type: "string",
+      default: "4141",
+      description: "Port to listen on",
+    },
+  },
+  async run({ args }) {
+    const portInt = parseInt(args.port, 10)
 
-const { port } = await initializeApp(options)
+    const port = await initializePort(portInt)
 
-serve({
-  fetch: server.fetch as ServerHandler,
-  port,
+    await initializeApp()
+
+    const serverUrl = `http://localhost:${port}`
+    consola.success(`Server started at ${serverUrl}`)
+
+    serve({
+      fetch: server.fetch as ServerHandler,
+      port,
+    })
+  },
 })
+
+await runMain(main)

src/main.ts

@@ -6,14 +6,17 @@ import {
   GITHUB_API_CONFIG,
   GITHUB_WEB_API_CONFIG,
 } from "~/lib/constants"
-import { TOKENS } from "~/lib/tokens"
+import { tokenService } from "~/services/token-service"
 
 export const copilot = ofetch.create({
   baseURL: COPILOT_API_CONFIG.baseURL,
   headers: COPILOT_API_CONFIG.headers,
 
   onRequest({ options }) {
-    options.headers.set("authorization", `Bearer ${TOKENS.COPILOT_TOKEN}`)
+    options.headers.set(
+      "authorization",
+      `Bearer ${tokenService.getCopilotToken()}`,
+    )
   },
 
   onRequestError({ error, options }) {
@@ -38,8 +41,9 @@ export const copilot = ofetch.create({
 export const github = ofetch.create({
   baseURL: GITHUB_API_CONFIG.baseURL,
 
-  onRequest({ options }) {
-    options.headers.set("authorization", `token ${TOKENS.GITHUB_TOKEN}`)
+  async onRequest({ options }) {
+    const token = await tokenService.getGithubToken()
+    options.headers.set("authorization", `token ${token}`)
   },
 })
 

src/services/api-instance.ts

@@ -0,0 +1,65 @@
+import consola from "consola"
+import fs from "node:fs/promises"
+
+import { PATHS } from "~/lib/paths"
+import { getCopilotToken } from "~/services/copilot/get-token/copilot-token"
+
+// Simple token manager with basic encapsulation
+export const tokenService = {
+  // Private token storage
+  _tokens: {
+    github: undefined as string | undefined,
+    copilot: undefined as string | undefined,
+  },
+
+  // Get GitHub token
+  async getGithubToken(): Promise<string | undefined> {
+    if (!this._tokens.github) {
+      try {
+        this._tokens.github = await fs.readFile(
+          PATHS.GITHUB_TOKEN_PATH,
+          "utf-8",
+        )
+      } catch (error) {
+        consola.warn("Failed to load GitHub token", error)
+      }
+    }
+
+    return this._tokens.github
+  },
+
+  // Set GitHub token
+  async setGithubToken(token: string): Promise<void> {
+    this._tokens.github = token
+    await fs.writeFile(PATHS.GITHUB_TOKEN_PATH, token)
+  },
+
+  // Get Copilot token
+  getCopilotToken(): string | undefined {
+    return this._tokens.copilot
+  },
+
+  // Set Copilot token
+  setCopilotToken(token: string): void {
+    this._tokens.copilot = token
+  },
+
+  // Initialize Copilot token with auto-refresh
+  async initCopilotToken(): Promise<void> {
+    const { token, refresh_in } = await getCopilotToken()
+    this.setCopilotToken(token)
+
+    // Set up refresh timer
+    const refreshInterval = (refresh_in - 60) * 1000
+    setInterval(async () => {
+      consola.start("Refreshing Copilot token")
+      try {
+        const { token: newToken } = await getCopilotToken()
+        this.setCopilotToken(newToken)
+        consola.success("Copilot token refreshed")
+      } catch (error) {
+        consola.error("Failed to refresh Copilot token:", error)
+      }
+    }, refreshInterval)
+  },
+}