1 files changed, 15 insertions, 63 deletions

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