chrome-devtools-frontend 1.0.1386602 → 1.0.1418433

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

@@ -2,19 +2,19 @@
 // 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 {mockAidaClient} from '../../../testing/AiAssistanceHelpers.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) {
@@ -38,29 +38,44 @@ describeWithEnvironment('FreestylerAgent', () => {
       },
     };
   }
+
+  let element: sinon.SinonStubbedInstance<SDK.DOMModel.DOMNode>;
+  beforeEach(() => {
+    mockHostConfig();
+    element = sinon.createStubInstance(SDK.DOMModel.DOMNode);
+  });
+
   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 +84,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 +110,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 +147,8 @@ c`;
             suggestions: undefined,
           },
       );
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `ANSWER: ${payload}\nACTION\naction\nSTOP`,
               ),
           {
@@ -147,24 +162,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 +187,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 +201,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 +213,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 +226,8 @@ c`;
       const payload = `const data = {
   someKey: "value",
 }`;
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `ACTION\n\`\`\`\n${payload}\n\`\`\`\nSTOP`,
               ),
           {
@@ -227,8 +242,8 @@ c`;
       const payload = `const data = {
   someKey: "value",
 }`;
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `ACTION\n\`\`\`\`\`\njs\n${payload}\n\`\`\`\`\`\nSTOP`,
               ),
           {
@@ -244,8 +259,8 @@ c`;
   someKey: "value",
 }`;
       const thoughtPayload = 'thought';
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               `THOUGHT:${thoughtPayload}\nACTION\n${actionPayload}\nSTOP`,
               ),
           {
@@ -259,8 +274,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 +289,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 +307,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 +325,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 +340,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 +355,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 +366,8 @@ c`;
     });
 
     it('parses a response as an answer', async () => {
-      assert.deepStrictEqual(
-          agent.parseResponse(
+      assert.deepEqual(
+          getParsedTextResponse(
               'This is also an answer',
               ),
           {
@@ -363,8 +378,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 +391,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,
@@ -387,21 +402,11 @@ c`;
   });
 
   describe('describeElement', () => {
-    let element: sinon.SinonStubbedInstance<SDK.DOMModel.DOMNode>;
-
-    beforeEach(() => {
-      element = sinon.createStubInstance(SDK.DOMModel.DOMNode);
-    });
-
-    afterEach(() => {
-      sinon.restore();
-    });
-
     it('should describe an element with no children, siblings, or parent', async () => {
       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 +429,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,11 +460,12 @@ 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
 * Its parent's selector is \`div#grandparentElement\`
+* Its parent is a non element node
 * Its parent has only 1 child element node
 * Its parent has only 1 child text node`;
 
@@ -478,75 +484,66 @@ 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'}, Host.AidaClient.Role.USER).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'}, Host.AidaClient.Role.USER).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'}, Host.AidaClient.Role.USER).metadata?.user_tier,
           3,
       );
     });
 
-    it('structure matches the snapshot', () => {
+    it('structure matches the snapshot', async () => {
       mockHostConfig('test model');
-
-      const agent = new FreestylerAgent({
-        aidaClient: {} as Host.AidaClient.AidaClient,
+      const agent = new StylingAgent({
+        aidaClient: mockAidaClient([[{
+          explanation: 'answer',
+        }]]),
         serverSideLoggingEnabled: true,
       });
       sinon.stub(agent, 'preamble').value('preamble');
-      agent.chatNewHistoryForTesting = [
-        {
-          type: Freestyler.ResponseType.USER_QUERY,
-          query: 'question',
-        },
-        {
-          type: Freestyler.ResponseType.QUERYING,
-          query: 'question',
-        },
-        {
-          type: Freestyler.ResponseType.ANSWER,
-          text: 'answer',
-        },
-      ];
-      assert.deepStrictEqual(
-          agent.buildRequest({
-            input: 'test input',
-          }),
+      await Array.fromAsync(agent.run('question', {selected: null}));
+
+      assert.deepEqual(
+          agent.buildRequest(
+              {
+                text: 'test input',
+              },
+              Host.AidaClient.Role.USER),
           {
-            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: 'QUERY: question'}],
               },
               {
-                entity: 2,
-                text: 'ANSWER: answer',
+                role: 2,
+                parts: [{text: 'ANSWER: answer'}],
               },
             ],
             metadata: {
@@ -563,54 +560,74 @@ c`;
           },
       );
     });
-  });
 
