From 7ae00320dcc6b8a3f9d5f1227842662bd2d5b65e Mon Sep 17 00:00:00 2001 From: Ben Vargas Date: Wed, 19 Nov 2025 15:44:55 -0700 Subject: [PATCH] fix(amp): enable OAuth fallback for Gemini v1beta1 routes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AMP CLI sends Gemini requests to non-standard paths that were being directly proxied to ampcode.com without checking for local OAuth. This fix adds: - GeminiBridge handler to transform AMP CLI paths to standard format - Enhanced model extraction from AMP's /publishers/google/models/* paths - FallbackHandler wrapper to check for local OAuth before proxying Flow: - If user has local Google OAuth → use it (free tier) - If no local OAuth → fallback to ampcode.com (charges credits) Fixes issue where gemini-3-pro-preview requests always charged AMP credits even when user had valid Google Cloud OAuth configured. --- internal/api/modules/amp/amp.go | 2 +- internal/api/modules/amp/fallback_handlers.go | 17 ++++++- internal/api/modules/amp/gemini_bridge.go | 45 +++++++++++++++++++ internal/api/modules/amp/routes.go | 14 ++++-- 4 files changed, 72 insertions(+), 6 deletions(-) create mode 100644 internal/api/modules/amp/gemini_bridge.go diff --git a/internal/api/modules/amp/amp.go b/internal/api/modules/amp/amp.go index 07e52e46..f95218ae 100644 --- a/internal/api/modules/amp/amp.go +++ b/internal/api/modules/amp/amp.go @@ -131,7 +131,7 @@ func (m *AmpModule) Register(ctx modules.Context) error { // Register management proxy routes (requires upstream) // Restrict to localhost by default for security (prevents drive-by browser attacks) handler := proxyHandler(proxy) - m.registerManagementRoutes(ctx.Engine, handler, ctx.Config.AmpRestrictManagementToLocalhost) + m.registerManagementRoutes(ctx.Engine, ctx.BaseHandler, handler, ctx.Config.AmpRestrictManagementToLocalhost) log.Infof("Amp upstream proxy enabled for: %s", upstreamURL) log.Debug("Amp provider alias routes registered") diff --git a/internal/api/modules/amp/fallback_handlers.go b/internal/api/modules/amp/fallback_handlers.go index d8c140ad..e7b28986 100644 --- a/internal/api/modules/amp/fallback_handlers.go +++ b/internal/api/modules/amp/fallback_handlers.go @@ -102,8 +102,8 @@ func extractModelFromRequest(body []byte, c *gin.Context) string { } } - // For Gemini requests, model is in the URL path: /models/{model}:generateContent - // Extract from :action parameter (e.g., "gemini-pro:generateContent") + // For Gemini requests, model is in the URL path + // Standard format: /models/{model}:generateContent -> :action parameter if action := c.Param("action"); action != "" { // Split by colon to get model name (e.g., "gemini-pro:generateContent" -> "gemini-pro") parts := strings.Split(action, ":") @@ -112,5 +112,18 @@ func extractModelFromRequest(body []byte, c *gin.Context) string { } } + // AMP CLI format: /publishers/google/models/{model}:method -> *path parameter + // Example: /publishers/google/models/gemini-3-pro-preview:streamGenerateContent + if path := c.Param("path"); path != "" { + // Look for /models/{model}:method pattern + if idx := strings.Index(path, "/models/"); idx >= 0 { + modelPart := path[idx+8:] // Skip "/models/" + // Split by colon to get model name + if colonIdx := strings.Index(modelPart, ":"); colonIdx > 0 { + return modelPart[:colonIdx] + } + } + } + return "" } diff --git a/internal/api/modules/amp/gemini_bridge.go b/internal/api/modules/amp/gemini_bridge.go new file mode 100644 index 00000000..3b3d8374 --- /dev/null +++ b/internal/api/modules/amp/gemini_bridge.go @@ -0,0 +1,45 @@ +package amp + +import ( + "strings" + + "github.com/gin-gonic/gin" + "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/gemini" +) + +// createGeminiBridgeHandler creates a handler that bridges AMP CLI's non-standard Gemini paths +// to our standard Gemini handler by rewriting the request context. +// +// AMP CLI format: /publishers/google/models/gemini-3-pro-preview:streamGenerateContent +// Standard format: /models/gemini-3-pro-preview:streamGenerateContent +// +// This extracts the model+method from the AMP path and sets it as the :action parameter +// so the standard Gemini handler can process it. +func createGeminiBridgeHandler(geminiHandler *gemini.GeminiAPIHandler) gin.HandlerFunc { + return func(c *gin.Context) { + // Get the full path from the catch-all parameter + path := c.Param("path") + + // Extract model:method from AMP CLI path format + // Example: /publishers/google/models/gemini-3-pro-preview:streamGenerateContent + if idx := strings.Index(path, "/models/"); idx >= 0 { + // Extract everything after "/models/" + actionPart := path[idx+8:] // Skip "/models/" + + // Set this as the :action parameter that the Gemini handler expects + c.Params = append(c.Params, gin.Param{ + Key: "action", + Value: actionPart, + }) + + // Call the standard Gemini handler + geminiHandler.GeminiHandler(c) + return + } + + // If we can't parse the path, return 400 + c.JSON(400, gin.H{ + "error": "Invalid Gemini API path format", + }) + } +} diff --git a/internal/api/modules/amp/routes.go b/internal/api/modules/amp/routes.go index 5231ec86..e0c11b5a 100644 --- a/internal/api/modules/amp/routes.go +++ b/internal/api/modules/amp/routes.go @@ -65,7 +65,7 @@ func noCORSMiddleware() gin.HandlerFunc { // registerManagementRoutes registers Amp management proxy routes // These routes proxy through to the Amp control plane for OAuth, user management, etc. // If restrictToLocalhost is true, routes will only accept connections from 127.0.0.1/::1. -func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, proxyHandler gin.HandlerFunc, restrictToLocalhost bool) { +func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *handlers.BaseAPIHandler, proxyHandler gin.HandlerFunc, restrictToLocalhost bool) { ampAPI := engine.Group("/api") // Always disable CORS for management routes to prevent browser-based attacks @@ -96,8 +96,16 @@ func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, proxyHandler gi ampAPI.Any("/otel", proxyHandler) ampAPI.Any("/otel/*path", proxyHandler) - // Google v1beta1 passthrough (Gemini native API) - ampAPI.Any("/provider/google/v1beta1/*path", proxyHandler) + // Google v1beta1 passthrough with OAuth fallback + // AMP CLI uses non-standard paths like /publishers/google/models/... + // We bridge these to our standard Gemini handler to enable local OAuth. + // If no local OAuth is available, falls back to ampcode.com proxy. + geminiHandlers := gemini.NewGeminiAPIHandler(baseHandler) + geminiBridge := createGeminiBridgeHandler(geminiHandlers) + geminiV1Beta1Fallback := NewFallbackHandler(func() *httputil.ReverseProxy { + return m.proxy + }) + ampAPI.POST("/provider/google/v1beta1/*path", geminiV1Beta1Fallback.WrapHandler(geminiBridge)) } // registerProviderAliases registers /api/provider/{provider}/... routes