chrome-devtools-frontend 1.0.1402249 → 1.0.1418433

package/front_end/panels/ai_assistance/agents/StylingAgent.test.ts

@@ -5,6 +5,7 @@
 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,
@@ -37,6 +38,13 @@ describeWithEnvironment('StylingAgent', () => {
       },
     };
   }
+
+  let element: sinon.SinonStubbedInstance<SDK.DOMModel.DOMNode>;
+  beforeEach(() => {
+    mockHostConfig();
+    element = sinon.createStubInstance(SDK.DOMModel.DOMNode);
+  });
+
   describe('parseResponse', () => {
     const agent = new StylingAgent({
       aidaClient: {} as Host.AidaClient.AidaClient,
@@ -394,16 +402,6 @@ 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);
@@ -467,6 +465,7 @@ c`;
 * 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`;
 
@@ -489,7 +488,7 @@ c`;
         aidaClient: {} as Host.AidaClient.AidaClient,
       });
       assert.strictEqual(
-          agent.buildRequest({text: 'test input'}).options?.model_id,
+          agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER).options?.model_id,
           'test model',
       );
     });
@@ -500,7 +499,7 @@ c`;
         aidaClient: {} as Host.AidaClient.AidaClient,
       });
       assert.strictEqual(
-          agent.buildRequest({text: 'test input'}).options?.temperature,
+          agent.buildRequest({text: 'test input'}, Host.AidaClient.Role.USER).options?.temperature,
           1,
       );
     });
