Propagate user_prompt_id to GenerateConentRequest for logging (#4741)

1 files changed, 63 insertions, 15 deletions

diff --git a/packages/core/src/code_assist/server.test.ts b/packages/core/src/code_assist/server.test.ts
index 6246fd4e..3fc1891f 100644
--- a/packages/core/src/code_assist/server.test.ts
+++ b/packages/core/src/code_assist/server.test.ts
@@ -14,13 +14,25 @@ vi.mock('google-auth-library');
 describe('CodeAssistServer', () => {
   it('should be able to be constructed', () => {
     const auth = new OAuth2Client();
-    const server = new CodeAssistServer(auth, 'test-project');
+    const server = new CodeAssistServer(
+      auth,
+      'test-project',
+      {},
+      'test-session',
+      UserTierId.FREE,
+    );
     expect(server).toBeInstanceOf(CodeAssistServer);
   });
 
   it('should call the generateContent endpoint', async () => {
     const client = new OAuth2Client();
-    const server = new CodeAssistServer(client, 'test-project');
+    const server = new CodeAssistServer(
+      client,
+      'test-project',
+      {},
+      'test-session',
+      UserTierId.FREE,
+    );
     const mockResponse = {
       response: {
         candidates: [
@@ -38,10 +50,13 @@ describe('CodeAssistServer', () => {
     };
     vi.spyOn(server, 'requestPost').mockResolvedValue(mockResponse);
 
-    const response = await server.generateContent({
-      model: 'test-model',
-      contents: [{ role: 'user', parts: [{ text: 'request' }] }],
-    });
+    const response = await server.generateContent(
+      {
+        model: 'test-model',
+        contents: [{ role: 'user', parts: [{ text: 'request' }] }],
+      },
+      'user-prompt-id',
+    );
 
     expect(server.requestPost).toHaveBeenCalledWith(
       'generateContent',
@@ -55,7 +70,13 @@ describe('CodeAssistServer', () => {
 
   it('should call the generateContentStream endpoint', async () => {
     const client = new OAuth2Client();
-    const server = new CodeAssistServer(client, 'test-project');
+    const server = new CodeAssistServer(
+      client,
+      'test-project',
+      {},
+      'test-session',
+      UserTierId.FREE,
+    );
     const mockResponse = (async function* () {
       yield {
         response: {
@@ -75,10 +96,13 @@ describe('CodeAssistServer', () => {
     })();
     vi.spyOn(server, 'requestStreamingPost').mockResolvedValue(mockResponse);
 
-    const stream = await server.generateContentStream({
-      model: 'test-model',
-      contents: [{ role: 'user', parts: [{ text: 'request' }] }],
-    });
+    const stream = await server.generateContentStream(
+      {
+        model: 'test-model',
+        contents: [{ role: 'user', parts: [{ text: 'request' }] }],
+      },
+      'user-prompt-id',
+    );
 
     for await (const res of stream) {
       expect(server.requestStreamingPost).toHaveBeenCalledWith(
@@ -92,7 +116,13 @@ describe('CodeAssistServer', () => {
 
   it('should call the onboardUser endpoint', async () => {
     const client = new OAuth2Client();
-    const server = new CodeAssistServer(client, 'test-project');
+    const server = new CodeAssistServer(
+      client,
+      'test-project',
+      {},
+      'test-session',
+      UserTierId.FREE,
+    );
     const mockResponse = {
       name: 'operations/123',
       done: true,
@@ -114,7 +144,13 @@ describe('CodeAssistServer', () => {
 
   it('should call the loadCodeAssist endpoint', async () => {
     const client = new OAuth2Client();
-    const server = new CodeAssistServer(client, 'test-project');
+    const server = new CodeAssistServer(
+      client,
+      'test-project',
+      {},
+      'test-session',
+      UserTierId.FREE,
+    );
     const mockResponse = {
       currentTier: {
         id: UserTierId.FREE,
@@ -140,7 +176,13 @@ describe('CodeAssistServer', () => {
 
   it('should return 0 for countTokens', async () => {
     const client = new OAuth2Client();
-    const server = new CodeAssistServer(client, 'test-project');
+    const server = new CodeAssistServer(
+      client,
+      'test-project',
+      {},
+      'test-session',
+      UserTierId.FREE,
+    );
     const mockResponse = {
       totalTokens: 100,
     };
@@ -155,7 +197,13 @@ describe('CodeAssistServer', () => {
 
   it('should throw an error for embedContent', async () => {
     const client = new OAuth2Client();
-    const server = new CodeAssistServer(client, 'test-project');
+    const server = new CodeAssistServer(
+      client,
+      'test-project',
+      {},
+      'test-session',
+      UserTierId.FREE,
+    );
     await expect(
       server.embedContent({
         model: 'test-model',