chrome-devtools-frontend 1.0.1402249 → 1.0.1418433

package/front_end/panels/ai_assistance/agents/AiAgent.ts

@@ -3,8 +3,8 @@
 // found in the LICENSE file.
 
 import * as Host from '../../../core/host/host.js';
-import type * as LitHtml from '../../../ui/lit-html/lit-html.js';
-import {AiHistoryStorage} from '../AiHistoryStorage.js';
+import type * as Lit from '../../../ui/lit/lit.js';
+import {debugLog, isDebugMode} from '../debug.js';
 
 export const enum ResponseType {
   CONTEXT = 'context',
@@ -12,6 +12,7 @@ export const enum ResponseType {
   THOUGHT = 'thought',
   ACTION = 'action',
   SIDE_EFFECT = 'side-effect',
+  SUGGESTIONS = 'suggestions',
   ANSWER = 'answer',
   ERROR = 'error',
   QUERYING = 'querying',
@@ -32,6 +33,11 @@ export interface AnswerResponse {
   suggestions?: [string, ...string[]];
 }
 
+export interface SuggestionsResponse {
+  type: ResponseType.SUGGESTIONS;
+  suggestions: [string, ...string[]];
+}
+
 export interface ErrorResponse {
   type: ResponseType.ERROR;
   error: ErrorType;
@@ -62,20 +68,20 @@ export interface ThoughtResponse {
 
 export interface SideEffectResponse {
   type: ResponseType.SIDE_EFFECT;
-  code: string;
+  code?: string;
   confirm: (confirm: boolean) => void;
 }
 
 export interface ActionResponse {
   type: ResponseType.ACTION;
-  code: string;
-  output: string;
+  code?: string;
+  output?: string;
   canceled: boolean;
 }
 
 export interface QueryResponse {
   type: ResponseType.QUERYING;
-  query: string;
+  query?: string;
 }
 
 export interface UserQuery {
@@ -83,8 +89,11 @@ export interface UserQuery {
   query: string;
 }
 
-export type ResponseData = AnswerResponse|ErrorResponse|ActionResponse|SideEffectResponse|ThoughtResponse|TitleResponse|
-    QueryResponse|ContextResponse|UserQuery;
+export type ResponseData = AnswerResponse|SuggestionsResponse|ErrorResponse|ActionResponse|SideEffectResponse|
+    ThoughtResponse|TitleResponse|QueryResponse|ContextResponse|UserQuery;
+
+export type FunctionCallResponseData =
+    TitleResponse|ThoughtResponse|ActionResponse|SideEffectResponse|SuggestionsResponse;
 
 export interface BuildRequestOptions {
   text: string;
@@ -98,6 +107,7 @@ export interface RequestOptions {
 export interface AgentOptions {
   aidaClient: Host.AidaClient.AidaClient;
   serverSideLoggingEnabled?: boolean;
+  confirmSideEffectForTest?: typeof Promise.withResolvers;
 }
 
 export interface ParsedAnswer {
@@ -118,23 +128,17 @@ export const enum AgentType {
   FILE = 'drjones-file',
   NETWORK = 'drjones-network-request',
   PERFORMANCE = 'drjones-performance',
+  PERFORMANCE_INSIGHT = 'performance-insight',
   PATCH = 'patch',
 }
 
-export interface SerializedAgent {
-  id: string;
-  type: AgentType;
-  history: HistoryEntryStorage;
-}
-
-const MAX_STEP = 10;
-export type HistoryEntryStorage = ResponseData[];
+export const MAX_STEPS = 10;
 
 export abstract class ConversationContext<T> {
   abstract getOrigin(): string;
   abstract getItem(): T;
   abstract getIcon(): HTMLElement;
-  abstract getTitle(): string|ReturnType<typeof LitHtml.Directives.until>;
+  abstract getTitle(): string|ReturnType<typeof Lit.Directives.until>;
 
   isOriginAllowed(agentOrigin: string|undefined): boolean {
     if (!agentOrigin) {
@@ -149,188 +153,129 @@ export abstract class ConversationContext<T> {
 
   /**
    * This method is called at the start of `AiAgent.run`.
-   * It will be overriden in subclasses to fetch data related to the context item.
+   * It will be overridden in subclasses to fetch data related to the context item.
    */
   async refresh(): Promise<void> {
     return;
   }
 }
 
-export type FunctionDeclaration<Args, ReturnType> = {
-  description: string,
-  parameters: Host.AidaClient.FunctionObjectParam,
-  handler: (args: Args) => Promise<ReturnType>,
-};
+export type FunctionCallHandlerResult<Result> = {
+  result: Result,
+}|{
+  requiresApproval: true,
+}|{error: string};
 
-export abstract class AiAgent<T> {
-  static validTemperature(temperature: number|undefined): number|undefined {
-    return typeof temperature === 'number' && temperature >= 0 ? temperature : undefined;
-  }
+export interface FunctionDeclaration<Args extends Record<string, unknown>, ReturnType> {
+  /**
+   * Description of function, this is send to the LLM
+   * to explain what will the function do.
+   */
+  description: string;
+  /**
+   * JSON schema like representation of the parameters
+   * the function needs to be called with.
+   * Provide description to all parameters as this is
+   * send to the LLM.
+   */
+  parameters: Host.AidaClient.FunctionObjectParam<keyof Args>;
+  /**
+   * Provided a way to give information back to
+   * the UI before running the the handler
+   */
+  displayInfoFromArgs?: (
+      args: Args,
+      ) => {
+    title?: string, thought?: string, code?: string, suggestions?: [string, ...string[]],
+  };
+  /**
+   * Function implementation that the LLM will try to execute,
+   */
+  handler: (args: Args, options?: {
+    /**
+     * Shows that the user approved
+     * the execution if it was required
+     */
+    approved?: boolean,
+    signal?: AbortSignal,
+  }) => Promise<FunctionCallHandlerResult<ReturnType>>;
+}
 
-  abstract type: AgentType;
-  #id: string = crypto.randomUUID();
-  readonly #sessionId: string = crypto.randomUUID();
-  #aidaClient: Host.AidaClient.AidaClient;
-  #serverSideLoggingEnabled: boolean;
-  abstract readonly preamble: string;
+const OBSERVATION_PREFIX = 'OBSERVATION:';
+
+export abstract class AiAgent<T> {
+  /** Subclasses need to define these. */
+  abstract readonly type: AgentType;
+  abstract readonly preamble: string|undefined;
   abstract readonly options: RequestOptions;
   abstract readonly clientFeature: Host.AidaClient.ClientFeature;
   abstract readonly userTier: string|undefined;
   abstract handleContextDetails(select: ConversationContext<T>|null): AsyncGenerator<ContextResponse, void, void>;
-  #functionDeclarations = new Map<string, FunctionDeclaration<unknown, unknown>>();
-  #generatedFromHistory = false;
+
+  readonly #sessionId: string = crypto.randomUUID();
+  readonly #aidaClient: Host.AidaClient.AidaClient;
+  readonly #serverSideLoggingEnabled: boolean;
+  readonly confirmSideEffect: typeof Promise.withResolvers;
+  readonly #functionDeclarations = new Map<string, FunctionDeclaration<Record<string, unknown>, unknown>>();
 
   /**
-   * Historical responses.
+   * Used in the debug mode and evals.
    */
-  #history: HistoryEntryStorage = [];
+  readonly #structuredLog: Array<{
+    request: Host.AidaClient.AidaRequest,
+    response: string,
+    aidaResponse?: Host.AidaClient.AidaResponse,
+  }> = [];
+
   /**
    * Might need to be part of history in case we allow chatting in
    * historical conversations.
    */
   #origin?: string;
   #context?: ConversationContext<T>;
-  #structuredLog: Array<{
-    request: Host.AidaClient.AidaRequest,
-    response: string,
-    aidaResponse?: Host.AidaClient.AidaResponse,
-  }> = [];
+  #id: string = crypto.randomUUID();
+  #history: Host.AidaClient.Content[] = [];
 
   constructor(opts: AgentOptions) {
     this.#aidaClient = opts.aidaClient;
     this.#serverSideLoggingEnabled = opts.serverSideLoggingEnabled ?? false;
+    this.confirmSideEffect = opts.confirmSideEffectForTest ?? (() => Promise.withResolvers());
   }
 
-  get chatHistoryForTesting(): Array<Host.AidaClient.Content> {
-    return this.buildChatHistoryForAida();
-  }
-
-  declareFunction<Args, ReturnType>(name: string, declaration: FunctionDeclaration<Args, ReturnType>): void {
-    if (this.#functionDeclarations.has(name)) {
-      throw new Error(`Duplicate function declaration ${name}`);
-    }
-    this.#functionDeclarations.set(name, declaration as FunctionDeclaration<unknown, unknown>);
-  }
-
-  async callFunction(name: string, args: unknown): Promise<Record<string, unknown>> {
-    const call = this.#functionDeclarations.get(name);
-    return (call ? await call.handler(args) : {
-             error: `Function ${name} is not found.`,
-           }) as Record<string, unknown>;
-  }
-
-  set chatNewHistoryForTesting(history: HistoryEntryStorage) {
-    this.#history = history;
-  }
-
-  get id(): string {
-    return this.#id;
-  }
-
-  get isEmpty(): boolean {
-    return this.#history.length === 0;
-  }
-
-  get origin(): string|undefined {
-    return this.#origin;
-  }
-
-  get context(): ConversationContext<T>|undefined {
-    return this.#context;
-  }
-
-  get title(): string|undefined {
-    return this.#history
-        .filter(response => {
-          return response.type === ResponseType.USER_QUERY;
-        })
-        .at(0)
-        ?.query;
-  }
-
-  get isHistoryEntry(): boolean {
-    return this.#generatedFromHistory;
+  async enhanceQuery(query: string, selected: ConversationContext<T>|null): Promise<string>;
+  async enhanceQuery(query: string): Promise<string> {
+    return query;
   }
 
-  #buildFunctionDeclarationsForAida(): Host.AidaClient.FunctionDeclaration[] {
-    const result: Host.AidaClient.FunctionDeclaration[] = [];
+  buildRequest(part: Host.AidaClient.Part, role: Host.AidaClient.Role.USER|Host.AidaClient.Role.ROLE_UNSPECIFIED):
+      Host.AidaClient.AidaRequest {
+    const currentMessage: Host.AidaClient.Content = {
+      parts: [part],
+      role,
+    };
+    const history = [...this.#history];
+    const declarations: Host.AidaClient.FunctionDeclaration[] = [];
     for (const [name, definition] of this.#functionDeclarations.entries()) {
-      result.push({
+      declarations.push({
         name,
         description: definition.description,
         parameters: definition.parameters,
       });
     }
-    return result;
-  }
-
-  serialized(): SerializedAgent {
-    return {
-      id: this.id,
-      type: this.type,
-      history: this.#history,
-    };
-  }
-
-  populateHistoryFromStorage(entry: SerializedAgent): void {
-    this.#id = entry.id;
-    this.#history = entry.history;
-    this.#generatedFromHistory = true;
-  }
-
-  async *
-      aidaFetch(
-          request: Host.AidaClient.AidaRequest,
-          options?: {signal?: AbortSignal},
-          ):
-          AsyncGenerator<
-              {parsedResponse: ParsedResponse, completed: boolean, rpcId?: Host.AidaClient.RpcGlobalId}, void, void> {
-    let aidaResponse: Host.AidaClient.AidaResponse|undefined = undefined;
-    let response = '';
-    let rpcId: Host.AidaClient.RpcGlobalId|undefined;
-    for await (aidaResponse of this.#aidaClient.fetch(request, options)) {
-      response = aidaResponse.explanation;
-      rpcId = aidaResponse.metadata.rpcGlobalId ?? rpcId;
-      const parsedResponse = this.parseResponse(aidaResponse);
-      yield {
-        rpcId,
-        parsedResponse,
-        completed: aidaResponse.completed,
-      };
+    function validTemperature(temperature: number|undefined): number|undefined {
+      return typeof temperature === 'number' && temperature >= 0 ? temperature : undefined;
     }
-
-    debugLog({
-      request,
-      response: aidaResponse,
-    });
-
-    this.#structuredLog.push({
-      request: structuredClone(request),
-      response,
-      aidaResponse,
-    });
-    localStorage.setItem('freestylerStructuredLog', JSON.stringify(this.#structuredLog));
-  }
-
-  buildRequest(part: Host.AidaClient.Part): Host.AidaClient.AidaRequest {
-    const currentMessage: Host.AidaClient.Content = {
-      parts: [part],
-      role: Host.AidaClient.Role.USER,
-    };
-    const history = this.buildChatHistoryForAida();
-    const declarations = this.#buildFunctionDeclarationsForAida();
     const request: Host.AidaClient.AidaRequest = {
       client: Host.AidaClient.CLIENT_NAME,
-      // eslint-disable-next-line @typescript-eslint/naming-convention
+
       current_message: currentMessage,
       preamble: this.preamble,
-      // eslint-disable-next-line @typescript-eslint/naming-convention
+
       historical_contexts: history.length ? history : undefined,
-      // eslint-disable-next-line @typescript-eslint/naming-convention
+
       ...(declarations.length ? {function_declarations: declarations} : {}),
       options: {
-        temperature: AiAgent.validTemperature(this.options.temperature),
-        // eslint-disable-next-line @typescript-eslint/naming-convention
+        temperature: validTemperature(this.options.temperature),
         model_id: this.options.modelId,
       },
       metadata: {
@@ -338,24 +283,25 @@ export abstract class AiAgent<T> {
         string_session_id: this.#sessionId,
         user_tier: Host.AidaClient.convertToUserTierEnum(this.userTier),
       },
-      // eslint-disable-next-line @typescript-eslint/naming-convention
+
       functionality_type: declarations.length ? Host.AidaClient.FunctionalityType.AGENTIC_CHAT :
                                                 Host.AidaClient.FunctionalityType.CHAT,
-      // eslint-disable-next-line @typescript-eslint/naming-convention
+
       client_feature: this.clientFeature,
     };
     return request;
   }
 
-  handleAction(action: string, options?: {signal?: AbortSignal}):
-      AsyncGenerator<SideEffectResponse, ActionResponse, void>;
-  handleAction(): never {
-    throw new Error('Unexpected action found');
+  get id(): string {
+    return this.#id;
   }
 
-  async enhanceQuery(query: string, selected: ConversationContext<T>|null): Promise<string>;
-  async enhanceQuery(query: string): Promise<string> {
-    return query;
+  get isEmpty(): boolean {
+    return this.#history.length === 0;
+  }
+
+  get origin(): string|undefined {
+    return this.#origin;
   }
 
   parseResponse(response: Host.AidaClient.AidaResponse): ParsedResponse {
@@ -367,100 +313,38 @@ export abstract class AiAgent<T> {
     };
   }
 
-  formatParsedAnswer({answer}: ParsedAnswer): string {
-    return answer;
-  }
-
-  formatParsedStep(step: ParsedStep): string {
-    let text = '';
-    if (step.thought) {
-      text = `THOUGHT: ${step.thought}`;
-    }
-    if (step.title) {
-      text += `\nTITLE: ${step.title}`;
-    }
-    if (step.action) {
-      text += `\nACTION
-${step.action}
-STOP`;
+  /**
+   * Declare a function that the AI model can call.
+   * @param name - The name of the function
+   * @param declaration - the function declaration. Currently functions must:
+   * 1. Return an object of serializable key/value pairs. You cannot return
+   *    anything other than a plain JavaScript object that can be serialized.
+   * 2. Take one parameter which is an object that can have
+   *    multiple keys and values. For example, rather than a function being called
+   *    with two args, `foo` and `bar`, you should instead have the function be
+   *    called with one object with `foo` and `bar` keys.
+   */
+  protected declareFunction<Args extends Record<string, unknown>, ReturnType = unknown>(
+      name: string, declaration: FunctionDeclaration<Args, ReturnType>): void {
+    if (this.#functionDeclarations.has(name)) {
+      throw new Error(`Duplicate function declaration ${name}`);
     }
-
-    return text;
+    this.#functionDeclarations.set(name, declaration as FunctionDeclaration<Record<string, unknown>, ReturnType>);
   }
 
-  buildChatHistoryForAida(): Host.AidaClient.Content[] {
-    const history: Array<Host.AidaClient.Content> = [];
-    let currentParsedStep: ParsedStep = {};
-    let lastRunStartIdx = 0;
-    const flushCurrentStep = (): void => {
-      const text = this.formatParsedStep(currentParsedStep);
-      if (text) {
-        history.push({
-          role: Host.AidaClient.Role.MODEL,
-          parts: [{text}],
-        });
-        currentParsedStep = {};
-      }
-    };
-    for (const data of this.#history) {
-      switch (data.type) {
-        case ResponseType.CONTEXT:
-        case ResponseType.SIDE_EFFECT:
-          break;
-        case ResponseType.USER_QUERY:
-          lastRunStartIdx = history.length;
-          break;
-        case ResponseType.QUERYING: {
-          flushCurrentStep();
-          history.push({
-            role: Host.AidaClient.Role.USER,
-            parts: [{
-              text: data.query,
-            }],
-          });
-          break;
-        }
-        case ResponseType.ANSWER:
-          history.push({
-            role: Host.AidaClient.Role.MODEL,
-            parts: [{
-              text: this.formatParsedAnswer({answer: data.text}),
-            }],
-          });
-          break;
-        case ResponseType.TITLE:
-          currentParsedStep.title = data.title;
-          break;
-        case ResponseType.THOUGHT:
-          currentParsedStep.thought = data.thought;
-          break;
-        case ResponseType.ACTION:
-          currentParsedStep.action = data.code;
-          break;
-        case ResponseType.ERROR:
-          // Delete the end of history.
-          history.splice(lastRunStartIdx);
-          currentParsedStep = {};
-          break;
-      }
-    }
-    // Flush remaining step data into history.
-    flushCurrentStep();
-    return history;
+  protected formatParsedAnswer({answer}: ParsedAnswer): string {
+    return answer;
   }
 
-  #addHistory(data: ResponseData): void {
-    this.#history.push(data);
-    void AiHistoryStorage.instance().upsertHistoryEntry(this.serialized());
+  protected handleAction(action: string, options?: {signal?: AbortSignal}):
+      AsyncGenerator<SideEffectResponse, ActionResponse, void>;
+  protected handleAction(): never {
+    throw new Error('Unexpected action found');
   }
 
-  async * run(query: string, options: {
+  async * run(initialQuery: string, options: {
     signal?: AbortSignal, selected: ConversationContext<T>|null,
   }): AsyncGenerator<ResponseData, void, void> {
-    if (this.#generatedFromHistory) {
-      throw new Error('History entries are read-only.');
-    }
-
     await options.selected?.refresh();
 
     // First context set on the agent determines its origin from now on.
@@ -472,43 +356,38 @@ STOP`;
       this.#context = options.selected;
     }
 
-    const enhancedQuery = await this.enhanceQuery(query, options.selected);
+    const enhancedQuery = await this.enhanceQuery(initialQuery, options.selected);
 
     Host.userMetrics.freestylerQueryLength(enhancedQuery.length);
 
+    let query: Host.AidaClient.Part = {text: enhancedQuery};
     // Request is built here to capture history up to this point.
-    let request = this.buildRequest({text: enhancedQuery});
+    let request = this.buildRequest(query, Host.AidaClient.Role.USER);
 
-    const response = {
+    yield {
       type: ResponseType.USER_QUERY,
-      query,
-    } as const;
-    this.#addHistory(response);
-    yield response;
-
-    for await (const response of this.handleContextDetails(options.selected)) {
-      this.#addHistory(response);
-      yield response;
-    }
+      query: initialQuery,
+    };
 
-    query = enhancedQuery;
+    yield* this.handleContextDetails(options.selected);
 
-    for (let i = 0; i < MAX_STEP; i++) {
-      const queryResponse = {
+    for (let i = 0; i < MAX_STEPS; i++) {
+      yield {
         type: ResponseType.QUERYING,
-        query,
-      } as const;
-      this.#addHistory(queryResponse);
-      yield queryResponse;
+      };
+
       let rpcId: Host.AidaClient.RpcGlobalId|undefined;
       let parsedResponse: ParsedResponse|undefined = undefined;
+      let functionCall: Host.AidaClient.AidaFunctionCallResponse|undefined = undefined;
       try {
-        for await (const fetchResult of this.aidaFetch(request, {signal: options.signal})) {
+        for await (const fetchResult of this.#aidaFetch(request, {signal: options.signal})) {
           rpcId = fetchResult.rpcId;
           parsedResponse = fetchResult.parsedResponse;
+          functionCall = fetchResult.functionCall;
 
           // Only yield partial responses here and do not add partial answers to the history.
-          if (!fetchResult.completed && 'answer' in parsedResponse && parsedResponse.answer) {
+          if (!fetchResult.completed && !fetchResult.functionCall && 'answer' in parsedResponse &&
+              parsedResponse.answer) {
             yield {
               type: ResponseType.ANSWER,
               text: parsedResponse.answer,
@@ -518,159 +397,287 @@ STOP`;
       } catch (err) {
         debugLog('Error calling the AIDA API', err);
 
+        let error = ErrorType.UNKNOWN;
         if (err instanceof Host.AidaClient.AidaAbortError) {
-          const response = this.#createAbortResponse();
-          this.#addHistory(response);
-          yield response;
-          break;
+          error = ErrorType.ABORT;
+        } else if (err instanceof Host.AidaClient.AidaBlockError) {
+          error = ErrorType.BLOCK;
         }
-
-        const error = (err instanceof Host.AidaClient.AidaBlockError) ? ErrorType.BLOCK : ErrorType.UNKNOWN;
-        const response = {
-          type: ResponseType.ERROR,
-          error,
-        } as const;
-        Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceError);
-        this.#addHistory(response);
-        yield response;
+        yield this.#createErrorResponse(error);
 
         break;
       }
 
-      if (!parsedResponse) {
-        const response = {
-          type: ResponseType.ERROR,
-          error: ErrorType.UNKNOWN,
-        } as const;
-        Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceError);
-        this.#addHistory(response);
-        yield response;
-        break;
-      }
+      this.#history.push(request.current_message);
 
-      if ('answer' in parsedResponse) {
+      if (parsedResponse && 'answer' in parsedResponse && Boolean(parsedResponse.answer)) {
+        this.#history.push({
+          parts: [{
+            text: this.formatParsedAnswer(parsedResponse),
+          }],
+          role: Host.AidaClient.Role.MODEL,
+        });
+        Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceAnswerReceived);
+        yield {
+          type: ResponseType.ANSWER,
+          text: parsedResponse.answer,
+          suggestions: parsedResponse.suggestions,
+          rpcId,
+        };
+        break;
+      } else if (parsedResponse && !('answer' in parsedResponse)) {
         const {
-          answer,
-          suggestions,
+          title,
+          thought,
+          action,
         } = parsedResponse;
-        if (answer) {
-          const response = {
-            type: ResponseType.ANSWER,
-            text: answer,
+
+        if (title) {
+          yield {
+            type: ResponseType.TITLE,
+            title,
+            rpcId,
+          };
+        }
+
+        if (thought) {
+          yield {
+            type: ResponseType.THOUGHT,
+            thought,
             rpcId,
-            suggestions,
-          } as const;
-          Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceAnswerReceived);
-          this.#addHistory(response);
-          yield response;
-        } else {
-          const response = {
-            type: ResponseType.ERROR,
-            error: ErrorType.UNKNOWN,
-          } as const;
-          Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceError);
-          this.#addHistory(response);
-          yield response;
+          };
         }
 
+        this.#history.push({
+          parts: [{
+            text: this.#formatParsedStep(parsedResponse),
+          }],
+          role: Host.AidaClient.Role.MODEL,
+        });
+
+        if (action) {
+          const result = yield* this.handleAction(action, {signal: options.signal});
+          if (options?.signal?.aborted) {
+            yield this.#createErrorResponse(ErrorType.ABORT);
+            break;
+          }
+          query = {text: `${OBSERVATION_PREFIX} ${result.output}`};
+          // Capture history state for the next iteration query.
+          request = this.buildRequest(query, Host.AidaClient.Role.USER);
+          yield result;
+        }
+      } else if (functionCall) {
+        try {
+          const result = yield* this.#callFunction(functionCall.name, functionCall.args);
+
+          if (result.result) {
+            yield {
+              type: ResponseType.ACTION,
+              output: JSON.stringify(result.result),
+              canceled: false,
+            };
+          }
+
+          query = {
+            functionResponse: {
+              name: functionCall.name,
+              response: result,
+            },
+          };
+          request = this.buildRequest(query, Host.AidaClient.Role.ROLE_UNSPECIFIED);
+        } catch {
+          yield this.#createErrorResponse(ErrorType.UNKNOWN);
+          break;
+        }
+      } else {
+        yield this.#createErrorResponse(i - 1 === MAX_STEPS ? ErrorType.MAX_STEPS : ErrorType.UNKNOWN);
         break;
       }
+    }
 
-      const {
-        title,
-        thought,
-        action,
-      } = parsedResponse;
+    if (isDebugMode()) {
+      window.dispatchEvent(new CustomEvent('aiassistancedone'));
+    }
+  }
 
+  async * #callFunction(name: string, args: Record<string, unknown>, options?: {
+    signal?: AbortSignal,
+    approved?: boolean,
+  }): AsyncGenerator<FunctionCallResponseData, {result: unknown}> {
+    const call = this.#functionDeclarations.get(name);
+    if (!call) {
+      throw new Error(`Function ${name} is not found.`);
+    }
+    this.#history.push({
+      parts: [{
+        functionCall: {
+          name,
+          args,
+        },
+      }],
+      role: Host.AidaClient.Role.MODEL,
+    });
+
+    if (call.displayInfoFromArgs) {
+      const {title, thought, code, suggestions} = call.displayInfoFromArgs(args);
       if (title) {
-        const response = {
+        yield {
           type: ResponseType.TITLE,
           title,
-          rpcId,
-        } as const;
-        this.#addHistory(response);
-        yield response;
+        };
       }
 
       if (thought) {
-        const response = {
+        yield {
           type: ResponseType.THOUGHT,
           thought,
-          rpcId,
-        } as const;
-        this.#addHistory(response);
-        yield response;
+        };
       }
 
-      if (action) {
-        const result = yield* this.handleAction(action, {signal: options.signal});
-        if (options?.signal?.aborted) {
-          const response = this.#createAbortResponse();
-          this.#addHistory(response);
-          yield response;
-          break;
-        }
-        this.#addHistory(result);
-        query = `OBSERVATION: ${result.output}`;
-        // Capture history state for the next iteration query.
-        request = this.buildRequest({text: query});
-        yield result;
+      if (code) {
+        yield {
+          type: ResponseType.ACTION,
+          code,
+          canceled: false,
+        };
       }
 
-      if (i === MAX_STEP - 1) {
-        const response = {
-          type: ResponseType.ERROR,
-          error: ErrorType.MAX_STEPS,
-        } as const;
-        Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceError);
-        this.#addHistory(response);
-        yield response;
-        break;
+      if (suggestions) {
+        yield {
+          type: ResponseType.SUGGESTIONS,
+          suggestions,
+        };
       }
     }
-    if (isDebugMode()) {
-      window.dispatchEvent(new CustomEvent('freestylerdone'));
+
+    let result = await call.handler(args, options);
+
+    if ('requiresApproval' in result) {
+      const sideEffectConfirmationPromiseWithResolvers = this.confirmSideEffect<boolean>();
+
+      void sideEffectConfirmationPromiseWithResolvers.promise.then(result => {
+        Host.userMetrics.actionTaken(
+            result ? Host.UserMetrics.Action.AiAssistanceSideEffectConfirmed :
+                     Host.UserMetrics.Action.AiAssistanceSideEffectRejected,
+        );
+      });
+
+      if (options?.signal?.aborted) {
+        sideEffectConfirmationPromiseWithResolvers.resolve(false);
+      }
+
+      options?.signal?.addEventListener('abort', () => {
+        sideEffectConfirmationPromiseWithResolvers.resolve(false);
+      }, {once: true});
+
+      yield {
+        type: ResponseType.SIDE_EFFECT,
+        confirm: (result: boolean) => {
+          sideEffectConfirmationPromiseWithResolvers.resolve(result);
+        },
+      };
+
+      const approvedRun = await sideEffectConfirmationPromiseWithResolvers.promise;
+      if (!approvedRun) {
+        yield {
+          type: ResponseType.ACTION,
+          code: '',
+          canceled: true,
+        };
+        return {
+          result: 'Error: User denied code execution with side effects.',
+        };
+      }
+
+      result = await call.handler(args, {
+        ...options,
+        approved: approvedRun,
+      });
     }
+
+    return result as {result: unknown};
   }
 
-  async * runFromHistory(): AsyncGenerator<ResponseData, void, void> {
-    if (this.isEmpty) {
-      return;
+  async *
+      #aidaFetch(request: Host.AidaClient.AidaRequest, options?: {signal?: AbortSignal}): AsyncGenerator<
+          {
+            parsedResponse: ParsedResponse,
+            functionCall?: Host.AidaClient.AidaFunctionCallResponse, completed: boolean,
+            rpcId?: Host.AidaClient.RpcGlobalId,
+          },
+          void, void> {
+    let aidaResponse: Host.AidaClient.AidaResponse|undefined = undefined;
+    let response = '';
+    let rpcId: Host.AidaClient.RpcGlobalId|undefined;
+
+    for await (aidaResponse of this.#aidaClient.fetch(request, options)) {
+      if (aidaResponse.functionCalls?.length) {
+        debugLog('functionCalls.length', aidaResponse.functionCalls.length);
+        yield {
+          rpcId,
+          parsedResponse: {answer: ''},
+          functionCall: aidaResponse.functionCalls[0],
+          completed: true,
+        };
+        break;
+      }
+
+      response = aidaResponse.explanation;
+      rpcId = aidaResponse.metadata.rpcGlobalId ?? rpcId;
+      const parsedResponse = this.parseResponse(aidaResponse);
+      yield {
+        rpcId,
+        parsedResponse,
+        completed: aidaResponse.completed,
+      };
     }
 
-    this.#generatedFromHistory = true;
-    for (const entry of this.#history) {
-      yield entry;
+    debugLog({
+      request,
+      response: aidaResponse,
+    });
+    if (isDebugMode()) {
+      this.#structuredLog.push({
+        request: structuredClone(request),
+        response,
+        aidaResponse,
+      });
+      localStorage.setItem('aiAssistanceStructuredLog', JSON.stringify(this.#structuredLog));
     }
   }
 
-  #createAbortResponse(): ResponseData {
-    return {
-      type: ResponseType.ERROR,
-      error: ErrorType.ABORT,
-    };
-  }
-}
+  #formatParsedStep(step: ParsedStep): string {
+    let text = '';
+    if (step.thought) {
+      text = `THOUGHT: ${step.thought}`;
+    }
+    if (step.title) {
+      text += `\nTITLE: ${step.title}`;
+    }
+    if (step.action) {
+      text += `\nACTION
+${step.action}
+STOP`;
+    }
 
-export function isDebugMode(): boolean {
-  return Boolean(localStorage.getItem('debugFreestylerEnabled'));
-}
+    return text;
+  }
 
-export function debugLog(...log: unknown[]): void {
-  if (!isDebugMode()) {
-    return;
+  #removeLastRunParts(): void {
+    this.#history.splice(this.#history.findLastIndex(item => {
+      return item.role === Host.AidaClient.Role.USER;
+    }));
   }
 
-  // eslint-disable-next-line no-console
-  console.log(...log);
-}
+  #createErrorResponse(error: ErrorType): ResponseData {
+    this.#removeLastRunParts();
+    if (error !== ErrorType.ABORT) {
+      Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceError);
+    }
 
-function setDebugFreestylerEnabled(enabled: boolean): void {
-  if (enabled) {
-    localStorage.setItem('debugFreestylerEnabled', 'true');
-  } else {
-    localStorage.removeItem('debugFreestylerEnabled');
+    return {
+      type: ResponseType.ERROR,
+      error,
+    };
   }
 }
-// @ts-ignore
-globalThis.setDebugFreestylerEnabled = setDebugFreestylerEnabled;