@@ -511,37 +510,28 @@ c`;
         aidaClient: {} as Host.AidaClient.AidaClient,
       });
       assert.strictEqual(
-          agent.buildRequest({text: '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 StylingAgent({
-        aidaClient: {} as Host.AidaClient.AidaClient,
+        aidaClient: mockAidaClient([[{
+          explanation: 'answer',
+        }]]),
         serverSideLoggingEnabled: true,
       });
       sinon.stub(agent, 'preamble').value('preamble');
-      agent.chatNewHistoryForTesting = [
-        {
-          type: AiAssistance.ResponseType.USER_QUERY,
-          query: 'question',
-        },
-        {
-          type: AiAssistance.ResponseType.QUERYING,
-          query: 'question',
-        },
-        {
-          type: AiAssistance.ResponseType.ANSWER,
-          text: 'answer',
-        },
-      ];
+      await Array.fromAsync(agent.run('question', {selected: null}));
+
       assert.deepEqual(
-          agent.buildRequest({
-            text: 'test input',
-          }),
+          agent.buildRequest(
+              {
+                text: 'test input',
+              },
+              Host.AidaClient.Role.USER),
           {
             current_message: {role: Host.AidaClient.Role.USER, parts: [{text: 'test input'}]},
             client: 'CHROME_DEVTOOLS',
@@ -549,7 +539,7 @@ c`;
             historical_contexts: [
               {
                 role: 1,
-                parts: [{text: 'question'}],
+                parts: [{text: 'QUERY: question'}],
               },
               {
                 role: 2,
@@ -570,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 {
-              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().throws(new AiAssistance.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
         const agent = new StylingAgent({
-          aidaClient: mockAidaClient(generateActionAndAnswer),
+          aidaClient: mockAidaClient([
+            [{
+              explanation: `ACTION
+$0.style.backgroundColor = 'red'
+STOP`,
+            }],
+            [{
+              explanation: 'ANSWER: This is the answer',
+            }]
+          ]),
           createExtensionScope,
           confirmSideEffectForTest: stub,
           execJs,
@@ -633,31 +643,20 @@ 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 AiAssistance.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
         execJs.onCall(1).resolves('value');
         const agent = new StylingAgent({
-          aidaClient: mockAidaClient(generateActionAndAnswer),
+          aidaClient: mockAidaClient([
+            [{
+              explanation: `ACTION
+            $0.style.backgroundColor = 'red'
+            STOP`,
+            }],
+            [{
+              explanation: 'ANSWER: This is the answer',
+            }]
+          ]),
           createExtensionScope,
           confirmSideEffectForTest: stub,
           execJs,
@@ -673,30 +672,19 @@ c`;
       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 {
-              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().once();
         execJs.onCall(0).throws(new AiAssistance.SideEffectError('EvalError: Possible side-effect in debug-evaluate'));
         const agent = new StylingAgent({
-          aidaClient: mockAidaClient(generateActionAndAnswer),
+          aidaClient: mockAidaClient([
+            [{
+              explanation: `ACTION
+$0.style.backgroundColor = 'red'
+STOP`,
+            }],
+            [{
+              explanation: 'ANSWER: This is the answer',
+            }]
+          ]),
           createExtensionScope,
           confirmSideEffectForTest: stub,
           execJs,
@@ -713,20 +701,15 @@ c`;
 
       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),
+          aidaClient: mockAidaClient([[{
+            explanation: `ACTION
+$0.style.backgroundColor = 'red'
+STOP`,
+          }]]),
           createExtensionScope,
           confirmSideEffectForTest: sinon.stub().returns(sideEffectConfirmationPromise),
           execJs,
@@ -754,31 +737,20 @@ c`;
 
     describe('long `Observation` text handling', () => {
       it('errors with too long input', async () => {
-        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().returns(new Array(10_000).fill('<div>...</div>').join());
         const agent = new StylingAgent({
-          aidaClient: mockAidaClient(generateActionAndAnswer),
+          aidaClient: mockAidaClient([
+            [{
+              explanation: `ACTION
+$0.style.backgroundColor = 'red';
+STOP`,
+            }],
+            [{
+              explanation: 'ANSWER: This is the answer',
+            }]
+          ]),
           createExtensionScope,
           execJs,
-
         });
 
         const result = await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
@@ -787,22 +759,14 @@ c`;
         });
         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 StylingAgent({
-        aidaClient: mockAidaClient(generateAnswer),
+        aidaClient: mockAidaClient([[{explanation: 'ANSWER: this is the answer'}]]),
         execJs,
       });
 
@@ -824,7 +788,6 @@ c`;
         },
         {
           type: AiAssistance.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
           type: AiAssistance.ResponseType.ANSWER,
@@ -834,7 +797,7 @@ c`;
         },
       ]);
       sinon.assert.notCalled(execJs);
-      assert.deepEqual(agent.chatHistoryForTesting, [
+      assert.deepEqual(agent.buildRequest({text: ''}, Host.AidaClient.Role.USER).historical_contexts, [
         {
           role: 1,
           parts: [{text: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test'}],
@@ -847,48 +810,38 @@ c`;
     });
 
     it('correctly handles historical_contexts 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 {
+      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');
+        }],
+        [{
+          explanation: 'ANSWER: this is the actual answer',
+        }]
+      ]);
       const agent = new StylingAgent({
-        aidaClient: mockAidaClient(generateAnswer),
+        aidaClient,
         createExtensionScope,
         execJs,
       });
 
       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].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'},
+          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(
@@ -915,18 +868,13 @@ STOP`,
     });
 
     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 StylingAgent({
-        aidaClient: mockAidaClient(generateAnswer),
+        }]]),
         execJs: sinon.spy(),
 
       });
@@ -949,7 +897,6 @@ STOP`,
         },
         {
           type: AiAssistance.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
           type: AiAssistance.ResponseType.ANSWER,
@@ -961,19 +908,21 @@ 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,
-          },
-          completed: false,
-        };
-        throw new Host.AidaClient.AidaBlockError();
-      }
-
       const agent = new StylingAgent({
-        aidaClient: mockAidaClient(generateAnswer),
+        aidaClient: mockAidaClient([[
+          {
+            explanation: 'ANSWER: this is the answer',
+          },
+          {
+            explanation: 'ANSWER: this is another answer',
+            metadata: {
+              attributionMetadata: {
+                attributionAction: Host.AidaClient.RecitationAction.BLOCK,
+                citations: [],
+              },
+            },
+          }
+        ]]),
         execJs: sinon.spy(),
       });
 
@@ -995,7 +944,6 @@ STOP`,
         },
         {
           type: AiAssistance.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
           text: 'this is the answer',
@@ -1009,8 +957,8 @@ STOP`,
     });
 
     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,
@@ -1019,12 +967,7 @@ STOP`,
               citations: [],
             },
           },
-          completed: true,
-        };
-      }
-
-      const agent = new StylingAgent({
-        aidaClient: mockAidaClient(generateAnswer),
+        }]]),
         execJs: sinon.spy(),
 
       });
