From 59847b17d683f718ef793aa5d1fca280f22afa86 Mon Sep 17 00:00:00 2001
From: Imanol Maiztegui <imanol.mzd@gmail.com>
Date: Mon, 2 Mar 2026 18:59:49 +0100
Subject: [PATCH] fix: type FIM stream response instead of z.any()/unknown

---
 packages/kilo-gateway/src/server/routes.ts    | 23 ++++++++++++++++++-
 .../kilo-vscode/src/kilo-provider-utils.ts    |  2 +-
 .../autocomplete/AutocompleteModel.ts         | 15 ++++--------
 packages/sdk/js/src/v2/gen/types.gen.ts       | 15 +++++++++++-
 4 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/packages/kilo-gateway/src/server/routes.ts b/packages/kilo-gateway/src/server/routes.ts
index 2c162db9a40..2b25b4ba790 100644
--- a/packages/kilo-gateway/src/server/routes.ts
+++ b/packages/kilo-gateway/src/server/routes.ts
@@ -96,6 +96,27 @@ export function createKiloRoutes(deps: KiloRoutesDeps) {
     currentOrgId: z.string().nullable(),
   })
 
+  const FimStreamChunk = z.object({
+    choices: z
+      .array(
+        z.object({
+          delta: z
+            .object({
+              content: z.string().optional(),
+            })
+            .optional(),
+        }),
+      )
+      .optional(),
+    usage: z
+      .object({
+        prompt_tokens: z.number().optional(),
+        completion_tokens: z.number().optional(),
+      })
+      .optional(),
+    cost: z.number().optional(),
+  })
+
   return new Hono()
     .get(
       "/profile",
@@ -193,7 +214,7 @@ export function createKiloRoutes(deps: KiloRoutesDeps) {
             description: "Streaming FIM completion response",
             content: {
               "text/event-stream": {
-                schema: resolver(z.any()),
+                schema: resolver(FimStreamChunk),
               },
             },
           },
diff --git a/packages/kilo-vscode/src/kilo-provider-utils.ts b/packages/kilo-vscode/src/kilo-provider-utils.ts
index 2b25c57e6ae..261ca3356a7 100644
--- a/packages/kilo-vscode/src/kilo-provider-utils.ts
+++ b/packages/kilo-vscode/src/kilo-provider-utils.ts
@@ -1,4 +1,4 @@
-import type { Session, Agent, Provider, Event, ProviderListResponse } from "@kilocode/sdk/v2/client"
+import type { Session, Agent, Event, ProviderListResponse } from "@kilocode/sdk/v2/client"
 
 /** A single provider entry as returned by the /provider list endpoint. */
 export type ProviderInfo = ProviderListResponse["all"][number]
diff --git a/packages/kilo-vscode/src/services/autocomplete/AutocompleteModel.ts b/packages/kilo-vscode/src/services/autocomplete/AutocompleteModel.ts
index d123cd9a44f..71adb447a9b 100644
--- a/packages/kilo-vscode/src/services/autocomplete/AutocompleteModel.ts
+++ b/packages/kilo-vscode/src/services/autocomplete/AutocompleteModel.ts
@@ -95,18 +95,13 @@ export class AutocompleteModel {
     )
 
     for await (const chunk of stream) {
-      const event = chunk as {
-        choices?: Array<{ delta?: { content?: string } }>
-        usage?: { prompt_tokens?: number; completion_tokens?: number }
-        cost?: number
-      }
-      const content = event.choices?.[0]?.delta?.content
+      const content = chunk.choices?.[0]?.delta?.content
       if (content) onChunk(content)
-      if (event.usage) {
-        inputTokens = event.usage.prompt_tokens ?? 0
-        outputTokens = event.usage.completion_tokens ?? 0
+      if (chunk.usage) {
+        inputTokens = chunk.usage.prompt_tokens ?? 0
+        outputTokens = chunk.usage.completion_tokens ?? 0
       }
-      if (event.cost !== undefined) cost = event.cost
+      if (chunk.cost !== undefined) cost = chunk.cost
     }
 
     return {
diff --git a/packages/sdk/js/src/v2/gen/types.gen.ts b/packages/sdk/js/src/v2/gen/types.gen.ts
index 0a39735a74e..debcdfcc475 100644
--- a/packages/sdk/js/src/v2/gen/types.gen.ts
+++ b/packages/sdk/js/src/v2/gen/types.gen.ts
@@ -4578,9 +4578,22 @@ export type KiloFimResponses = {
   /**
    * Streaming FIM completion response
    */
-  200: unknown
+  200: {
+    choices?: Array<{
+      delta?: {
+        content?: string
+      }
+    }>
+    usage?: {
+      prompt_tokens?: number
+      completion_tokens?: number
+    }
+    cost?: number
+  }
 }
 
+export type KiloFimResponse = KiloFimResponses[keyof KiloFimResponses]
+
 export type KiloNotificationsData = {
   body?: never
   path?: never