FF: fix crashes on undefined tokens

src/modules/aix/server/api/aix.wiretypes.ts

@@ -412,6 +412,10 @@ export namespace AixWire_API_ChatGenerate {
 /**
  * This is the protocol for both the control objects sent by the tRPC streaming procedures,
  * and the thePartTransmitter/PartReassembler.
+ *
+ * VITAL: when transmitting anything that's "undefined", leave it out of the
+ * object rather than setting it as 'undefined' as 'superjson' will mess it up
+ * and tRPC decoding will be broken (very important!)
  */
 export namespace AixAPI_Particles {
 
@@ -438,19 +442,19 @@ export namespace AixAPI_Particles {
 
   export type ParticleOp =
     | { p: 't_', t: string /* incremental text, despite not having the 'i_' prefix for brevity */ }
-    | { p: 'inline-image', mimeType: string, i_b64?: string }
+    | { p: 'inline-image', mimeType: string, i_b64?: string /* never undefined */ }
     | { p: 'ii_', i_b64: string }
-    | { p: 'inline-doc', type: string, ref: string, l1Title: string, i_text?: string }
+    | { p: 'inline-doc', type: string, ref: string, l1Title: string, i_text?: string /* never undefined */ }
     | { p: 'id_', i_text: string }
-    | { p: 'function-call', id: string, name: string, i_args?: string }
+    | { p: 'function-call', id: string, name: string, i_args?: string /* never undefined */ }
     | { p: 'fc_', i_args: string }
     | { p: 'code-call', id: string, language: string, code: string }
-    | { p: 'code-response', id: string, output: string, error?: string }
+    | { p: 'code-response', id: string, output: string, error?: string /* never undefined */ }
 
+  // NOTE: see the Vital notice
   export type ChatGenerateCounts = {
     chatIn?: number,
     chatOut?: number,
-    chatTotal?: number,
     chatOutRate?: number,
     chatTimeInner?: number,
   };

src/modules/aix/server/dispatch/chatGenerate/ChatGenerateTransmitter.ts

@@ -239,8 +239,15 @@ export class ChatGenerateTransmitter implements IPartTransmitter {
   /** Update the counters, sent twice (after the first call, and then at the end of the transmission) */
   setCounters(counts: AixAPI_Particles.ChatGenerateCounts) {
     if (!this.accCounts)
-      this.accCounts = {};
-    Object.assign(this.accCounts, counts);
+      this.accCounts = {} as AixAPI_Particles.ChatGenerateCounts;
+
+    // similar to Object.assign, but takes care of removing the "undefined" entries
+    for (const key in counts) {
+      const value = (counts as any)[key] as number | undefined;
+      if (value !== undefined)
+        (this.accCounts as any)[key] = value;
+    }
+
     this.freshCounts = true;
   }
 

src/modules/aix/server/dispatch/chatGenerate/parsers/gemini.parser.ts

@@ -118,7 +118,10 @@ export function createGeminiGenerateContentResponseParser(modelId: string): Chat
 
     // -> Stats
     if (generationChunk.usageMetadata)
-      pt.setCounters({ chatIn: generationChunk.usageMetadata.promptTokenCount, chatOut: generationChunk.usageMetadata.candidatesTokenCount });
+      pt.setCounters({
+        chatIn: generationChunk.usageMetadata.promptTokenCount,
+        chatOut: generationChunk.usageMetadata.candidatesTokenCount,
+      });
 
   };
 }

src/modules/aix/server/dispatch/chatGenerate/parsers/openai.parser.ts

@@ -82,7 +82,10 @@ export function createOpenAIChatCompletionsChunkParser(): ChatGenerateParseFunct
     // -> Stats
     if (json.usage) {
       if (json.usage.completion_tokens !== undefined)
-        pt.setCounters({ chatIn: json.usage.prompt_tokens || -1, chatOut: json.usage.completion_tokens });
+        pt.setCounters({
+          chatIn: json.usage.prompt_tokens || -1,
+          chatOut: json.usage.completion_tokens,
+        });
 
       // [OpenAI] Expected correct case: the last object has usage, but an empty choices array
       if (!json.choices.length)
@@ -196,7 +199,10 @@ export function createOpenAIChatCompletionsParserNS(): ChatGenerateParseFunction
 
     // -> Stats
     if (json.usage)
-      pt.setCounters({ chatIn: json.usage.prompt_tokens, chatOut: json.usage.completion_tokens, chatTotal: json.usage.total_tokens });
+      pt.setCounters({
+        chatIn: json.usage.prompt_tokens,
+        chatOut: json.usage.completion_tokens,
+      });
 
     // Assumption/validate: expect 1 completion, or stop
     if (json.choices.length !== 1)

src/modules/aix/server/dispatch/wiretypes/gemini.wiretypes.ts

@@ -469,7 +469,7 @@ export namespace GeminiWire_API_Generate_Content {
   const UsageMetadata_schema = z.object({
     promptTokenCount: z.number(),
     candidatesTokenCount: z.number().optional(), // .optional: in case the first message is 'RECITATION' there could be no output token count
-    totalTokenCount: z.number(),
+    // totalTokenCount: z.number(),
     // cachedContentTokenCount: z.number().optional(), // Not supported for now, hence disabled
   });