@@ -1047,7 +990,6 @@ STOP`,
         },
         {
           type: AiAssistance.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
           type: AiAssistance.ResponseType.ANSWER,
@@ -1060,32 +1002,24 @@ 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 StylingAgent({
-        aidaClient: mockAidaClient(generatePartialAndFullAction),
+          }
+        ]]),
         createExtensionScope,
         execJs,
       });
@@ -1096,17 +1030,9 @@ STOP
     });
 
     it('generates a response if nothing is returned', async () => {
-      async function* generateNothing() {
-        yield {
-          explanation: '',
-          metadata: {},
-          completed: true,
-        };
-      }
-
       const execJs = sinon.spy();
       const agent = new StylingAgent({
-        aidaClient: mockAidaClient(generateNothing),
+        aidaClient: mockAidaClient([[{explanation: ''}]]),
         execJs,
       });
       const responses = await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
@@ -1127,7 +1053,6 @@ STOP
         },
         {
           type: AiAssistance.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
           type: AiAssistance.ResponseType.ERROR,
@@ -1135,38 +1060,28 @@ STOP
         },
       ]);
       sinon.assert.notCalled(execJs);
-      assert.deepEqual(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 StylingAgent({
-        aidaClient: mockAidaClient(generateNothing),
+        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,
 
@@ -1189,7 +1104,6 @@ ANSWER: this is the answer`,
         },
         {
           type: AiAssistance.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
         },
         {
           type: AiAssistance.ResponseType.THOUGHT,
@@ -1204,7 +1118,6 @@ ANSWER: this is the answer`,
         },
         {
           type: AiAssistance.ResponseType.QUERYING,
-          query: 'OBSERVATION: hello',
         },
         {
           type: AiAssistance.ResponseType.ANSWER,
@@ -1217,27 +1130,22 @@ 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 StylingAgent({
-        aidaClient: mockAidaClient(generateMultipleTimes),
+        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,
 
@@ -1245,7 +1153,7 @@ ANSWER: this is the answer`,
 
       await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
 
-      assert.deepEqual(agent.chatHistoryForTesting, [
+      assert.deepEqual(agent.buildRequest({text: ''}, Host.AidaClient.Role.USER).historical_contexts, [
         {
           role: 1,
           parts: [{text: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test'}],
@@ -1282,30 +1190,22 @@ 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 StylingAgent({
-        aidaClient: mockAidaClient(generateAndAbort),
+        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,
       });
@@ -1315,120 +1215,22 @@ ANSWER: this is the answer`,
       await Array.fromAsync(
           agent.run('test', {selected: new AiAssistance.NodeContext(element), signal: controller.signal}));
 
-      assert.deepEqual(agent.chatHistoryForTesting, []);
-    });
-  });
-
-  describe('history', () => {
-    let element: sinon.SinonStubbedInstance<SDK.DOMModel.DOMNode>;
-    beforeEach(() => {
-      mockHostConfig();
-      element = sinon.createStubInstance(SDK.DOMModel.DOMNode);
-    });
-
-    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(AiAssistance.AiHistoryStorage.instance(), 'upsertHistoryEntry');
-      const execJs = sinon.spy(async () => 'undefined');
-      const agent = new StylingAgent({
-        aidaClient: mockAidaClient(generateMultipleTimes),
-        createExtensionScope,
-        execJs,
-      });
-
-      await Array.fromAsync(agent.run('test', {selected: new AiAssistance.NodeContext(element)}));
-
-      assert.deepEqual(historyStub.lastCall.args[0].history, [
-        {
-          type: AiAssistance.ResponseType.USER_QUERY,
-          query: 'test',
-        },
-        {
-          type: AiAssistance.ResponseType.CONTEXT,
-          title: 'Analyzing the prompt',
-          details: [{
-            text: '* Its selector is `undefined`',
-            title: 'Data used',
-          }],
-        },
-        {
-          type: AiAssistance.ResponseType.QUERYING,
-          query: '# Inspected element\n\n* Its selector is `undefined`\n\n# User request\n\nQUERY: test',
-        },
-        {
-          type: AiAssistance.ResponseType.TITLE,
-          title: 'test',
-          rpcId: undefined,
-        },
-        {
-          type: AiAssistance.ResponseType.THOUGHT,
-          thought: 'thought 1',
-          rpcId: undefined,
-        },
-        {
-          type: AiAssistance.ResponseType.ACTION,
-          code: 'console.log(\'test\')',
-          output: 'undefined',
-          canceled: false,
-        },
-        {
-          type: AiAssistance.ResponseType.QUERYING,
-          query: 'OBSERVATION: undefined',
-        },
-        {
-          type: AiAssistance.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', () => {