all versioned tests ok except for streaming

Author: amychisholm03

lib/instrumentation/@google/genai.js

@@ -154,28 +154,31 @@ function instrumentStream ({ agent, shim, request, response, segment, transactio
         // time it took to handle the stream
         segment.touch()
 
+        // result will be {value: undefined, done: true}
+        // when the stream is done, so we need to create
+        // a mock GenerateContentResponse object with
+        // the entire message
+        //
+        // also need to enter this block if there was an
+        // error, so we can record it
         if (result?.done || err) {
-          // result will be {value: undefined, done: true}
-          // when the stream is done, so we need to create
-          // a mock GenerateContentResponse object with
-          // the entire message
-          //
-          // also need to enter this block if there was an
-          // error, so we can record it
-          content.parts[0].text = entireMessage
-          result.value = {
-            candidates: [
-              { content, finishReason }
-            ],
-            modelVersion
+          if (content) {
+            content.parts[0].text = entireMessage
+            result.value = {
+              candidates: [
+                { content, finishReason }
+              ],
+              modelVersion
+            }
           }
+
           recordChatCompletionMessages({
             agent: shim.agent,
             shim,
             segment,
             transaction,
             request,
-            response: result.value,
+            response: result?.value,
             err
           })
         }

test/versioned/google-genai/chat-completions.test.js

@@ -59,7 +59,7 @@ test.afterEach((ctx) => {
 })
 
 test('should create span on successful models generateContent', (t, end) => {
-  const { client, agent, host, port } = t.nr
+  const { client, agent } = t.nr
   helper.runInTransaction(agent, async (tx) => {
     const result = await client.models.generateContent({
       model: 'gemini-2.0-flash',
@@ -69,20 +69,18 @@ test('should create span on successful models generateContent', (t, end) => {
     assert.equal(result.headers, undefined, 'should remove response headers from user result')
     assert.equal(result.candidates[0].content.parts[0].text, '1 plus 2 is 3.')
 
-    const name = `External/${host}:${port}/chat/completions`
     assertSegments(
       tx.trace,
       tx.trace.root,
-      [GEMINI.COMPLETION, [name]],
+      [GEMINI.COMPLETION],
       { exact: false }
     )
 
     tx.end()
     assertSpanKind({
       agent,
       segments: [
-        { name: GEMINI.COMPLETION, kind: 'internal' },
-        { name, kind: 'client' }
+        { name: GEMINI.COMPLETION, kind: 'internal' }
       ]
     })
     end()
@@ -151,12 +149,14 @@ test('should create span on successful models generateContentStream', (t, end) =
     let chunk = {}
     let res = ''
     for await (chunk of stream) {
-      res += chunk?.text
+      assert.ok(chunk.text, 'should have text in chunk')
+      res += chunk.text
     }
+
     assert.equal(chunk.headers, undefined, 'should remove response headers from user result')
     assert.equal(chunk.candidates[0].content.role, 'model')
     const expectedRes = responses.get(content)
-    assert.equal(chunk.candidates[0].content.parts[0].text, expectedRes.streamData)
+    assert.equal(chunk.candidates[0].content.parts[0].text, expectedRes.candidates[0].content.parts[0].text)
     assert.equal(chunk.candidates[0].content.parts[0].text, res)
 
     assertSegments(
@@ -186,16 +186,9 @@ test('should create chat completion message and summary for every message sent i
     })
 
     let res = ''
-
-    let i = 0
     for await (const chunk of stream) {
-      res += chunk?.text
-
-      // break stream
-      if (i === 10) {
-        break
-      }
-      i++
+      assert.ok(chunk.text, 'should have text in chunk')
+      res += chunk.text
     }
 
     const events = agent.customEventAggregator.events.toArray()
@@ -246,7 +239,8 @@ test('should call the tokenCountCallback in streaming', (t, end) => {
     })
 
     for await (const chunk of stream) {
-      res += chunk.choices[0]?.delta?.content
+      assert.ok(chunk.text, 'should have text in chunk')
+      res += chunk.text
     }
 
     const events = agent.customEventAggregator.events.toArray()
@@ -272,20 +266,15 @@ test('handles error in stream', (t, end) => {
     const content = 'bad stream'
     const model = 'gemini-2.0-flash'
     const stream = await client.models.generateContentStream({
-      config: {
-        maxOutputTokens: 100,
-        temperature: 0.5
-      },
       model,
-      contents: [content, 'What does 1 plus 1 equal?'],
-      stream: true
+      contents: [content, 'What does 1 plus 1 equal?']
     })
 
     let res = ''
 
     try {
       for await (const chunk of stream) {
-        res += chunk.choices[0]?.delta?.content
+        if (chunk.text) res += chunk?.text
       }
     } catch (err) {
       assert.ok(res)

test/versioned/google-genai/embeddings.test.js

@@ -56,7 +56,7 @@ test.afterEach((ctx) => {
 })
 
 test('should create span on successful embedding create', (t, end) => {
-  const { client, agent, host, port } = t.nr
+  const { client, agent } = t.nr
   helper.runInTransaction(agent, async (tx) => {
     const results = await client.models.embedContent({
       contents: 'This is an embedding test.',
@@ -65,11 +65,10 @@ test('should create span on successful embedding create', (t, end) => {
 
     assert.equal(results.headers, undefined, 'should remove response headers from user result')
 
-    const name = `External/${host}:${port}/embeddings`
     assertSegments(
       tx.trace,
       tx.trace.root,
-      [GEMINI.EMBEDDING, [name]],
+      [GEMINI.EMBEDDING],
       {
         exact: false
       }
@@ -78,8 +77,7 @@ test('should create span on successful embedding create', (t, end) => {
     assertSpanKind({
       agent,
       segments: [
-        { name: GEMINI.EMBEDDING, kind: 'internal' },
-        { name, kind: 'client' }
+        { name: GEMINI.EMBEDDING, kind: 'internal' }
       ]
     })
     end()

test/versioned/google-genai/mock-server.js

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 New Relic Corporation. All rights reserved.
+ * Copyright 2025 New Relic Corporation. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -8,9 +8,7 @@
 module.exports = GoogleGenAIMockServer
 
 const http = require('node:http')
-const { Readable } = require('node:stream')
 const RESPONSES = require('./mock-responses')
-const crypto = require('crypto')
 
 /**
  * Build a mock server that listens on a 127.0.0.1 and a random port that
@@ -68,80 +66,13 @@ function handler(req, res) {
 
     const { code, body } = RESPONSES.get(prompt)
     res.statusCode = code
+    res.write(JSON.stringify(body))
+    res.end()
 
-    if (prompt.toLowerCase().includes('stream')) {
-      const streamData = body.candidates[0].content.parts[0].text
-      let asyncGen
-      if (streamData !== 'do random') {
-        asyncGen = finiteAsyncGen(streamData, { ...body })
-      } else {
-        asyncGen = randomAsyncGen({ ...body })
-      }
-      // return asyncGen
-
-      // Write each chunk from the async iterator to the response
-      (async () => {
-        try {
-          for await (const chunk of asyncGen) {
-            if (chunk?.text) res.write(chunk.text)
-          }
-          res.end()
-        } catch (err) {
-          res.destroy(err)
-        }
-      })()
-    } else {
-      res.write(JSON.stringify(body))
-      res.end()
-    }
+    // TODO: Mock streaming responses
   })
 }
 
-/**
- * Mocks the Google GenAI streaming API by returning a Promise that resolves
- * to an async generator, which yields response chunks as the real API would.
- *
- * @param {string} dataToStream The string to split and stream.
- * @param {object} chunkTemplate The template for each chunk.
- * @returns {object} An async generator that yields chunks of data.
- */
-function finiteAsyncGen(dataToStream, chunkTemplate) {
-  const parts = dataToStream.split(' ')
-  const asyncGen = (async function * () {
-    for (let i = 0; i < parts.length; i++) {
-      const content = i === parts.length - 1 ? parts[i] : `${parts[i]} `
-      const chunk = chunkTemplate
-      chunk.candidates[0].content.parts[0].text = content
-      chunk.text = content
-      yield chunk
-    }
-    // End the stream
-    yield undefined
-  })()
-  return asyncGen
-}
-
-/**
- * Creates a stream that will stream an infinite number of GoogleGenAI stream data
- * chunks using an async generator.
- *
- * @param {object} chunkTemplate An object that is shaped like a GoogleGenAI stream
- * data object.
- * @returns {object} An async generator that yields chunks of data.
- */
-function randomAsyncGen(chunkTemplate) {
-  const asyncGen = (async function * () {
-    while (true) {
-      const data = crypto.randomBytes(16)
-      // Deep clone to avoid mutating the original template
-      const chunk = JSON.parse(JSON.stringify(chunkTemplate))
-      chunk.value.candidates[0].content.parts[0].text = data.toString('base64')
-      yield chunk
-    }
-  })()
-  return asyncGen
-}
-
 function getShortenedPrompt(reqBody) {
   try {
     const prompt = reqBody.contents?.[0]?.parts?.[0]?.text ||

test/versioned/google-genai/package.json

@@ -9,15 +9,15 @@
   "version": "0.0.0",
   "private": true,
   "engines": {
-    "node": ">=18"
+    "node": ">=20"
   },
   "tests": [
     {
       "engines": {
-        "node": ">=18"
+        "node": ">=20"
       },
       "dependencies": {
-        "@google/genai": "^1.1.0"
+        "@google/genai": "^1.2.0"
       },
       "files": [
         "chat-completions.test.js",