@agentuity/pi 1.0.36 → 2.0.16

package/src/index.ts

@@ -1,243 +1,465 @@
+/**
+ * Agentuity AI Gateway Custom Provider Extension
+ *
+ * Registers models from the Agentuity AI Gateway using the appropriate API type
+ * based on model ID patterns. Models are loaded dynamically from the gateway's /models endpoint.
+ *
+ * Usage:
+ *   Use /model to switch to aigateway models
+ */
+import { execFileSync } from 'node:child_process';
+import { existsSync } from 'node:fs';
+import { delimiter, join } from 'node:path';
+import { createMinimalLogger, StructuredError } from '@agentuity/core';
+import {
+	AIGatewayService,
+	type AIGatewayModel,
+	type AIGatewayModels,
+} from '@agentuity/core/aigateway';
+import { createServerFetchAdapter } from '@agentuity/server';
 import type {
-	AgentToolResult,
 	ExtensionAPI,
-	ExtensionContext,
 	ExtensionCommandContext,
+	ProviderModelConfig,
 } from '@mariozechner/pi-coding-agent';
-import type { TSchema } from '@sinclair/typebox';
-import { HubClient } from './client.ts';
-import { processActions } from './handlers.ts';
-import type { HubResponse, InitMessage } from './protocol.ts';
-
-const HUB_URL_ENV = 'AGENTUITY_CODER_HUB_URL';
-
-// All Pi events we subscribe to (order matters — session_start is handled separately)
-const PROXY_EVENTS = [
-	'session_shutdown',
-	'session_before_switch',
-	'session_switch',
-	'session_before_fork',
-	'session_fork',
-	'session_before_compact',
-	'session_compact',
-	'before_agent_start',
-	'agent_start',
-	'agent_end',
-	'turn_start',
-	'turn_end',
-	'tool_call',
-	'tool_result',
-	'tool_execution_start',
-	'tool_execution_update',
-	'tool_execution_end',
-	'message_start',
-	'message_update',
-	'message_end',
-	'input',
-	'model_select',
-	'context',
-] as const;
-
-// Generic event handler type for the iteration loop
-type GenericEventHandler = (
-	event: string,
-	handler: (event: unknown, ctx: ExtensionContext) => Promise<unknown>
-) => void;
-
-function log(msg: string): void {
-	console.error(`[agentuity-pi] ${msg}`);
-}
-
-export function agentuityCoderHub(pi: ExtensionAPI) {
-	const hubUrl = process.env[HUB_URL_ENV];
-
-	// No-op if not configured
-	if (!hubUrl) {
-		return;
-	}
-
-	log(`Hub URL: ${hubUrl}`);
-
-	const client = new HubClient();
 
-	// Connect to the Hub and register tools/commands
-	async function connectAndRegister(): Promise<void> {
-		if (client.connected) return;
-
-		log('Connecting to Hub...');
-		let initMsg: InitMessage;
+export type KnownApi =
+	| 'openai-completions'
+	| 'mistral-conversations'
+	| 'openai-responses'
+	| 'azure-openai-responses'
+	| 'openai-codex-responses'
+	| 'anthropic-messages'
+	| 'bedrock-converse-stream'
+	| 'google-generative-ai'
+	| 'google-gemini-cli'
+	| 'google-vertex';
+
+const KNOWN_APIS = new Set<string>([
+	'openai-completions',
+	'mistral-conversations',
+	'openai-responses',
+	'azure-openai-responses',
+	'openai-codex-responses',
+	'anthropic-messages',
+	'bedrock-converse-stream',
+	'google-generative-ai',
+	'google-gemini-cli',
+	'google-vertex',
+] satisfies KnownApi[]);
+
+const AIGatewayModelFetchError = StructuredError('AIGatewayModelFetchError')<{
+	cause?: unknown;
+}>();
+
+type AgentuityOrganization = {
+	id: string;
+	name: string;
+};
+
+type AgentuityWhoami = {
+	organizations?: AgentuityOrganization[];
+};
+
+type AgentuityRegion = {
+	region: string;
+	description: string;
+	default?: boolean;
+};
+
+function parseFirstJsonObject(value: string): unknown {
+	const start = value.indexOf('{');
+	if (start === -1) {
+		throw new SyntaxError('No JSON object found');
+	}
 
-		try {
-			initMsg = await client.connect(hubUrl!);
-		} catch (err) {
-			const msg = err instanceof Error ? err.message : String(err);
-			log(`Failed to connect: ${msg}`);
-			return;
+	let depth = 0;
+	let inString = false;
+	let escaped = false;
+	for (let i = start; i < value.length; i++) {
+		const char = value[i];
+		if (escaped) {
+			escaped = false;
+			continue;
 		}
-
-		log(
-			`Connected. Init: ${initMsg.tools?.length ?? 0} tools, ${initMsg.commands?.length ?? 0} commands`
-		);
-
-		// Register tools from Hub
-		if (initMsg.tools) {
-			for (const toolDef of initMsg.tools) {
-				log(`Registering tool: ${toolDef.name}`);
-				pi.registerTool({
-					name: toolDef.name,
-					label: toolDef.label,
-					description: toolDef.description,
-					parameters: toolDef.parameters as unknown as TSchema,
-					async execute(
-						toolCallId: string,
-						params: unknown,
-						_signal: AbortSignal | undefined,
-						_onUpdate: unknown,
-						ctx: ExtensionContext
-					) {
-						log(`Tool execute: ${toolDef.name}`);
-						const id = client.nextId();
-						let response: HubResponse;
-
-						try {
-							response = await client.send({
-								id,
-								type: 'tool',
-								name: toolDef.name,
-								toolCallId,
-								params: params as Record<string, unknown>,
-							});
-						} catch {
-							return {
-								content: [
-									{
-										type: 'text' as const,
-										text: 'Error: Hub connection lost',
-									},
-								],
-								details: {},
-							};
-						}
-
-						const result = await processActions(response.actions, ctx);
-
-						if (result.returnValue !== undefined) {
-							return result.returnValue as AgentToolResult<unknown>;
-						}
-
-						return {
-							content: [{ type: 'text' as const, text: 'Done' }],
-							details: {},
-						};
-					},
-				});
+		if (char === '\\') {
+			escaped = true;
+			continue;
+		}
+		if (char === '"') {
+			inString = !inString;
+			continue;
+		}
+		if (inString) {
+			continue;
+		}
+		if (char === '{') {
+			depth++;
+		} else if (char === '}') {
+			depth--;
+			if (depth === 0) {
+				return JSON.parse(value.slice(start, i + 1));
 			}
 		}
+	}
 
-		// Register commands from Hub
-		if (initMsg.commands) {
-			for (const cmdDef of initMsg.commands) {
-				log(`Registering command: /${cmdDef.name}`);
-				pi.registerCommand(cmdDef.name, {
-					description: cmdDef.description,
-					handler: async (args: string, ctx: ExtensionCommandContext) => {
-						log(`Command execute: /${cmdDef.name}`);
-						const id = client.nextId();
-						let response: HubResponse;
-
-						try {
-							response = await client.send({
-								id,
-								type: 'command',
-								name: cmdDef.name,
-								args,
-							});
-						} catch {
-							ctx.ui.notify('Hub connection lost', 'error');
-							return;
-						}
-
-						await processActions(response.actions, ctx);
-					},
-				});
-			}
+	throw new SyntaxError('Unterminated JSON object');
+}
+
+function parseJson(value: string): unknown {
+	const trimmed = value.trim();
+	if (trimmed.startsWith('[')) {
+		return JSON.parse(trimmed);
+	}
+	return parseFirstJsonObject(trimmed);
+}
+
+function getEnv(...keys: string[]): string | undefined {
+	for (const key of keys) {
+		if (process.env[key]) {
+			return process.env[key];
 		}
+	}
+}
 
-		log('Registration complete');
+function normalizeCredential(value: unknown): string | undefined {
+	if (value === undefined || value === null) {
+		return undefined;
 	}
+	const normalized = String(value).trim();
+	return normalized.length > 0 ? normalized : undefined;
+}
 
-	// Helper to send event and process response actions
-	async function sendEvent(
-		eventName: string,
-		data: Record<string, unknown>,
-		ctx: ExtensionContext
-	): Promise<unknown> {
-		if (!client.connected) return undefined;
+function isKnownApi(api: unknown): api is KnownApi {
+	return typeof api === 'string' && KNOWN_APIS.has(api);
+}
 
-		const id = client.nextId();
-		let response: HubResponse;
+function getRegion(): string {
+	return getEnv('AGENTUITY_REGION') ?? 'usc';
+}
 
-		try {
-			response = await client.send({
-				id,
-				type: 'event',
-				event: eventName,
-				data,
-			});
-		} catch {
-			return undefined;
+function getBaseUrl(): string {
+	const region = getRegion();
+	return `https://aigateway-${region}.agentuity.cloud`;
+}
+
+function getAgentuityCliPath(): string | undefined {
+	const path = process.env.PATH?.split(delimiter) ?? [];
+	for (const dir of path) {
+		const fn = join(dir, 'agentuity');
+		if (existsSync(fn)) {
+			return fn;
 		}
+	}
+}
 
-		const result = await processActions(response.actions, ctx);
+function fetchOrganizations(): AgentuityOrganization[] {
+	const agentuity = getAgentuityCliPath();
+	if (!agentuity) {
+		return [];
+	}
 
-		if (result.block) return result.block;
-		if (result.returnValue !== undefined) return result.returnValue;
-		return undefined;
+	const res = execFileSync(agentuity, ['auth', 'whoami', '--json']);
+	const whoami = parseJson(res.toString()) as AgentuityWhoami;
+	return (whoami.organizations ?? []).filter(
+		(org): org is AgentuityOrganization =>
+			typeof org?.id === 'string' &&
+			org.id.length > 0 &&
+			typeof org?.name === 'string' &&
+			org.name.length > 0
+	);
+}
+
+function fetchRegions(): AgentuityRegion[] {
+	const agentuity = getAgentuityCliPath();
+	if (!agentuity) {
+		return [];
 	}
 
-	// Serialize event data — strip non-serializable fields
-	function serializeEvent(event: unknown): Record<string, unknown> {
-		const data: Record<string, unknown> = {};
-		if (event && typeof event === 'object') {
-			for (const [key, value] of Object.entries(event)) {
-				if (typeof value !== 'function' && key !== 'signal') {
-					try {
-						JSON.stringify(value);
-						data[key] = value;
-					} catch {
-						// Skip non-serializable values
+	const res = execFileSync(agentuity, ['cloud', 'region', 'list', '--json']);
+	const regions = parseJson(res.toString()) as AgentuityRegion[];
+	return (Array.isArray(regions) ? regions : []).filter(
+		(region): region is AgentuityRegion =>
+			typeof region?.region === 'string' &&
+			region.region.length > 0 &&
+			typeof region?.description === 'string' &&
+			region.description.length > 0
+	);
+}
+
+function getCurrentOrgId(): string | undefined {
+	return normalizeCredential(
+		getEnv(
+			'AGENTUITY_AIGATEWAY_ORGID',
+			'AGENTUITY_ORGID',
+			'AGENTUITY_CLOUD_ORG_ID',
+			'AGENTUITY_ORG_ID'
+		)
+	);
+}
+
+async function fetchModels(): Promise<AIGatewayModels> {
+	const baseUrl = getBaseUrl();
+	let apiKey = normalizeCredential(
+		getEnv(
+			'AGENTUITY_CODER_API_KEY',
+			'AGENTUITY_SDK_KEY',
+			'AGENTUITY_CLI_API_KEY',
+			'AGENTUITY_CLI_KEY'
+		)
+	);
+	let orgId = normalizeCredential(
+		getEnv('AGENTUITY_ORGID', 'AGENTUITY_CLOUD_ORG_ID', 'AGENTUITY_ORG_ID')
+	);
+
+	if (!apiKey) {
+		let found = false;
+		const fn = getAgentuityCliPath();
+		if (fn) {
+			try {
+				const res = execFileSync(fn, ['auth', 'apikey', '--json']);
+				const apiKeyResult = parseJson(res.toString()) as { apiKey: string };
+				apiKey = normalizeCredential(apiKeyResult.apiKey);
+				found = true;
+				if (!orgId) {
+					const ores = execFileSync(fn, ['auth', 'org', 'current']);
+					orgId = normalizeCredential(ores);
+					if (!orgId) {
+						return {};
 					}
 				}
+			} catch (error) {
+				throw new AIGatewayModelFetchError({
+					message: 'Failed to fetch models from AI Gateway',
+					cause: error,
+				});
 			}
 		}
-		return data;
+		if (!found) {
+			console.warn(
+				'AGENTUITY_SDK_KEY, AGENTUITY_CLI_API_KEY or AGENTUITY_CLI_KEY not set and cannot find the agentuity cli, cannot fetch models from AI Gateway'
+			);
+			return {};
+		}
+	}
+
+	if (!apiKey) {
+		console.warn('Cannot determine the API key, cannot fetch models from AI Gateway');
+		return {};
+	}
+
+	process.env.AGENTUITY_AIGATEWAY_KEY = apiKey;
+	if (orgId) {
+		process.env.AGENTUITY_AIGATEWAY_ORGID = orgId;
+	}
+
+	try {
+		const service = new AIGatewayService(
+			baseUrl,
+			createServerFetchAdapter({ headers: {} }, createMinimalLogger())
+		);
+		return await service.listModels();
+	} catch (error) {
+		throw new AIGatewayModelFetchError({
+			message: 'Failed to fetch models from AI Gateway',
+			cause: error,
+		});
 	}
+}
+
+function sanitizeModalities(modalities: string[] | undefined): ('text' | 'image')[] {
+	const sanitized = (modalities ?? []).filter(
+		(modality): modality is 'text' | 'image' => modality === 'text' || modality === 'image'
+	);
+	return sanitized.length > 0 ? sanitized : ['text'];
+}
 
-	// Cast pi.on to generic handler since we iterate over a union of event names
-	// TypeScript overloads require exact string literal matching which doesn't work in loops
-	const onEvent = pi.on.bind(pi) as GenericEventHandler;
+function toPiModel(m: AIGatewayModel): ProviderModelConfig {
+	return {
+		id: m.id,
+		name: m.name,
+		reasoning: m.reasoning ?? false,
+		input: sanitizeModalities(m.input_modalities),
+		contextWindow: m.context_window ?? 40000,
+		maxTokens: m.max_output_tokens ?? 64000,
+		cost: {
+			input: m.pricing?.input ?? 0,
+			output: m.pricing?.output ?? 0,
+			cacheRead: m.pricing?.cached_input ?? 0,
+			cacheWrite: 0,
+		},
+		compat: {
+			supportsDeveloperRole: false,
+		},
+	};
+}
 
-	// ── session_start: connect + register BEFORE Pi builds the tool list ──
-	onEvent('session_start', async (event: unknown, ctx: ExtensionContext) => {
-		await connectAndRegister();
+async function registerAIGatewayProviders(pi: ExtensionAPI) {
+	const models = await fetchModels();
+	const baseUrl = getBaseUrl();
 
-		if (client.connected) {
-			return sendEvent('session_start', serializeEvent(event), ctx);
+	const allModels: AIGatewayModel[] = [];
+	for (const providerModels of Object.values(models)) {
+		if (providerModels) {
+			allModels.push(...providerModels);
 		}
-	});
+	}
+	if (allModels.length === 0) {
+		return;
+	}
+
+	const modelsByApi = new Map<KnownApi, ProviderModelConfig[]>();
+
+	for (const m of allModels) {
+		const apiType = m.api;
+		if (!isKnownApi(apiType)) {
+			continue;
+		}
+		const existing = modelsByApi.get(apiType) ?? [];
+		existing.push(toPiModel(m));
+		modelsByApi.set(apiType, existing);
+	}
+
+	const headers: Record<string, string> = {};
+	if (process.env.AGENTUITY_AIGATEWAY_ORGID) {
+		headers['x-agentuity-orgid'] = process.env.AGENTUITY_AIGATEWAY_ORGID;
+	}
 
-	for (const eventName of PROXY_EVENTS) {
-		onEvent(eventName, async (event: unknown, ctx: ExtensionContext) => {
-			if (!client.connected) return undefined;
-			return sendEvent(eventName, serializeEvent(event), ctx);
+	for (const [apiType, providerModels] of modelsByApi) {
+		const apitok = apiType.split('-');
+		const name = apitok.length >= 2 ? apitok.slice(0, 2).join('-') : apitok[0];
+		const providerName = `agentuity/${name}`;
+		pi.registerProvider(providerName, {
+			baseUrl,
+			apiKey: 'AGENTUITY_AIGATEWAY_KEY',
+			headers,
+			authHeader: true,
+			api: apiType,
+			models: providerModels,
 		});
 	}
+}
+
+function registerRegionCommand(pi: ExtensionAPI): void {
+	pi.registerCommand('region', {
+		description: 'Select active Agentuity region',
+		handler: async (_args: string, ctx: ExtensionCommandContext) => {
+			if (!ctx.hasUI) {
+				return;
+			}
+
+			let regions: AgentuityRegion[];
+			try {
+				regions = fetchRegions();
+			} catch (error) {
+				ctx.ui.notify(`Failed to load Agentuity regions: ${String(error)}`, 'error');
+				return;
+			}
+
+			if (regions.length === 0) {
+				ctx.ui.notify('No Agentuity regions found.', 'warning');
+				return;
+			}
+
+			const currentRegion = getRegion();
+			const labels = regions.map((region) => {
+				const marker = region.region === currentRegion ? '* ' : '';
+				const defaultLabel = region.default ? ' default' : '';
+				return `${marker}${region.description} (${region.region})${defaultLabel}`;
+			});
+			const selected = await ctx.ui.select('Select Agentuity Region', labels);
+			if (!selected) {
+				return;
+			}
+
+			const selectedIndex = labels.indexOf(selected);
+			const region = regions[selectedIndex];
+			if (!region) {
+				return;
+			}
 
-	// Clean up on shutdown
-	pi.on('session_shutdown', async () => {
-		log('Shutting down');
-		client.close();
+			process.env.AGENTUITY_REGION = region.region;
+			ctx.ui.notify(`Using Agentuity region: ${region.description}`, 'info');
+
+			try {
+				await registerAIGatewayProviders(pi);
+				ctx.ui.notify('Agentuity AI Gateway models refreshed.', 'info');
+			} catch (error) {
+				ctx.ui.notify(`Failed to refresh AI Gateway models: ${String(error)}`, 'error');
+			}
+		},
+	});
+}
+
+function registerOrganizationCommand(pi: ExtensionAPI): void {
+	pi.registerCommand('organization', {
+		description: 'Select active Agentuity organization',
+		handler: async (_args: string, ctx: ExtensionCommandContext) => {
+			if (!ctx.hasUI) {
+				return;
+			}
+
+			let organizations: AgentuityOrganization[];
+			try {
+				organizations = fetchOrganizations();
+			} catch (error) {
+				ctx.ui.notify(`Failed to load Agentuity organizations: ${String(error)}`, 'error');
+				return;
+			}
+
+			if (organizations.length === 0) {
+				ctx.ui.notify('No Agentuity organizations found for the current CLI login.', 'warning');
+				return;
+			}
+
+			const currentOrgId = getCurrentOrgId();
+			const labels = organizations.map((org) => {
+				const marker = org.id === currentOrgId ? '* ' : '';
+				return `${marker}${org.name} (${org.id})`;
+			});
+			const selected = await ctx.ui.select('Select Agentuity Organization', labels);
+			if (!selected) {
+				return;
+			}
+
+			const selectedIndex = labels.indexOf(selected);
+			const organization = organizations[selectedIndex];
+			if (!organization) {
+				return;
+			}
+
+			process.env.AGENTUITY_AIGATEWAY_ORGID = organization.id;
+			process.env.AGENTUITY_ORGID = organization.id;
+			process.env.AGENTUITY_CLOUD_ORG_ID = organization.id;
+			process.env.AGENTUITY_ORG_ID = organization.id;
+
+			ctx.ui.notify(`Using Agentuity organization: ${organization.name}`, 'info');
+
+			try {
+				await registerAIGatewayProviders(pi);
+				ctx.ui.notify('Agentuity AI Gateway models refreshed.', 'info');
+			} catch (error) {
+				ctx.ui.notify(`Failed to refresh AI Gateway models: ${String(error)}`, 'error');
+			}
+		},
+	});
+}
+
+export async function setupAIGateway(pi: ExtensionAPI) {
+	registerOrganizationCommand(pi);
+	registerRegionCommand(pi);
+	let showedOrganizationPrompt = false;
+	pi.on('session_start', (_event, ctx) => {
+		if (showedOrganizationPrompt || !ctx.hasUI || getCurrentOrgId()) {
+			return;
+		}
+		showedOrganizationPrompt = true;
+		ctx.ui.notify(
+			'Use /organization to select an Agentuity organization for AI Gateway models.',
+			'info'
+		);
 	});
+	await registerAIGatewayProviders(pi);
 }
 
-export default agentuityCoderHub;
+export default setupAIGateway;

package/dist/client.d.ts

@@ -1,11 +0,0 @@
-import type { InitMessage, HubRequest, HubResponse } from './protocol.ts';
-export declare class HubClient {
-    private ws;
-    private pending;
-    connect(url: string): Promise<InitMessage>;
-    nextId(): string;
-    send(request: HubRequest): Promise<HubResponse>;
-    close(): void;
-    get connected(): boolean;
-}
-//# sourceMappingURL=client.d.ts.map

package/dist/client.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"client.d.ts","sourceRoot":"","sources":["../src/client.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,WAAW,EAAE,UAAU,EAAE,WAAW,EAAE,MAAM,eAAe,CAAC;AAQ1E,qBAAa,SAAS;IACrB,OAAO,CAAC,EAAE,CAA0B;IACpC,OAAO,CAAC,OAAO,CAOX;IACE,OAAO,CAAC,GAAG,EAAE,MAAM,GAAG,OAAO,CAAC,WAAW,CAAC;IAuFhD,MAAM,IAAI,MAAM;IAIV,IAAI,CAAC,OAAO,EAAE,UAAU,GAAG,OAAO,CAAC,WAAW,CAAC;IAwBrD,KAAK,IAAI,IAAI;IAOb,IAAI,SAAS,IAAI,OAAO,CAEvB;CACD"}