Update aiEngine.js

aiEngine.js

@@ -38,7 +38,13 @@ export const AIEngine = {
         config,
         contents,
       });
-      return response.text; 
+      
+      // Return both text and usage metadata
+      return { 
+        text: response.text, 
+        usage: response.usageMetadata 
+      };
+
     } catch (error) {
       console.error("PM AI Error:", error);
       throw error;
@@ -87,7 +93,13 @@ export const AIEngine = {
         config,
         contents,
       });
-      return response.text;
+
+      // Return both text and usage metadata
+      return { 
+        text: response.text, 
+        usage: response.usageMetadata 
+      };
+
     } catch (error) {
       console.error("Worker AI Error:", error);
       throw error;
@@ -114,9 +126,17 @@ export const AIEngine = {
       });
       
       const text = response.text;
-      return JSON.parse(text);
+      const parsed = JSON.parse(text);
+      
+      // Attach usage to the JSON object
+      return { 
+        ...parsed, 
+        usage: response.usageMetadata 
+      };
+
     } catch (e) {
       console.error("Analyst Error:", e);
+      // On failure, we don't return usage, so no charge applies
       return { status: "ACCEPTED", questions: [{ id: "fallback", label: "Please describe the core gameplay loop in detail.", type: "textarea" }] };
     }
   },
@@ -139,9 +159,18 @@ export const AIEngine = {
         },
         contents: [{ role: 'user', parts: [{ text: input }] }]
       });
-      return JSON.parse(response.text);
+      
+      const parsed = JSON.parse(response.text);
+      
+      // Attach usage to the JSON object
+      return { 
+        ...parsed, 
+        usage: response.usageMetadata 
+      };
+
     } catch (e) {
       console.error("Grading Error:", e);
+      // On failure, no usage returned
       return { feasibility: 80, rating: "B", title: "Untitled Project", summary: "Standard project structure detected." };
     }
   },
@@ -172,37 +201,38 @@ export const AIEngine = {
         contents,
       });
 
+      let finalDataUrl = null;
+
       for await (const chunk of response) {
         if (!chunk.candidates || !chunk.candidates[0].content || !chunk.candidates[0].content.parts) {
           continue;
         }
+        
+        // Capture image data if present
         if (chunk.candidates?.[0]?.content?.parts?.[0]?.inlineData) {
-        /*  const inlineData = chunk.candidates[0].content.parts[0].inlineData;
-          const rawB64 = (inlineData.data || "").replace(/\s+/g, ""); 
-          // Return raw Base64 
-          return rawB64; 
-          */
-          
-          // inlineData is the object from the model: { mimeType: 'image/png', data: '...base64...' }
-const inlineData = chunk.candidates[0].content.parts[0].inlineData;
-const rawB64 = (inlineData.data || "").replace(/\s+/g, ""); // remove whitespace/newlines just in case
-const mimeType = inlineData.mimeType || "image/png";
-
-// ensure it's a Buffer (optional, for validation)
-const buffer = Buffer.from(rawB64, "base64");
-
-// produce a canonical base64 string (this also validates)
-const base64 = buffer.toString("base64");
-
-// build a browser-friendly data URL and return it
-const dataUrl = `data:${mimeType};base64,${base64}`;
-return dataUrl;
-          
+            const inlineData = chunk.candidates[0].content.parts[0].inlineData;
+            const rawB64 = (inlineData.data || "").replace(/\s+/g, ""); 
+            const mimeType = inlineData.mimeType || "image/png";
+            const buffer = Buffer.from(rawB64, "base64");
+            const base64 = buffer.toString("base64");
+            
+            finalDataUrl = `data:${mimeType};base64,${base64}`;
+            // We do NOT return here immediately, we continue to allow the stream to finish 
+            // so we can access the aggregated usage metadata at the end.
         } 
       }
-      return null;
+      
+      // Retrieve the full response object to get usage metadata
+      const aggregatedResponse = await response.response;
+
+      return { 
+          image: finalDataUrl, 
+          usage: aggregatedResponse.usageMetadata 
+      };
+
     } catch (error) {
       console.error("Image Gen Error:", error);
+      // On failure, return null (logic in backend handles null as no-op)
       return null; 
     }
   }