feat(client): support extra headers from model resolution

.github/workflows/pullfrog.yml

@@ -34,22 +34,21 @@ jobs:
           ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
           CLAUDE_CODE_OAUTH_TOKEN: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-          GOOGLE_GENERATIVE_AI_API_KEY:
-            ${{ secrets.GOOGLE_GENERATIVE_AI_API_KEY }}
+          GOOGLE_GENERATIVE_AI_API_KEY: ${{ secrets.GOOGLE_GENERATIVE_AI_API_KEY }}
           GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
           XAI_API_KEY: ${{ secrets.XAI_API_KEY }}
           DEEPSEEK_API_KEY: ${{ secrets.DEEPSEEK_API_KEY }}
           MOONSHOT_API_KEY: ${{ secrets.MOONSHOT_API_KEY }}
           OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }}
           OPENCODE_API_KEY: ${{ secrets.OPENCODE_API_KEY }}
-          
+
           # for Amazon Bedrock (https://docs.pullfrog.com/bedrock)
           # AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.AWS_BEARER_TOKEN_BEDROCK }}
           # AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
           # AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
           # AWS_REGION: us-east-1
           # BEDROCK_MODEL_ID: <bedrock-model-id>
-          
+
           # for Google Vertex AI (https://docs.pullfrog.com/vertex)
           # VERTEX_SERVICE_ACCOUNT_JSON: ${{ secrets.VERTEX_SERVICE_ACCOUNT_JSON }}
           # GOOGLE_CLOUD_PROJECT: my-project

README.md

@@ -452,7 +452,7 @@ Add custom AI providers:
   get_models?(headers: table): table<CopilotChat.Provider.model>,
 
   -- Optional: Resolve a user-facing model id to a provider model id
-  resolve_model?(headers: table, model: string): string,
+  resolve_model?(headers: table, model: string): string, table<string, string>?,
 }
 ```
 

lua/CopilotChat/client.lua

@@ -319,13 +319,20 @@ function Client:ask(opts)
     error('Provider not found: ' .. provider_name)
   end
 
+  local requested_model = opts.model
+  local resolved_headers = nil
   if provider.resolve_model then
     local headers = self:authenticate(provider_name)
-    local resolved_model = provider.resolve_model(headers, opts.model)
+    local resolved_model, extra_headers = provider.resolve_model(headers, opts.model)
+    resolved_headers = extra_headers
     opts.model = resolved_model
     model_config = models[opts.model]
     if not model_config then
-      error('Resolved model not found: ' .. opts.model)
+      if requested_model == 'auto' then
+        model_config = vim.tbl_extend('force', models[requested_model], { id = opts.model })
+      else
+        error('Resolved model not found: ' .. opts.model)
+      end
     end
   end
 
@@ -538,6 +545,10 @@ function Client:ask(opts)
 
   local headers = self:authenticate(provider_name)
 
+  if resolved_headers then
+    headers = vim.tbl_extend('force', headers, resolved_headers)
+  end
+
   local request, extra_headers =
     provider.prepare_input(generate_ask_request(opts.system_prompt, history, generated_messages), options)
 

lua/CopilotChat/config/providers.lua

@@ -530,7 +530,7 @@ end
 ---@field get_headers nil|fun():table<string, string>,number?
 ---@field get_info nil|fun(headers:table):string[]
 ---@field get_models nil|fun(headers:table):table<CopilotChat.client.Model>
----@field resolve_model nil|fun(headers:table, model: string):string
+---@field resolve_model nil|fun(headers:table, model: string):string, table<string, string>?
 ---@field prepare_input nil|fun(inputs:CopilotChat.client.Message[], opts:CopilotChat.config.providers.Options):table,table?
 ---@field prepare_output nil|fun(output:table, opts:CopilotChat.config.providers.Options):CopilotChat.config.providers.Output
 ---@field get_url nil|fun(opts:CopilotChat.config.providers.Options):string
@@ -714,7 +714,7 @@ M.copilot = {
       error(err)
     end
 
-    return response.body.selected_model
+    return response.body.selected_model, { ['Copilot-Session-Token'] = response.body.session_token }
   end,
 
   prepare_input = function(inputs, opts)