Skip to content

Add support for custom session/client injection for proxy authentication#156

Open
Barneyjm wants to merge 4 commits intotavily-ai:masterfrom
Barneyjm:custom-session
Open

Add support for custom session/client injection for proxy authentication#156
Barneyjm wants to merge 4 commits intotavily-ai:masterfrom
Barneyjm:custom-session

Conversation

@Barneyjm
Copy link

@Barneyjm Barneyjm commented Feb 25, 2026

Summary

Adds the ability to pass a custom requests.Session (sync) or httpx.AsyncClient (async) to Tavily clients, enabling integration with authentication proxies without distributing API keys to each application team.

Motivation

When using Tavily behind an authentication proxy (e.g., client → some API Management → Tavily API), organizations need to:

  • Centralize API key management at the proxy level
  • Leverage existing proxy authentication mechanisms (subscription keys, OAuth, etc.)
  • Avoid distributing Tavily API keys to individual application teams

Changes

  • TavilyClient: Added optional session parameter to accept a pre-configured requests.Session
  • AsyncTavilyClient: Added optional client parameter to accept a pre-configured httpx.AsyncClient
  • Both implementations properly merge custom headers and proxy settings
  • Maintains backward compatibility - creates default session/client if none provided

Usage Example

import requests
from tavily import TavilyClient

# Create custom session with proxy auth
custom_session = requests.Session()
custom_session.headers.update({
    "custom-header": "your-key"
})

# Use with Tavily client
client = TavilyClient(
    api_key="dummy-key",  # Can be placeholder when using proxy auth
    api_base_url="https://your-api.net/tavily",
    session=custom_session
)

response = client.search("query")

<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> **Medium Risk**
> Changes how `TavilyClient` creates/manages HTTP sessions and merges headers/proxies, which can affect request authentication/routing for existing integrations. Behavior is covered by new tests but could surface subtle precedence/regression issues in real network environments.
> 
> **Overview**
> **Adds support for injecting a custom `requests.Session` into `TavilyClient`.** The client now merges default Tavily headers and proxy settings into the provided session *without overriding* any existing session `headers` (including `Authorization`) or `proxies`, and tracks whether the session is external.
> 
> `close()`/context manager cleanup is updated to only close internally-created sessions, and `self.headers` is now kept as Tavily-specific defaults rather than mirroring the full session headers. A new `tests/test_custom_session.py` suite exercises header/proxy precedence and lifecycle behavior.
> 
> <sup>Written by [Cursor Bugbot](https://cursor.com/dashboard?tab=bugbot) for commit 0a86fd8dcbd1303d1c2fc764c73470ee2176d6ef. This will update automatically on new commits. Configure [here](https://cursor.com/dashboard?tab=bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

- Add optional session parameter to __init__
- Preserve custom headers from provided session (e.g., Authorization)
- Allows integration with enterprise API gateways requiring custom auth
- Add tests for custom session functionality
- Track whether session is external with _external_session flag
- Only close internally-created sessions in close() method
- Prevents closing shared sessions used across multiple clients
- Add comprehensive tests for session lifecycle management
- Change proxy handling from unconditional .update() to conditional merge
- Only add proxy for protocol if not already in session.proxies
- Matches existing header preservation behavior for consistency
- Custom session proxies now take precedence over:
  - proxies parameter
  - TAVILY_*_PROXY environment variables
- Add 4 comprehensive tests for proxy preservation scenarios
Copy link

@cursor cursor bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes and found 1 potential issue.

Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

@Barneyjm
Copy link
Author

Barneyjm commented Mar 2, 2026

this PR also allows for the functionality requested in #128 and allows more flexibility on top.

- self.headers previously leaked session defaults (User-Agent, Accept-Encoding, etc.)
- For custom sessions, also leaked user-defined headers
- Now self.headers contains only Tavily-specific headers (Content-Type, Authorization, X-Client-Source, X-Project-ID)
- This restores the original public API contract for self.headers attribute
- Add comprehensive test to verify headers attribute semantics
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant