Merge pull request #163 from yg1996/add-custom-headers-support

Add custom headers support for API Management integration
2026-01-19 17:28:08 +11:00
parent 63993128d7 15ae50ae2f
commit a6101c89c6
10 changed files with 648 additions and 102 deletions
@@ -156,6 +156,54 @@ steps:
      cat "${{ steps.inference.outputs.response-file }}"
 ```

+### Using custom headers
+
+You can include custom HTTP headers in your API requests, which is useful for integrating with API Management platforms, adding tracking information, or routing requests through custom gateways.
+
+#### YAML format (recommended for multiple headers)
+
+```yaml
+steps:
+  - name: AI Inference with Azure APIM
+    id: inference
+    uses: actions/ai-inference@v1
+    with:
+      prompt: 'Analyze this code for security issues...'
+      endpoint: ${{ secrets.APIM_ENDPOINT }}
+      token: ${{ secrets.APIM_KEY }}
+      custom-headers: |
+        Ocp-Apim-Subscription-Key: ${{ secrets.APIM_SUBSCRIPTION_KEY }}
+        serviceName: code-review-workflow
+        env: production
+        team: security
+        computer: github-actions
+```
+
+#### JSON format (alternative for compact syntax)
+
+```yaml
+steps:
+  - name: AI Inference with Custom Headers
+    id: inference
+    uses: actions/ai-inference@v1
+    with:
+      prompt: 'Hello!'
+      custom-headers: '{"X-Custom-Header": "value", "X-Team": "engineering", "X-Request-ID": "${{ github.run_id }}"}'
+```
+
+#### Use cases for custom headers
+
+- **API Management**: Integrate with Azure APIM, AWS API Gateway, Kong, or other API management platforms
+- **Request tracking**: Add correlation IDs, request IDs, or workflow identifiers
+- **Rate limiting**: Include quota or tier information for custom rate limiting
+- **Multi-tenancy**: Identify teams, services, or environments
+- **Observability**: Add metadata for logging, monitoring, and debugging
+- **Routing**: Control request routing through custom gateways or load balancers
+
+**Header name requirements**: Header names must follow the HTTP token syntax defined in RFC 7230 (which permits underscores). For maximum compatibility with intermediaries and tooling, we recommend using only alphanumeric characters and hyphens.
+
+**Security note**: Always use GitHub secrets for sensitive header values like API keys, tokens, or passwords. The action automatically masks common sensitive headers (containing `key`, `token`, `secret`, `password`, or `authorization`) in logs.
+
 ### GitHub MCP Integration (Model Context Protocol)

 This action now supports **read-only** integration with the GitHub-hosted Model
@@ -228,20 +276,21 @@ perform actions like searching issues and PRs.
 Various inputs are defined in [`action.yml`](action.yml) to let you configure
 the action:

-| Name                 | Description                                                                                                                                                   | Default                              |
-| -------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------ |
-| `token`              | Token to use for inference. Typically the GITHUB_TOKEN secret                                                                                                 | `github.token`                       |
-| `prompt`             | The prompt to send to the model                                                                                                                               | N/A                                  |
-| `prompt-file`        | Path to a file containing the prompt (supports .txt and .prompt.yml formats). If both `prompt` and `prompt-file` are provided, `prompt-file` takes precedence | `""`                                 |
-| `input`              | Template variables in YAML format for .prompt.yml files (e.g., `var1: value1` on separate lines)                                                              | `""`                                 |
-| `file_input`         | Template variables in YAML where values are file paths. The file contents are read and used for templating                                                    | `""`                                 |
-| `system-prompt`      | The system prompt to send to the model                                                                                                                        | `"You are a helpful assistant"`      |
-| `system-prompt-file` | Path to a file containing the system prompt. If both `system-prompt` and `system-prompt-file` are provided, `system-prompt-file` takes precedence             | `""`                                 |
-| `model`              | The model to use for inference. Must be available in the [GitHub Models](https://github.com/marketplace?type=models) catalog                                  | `openai/gpt-4o`                      |
-| `endpoint`           | The endpoint to use for inference. If you're running this as part of an org, you should probably use the org-specific Models endpoint                         | `https://models.github.ai/inference` |
-| `max-tokens`         | The max number of tokens to generate                                                                                                                          | 200                                  |
-| `enable-github-mcp`  | Enable Model Context Protocol integration with GitHub tools                                                                                                   | `false`                              |
-| `github-mcp-token`   | Token to use for GitHub MCP server (defaults to the main token if not specified).                                                                             | `""`                                 |
+| Name                 | Description                                                                                                                                                                                                        | Default                              |
+| -------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------ |
+| `token`              | Token to use for inference. Typically the GITHUB_TOKEN secret                                                                                                                                                      | `github.token`                       |
+| `prompt`             | The prompt to send to the model                                                                                                                                                                                    | N/A                                  |
+| `prompt-file`        | Path to a file containing the prompt (supports .txt and .prompt.yml formats). If both `prompt` and `prompt-file` are provided, `prompt-file` takes precedence                                                      | `""`                                 |
+| `input`              | Template variables in YAML format for .prompt.yml files (e.g., `var1: value1` on separate lines)                                                                                                                   | `""`                                 |
+| `file_input`         | Template variables in YAML where values are file paths. The file contents are read and used for templating                                                                                                         | `""`                                 |
+| `system-prompt`      | The system prompt to send to the model                                                                                                                                                                             | `"You are a helpful assistant"`      |
+| `system-prompt-file` | Path to a file containing the system prompt. If both `system-prompt` and `system-prompt-file` are provided, `system-prompt-file` takes precedence                                                                  | `""`                                 |
+| `model`              | The model to use for inference. Must be available in the [GitHub Models](https://github.com/marketplace?type=models) catalog                                                                                       | `openai/gpt-4o`                      |
+| `endpoint`           | The endpoint to use for inference. If you're running this as part of an org, you should probably use the org-specific Models endpoint                                                                              | `https://models.github.ai/inference` |
+| `max-tokens`         | The max number of tokens to generate                                                                                                                                                                               | 200                                  |
+| `enable-github-mcp`  | Enable Model Context Protocol integration with GitHub tools                                                                                                                                                        | `false`                              |
+| `github-mcp-token`   | Token to use for GitHub MCP server (defaults to the main token if not specified).                                                                                                                                  | `""`                                 |
+| `custom-headers`     | Custom HTTP headers to include in API requests. Supports both YAML format (`header1: value1`) and JSON format (`{"header1": "value1"}`). Useful for API Management platforms, rate limiting, and request tracking. | `""`                                 |

 ## Outputs

@@ -11,7 +11,7 @@ vi.mock('fs', () => ({

 vi.mock('@actions/core', () => core)

-const {loadContentFromFileOrInput} = await import('../src/helpers.js')
+const {loadContentFromFileOrInput, parseCustomHeaders} = await import('../src/helpers.js')

 describe('helpers.ts', () => {
  beforeEach(() => {
@@ -132,4 +132,243 @@ describe('helpers.ts', () => {
      expect(result).toBe(defaultValue)
    })
  })
+
+  describe('parseCustomHeaders', () => {
+    beforeEach(() => {
+      vi.clearAllMocks()
+    })
+
+    it('parses YAML format headers correctly', () => {
+      const yamlInput = `header1: value1
+header2: value2
+X-Custom-Header: custom-value`
+
+      const result = parseCustomHeaders(yamlInput)
+
+      expect(result).toEqual({
+        header1: 'value1',
+        header2: 'value2',
+        'X-Custom-Header': 'custom-value',
+      })
+      expect(core.info).toHaveBeenCalledWith('Custom header added: header1: value1')
+      expect(core.info).toHaveBeenCalledWith('Custom header added: header2: value2')
+      expect(core.info).toHaveBeenCalledWith('Custom header added: X-Custom-Header: custom-value')
+    })
+
+    it('parses JSON format headers correctly', () => {
+      const jsonInput = '{"header1": "value1", "header2": "value2", "X-Team": "engineering"}'
+
+      const result = parseCustomHeaders(jsonInput)
+
+      expect(result).toEqual({
+        header1: 'value1',
+        header2: 'value2',
+        'X-Team': 'engineering',
+      })
+      expect(core.info).toHaveBeenCalledWith('Custom header added: header1: value1')
+      expect(core.info).toHaveBeenCalledWith('Custom header added: header2: value2')
+      expect(core.info).toHaveBeenCalledWith('Custom header added: X-Team: engineering')
+    })
+
+    it('returns empty object for empty input', () => {
+      expect(parseCustomHeaders('')).toEqual({})
+      expect(parseCustomHeaders('  ')).toEqual({})
+      expect(core.warning).not.toHaveBeenCalled()
+    })
+
+    it('masks sensitive header values in logs', () => {
+      const yamlInput = `Ocp-Apim-Subscription-Key: secret123
+X-Api-Token: token456
+Authorization: Bearer abc123
+serviceName: my-service
+password: pass123`
+
+      const result = parseCustomHeaders(yamlInput)
+
+      expect(result).toEqual({
+        'Ocp-Apim-Subscription-Key': 'secret123',
+        'X-Api-Token': 'token456',
+        Authorization: 'Bearer abc123',
+        serviceName: 'my-service',
+        password: 'pass123',
+      })
+
+      // Sensitive headers should be masked
+      expect(core.info).toHaveBeenCalledWith('Custom header added: Ocp-Apim-Subscription-Key: ***MASKED***')
+      expect(core.info).toHaveBeenCalledWith('Custom header added: X-Api-Token: ***MASKED***')
+      expect(core.info).toHaveBeenCalledWith('Custom header added: Authorization: ***MASKED***')
+      expect(core.info).toHaveBeenCalledWith('Custom header added: password: ***MASKED***')
+
+      // Non-sensitive headers should not be masked
+      expect(core.info).toHaveBeenCalledWith('Custom header added: serviceName: my-service')
+    })
+
+    it('validates header names and skips invalid ones', () => {
+      const yamlInput = `valid-header: value1
+invalid header: value2
+invalid_underscore: value3
+invalid@header: value4
+valid123: value5`
+
+      const result = parseCustomHeaders(yamlInput)
+
+      expect(result).toEqual({
+        'valid-header': 'value1',
+        valid123: 'value5',
+      })
+
+      expect(core.warning).toHaveBeenCalledWith(expect.stringContaining('Skipping invalid header name: invalid header'))
+      expect(core.warning).toHaveBeenCalledWith(
+        expect.stringContaining('Skipping invalid header name: invalid_underscore'),
+      )
+      expect(core.warning).toHaveBeenCalledWith(expect.stringContaining('Skipping invalid header name: invalid@header'))
+    })
+
+    it('warns and returns empty object for invalid JSON', () => {
+      const invalidJson = '{invalid json}'
+
+      const result = parseCustomHeaders(invalidJson)
+
+      expect(result).toEqual({})
+      expect(core.warning).toHaveBeenCalledWith(expect.stringContaining('Failed to parse custom headers'))
+    })
+
+    it('warns and returns empty object for invalid YAML', () => {
+      const invalidYaml = 'invalid: yaml: structure: bad'
+
+      const result = parseCustomHeaders(invalidYaml)
+
+      expect(result).toEqual({})
+      expect(core.warning).toHaveBeenCalledWith(expect.stringContaining('Failed to parse custom headers'))
+    })
+
+    it('warns and returns empty object for JSON array', () => {
+      const jsonArray = '["header1", "header2"]'
+
+      const result = parseCustomHeaders(jsonArray)
+
+      expect(result).toEqual({})
+      expect(core.warning).toHaveBeenCalledWith('Custom headers JSON must be an object, not null or an array')
+    })
+
+    it('warns and returns empty object for null value', () => {
+      // The string 'null' is valid YAML and gets parsed as null
+      const nullValue = 'null'
+
+      const result = parseCustomHeaders(nullValue)
+
+      expect(result).toEqual({})
+      expect(core.warning).toHaveBeenCalledWith('Custom headers YAML must be an object')
+    })
+
+    it('warns and returns empty object for YAML array', () => {
+      const yamlArray = `- header1
+- header2`
+
+      const result = parseCustomHeaders(yamlArray)
+
+      expect(result).toEqual({})
+      expect(core.warning).toHaveBeenCalledWith('Custom headers YAML must be an object')
+    })
+
+    it('converts non-string values to strings', () => {
+      const jsonInput = '{"numericHeader": 123, "boolHeader": true, "nullHeader": null}'
+
+      const result = parseCustomHeaders(jsonInput)
+
+      expect(result).toEqual({
+        numericHeader: '123',
+        boolHeader: 'true',
+        nullHeader: 'null',
+      })
+    })
+
+    it('rejects header values with newline characters (LF)', () => {
+      const jsonInput = '{"X-Custom-Header": "value\\nwith\\nnewline", "header1": "safe-value"}'
+
+      const result = parseCustomHeaders(jsonInput)
+
+      // Only the safe header should be accepted
+      expect(result).toEqual({
+        header1: 'safe-value',
+      })
+
+      expect(core.warning).toHaveBeenCalledWith(
+        'Skipping header "X-Custom-Header" because its value contains newline characters, which are not allowed in HTTP header values.',
+      )
+    })
+
+    it('rejects header values with carriage return characters (CR)', () => {
+      const jsonInput = '{"X-Injected": "value\\rwith\\rcarriage", "X-Safe": "safe-value"}'
+
+      const result = parseCustomHeaders(jsonInput)
+
+      // Only the safe header should be accepted
+      expect(result).toEqual({
+        'X-Safe': 'safe-value',
+      })
+
+      expect(core.warning).toHaveBeenCalledWith(
+        'Skipping header "X-Injected" because its value contains newline characters, which are not allowed in HTTP header values.',
+      )
+    })
+
+    it('rejects header values with CRLF sequences', () => {
+      const jsonInput = '{"X-Attack": "value\\r\\nInjected-Header: malicious", "X-Valid": "normal"}'
+
+      const result = parseCustomHeaders(jsonInput)
+
+      // Only the valid header should be accepted
+      expect(result).toEqual({
+        'X-Valid': 'normal',
+      })
+
+      expect(core.warning).toHaveBeenCalledWith(
+        'Skipping header "X-Attack" because its value contains newline characters, which are not allowed in HTTP header values.',
+      )
+    })
+
+    it('rejects multiline YAML values for security', () => {
+      const yamlInput = `header1: value1
+header2: |
+  multiline
+  value
+  here`
+
+      const result = parseCustomHeaders(yamlInput)
+
+      // header2 should be rejected because it contains newlines
+      expect(result).toEqual({
+        header1: 'value1',
+      })
+
+      expect(core.warning).toHaveBeenCalledWith(
+        'Skipping header "header2" because its value contains newline characters, which are not allowed in HTTP header values.',
+      )
+    })
+
+    it('handles complex real-world Azure APIM example', () => {
+      const apimHeaders = `Ocp-Apim-Subscription-Key: my-subscription-key-123
+serviceName: terraform-plan-workflow
+env: prod
+team: infrastructure
+computer: github-actions
+systemID: terraform-ci`
+
+      const result = parseCustomHeaders(apimHeaders)
+
+      expect(result).toEqual({
+        'Ocp-Apim-Subscription-Key': 'my-subscription-key-123',
+        serviceName: 'terraform-plan-workflow',
+        env: 'prod',
+        team: 'infrastructure',
+        computer: 'github-actions',
+        systemID: 'terraform-ci',
+      })
+
+      // Only the subscription key should be masked
+      expect(core.info).toHaveBeenCalledWith('Custom header added: Ocp-Apim-Subscription-Key: ***MASKED***')
+      expect(core.info).toHaveBeenCalledWith('Custom header added: serviceName: terraform-plan-workflow')
+    })
+  })
 })
@@ -75,6 +75,49 @@ describe('inference.ts', () => {
        max_tokens: 100,
        model: 'gpt-4',
      })
+
+      // Verify OpenAI client was initialized with empty custom headers
+      expect(mockOpenAIClient).toHaveBeenCalledWith({
+        apiKey: 'test-token',
+        baseURL: 'https://api.test.com',
+        defaultHeaders: {},
+      })
+    })
+
+    it('includes custom headers in OpenAI client', async () => {
+      const requestWithHeaders = {
+        ...mockRequest,
+        customHeaders: {
+          'X-Custom-Header': 'custom-value',
+          'Ocp-Apim-Subscription-Key': 'secret123',
+        },
+      }
+
+      const mockResponse = {
+        choices: [
+          {
+            message: {
+              content: 'Response with headers',
+            },
+          },
+        ],
+      }
+
+      mockCreate.mockResolvedValue(mockResponse)
+
+      const result = await simpleInference(requestWithHeaders)
+
+      expect(result).toBe('Response with headers')
+
+      // Verify OpenAI client was initialized with custom headers
+      expect(mockOpenAIClient).toHaveBeenCalledWith({
+        apiKey: 'test-token',
+        baseURL: 'https://api.test.com',
+        defaultHeaders: {
+          'X-Custom-Header': 'custom-value',
+          'Ocp-Apim-Subscription-Key': 'secret123',
+        },
+      })
    })

    it('handles null response content', async () => {
@@ -186,6 +229,50 @@ describe('inference.ts', () => {
      expect(callArgs.response_format).toBeUndefined()
      expect(callArgs.model).toBe('gpt-4')
      expect(callArgs.max_tokens).toBe(100)
+
+      // Verify OpenAI client was initialized with empty custom headers
+      expect(mockOpenAIClient).toHaveBeenCalledWith({
+        apiKey: 'test-token',
+        baseURL: 'https://api.test.com',
+        defaultHeaders: {},
+      })
+    })
+
+    it('includes custom headers in MCP inference', async () => {
+      const requestWithHeaders = {
+        ...mockRequest,
+        customHeaders: {
+          serviceName: 'test-service',
+          'X-Team': 'engineering',
+        },
+      }
+
+      const mockResponse = {
+        choices: [
+          {
+            message: {
+              content: 'MCP response with headers',
+              tool_calls: null,
+            },
+          },
+        ],
+      }
+
+      mockCreate.mockResolvedValue(mockResponse)
+
+      const result = await mcpInference(requestWithHeaders, mockMcpClient)
+
+      expect(result).toBe('MCP response with headers')
+
+      // Verify OpenAI client was initialized with custom headers
+      expect(mockOpenAIClient).toHaveBeenCalledWith({
+        apiKey: 'test-token',
+        baseURL: 'https://api.test.com',
+        defaultHeaders: {
+          serviceName: 'test-service',
+          'X-Team': 'engineering',
+        },
+      })
    })

    it('executes tool calls and continues conversation', async () => {
@@ -171,6 +171,9 @@ describe('main.ts', () => {
      endpoint: 'https://api.test.com',
      token: 'fake-token',
      responseFormat: undefined,
+      temperature: undefined,
+      topP: undefined,
+      customHeaders: {},
    })
    expect(mockConnectToGitHubMCP).not.toHaveBeenCalled()
    expect(mockMcpInference).not.toHaveBeenCalled()
@@ -259,6 +262,9 @@ describe('main.ts', () => {
      endpoint: 'https://api.test.com',
      token: 'fake-token',
      responseFormat: undefined,
+      temperature: undefined,
+      topP: undefined,
+      customHeaders: {},
    })
    verifyStandardResponse()
    expect(mockProcessExit).toHaveBeenCalledWith(0)
@@ -62,6 +62,10 @@ inputs:
    description: 'Comma-separated list of toolsets to enable for GitHub MCP (e.g., "repos,issues,pull_requests,actions"). Use "all" for all toolsets, "default" for default set. If not specified, uses default toolsets (context,repos,issues,pull_requests,users).'
    required: false
    default: ''
+  custom-headers:
+    description: 'Custom HTTP headers to include in API requests. Supports both YAML format (header1: value1) and JSON format ({"header1": "value1"}). Useful for API Management platforms, rate limiting, and request tracking.'
+    required: false
+    default: ''

 # Define your outputs here.
 outputs:
@@ -58308,6 +58308,7 @@ async function simpleInference(request) {
    const client = new OpenAI({
        apiKey: request.token,
        baseURL: request.endpoint,
+        defaultHeaders: request.customHeaders || {},
    });
    const chatCompletionRequest = {
        messages: request.messages,
@@ -58334,6 +58335,7 @@ async function mcpInference(request, githubMcpClient) {
    const client = new OpenAI({
        apiKey: request.token,
        baseURL: request.endpoint,
+        defaultHeaders: request.customHeaders || {},
    });
    // Start with the pre-processed messages
    const messages = [...request.messages];
@@ -58437,90 +58439,6 @@ async function chatCompletion(client, params, context) {
    }
 }

-/**
- * Helper function to load content from a file or use fallback input
- * @param filePathInput - Input name for the file path
- * @param contentInput - Input name for the direct content
- * @param defaultValue - Default value to use if neither file nor content is provided
- * @returns The loaded content
- */
-function loadContentFromFileOrInput(filePathInput, contentInput, defaultValue) {
-    const filePath = coreExports.getInput(filePathInput);
-    const contentString = coreExports.getInput(contentInput);
-    if (filePath !== undefined && filePath !== '') {
-        if (!fs.existsSync(filePath)) {
-            throw new Error(`File for ${filePathInput} was not found: ${filePath}`);
-        }
-        return fs.readFileSync(filePath, 'utf-8');
-    }
-    else if (contentString !== undefined && contentString !== '') {
-        return contentString;
-    }
-    else if (defaultValue !== undefined) {
-        return defaultValue;
-    }
-    else {
-        throw new Error(`Neither ${filePathInput} nor ${contentInput} was set`);
-    }
-}
-/**
- * Build messages array from either prompt config or legacy format
- */
-function buildMessages(promptConfig, systemPrompt, prompt) {
-    if (promptConfig?.messages && promptConfig.messages.length > 0) {
-        // Use new message format
-        return promptConfig.messages.map(msg => ({
-            role: msg.role,
-            content: msg.content,
-        }));
-    }
-    else {
-        // Use legacy format
-        return [
-            {
-                role: 'system',
-                content: systemPrompt || 'You are a helpful assistant',
-            },
-            { role: 'user', content: prompt || '' },
-        ];
-    }
-}
-/**
- * Build response format object for API from prompt config
- */
-function buildResponseFormat(promptConfig) {
-    if (promptConfig?.responseFormat === 'json_schema' && promptConfig.jsonSchema) {
-        try {
-            const schema = JSON.parse(promptConfig.jsonSchema);
-            return {
-                type: 'json_schema',
-                json_schema: schema,
-            };
-        }
-        catch (error) {
-            throw new Error(`Invalid JSON schema: ${error instanceof Error ? error.message : 'Unknown error'}`);
-        }
-    }
-    return undefined;
-}
-/**
- * Build complete InferenceRequest from prompt config and inputs
- */
-function buildInferenceRequest(promptConfig, systemPrompt, prompt, modelName, temperature, topP, maxTokens, endpoint, token) {
-    const messages = buildMessages(promptConfig, systemPrompt, prompt);
-    const responseFormat = buildResponseFormat(promptConfig);
-    return {
-        messages,
-        modelName,
-        temperature,
-        topP,
-        maxTokens,
-        endpoint,
-        token,
-        responseFormat,
-    };
-}
-
 /*! js-yaml 4.1.1 https://github.com/nodeca/js-yaml @license MIT */
 function isNothing(subject) {
  return (typeof subject === 'undefined') || (subject === null);
@@ -61328,6 +61246,158 @@ var loader = {
 };
 var load                = loader.load;

+/**
+ * Helper function to load content from a file or use fallback input
+ * @param filePathInput - Input name for the file path
+ * @param contentInput - Input name for the direct content
+ * @param defaultValue - Default value to use if neither file nor content is provided
+ * @returns The loaded content
+ */
+function loadContentFromFileOrInput(filePathInput, contentInput, defaultValue) {
+    const filePath = coreExports.getInput(filePathInput);
+    const contentString = coreExports.getInput(contentInput);
+    if (filePath !== undefined && filePath !== '') {
+        if (!fs.existsSync(filePath)) {
+            throw new Error(`File for ${filePathInput} was not found: ${filePath}`);
+        }
+        return fs.readFileSync(filePath, 'utf-8');
+    }
+    else if (contentString !== undefined && contentString !== '') {
+        return contentString;
+    }
+    else if (defaultValue !== undefined) {
+        return defaultValue;
+    }
+    else {
+        throw new Error(`Neither ${filePathInput} nor ${contentInput} was set`);
+    }
+}
+/**
+ * Build messages array from either prompt config or legacy format
+ */
+function buildMessages(promptConfig, systemPrompt, prompt) {
+    if (promptConfig?.messages && promptConfig.messages.length > 0) {
+        // Use new message format
+        return promptConfig.messages.map(msg => ({
+            role: msg.role,
+            content: msg.content,
+        }));
+    }
+    else {
+        // Use legacy format
+        return [
+            {
+                role: 'system',
+                content: systemPrompt || 'You are a helpful assistant',
+            },
+            { role: 'user', content: prompt || '' },
+        ];
+    }
+}
+/**
+ * Build response format object for API from prompt config
+ */
+function buildResponseFormat(promptConfig) {
+    if (promptConfig?.responseFormat === 'json_schema' && promptConfig.jsonSchema) {
+        try {
+            const schema = JSON.parse(promptConfig.jsonSchema);
+            return {
+                type: 'json_schema',
+                json_schema: schema,
+            };
+        }
+        catch (error) {
+            throw new Error(`Invalid JSON schema: ${error instanceof Error ? error.message : 'Unknown error'}`);
+        }
+    }
+    return undefined;
+}
+/**
+ * Parse custom headers from YAML or JSON format
+ * @param input - String in YAML or JSON format containing headers
+ * @returns Record of header names to values, or empty object if invalid
+ */
+function parseCustomHeaders(input) {
+    if (!input || input.trim() === '') {
+        return {};
+    }
+    const trimmedInput = input.trim();
+    try {
+        // Try JSON first (check if it starts with { or [)
+        if (trimmedInput.startsWith('{') || trimmedInput.startsWith('[')) {
+            const parsed = JSON.parse(trimmedInput);
+            if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) {
+                coreExports.warning('Custom headers JSON must be an object, not null or an array');
+                return {};
+            }
+            return validateAndMaskHeaders(parsed);
+        }
+        // Try YAML
+        const parsed = load(trimmedInput);
+        if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) {
+            coreExports.warning('Custom headers YAML must be an object');
+            return {};
+        }
+        return validateAndMaskHeaders(parsed);
+    }
+    catch (error) {
+        coreExports.warning(`Failed to parse custom headers: ${error instanceof Error ? error.message : 'Unknown error'}`);
+        return {};
+    }
+}
+/**
+ * Validate header names and mask sensitive values in logs
+ * @param headers - Raw headers object
+ * @returns Validated headers with string values
+ */
+function validateAndMaskHeaders(headers) {
+    const validHeaders = {};
+    const sensitivePatterns = ['key', 'token', 'secret', 'password', 'authorization'];
+    for (const [name, value] of Object.entries(headers)) {
+        // Validate header name (basic HTTP header name validation, RFC 7230: letters, digits, and hyphens)
+        if (!/^[A-Za-z0-9-]+$/.test(name)) {
+            coreExports.warning(`Skipping invalid header name: ${name} (only alphanumeric characters and hyphens allowed)`);
+            continue;
+        }
+        // Convert value to string
+        const stringValue = String(value);
+        // Validate header value to prevent CRLF/header injection
+        if (stringValue.includes('\r') || stringValue.includes('\n')) {
+            coreExports.warning(`Skipping header "${name}" because its value contains newline characters, which are not allowed in HTTP header values.`);
+            continue;
+        }
+        validHeaders[name] = stringValue;
+        // Mask sensitive headers in logs
+        const lowerName = name.toLowerCase();
+        const isSensitive = sensitivePatterns.some(pattern => lowerName.includes(pattern));
+        if (isSensitive) {
+            coreExports.info(`Custom header added: ${name}: ***MASKED***`);
+        }
+        else {
+            coreExports.info(`Custom header added: ${name}: ${stringValue}`);
+        }
+    }
+    return validHeaders;
+}
+/**
+ * Build complete InferenceRequest from prompt config and inputs
+ */
+function buildInferenceRequest(promptConfig, systemPrompt, prompt, modelName, temperature, topP, maxTokens, endpoint, token, customHeaders) {
+    const messages = buildMessages(promptConfig, systemPrompt, prompt);
+    const responseFormat = buildResponseFormat(promptConfig);
+    return {
+        messages,
+        modelName,
+        temperature,
+        topP,
+        maxTokens,
+        endpoint,
+        token,
+        responseFormat,
+        customHeaders,
+    };
+}
+
 /**
 * Parse template variables from YAML input string
 */
@@ -61478,8 +61548,11 @@ async function run() {
        const githubMcpToken = coreExports.getInput('github-mcp-token') || token;
        const githubMcpToolsets = coreExports.getInput('github-mcp-toolsets');
        const endpoint = coreExports.getInput('endpoint');
+        // Parse custom headers
+        const customHeadersInput = coreExports.getInput('custom-headers');
+        const customHeaders = parseCustomHeaders(customHeadersInput);
        // Build the inference request with pre-processed messages and response format
-        const inferenceRequest = buildInferenceRequest(promptConfig, systemPrompt, prompt, modelName, promptConfig?.modelParameters?.temperature, promptConfig?.modelParameters?.topP, maxTokens, endpoint, token);
+        const inferenceRequest = buildInferenceRequest(promptConfig, systemPrompt, prompt, modelName, promptConfig?.modelParameters?.temperature, promptConfig?.modelParameters?.topP, maxTokens, endpoint, token, customHeaders);
        const enableMcp = coreExports.getBooleanInput('enable-github-mcp') || false;
        let modelResponse = null;
        if (enableMcp) {
@@ -1,5 +1,6 @@
 import * as core from '@actions/core'
 import * as fs from 'fs'
+import * as yaml from 'js-yaml'
 import {PromptConfig} from './prompt.js'
 import {InferenceRequest} from './inference.js'

@@ -74,6 +75,83 @@ export function buildResponseFormat(
  return undefined
 }

+/**
+ * Parse custom headers from YAML or JSON format
+ * @param input - String in YAML or JSON format containing headers
+ * @returns Record of header names to values, or empty object if invalid
+ */
+export function parseCustomHeaders(input: string): Record<string, string> {
+  if (!input || input.trim() === '') {
+    return {}
+  }
+
+  const trimmedInput = input.trim()
+
+  try {
+    // Try JSON first (check if it starts with { or [)
+    if (trimmedInput.startsWith('{') || trimmedInput.startsWith('[')) {
+      const parsed = JSON.parse(trimmedInput)
+      if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) {
+        core.warning('Custom headers JSON must be an object, not null or an array')
+        return {}
+      }
+      return validateAndMaskHeaders(parsed as Record<string, unknown>)
+    }
+
+    // Try YAML
+    const parsed = yaml.load(trimmedInput)
+    if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) {
+      core.warning('Custom headers YAML must be an object')
+      return {}
+    }
+    return validateAndMaskHeaders(parsed as Record<string, unknown>)
+  } catch (error) {
+    core.warning(`Failed to parse custom headers: ${error instanceof Error ? error.message : 'Unknown error'}`)
+    return {}
+  }
+}
+
+/**
+ * Validate header names and mask sensitive values in logs
+ * @param headers - Raw headers object
+ * @returns Validated headers with string values
+ */
+function validateAndMaskHeaders(headers: Record<string, unknown>): Record<string, string> {
+  const validHeaders: Record<string, string> = {}
+  const sensitivePatterns = ['key', 'token', 'secret', 'password', 'authorization']
+
+  for (const [name, value] of Object.entries(headers)) {
+    // Validate header name (basic HTTP header name validation, RFC 7230: letters, digits, and hyphens)
+    if (!/^[A-Za-z0-9-]+$/.test(name)) {
+      core.warning(`Skipping invalid header name: ${name} (only alphanumeric characters and hyphens allowed)`)
+      continue
+    }
+
+    // Convert value to string
+    const stringValue = String(value)
+
+    // Validate header value to prevent CRLF/header injection
+    if (stringValue.includes('\r') || stringValue.includes('\n')) {
+      core.warning(
+        `Skipping header "${name}" because its value contains newline characters, which are not allowed in HTTP header values.`,
+      )
+      continue
+    }
+    validHeaders[name] = stringValue
+
+    // Mask sensitive headers in logs
+    const lowerName = name.toLowerCase()
+    const isSensitive = sensitivePatterns.some(pattern => lowerName.includes(pattern))
+    if (isSensitive) {
+      core.info(`Custom header added: ${name}: ***MASKED***`)
+    } else {
+      core.info(`Custom header added: ${name}: ${stringValue}`)
+    }
+  }
+
+  return validHeaders
+}
+
 /**
 * Build complete InferenceRequest from prompt config and inputs
 */
@@ -87,6 +165,7 @@ export function buildInferenceRequest(
  maxTokens: number,
  endpoint: string,
  token: string,
+  customHeaders?: Record<string, string>,
 ): InferenceRequest {
  const messages = buildMessages(promptConfig, systemPrompt, prompt)
  const responseFormat = buildResponseFormat(promptConfig)
@@ -100,5 +179,6 @@ export function buildInferenceRequest(
    endpoint,
    token,
    responseFormat,
+    customHeaders,
  }
 }
@@ -18,6 +18,7 @@ export interface InferenceRequest {
  temperature?: number
  topP?: number
  responseFormat?: {type: 'json_schema'; json_schema: unknown} // Processed response format for the API
+  customHeaders?: Record<string, string> // Custom HTTP headers to include in API requests
 }

 export interface InferenceResponse {
@@ -41,6 +42,7 @@ export async function simpleInference(request: InferenceRequest): Promise<string
  const client = new OpenAI({
    apiKey: request.token,
    baseURL: request.endpoint,
+    defaultHeaders: request.customHeaders || {},
  })

  const chatCompletionRequest: OpenAI.Chat.Completions.ChatCompletionCreateParams = {
@@ -75,6 +77,7 @@ export async function mcpInference(
  const client = new OpenAI({
    apiKey: request.token,
    baseURL: request.endpoint,
+    defaultHeaders: request.customHeaders || {},
  })

  // Start with the pre-processed messages
@@ -3,7 +3,7 @@ import * as fs from 'fs'
 import * as tmp from 'tmp'
 import {connectToGitHubMCP} from './mcp.js'
 import {simpleInference, mcpInference} from './inference.js'
-import {loadContentFromFileOrInput, buildInferenceRequest} from './helpers.js'
+import {loadContentFromFileOrInput, buildInferenceRequest, parseCustomHeaders} from './helpers.js'
 import {
  loadPromptFile,
  parseTemplateVariables,
@@ -65,6 +65,10 @@ export async function run(): Promise<void> {

    const endpoint = core.getInput('endpoint')

+    // Parse custom headers
+    const customHeadersInput = core.getInput('custom-headers')
+    const customHeaders = parseCustomHeaders(customHeadersInput)
+
    // Build the inference request with pre-processed messages and response format
    const inferenceRequest = buildInferenceRequest(
      promptConfig,
@@ -76,6 +80,7 @@ export async function run(): Promise<void> {
      maxTokens,
      endpoint,
      token,
+      customHeaders,
    )

    const enableMcp = core.getBooleanInput('enable-github-mcp') || false