⚡️ Speed up method TimeoutInterceptor.intercept_unary_unary by 90%

[codeflash-ai]

📄 90% (0.90x) speedup for TimeoutInterceptor.intercept_unary_unary in src/xai_sdk/client.py

⏱️ Runtime : 1.76 milliseconds → 926 microseconds (best of 220 runs)

📝 Explanation and details

The optimization introduces a memoization cache to avoid repeatedly calling the expensive _replace() method on the same client_call_details objects.

Key changes:

• Added self._timeout_details_cache dictionary to cache modified client call details
• Cache key uses (id(client_call_details), self.timeout) to uniquely identify combinations
• Only calls _replace() when encountering a new client_call_details object

Why this speeds up the code:
The line profiler shows that client_call_details._replace(timeout=self.timeout) was consuming 79.2% of execution time (4.37ms out of 5.51ms). The _replace() method on namedtuples is expensive because it creates a new object and copies all fields.

In the optimized version, _replace() is only called 45 times instead of 2043 times, reducing its impact from 79.2% to just 5.7% of total execution time. The cache lookup operations are much faster than object creation.

Test case performance patterns:

• Single requests: 20-28% slower due to cache overhead with no reuse benefit
• Bulk/repeated requests: Up to 119% faster when the same client_call_details objects are reused (as shown in test_intercept_unary_unary_stress_many_requests and test_intercept_large_number_of_requests)

This optimization is particularly effective for high-throughput gRPC applications where the same client call details objects are frequently reused across multiple requests.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 2088 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

from collections import namedtuple

import grpc
# imports
import pytest  # used for our unit tests
from xai_sdk.client import TimeoutInterceptor

# Helper: a mock client_call_details object with _replace
ClientCallDetails = namedtuple(
    "ClientCallDetails",
    ["timeout", "method", "metadata", "credentials", "wait_for_ready", "compression"]
)

def make_client_call_details(timeout=None):
    """Helper to create a ClientCallDetails-like object."""
    return ClientCallDetails(
        timeout=timeout,
        method="TestMethod",
        metadata=(("key", "value"),),
        credentials=None,
        wait_for_ready=False,
        compression=None,
    )

# Helper: a mock continuation function
def mock_continuation(expected_timeout):
    """Returns a continuation that asserts the timeout is as expected, and returns a marker value."""
    def continuation(client_call_details, request):
        return (client_call_details, request)
    return continuation

# ---- UNIT TESTS ----

# 1. BASIC TEST CASES

def test_intercept_unary_unary_sets_timeout_basic():
    """Test that intercept_unary_unary sets the timeout correctly for a normal request."""
    interceptor = TimeoutInterceptor(timeout=10.0)
    client_call_details = make_client_call_details(timeout=None)
    request = "test-request"
    continuation = mock_continuation(expected_timeout=10.0)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 3.29μs -> 4.38μs (25.0% slower)

def test_intercept_unary_unary_overwrites_existing_timeout():
    """Test that an existing timeout is overwritten by the interceptor."""
    interceptor = TimeoutInterceptor(timeout=5.0)
    client_call_details = make_client_call_details(timeout=20.0)
    request = "another-request"
    continuation = mock_continuation(expected_timeout=5.0)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 3.09μs -> 3.92μs (21.2% slower)

def test_intercept_unary_unary_with_zero_timeout():
    """Test that a zero timeout is set correctly."""
    interceptor = TimeoutInterceptor(timeout=0.0)
    client_call_details = make_client_call_details(timeout=None)
    request = "zero-request"
    continuation = mock_continuation(expected_timeout=0.0)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.98μs -> 3.85μs (22.7% slower)

def test_intercept_unary_unary_with_negative_timeout():
    """Test that a negative timeout is set correctly (even if not recommended)."""
    interceptor = TimeoutInterceptor(timeout=-1.0)
    client_call_details = make_client_call_details(timeout=100.0)
    request = "negative-request"
    continuation = mock_continuation(expected_timeout=-1.0)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.79μs -> 3.89μs (28.3% slower)

def test_intercept_unary_unary_with_float_timeout():
    """Test that a float timeout is set correctly."""
    interceptor = TimeoutInterceptor(timeout=2.345)
    client_call_details = make_client_call_details(timeout=None)
    request = "float-request"
    continuation = mock_continuation(expected_timeout=2.345)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.90μs -> 3.67μs (20.9% slower)

# 2. EDGE TEST CASES

def test_intercept_unary_unary_with_large_timeout():
    """Test that a very large timeout is set correctly."""
    large_timeout = 1e9
    interceptor = TimeoutInterceptor(timeout=large_timeout)
    client_call_details = make_client_call_details(timeout=None)
    request = "large-timeout"
    continuation = mock_continuation(expected_timeout=large_timeout)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.85μs -> 3.97μs (28.1% slower)

