feat(copilot): add GitHub Copilot quota management endpoints and response enrichment

2026-03-21 16:40:22 +00:00 · 2026-02-03 13:02:51 +08:00
parent b9cdc2f54c
commit 92791194e5
1 changed files with 354 additions and 2 deletions
--- a/internal/api/handlers/management/api_tools.go
+++ b/internal/api/handlers/management/api_tools.go
@@ -13,12 +13,13 @@ import (

 	"github.com/fxamacker/cbor/v2"
 	"github.com/gin-gonic/gin"
-	"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/geminicli"
-	coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/net/proxy"
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/google"
+
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/geminicli"
+	coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
 )

 const defaultAPICallTimeout = 60 * time.Second
@@ -55,6 +56,7 @@ type apiCallResponse struct {
 	StatusCode int                 `json:"status_code"`
 	Header     map[string][]string `json:"header"`
 	Body       string              `json:"body"`
+	Quota      *QuotaSnapshots     `json:"quota,omitempty"`
 }

 // APICall makes a generic HTTP request on behalf of the management API caller.
@@ -97,6 +99,8 @@ type apiCallResponse struct {
 //	- status_code: Upstream HTTP status code.
 //	- header: Upstream response headers.
 //	- body: Upstream response body as string.
+//	- quota (optional): For GitHub Copilot enterprise accounts, contains quota_snapshots
+//	  with details for chat, completions, and premium_interactions.
 //
 // Example:
 //
@@ -236,6 +240,13 @@ func (h *Handler) APICall(c *gin.Context) {
 		Body:       string(respBody),
 	}

+	// If this is a GitHub Copilot token endpoint response, try to enrich with quota information
+	if resp.StatusCode == http.StatusOK &&
+		strings.Contains(urlStr, "copilot_internal") &&
+		strings.Contains(urlStr, "/token") {
+		response = h.enrichCopilotTokenResponse(c.Request.Context(), response, auth, urlStr)
+	}
+
 	// Return response in the same format as the request
 	if isCBOR {
 		cborData, errMarshal := cbor.Marshal(response)
@@ -735,3 +746,344 @@ func buildProxyTransport(proxyStr string) *http.Transport {
 	log.Debugf("unsupported proxy scheme: %s", proxyURL.Scheme)
 	return nil
 }
+
+// QuotaDetail represents quota information for a specific resource type
+type QuotaDetail struct {
+	Entitlement      float64 `json:"entitlement"`
+	OverageCount     float64 `json:"overage_count"`
+	OveragePermitted bool    `json:"overage_permitted"`
+	PercentRemaining float64 `json:"percent_remaining"`
+	QuotaID          string  `json:"quota_id"`
+	QuotaRemaining   float64 `json:"quota_remaining"`
+	Remaining        float64 `json:"remaining"`
+	Unlimited        bool    `json:"unlimited"`
+}
+
+// QuotaSnapshots contains quota details for different resource types
+type QuotaSnapshots struct {
+	Chat                QuotaDetail `json:"chat"`
+	Completions         QuotaDetail `json:"completions"`
+	PremiumInteractions QuotaDetail `json:"premium_interactions"`
+}
+
+// CopilotUsageResponse represents the GitHub Copilot usage information
+type CopilotUsageResponse struct {
+	AccessTypeSKU         string         `json:"access_type_sku"`
+	AnalyticsTrackingID   string         `json:"analytics_tracking_id"`
+	AssignedDate          string         `json:"assigned_date"`
+	CanSignupForLimited   bool           `json:"can_signup_for_limited"`
+	ChatEnabled           bool           `json:"chat_enabled"`
+	CopilotPlan           string         `json:"copilot_plan"`
+	OrganizationLoginList []interface{}  `json:"organization_login_list"`
+	OrganizationList      []interface{}  `json:"organization_list"`
+	QuotaResetDate        string         `json:"quota_reset_date"`
+	QuotaSnapshots        QuotaSnapshots `json:"quota_snapshots"`
+}
+
+type copilotQuotaRequest struct {
+	AuthIndexSnake  *string `json:"auth_index"`
+	AuthIndexCamel  *string `json:"authIndex"`
+	AuthIndexPascal *string `json:"AuthIndex"`
+}
+
+// GetCopilotQuota fetches GitHub Copilot quota information from the /copilot_internal/user endpoint.
+//
+// Endpoint:
+//
+//	GET /v0/management/copilot-quota
+//
+// Query Parameters (optional):
+//   - auth_index: The credential "auth_index" from GET /v0/management/auth-files.
+//     If omitted, uses the first available GitHub Copilot credential.
+//
+// Response:
+//
+//	Returns the CopilotUsageResponse with quota_snapshots containing detailed quota information
+//	for chat, completions, and premium_interactions.
+//
+// Example:
+//
+//	curl -sS -X GET "http://127.0.0.1:8317/v0/management/copilot-quota?auth_index=<AUTH_INDEX>" \
+//	  -H "Authorization: Bearer <MANAGEMENT_KEY>"
+func (h *Handler) GetCopilotQuota(c *gin.Context) {
+	authIndex := strings.TrimSpace(c.Query("auth_index"))
+	if authIndex == "" {
+		authIndex = strings.TrimSpace(c.Query("authIndex"))
+	}
+	if authIndex == "" {
+		authIndex = strings.TrimSpace(c.Query("AuthIndex"))
+	}
+
+	auth := h.findCopilotAuth(authIndex)
+	if auth == nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "no github copilot credential found"})
+		return
+	}
+
+	token, tokenErr := h.resolveTokenForAuth(c.Request.Context(), auth)
+	if tokenErr != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "failed to refresh copilot token"})
+		return
+	}
+	if token == "" {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "copilot token not found"})
+		return
+	}
+
+	apiURL := "https://api.github.com/copilot_internal/user"
+	req, errNewRequest := http.NewRequestWithContext(c.Request.Context(), http.MethodGet, apiURL, nil)
+	if errNewRequest != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to build request"})
+		return
+	}
+
+	req.Header.Set("Authorization", "Bearer "+token)
+	req.Header.Set("User-Agent", "CLIProxyAPIPlus")
+	req.Header.Set("Accept", "application/json")
+
+	httpClient := &http.Client{
+		Timeout:   defaultAPICallTimeout,
+		Transport: h.apiCallTransport(auth),
+	}
+
+	resp, errDo := httpClient.Do(req)
+	if errDo != nil {
+		log.WithError(errDo).Debug("copilot quota request failed")
+		c.JSON(http.StatusBadGateway, gin.H{"error": "request failed"})
+		return
+	}
+	defer func() {
+		if errClose := resp.Body.Close(); errClose != nil {
+			log.Errorf("response body close error: %v", errClose)
+		}
+	}()
+
+	respBody, errReadAll := io.ReadAll(resp.Body)
+	if errReadAll != nil {
+		c.JSON(http.StatusBadGateway, gin.H{"error": "failed to read response"})
+		return
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		c.JSON(http.StatusBadGateway, gin.H{
+			"error":       "github api request failed",
+			"status_code": resp.StatusCode,
+			"body":        string(respBody),
+		})
+		return
+	}
+
+	var usage CopilotUsageResponse
+	if errUnmarshal := json.Unmarshal(respBody, &usage); errUnmarshal != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to parse response"})
+		return
+	}
+
+	c.JSON(http.StatusOK, usage)
+}
+
+// findCopilotAuth locates a GitHub Copilot credential by auth_index or returns the first available one
+func (h *Handler) findCopilotAuth(authIndex string) *coreauth.Auth {
+	if h == nil || h.authManager == nil {
+		return nil
+	}
+
+	auths := h.authManager.List()
+	var firstCopilot *coreauth.Auth
+
+	for _, auth := range auths {
+		if auth == nil {
+			continue
+		}
+
+		provider := strings.ToLower(strings.TrimSpace(auth.Provider))
+		if provider != "copilot" && provider != "github" && provider != "github-copilot" {
+			continue
+		}
+
+		if firstCopilot == nil {
+			firstCopilot = auth
+		}
+
+		if authIndex != "" {
+			auth.EnsureIndex()
+			if auth.Index == authIndex {
+				return auth
+			}
+		}
+	}
+
+	return firstCopilot
+}
+
+// enrichCopilotTokenResponse fetches quota information and adds it to the Copilot token response body
+func (h *Handler) enrichCopilotTokenResponse(ctx context.Context, response apiCallResponse, auth *coreauth.Auth, originalURL string) apiCallResponse {
+	if auth == nil || response.Body == "" {
+		return response
+	}
+
+	// Parse the token response to check if it's enterprise (null limited_user_quotas)
+	var tokenResp map[string]interface{}
+	if err := json.Unmarshal([]byte(response.Body), &tokenResp); err != nil {
+		log.WithError(err).Debug("enrichCopilotTokenResponse: failed to parse copilot token response")
+		return response
+	}
+
+	// Check if this is an enterprise account (limited_user_quotas is null)
+	// limitedQuotas, hasLimitedQuotas := tokenResp["limited_user_quotas"]
+	// isEnterprise := !hasLimitedQuotas || limitedQuotas == nil
+
+	// // Only fetch additional quota info for enterprise accounts
+	// if !isEnterprise {
+	// 	return response
+	// }
+
+	// Get the GitHub token to call the copilot_internal/user endpoint
+	token, tokenErr := h.resolveTokenForAuth(ctx, auth)
+	if tokenErr != nil {
+		log.WithError(tokenErr).Debug("enrichCopilotTokenResponse: failed to resolve token")
+		return response
+	}
+	if token == "" {
+		return response
+	}
+
+	// Fetch quota information from /copilot_internal/user
+	// Derive the base URL from the original token request to support proxies and test servers
+	parsedURL, errParse := url.Parse(originalURL)
+	if errParse != nil {
+		log.WithError(errParse).Debug("enrichCopilotTokenResponse: failed to parse URL")
+		return response
+	}
+	quotaURL := fmt.Sprintf("%s://%s/copilot_internal/user", parsedURL.Scheme, parsedURL.Host)
+
+	req, errNewRequest := http.NewRequestWithContext(ctx, http.MethodGet, quotaURL, nil)
+	if errNewRequest != nil {
+		log.WithError(errNewRequest).Debug("enrichCopilotTokenResponse: failed to build request")
+		return response
+	}
+
+	req.Header.Set("Authorization", "Bearer "+token)
+	req.Header.Set("User-Agent", "CLIProxyAPIPlus")
+	req.Header.Set("Accept", "application/json")
+
+	httpClient := &http.Client{
+		Timeout:   defaultAPICallTimeout,
+		Transport: h.apiCallTransport(auth),
+	}
+
+	quotaResp, errDo := httpClient.Do(req)
+	if errDo != nil {
+		log.WithError(errDo).Debug("enrichCopilotTokenResponse: quota fetch HTTP request failed")
+		return response
+	}
+
+	defer func() {
+		if errClose := quotaResp.Body.Close(); errClose != nil {
+			log.Errorf("quota response body close error: %v", errClose)
+		}
+	}()
+
+	if quotaResp.StatusCode != http.StatusOK {
+		return response
+	}
+
+	quotaBody, errReadAll := io.ReadAll(quotaResp.Body)
+	if errReadAll != nil {
+		log.WithError(errReadAll).Debug("enrichCopilotTokenResponse: failed to read response")
+		return response
+	}
+
+	log.Debugf("enrichCopilotTokenResponse: %s", string(quotaBody))
+
+	// Parse the quota response
+	var quotaData CopilotUsageResponse
+	if err := json.Unmarshal(quotaBody, &quotaData); err != nil {
+		log.WithError(err).Debug("enrichCopilotTokenResponse: failed to parse response")
+		return response
+	}
+
+	// Check if this is an enterprise account by looking for quota_snapshots in the response
+	// Enterprise accounts have quota_snapshots, non-enterprise have limited_user_quotas
+	var quotaRaw map[string]interface{}
+	if err := json.Unmarshal(quotaBody, &quotaRaw); err == nil {
+		if _, hasQuotaSnapshots := quotaRaw["quota_snapshots"]; hasQuotaSnapshots {
+			// Enterprise account - has quota_snapshots
+			tokenResp["quota_snapshots"] = quotaData.QuotaSnapshots
+			tokenResp["quota_reset_date"] = quotaData.QuotaResetDate
+			tokenResp["access_type_sku"] = quotaData.AccessTypeSKU
+			tokenResp["copilot_plan"] = quotaData.CopilotPlan
+		} else {
+			// Non-enterprise account - build quota from limited_user_quotas and monthly_quotas
+			var quotaSnapshots QuotaSnapshots
+			
+			// Get monthly quotas (total entitlement) and limited_user_quotas (remaining)
+			monthlyQuotas, hasMonthly := quotaRaw["monthly_quotas"].(map[string]interface{})
+			limitedQuotas, hasLimited := quotaRaw["limited_user_quotas"].(map[string]interface{})
+			
+			// Process chat quota
+			if hasMonthly && hasLimited {
+				if chatTotal, ok := monthlyQuotas["chat"].(float64); ok {
+					chatRemaining := chatTotal // default to full if no limited quota
+					if chatLimited, ok := limitedQuotas["chat"].(float64); ok {
+						chatRemaining = chatLimited
+					}
+					percentRemaining := 0.0
+					if chatTotal > 0 {
+						percentRemaining = (chatRemaining / chatTotal) * 100.0
+					}
+					quotaSnapshots.Chat = QuotaDetail{
+						Entitlement:      chatTotal,
+						Remaining:        chatRemaining,
+						QuotaRemaining:   chatRemaining,
+						PercentRemaining: percentRemaining,
+						QuotaID:          "chat",
+						Unlimited:        false,
+					}
+				}
+				
+				// Process completions quota
+				if completionsTotal, ok := monthlyQuotas["completions"].(float64); ok {
+					completionsRemaining := completionsTotal // default to full if no limited quota
+					if completionsLimited, ok := limitedQuotas["completions"].(float64); ok {
+						completionsRemaining = completionsLimited
+					}
+					percentRemaining := 0.0
+					if completionsTotal > 0 {
+						percentRemaining = (completionsRemaining / completionsTotal) * 100.0
+					}
+					quotaSnapshots.Completions = QuotaDetail{
+						Entitlement:      completionsTotal,
+						Remaining:        completionsRemaining,
+						QuotaRemaining:   completionsRemaining,
+						PercentRemaining: percentRemaining,
+						QuotaID:          "completions",
+						Unlimited:        false,
+					}
+				}
+			}
+			
+			// Premium interactions don't exist for non-enterprise, leave as zero values
+			quotaSnapshots.PremiumInteractions = QuotaDetail{
+				QuotaID:   "premium_interactions",
+				Unlimited: false,
+			}
+			
+			// Add quota_snapshots to the token response
+			tokenResp["quota_snapshots"] = quotaSnapshots
+			tokenResp["access_type_sku"] = quotaData.AccessTypeSKU
+			tokenResp["copilot_plan"] = quotaData.CopilotPlan
+		}
+	}
+
+	// Re-serialize the enriched response
+	enrichedBody, errMarshal := json.Marshal(tokenResp)
+	if errMarshal != nil {
+		log.WithError(errMarshal).Debug("failed to marshal enriched response")
+		return response
+	}
+
+	response.Body = string(enrichedBody)
+
+	return response
+}