-  function mockAidaClient(
-      fetch: (_: Host.AidaClient.AidaRequest, options?: {signal: AbortSignal}) =>
-          AsyncGenerator<Host.AidaClient.AidaResponse, void, void>,
-      ): Host.AidaClient.AidaClient {
-    return {
-      fetch,
-      registerClientEvent: () => Promise.resolve({}),
-    };
-  }
+    it('builds a request with aborted query in history before a real request', async () => {
+      const execJs = sinon.mock().once();
+      execJs.onCall(0).returns('result2');
+      const agent = new StylingAgent({
+        aidaClient: mockAidaClient([
+          [{
+            explanation: `THOUGHT: thought2
+TITLE: title2
+ACTION
+action2
+STOP`
+          }],
+          [{explanation: 'answer2'}]
+        ]),
+        createExtensionScope,
+        execJs,
+      });
 
-  describe('run', () => {
-    let element: sinon.SinonStubbedInstance<SDK.DOMModel.DOMNode>;
-    beforeEach(() => {
-      mockHostConfig();
-      element = sinon.createStubInstance(SDK.DOMModel.DOMNode);
+      const controller = new AbortController();
+      controller.abort();
+      await Array.fromAsync(agent.run('test', {
+        selected: null,
+        signal: controller.signal,
+      }));
+      await Array.fromAsync(agent.run('test2', {selected: null}));
+
+      const request = agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER);
+      assert.deepEqual(request.current_message?.parts[0], {text: 'test input'});
+      assert.deepEqual(request.historical_contexts, [
+        {
+          parts: [{text: 'QUERY: test2'}],
+          role: 1,
+        },
+        {
+          role: 2,
+          parts: [{text: 'THOUGHT: thought2\nTITLE: title2\nACTION\naction2\nSTOP'}],
+        },
+        {
+          role: 1,
+          parts: [{text: 'OBSERVATION: result2'}],
+        },
+        {
+          role: 2,
+          parts: [{text: 'ANSWER: answer2'}],
+        },
+      ]);
     });
+  });
 
+  describe('run', () => {
     describe('side effect handling', () => {
       it('calls confirmSideEffect when the code execution contains a side effect', async () => {
         const promise = Promise.withResolvers();
         const stub = sinon.stub().returns(promise);
-
-        let count = 0;
-        async function* generateActionAndAnswer() {
-          if (count === 0) {
-            yield {
+        const execJs =
+            sinon.mock().throws(new AiAssistance.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
+        const agent = new StylingAgent({
+          aidaClient: mockAidaClient([
+            [{
               explanation: `ACTION
-              $0.style.backgroundColor = 'red'
-              STOP`,
-              metadata: {},
-              completed: true,
-            };
-          } else {
-            yield {
+$0.style.backgroundColor = 'red'
+STOP`,
+            }],
+            [{
               explanation: 'ANSWER: This is the answer',
-              metadata: {},
-              completed: true,
-            };
-          }
-
-          count++;
-        }
-        const execJs =
-            sinon.mock().throws(new Freestyler.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
-        const agent = new FreestylerAgent({
-          aidaClient: mockAidaClient(generateActionAndAnswer),
+            }]
+          ]),
           createExtensionScope,
           confirmSideEffectForTest: stub,
           execJs,
@@ -618,7 +635,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}));
       });
@@ -626,147 +643,141 @@ c`;
       it('calls execJs with allowing side effects when confirmSideEffect resolves to true', async () => {
         const promise = Promise.withResolvers();
         const stub = sinon.stub().returns(promise);
-        let count = 0;
-        async function* generateActionAndAnswer() {
-          if (count === 0) {
-            yield {
-              explanation: `ACTION
-              $0.style.backgroundColor = 'red'
-              STOP`,
-              metadata: {},
-              completed: true,
-            };
-          } else {
-            yield {
-              explanation: 'ANSWER: This is the answer',
-              metadata: {},
-              completed: true,
-            };
-          }
-
-          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({
-          aidaClient: mockAidaClient(generateActionAndAnswer),
+        const agent = new StylingAgent({
+          aidaClient: mockAidaClient([
+            [{
+              explanation: `ACTION
+            $0.style.backgroundColor = 'red'
+            STOP`,
+            }],
+            [{
+              explanation: 'ANSWER: This is the answer',
+            }]
+          ]),
           createExtensionScope,
           confirmSideEffectForTest: stub,
           execJs,
 
         });
         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}));
       });
 
       it('returns side effect error when confirmSideEffect resolves to false', async () => {
         const promise = Promise.withResolvers();
         const stub = sinon.stub().returns(promise);
-        let count = 0;
-        async function* generateActionAndAnswer() {
-          if (count === 0) {
-            yield {
+        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([
+            [{
               explanation: `ACTION
-              $0.style.backgroundColor = 'red'
-              STOP`,
-              metadata: {},
-              completed: true,
-            };
-          } else {
-            yield {
+$0.style.backgroundColor = 'red'
+STOP`,
+            }],
+            [{
               explanation: 'ANSWER: This is the answer',
-              metadata: {},
-              completed: true,
-            };
-          }
-
-          count++;
-        }
-        const execJs = sinon.mock().twice();
-        execJs.onCall(0).throws(new Freestyler.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
-        const agent = new FreestylerAgent({
-          aidaClient: mockAidaClient(generateActionAndAnswer),
+            }]
+          ]),
           createExtensionScope,
           confirmSideEffectForTest: stub,
           execJs,
 
         });
         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);
+        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([[{
+            explanation: `ACTION
+$0.style.backgroundColor = 'red'
+STOP`,
+          }]]),
+          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);
       });
     });
 
     describe('long `Observation` text handling', () => {
       it('errors with too long input', async () => {
-        let count = 0;
-        async function* generateActionAndAnswer() {
-          if (count === 0) {
-            yield {
+        const execJs = sinon.mock().returns(new Array(10_000).fill('<div>...</div>').join());
+        const agent = new StylingAgent({
+          aidaClient: mockAidaClient([
+            [{
               explanation: `ACTION
-              $0.style.backgroundColor = 'red'
-              STOP`,
-              metadata: {},
-              completed: true,
-            };
-          } else {
-            yield {
+$0.style.backgroundColor = 'red';
+STOP`,
+            }],
+            [{
               explanation: 'ANSWER: This is the answer',
-              metadata: {},
-              completed: true,
-            };
-          }
-          count++;
-        }
-        const execJs = sinon.mock().returns(new Array(10_000).fill('<div>...</div>').join());
-        const agent = new FreestylerAgent({
-          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)!;
-        assert(actionStep.output.includes('Error: Output exceeded the maximum allowed length.'));
+        assert(actionStep.output!.includes('Error: Output exceeded the maximum allowed length.'));
       });
     });
 
     it('generates an answer immediately', async () => {
-      async function* generateAnswer() {
-        yield {
-          explanation: 'ANSWER: this is the answer',
-          metadata: {},
-          completed: true,
-        };
-      }
-
       const execJs = sinon.spy();
-      const agent = new FreestylerAgent({
-        aidaClient: mockAidaClient(generateAnswer),
+      const agent = new StylingAgent({
+        aidaClient: mockAidaClient([[{explanation: 'ANSWER: this is the answer'}]]),
         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,111 +787,106 @@ c`;
           ],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
+          type: AiAssistance.ResponseType.QUERYING,
         },
         {
-          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.buildRequest({text: ''}, Host.AidaClient.Role.USER).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: 'ANSWER: this is the answer',
+          role: 2,
+          parts: [{text: 'ANSWER: this is the answer'}],
         },
       ]);
     });
 
-    it('correctly handles chat_history in AIDA requests', async () => {
-      const requests: Host.AidaClient.AidaRequest[] = [];
-
-      let i = 0;
-      async function* generateAnswer(request: Host.AidaClient.AidaRequest) {
-        requests.push(request);
-        if (i !== 0) {
-          yield {
-            explanation: 'ANSWER: this is the actual answer',
-            metadata: {},
-            completed: true,
-          };
-          return;
-        }
-        yield {
+    it('correctly handles historical_contexts in AIDA requests', async () => {
+      const execJs = sinon.mock().once();
+      execJs.onCall(0).returns('test data');
+      const aidaClient = mockAidaClient([
+        [{
           explanation: `THOUGHT: I am thinking.
 TITLE: thinking
 ACTION
 const data = {"test": "observation"};
 STOP`,
-          metadata: {},
-          completed: false,
-        };
-        i++;
-      }
-
-      const execJs = sinon.mock().once();
-      execJs.onCall(0).returns('test data');
-      const agent = new FreestylerAgent({
-        aidaClient: mockAidaClient(generateAnswer),
+        }],
+        [{
+          explanation: 'ANSWER: this is the actual answer',
+        }]
+      ]);
+      const agent = new StylingAgent({
+        aidaClient,
         createExtensionScope,
         execJs,
       });
 
-      await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
+      await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
+
+      const requests: Host.AidaClient.AidaRequest[] = (aidaClient.fetch as sinon.SinonStub).args.map(arg => arg[0]);
 
       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 () => {
-      async function* generateAnswer() {
-        yield {
+      const agent = new StylingAgent({
+        aidaClient: mockAidaClient([[{
           explanation: 'ANSWER: this is the answer',
           metadata: {
             rpcGlobalId: 123,
           },
-          completed: true,
-        };
-      }
-
-      const agent = new FreestylerAgent({
-        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 +896,10 @@ STOP`,
           ],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
+          type: AiAssistance.ResponseType.QUERYING,
         },
         {
-          type: Freestyler.ResponseType.ANSWER,
+          type: AiAssistance.ResponseType.ANSWER,
           text: 'this is the answer',
           suggestions: undefined,
           rpcId: 123,
@@ -903,34 +908,32 @@ STOP`,
     });
 
     it('throws an error based on the attribution metadata including RecitationAction.BLOCK', async () => {
-      async function* generateAnswer() {
-        yield {
-          explanation: 'ANSWER: this is the answer',
-          metadata: {
-            rpcGlobalId: 123,
-            attributionMetadata: [{
-              attributionAction: Host.AidaClient.RecitationAction.BLOCK,
-              citations: [],
-            }],
+      const agent = new StylingAgent({
+        aidaClient: mockAidaClient([[
+          {
+            explanation: 'ANSWER: this is the answer',
           },
-          completed: true,
-        };
-      }
-
-      const agent = new FreestylerAgent({
-        aidaClient: mockAidaClient(generateAnswer),
+          {
+            explanation: 'ANSWER: this is another answer',
+            metadata: {
+              attributionMetadata: {
+                attributionAction: Host.AidaClient.RecitationAction.BLOCK,
+                citations: [],
+              },
+            },
+          }
+        ]]),
         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,46 +943,43 @@ STOP`,
           ],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
+          type: AiAssistance.ResponseType.QUERYING,
         },
         {
-          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,
         },
       ]);
     });
 
     it('does not throw an error based on attribution metadata not including RecitationAction.BLOCK', async () => {
-      async function* generateAnswer() {
-        yield {
+      const agent = new StylingAgent({
+        aidaClient: mockAidaClient([[{
           explanation: 'ANSWER: this is the answer',
           metadata: {
             rpcGlobalId: 123,
-            attributionMetadata: [{
+            attributionMetadata: {
               attributionAction: Host.AidaClient.RecitationAction.ACTION_UNSPECIFIED,
               citations: [],
-            }],
+            },
           },
-          completed: true,
-        };
-      }
-
-      const agent = new FreestylerAgent({
-        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 +989,10 @@ STOP`,
           ],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
+          type: AiAssistance.ResponseType.QUERYING,
         },
         {
-          type: Freestyler.ResponseType.ANSWER,
+          type: AiAssistance.ResponseType.ANSWER,
           text: 'this is the answer',
           suggestions: undefined,
           rpcId: 123,
@@ -1003,63 +1002,47 @@ STOP`,
 
     it('should execute an action only once even when the partial response contains an action', async () => {
       const execJs = sinon.spy();
-      async function* generatePartialAndFullAction() {
-        yield {
-          explanation: `THOUGHT: I am thinking.
+      const agent = new StylingAgent({
+        aidaClient: mockAidaClient([[
+          {
+            explanation: `THOUGHT: I am thinking.
 
 ACTION
 console.log('hel
           `,
-          metadata: {},
-          completed: false,
-        };
-
-        sinon.assert.notCalled(execJs);
-        yield {
-          explanation: `THOUGHT: I am thinking.
+          },
+          {
+            explanation: `THOUGHT: I am thinking.
 
 ACTION
 console.log('hello');
 STOP
           `,
-          metadata: {},
-          completed: true,
-        };
-      }
-
-      const agent = new FreestylerAgent({
-        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\');');
     });
 
     it('generates a response if nothing is returned', async () => {
-      async function* generateNothing() {
-        yield {
-          explanation: '',
-          metadata: {},
-          completed: true,
-        };
-      }
-
       const execJs = sinon.spy();
-      const agent = new FreestylerAgent({
-        aidaClient: mockAidaClient(generateNothing),
+      const agent = new StylingAgent({
+        aidaClient: mockAidaClient([[{explanation: ''}]]),
         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,60 +1052,48 @@ STOP
           ],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
+          type: AiAssistance.ResponseType.QUERYING,
         },
         {
-          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.isUndefined(agent.buildRequest({text: ''}, Host.AidaClient.Role.USER).historical_contexts);
     });
 
     it('generates an action response if action and answer both present', async () => {
-      let i = 0;
-      async function* generateNothing() {
-        if (i !== 0) {
-          yield {
-            explanation: 'ANSWER: this is the actual answer',
-            metadata: {},
-            completed: true,
-          };
-          return;
-        }
-        yield {
-          explanation: `THOUGHT: I am thinking.
-
-ACTION
-console.log('hello');
-STOP
-
-ANSWER: this is the answer`,
-          metadata: {},
-          completed: false,
-        };
-        i++;
-      }
-
       const execJs = sinon.mock().once();
       execJs.onCall(0).returns('hello');
-      const agent = new FreestylerAgent({
-        aidaClient: mockAidaClient(generateNothing),
+      const agent = new StylingAgent({
+        aidaClient: mockAidaClient([
+          [{
+            explanation: `THOUGHT: I am thinking.
+
+          ACTION
+          console.log('hello');
+          STOP
+
+          ANSWER: this is the answer`,
+          }],
+          [{
+            explanation: 'ANSWER: this is the actual answer',
+            metadata: {},
+          }]
+        ]),
         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 +1103,24 @@ ANSWER: this is the answer`,
           ],
         },
         {
-          type: Freestyler.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
+          type: AiAssistance.ResponseType.QUERYING,
         },
         {
-          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,
-          query: 'OBSERVATION: hello',
+          type: AiAssistance.ResponseType.QUERYING,
         },
         {
-          type: Freestyler.ResponseType.ANSWER,
+          type: AiAssistance.ResponseType.ANSWER,
           text: 'this is the actual answer',
           suggestions: undefined,
           rpcId: undefined,
@@ -1162,95 +1130,82 @@ ANSWER: this is the answer`,
     });
 
     it('generates history for multiple actions', async () => {
-      let count = 0;
-      async function* generateMultipleTimes() {
-        if (count === 3) {
-          yield {
-            explanation: 'ANSWER: this is the answer',
-            metadata: {},
-            completed: true,
-          };
-          return;
-        }
-        count++;
-        yield {
-          explanation: `THOUGHT: thought ${count}\nTITLE:test\nACTION\nconsole.log('test')\nSTOP\n`,
-          metadata: {},
-          completed: false,
-        };
-      }
-
       const execJs = sinon.spy(async () => 'undefined');
-      const agent = new FreestylerAgent({
-        aidaClient: mockAidaClient(generateMultipleTimes),
+      const agent = new StylingAgent({
+        aidaClient: mockAidaClient([
+          [{
+            explanation: 'THOUGHT: thought 1\nTITLE:test\nACTION\nconsole.log(\'test\')\nSTOP\n',
+          }],
+          [{
+            explanation: 'THOUGHT: thought 2\nTITLE:test\nACTION\nconsole.log(\'test\')\nSTOP\n',
+          }],
+          [{
+            explanation: 'THOUGHT: thought 3\nTITLE:test\nACTION\nconsole.log(\'test\')\nSTOP\n',
+          }],
+          [{
+            explanation: 'ANSWER: this is the answer',
+          }]
+        ]),
         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.buildRequest({text: ''}, Host.AidaClient.Role.USER).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: 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'}],
         },
       ]);
     });
 
     it('stops when aborted', async () => {
-      let count = 0;
-      async function* generateAndAbort(_: unknown, options?: {signal: AbortSignal}) {
-        if (options?.signal.aborted) {
-          throw new Host.AidaClient.AidaAbortError();
-        }
-        if (count === 3) {
-          yield {
-            explanation: 'ANSWER: this is the answer',
-            metadata: {},
-            completed: true,
-          };
-          return;
-        }
-        count++;
-        yield {
-          explanation: `THOUGHT: thought ${count}\nACTION\nconsole.log('test')\nSTOP\n`,
-          metadata: {},
-          completed: false,
-        };
-      }
-
       const execJs = sinon.spy();
-      const agent = new FreestylerAgent({
-        aidaClient: mockAidaClient(generateAndAbort),
+      const agent = new StylingAgent({
+        aidaClient: mockAidaClient([
+          [{
+            explanation: 'THOUGHT: thought 1\nTITLE:test\nACTION\nconsole.log(\'test\')\nSTOP\n',
+          }],
+          [{
+            explanation: 'THOUGHT: thought 2\nTITLE:test\nACTION\nconsole.log(\'test\')\nSTOP\n',
+          }],
+          [{
+            explanation: 'THOUGHT: thought 3\nTITLE:test\nACTION\nconsole.log(\'test\')\nSTOP\n',
+          }],
+          [{
+            explanation: 'ANSWER: this is the answer',
+          }]
+        ]),
         createExtensionScope,
         execJs,
       });
@@ -1258,125 +1213,24 @@ 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, []);
-    });
-  });
-
-  describe('history', () => {
-    let element: sinon.SinonStubbedInstance<SDK.DOMModel.DOMNode>;
-    beforeEach(() => {
-      mockHostConfig();
-      element = sinon.createStubInstance(SDK.DOMModel.DOMNode);
-      // @ts-ignore
-      setAiAssistancePersistentHistory(true);
-    });
-
-    it('stores history via AiHistoryStorage', async () => {
-      let count = 0;
-      async function* generateMultipleTimes() {
-        if (count === 1) {
-          yield {
-            explanation: 'ANSWER: this is the answer',
-            metadata: {},
-            completed: true,
-          };
-          return;
-        }
-        count++;
-        yield {
-          explanation: `THOUGHT: thought ${count}\nTITLE:test\nACTION\nconsole.log('test')\nSTOP\n`,
-          metadata: {},
-          completed: false,
-        };
-      }
-      const historyStub = sinon.stub(Freestyler.AiHistoryStorage.instance(), 'upsertHistoryEntry');
-      const execJs = sinon.spy(async () => 'undefined');
-      const agent = new FreestylerAgent({
-        aidaClient: mockAidaClient(generateMultipleTimes),
-        createExtensionScope,
-        execJs,
-      });
-
-      await Array.fromAsync(agent.run('test', {selected: new Freestyler.NodeContext(element)}));
-
-      assert.deepStrictEqual(historyStub.lastCall.args[0].history, [
-        {
-          type: Freestyler.ResponseType.USER_QUERY,
-          query: 'test',
-        },
-        {
-          type: Freestyler.ResponseType.CONTEXT,
-          title: 'Analyzing the prompt',
-          details: [{
-            text: '* Its selector is `undefined`',
-            title: 'Data used',
-          }],
-        },
-        {
-          type: Freestyler.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
-        },
-        {
-          type: Freestyler.ResponseType.TITLE,
-          title: 'test',
-          rpcId: undefined,
-        },
-        {
-          type: Freestyler.ResponseType.THOUGHT,
-          thought: 'thought 1',
-          rpcId: undefined,
-        },
-        {
-          type: Freestyler.ResponseType.ACTION,
-          code: 'console.log(\'test\')',
-          output: 'undefined',
-          canceled: false,
-          rpcId: undefined,
-        },
-        {
-          type: Freestyler.ResponseType.QUERYING,
-          query: 'OBSERVATION: undefined',
-        },
-        {
-          type: Freestyler.ResponseType.ANSWER,
-          text: 'this is the answer',
-          suggestions: undefined,
-          rpcId: undefined,
-        },
-      ]);
+      assert.isUndefined(agent.buildRequest({text: ''}, Host.AidaClient.Role.USER).historical_contexts);
     });
   });
 
   describe('HostConfigFreestylerExecutionMode', () => {
-    let element: sinon.SinonStubbedInstance<SDK.DOMModel.DOMNode>;
-    beforeEach(() => {
-      element = sinon.createStubInstance(SDK.DOMModel.DOMNode);
-    });
-
     function getMockClient() {
-      let count = 0;
-      async function* generateActionAndAnswer() {
-        if (count === 0) {
-          yield {
-            explanation: `ACTION
-            $0.style.backgroundColor = 'red'
-            STOP`,
-            metadata: {},
-            completed: true,
-          };
-        } else {
-          yield {
-            explanation: 'ANSWER: This is the answer',
-            metadata: {},
-            completed: true,
-          };
-        }
-
-        count++;
-      }
-      return mockAidaClient(generateActionAndAnswer);
+      return mockAidaClient([
+        [{
+          explanation: `ACTION
+$0.style.backgroundColor = 'red'
+STOP`,
+        }],
+        [{
+          explanation: 'ANSWER: This is the answer',
+        }]
+      ]);
     }
 
     describe('NO_SCRIPTS', () => {
@@ -1386,15 +1240,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 +1261,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);
       });
     });
   });