Use 'user credits' header to update balance and show user under task header

src/ClaudeDev.ts

@@ -1251,7 +1251,12 @@ ${this.customInstructions.trim()}
 				this.apiConversationHistory,
 				tools
 			)
-			return await this.api.createMessage(systemPrompt, adjustedMessages, tools)
+			const { message, userCredits } = await this.api.createMessage(systemPrompt, adjustedMessages, tools)
+			if (userCredits !== undefined) {
+				console.log("Updating kodu credits", userCredits)
+				this.providerRef.deref()?.updateKoduCredits(userCredits)
+			}
+			return message
 		} catch (error) {
 			const { response } = await this.ask(
 				"api_req_failed",

src/api/anthropic.ts

@@ -1,5 +1,5 @@
 import { Anthropic } from "@anthropic-ai/sdk"
-import { ApiHandler, withoutImageData } from "."
+import { ApiHandler, ApiHandlerMessageResponse, withoutImageData } from "."
 import { anthropicDefaultModelId, AnthropicModelId, anthropicModels, ApiHandlerOptions, ModelInfo } from "../shared/api"
 
 export class AnthropicHandler implements ApiHandler {
@@ -15,12 +15,12 @@ export class AnthropicHandler implements ApiHandler {
 		systemPrompt: string,
 		messages: Anthropic.Messages.MessageParam[],
 		tools: Anthropic.Messages.Tool[]
-	): Promise<Anthropic.Messages.Message> {
+	): Promise<ApiHandlerMessageResponse> {
 		const modelId = this.getModel().id
 		switch (modelId) {
 			case "claude-3-5-sonnet-20240620":
 			case "claude-3-opus-20240229":
-			case "claude-3-haiku-20240307":
+			case "claude-3-haiku-20240307": {
 				/*
 				The latest message will be the new user message, one before will be the assistant message from a previous request, and the user message before that will be a previously cached user message. So we need to mark the latest user message as ephemeral to cache it for the next request, and mark the second to last user message as ephemeral to let the server know the last message to retrieve from the cache for the current request..
 				*/
@@ -30,7 +30,7 @@ export class AnthropicHandler implements ApiHandler {
 				)
 				const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
 				const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
-				return await this.client.beta.promptCaching.messages.create(
+				const message = await this.client.beta.promptCaching.messages.create(
 					{
 						model: modelId,
 						max_tokens: this.getModel().info.maxTokens,
@@ -80,8 +80,10 @@ export class AnthropicHandler implements ApiHandler {
 						}
 					})()
 				)
-			default:
+				return { message }
-				return await this.client.messages.create({
+			}
+			default: {
+				const message = await this.client.messages.create({
 					model: modelId,
 					max_tokens: this.getModel().info.maxTokens,
 					system: [{ text: systemPrompt, type: "text" }],
@@ -89,6 +91,8 @@ export class AnthropicHandler implements ApiHandler {
 					tools,
 					tool_choice: { type: "auto" },
 				})
+				return { message }
+			}
 		}
 	}
 

src/api/bedrock.ts

@@ -1,6 +1,6 @@
 import AnthropicBedrock from "@anthropic-ai/bedrock-sdk"
 import { Anthropic } from "@anthropic-ai/sdk"
-import { ApiHandler, withoutImageData } from "."
+import { ApiHandler, ApiHandlerMessageResponse, withoutImageData } from "."
 import { ApiHandlerOptions, bedrockDefaultModelId, BedrockModelId, bedrockModels, ModelInfo } from "../shared/api"
 
 // https://docs.anthropic.com/en/api/claude-on-amazon-bedrock
@@ -26,8 +26,8 @@ export class AwsBedrockHandler implements ApiHandler {
 		systemPrompt: string,
 		messages: Anthropic.Messages.MessageParam[],
 		tools: Anthropic.Messages.Tool[]
-	): Promise<Anthropic.Messages.Message> {
+	): Promise<ApiHandlerMessageResponse> {
-		return await this.client.messages.create({
+		const message = await this.client.messages.create({
 			model: this.getModel().id,
 			max_tokens: this.getModel().info.maxTokens,
 			system: systemPrompt,
@@ -35,6 +35,7 @@ export class AwsBedrockHandler implements ApiHandler {
 			tools,
 			tool_choice: { type: "auto" },
 		})
+		return { message }
 	}
 
 	createUserReadableRequest(

src/api/index.ts

@@ -5,12 +5,17 @@ import { AwsBedrockHandler } from "./bedrock"
 import { OpenRouterHandler } from "./openrouter"
 import { KoduHandler } from "./kodu"
 
+export interface ApiHandlerMessageResponse {
+	message: Anthropic.Messages.Message
+	userCredits?: number
+}
+
 export interface ApiHandler {
 	createMessage(
 		systemPrompt: string,
 		messages: Anthropic.Messages.MessageParam[],
 		tools: Anthropic.Messages.Tool[]
-	): Promise<Anthropic.Messages.Message>
+	): Promise<ApiHandlerMessageResponse>
 
 	createUserReadableRequest(
 		userContent: Array<

src/api/kodu.ts

@@ -1,6 +1,6 @@
 import { Anthropic } from "@anthropic-ai/sdk"
 import axios from "axios"
-import { ApiHandler, withoutImageData } from "."
+import { ApiHandler, ApiHandlerMessageResponse, withoutImageData } from "."
 import { ApiHandlerOptions, koduDefaultModelId, KoduModelId, koduModels, ModelInfo } from "../shared/api"
 import { getKoduCreditsUrl, getKoduInferenceUrl } from "../shared/kodu"
 
@@ -24,7 +24,7 @@ export class KoduHandler implements ApiHandler {
 		systemPrompt: string,
 		messages: Anthropic.Messages.MessageParam[],
 		tools: Anthropic.Messages.Tool[]
-	): Promise<Anthropic.Messages.Message> {
+	): Promise<ApiHandlerMessageResponse> {
 		const modelId = this.getModel().id
 		let requestBody: Anthropic.Beta.PromptCaching.Messages.MessageCreateParamsNonStreaming
 		switch (modelId) {
@@ -82,7 +82,9 @@ export class KoduHandler implements ApiHandler {
 				"x-api-key": this.options.koduApiKey,
 			},
 		})
-		return response.data
+		const message = response.data
+		const userCredits = response.headers["user-credits"]
+		return { message, userCredits: userCredits !== undefined ? parseFloat(userCredits) : undefined }
 	}
 
 	createUserReadableRequest(

src/api/openrouter.ts

@@ -1,6 +1,6 @@
 import { Anthropic } from "@anthropic-ai/sdk"
 import OpenAI from "openai"
-import { ApiHandler, withoutImageData } from "."
+import { ApiHandler, ApiHandlerMessageResponse, withoutImageData } from "."
 import {
 	ApiHandlerOptions,
 	ModelInfo,
@@ -30,7 +30,7 @@ export class OpenRouterHandler implements ApiHandler {
 		systemPrompt: string,
 		messages: Anthropic.Messages.MessageParam[],
 		tools: Anthropic.Messages.Tool[]
-	): Promise<Anthropic.Messages.Message> {
+	): Promise<ApiHandlerMessageResponse> {
 		// Convert Anthropic messages to OpenAI format
 		const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [
 			{ role: "system", content: systemPrompt },
@@ -120,7 +120,7 @@ export class OpenRouterHandler implements ApiHandler {
 			)
 		}
 
-		return anthropicMessage
+		return { message: anthropicMessage }
 	}
 
 	/*

src/providers/ClaudeDevProvider.ts

@@ -696,6 +696,10 @@ export class ClaudeDevProvider implements vscode.WebviewViewProvider {
 		return history
 	}
 
+	async updateKoduCredits(credits: number) {
+		await this.updateGlobalState("koduCredits", credits)
+	}
+
 	// global
 
 	private async updateGlobalState(key: GlobalStateKey, value: any) {

src/shared/kodu.ts

@@ -1,4 +1,4 @@
-const KODU_BASE_URL = "https://claude-dev.com"
+const KODU_BASE_URL = "https://kodu.ai"
 
 export function getKoduSignInUrl(uriScheme?: string) {
 	return `${KODU_BASE_URL}/auth/login?redirectTo=${uriScheme}://saoudrizwan.claude-dev&ext=1`

webview-ui/src/App.tsx

@@ -169,6 +169,7 @@ const App: React.FC = () => {
 						apiConfiguration={apiConfiguration}
 						vscodeUriScheme={vscodeUriScheme}
 						shouldShowKoduPromo={shouldShowKoduPromo}
+						koduCredits={koduCredits}
 					/>
 				</>
 			)}

webview-ui/src/components/ApiOptions.tsx

@@ -36,7 +36,7 @@ const ApiOptions: React.FC<ApiOptionsProps> = ({
 	apiErrorMessage,
 	vscodeUriScheme,
 }) => {
-	const [didFetchKoduCredits, setDidFetchKoduCredits] = useState(false)
+	const [, setDidFetchKoduCredits] = useState(false)
 	const handleInputChange = (field: keyof ApiConfiguration) => (event: any) => {
 		setApiConfiguration((prev) => ({ ...prev, [field]: event.target.value }))
 	}
@@ -78,11 +78,11 @@ const ApiOptions: React.FC<ApiOptionsProps> = ({
 	}
 
 	useEffect(() => {
-		if (selectedProvider === "kodu" && apiConfiguration?.koduApiKey) {
+		if (selectedProvider === "kodu" && apiConfiguration?.koduApiKey && koduCredits === undefined) {
 			setDidFetchKoduCredits(false)
 			vscode.postMessage({ type: "fetchKoduCredits" })
 		}
-	}, [selectedProvider, apiConfiguration?.koduApiKey])
+	}, [selectedProvider, apiConfiguration?.koduApiKey, koduCredits])
 
 	const handleMessage = useCallback((e: MessageEvent) => {
 		const message: ExtensionMessage = e.data
@@ -178,7 +178,7 @@ const ApiOptions: React.FC<ApiOptionsProps> = ({
 							</div>
 							<div style={{ marginBottom: 7 }}>
 								Credits remaining:{" "}
-								<span style={{ fontWeight: 500, opacity: didFetchKoduCredits ? 1 : 0.6 }}>
+								<span style={{ fontWeight: 500, opacity: koduCredits !== undefined ? 1 : 0.6 }}>
 									{formatPrice(koduCredits || 0)}
 								</span>
 							</div>
@@ -307,7 +307,7 @@ const ApiOptions: React.FC<ApiOptionsProps> = ({
 	)
 }
 
-const formatPrice = (price: number) => {
+export const formatPrice = (price: number) => {
 	return new Intl.NumberFormat("en-US", {
 		style: "currency",
 		currency: "USD",

webview-ui/src/components/ChatView.tsx

@@ -33,6 +33,7 @@ interface ChatViewProps {
 	apiConfiguration?: ApiConfiguration
 	vscodeUriScheme?: string
 	shouldShowKoduPromo: boolean
+	koduCredits?: number
 }
 
 const MAX_IMAGES_PER_MESSAGE = 20 // Anthropic limits to 20 images
@@ -51,6 +52,7 @@ const ChatView = ({
 	apiConfiguration,
 	vscodeUriScheme,
 	shouldShowKoduPromo,
+	koduCredits,
 }: ChatViewProps) => {
 	//const task = messages.length > 0 ? (messages[0].say === "task" ? messages[0] : undefined) : undefined) : undefined
 	const task = messages.length > 0 ? messages[0] : undefined // leaving this less safe version here since if the first message is not a task, then the extension is in a bad state and needs to be debugged (see ClaudeDev.abort)
@@ -487,6 +489,9 @@ const ChatView = ({
 					totalCost={apiMetrics.totalCost}
 					onClose={handleTaskCloseButtonClick}
 					isHidden={isHidden}
+					koduCredits={koduCredits}
+					vscodeUriScheme={vscodeUriScheme}
+					apiProvider={apiConfiguration?.apiProvider}
 				/>
 			) : (
 				<>

webview-ui/src/components/TaskHeader.tsx

@@ -1,9 +1,12 @@
-import { VSCodeButton } from "@vscode/webview-ui-toolkit/react"
+import { VSCodeButton, VSCodeLink } from "@vscode/webview-ui-toolkit/react"
 import React, { useEffect, useRef, useState } from "react"
 import { useWindowSize } from "react-use"
 import { ClaudeMessage } from "../../../src/shared/ExtensionMessage"
 import { vscode } from "../utils/vscode"
 import Thumbnails from "./Thumbnails"
+import { formatPrice } from "./ApiOptions"
+import { getKoduAddCreditsUrl } from "../../../src/shared/kodu"
+import { ApiProvider } from "../../../src/shared/api"
 
 interface TaskHeaderProps {
 	task: ClaudeMessage
@@ -15,6 +18,9 @@ interface TaskHeaderProps {
 	totalCost: number
 	onClose: () => void
 	isHidden: boolean
+	koduCredits?: number
+	vscodeUriScheme?: string
+	apiProvider?: ApiProvider
 }
 
 const TaskHeader: React.FC<TaskHeaderProps> = ({
@@ -27,6 +33,9 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 	totalCost,
 	onClose,
 	isHidden,
+	koduCredits,
+	vscodeUriScheme,
+	apiProvider,
 }) => {
 	const [isExpanded, setIsExpanded] = useState(false)
 	const [showSeeMore, setShowSeeMore] = useState(false)
@@ -107,6 +116,7 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 					flexDirection: "column",
 					gap: "8px",
 					position: "relative",
+					zIndex: 1,
 				}}>
 				<div
 					style={{
@@ -248,6 +258,34 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 					</div>
 				</div>
 			</div>
+			{apiProvider === "kodu" && (
+				<div
+					style={{
+						backgroundColor: "color-mix(in srgb, var(--vscode-badge-background) 50%, transparent)",
+						color: "var(--vscode-badge-foreground)",
+						borderRadius: "0 0 3px 3px",
+						display: "flex",
+						justifyContent: "space-between",
+						alignItems: "center",
+						padding: "4px 12px 6px 12px",
+						fontSize: "0.9em",
+						marginLeft: "10px",
+						marginRight: "10px",
+					}}>
+					<div style={{ fontWeight: "500" }}>Credits Remaining:</div>
+					<div>
+						{formatPrice(koduCredits || 0)}
+						{(koduCredits || 0) < 1 && (
+							<>
+								{" "}
+								<VSCodeLink style={{ fontSize: "0.9em" }} href={getKoduAddCreditsUrl(vscodeUriScheme)}>
+									(get more?)
+								</VSCodeLink>
+							</>
+						)}
+					</div>
+				</div>
+			)}
 		</div>
 	)
 }