chrome-devtools-frontend 1.0.1386602 → 1.0.1402249

package/front_end/panels/{freestyler/FreestylerAgent.test.ts → ai_assistance/agents/StylingAgent.test.ts}

@@ -2,19 +2,18 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-import * as Host from '../../core/host/host.js';
-import * as Root from '../../core/root/root.js';
-import * as SDK from '../../core/sdk/sdk.js';
+import * as Host from '../../../core/host/host.js';
+import * as Root from '../../../core/root/root.js';
+import * as SDK from '../../../core/sdk/sdk.js';
 import {
   describeWithEnvironment,
   getGetHostConfigStub,
-} from '../../testing/EnvironmentHelpers.js';
+} from '../../../testing/EnvironmentHelpers.js';
+import * as AiAssistance from '../ai_assistance.js';
 
-import * as Freestyler from './freestyler.js';
+const {StylingAgent, ErrorType} = AiAssistance;
 
-const {FreestylerAgent} = Freestyler;
-
-describeWithEnvironment('FreestylerAgent', () => {
+describeWithEnvironment('StylingAgent', () => {
   function mockHostConfig(
       modelId?: string, temperature?: number, userTier?: string,
       executionMode?: Root.Runtime.HostConfigFreestylerExecutionMode) {
@@ -39,28 +38,36 @@ describeWithEnvironment('FreestylerAgent', () => {
     };
   }
   describe('parseResponse', () => {
-    const agent = new FreestylerAgent({
+    const agent = new StylingAgent({
       aidaClient: {} as Host.AidaClient.AidaClient,
     });
 
+    function getParsedTextResponse(explanation: string): AiAssistance.ParsedResponse {
+      return agent.parseResponse({
+        explanation,
+        metadata: {},
+        completed: false,
+      });
+    }
+
     it('parses a thought', async () => {
       const payload = 'some response';
-      assert.deepStrictEqual(
-          agent.parseResponse(`THOUGHT: ${payload}`),
+      assert.deepEqual(
+          getParsedTextResponse(`THOUGHT: ${payload}`),
           {
             title: undefined,
             thought: payload,
           },
       );
-      assert.deepStrictEqual(
-          agent.parseResponse(`   THOUGHT: ${payload}`),
+      assert.deepEqual(
+          getParsedTextResponse(`   THOUGHT: ${payload}`),
           {
             title: undefined,
             thought: payload,
           },
       );
-      assert.deepStrictEqual(
-          agent.parseResponse(`Something\n   THOUGHT: ${payload}`),
+      assert.deepEqual(
+          getParsedTextResponse(`Something\n   THOUGHT: ${payload}`),
           {
             title: undefined,
             thought: payload,
@@ -69,22 +76,22 @@ describeWithEnvironment('FreestylerAgent', () => {
     });
     it('parses a answer', async () => {
       const payload = 'some response';
-      assert.deepStrictEqual(
-          agent.parseResponse(`ANSWER: ${payload}`),
+      assert.deepEqual(
+          getParsedTextResponse(`ANSWER: ${payload}`),
           {
             answer: payload,
             suggestions: undefined,
           },
       );
-      assert.deepStrictEqual(
-          agent.parseResponse(`   ANSWER: ${payload}`),
+      assert.deepEqual(
+          getParsedTextResponse(`   ANSWER: ${payload}`),
           {
             answer: payload,
             suggestions: undefined,
           },
       );
-      assert.deepStrictEqual(
-          agent.parseResponse(`Something\n   ANSWER: ${payload}`),
+      assert.deepEqual(
+          getParsedTextResponse(`Something\n   ANSWER: ${payload}`),
           {
             answer: payload,
             suggestions: undefined,
@@ -95,36 +102,36 @@ describeWithEnvironment('FreestylerAgent', () => {
       const payload = `a
 b
 c`;
-      assert.deepStrictEqual(
-          agent.parseResponse(`ANSWER: ${payload}`),
+      assert.deepEqual(
+          getParsedTextResponse(`ANSWER: ${payload}`),
           {
             answer: payload,
             suggestions: undefined,
           },
       );
-      assert.deepStrictEqual(
-          agent.parseResponse(`   ANSWER: ${payload}`),
+      assert.deepEqual(
+          getParsedTextResponse(`   ANSWER: ${payload}`),
           {
             answer: payload,
             suggestions: undefined,
           },
       );
-      assert.deepStrictEqual(
-          agent.parseResponse(`Something\n   ANSWER: ${payload}`),
+      assert.deepEqual(
+          getParsedTextResponse(`Something\n   ANSWER: ${payload}`),
           {
             answer: payload,
             suggestions: undefined,
           },
       );
-      assert.deepStrictEqual(
-          agent.parseResponse(`ANSWER: ${payload}\nTHOUGHT: thought`),
+      assert.deepEqual(
+          getParsedTextResponse(`ANSWER: ${payload}\nTHOUGHT: thought`),
           {
             answer: payload,
             suggestions: undefined,
           },
       );
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `ANSWER: ${payload}\nOBSERVATION: observation`,
               ),
           {
@@ -132,8 +139,8 @@ c`;
             suggestions: undefined,
           },
       );
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `ANSWER: ${payload}\nACTION\naction\nSTOP`,
               ),
           {
@@ -147,24 +154,24 @@ c`;
       const payload = `const data = {
   someKey: "value",
 }`;
-      assert.deepStrictEqual(
-          agent.parseResponse(`ACTION\n${payload}\nSTOP`),
+      assert.deepEqual(
+          getParsedTextResponse(`ACTION\n${payload}\nSTOP`),
           {
             action: payload,
             title: undefined,
             thought: undefined,
           },
       );
-      assert.deepStrictEqual(
-          agent.parseResponse(`ACTION\n${payload}`),
+      assert.deepEqual(
+          getParsedTextResponse(`ACTION\n${payload}`),
           {
             action: payload,
             title: undefined,
             thought: undefined,
           },
       );
-      assert.deepStrictEqual(
-          agent.parseResponse(`ACTION\n\n${payload}\n\nSTOP`),
+      assert.deepEqual(
+          getParsedTextResponse(`ACTION\n\n${payload}\n\nSTOP`),
           {
             action: payload,
             title: undefined,
@@ -172,8 +179,8 @@ c`;
           },
       );
 
-      assert.deepStrictEqual(
-          agent.parseResponse(`ACTION\n\n${payload}\n\nANSWER: answer`),
+      assert.deepEqual(
+          getParsedTextResponse(`ACTION\n\n${payload}\n\nANSWER: answer`),
           {
             action: payload,
             title: undefined,
@@ -186,8 +193,8 @@ c`;
         const data = {
           styles
         };`;
-      assert.deepStrictEqual(
-          agent.parseResponse(`ACTION\n${payload}STOP`),
+      assert.deepEqual(
+          getParsedTextResponse(`ACTION\n${payload}STOP`),
           {
             action: payload,
             title: undefined,
@@ -198,8 +205,8 @@ c`;
     it('parses a thought and title', async () => {
       const payload = 'some response';
       const title = 'this is the title';
-      assert.deepStrictEqual(
-          agent.parseResponse(`THOUGHT: ${payload}\nTITLE: ${title}`),
+      assert.deepEqual(
+          getParsedTextResponse(`THOUGHT: ${payload}\nTITLE: ${title}`),
           {
             thought: payload,
             title,
@@ -211,8 +218,8 @@ c`;
       const payload = `const data = {
   someKey: "value",
 }`;
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `ACTION\n\`\`\`\n${payload}\n\`\`\`\nSTOP`,
               ),
           {
@@ -227,8 +234,8 @@ c`;
       const payload = `const data = {
   someKey: "value",
 }`;
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `ACTION\n\`\`\`\`\`\njs\n${payload}\n\`\`\`\`\`\nSTOP`,
               ),
           {
@@ -244,8 +251,8 @@ c`;
   someKey: "value",
 }`;
       const thoughtPayload = 'thought';
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `THOUGHT:${thoughtPayload}\nACTION\n${actionPayload}\nSTOP`,
               ),
           {
@@ -259,8 +266,8 @@ c`;
     it('parses a thought and an answer', async () => {
       const answerPayload = 'answer';
       const thoughtPayload = 'thought';
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `THOUGHT:${thoughtPayload}\nANSWER:${answerPayload}`,
               ),
           {
@@ -274,8 +281,8 @@ c`;
       const answerPayload = 'answer';
       const suggestions = ['suggestion'] as [string];
       const suggestionsText = JSON.stringify(suggestions);
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `ANSWER:${answerPayload}\nSUGGESTIONS: ${suggestionsText}`,
               ),
           {
@@ -292,8 +299,8 @@ c`;
   someKey: "value",
 }`;
       const title = 'title';
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `THOUGHT: ${thoughtPayload}\nTITLE: ${title}\nACTION\n${actionPayload}\nSTOP\nANSWER:${answerPayload}`,
               ),
           {
@@ -310,8 +317,8 @@ c`;
         const data = {
           styles
         };`;
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `ACTION\n${payload}STOP\nANSWER:${answerPayload}\nSUGGESTIONS: ${JSON.stringify(suggestions)}`),
           {
             action: payload,
@@ -325,8 +332,8 @@ c`;
         const data = {
           styles
         };`;
-      assert.deepStrictEqual(
-          agent.parseResponse(`ACTION\n${payload}STOP\nOBSERVATION:{styles: {}}`),
+      assert.deepEqual(
+          getParsedTextResponse(`ACTION\n${payload}STOP\nOBSERVATION:{styles: {}}`),
           {
             action: payload,
             thought: undefined,
@@ -340,8 +347,8 @@ c`;
           styles
         };`;
       const thoughtPayload = 'thought';
-      assert.deepStrictEqual(
-          agent.parseResponse(`ACTION\n${payload}STOP\nTHOUGHT:${thoughtPayload}`),
+      assert.deepEqual(
+          getParsedTextResponse(`ACTION\n${payload}STOP\nTHOUGHT:${thoughtPayload}`),
           {
             action: payload,
             thought: thoughtPayload,
@@ -351,8 +358,8 @@ c`;
     });
 
     it('parses a response as an answer', async () => {
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               'This is also an answer',
               ),
           {
@@ -363,8 +370,8 @@ c`;
     });
 
     it('parses a response with no instruction tags as an answer and correctly parses suggestions', async () => {
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               'This is also an answer\nSUGGESTIONS: [\"suggestion\"]',
               ),
           {
@@ -376,8 +383,8 @@ c`;
 
     it('parses multi line thoughts', () => {
       const thoughtText = 'first line\nsecond line';
-      assert.deepStrictEqual(
-          agent.parseResponse(`THOUGHT: ${thoughtText}`),
+      assert.deepEqual(
+          getParsedTextResponse(`THOUGHT: ${thoughtText}`),
           {
             thought: thoughtText,
             title: undefined,
@@ -401,7 +408,7 @@ c`;
       element.simpleSelector.returns('div#myElement');
       element.getChildNodesPromise.resolves(null);
 
-      const result = await FreestylerAgent.describeElement(element);
+      const result = await StylingAgent.describeElement(element);
 
       assert.strictEqual(result, '* Its selector is `div#myElement`');
     });
@@ -424,7 +431,7 @@ c`;
       element.previousSibling = null;
       element.parentNode = null;
 
-      const result = await FreestylerAgent.describeElement(element);
+      const result = await StylingAgent.describeElement(element);
       const expectedOutput = `* Its selector is \`div#parentElement\`
 * It has 2 child element nodes: \`span.child1\`, \`span.child2\`
 * It only has 1 child text node`;
@@ -455,7 +462,7 @@ c`;
       element.previousSibling = previousSibling;
       element.parentNode = parentNode;
 
-      const result = await FreestylerAgent.describeElement(element);
+      const result = await StylingAgent.describeElement(element);
       const expectedOutput = `* Its selector is \`div#parentElement\`
 * It has a next sibling and it is an element node
 * It has a previous sibling and it is a non element node
@@ -478,33 +485,33 @@ c`;
 
     it('builds a request with a model id', async () => {
       mockHostConfig('test model');
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: {} as Host.AidaClient.AidaClient,
       });
       assert.strictEqual(
-          agent.buildRequest({input: 'test input'}).options?.model_id,
+          agent.buildRequest({text: 'test input'}).options?.model_id,
           'test model',
       );
     });
 
     it('builds a request with a temperature', async () => {
       mockHostConfig('test model', 1);
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: {} as Host.AidaClient.AidaClient,
       });
       assert.strictEqual(
-          agent.buildRequest({input: 'test input'}).options?.temperature,
+          agent.buildRequest({text: 'test input'}).options?.temperature,
           1,
       );
     });
 
     it('builds a request with a user tier', async () => {
       mockHostConfig('test model', 1, 'PUBLIC');
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: {} as Host.AidaClient.AidaClient,
       });
       assert.strictEqual(
-          agent.buildRequest({input: 'test input'}).metadata?.user_tier,
+          agent.buildRequest({text: 'test input'}).metadata?.user_tier,
           3,
       );
     });
@@ -512,41 +519,41 @@ c`;
     it('structure matches the snapshot', () => {
       mockHostConfig('test model');
 
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: {} as Host.AidaClient.AidaClient,
         serverSideLoggingEnabled: true,
       });
       sinon.stub(agent, 'preamble').value('preamble');
       agent.chatNewHistoryForTesting = [
         {
-          type: Freestyler.ResponseType.USER_QUERY,
+          type: AiAssistance.ResponseType.USER_QUERY,
           query: 'question',
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
+          type: AiAssistance.ResponseType.QUERYING,
           query: 'question',
         },
         {
-          type: Freestyler.ResponseType.ANSWER,
+          type: AiAssistance.ResponseType.ANSWER,
           text: 'answer',
         },
       ];
-      assert.deepStrictEqual(
+      assert.deepEqual(
           agent.buildRequest({
-            input: 'test input',
+            text: 'test input',
           }),
           {
-            input: 'test input',
+            current_message: {role: Host.AidaClient.Role.USER, parts: [{text: 'test input'}]},
             client: 'CHROME_DEVTOOLS',
             preamble: 'preamble',
-            chat_history: [
+            historical_contexts: [
               {
-                entity: 1,
-                text: 'question',
+                role: 1,
+                parts: [{text: 'question'}],
               },
               {
-                entity: 2,
-                text: 'ANSWER: answer',
+                role: 2,
+                parts: [{text: 'ANSWER: answer'}],
               },
             ],
             metadata: {
@@ -608,8 +615,8 @@ c`;
           count++;
         }
         const execJs =
-            sinon.mock().throws(new Freestyler.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
-        const agent = new FreestylerAgent({
+            sinon.mock().throws(new AiAssistance.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
+        const agent = new StylingAgent({
           aidaClient: mockAidaClient(generateActionAndAnswer),
           createExtensionScope,
           confirmSideEffectForTest: stub,
@@ -618,7 +625,7 @@ c`;
         });
 
         promise.resolve(true);
-        await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
+        await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
 
         sinon.assert.match(execJs.getCall(0).args[1], sinon.match({throwOnSideEffect: true}));
       });
@@ -647,9 +654,9 @@ c`;
           count++;
         }
         const execJs = sinon.mock().twice();
-        execJs.onCall(0).throws(new Freestyler.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
+        execJs.onCall(0).throws(new AiAssistance.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
         execJs.onCall(1).resolves('value');
-        const agent = new FreestylerAgent({
+        const agent = new StylingAgent({
           aidaClient: mockAidaClient(generateActionAndAnswer),
           createExtensionScope,
           confirmSideEffectForTest: stub,
@@ -657,9 +664,9 @@ c`;
 
         });
         promise.resolve(true);
-        await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
+        await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
 
-        assert.strictEqual(execJs.getCalls().length, 2);
+        assert.lengthOf(execJs.getCalls(), 2);
         sinon.assert.match(execJs.getCall(1).args[1], sinon.match({throwOnSideEffect: false}));
       });
 
@@ -686,9 +693,9 @@ c`;
 
           count++;
         }
-        const execJs = sinon.mock().twice();
-        execJs.onCall(0).throws(new Freestyler.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
-        const agent = new FreestylerAgent({
+        const execJs = sinon.mock().once();
+        execJs.onCall(0).throws(new AiAssistance.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
+        const agent = new StylingAgent({
           aidaClient: mockAidaClient(generateActionAndAnswer),
           createExtensionScope,
           confirmSideEffectForTest: stub,
@@ -696,12 +703,52 @@ c`;
 
         });
         promise.resolve(false);
-        const responses = await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
+        const responses = await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
 
-        const actionStep = responses.find(response => response.type === Freestyler.ResponseType.ACTION)!;
+        const actionStep = responses.find(response => response.type === AiAssistance.ResponseType.ACTION)!;
 
         assert.strictEqual(actionStep.output, 'Error: User denied code execution with side effects.');
-        assert.strictEqual(execJs.getCalls().length, 1);
+        assert.lengthOf(execJs.getCalls(), 1);
+      });
+
+      it('returns error when side effect is aborted', async () => {
+        const selected = new AiAssistance.NodeContext(element);
+        async function* generateAction() {
+          yield {
+            explanation: `ACTION
+            $0.style.backgroundColor = 'red'
+            STOP`,
+            metadata: {},
+            completed: true,
+          };
+        }
+        const execJs = sinon.mock().once().throws(
+            new AiAssistance.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
+        const sideEffectConfirmationPromise = Promise.withResolvers();
+        const agent = new StylingAgent({
+          aidaClient: mockAidaClient(generateAction),
+          createExtensionScope,
+          confirmSideEffectForTest: sinon.stub().returns(sideEffectConfirmationPromise),
+          execJs,
+        });
+
+        const responses: AiAssistance.ResponseData[] = [];
+        const controller = new AbortController();
+        for await (const result of agent.run('test', {selected, signal: controller.signal})) {
+          responses.push(result);
+          if (result.type === 'side-effect') {
+            // Initial code invocation resulting in a side-effect
+            // happened.
+            assert.isTrue(execJs.calledOnce);
+            // Emulate abort when waiting for the side-effect confirmation.
+            controller.abort();
+          }
+        }
+
+        const errorStep = responses.at(-1) as AiAssistance.ErrorResponse;
+        assert.exists(errorStep);
+        assert.strictEqual(errorStep.error, ErrorType.ABORT);
+        assert.isFalse(await sideEffectConfirmationPromise.promise);
       });
     });
 
@@ -727,16 +774,16 @@ c`;
           count++;
         }
         const execJs = sinon.mock().returns(new Array(10_000).fill('<div>...</div>').join());
-        const agent = new FreestylerAgent({
+        const agent = new StylingAgent({
           aidaClient: mockAidaClient(generateActionAndAnswer),
           createExtensionScope,
           execJs,
 
         });
 
-        const result = await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
+        const result = await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
         const actionSteps = result.filter(step => {
-          return step.type === Freestyler.ResponseType.ACTION;
+          return step.type === AiAssistance.ResponseType.ACTION;
         });
         assert(actionSteps.length === 1, 'Found non or multiple action steps');
         const actionStep = actionSteps.at(0)!;
@@ -754,19 +801,19 @@ c`;
       }
 
       const execJs = sinon.spy();
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: mockAidaClient(generateAnswer),
         execJs,
       });
 
-      const responses = await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
-      assert.deepStrictEqual(responses, [
+      const responses = await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
+      assert.deepEqual(responses, [
         {
-          type: Freestyler.ResponseType.USER_QUERY,
+          type: AiAssistance.ResponseType.USER_QUERY,
           query: 'test',
         },
         {
-          type: Freestyler.ResponseType.CONTEXT,
+          type: AiAssistance.ResponseType.CONTEXT,
           title: 'Analyzing the prompt',
           details: [
             {
@@ -776,30 +823,30 @@ c`;
           ],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
+          type: AiAssistance.ResponseType.QUERYING,
           query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
-          type: Freestyler.ResponseType.ANSWER,
+          type: AiAssistance.ResponseType.ANSWER,
           text: 'this is the answer',
           suggestions: undefined,
           rpcId: undefined,
         },
       ]);
       sinon.assert.notCalled(execJs);
-      assert.deepStrictEqual(agent.chatHistoryForTesting, [
+      assert.deepEqual(agent.chatHistoryForTesting, [
         {
-          entity: 1,
-          text: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
+          role: 1,
+          parts: [{text: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test'}],
         },
         {
-          entity: 2,
-          text: 'ANSWER: this is the answer',
+          role: 2,
+          parts: [{text: 'ANSWER: this is the answer'}],
         },
       ]);
     });
 
-    it('correctly handles chat_history in AIDA requests', async () => {
+    it('correctly handles historical_contexts in AIDA requests', async () => {
       const requests: Host.AidaClient.AidaRequest[] = [];
 
       let i = 0;
@@ -827,33 +874,44 @@ STOP`,
 
       const execJs = sinon.mock().once();
       execJs.onCall(0).returns('test data');
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: mockAidaClient(generateAnswer),
         createExtensionScope,
         execJs,
       });
 
-      await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
+      await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
 
       assert.lengthOf(requests, 2, 'Unexpected number of AIDA requests');
-      assert.isUndefined(requests[0].chat_history, 'Unexpected chat history in the initial request');
-      assert.strictEqual(
-          requests[0].input, '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
-          'Unexpected input in the initial request');
-      assert.deepStrictEqual(
-          requests[1].chat_history,
+      assert.isUndefined(requests[0].historical_contexts, 'Unexpected historical contexts in the initial request');
+      assert.exists(requests[0].current_message);
+      assert.lengthOf(requests[0].current_message.parts, 1);
+      assert.deepEqual(
+          requests[0].current_message.parts[0],
+          {text: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test'},
+          'Unexpected input text in the initial request');
+      assert.strictEqual(requests[0].current_message.role, Host.AidaClient.Role.USER);
+      assert.deepEqual(
+          requests[1].historical_contexts,
           [
             {
-              entity: 1,
-              text: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
+              role: 1,
+              parts: [{text: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test'}],
             },
             {
-              entity: 2,
-              text: 'THOUGHT: I am thinking.\nTITLE: thinking\nACTION\nconst data = {\"test\": \"observation\"};\nSTOP',
+              role: 2,
+              parts: [{
+                text:
+                    'THOUGHT: I am thinking.\nTITLE: thinking\nACTION\nconst data = {\"test\": \"observation\"};\nSTOP',
+              }],
             },
           ],
-          'Unexpected chat history in the follow-up request');
-      assert.strictEqual(requests[1].input, 'OBSERVATION: test data', 'Unexpected input in the follow-up request');
+          'Unexpected historical contexts in the follow-up request');
+      assert.exists(requests[1].current_message);
+      assert.lengthOf(requests[1].current_message.parts, 1);
+      assert.deepEqual(
+          requests[1].current_message.parts[0], {text: 'OBSERVATION: test data'},
+          'Unexpected input in the follow-up request');
     });
 
     it('generates an rpcId for the answer', async () => {
@@ -867,20 +925,20 @@ STOP`,
         };
       }
 
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: mockAidaClient(generateAnswer),
         execJs: sinon.spy(),
 
       });
 
-      const responses = await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
-      assert.deepStrictEqual(responses, [
+      const responses = await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
+      assert.deepEqual(responses, [
         {
-          type: Freestyler.ResponseType.USER_QUERY,
+          type: AiAssistance.ResponseType.USER_QUERY,
           query: 'test',
         },
         {
-          type: Freestyler.ResponseType.CONTEXT,
+          type: AiAssistance.ResponseType.CONTEXT,
           title: 'Analyzing the prompt',
           details: [
             {
@@ -890,11 +948,11 @@ STOP`,
           ],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
+          type: AiAssistance.ResponseType.QUERYING,
           query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
-          type: Freestyler.ResponseType.ANSWER,
+          type: AiAssistance.ResponseType.ANSWER,
           text: 'this is the answer',
           suggestions: undefined,
           rpcId: 123,
@@ -908,29 +966,25 @@ STOP`,
           explanation: 'ANSWER: this is the answer',
           metadata: {
             rpcGlobalId: 123,
-            attributionMetadata: [{
-              attributionAction: Host.AidaClient.RecitationAction.BLOCK,
-              citations: [],
-            }],
           },
-          completed: true,
+          completed: false,
         };
+        throw new Host.AidaClient.AidaBlockError();
       }
 
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: mockAidaClient(generateAnswer),
         execJs: sinon.spy(),
-
       });
 
-      const responses = await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
-      assert.deepStrictEqual(responses, [
+      const responses = await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
+      assert.deepEqual(responses, [
         {
-          type: Freestyler.ResponseType.USER_QUERY,
+          type: AiAssistance.ResponseType.USER_QUERY,
           query: 'test',
         },
         {
-          type: Freestyler.ResponseType.CONTEXT,
+          type: AiAssistance.ResponseType.CONTEXT,
           title: 'Analyzing the prompt',
           details: [
             {
@@ -940,13 +994,16 @@ STOP`,
           ],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
+          type: AiAssistance.ResponseType.QUERYING,
           query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
-          rpcId: undefined,
-          type: Freestyler.ResponseType.ERROR,
-          error: Freestyler.ErrorType.UNKNOWN,
+          text: 'this is the answer',
+          type: AiAssistance.ResponseType.ANSWER,
+        },
+        {
+          type: AiAssistance.ResponseType.ERROR,
+          error: AiAssistance.ErrorType.BLOCK,
         },
       ]);
     });
@@ -957,29 +1014,29 @@ STOP`,
           explanation: 'ANSWER: this is the answer',
           metadata: {
             rpcGlobalId: 123,
-            attributionMetadata: [{
+            attributionMetadata: {
               attributionAction: Host.AidaClient.RecitationAction.ACTION_UNSPECIFIED,
               citations: [],
-            }],
+            },
           },
           completed: true,
         };
       }
 
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: mockAidaClient(generateAnswer),
         execJs: sinon.spy(),
 
       });
 
-      const responses = await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
-      assert.deepStrictEqual(responses, [
+      const responses = await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
+      assert.deepEqual(responses, [
         {
-          type: Freestyler.ResponseType.USER_QUERY,
+          type: AiAssistance.ResponseType.USER_QUERY,
           query: 'test',
         },
         {
-          type: Freestyler.ResponseType.CONTEXT,
+          type: AiAssistance.ResponseType.CONTEXT,
           title: 'Analyzing the prompt',
           details: [
             {
@@ -989,11 +1046,11 @@ STOP`,
           ],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
+          type: AiAssistance.ResponseType.QUERYING,
           query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
-          type: Freestyler.ResponseType.ANSWER,
+          type: AiAssistance.ResponseType.ANSWER,
           text: 'this is the answer',
           suggestions: undefined,
           rpcId: 123,
@@ -1027,12 +1084,12 @@ STOP
         };
       }
 
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: mockAidaClient(generatePartialAndFullAction),
         createExtensionScope,
         execJs,
       });
-      await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
+      await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
 
       sinon.assert.calledOnce(execJs);
       assert.include(execJs.lastCall.args[0], 'console.log(\'hello\');');
@@ -1048,18 +1105,18 @@ STOP
       }
 
       const execJs = sinon.spy();
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: mockAidaClient(generateNothing),
         execJs,
       });
-      const responses = await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
-      assert.deepStrictEqual(responses, [
+      const responses = await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
+      assert.deepEqual(responses, [
         {
-          type: Freestyler.ResponseType.USER_QUERY,
+          type: AiAssistance.ResponseType.USER_QUERY,
           query: 'test',
         },
         {
-          type: Freestyler.ResponseType.CONTEXT,
+          type: AiAssistance.ResponseType.CONTEXT,
           title: 'Analyzing the prompt',
           details: [
             {
@@ -1069,17 +1126,16 @@ STOP
           ],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
+          type: AiAssistance.ResponseType.QUERYING,
           query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
-          type: Freestyler.ResponseType.ERROR,
-          error: Freestyler.ErrorType.UNKNOWN,
-          rpcId: undefined,
+          type: AiAssistance.ResponseType.ERROR,
+          error: AiAssistance.ErrorType.UNKNOWN,
         },
       ]);
       sinon.assert.notCalled(execJs);
-      assert.deepStrictEqual(agent.chatHistoryForTesting, []);
+      assert.deepEqual(agent.chatHistoryForTesting, []);
     });
 
     it('generates an action response if action and answer both present', async () => {
@@ -1109,20 +1165,20 @@ ANSWER: this is the answer`,
 
       const execJs = sinon.mock().once();
       execJs.onCall(0).returns('hello');
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: mockAidaClient(generateNothing),
         createExtensionScope,
         execJs,
 
       });
-      const responses = await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
-      assert.deepStrictEqual(responses, [
+      const responses = await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
+      assert.deepEqual(responses, [
         {
-          type: Freestyler.ResponseType.USER_QUERY,
+          type: AiAssistance.ResponseType.USER_QUERY,
           query: 'test',
         },
         {
-          type: Freestyler.ResponseType.CONTEXT,
+          type: AiAssistance.ResponseType.CONTEXT,
           title: 'Analyzing the prompt',
           details: [
             {
@@ -1132,27 +1188,26 @@ ANSWER: this is the answer`,
           ],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
+          type: AiAssistance.ResponseType.QUERYING,
           query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
-          type: Freestyler.ResponseType.THOUGHT,
+          type: AiAssistance.ResponseType.THOUGHT,
           thought: 'I am thinking.',
           rpcId: undefined,
         },
         {
-          type: Freestyler.ResponseType.ACTION,
+          type: AiAssistance.ResponseType.ACTION,
           code: 'console.log(\'hello\');',
           output: 'hello',
           canceled: false,
-          rpcId: undefined,
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
+          type: AiAssistance.ResponseType.QUERYING,
           query: 'OBSERVATION: hello',
         },
         {
-          type: Freestyler.ResponseType.ANSWER,
+          type: AiAssistance.ResponseType.ANSWER,
           text: 'this is the actual answer',
           suggestions: undefined,
           rpcId: undefined,
@@ -1181,47 +1236,47 @@ ANSWER: this is the answer`,
       }
 
       const execJs = sinon.spy(async () => 'undefined');
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: mockAidaClient(generateMultipleTimes),
         createExtensionScope,
         execJs,
 
       });
 
-      await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
+      await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
 
-      assert.deepStrictEqual(agent.chatHistoryForTesting, [
+      assert.deepEqual(agent.chatHistoryForTesting, [
         {
-          entity: 1,
-          text: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
+          role: 1,
+          parts: [{text: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test'}],
         },
         {
-          entity: 2,
-          text: 'THOUGHT: thought 1\nTITLE: test\nACTION\nconsole.log(\'test\')\nSTOP',
+          role: 2,
+          parts: [{text: 'THOUGHT: thought 1\nTITLE: test\nACTION\nconsole.log(\'test\')\nSTOP'}],
         },
         {
-          entity: 1,
-          text: 'OBSERVATION: undefined',
+          role: 1,
+          parts: [{text: 'OBSERVATION: undefined'}],
         },
         {
-          entity: 2,
-          text: 'THOUGHT: thought 2\nTITLE: test\nACTION\nconsole.log(\'test\')\nSTOP',
+          role: 2,
+          parts: [{text: 'THOUGHT: thought 2\nTITLE: test\nACTION\nconsole.log(\'test\')\nSTOP'}],
         },
         {
-          entity: 1,
-          text: 'OBSERVATION: undefined',
+          role: 1,
+          parts: [{text: 'OBSERVATION: undefined'}],
         },
         {
-          entity: 2,
-          text: 'THOUGHT: thought 3\nTITLE: test\nACTION\nconsole.log(\'test\')\nSTOP',
+          role: 2,
+          parts: [{text: 'THOUGHT: thought 3\nTITLE: test\nACTION\nconsole.log(\'test\')\nSTOP'}],
         },
         {
-          entity: 1,
-          text: 'OBSERVATION: undefined',
+          role: 1,
+          parts: [{text: 'OBSERVATION: undefined'}],
         },
         {
-          entity: 2,
-          text: 'ANSWER: this is the answer',
+          role: 2,
+          parts: [{text: 'ANSWER: this is the answer'}],
         },
       ]);
     });
@@ -1249,7 +1304,7 @@ ANSWER: this is the answer`,
       }
 
       const execJs = sinon.spy();
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: mockAidaClient(generateAndAbort),
         createExtensionScope,
         execJs,
@@ -1258,9 +1313,9 @@ ANSWER: this is the answer`,
       const controller = new AbortController();
       controller.abort();
       await Array.fromAsync(
-          agent.run('test', {selected: new Freestyler.NodeContext(element), signal: controller.signal}));
+          agent.run('test', {selected: new AiAssistance.NodeContext(element), signal: controller.signal}));
 
-      assert.deepStrictEqual(agent.chatHistoryForTesting, []);
+      assert.deepEqual(agent.chatHistoryForTesting, []);
     });
   });
 
@@ -1269,8 +1324,6 @@ ANSWER: this is the answer`,
     beforeEach(() => {
       mockHostConfig();
       element = sinon.createStubInstance(SDK.DOMModel.DOMNode);
-      // @ts-ignore
-      setAiAssistancePersistentHistory(true);
     });
 
     it('stores history via AiHistoryStorage', async () => {
@@ -1291,23 +1344,23 @@ ANSWER: this is the answer`,
           completed: false,
         };
       }
-      const historyStub = sinon.stub(Freestyler.AiHistoryStorage.instance(), 'upsertHistoryEntry');
+      const historyStub = sinon.stub(AiAssistance.AiHistoryStorage.instance(), 'upsertHistoryEntry');
       const execJs = sinon.spy(async () => 'undefined');
-      const agent = new FreestylerAgent({
+      const agent = new StylingAgent({
         aidaClient: mockAidaClient(generateMultipleTimes),
         createExtensionScope,
         execJs,
       });
 
-      await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
+      await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
 
-      assert.deepStrictEqual(historyStub.lastCall.args[0].history, [
+      assert.deepEqual(historyStub.lastCall.args[0].history, [
         {
-          type: Freestyler.ResponseType.USER_QUERY,
+          type: AiAssistance.ResponseType.USER_QUERY,
           query: 'test',
         },
         {
-          type: Freestyler.ResponseType.CONTEXT,
+          type: AiAssistance.ResponseType.CONTEXT,
           title: 'Analyzing the prompt',
           details: [{
             text: '* Its selector is `undefined`',
@@ -1315,32 +1368,31 @@ ANSWER: this is the answer`,
           }],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
+          type: AiAssistance.ResponseType.QUERYING,
           query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
-          type: Freestyler.ResponseType.TITLE,
+          type: AiAssistance.ResponseType.TITLE,
           title: 'test',
           rpcId: undefined,
         },
         {
-          type: Freestyler.ResponseType.THOUGHT,
+          type: AiAssistance.ResponseType.THOUGHT,
           thought: 'thought 1',
           rpcId: undefined,
         },
         {
-          type: Freestyler.ResponseType.ACTION,
+          type: AiAssistance.ResponseType.ACTION,
           code: 'console.log(\'test\')',
           output: 'undefined',
           canceled: false,
-          rpcId: undefined,
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
+          type: AiAssistance.ResponseType.QUERYING,
           query: 'OBSERVATION: undefined',
         },
         {
-          type: Freestyler.ResponseType.ANSWER,
+          type: AiAssistance.ResponseType.ANSWER,
           text: 'this is the answer',
           suggestions: undefined,
           rpcId: undefined,
@@ -1386,15 +1438,15 @@ ANSWER: this is the answer`,
 
       it('returns an error if scripts are disabled', async () => {
         const execJs = sinon.mock();
-        const agent = new FreestylerAgent({
+        const agent = new StylingAgent({
           aidaClient: getMockClient(),
           createExtensionScope,
           execJs,
         });
-        const responses = await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
-        const actionStep = responses.find(response => response.type === Freestyler.ResponseType.ACTION)!;
+        const responses = await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
+        const actionStep = responses.find(response => response.type === AiAssistance.ResponseType.ACTION)!;
         assert.strictEqual(actionStep.output, 'Error: JavaScript execution is currently disabled.');
-        assert.strictEqual(execJs.getCalls().length, 0);
+        assert.lengthOf(execJs.getCalls(), 0);
       });
     });
 
@@ -1407,17 +1459,17 @@ ANSWER: this is the answer`,
 
       it('returns an error if a script causes a side effect', async () => {
         const execJs =
-            sinon.mock().throws(new Freestyler.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
-        const agent = new FreestylerAgent({
+            sinon.mock().throws(new AiAssistance.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
+        const agent = new StylingAgent({
           aidaClient: getMockClient(),
           createExtensionScope,
           execJs,
         });
-        const responses = await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
-        const actionStep = responses.find(response => response.type === Freestyler.ResponseType.ACTION)!;
+        const responses = await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
+        const actionStep = responses.find(response => response.type === AiAssistance.ResponseType.ACTION)!;
         assert.strictEqual(
             actionStep.output, 'Error: JavaScript execution that modifies the page is currently disabled.');
-        assert.strictEqual(execJs.getCalls().length, 1);
+        assert.lengthOf(execJs.getCalls(), 1);
       });
     });
   });