def test_intercept_unary_unary_with_none_timeout():
    """Test that None timeout in input is replaced by interceptor's timeout."""
    interceptor = TimeoutInterceptor(timeout=3.5)
    client_call_details = make_client_call_details(timeout=None)
    request = "none-timeout"
    continuation = mock_continuation(expected_timeout=3.5)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.79μs -> 3.86μs (27.6% slower)

def test_intercept_unary_unary_preserves_other_fields():
    """Test that all fields except timeout are preserved after interception."""
    interceptor = TimeoutInterceptor(timeout=7.0)
    client_call_details = make_client_call_details(timeout=1.0)
    request = "preserve-fields"
    continuation = mock_continuation(expected_timeout=7.0)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.77μs -> 3.59μs (22.9% slower)
    # All fields except timeout should be unchanged
    for field in ClientCallDetails._fields:
        if field == "timeout":
            pass
        else:
            pass

def test_intercept_unary_unary_with_empty_metadata():
    """Test that empty metadata is preserved."""
    interceptor = TimeoutInterceptor(timeout=8.0)
    # Create a client_call_details with empty metadata
    client_call_details = ClientCallDetails(
        timeout=0.5,
        method="EmptyMetadataMethod",
        metadata=(),
        credentials=None,
        wait_for_ready=True,
        compression=None,
    )
    request = "empty-metadata"
    continuation = mock_continuation(expected_timeout=8.0)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.72μs -> 3.65μs (25.5% slower)

def test_intercept_unary_unary_with_non_string_request():
    """Test that non-string request objects are preserved."""
    interceptor = TimeoutInterceptor(timeout=2.0)
    client_call_details = make_client_call_details(timeout=0.1)
    request = {"foo": "bar", "baz": [1, 2, 3]}
    continuation = mock_continuation(expected_timeout=2.0)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.74μs -> 3.43μs (20.2% slower)

def test_intercept_unary_unary_with_long_method_name():
    """Test that long method names are preserved."""
    interceptor = TimeoutInterceptor(timeout=1.1)
    long_method = "x" * 256
    client_call_details = ClientCallDetails(
        timeout=0.0,
        method=long_method,
        metadata=(("a", "b"),),
        credentials=None,
        wait_for_ready=False,
        compression=None,
    )
    request = "long-method"
    continuation = mock_continuation(expected_timeout=1.1)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.72μs -> 3.59μs (24.2% slower)

def test_intercept_unary_unary_with_custom_compression():
    """Test that custom compression is preserved."""
    interceptor = TimeoutInterceptor(timeout=6.0)
    client_call_details = ClientCallDetails(
        timeout=0.0,
        method="CompressionMethod",
        metadata=(("x", "y"),),
        credentials=None,
        wait_for_ready=True,
        compression="gzip",
    )
    request = "compression-request"
    continuation = mock_continuation(expected_timeout=6.0)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.62μs -> 3.57μs (26.7% slower)

def test_intercept_unary_unary_with_credentials_object():
    """Test that credentials object is preserved."""
    class DummyCreds:
        pass
    creds = DummyCreds()
    interceptor = TimeoutInterceptor(timeout=4.0)
    client_call_details = ClientCallDetails(
        timeout=0.0,
        method="CredsMethod",
        metadata=(("x", "y"),),
        credentials=creds,
        wait_for_ready=True,
        compression=None,
    )
    request = "creds-request"
    continuation = mock_continuation(expected_timeout=4.0)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.73μs -> 3.60μs (24.0% slower)

def test_intercept_unary_unary_with_wait_for_ready_true_false():
    """Test that wait_for_ready is preserved (True and False)."""
    for value in [True, False]:
        interceptor = TimeoutInterceptor(timeout=9.0)
        client_call_details = ClientCallDetails(
            timeout=0.0,
            method="WaitForReadyMethod",
            metadata=(("x", "y"),),
            credentials=None,
            wait_for_ready=value,
            compression=None,
        )
        request = "wait-for-ready"
        continuation = mock_continuation(expected_timeout=9.0)
        codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 4.24μs -> 5.49μs (22.7% slower)

# 3. LARGE SCALE TEST CASES

def test_intercept_unary_unary_large_metadata():
    """Test that large metadata is preserved and timeout is set."""
    large_metadata = tuple((f"key{i}", f"value{i}") for i in range(1000))
    interceptor = TimeoutInterceptor(timeout=12.0)
    client_call_details = ClientCallDetails(
        timeout=0.0,
        method="LargeMetadataMethod",
        metadata=large_metadata,
        credentials=None,
        wait_for_ready=True,
        compression=None,
    )
    request = "large-metadata"
    continuation = mock_continuation(expected_timeout=12.0)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.84μs -> 3.84μs (26.0% slower)

