feat: token counting and manual approval

Author: ericc-ch

src/lib/state.ts

@@ -8,12 +8,10 @@ export interface State {
   models?: ModelsResponse
   vsCodeVersion?: string
 
-  estimateToken: boolean
   manualApprove: boolean
 }
 
 export const state: State = {
   accountType: "individual",
-  estimateToken: true,
   manualApprove: false,
 }

src/lib/tokenizer.ts

@@ -2,7 +2,7 @@ import { countTokens } from "gpt-tokenizer/model/gpt-4o"
 
 import type { Message } from "~/services/copilot/create-chat-completions"
 
-export const getTokenLength = (messages: Array<Message>) => {
+export const getTokenCount = (messages: Array<Message>) => {
   const input = messages.filter((m) => m.role !== "assistant")
   const output = messages.filter((m) => m.role === "assistant")
 

src/main.ts

@@ -15,6 +15,7 @@ interface RunServerOptions {
   port: number
   verbose: boolean
   business: boolean
+  manual: boolean
 }
 
 export async function runServer(options: RunServerOptions): Promise<void> {
@@ -28,6 +29,8 @@ export async function runServer(options: RunServerOptions): Promise<void> {
     consola.info("Using business plan GitHub account")
   }
 
+  state.manualApprove = options.manual
+
   await ensurePaths()
   await cacheVSCodeVersion()
   await setupGitHubToken()
@@ -62,6 +65,11 @@ const main = defineCommand({
       default: false,
       description: "Use a business plan GitHub Account",
     },
+    manual: {
+      type: "boolean",
+      default: false,
+      description: "Enable manual request approval",
+    },
   },
   run({ args }) {
     const port = Number.parseInt(args.port, 10)
@@ -70,6 +78,7 @@ const main = defineCommand({
       port,
       verbose: args.verbose,
       business: args.business,
+      manual: args.manual,
     })
   },
 })

src/routes/chat-completions/handler.ts

@@ -1,9 +1,11 @@
 import type { Context } from "hono"
 
+import consola from "consola"
 import { streamSSE, type SSEMessage } from "hono/streaming"
 
 import { isNullish } from "~/lib/is-nullish"
 import { state } from "~/lib/state"
+import { getTokenCount } from "~/lib/tokenizer"
 import {
   createChatCompletions,
   type ChatCompletionResponse,
@@ -13,6 +15,16 @@ import {
 export async function handleCompletion(c: Context) {
   let payload = await c.req.json<ChatCompletionsPayload>()
 
+  consola.info("Current token count:", getTokenCount(payload.messages))
+
+  if (state.manualApprove) {
+    const response = await consola.prompt(`Accept incoming request?`, {
+      type: "confirm",
+    })
+
+    if (!response) throw new Error("Request cancelled")
+  }
+
   if (isNullish(payload.max_tokens)) {
     const selectedModel = state.models?.data.find(
       (model) => model.id === payload.model,