ak-gemini 2.0.7 → 2.1.1

package/package.json

@@ -2,7 +2,7 @@
 	"name": "ak-gemini",
 	"author": "ak@mixpanel.com",
 	"description": "AK's Generative AI Helper for doing... everything",
-	"version": "2.0.7",
+	"version": "2.1.1",
 	"main": "index.js",
 	"files": [
 		"index.js",

package/tool-agent.js

@@ -13,6 +13,31 @@ import log from './logger.js';
  * @typedef {import('./types').AgentStreamEvent} AgentStreamEvent
  */
 
+/**
+ * Execute async task factories with a concurrency limit.
+ * @param {Array<() => Promise<any>>} tasks
+ * @param {number} concurrency - Infinity for unlimited, 1 for sequential
+ * @returns {Promise<any[]>} Results in same order as tasks
+ */
+async function runWithConcurrency(tasks, concurrency) {
+	if (concurrency === Infinity) return Promise.all(tasks.map(t => t()));
+	if (concurrency === 1) {
+		const results = [];
+		for (const t of tasks) results.push(await t());
+		return results;
+	}
+	const results = new Array(tasks.length);
+	let next = 0;
+	async function worker() {
+		while (next < tasks.length) {
+			const i = next++;
+			results[i] = await tasks[i]();
+		}
+	}
+	await Promise.all(Array.from({ length: Math.min(concurrency, tasks.length) }, () => worker()));
+	return results;
+}
+
 /**
  * AI agent that uses user-provided tools to accomplish tasks.
  * Automatically manages the tool-use loop: when the model decides to call
@@ -75,6 +100,13 @@ class ToolAgent extends BaseGemini {
 			throw new Error("ToolAgent: toolExecutor provided without tools. Provide tool declarations so the model knows what tools are available.");
 		}
 
+		// ── Parallel execution ──
+		this.parallelToolCalls = options.parallelToolCalls ?? true;
+		/** @private */
+		this._concurrency = this.parallelToolCalls === true ? Infinity
+			: this.parallelToolCalls === false ? 1
+			: this.parallelToolCalls;
+
 		// ── Tool loop config ──
 		this.maxToolRounds = options.maxToolRounds || 10;
 		this.onToolCall = options.onToolCall || null;
@@ -116,41 +148,39 @@ class ToolAgent extends BaseGemini {
 			const functionCalls = response.functionCalls;
 			if (!functionCalls || functionCalls.length === 0) break;
 
-			const toolResults = await Promise.all(
-				functionCalls.map(async (call) => {
-					// Fire onToolCall callback
-					if (this.onToolCall) {
-						try { this.onToolCall(call.name, call.args); }
-						catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
-					}
+			const tasks = functionCalls.map(call => async () => {
+				// Fire onToolCall callback
+				if (this.onToolCall) {
+					try { this.onToolCall(call.name, call.args); }
+					catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
+				}
 
-					// Check onBeforeExecution gate
-					if (this.onBeforeExecution) {
-						try {
-							const allowed = await this.onBeforeExecution(call.name, call.args);
-							if (allowed === false) {
-								const result = { error: 'Execution denied by onBeforeExecution callback' };
-								allToolCalls.push({ name: call.name, args: call.args, result });
-								return { id: call.id, name: call.name, result };
-							}
-						} catch (e) {
-							log.warn(`onBeforeExecution callback error: ${e.message}`);
+				// Check onBeforeExecution gate
+				if (this.onBeforeExecution) {
+					try {
+						const allowed = await this.onBeforeExecution(call.name, call.args);
+						if (allowed === false) {
+							const result = { error: 'Execution denied by onBeforeExecution callback' };
+							return { id: call.id, name: call.name, args: call.args, result };
 						}
+					} catch (e) {
+						log.warn(`onBeforeExecution callback error: ${e.message}`);
 					}
+				}
 
-					let result;
-					try {
-						result = await this.toolExecutor(call.name, call.args);
-					} catch (err) {
-						log.warn(`Tool ${call.name} failed: ${err.message}`);
-						result = { error: err.message };
-					}
+				let result;
+				try {
+					result = await this.toolExecutor(call.name, call.args);
+				} catch (err) {
+					log.warn(`Tool ${call.name} failed: ${err.message}`);
+					result = { error: err.message };
+				}
 
-					allToolCalls.push({ name: call.name, args: call.args, result });
+				return { id: call.id, name: call.name, args: call.args, result };
+			});
 
-					return { id: call.id, name: call.name, result };
-				})
-			);
+			const toolResults = await runWithConcurrency(tasks, this._concurrency);
+			for (const r of toolResults) allToolCalls.push({ name: r.name, args: r.args, result: r.result });
 
 			// Send function responses back to the model
 			response = await this._withRetry(() => this.chatSession.sendMessage({
@@ -234,46 +264,90 @@ class ToolAgent extends BaseGemini {
 				return;
 			}
 
-			// Execute tools sequentially so we can yield events
+			// Execute tools and yield events
 			const toolResults = [];
-			for (const call of functionCalls) {
-				if (this._stopped) break;
+			if (this._concurrency === 1) {
+				// Sequential: yield tool_call, execute, yield tool_result for each
+				for (const call of functionCalls) {
+					if (this._stopped) break;
 
-				yield { type: 'tool_call', toolName: call.name, args: call.args };
+					yield { type: 'tool_call', toolName: call.name, args: call.args };
 
-				// Fire onToolCall callback
-				if (this.onToolCall) {
-					try { this.onToolCall(call.name, call.args); }
-					catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
-				}
+					if (this.onToolCall) {
+						try { this.onToolCall(call.name, call.args); }
+						catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
+					}
 
-				// Check onBeforeExecution gate
-				let denied = false;
-				if (this.onBeforeExecution) {
-					try {
-						const allowed = await this.onBeforeExecution(call.name, call.args);
-						if (allowed === false) denied = true;
-					} catch (e) {
-						log.warn(`onBeforeExecution callback error: ${e.message}`);
+					let denied = false;
+					if (this.onBeforeExecution) {
+						try {
+							const allowed = await this.onBeforeExecution(call.name, call.args);
+							if (allowed === false) denied = true;
+						} catch (e) {
+							log.warn(`onBeforeExecution callback error: ${e.message}`);
+						}
 					}
-				}
 
-				let result;
-				if (denied) {
-					result = { error: 'Execution denied by onBeforeExecution callback' };
-				} else {
-					try {
-						result = await this.toolExecutor(call.name, call.args);
-					} catch (err) {
-						log.warn(`Tool ${call.name} failed: ${err.message}`);
-						result = { error: err.message };
+					let result;
+					if (denied) {
+						result = { error: 'Execution denied by onBeforeExecution callback' };
+					} else {
+						try {
+							result = await this.toolExecutor(call.name, call.args);
+						} catch (err) {
+							log.warn(`Tool ${call.name} failed: ${err.message}`);
+							result = { error: err.message };
+						}
 					}
+
+					allToolCalls.push({ name: call.name, args: call.args, result });
+					yield { type: 'tool_result', toolName: call.name, result };
+
+					toolResults.push({ id: call.id, name: call.name, result });
+				}
+			} else {
+				// Parallel: yield all tool_call events, execute all, yield all tool_result events
+				for (const call of functionCalls) {
+					yield { type: 'tool_call', toolName: call.name, args: call.args };
 				}
 
-				allToolCalls.push({ name: call.name, args: call.args, result });
-				yield { type: 'tool_result', toolName: call.name, result };
+				const tasks = functionCalls.map(call => async () => {
+					if (this.onToolCall) {
+						try { this.onToolCall(call.name, call.args); }
+						catch (e) { log.warn(`onToolCall callback error: ${e.message}`); }
+					}
 
-				toolResults.push({ id: call.id, name: call.name, result });
+					let denied = false;
+					if (this.onBeforeExecution) {
+						try {
+							const allowed = await this.onBeforeExecution(call.name, call.args);
+							if (allowed === false) denied = true;
+						} catch (e) {
+							log.warn(`onBeforeExecution callback error: ${e.message}`);
+						}
+					}
+
+					let result;
+					if (denied) {
+						result = { error: 'Execution denied by onBeforeExecution callback' };
+					} else {
+						try {
+							result = await this.toolExecutor(call.name, call.args);
+						} catch (err) {
+							log.warn(`Tool ${call.name} failed: ${err.message}`);
+							result = { error: err.message };
+						}
+					}
+
+					return { id: call.id, name: call.name, args: call.args, result };
+				});
+
+				const results = await runWithConcurrency(tasks, this._concurrency);
+				for (const r of results) {
+					allToolCalls.push({ name: r.name, args: r.args, result: r.result });
+					yield { type: 'tool_result', toolName: r.name, result: r.result };
+					toolResults.push({ id: r.id, name: r.name, result: r.result });
+				}
 			}
 
 			// Send function responses back and get next stream

package/types.d.ts

@@ -279,6 +279,8 @@ export interface ToolAgentOptions extends BaseGeminiOptions {
   onBeforeExecution?: (toolName: string, args: Record<string, any>) => Promise<boolean>;
   /** Directory for tool-written files (pass-through for toolExecutor use) */
   writeDir?: string;
+  /** Parallel tool execution: false = sequential, true = unlimited parallel, number = concurrency limit (default: true) */
+  parallelToolCalls?: boolean | number;
 }
 
 export interface LocalDataEntry {
@@ -304,9 +306,9 @@ export interface CodeAgentOptions extends BaseGeminiOptions {
   maxRounds?: number;
   /** Per-execution timeout in milliseconds (default: 30000) */
   timeout?: number;
-  /** Async callback before code execution; return false to deny */
-  onBeforeExecution?: (code: string) => Promise<boolean>;
-  /** Notification callback after code execution */
+  /** Async callback before code/bash execution; return false to deny. Receives (content, toolName). */
+  onBeforeExecution?: (content: string, toolName: string) => Promise<boolean> | boolean;
+  /** Notification callback after code/bash execution */
   onCodeExecution?: (code: string, output: { stdout: string; stderr: string; exitCode: number }) => void;
   /** Files whose contents are included in the system prompt for project context */
   importantFiles?: string[];
@@ -318,6 +320,10 @@ export interface CodeAgentOptions extends BaseGeminiOptions {
   comments?: boolean;
   /** Max consecutive failed executions before stopping (default: 3) */
   maxRetries?: number;
+  /** Paths to skill files (markdown) loaded dynamically via the use_skill tool */
+  skills?: string[];
+  /** Plain text environment overview appended to the system prompt — describe the project, stack, conventions, etc. */
+  envOverview?: string;
 }
 
 export interface CodeExecution {
@@ -333,35 +339,57 @@ export interface CodeExecution {
   exitCode: number;
 }
 
+export interface ToolCallResult {
+  tool: 'write_code' | 'execute_code' | 'write_and_run_code' | 'fix_code' | 'run_bash' | 'use_skill';
+  code?: string;
+  purpose?: string;
+  language?: string;
+  originalCode?: string;
+  fixedCode?: string;
+  explanation?: string;
+  executed?: boolean;
+  command?: string;
+  skillName?: string;
+  content?: string;
+  found?: boolean;
+  stdout?: string;
+  stderr?: string;
+  exitCode?: number;
+  denied?: boolean;
+}
+
 export interface CodeAgentResponse {
   /** The agent's final text response */
   text: string;
-  /** All code executions during this interaction */
+  /** Backward-compatible: only code executions (execute_code, write_and_run_code, fix_code with execute) */
   codeExecutions: CodeExecution[];
+  /** All tool calls made during this chat turn */
+  toolCalls: ToolCallResult[];
   /** Token usage data */
   usage: UsageData | null;
 }
 
 export interface CodeAgentStreamEvent {
-  type: 'text' | 'code' | 'output' | 'done';
-  /** For 'text' events: the text chunk */
+  type: 'text' | 'code' | 'output' | 'write' | 'fix' | 'bash' | 'skill' | 'done';
   text?: string;
-  /** For 'code' events: the code about to be executed */
   code?: string;
-  /** For 'output' events: stdout from execution */
   stdout?: string;
-  /** For 'output' events: stderr from execution */
   stderr?: string;
-  /** For 'output' events: process exit code */
   exitCode?: number;
-  /** For 'done' events: complete accumulated text */
   fullText?: string;
-  /** For 'done' events: all code executions */
   codeExecutions?: CodeExecution[];
-  /** For 'done' events: token usage */
+  toolCalls?: ToolCallResult[];
   usage?: UsageData | null;
-  /** For 'done' events: e.g. "Max tool rounds reached" or "Agent was stopped" */
   warning?: string;
+  purpose?: string;
+  language?: string;
+  originalCode?: string;
+  fixedCode?: string;
+  explanation?: string;
+  command?: string;
+  skillName?: string;
+  content?: string;
+  found?: boolean;
 }
 
 // ── Per-Message Options ──────────────────────────────────────────────────────
@@ -539,6 +567,7 @@ export declare class ToolAgent extends BaseGemini {
   onBeforeExecution: ((toolName: string, args: Record<string, any>) => Promise<boolean>) | null;
   /** Directory for tool-written files (pass-through for toolExecutor use) */
   writeDir: string | null;
+  parallelToolCalls: boolean | number;
 
   chat(message: string, opts?: { labels?: Record<string, string> }): Promise<AgentResponse>;
   stream(message: string, opts?: { labels?: Record<string, string> }): AsyncGenerator<AgentStreamEvent, void, unknown>;
@@ -579,25 +608,20 @@ export declare class CodeAgent extends BaseGemini {
   workingDirectory: string;
   maxRounds: number;
   timeout: number;
-  onBeforeExecution: ((code: string) => Promise<boolean>) | null;
+  onBeforeExecution: ((content: string, toolName: string) => Promise<boolean> | boolean) | null;
   onCodeExecution: ((code: string, output: { stdout: string; stderr: string; exitCode: number }) => void) | null;
-  /** Files whose contents are included in the system prompt */
   importantFiles: string[];
-  /** Directory for writing script files */
   writeDir: string;
-  /** Keep script files on disk after execution */
   keepArtifacts: boolean;
-  /** Whether the model writes comments in generated code */
   comments: boolean;
-  /** Max consecutive failed executions before stopping */
   maxRetries: number;
+  skills: string[];
+  envOverview: string;
 
   init(force?: boolean): Promise<void>;
   chat(message: string, opts?: { labels?: Record<string, string> }): Promise<CodeAgentResponse>;
   stream(message: string, opts?: { labels?: Record<string, string> }): AsyncGenerator<CodeAgentStreamEvent, void, unknown>;
-  /** Returns all code scripts written across all chat/stream calls. */
-  dump(): Array<{ fileName: string; purpose: string | null; script: string; filePath: string | null }>;
-  /** Stop the agent before the next code execution. Kills any running child process. */
+  dump(): Array<{ fileName: string; purpose: string | null; script: string; filePath: string | null; tool: string }>;
   stop(): void;
 }