def test_intercept_unary_unary_large_request_object():
    """Test that a large request object is preserved."""
    large_request = {"data": [i for i in range(1000)]}
    interceptor = TimeoutInterceptor(timeout=15.0)
    client_call_details = make_client_call_details(timeout=0.0)
    continuation = mock_continuation(expected_timeout=15.0)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, large_request); result = codeflash_output # 2.74μs -> 3.52μs (22.2% slower)

def test_intercept_unary_unary_multiple_calls_consistency():
    """Test that multiple calls with different timeouts are handled correctly."""
    for i in range(10):
        timeout = float(i)
        interceptor = TimeoutInterceptor(timeout=timeout)
        client_call_details = make_client_call_details(timeout=100.0)
        request = f"request-{i}"
        continuation = mock_continuation(expected_timeout=timeout)
        codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 11.4μs -> 14.1μs (19.7% slower)

def test_intercept_unary_unary_stress_many_metadata_entries():
    """Test with maximum allowed metadata entries."""
    max_entries = 1000
    metadata = tuple((str(i), str(i)) for i in range(max_entries))
    interceptor = TimeoutInterceptor(timeout=20.0)
    client_call_details = ClientCallDetails(
        timeout=0.0,
        method="StressMethod",
        metadata=metadata,
        credentials=None,
        wait_for_ready=True,
        compression=None,
    )
    request = "stress-request"
    continuation = mock_continuation(expected_timeout=20.0)
    codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 2.78μs -> 3.65μs (23.8% slower)

def test_intercept_unary_unary_stress_many_requests():
    """Test interceptor with many sequential requests to check for state leakage."""
    interceptor = TimeoutInterceptor(timeout=21.0)
    client_call_details = make_client_call_details(timeout=0.0)
    for i in range(1000):
        request = f"req-{i}"
        continuation = mock_continuation(expected_timeout=21.0)
        codeflash_output = interceptor.intercept_unary_unary(continuation, client_call_details, request); result = codeflash_output # 858μs -> 392μs (119% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from collections import namedtuple

import grpc
# imports
import pytest
from xai_sdk.client import TimeoutInterceptor

# Helper: client_call_details mock using namedtuple
ClientCallDetails = namedtuple('ClientCallDetails', ['timeout', 'other_field'])

# Helper: continuation function that returns its inputs for inspection
def continuation_return(client_call_details, request):
    # Returns a tuple for inspection in tests
    return (client_call_details, request)

# -------------------------
# Unit tests start here
# -------------------------

# 1. Basic Test Cases

def test_intercept_sets_timeout_basic():
    """Test that the interceptor sets the timeout correctly for a simple request."""
    interceptor = TimeoutInterceptor(timeout=10.0)
    details = ClientCallDetails(timeout=None, other_field="foo")
    request = "bar"
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, request) # 4.00μs -> 5.18μs (22.8% slower)

def test_intercept_overwrites_existing_timeout():
    """Test that the interceptor overwrites an existing timeout value."""
    interceptor = TimeoutInterceptor(timeout=5.5)
    details = ClientCallDetails(timeout=99.9, other_field="baz")
    request = "request_data"
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, request) # 3.18μs -> 4.04μs (21.3% slower)

def test_intercept_with_zero_timeout():
    """Test that a zero timeout is set correctly."""
    interceptor = TimeoutInterceptor(timeout=0.0)
    details = ClientCallDetails(timeout=None, other_field="zero")
    request = "req"
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, request) # 3.02μs -> 3.94μs (23.1% slower)

def test_intercept_with_negative_timeout():
    """Test that a negative timeout is set correctly (even if not recommended)."""
    interceptor = TimeoutInterceptor(timeout=-1.0)
    details = ClientCallDetails(timeout=3.0, other_field="neg")
    request = "negative"
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, request) # 2.97μs -> 3.80μs (22.0% slower)

# 2. Edge Test Cases

def test_intercept_with_large_timeout():
    """Test that a very large timeout value is set correctly."""
    large_timeout = 1e9
    interceptor = TimeoutInterceptor(timeout=large_timeout)
    details = ClientCallDetails(timeout=0.1, other_field="large")
    request = "data"
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, request) # 2.90μs -> 3.76μs (22.9% slower)

def test_intercept_with_none_fields():
    """Test that fields other than timeout can be None and are preserved."""
    interceptor = TimeoutInterceptor(timeout=2.0)
    details = ClientCallDetails(timeout=None, other_field=None)
    request = None
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, request) # 2.92μs -> 3.70μs (21.3% slower)

