[AI] Add getFunctionCode to PerformanceAgent

Bug: 463466660
Change-Id: I3d1beec6befcbb7817a2a081fedc0f58485fea2f
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/7201966
Auto-Submit: Connor Clark <cjamcl@chromium.org>
Commit-Queue: Connor Clark <cjamcl@chromium.org>
Reviewed-by: Paul Irish <paulirish@chromium.org>

Author: Connor Clark

front_end/models/ai_assistance/agents/PerformanceAgent.ts

@@ -9,6 +9,7 @@ import * as Platform from '../../../core/platform/platform.js';
 import * as Root from '../../../core/root/root.js';
 import * as SDK from '../../../core/sdk/sdk.js';
 import * as Tracing from '../../../services/tracing/tracing.js';
+import * as SourceMapScopes from '../../source_map_scopes/source_map_scopes.js';
 import * as Trace from '../../trace/trace.js';
 import {
   PerformanceInsightFormatter,
@@ -883,12 +884,77 @@ export class PerformanceAgent extends AiAgent<AgentFocus> {
 
     });
 
+    this.declareFunction<{scriptUrl: string, line: number, column: number}, {result: string}>('getFunctionCode', {
+      description: 'Returns the code for a function defined at the given location.',
+      parameters: {
+        type: Host.AidaClient.ParametersTypes.OBJECT,
+        description: '',
+        nullable: false,
+        properties: {
+          scriptUrl: {
+            type: Host.AidaClient.ParametersTypes.STRING,
+            description: 'The url of the function.',
+            nullable: false,
+          },
+          line: {
+            type: Host.AidaClient.ParametersTypes.INTEGER,
+            description: 'The line number where the function is defined.',
+            nullable: false,
+          },
+          column: {
+            type: Host.AidaClient.ParametersTypes.INTEGER,
+            description: 'The column number where the function is defined.',
+            nullable: false,
+          },
+        },
+      },
+      displayInfoFromArgs: args => {
+        return {
+          title: lockedString('Looking up function code…'),
+          action: `getFunctionCode('${args.scriptUrl}', ${args.line}, ${args.column})`
+        };
+      },
+      handler: async args => {
+        debugLog('Function call: getFunctionCode');
+
+        if (args.line === undefined) {
+          return {error: 'Missing arg: line'};
+        }
+
+        if (args.column === undefined) {
+          return {error: 'Missing arg: column'};
+        }
+
+        if (!this.#formatter) {
+          throw new Error('missing formatter');
+        }
+
+        const target = SDK.TargetManager.TargetManager.instance().primaryPageTarget();
+        if (!target) {
+          throw new Error('missing target');
+        }
+
+        const url = args.scriptUrl as Platform.DevToolsPath.UrlString;
+        const code = await SourceMapScopes.FunctionCodeResolver.getFunctionCodeFromLocation(
+            target, url, args.line, args.column, {contextLength: 200, contextLineLength: 5});
+        if (!code) {
+          return {error: 'Could not find code'};
+        }
+
+        const result = this.#formatter.formatFunctionCode(code);
+
+        const key = `getFunctionCode('${args.scriptUrl}', ${args.line}, ${args.column})`;
+        this.#cacheFunctionResult(focus, key, result);
+        return {result: {result}};
+      },
+    });
+
     const isFresh = Tracing.FreshRecording.Tracker.instance().recordingIsFresh(parsedTrace);
     const isTraceApp = Root.Runtime.Runtime.isTraceApp();
 
     this.declareFunction<{url: string}, {content: string}>('getResourceContent', {
       description:
-          'Returns the content of the resource with the given url. Only use this for text resource types. This function is helpful for getting script contents in order to further analyze main thread activity and suggest code improvements. When analyzing the main thread activity, always call this function to get more detail. Always call this function when asked to provide specifics about what is happening in the code. Never ask permission to call this function, just do it.',
+          'Returns the content of the resource with the given url. Only use this for text resource types. Prefer getFunctionCode when possible.',
       parameters: {
         type: Host.AidaClient.ParametersTypes.OBJECT,
         description: '',

front_end/models/ai_assistance/data_formatters/PerformanceTraceFormatter.ts

@@ -3,6 +3,7 @@
 // found in the LICENSE file.
 
 import * as CrUXManager from '../../crux-manager/crux-manager.js';
+import type * as SourceMapScopes from '../../source_map_scopes/source_map_scopes.js';
 import * as Trace from '../../trace/trace.js';
 import type {AICallTree} from '../performance/AICallTree.js';
 import type {AgentFocus} from '../performance/AIContext.js';
@@ -813,4 +814,27 @@ The order of headers corresponds to an internal fixed list. If a header is not p
     ];
     return parts.join(';');
   }
+
+  formatFunctionCode(code: SourceMapScopes.FunctionCodeResolver.FunctionCode): string {
+    const {startLine, startColumn} = code.range;
+    const {
+      startLine: contextStartLine,
+      startColumn: contextStartColumn,
+      endLine: contextEndLine,
+      endColumn: contextEndColumn
+    } = code.rangeWithContext;
+    const name = code.functionBounds.name;
+    const url = code.functionBounds.uiSourceCode.url();
+
+    const parts = [];
+    parts.push(`${name} @ ${url}:${startLine}:${startColumn}. With added context, chunk is from ${contextStartLine}:${
+        contextStartColumn} to ${contextEndLine}:${contextEndColumn}`);
+    parts.push(
+        '\nThe following is a markdown block of JavaScript. <FUNCTION_START> and <FUNCTION_END> marks the exact function declaration, and everything outside that is provided for additional context. Do not show the user the function markers or the additional context.\n');
+    parts.push('```');
+    parts.push(code.codeWithContext);
+    parts.push('```');
+
+    return parts.join('\n');
+  }
 }