Pre-factoring for cross-origin restrictions

- introduces a context class for checking origins - makes the agent instance assign an origin to itself based on the context Bug: 377227220 Change-Id: If4b629d8bbf2427f0df35dd6807235459a257338 Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/5987650 Reviewed-by: Nikolay Vitkov <[email protected]> Commit-Queue: Alex Rudenko <[email protected]>
ChromeDevTools · Nov 4, 2024 · bd344de · bd344de
1 parent f31e4cc
commit bd344de
Show file tree

Hide file tree

Showing 15 changed files with 279 additions and 112 deletions.
diff --git a/front_end/panels/freestyler/AiAgent.ts b/front_end/panels/freestyler/AiAgent.ts
@@ -127,6 +127,21 @@ export const enum AgentType {
 
 const MAX_STEP = 10;
 
+export abstract class ConversationContext<T> {
+  abstract getOrigin(): string;
+  abstract getItem(): T;
+  isOriginAllowed(agentOrigin: string|undefined): boolean {
+    if (!agentOrigin) {
+      return true;
+    }
+    // Currently does not handle opaque origins because they
+    // are not available to DevTools, instead checks
+    // that serialization of the origin is the same
+    // https://html.spec.whatwg.org/#ascii-serialisation-of-an-origin.
+    return this.getOrigin() === agentOrigin;
+  }
+}
+
 export abstract class AiAgent<T> {
   static validTemperature(temperature: number|undefined): number|undefined {
     return typeof temperature === 'number' && temperature >= 0 ? temperature : undefined;
@@ -136,11 +151,12 @@ export abstract class AiAgent<T> {
   readonly #sessionId: string = crypto.randomUUID();
   #aidaClient: Host.AidaClient.AidaClient;
   #serverSideLoggingEnabled: boolean;
+  #origin?: string;
   abstract readonly preamble: string;
   abstract readonly options: AidaRequestOptions;
   abstract readonly clientFeature: Host.AidaClient.ClientFeature;
   abstract readonly userTier: string|undefined;
-  abstract handleContextDetails(select: T|null): AsyncGenerator<ContextResponse, void, void>;
+  abstract handleContextDetails(select: ConversationContext<T>|null): AsyncGenerator<ContextResponse, void, void>;
 
   /**
    * Mapping between the unique request id and
@@ -165,6 +181,10 @@ export abstract class AiAgent<T> {
     return this.#history.size <= 0;
   }
 
+  get origin(): string|undefined {
+    return this.#origin;
+  }
+
   get title(): string|undefined {
     return [...this.#history.values()]
         .flat()
@@ -248,7 +268,7 @@ export abstract class AiAgent<T> {
     throw new Error('Unexpected action found');
   }
 
-  async enhanceQuery(query: string, selected: T|null): Promise<string>;
+  async enhanceQuery(query: string, selected: ConversationContext<T>|null): Promise<string>;
   async enhanceQuery(query: string): Promise<string> {
     return query;
   }
@@ -359,8 +379,12 @@ STOP`;
 
   #runId = 0;
   async * run(query: string, options: {
-    signal?: AbortSignal, selected: T|null,
+    signal?: AbortSignal, selected: ConversationContext<T>|null,
   }): AsyncGenerator<ResponseData, void, void> {
+    // 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();
+    }
     const id = this.#runId++;
 
     const response = {

diff --git a/front_end/panels/freestyler/DrJonesFileAgent.test.ts b/front_end/panels/freestyler/DrJonesFileAgent.test.ts
@@ -16,7 +16,7 @@ import {describeWithMockConnection} from '../../testing/MockConnection.js';
 import {loadBasicSourceMapExample} from '../../testing/SourceMapHelpers.js';
 import {createContentProviderUISourceCodes} from '../../testing/UISourceCodeHelpers.js';
 
-import {DrJonesFileAgent, formatSourceMapDetails, ResponseType} from './freestyler.js';
+import {DrJonesFileAgent, FileContext, formatSourceMapDetails, ResponseType} from './freestyler.js';
 
 describeWithMockConnection('DrJonesFileAgent', () => {
   function mockHostConfig(modelId?: string, temperature?: number) {
@@ -162,7 +162,8 @@ describeWithMockConnection('DrJonesFileAgent', () => {
       });
 
       const uiSourceCode = project.uiSourceCodeForURL(url);
-      const responses = await Array.fromAsync(agent.run('test', {selected: uiSourceCode}));
+      const responses =
+          await Array.fromAsync(agent.run('test', {selected: uiSourceCode ? new FileContext(uiSourceCode) : null}));
 
       assert.deepStrictEqual(responses, [
         {

diff --git a/front_end/panels/freestyler/DrJonesFileAgent.ts b/front_end/panels/freestyler/DrJonesFileAgent.ts
@@ -14,6 +14,7 @@ import {
   type AidaRequestOptions,
   type ContextDetail,
   type ContextResponse,
+  ConversationContext,
   type ParsedResponse,
   ResponseType,
 } from './AiAgent.js';
@@ -69,6 +70,23 @@ const lockedString = i18n.i18n.lockedString;
 
 const MAX_FILE_SIZE = 10000;
 
+export class FileContext extends ConversationContext<Workspace.UISourceCode.UISourceCode> {
+  #file: Workspace.UISourceCode.UISourceCode;
+
+  constructor(file: Workspace.UISourceCode.UISourceCode) {
+    super();
+    this.#file = file;
+  }
+
+  getOrigin(): string {
+    return new URL(this.#file.url()).origin;
+  }
+
+  getItem(): Workspace.UISourceCode.UISourceCode {
+    return this.#file;
+  }
+}
+
 /**
  * One agent instance handles one conversation. Create a new agent
  * instance for a new conversation.
@@ -93,7 +111,7 @@ export class DrJonesFileAgent extends AiAgent<Workspace.UISourceCode.UISourceCod
   }
 
   async *
-      handleContextDetails(selectedFile: Workspace.UISourceCode.UISourceCode|null):
+      handleContextDetails(selectedFile: ConversationContext<Workspace.UISourceCode.UISourceCode>|null):
           AsyncGenerator<ContextResponse, void, void> {
     if (!selectedFile) {
       return;
@@ -106,9 +124,10 @@ export class DrJonesFileAgent extends AiAgent<Workspace.UISourceCode.UISourceCod
     };
   }
 
-  override async enhanceQuery(query: string, selectedFile: Workspace.UISourceCode.UISourceCode|null): Promise<string> {
+  override async enhanceQuery(
+      query: string, selectedFile: ConversationContext<Workspace.UISourceCode.UISourceCode>|null): Promise<string> {
     const fileEnchantmentQuery =
-        selectedFile ? `# Selected file\n${formatFile(selectedFile)}\n\n# User request\n\n` : '';
+        selectedFile ? `# Selected file\n${formatFile(selectedFile.getItem())}\n\n# User request\n\n` : '';
     return `${fileEnchantmentQuery}${query}`;
   }
 
@@ -119,12 +138,12 @@ export class DrJonesFileAgent extends AiAgent<Workspace.UISourceCode.UISourceCod
   }
 }
 
-function createContextDetailsForDrJonesFileAgent(selectedFile: Workspace.UISourceCode.UISourceCode):
-    [ContextDetail, ...ContextDetail[]] {
+function createContextDetailsForDrJonesFileAgent(
+    selectedFile: ConversationContext<Workspace.UISourceCode.UISourceCode>): [ContextDetail, ...ContextDetail[]] {
   return [
     {
       title: 'Selected file',
-      text: formatFile(selectedFile),
+      text: formatFile(selectedFile.getItem()),
     },
   ];
 }

diff --git a/front_end/panels/freestyler/DrJonesNetworkAgent.test.ts b/front_end/panels/freestyler/DrJonesNetworkAgent.test.ts
@@ -15,7 +15,14 @@ import {describeWithMockConnection} from '../../testing/MockConnection.js';
 import {createNetworkPanelForMockConnection} from '../../testing/NetworkHelpers.js';
 import * as Coordinator from '../../ui/components/render_coordinator/render_coordinator.js';
 
-import {allowHeader, DrJonesNetworkAgent, formatHeaders, formatInitiatorUrl, ResponseType} from './freestyler.js';
+import {
+  allowHeader,
+  DrJonesNetworkAgent,
+  formatHeaders,
+  formatInitiatorUrl,
+  RequestContext,
+  ResponseType,
+} from './freestyler.js';
 
 const coordinator = Coordinator.RenderCoordinator.RenderCoordinator.instance();
 
@@ -215,7 +222,8 @@ describeWithMockConnection('DrJonesNetworkAgent', () => {
         aidaClient: mockAidaClient(generateAnswer),
       });
 
-      const responses = await Array.fromAsync(agent.run('test', {selected: selectedNetworkRequest}));
+      const responses =
+          await Array.fromAsync(agent.run('test', {selected: new RequestContext(selectedNetworkRequest)}));
       assert.deepStrictEqual(responses, [
         {
           type: ResponseType.USER_QUERY,

diff --git a/front_end/panels/freestyler/DrJonesNetworkAgent.ts b/front_end/panels/freestyler/DrJonesNetworkAgent.ts
@@ -15,6 +15,7 @@ import {
   type AidaRequestOptions,
   type ContextDetail,
   type ContextResponse,
+  ConversationContext,
   type ParsedResponse,
   ResponseType,
 } from './AiAgent.js';
@@ -97,6 +98,23 @@ const UIStringsNotTranslate = {
 
 const lockedString = i18n.i18n.lockedString;
 
+export class RequestContext extends ConversationContext<SDK.NetworkRequest.NetworkRequest> {
+  #request: SDK.NetworkRequest.NetworkRequest;
+
+  constructor(request: SDK.NetworkRequest.NetworkRequest) {
+    super();
+    this.#request = request;
+  }
+
+  getOrigin(): string {
+    return new URL(this.#request.url()).origin;
+  }
+
+  getItem(): SDK.NetworkRequest.NetworkRequest {
+    return this.#request;
+  }
+}
+
 /**
  * One agent instance handles one conversation. Create a new agent
  * instance for a new conversation.
@@ -121,7 +139,7 @@ export class DrJonesNetworkAgent extends AiAgent<SDK.NetworkRequest.NetworkReque
   }
 
   async *
-      handleContextDetails(selectedNetworkRequest: SDK.NetworkRequest.NetworkRequest|null):
+      handleContextDetails(selectedNetworkRequest: ConversationContext<SDK.NetworkRequest.NetworkRequest>|null):
           AsyncGenerator<ContextResponse, void, void> {
     if (!selectedNetworkRequest) {
       return;
@@ -130,14 +148,15 @@ export class DrJonesNetworkAgent extends AiAgent<SDK.NetworkRequest.NetworkReque
     yield {
       type: ResponseType.CONTEXT,
       title: lockedString(UIStringsNotTranslate.analyzingNetworkData),
-      details: createContextDetailsForDrJonesNetworkAgent(selectedNetworkRequest),
+      details: createContextDetailsForDrJonesNetworkAgent(selectedNetworkRequest.getItem()),
     };
   }
 
-  override async enhanceQuery(query: string, selectedNetworkRequest: SDK.NetworkRequest.NetworkRequest|null):
-      Promise<string> {
+  override async enhanceQuery(
+      query: string,
+      selectedNetworkRequest: ConversationContext<SDK.NetworkRequest.NetworkRequest>|null): Promise<string> {
     const networkEnchantmentQuery = selectedNetworkRequest ?
-        `# Selected network request \n${formatNetworkRequest(selectedNetworkRequest)}\n\n# User request\n\n` :
+        `# Selected network request \n${formatNetworkRequest(selectedNetworkRequest.getItem())}\n\n# User request\n\n` :
         '';
     return `${networkEnchantmentQuery}${query}`;
   }

diff --git a/front_end/panels/freestyler/DrJonesPerformanceAgent.test.ts b/front_end/panels/freestyler/DrJonesPerformanceAgent.test.ts
@@ -7,7 +7,7 @@ import {describeWithEnvironment, getGetHostConfigStub} from '../../testing/Envir
 import {TraceLoader} from '../../testing/TraceLoader.js';
 import * as TimelineUtils from '../timeline/utils/utils.js';
 
-import {DrJonesPerformanceAgent, ResponseType} from './freestyler.js';
+import {CallTreeContext, DrJonesPerformanceAgent, ResponseType} from './freestyler.js';
 
 describeWithEnvironment('DrJonesPerformanceAgent', () => {
   function mockHostConfig(modelId?: string, temperature?: number) {
@@ -133,7 +133,7 @@ describeWithEnvironment('DrJonesPerformanceAgent', () => {
         aidaClient: mockAidaClient(generateAnswer),
       });
 
-      const responses = await Array.fromAsync(agent.run('test', {selected: aiCallTree}));
+      const responses = await Array.fromAsync(agent.run('test', {selected: new CallTreeContext(aiCallTree)}));
       const expectedData = '\n\n' +
           `
 
@@ -198,7 +198,7 @@ self: 3
         serialize: () => 'Mock call tree',
       } as unknown as TimelineUtils.AICallTree.AICallTree;
 
-      const enhancedQuery1 = await agent.enhanceQuery('What is this?', mockAiCallTree);
+      const enhancedQuery1 = await agent.enhanceQuery('What is this?', new CallTreeContext(mockAiCallTree));
       assert.strictEqual(enhancedQuery1, 'Mock call tree\n\n# User request\n\nWhat is this?');
 
       // Create history state of the above query
@@ -227,13 +227,13 @@ self: 3
       ]]);
 
       const query2 = 'But what about this follow-up question?';
-      const enhancedQuery2 = await agent.enhanceQuery(query2, mockAiCallTree);
+      const enhancedQuery2 = await agent.enhanceQuery(query2, new CallTreeContext(mockAiCallTree));
       assert.strictEqual(enhancedQuery2, query2);
       assert.isFalse(enhancedQuery2.includes(mockAiCallTree.serialize()));
 
       // Just making sure any subsequent chat doesnt include it either.
       const query3 = 'And this 3rd question?';
-      const enhancedQuery3 = await agent.enhanceQuery(query3, mockAiCallTree);
+      const enhancedQuery3 = await agent.enhanceQuery(query3, new CallTreeContext(mockAiCallTree));
       assert.strictEqual(enhancedQuery3, query3);
       assert.isFalse(enhancedQuery3.includes(mockAiCallTree.serialize()));
     });

diff --git a/front_end/panels/freestyler/DrJonesPerformanceAgent.ts b/front_end/panels/freestyler/DrJonesPerformanceAgent.ts
@@ -12,6 +12,7 @@ import {
   AiAgent,
   type AidaRequestOptions,
   type ContextResponse,
+  ConversationContext,
   type ParsedResponse,
   ResponseType,
 } from './AiAgent.js';
@@ -120,6 +121,24 @@ const UIStringsNotTranslate = {
 
 const lockedString = i18n.i18n.lockedString;
 
+export class CallTreeContext extends ConversationContext<TimelineUtils.AICallTree.AICallTree> {
+  #callTree: TimelineUtils.AICallTree.AICallTree;
+
+  constructor(callTree: TimelineUtils.AICallTree.AICallTree) {
+    super();
+    this.#callTree = callTree;
+  }
+
+  getOrigin(): string {
+    // TODO: implement cross-origin checks for the PerformanceAgent.
+    return '';
+  }
+
+  getItem(): TimelineUtils.AICallTree.AICallTree {
+    return this.#callTree;
+  }
+}
+
 /**
  * One agent instance handles one conversation. Create a new agent
  * instance for a new conversation.
@@ -144,22 +163,23 @@ export class DrJonesPerformanceAgent extends AiAgent<TimelineUtils.AICallTree.AI
   }
 
   async *
-      handleContextDetails(aiCallTree: TimelineUtils.AICallTree.AICallTree|null):
+      handleContextDetails(aiCallTree: ConversationContext<TimelineUtils.AICallTree.AICallTree>|null):
           AsyncGenerator<ContextResponse, void, void> {
     yield {
       type: ResponseType.CONTEXT,
       title: lockedString(UIStringsNotTranslate.analyzingCallTree),
       details: [
         {
           title: 'Selected call tree',
-          text: aiCallTree?.serialize() ?? '',
+          text: aiCallTree?.getItem().serialize() ?? '',
         },
       ],
     };
   }
 
-  override async enhanceQuery(query: string, aiCallTree: TimelineUtils.AICallTree.AICallTree|null): Promise<string> {
-    const treeStr = aiCallTree?.serialize();
+  override async enhanceQuery(query: string, aiCallTree: ConversationContext<TimelineUtils.AICallTree.AICallTree>|null):
+      Promise<string> {
+    const treeStr = aiCallTree?.getItem().serialize();
 
     // Collect the queries from previous messages in this session
     const prevQueries: string[] = [];