Store context in AiAgent

The StylingAgent needs to evaluate JS code in the correct execution
context. Before this CL the currently selected element was used both
for determining the execution context, and for evaluating `$0`. This can
be wrong for external requests with cross-proccess iframes, and also if the user selects a different UI element while the conversation is actively running.

Both problems are fixed by keeping the context constant while the
conversation is being executed. Selecting a new element in the UI and
restarting the conversation with a new context still works after this
CL.

Bug: 422019184, 428690606, 419246352
Change-Id: Ica9abdc2e7c165f8b6577817b78cec733620e7a1
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/6687895
Reviewed-by: Alex Rudenko <[email protected]>
Reviewed-by: Ergün Erdoğmuş <[email protected]>
Commit-Queue: Wolfgang Beyer <[email protected]>

Author: wolfib

front_end/models/ai_assistance/agents/AiAgent.ts

@@ -268,7 +268,14 @@ export abstract class AiAgent<T> {
    * historical conversations.
    */
   #origin?: string;
-  #context?: ConversationContext<T>;
+
+  /**
+   * `context` does not change during `AiAgent.run()`, ensuring that calls to JS
+   * have the correct `context`. We don't want element selection by the user to
+   * change the `context` during an `AiAgent.run()`.
+   */
+  protected context?: ConversationContext<T>;
+
   #id: string = crypto.randomUUID();
   #history: Host.AidaClient.Content[] = [];
 
@@ -418,13 +425,14 @@ export abstract class AiAgent<T> {
           multimodalInput?: MultimodalInput): AsyncGenerator<ResponseData, void, void> {
     await options.selected?.refresh();
 
-    // First context set on the agent determines its origin from now on.
-    if (options.selected && this.#origin === undefined && options.selected) {
-      this.#origin = options.selected.getOrigin();
-    }
-    // Remember if the context that is set.
-    if (options.selected && !this.#context) {
-      this.#context = options.selected;
+    if (options.selected) {
+      // First context set on the agent determines its origin from now on.
+      if (this.#origin === undefined) {
+        this.#origin = options.selected.getOrigin();
+      }
+      if (options.selected.isOriginAllowed(this.#origin)) {
+        this.context = options.selected;
+      }
     }
 
     const enhancedQuery = await this.enhanceQuery(initialQuery, options.selected, multimodalInput?.type);

front_end/models/ai_assistance/agents/StylingAgent.ts

@@ -182,16 +182,20 @@ const MULTIMODAL_ENHANCEMENT_PROMPTS: Record<MultimodalInputType, string> = {
 };
 
 async function executeJsCode(
-    functionDeclaration: string, {throwOnSideEffect}: {throwOnSideEffect: boolean}): Promise<string> {
-  const selectedNode = UI.Context.Context.instance().flavor(SDK.DOMModel.DOMNode);
-  const target = selectedNode?.domModel().target() ?? UI.Context.Context.instance().flavor(SDK.Target.Target);
+    functionDeclaration: string,
+    {throwOnSideEffect, contextNode}: {throwOnSideEffect: boolean, contextNode: SDK.DOMModel.DOMNode|null}):
+    Promise<string> {
+  if (!contextNode) {
+    throw new Error('Cannot execute JavaScript because of missing context node');
+  }
+  const target = contextNode.domModel().target() ?? UI.Context.Context.instance().flavor(SDK.Target.Target);
 
   if (!target) {
     throw new Error('Target is not found for executing code');
   }
 
   const resourceTreeModel = target.model(SDK.ResourceTreeModel.ResourceTreeModel);
-  const frameId = selectedNode?.frameId() ?? resourceTreeModel?.mainFrame?.id;
+  const frameId = contextNode.frameId() ?? resourceTreeModel?.mainFrame?.id;
 
   if (!frameId) {
     throw new Error('Main frame is not found for executing code');
@@ -211,19 +215,12 @@ async function executeJsCode(
     return formatError('Cannot evaluate JavaScript because the execution is paused on a breakpoint.');
   }
 
-  const result = await executionContext.evaluate(
-      {
-        expression: '$0',
-        returnByValue: false,
-        includeCommandLineAPI: true,
-      },
-      false, false);
-
-  if ('error' in result) {
-    return formatError('Cannot find $0');
+  const remoteObject = await contextNode.resolveToObject(undefined, executionContextId);
+  if (!remoteObject) {
+    throw new Error('Cannot execute JavaScript because remote object cannot be resolved');
   }
 
-  return await EvaluateAction.execute(functionDeclaration, [result.object], executionContext, {throwOnSideEffect});
+  return await EvaluateAction.execute(functionDeclaration, [remoteObject], executionContext, {throwOnSideEffect});
 }
 
 const MAX_OBSERVATION_BYTE_LENGTH = 25_000;
@@ -640,7 +637,7 @@ const data = {
       const result = await Promise.race([
         this.#execJs(
             functionDeclaration,
-            {throwOnSideEffect},
+            {throwOnSideEffect, contextNode: this.context?.getItem() || null},
             ),
         new Promise<never>((_, reject) => {
           setTimeout(