ui/src/plugins/dev.perfetto.Intelletto/agent.ts - third_party/perfetto - Git at Google

 // Copyright (C) 2026 The Android Open Source Project
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 import {z} from 'zod';
 import type {LlmGateway, ModelPath} from '../dev.perfetto.Llm/gateway';
 import type {
   Message,
   ToolCall,
   ToolDef,
   ToolResult,
 } from '../dev.perfetto.Llm/protocol';
 import type {ToolRegistry} from './tools';

 // Hard bound on tool-use rounds per user turn. Stops a misbehaving model from
 // looping forever (and caps token spend). Hitting it surfaces to the user
 // rather than failing silently.
 const MAX_ITERATIONS = 50;

 // What the UI consumes as a turn streams.
 export type AgentEvent =
   | {readonly type: 'text'; readonly text: string}
   | {readonly type: 'thought'; readonly text: string}
   | {readonly type: 'toolcall'; readonly name: string; readonly args: unknown}
   | {
       readonly type: 'toolresult';
       readonly name: string;
       readonly isError: boolean;
       // The (possibly truncated) tool result string, so the UI can show a
       // summary of what the tool returned. This is the same string fed back to
       // the model.
       readonly result: string;
     }
   | {
       readonly type: 'usage';
       readonly inputTokens?: number;
       readonly outputTokens?: number;
       readonly totalTokens?: number;
     }
   // A user-facing error (backend failure, iteration cap, no model configured).
   // Distinct from a tool error, which is fed back to the model, not surfaced.
   | {readonly type: 'error'; readonly message: string};

 export class Agent {
   // The full conversation history (system prompt aside). Owned here because LLM
   // endpoints are stateless and we resend it every request.
   private history: Message[] = [];

   /**
    * @param gateway - The LLM gateway used to drive the model.
    * @param tools - The registry of tools available to the model.
    * @param systemPrompt - The system prompt sent on every request.
    * @param eagerTools - When true (the default), every tool definition is sent
    *     on every request and the list_tools/more_tools meta-tools are omitted.
    *     The lazy scheme (discover via list_tools, provision via more_tools) is
    *     an optimisation for *large* tool sets, but the indirection reliably
    *     trips up smaller / local models - they loop calling more_tools instead
    *     of the real tool. We only have a handful of core tools, so eager is
    *     both simpler and more robust. Flip this off once the tool count makes
    *     eager loading expensive.
    */
   constructor(
     private readonly gateway: LlmGateway,
     private readonly tools: ToolRegistry,
     private readonly systemPrompt: string,
     private readonly eagerTools: boolean = true,
   ) {}

   reset(): void {
     this.history = [];
   }

   /**
    * Run one user turn to completion, yielding events as they stream.
    *
    * @param userText - The user's message for this turn.
    * @param signal - Optional cancellation signal. An aborted turn stops
    *     cleanly and what completed stays in the history (the transcript stays
    *     truthful).
    * @param dequeueFollowUp - Implements queued follow-ups: it is polled once at
    *     the top of each tool-use round, and any message it returns is appended
    *     to the history as a user message, so something the user typed mid-loop
    *     ("oh wait, also check X") reaches the model at the next round without
    *     any interruption machinery.
    * @yields Events as the turn streams.
    */
   async *sendMessage(
     userText: string,
     signal?: AbortSignal,
     dequeueFollowUp?: () => string | undefined,
   ): AsyncGenerator<AgentEvent, void, void> {
     this.history.push({role: 'user', text: userText});

     // The model driving this turn: the first configured model advertising the
     // 'agentic' role. Looked up per turn so adding/removing a provider between
     // turns routes the next turn accordingly (history carries over).
     const modelDetails = this.gateway
       .listModels()
       .find((m) => m.model.roles.includes('agentic'));
     if (modelDetails === undefined) {
       yield {
         type: 'error',
         message:
           'No agentic model configured. Add a provider and a model with the ' +
           'agentic role in the LLM settings.',
       };
       return;
     }
     const modelPath: ModelPath = {
       providerId: modelDetails.provider.id,
       modelId: modelDetails.model.id,
     };

     // createStream requires a signal; use a never-aborting fallback when the
     // caller didn't supply one.
     const abortSignal = signal ?? new AbortController().signal;

     // The set of tool names exposed to the model this turn. In eager mode it's
     // every tool up front; in lazy mode it starts empty and grows as the model
     // pulls tools in via more_tools.
     const loaded = new Set<string>(
       this.eagerTools ? this.tools.list().map((t) => t.name) : [],
     );

     for (let iter = 0; iter < MAX_ITERATIONS; iter++) {
       if (signal?.aborted) return;

       // Pick up any follow-up the user queued while the loop was running.
       const followUp = dequeueFollowUp?.();
       if (followUp !== undefined) {
         this.history.push({role: 'user', text: followUp});
       }

       const toolDefs = this.buildToolDefs(loaded);
       const calls: ToolCall[] = [];
       let turnText = '';
       let backendError: string | undefined;

       for await (const evt of this.gateway.createStream(
         modelPath,
         {
           systemPrompt: this.systemPrompt,
           messages: this.history,
           tools: toolDefs,
         },
         abortSignal,
       )) {
         switch (evt.type) {
           case 'text':
             turnText += evt.text;
             yield {type: 'text', text: evt.text};
             break;
           case 'thought':
             yield {type: 'thought', text: evt.text};
             break;
           case 'tool-call':
             calls.push(evt.call);
             yield {
               type: 'toolcall',
               name: evt.call.name,
               args: evt.call.args,
             };
             break;
           case 'usage':
             yield {type: 'usage', ...evt.usage};
             break;
           case 'stop':
             if (evt.reason === 'error') {
               backendError = evt.error?.message ?? 'Unknown backend error';
             }
             break;
         }
       }

       if (signal?.aborted) return;

       if (backendError !== undefined) {
         yield {type: 'error', message: backendError};
         return;
       }

       // Record the model's turn (text and/or tool calls) in the history.
       if (turnText.length > 0) {
         this.history.push({role: 'model', text: turnText});
       }
       if (calls.length === 0) {
         // No tool calls -> the model is done with this turn.
         return;
       }
       this.history.push({role: 'tool-call', calls});

       // Execute every requested tool, threading results back. Tool errors
       // (bad SQL, invalid args) become tool results the model can self-correct
       // from - they are NOT surfaced to the user as errors.
       const results: ToolResult[] = [];
       for (const call of calls) {
         if (signal?.aborted) return;
         const {result, isError} = await this.runTool(call, loaded);
         yield {type: 'toolresult', name: call.name, isError, result};
         results.push({name: call.name, result, isError});
       }
       this.history.push({role: 'tool-result', results});
     }

     yield {
       type: 'error',
       message:
         `Stopped after ${MAX_ITERATIONS} tool-use rounds without a ` +
         `final answer. Try narrowing the question.`,
     };
   }

   // Execute one tool call. The meta-tools (list_tools / more_tools) are handled
   // here because they mutate the loop's loaded-tool set; everything else goes
   // through the registry (which validates args first).
   private async runTool(
     call: ToolCall,
     loaded: Set<string>,
   ): Promise<{result: string; isError: boolean}> {
     try {
       if (call.name === 'list_tools') {
         const list = this.tools.list().map((t) => ({
           name: t.name,
           description: firstLine(t.description),
         }));
         return {result: JSON.stringify(list), isError: false};
       }
       if (call.name === 'more_tools') {
         const names = z
           .object({names: z.array(z.string())})
           .parse(call.args).names;
         const unknown = names.filter((n) => this.tools.get(n) === undefined);
         names.forEach((n) => this.tools.get(n) && loaded.add(n));
         const msg =
           unknown.length > 0
             ? `Loaded. Unknown (ignored): ${unknown.join(', ')}.`
             : 'Loaded. The requested tools are now available.';
         return {result: msg, isError: unknown.length > 0};
       }
       const result = await this.tools.call(call.name, call.args);
       return {result, isError: false};
     } catch (e) {
       return {result: `Error: ${String(e)}`, isError: true};
     }
   }

   // The tool list sent on a request. In lazy mode this is the bootstrap
   // meta-tools (list_tools/more_tools) plus every tool the model has loaded so
   // far; in eager mode the meta-tools are omitted and `loaded` already holds
   // every tool. Order is stable and the set only grows within a turn,
   // preserving the cached prompt prefix.
   private buildToolDefs(loaded: Set<string>): ToolDef[] {
     const defs: ToolDef[] = [];
     if (!this.eagerTools) {
       defs.push(
         {
           name: 'list_tools',
           description:
             'List the names and one-line descriptions of every tool ' +
             'available. Call this first to discover what you can do, then ' +
             'more_tools to load the ones you need before calling them.',
           inputSchema: z.object({}),
         },
         {
           name: 'more_tools',
           description:
             'Load tool definitions by name so you can call them. You can only ' +
             'call a tool after loading it here. Discover names with list_tools.',
           inputSchema: z.object({
             names: z
               .array(z.string())
               .describe('Tool names to load, from list_tools.'),
           }),
         },
       );
     }
     for (const tool of this.tools.list()) {
       if (loaded.has(tool.name)) {
         defs.push({
           name: tool.name,
           description: tool.description,
           inputSchema: tool.inputSchema,
         });
       }
     }
     return defs;
   }
 }

 function firstLine(s: string): string {
   const trimmed = s.trim();
   const nl = trimmed.indexOf('\n');
   return nl === -1 ? trimmed : trimmed.slice(0, nl);
 }
	// Copyright (C) 2026 The Android Open Source Project
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	import {z} from 'zod';
	import type {LlmGateway, ModelPath} from '../dev.perfetto.Llm/gateway';
	import type {
	Message,
	ToolCall,
	ToolDef,
	ToolResult,
	} from '../dev.perfetto.Llm/protocol';
	import type {ToolRegistry} from './tools';

	// Hard bound on tool-use rounds per user turn. Stops a misbehaving model from
	// looping forever (and caps token spend). Hitting it surfaces to the user
	// rather than failing silently.
	const MAX_ITERATIONS = 50;

	// What the UI consumes as a turn streams.
	export type AgentEvent =
	\| {readonly type: 'text'; readonly text: string}
	\| {readonly type: 'thought'; readonly text: string}
	\| {readonly type: 'toolcall'; readonly name: string; readonly args: unknown}
	\| {
	readonly type: 'toolresult';
	readonly name: string;
	readonly isError: boolean;
	// The (possibly truncated) tool result string, so the UI can show a
	// summary of what the tool returned. This is the same string fed back to
	// the model.
	readonly result: string;
	}
	\| {
	readonly type: 'usage';
	readonly inputTokens?: number;
	readonly outputTokens?: number;
	readonly totalTokens?: number;
	}
	// A user-facing error (backend failure, iteration cap, no model configured).
	// Distinct from a tool error, which is fed back to the model, not surfaced.
	\| {readonly type: 'error'; readonly message: string};

	export class Agent {
	// The full conversation history (system prompt aside). Owned here because LLM
	// endpoints are stateless and we resend it every request.
	private history: Message[] = [];

	/**
	* @param gateway - The LLM gateway used to drive the model.
	* @param tools - The registry of tools available to the model.
	* @param systemPrompt - The system prompt sent on every request.
	* @param eagerTools - When true (the default), every tool definition is sent
	* on every request and the list_tools/more_tools meta-tools are omitted.
	* The lazy scheme (discover via list_tools, provision via more_tools) is
	* an optimisation for large tool sets, but the indirection reliably
	* trips up smaller / local models - they loop calling more_tools instead
	* of the real tool. We only have a handful of core tools, so eager is
	* both simpler and more robust. Flip this off once the tool count makes
	* eager loading expensive.
	*/
	constructor(
	private readonly gateway: LlmGateway,
	private readonly tools: ToolRegistry,
	private readonly systemPrompt: string,
	private readonly eagerTools: boolean = true,
	) {}

	reset(): void {
	this.history = [];
	}

	/**
	* Run one user turn to completion, yielding events as they stream.
	*
	* @param userText - The user's message for this turn.
	* @param signal - Optional cancellation signal. An aborted turn stops
	* cleanly and what completed stays in the history (the transcript stays
	* truthful).
	* @param dequeueFollowUp - Implements queued follow-ups: it is polled once at
	* the top of each tool-use round, and any message it returns is appended
	* to the history as a user message, so something the user typed mid-loop
	* ("oh wait, also check X") reaches the model at the next round without
	* any interruption machinery.
	* @yields Events as the turn streams.
	*/
	async *sendMessage(
	userText: string,
	signal?: AbortSignal,
	dequeueFollowUp?: () => string \| undefined,
	): AsyncGenerator<AgentEvent, void, void> {
	this.history.push({role: 'user', text: userText});

	// The model driving this turn: the first configured model advertising the
	// 'agentic' role. Looked up per turn so adding/removing a provider between
	// turns routes the next turn accordingly (history carries over).
	const modelDetails = this.gateway
	.listModels()
	.find((m) => m.model.roles.includes('agentic'));
	if (modelDetails === undefined) {
	yield {
	type: 'error',
	message:
	'No agentic model configured. Add a provider and a model with the ' +
	'agentic role in the LLM settings.',
	};
	return;
	}
	const modelPath: ModelPath = {
	providerId: modelDetails.provider.id,
	modelId: modelDetails.model.id,
	};

	// createStream requires a signal; use a never-aborting fallback when the
	// caller didn't supply one.
	const abortSignal = signal ?? new AbortController().signal;

	// The set of tool names exposed to the model this turn. In eager mode it's
	// every tool up front; in lazy mode it starts empty and grows as the model
	// pulls tools in via more_tools.
	const loaded = new Set<string>(
	this.eagerTools ? this.tools.list().map((t) => t.name) : [],
	);

	for (let iter = 0; iter < MAX_ITERATIONS; iter++) {
	if (signal?.aborted) return;

	// Pick up any follow-up the user queued while the loop was running.
	const followUp = dequeueFollowUp?.();
	if (followUp !== undefined) {
	this.history.push({role: 'user', text: followUp});
	}

	const toolDefs = this.buildToolDefs(loaded);
	const calls: ToolCall[] = [];
	let turnText = '';
	let backendError: string \| undefined;

	for await (const evt of this.gateway.createStream(
	modelPath,
	{
	systemPrompt: this.systemPrompt,
	messages: this.history,
	tools: toolDefs,
	},
	abortSignal,
	)) {
	switch (evt.type) {
	case 'text':
	turnText += evt.text;
	yield {type: 'text', text: evt.text};
	break;
	case 'thought':
	yield {type: 'thought', text: evt.text};
	break;
	case 'tool-call':
	calls.push(evt.call);
	yield {
	type: 'toolcall',
	name: evt.call.name,
	args: evt.call.args,
	};
	break;
	case 'usage':
	yield {type: 'usage', ...evt.usage};
	break;
	case 'stop':
	if (evt.reason === 'error') {
	backendError = evt.error?.message ?? 'Unknown backend error';
	}
	break;
	}
	}

	if (signal?.aborted) return;

	if (backendError !== undefined) {
	yield {type: 'error', message: backendError};
	return;
	}

	// Record the model's turn (text and/or tool calls) in the history.
	if (turnText.length > 0) {
	this.history.push({role: 'model', text: turnText});
	}
	if (calls.length === 0) {
	// No tool calls -> the model is done with this turn.
	return;
	}
	this.history.push({role: 'tool-call', calls});

	// Execute every requested tool, threading results back. Tool errors
	// (bad SQL, invalid args) become tool results the model can self-correct
	// from - they are NOT surfaced to the user as errors.
	const results: ToolResult[] = [];
	for (const call of calls) {
	if (signal?.aborted) return;
	const {result, isError} = await this.runTool(call, loaded);
	yield {type: 'toolresult', name: call.name, isError, result};
	results.push({name: call.name, result, isError});
	}
	this.history.push({role: 'tool-result', results});
	}

	yield {
	type: 'error',
	message:
	`Stopped after ${MAX_ITERATIONS} tool-use rounds without a ` +
	`final answer. Try narrowing the question.`,
	};
	}

	// Execute one tool call. The meta-tools (list_tools / more_tools) are handled
	// here because they mutate the loop's loaded-tool set; everything else goes
	// through the registry (which validates args first).
	private async runTool(
	call: ToolCall,
	loaded: Set<string>,
	): Promise<{result: string; isError: boolean}> {
	try {
	if (call.name === 'list_tools') {
	const list = this.tools.list().map((t) => ({
	name: t.name,
	description: firstLine(t.description),
	}));
	return {result: JSON.stringify(list), isError: false};
	}
	if (call.name === 'more_tools') {
	const names = z
	.object({names: z.array(z.string())})
	.parse(call.args).names;
	const unknown = names.filter((n) => this.tools.get(n) === undefined);
	names.forEach((n) => this.tools.get(n) && loaded.add(n));
	const msg =
	unknown.length > 0
	? `Loaded. Unknown (ignored): ${unknown.join(', ')}.`
	: 'Loaded. The requested tools are now available.';
	return {result: msg, isError: unknown.length > 0};
	}
	const result = await this.tools.call(call.name, call.args);
	return {result, isError: false};
	} catch (e) {
	return {result: `Error: ${String(e)}`, isError: true};
	}
	}

	// The tool list sent on a request. In lazy mode this is the bootstrap
	// meta-tools (list_tools/more_tools) plus every tool the model has loaded so
	// far; in eager mode the meta-tools are omitted and `loaded` already holds
	// every tool. Order is stable and the set only grows within a turn,
	// preserving the cached prompt prefix.
	private buildToolDefs(loaded: Set<string>): ToolDef[] {
	const defs: ToolDef[] = [];
	if (!this.eagerTools) {
	defs.push(
	{
	name: 'list_tools',
	description:
	'List the names and one-line descriptions of every tool ' +
	'available. Call this first to discover what you can do, then ' +
	'more_tools to load the ones you need before calling them.',
	inputSchema: z.object({}),
	},
	{
	name: 'more_tools',
	description:
	'Load tool definitions by name so you can call them. You can only ' +
	'call a tool after loading it here. Discover names with list_tools.',
	inputSchema: z.object({
	names: z
	.array(z.string())
	.describe('Tool names to load, from list_tools.'),
	}),
	},
	);
	}
	for (const tool of this.tools.list()) {
	if (loaded.has(tool.name)) {
	defs.push({
	name: tool.name,
	description: tool.description,
	inputSchema: tool.inputSchema,
	});
	}
	}
	return defs;
	}
	}

	function firstLine(s: string): string {
	const trimmed = s.trim();
	const nl = trimmed.indexOf('\n');
	return nl === -1 ? trimmed : trimmed.slice(0, nl);
	}