Add Maestro login button

src/api/index.ts

@@ -3,6 +3,7 @@ import { ApiConfiguration, ApiModelId, ModelInfo } from "../shared/api"
 import { AnthropicHandler } from "./anthropic"
 import { AwsBedrockHandler } from "./bedrock"
 import { OpenRouterHandler } from "./openrouter"
+import { MaestroHandler } from "./maestro"
 
 export interface ApiHandler {
 	createMessage(
@@ -32,6 +33,8 @@ export function buildApiHandler(configuration: ApiConfiguration): ApiHandler {
 			return new OpenRouterHandler(options)
 		case "bedrock":
 			return new AwsBedrockHandler(options)
+		case "maestro":
+			return new MaestroHandler(options)
 		default:
 			return new AnthropicHandler(options)
 	}

src/api/maestro.ts

@@ -0,0 +1,114 @@
+import { Anthropic } from "@anthropic-ai/sdk"
+import { ApiHandler, withoutImageData } from "."
+import { ApiHandlerOptions, maestroDefaultModelId, MaestroModelId, maestroModels, ModelInfo } from "../shared/api"
+
+export class MaestroHandler implements ApiHandler {
+	private options: ApiHandlerOptions
+	private client: Anthropic
+
+	constructor(options: ApiHandlerOptions) {
+		this.options = options
+		this.client = new Anthropic({ apiKey: this.options.apiKey })
+	}
+
+	async createMessage(
+		systemPrompt: string,
+		messages: Anthropic.Messages.MessageParam[],
+		tools: Anthropic.Messages.Tool[]
+	): Promise<Anthropic.Messages.Message> {
+		const modelId = this.getModel().id
+		switch (modelId) {
+			case "claude-3-5-sonnet-20240620":
+			case "claude-3-haiku-20240307":
+				const userMsgIndices = messages.reduce(
+					(acc, msg, index) => (msg.role === "user" ? [...acc, index] : acc),
+					[] as number[]
+				)
+				const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
+				const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
+				return await this.client.beta.promptCaching.messages.create(
+					{
+						model: modelId,
+						max_tokens: this.getModel().info.maxTokens,
+						system: [{ text: systemPrompt, type: "text", cache_control: { type: "ephemeral" } }],
+						messages: messages.map((message, index) => {
+							if (index === lastUserMsgIndex || index === secondLastMsgUserIndex) {
+								return {
+									...message,
+									content:
+										typeof message.content === "string"
+											? [
+													{
+														type: "text",
+														text: message.content,
+														cache_control: { type: "ephemeral" },
+													},
+											  ]
+											: message.content.map((content, contentIndex) =>
+													contentIndex === message.content.length - 1
+														? { ...content, cache_control: { type: "ephemeral" } }
+														: content
+											  ),
+								}
+							}
+							return message
+						}),
+						tools,
+						tool_choice: { type: "auto" },
+					},
+					(() => {
+						switch (modelId) {
+							case "claude-3-5-sonnet-20240620":
+								return {
+									headers: {
+										"anthropic-beta": "prompt-caching-2024-07-31",
+									},
+								}
+							case "claude-3-haiku-20240307":
+								return {
+									headers: { "anthropic-beta": "prompt-caching-2024-07-31" },
+								}
+							default:
+								return undefined
+						}
+					})()
+				)
+			default:
+				return await this.client.messages.create({
+					model: modelId,
+					max_tokens: this.getModel().info.maxTokens,
+					system: [{ text: systemPrompt, type: "text" }],
+					messages,
+					tools,
+					tool_choice: { type: "auto" },
+				})
+		}
+	}
+
+	createUserReadableRequest(
+		userContent: Array<
+			| Anthropic.TextBlockParam
+			| Anthropic.ImageBlockParam
+			| Anthropic.ToolUseBlockParam
+			| Anthropic.ToolResultBlockParam
+		>
+	): any {
+		return {
+			model: this.getModel().id,
+			max_tokens: this.getModel().info.maxTokens,
+			system: "(see SYSTEM_PROMPT in src/ClaudeDev.ts)",
+			messages: [{ conversation_history: "..." }, { role: "user", content: withoutImageData(userContent) }],
+			tools: "(see tools in src/ClaudeDev.ts)",
+			tool_choice: { type: "auto" },
+		}
+	}
+
+	getModel(): { id: MaestroModelId; info: ModelInfo } {
+		const modelId = this.options.apiModelId
+		if (modelId && modelId in maestroModels) {
+			const id = modelId as MaestroModelId
+			return { id, info: maestroModels[id] }
+		}
+		return { id: maestroDefaultModelId, info: maestroModels[maestroDefaultModelId] }
+	}
+}

src/extension.ts

@@ -108,6 +108,20 @@ export function activate(context: vscode.ExtensionContext) {
 	context.subscriptions.push(
 		vscode.workspace.registerTextDocumentContentProvider("claude-dev-diff", diffContentProvider)
 	)
+
+	// URI Handler
+	const handleUri = async (uri: vscode.Uri) => {
+		const query = new URLSearchParams(uri.query)
+		const token = query.get("token")
+		const fixedToken = token?.replaceAll("jwt?token=", "")
+		console.log(fixedToken)
+		console.log(uri)
+
+		if (fixedToken) {
+			await sidebarProvider.saveMaestroToken(fixedToken)
+		}
+	}
+	context.subscriptions.push(vscode.window.registerUriHandler({ handleUri }))
 }
 
 // This method is called when your extension is deactivated

src/maestro/auth.ts

@@ -0,0 +1,38 @@
+import axios from "axios"
+import * as vscode from "vscode"
+import { MaestroUser, MaestroUserSchema } from "../shared/maestro"
+
+const MAESTRO_BASE_URL = "https://maestro.im-ada.ai"
+
+export function didClickMaestroSignIn() {
+	const loginUrl = `${MAESTRO_BASE_URL}/auth/login?ext=1&redirectTo=${vscode.env.uriScheme}://saoudrizwan.claude-dev?token=jwt`
+	vscode.env.openExternal(vscode.Uri.parse(loginUrl))
+}
+
+export async function validateMaestroToken({
+	token,
+	showError = false,
+}: {
+	token: string
+	showError?: boolean
+}): Promise<MaestroUser> {
+	try {
+		const response = await axios.post(`${MAESTRO_BASE_URL}/api/extension/auth/callback`, { token })
+		const user = MaestroUserSchema.parse(response.data.user)
+		console.log("retrieved user", user)
+		return user
+	} catch (error) {
+		if (showError) {
+			if (axios.isAxiosError(error)) {
+				vscode.window.showErrorMessage(
+					"Failed to validate token:",
+					error.response?.status,
+					error.response?.data
+				)
+			} else {
+				vscode.window.showErrorMessage("An unexpected error occurred:", error)
+			}
+		}
+		throw error
+	}
+}

src/maestro/index.ts

@@ -0,0 +1 @@
+export * from "./auth"

src/providers/ClaudeDevProvider.ts

@@ -8,6 +8,8 @@ import { downloadTask, getNonce, getUri, selectImages } from "../utils"
 import * as path from "path"
 import fs from "fs/promises"
 import { HistoryItem } from "../shared/HistoryItem"
+import { didClickMaestroSignIn, validateMaestroToken } from "../maestro"
+import { MaestroUser } from "../shared/maestro"
 
 /*
 https://github.com/microsoft/vscode-webview-ui-toolkit-samples/blob/main/default/weather-webview/src/providers/WeatherViewProvider.ts
@@ -15,7 +17,7 @@ https://github.com/microsoft/vscode-webview-ui-toolkit-samples/blob/main/default
 https://github.com/KumarVariable/vscode-extension-sidebar-html/blob/master/src/customSidebarViewProvider.ts
 */
 
-type SecretKey = "apiKey" | "openRouterApiKey" | "awsAccessKey" | "awsSecretKey"
+type SecretKey = "apiKey" | "openRouterApiKey" | "awsAccessKey" | "awsSecretKey" | "maestroToken"
 type GlobalStateKey =
 	| "apiProvider"
 	| "apiModelId"
@@ -32,9 +34,11 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 	private view?: vscode.WebviewView | vscode.WebviewPanel
 	private claudeDev?: ClaudeDev
 	private latestAnnouncementId = "aug-17-2024" // update to some unique identifier when we add a new announcement
+	private maestroUser?: MaestroUser
 
 	constructor(readonly context: vscode.ExtensionContext, private readonly outputChannel: vscode.OutputChannel) {
 		this.outputChannel.appendLine("ClaudeDevProvider instantiated")
+		this.fetchMaestroUser({})
 	}
 
 	/*
@@ -340,6 +344,12 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 					case "exportTaskWithId":
 						this.exportTaskWithId(message.text!)
 						break
+					case "didClickMaestroSignIn":
+						didClickMaestroSignIn()
+						break
+					case "didClickMaestroSignOut":
+						await this.signOutMaestro()
+						break
 					// Add more switch case statements here as more webview message commands
 					// are created within the webview context (i.e. inside media/main.js)
 				}
@@ -349,6 +359,36 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 		)
 	}
 
+	// Maestro
+
+	async saveMaestroToken(token: string) {
+		await this.storeSecret("maestroToken", token)
+		await this.updateGlobalState("apiProvider", "maestro")
+		await this.fetchMaestroUser({ showError: true })
+		this.claudeDev?.updateApi({ apiProvider: "maestro", maestroToken: token })
+	}
+
+	async fetchMaestroUser({ showError = false }: { showError?: boolean }): Promise<MaestroUser | undefined> {
+		if (this.maestroUser) {
+			return this.maestroUser
+		}
+		const token = await this.getSecret("maestroToken")
+		if (!token) {
+			return undefined
+		}
+		const user = await validateMaestroToken({ token, showError })
+		this.maestroUser = user
+		await this.postStateToWebview()
+		return user
+	}
+
+	async signOutMaestro() {
+		await this.storeSecret("maestroToken", undefined)
+		this.claudeDev?.updateApi({ apiProvider: "maestro", maestroToken: undefined })
+		this.maestroUser = undefined
+		await this.postStateToWebview()
+	}
+
 	// Task history
 
 	async getTaskWithId(id: string): Promise<{
@@ -449,6 +489,7 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 				claudeMessages: this.claudeDev?.claudeMessages || [],
 				taskHistory: (taskHistory || []).filter((item) => item.ts && item.task).sort((a, b) => b.ts - a.ts),
 				shouldShowAnnouncement: lastShownAnnouncementId !== this.latestAnnouncementId,
+				maestroUser: this.maestroUser,
 			},
 		})
 	}

src/shared/ExtensionMessage.ts

@@ -2,6 +2,7 @@
 
 import { ApiConfiguration } from "./api"
 import { HistoryItem } from "./HistoryItem"
+import { MaestroUser } from "./maestro"
 
 // webview will hold state
 export interface ExtensionMessage {
@@ -21,6 +22,7 @@ export interface ExtensionState {
 	claudeMessages: ClaudeMessage[]
 	taskHistory: HistoryItem[]
 	shouldShowAnnouncement: boolean
+	maestroUser?: MaestroUser
 }
 
 export interface ClaudeMessage {

src/shared/WebviewMessage.ts

@@ -15,6 +15,8 @@ export interface WebviewMessage {
 		| "showTaskWithId"
 		| "deleteTaskWithId"
 		| "exportTaskWithId"
+		| "didClickMaestroSignIn"
+		| "didClickMaestroSignOut"
 	text?: string
 	askResponse?: ClaudeAskResponse
 	apiConfiguration?: ApiConfiguration

src/shared/api.ts

@@ -1,4 +1,4 @@
-export type ApiProvider = "anthropic" | "openrouter" | "bedrock"
+export type ApiProvider = "anthropic" | "openrouter" | "bedrock" | "maestro"
 
 export interface ApiHandlerOptions {
 	apiModelId?: ApiModelId
@@ -7,6 +7,7 @@ export interface ApiHandlerOptions {
 	awsAccessKey?: string
 	awsSecretKey?: string
 	awsRegion?: string
+	maestroToken?: string
 }
 
 export type ApiConfiguration = ApiHandlerOptions & {
@@ -232,3 +233,10 @@ export const openRouterModels = {
 	// 	outputPrice: 1.5,
 	// },
 } as const satisfies Record<string, ModelInfo>
+
+// Maestro
+export type MaestroModelId = keyof typeof maestroModels
+export const maestroDefaultModelId: MaestroModelId = "claude-3-5-sonnet-20240620"
+export const maestroModels = {
+	...anthropicModels,
+} as const satisfies Record<string, ModelInfo>

src/shared/maestro.ts

@@ -0,0 +1,10 @@
+import { z } from "zod"
+
+export const MaestroUserSchema = z.object({
+	id: z.string(),
+	image: z.string().nullable(),
+	email: z.string().email(),
+	name: z.string().nullable(),
+	emailVerified: z.coerce.date().nullable(),
+})
+export type MaestroUser = z.infer<typeof MaestroUserSchema>