feat: improve user experience

mcp/claude_client.go

@@ -99,6 +99,10 @@ func (c *ClaudeClient) parseMCPResponse(body []byte) (string, error) {
 			Type string `json:"type"`
 			Text string `json:"text"`
 		} `json:"content"`
+		Usage struct {
+			InputTokens  int `json:"input_tokens"`
+			OutputTokens int `json:"output_tokens"`
+		} `json:"usage"`
 		Error *struct {
 			Type    string `json:"type"`
 			Message string `json:"message"`
@@ -117,6 +121,18 @@ func (c *ClaudeClient) parseMCPResponse(body []byte) (string, error) {
 		return "", fmt.Errorf("Claude returned empty content, body: %s", string(body))
 	}
 
+	// Report token usage if callback is set
+	totalTokens := response.Usage.InputTokens + response.Usage.OutputTokens
+	if TokenUsageCallback != nil && totalTokens > 0 {
+		TokenUsageCallback(TokenUsage{
+			Provider:         c.Provider,
+			Model:            c.Model,
+			PromptTokens:     response.Usage.InputTokens,
+			CompletionTokens: response.Usage.OutputTokens,
+			TotalTokens:      totalTokens,
+		})
+	}
+
 	// Find text content
 	for _, content := range response.Content {
 		if content.Type == "text" {

mcp/client.go

@@ -31,8 +31,20 @@ var (
 		"stream error",   // HTTP/2 stream error
 		"INTERNAL_ERROR", // Server internal error
 	}
+
+	// TokenUsageCallback is called after each AI request with token usage info
+	TokenUsageCallback func(usage TokenUsage)
 )
 
+// TokenUsage represents token usage from AI API response
+type TokenUsage struct {
+	Provider         string
+	Model            string
+	PromptTokens     int
+	CompletionTokens int
+	TotalTokens      int
+}
+
 // Client AI API configuration
 type Client struct {
 	Provider   string
@@ -226,6 +238,11 @@ func (client *Client) parseMCPResponse(body []byte) (string, error) {
 				Content string `json:"content"`
 			} `json:"message"`
 		} `json:"choices"`
+		Usage struct {
+			PromptTokens     int `json:"prompt_tokens"`
+			CompletionTokens int `json:"completion_tokens"`
+			TotalTokens      int `json:"total_tokens"`
+		} `json:"usage"`
 	}
 
 	if err := json.Unmarshal(body, &result); err != nil {
@@ -236,6 +253,17 @@ func (client *Client) parseMCPResponse(body []byte) (string, error) {
 		return "", fmt.Errorf("API returned empty response")
 	}
 
+	// Report token usage if callback is set
+	if TokenUsageCallback != nil && result.Usage.TotalTokens > 0 {
+		TokenUsageCallback(TokenUsage{
+			Provider:         client.Provider,
+			Model:            client.Model,
+			PromptTokens:     result.Usage.PromptTokens,
+			CompletionTokens: result.Usage.CompletionTokens,
+			TotalTokens:      result.Usage.TotalTokens,
+		})
+	}
+
 	return result.Choices[0].Message.Content, nil
 }