def test_intercept_with_float_precision():
    """Test that float precision is preserved when setting timeout."""
    interceptor = TimeoutInterceptor(timeout=0.123456789)
    details = ClientCallDetails(timeout=0.987654321, other_field="precise")
    request = "precise_request"
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, request) # 2.76μs -> 3.61μs (23.5% slower)

def test_intercept_with_multiple_fields():
    """Test that additional fields in client_call_details are preserved."""
    ExtendedClientCallDetails = namedtuple('ExtendedClientCallDetails', ['timeout', 'field1', 'field2'])
    interceptor = TimeoutInterceptor(timeout=42.0)
    details = ExtendedClientCallDetails(timeout=0.1, field1="a", field2="b")
    request = "extended"
    # We need a continuation that works with ExtendedClientCallDetails
    def continuation_ext(client_call_details, request):
        return (client_call_details, request)
    result_details, result_request = interceptor.intercept_unary_unary(continuation_ext, details, request) # 3.04μs -> 4.00μs (24.1% slower)

def test_intercept_with_non_string_request():
    """Test that non-string request objects are passed through unchanged."""
    interceptor = TimeoutInterceptor(timeout=7.0)
    details = ClientCallDetails(timeout=1.0, other_field="obj")
    request = {"key": [1, 2, 3], "val": 99}
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, request) # 3.17μs -> 3.89μs (18.4% slower)

def test_intercept_with_tuple_request():
    """Test that tuple requests are handled correctly."""
    interceptor = TimeoutInterceptor(timeout=3.14)
    details = ClientCallDetails(timeout=2.71, other_field="tuple")
    request = (1, 2, 3)
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, request) # 2.90μs -> 3.51μs (17.4% slower)

def test_intercept_with_empty_string_request():
    """Test that empty string requests are handled correctly."""
    interceptor = TimeoutInterceptor(timeout=1.0)
    details = ClientCallDetails(timeout=0.0, other_field="empty")
    request = ""
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, request) # 2.80μs -> 3.57μs (21.6% slower)

# 3. Large Scale Test Cases

def test_intercept_large_number_of_requests():
    """Test interceptor with a large number of different requests."""
    interceptor = TimeoutInterceptor(timeout=123.456)
    details = ClientCallDetails(timeout=0.0, other_field="bulk")
    # Generate 1000 unique requests
    for i in range(1000):
        req = f"request_{i}"
        result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, req) # 795μs -> 391μs (103% faster)

def test_intercept_large_client_call_details():
    """Test interceptor with client_call_details containing many fields."""
    # Create a namedtuple with 100 fields, including 'timeout'
    fields = ['timeout'] + [f'f{i}' for i in range(99)]
    LargeClientCallDetails = namedtuple('LargeClientCallDetails', fields)
    # Set initial values
    values = [0.0] + [i for i in range(99)]
    details = LargeClientCallDetails(*values)
    interceptor = TimeoutInterceptor(timeout=999.999)
    request = "large_fields"
    def continuation_large(client_call_details, request):
        return (client_call_details, request)
    result_details, result_request = interceptor.intercept_unary_unary(continuation_large, details, request) # 5.74μs -> 7.27μs (21.0% slower)
    for i in range(1, 100):
        pass

def test_intercept_with_large_request_object():
    """Test that a large request object (e.g., large list) is passed through unchanged."""
    interceptor = TimeoutInterceptor(timeout=111.0)
    details = ClientCallDetails(timeout=222.0, other_field="large_req")
    large_request = list(range(1000))
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, large_request) # 3.14μs -> 4.24μs (25.9% slower)

def test_intercept_with_large_timeout_precision():
    """Test that a large timeout with high precision is set correctly."""
    interceptor = TimeoutInterceptor(timeout=123456.7890123456)
    details = ClientCallDetails(timeout=0.0, other_field="precise_large")
    request = "precision_test"
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, request) # 2.92μs -> 3.92μs (25.6% slower)

# Edge: Ensure mutation testing fails if _replace is not used
def test_intercept_does_not_mutate_original_details():
    """Test that the original client_call_details is not mutated (immutability)."""
    interceptor = TimeoutInterceptor(timeout=77.7)
    details = ClientCallDetails(timeout=1.1, other_field="immutable")
    request = "immut"
    # Save original details
    original_details = details
    result_details, result_request = interceptor.intercept_unary_unary(continuation_return, details, request) # 2.77μs -> 3.77μs (26.6% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-TimeoutInterceptor.intercept_unary_unary-mgu10u33 and push.