fix(amp): pass mapped model to gemini bridge via context

Gemini handler extracts model from URL path, not JSON body, so
rewriting the request body alone wasn't sufficient for model mapping.

- Add MappedModelContextKey constant for context passing
- Update routes.go to use NewFallbackHandlerWithMapper
- Add check for valid mapping before routing to local handler
- Add tests for gemini bridge model mapping
This commit is contained in:
huynguyen03.dev
2025-12-06 18:59:44 +07:00
parent 7ea14479fb
commit 08586334af
4 changed files with 150 additions and 3 deletions

View File

@@ -170,9 +170,9 @@ func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *ha
// If no local OAuth is available, falls back to ampcode.com proxy.
geminiHandlers := gemini.NewGeminiAPIHandler(baseHandler)
geminiBridge := createGeminiBridgeHandler(geminiHandlers)
geminiV1Beta1Fallback := NewFallbackHandler(func() *httputil.ReverseProxy {
geminiV1Beta1Fallback := NewFallbackHandlerWithMapper(func() *httputil.ReverseProxy {
return m.getProxy()
})
}, m.modelMapper)
geminiV1Beta1Handler := geminiV1Beta1Fallback.WrapHandler(geminiBridge)
// Route POST model calls through Gemini bridge when a local provider exists, otherwise proxy.
@@ -187,8 +187,18 @@ func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *ha
}
if modelPart != "" {
normalized, _ := util.NormalizeGeminiThinkingModel(modelPart)
// Only handle locally when we have a provider; otherwise fall back to proxy
// Only handle locally when we have a provider or a valid mapping; otherwise fall back to proxy
hasProvider := false
if providers := util.GetProviderName(normalized); len(providers) > 0 {
hasProvider = true
} else if m.modelMapper != nil {
// Check if mapped model has provider (MapModel returns target only if it has providers)
if mapped := m.modelMapper.MapModel(normalized); mapped != "" {
hasProvider = true
}
}
if hasProvider {
geminiV1Beta1Handler(c)
return
}