pi-gitlab-duo 0.1.0 → 0.1.1

package/index.ts

@@ -1,18 +1,12 @@
 /**
- * GitLab Duo Provider Extension for pi
+ * GitLab Duo Provider Extension
  *
  * Provides access to GitLab Duo AI models (Claude and GPT) through GitLab's AI Gateway.
  * Delegates to pi-ai's built-in Anthropic and OpenAI streaming implementations.
  *
- * Installation:
- *   pi install npm:pi-gitlab-duo
- *
  * Usage:
- *   # With OAuth (run /login gitlab-duo first)
- *   pi --provider gitlab-duo --model duo-chat-sonnet-4-5
- *
- *   # With PAT
- *   GITLAB_TOKEN=glpat-... pi --provider gitlab-duo --model duo-chat-sonnet-4-5
+ *   pi -e ./packages/coding-agent/examples/extensions/custom-provider-gitlab-duo
+ *   # Then /login gitlab-duo, or set GITLAB_TOKEN=glpat-...
  */
 
 import {
@@ -24,7 +18,8 @@ import {
 	type OAuthCredentials,
 	type OAuthLoginCallbacks,
 	type SimpleStreamOptions,
-	streamSimple,
+	streamSimpleAnthropic,
+	streamSimpleOpenAIResponses,
 } from "@mariozechner/pi-ai";
 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 
@@ -37,42 +32,101 @@ const AI_GATEWAY_URL = "https://cloud.gitlab.com";
 const ANTHROPIC_PROXY_URL = `${AI_GATEWAY_URL}/ai/v1/proxy/anthropic/`;
 const OPENAI_PROXY_URL = `${AI_GATEWAY_URL}/ai/v1/proxy/openai/v1`;
 
-// Bundled OAuth client ID for gitlab.com (from opencode-gitlab-auth, registered with localhost redirect)
 const BUNDLED_CLIENT_ID = "1d89f9fdb23ee96d4e603201f6861dab6e143c5c3c00469a018a2d94bdc03d4e";
 const OAUTH_SCOPES = ["api"];
 const REDIRECT_URI = "http://127.0.0.1:8080/callback";
-
-// Direct access token cache (25 min, tokens expire after 30 min)
 const DIRECT_ACCESS_TTL = 25 * 60 * 1000;
 
-// Model mappings: duo model ID -> backend config
-const MODEL_MAPPINGS: Record<
-	string,
-	{ api: "anthropic-messages" | "openai-completions"; backendModel: string; baseUrl: string }
-> = {
-	"duo-chat-opus-4-5": {
-		api: "anthropic-messages",
-		backendModel: "claude-opus-4-5-20251101",
+// =============================================================================
+// Models - exported for use by tests
+// =============================================================================
+
+type Backend = "anthropic" | "openai";
+
+interface GitLabModel {
+	id: string;
+	name: string;
+	backend: Backend;
+	baseUrl: string;
+	reasoning: boolean;
+	input: ("text" | "image")[];
+	cost: { input: number; output: number; cacheRead: number; cacheWrite: number };
+	contextWindow: number;
+	maxTokens: number;
+}
+
+export const MODELS: GitLabModel[] = [
+	// Anthropic
+	{
+		id: "claude-opus-4-5-20251101",
+		name: "Claude Opus 4.5",
+		backend: "anthropic",
 		baseUrl: ANTHROPIC_PROXY_URL,
+		reasoning: true,
+		input: ["text", "image"],
+		cost: { input: 15, output: 75, cacheRead: 1.5, cacheWrite: 18.75 },
+		contextWindow: 200000,
+		maxTokens: 32000,
 	},
-	"duo-chat-sonnet-4-5": {
-		api: "anthropic-messages",
-		backendModel: "claude-sonnet-4-5-20250929",
+	{
+		id: "claude-sonnet-4-5-20250929",
+		name: "Claude Sonnet 4.5",
+		backend: "anthropic",
 		baseUrl: ANTHROPIC_PROXY_URL,
+		reasoning: true,
+		input: ["text", "image"],
+		cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 },
+		contextWindow: 200000,
+		maxTokens: 16384,
 	},
-	"duo-chat-haiku-4-5": {
-		api: "anthropic-messages",
-		backendModel: "claude-haiku-4-5-20251001",
+	{
+		id: "claude-haiku-4-5-20251001",
+		name: "Claude Haiku 4.5",
+		backend: "anthropic",
 		baseUrl: ANTHROPIC_PROXY_URL,
+		reasoning: true,
+		input: ["text", "image"],
+		cost: { input: 1, output: 5, cacheRead: 0.1, cacheWrite: 1.25 },
+		contextWindow: 200000,
+		maxTokens: 8192,
+	},
+	// OpenAI (all use Responses API)
+	{
+		id: "gpt-5.1-2025-11-13",
+		name: "GPT-5.1",
+		backend: "openai",
+		baseUrl: OPENAI_PROXY_URL,
+		reasoning: true,
+		input: ["text", "image"],
+		cost: { input: 2.5, output: 10, cacheRead: 0, cacheWrite: 0 },
+		contextWindow: 128000,
+		maxTokens: 16384,
 	},
-	"duo-chat-gpt-5-1": { api: "openai-completions", backendModel: "gpt-5.1-2025-11-13", baseUrl: OPENAI_PROXY_URL },
-	"duo-chat-gpt-5-mini": {
-		api: "openai-completions",
-		backendModel: "gpt-5-mini-2025-08-07",
+	{
+		id: "gpt-5-mini-2025-08-07",
+		name: "GPT-5 Mini",
+		backend: "openai",
 		baseUrl: OPENAI_PROXY_URL,
+		reasoning: true,
+		input: ["text", "image"],
+		cost: { input: 0.15, output: 0.6, cacheRead: 0, cacheWrite: 0 },
+		contextWindow: 128000,
+		maxTokens: 16384,
 	},
-	"duo-chat-gpt-5-codex": { api: "openai-completions", backendModel: "gpt-5-codex", baseUrl: OPENAI_PROXY_URL },
-};
+	{
+		id: "gpt-5-codex",
+		name: "GPT-5 Codex",
+		backend: "openai",
+		baseUrl: OPENAI_PROXY_URL,
+		reasoning: true,
+		input: ["text", "image"],
+		cost: { input: 2.5, output: 10, cacheRead: 0, cacheWrite: 0 },
+		contextWindow: 128000,
+		maxTokens: 16384,
+	},
+];
+
+const MODEL_MAP = new Map(MODELS.map((m) => [m.id, m]));
 
 // =============================================================================
 // Direct Access Token Cache
@@ -92,13 +146,9 @@ async function getDirectAccessToken(gitlabAccessToken: string): Promise<DirectAc
 		return cachedDirectAccess;
 	}
 
-	const url = `${GITLAB_COM_URL}/api/v4/ai/third_party_agents/direct_access`;
-	const response = await fetch(url, {
+	const response = await fetch(`${GITLAB_COM_URL}/api/v4/ai/third_party_agents/direct_access`, {
 		method: "POST",
-		headers: {
-			Authorization: `Bearer ${gitlabAccessToken}`,
-			"Content-Type": "application/json",
-		},
+		headers: { Authorization: `Bearer ${gitlabAccessToken}`, "Content-Type": "application/json" },
 		body: JSON.stringify({ feature_flags: { DuoAgentPlatformNext: true } }),
 	});
 
@@ -113,11 +163,7 @@ async function getDirectAccessToken(gitlabAccessToken: string): Promise<DirectAc
 	}
 
 	const data = (await response.json()) as { token: string; headers: Record<string, string> };
-	cachedDirectAccess = {
-		token: data.token,
-		headers: data.headers,
-		expiresAt: now + DIRECT_ACCESS_TTL,
-	};
+	cachedDirectAccess = { token: data.token, headers: data.headers, expiresAt: now + DIRECT_ACCESS_TTL };
 	return cachedDirectAccess;
 }
 
@@ -126,7 +172,7 @@ function invalidateDirectAccessToken() {
 }
 
 // =============================================================================
-// OAuth Implementation
+// OAuth
 // =============================================================================
 
 async function generatePKCE(): Promise<{ verifier: string; challenge: string }> {
@@ -136,21 +182,16 @@ async function generatePKCE(): Promise<{ verifier: string; challenge: string }>
 		.replace(/\+/g, "-")
 		.replace(/\//g, "_")
 		.replace(/=+$/, "");
-
-	const encoder = new TextEncoder();
-	const data = encoder.encode(verifier);
-	const hash = await crypto.subtle.digest("SHA-256", data);
+	const hash = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(verifier));
 	const challenge = btoa(String.fromCharCode(...new Uint8Array(hash)))
 		.replace(/\+/g, "-")
 		.replace(/\//g, "_")
 		.replace(/=+$/, "");
-
 	return { verifier, challenge };
 }
 
 async function loginGitLab(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
 	const { verifier, challenge } = await generatePKCE();
-
 	const authParams = new URLSearchParams({
 		client_id: BUNDLED_CLIENT_ID,
 		redirect_uri: REDIRECT_URI,
@@ -163,9 +204,7 @@ async function loginGitLab(callbacks: OAuthLoginCallbacks): Promise<OAuthCredent
 
 	callbacks.onAuth({ url: `${GITLAB_COM_URL}/oauth/authorize?${authParams.toString()}` });
 	const callbackUrl = await callbacks.onPrompt({ message: "Paste the callback URL:" });
-
-	const urlObj = new URL(callbackUrl);
-	const code = urlObj.searchParams.get("code");
+	const code = new URL(callbackUrl).searchParams.get("code");
 	if (!code) throw new Error("No authorization code found in callback URL");
 
 	const tokenResponse = await fetch(`${GITLAB_COM_URL}/oauth/token`, {
@@ -181,14 +220,12 @@ async function loginGitLab(callbacks: OAuthLoginCallbacks): Promise<OAuthCredent
 	});
 
 	if (!tokenResponse.ok) throw new Error(`Token exchange failed: ${await tokenResponse.text()}`);
-
 	const data = (await tokenResponse.json()) as {
 		access_token: string;
 		refresh_token: string;
 		expires_in: number;
 		created_at: number;
 	};
-
 	invalidateDirectAccessToken();
 	return {
 		refresh: data.refresh_token,
@@ -207,16 +244,13 @@ async function refreshGitLabToken(credentials: OAuthCredentials): Promise<OAuthC
 			refresh_token: credentials.refresh,
 		}).toString(),
 	});
-
 	if (!response.ok) throw new Error(`Token refresh failed: ${await response.text()}`);
-
 	const data = (await response.json()) as {
 		access_token: string;
 		refresh_token: string;
 		expires_in: number;
 		created_at: number;
 	};
-
 	invalidateDirectAccessToken();
 	return {
 		refresh: data.refresh_token,
@@ -226,10 +260,10 @@ async function refreshGitLabToken(credentials: OAuthCredentials): Promise<OAuthC
 }
 
 // =============================================================================
-// Main Stream Function - Delegates to pi-ai's built-in implementations
+// Stream Function
 // =============================================================================
 
-function streamGitLabDuo(
+export function streamGitLabDuo(
 	model: Model<Api>,
 	context: Context,
 	options?: SimpleStreamOptions,
@@ -239,57 +273,22 @@ function streamGitLabDuo(
 	(async () => {
 		try {
 			const gitlabAccessToken = options?.apiKey;
-			if (!gitlabAccessToken) {
-				throw new Error("No GitLab access token. Run /login gitlab-duo or set GITLAB_TOKEN");
-			}
+			if (!gitlabAccessToken) throw new Error("No GitLab access token. Run /login gitlab-duo or set GITLAB_TOKEN");
 
-			const mapping = MODEL_MAPPINGS[model.id];
-			if (!mapping) throw new Error(`Unknown model: ${model.id}`);
+			const cfg = MODEL_MAP.get(model.id);
+			if (!cfg) throw new Error(`Unknown model: ${model.id}`);
 
-			// Get direct access token (cached)
 			const directAccess = await getDirectAccessToken(gitlabAccessToken);
+			const modelWithBaseUrl = { ...model, baseUrl: cfg.baseUrl };
+			const headers = { ...directAccess.headers, Authorization: `Bearer ${directAccess.token}` };
+			const streamOptions = { ...options, apiKey: "gitlab-duo", headers };
 
-			// Create a proxy model that uses the backend API
-			const proxyModel: Model<typeof mapping.api> = {
-				...model,
-				id: mapping.backendModel,
-				api: mapping.api,
-				baseUrl: mapping.baseUrl,
-			};
-
-			// Merge GitLab headers with Authorization bearer token
-			const headers = {
-				...directAccess.headers,
-				Authorization: `Bearer ${directAccess.token}`,
-			};
-
-			// Delegate to pi-ai's built-in streaming
-			const innerStream = streamSimple(proxyModel, context, {
-				...options,
-				apiKey: "gitlab-duo", // Dummy value to pass validation
-				headers,
-			});
+			const innerStream =
+				cfg.backend === "anthropic"
+					? streamSimpleAnthropic(modelWithBaseUrl as Model<"anthropic-messages">, context, streamOptions)
+					: streamSimpleOpenAIResponses(modelWithBaseUrl as Model<"openai-responses">, context, streamOptions);
 
-			// Forward all events
-			for await (const event of innerStream) {
-				// Patch the model info back to gitlab-duo
-				if ("partial" in event && event.partial) {
-					event.partial.api = model.api;
-					event.partial.provider = model.provider;
-					event.partial.model = model.id;
-				}
-				if ("message" in event && event.message) {
-					event.message.api = model.api;
-					event.message.provider = model.provider;
-					event.message.model = model.id;
-				}
-				if ("error" in event && event.error) {
-					event.error.api = model.api;
-					event.error.provider = model.provider;
-					event.error.model = model.id;
-				}
-				stream.push(event);
-			}
+			for await (const event of innerStream) stream.push(event);
 			stream.end();
 		} catch (error) {
 			stream.push({
@@ -330,73 +329,21 @@ export default function (pi: ExtensionAPI) {
 		baseUrl: AI_GATEWAY_URL,
 		apiKey: "GITLAB_TOKEN",
 		api: "gitlab-duo-api",
-
-		models: [
-			// Anthropic models
-			{
-				id: "duo-chat-opus-4-5",
-				name: "GitLab Duo Claude Opus 4.5",
-				reasoning: false,
-				input: ["text"],
-				cost: { input: 15, output: 75, cacheRead: 1.5, cacheWrite: 18.75 },
-				contextWindow: 200000,
-				maxTokens: 32000,
-			},
-			{
-				id: "duo-chat-sonnet-4-5",
-				name: "GitLab Duo Claude Sonnet 4.5",
-				reasoning: false,
-				input: ["text"],
-				cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 },
-				contextWindow: 200000,
-				maxTokens: 16384,
-			},
-			{
-				id: "duo-chat-haiku-4-5",
-				name: "GitLab Duo Claude Haiku 4.5",
-				reasoning: false,
-				input: ["text"],
-				cost: { input: 1, output: 5, cacheRead: 0.1, cacheWrite: 1.25 },
-				contextWindow: 200000,
-				maxTokens: 8192,
-			},
-			// OpenAI models
-			{
-				id: "duo-chat-gpt-5-1",
-				name: "GitLab Duo GPT-5.1",
-				reasoning: false,
-				input: ["text"],
-				cost: { input: 2.5, output: 10, cacheRead: 0, cacheWrite: 0 },
-				contextWindow: 128000,
-				maxTokens: 16384,
-			},
-			{
-				id: "duo-chat-gpt-5-mini",
-				name: "GitLab Duo GPT-5 Mini",
-				reasoning: false,
-				input: ["text"],
-				cost: { input: 0.15, output: 0.6, cacheRead: 0, cacheWrite: 0 },
-				contextWindow: 128000,
-				maxTokens: 16384,
-			},
-			{
-				id: "duo-chat-gpt-5-codex",
-				name: "GitLab Duo GPT-5 Codex",
-				reasoning: false,
-				input: ["text"],
-				cost: { input: 2.5, output: 10, cacheRead: 0, cacheWrite: 0 },
-				contextWindow: 128000,
-				maxTokens: 16384,
-			},
-		],
-
+		models: MODELS.map(({ id, name, reasoning, input, cost, contextWindow, maxTokens }) => ({
+			id,
+			name,
+			reasoning,
+			input,
+			cost,
+			contextWindow,
+			maxTokens,
+		})),
 		oauth: {
 			name: "GitLab Duo",
 			login: loginGitLab,
 			refreshToken: refreshGitLabToken,
 			getApiKey: (cred) => cred.access,
 		},
-
 		streamSimple: streamGitLabDuo,
 	});
 }

package/package.json

@@ -1,13 +1,13 @@
 {
   "name": "pi-gitlab-duo",
-  "version": "0.1.0",
+  "version": "0.1.1",
   "description": "GitLab Duo provider extension for pi",
   "type": "module",
   "license": "MIT",
   "author": "Mario Zechner",
   "repository": {
     "type": "git",
-    "url": "https://github.com/badlogic/pi-gitlab-duo.git"
+    "url": "git+https://github.com/badlogic/pi-gitlab-duo.git"
   },
   "keywords": [
     "pi",