diff --git a/internal/api/modules/amp/gemini_bridge.go b/internal/api/modules/amp/gemini_bridge.go index cc31b685..d6ad8f79 100644 --- a/internal/api/modules/amp/gemini_bridge.go +++ b/internal/api/modules/amp/gemini_bridge.go @@ -4,7 +4,6 @@ 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 @@ -15,16 +14,19 @@ import ( // // 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 { +// +// The handler parameter should be a Gemini-compatible handler that expects the :action param. +func createGeminiBridgeHandler(handler gin.HandlerFunc) 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/" + const modelsPrefix = "/models/" + if idx := strings.Index(path, modelsPrefix); idx >= 0 { + // Extract everything after modelsPrefix + actionPart := path[idx+len(modelsPrefix):] // Check if model was mapped by FallbackHandler if mappedModel, exists := c.Get(MappedModelContextKey); exists { @@ -44,8 +46,8 @@ func createGeminiBridgeHandler(geminiHandler *gemini.GeminiAPIHandler) gin.Handl Value: actionPart, }) - // Call the standard Gemini handler - geminiHandler.GeminiHandler(c) + // Call the handler + handler(c) return } diff --git a/internal/api/modules/amp/gemini_bridge_test.go b/internal/api/modules/amp/gemini_bridge_test.go index 88accbf4..347456c3 100644 --- a/internal/api/modules/amp/gemini_bridge_test.go +++ b/internal/api/modules/amp/gemini_bridge_test.go @@ -3,12 +3,9 @@ package amp import ( "net/http" "net/http/httptest" - "strings" "testing" "github.com/gin-gonic/gin" - "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers" - "github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers/gemini" ) func TestCreateGeminiBridgeHandler_ActionParameterExtraction(t *testing.T) { @@ -38,12 +35,6 @@ func TestCreateGeminiBridgeHandler_ActionParameterExtraction(t *testing.T) { mappedModel: "gemini-flash", expectedAction: "gemini-flash:streamGenerateContent", }, - { - name: "empty_mapped_model_ignored", - path: "/publishers/google/models/gemini-pro:generateContent", - mappedModel: "", - expectedAction: "gemini-pro:generateContent", - }, } for _, tt := range tests { @@ -55,27 +46,8 @@ func TestCreateGeminiBridgeHandler_ActionParameterExtraction(t *testing.T) { c.JSON(http.StatusOK, gin.H{"captured": capturedAction}) } - // Mirror the bridge logic from gemini_bridge.go - bridgeHandler := func(c *gin.Context) { - path := c.Param("path") - if idx := strings.Index(path, "/models/"); idx >= 0 { - actionPart := path[idx+8:] - - if mappedModel, exists := c.Get(MappedModelContextKey); exists { - if strModel, ok := mappedModel.(string); ok && strModel != "" { - if colonIdx := strings.Index(actionPart, ":"); colonIdx > 0 { - method := actionPart[colonIdx:] - actionPart = strModel + method - } - } - } - - c.Params = append(c.Params, gin.Param{Key: "action", Value: actionPart}) - mockGeminiHandler(c) - return - } - c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid path"}) - } + // Use the actual createGeminiBridgeHandler function + bridgeHandler := createGeminiBridgeHandler(mockGeminiHandler) r := gin.New() if tt.mappedModel != "" { @@ -103,9 +75,10 @@ func TestCreateGeminiBridgeHandler_ActionParameterExtraction(t *testing.T) { func TestCreateGeminiBridgeHandler_InvalidPath(t *testing.T) { gin.SetMode(gin.TestMode) - base := &handlers.BaseAPIHandler{} - geminiHandlers := gemini.NewGeminiAPIHandler(base) - bridgeHandler := createGeminiBridgeHandler(geminiHandlers) + mockHandler := func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + } + bridgeHandler := createGeminiBridgeHandler(mockHandler) r := gin.New() r.POST("/api/provider/google/v1beta1/*path", bridgeHandler) diff --git a/internal/api/modules/amp/routes.go b/internal/api/modules/amp/routes.go index 33525e61..8cb208fe 100644 --- a/internal/api/modules/amp/routes.go +++ b/internal/api/modules/amp/routes.go @@ -169,7 +169,7 @@ func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *ha // 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) + geminiBridge := createGeminiBridgeHandler(geminiHandlers.GeminiHandler) geminiV1Beta1Fallback := NewFallbackHandlerWithMapper(func() *httputil.ReverseProxy { return m.getProxy() }, m.modelMapper)