276 lines
13 KiB
Python
276 lines
13 KiB
Python
import http.server
|
|
import json
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
import urllib.error
|
|
import urllib.request
|
|
from unittest.mock import Mock, patch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
from codex_http_transport import ModelTransport, RequestCancelled
|
|
|
|
|
|
class TransportTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.requests = []
|
|
self.responses = []
|
|
case = self
|
|
|
|
class Upstream(http.server.BaseHTTPRequestHandler):
|
|
def log_message(self, *_):
|
|
pass
|
|
|
|
def do_POST(self):
|
|
body = self.rfile.read(int(self.headers['Content-Length']))
|
|
case.requests.append((self.path, body, dict(self.headers)))
|
|
scripted = case.responses.pop(0)
|
|
if callable(scripted):
|
|
scripted(self)
|
|
return
|
|
status, headers, response = scripted
|
|
self.send_response(status)
|
|
for key, value in headers.items():
|
|
self.send_header(key, value)
|
|
self.end_headers()
|
|
self.wfile.write(response)
|
|
|
|
self.server = http.server.ThreadingHTTPServer(('127.0.0.1', 0), Upstream)
|
|
self.server.daemon_threads = True
|
|
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
|
|
self.thread.start()
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.log = Path(self.temp.name) / 'events.jsonl'
|
|
self.transport = ModelTransport('OpenAI', f'http://127.0.0.1:{self.server.server_port}/v1',
|
|
timeout=2, retries=2, retry_delay=0, log_path=self.log)
|
|
|
|
def tearDown(self):
|
|
self.transport.close()
|
|
self.server.shutdown()
|
|
self.server.server_close()
|
|
self.thread.join()
|
|
self.temp.cleanup()
|
|
|
|
def send(self, body=b'{"private_prompt":"DO_NOT_LOG"}', path='/responses'):
|
|
request = urllib.request.Request(self.transport.url + path, data=body,
|
|
headers={'Authorization': 'Bearer DO_NOT_LOG_KEY', 'Content-Type': 'application/json'})
|
|
return urllib.request.urlopen(request, timeout=5)
|
|
|
|
def test_payload_auth_and_stream_preserved_without_sensitive_logging(self):
|
|
response = b'data: {"type":"response.created"}\n\ndata: {"type":"response.completed"}\n\n'
|
|
self.responses.append((200, {'Content-Type': 'text/event-stream'}, response))
|
|
body = b'{"input":[{"type":"additional_tools","tools":[{"type":"namespace"}]}],"private":"DO_NOT_LOG"}'
|
|
with self.send(body) as result:
|
|
self.assertEqual(result.read(), response)
|
|
path, received, headers = self.requests[0]
|
|
self.assertEqual(path, '/v1/responses')
|
|
self.assertEqual(received, body)
|
|
self.assertEqual(headers['Authorization'], 'Bearer DO_NOT_LOG_KEY')
|
|
self.assertEqual(headers['Accept-Encoding'], 'identity')
|
|
log = self.log.read_text()
|
|
self.assertNotIn('DO_NOT_LOG', log)
|
|
self.assertIn('response.completed', log)
|
|
|
|
def test_transient_failure_retries_only_before_stream_headers(self):
|
|
self.responses.extend([(503, {}, b'unavailable'), (200, {'Content-Type': 'text/event-stream'}, b'data: {"type":"response.completed"}\n\n')])
|
|
with self.send() as response:
|
|
response.read()
|
|
self.assertEqual(len(self.requests), 2)
|
|
self.assertEqual(self.requests[0][1], self.requests[1][1])
|
|
self.assertIn('"retry"', self.log.read_text())
|
|
|
|
def test_header_timeout_retries_the_same_request_with_a_bound(self):
|
|
request = urllib.request.Request(self.transport.upstream + '/responses', data=b'body')
|
|
response = object()
|
|
opener = Mock()
|
|
opener.open.side_effect = [TimeoutError(), response]
|
|
with patch('codex_http_transport.urllib.request.build_opener', return_value=opener):
|
|
self.assertIs(self.transport.open_request(request, 1), response)
|
|
self.assertEqual(opener.open.call_count, 2)
|
|
self.assertTrue(all(call.args[0] is request for call in opener.open.call_args_list))
|
|
opener.open.side_effect = TimeoutError()
|
|
opener.open.reset_mock()
|
|
with patch('codex_http_transport.urllib.request.build_opener', return_value=opener), self.assertRaises(TimeoutError):
|
|
self.transport.open_request(request, 2)
|
|
self.assertEqual(opener.open.call_count, 3)
|
|
|
|
def test_body_timeout_does_not_repeat_a_delivered_tool_call(self):
|
|
partial = b'data: {"type":"response.output_item.done","item":{"type":"function_call"}}\n\n'
|
|
class Response:
|
|
status = 200
|
|
headers = {'Content-Type': 'text/event-stream'}
|
|
def __enter__(self):
|
|
return self
|
|
def __exit__(self, *_):
|
|
pass
|
|
def __iter__(self):
|
|
yield partial
|
|
raise TimeoutError()
|
|
with patch.object(self.transport, 'open_request', return_value=Response()) as opener:
|
|
with self.send() as result:
|
|
self.assertEqual(result.read(), partial)
|
|
self.assertEqual(opener.call_count, 1)
|
|
self.assertNotIn('response.completed', self.log.read_text())
|
|
self.assertIn('"afterHeaders": true', self.log.read_text())
|
|
|
|
def test_partial_stream_is_not_retried_or_given_fake_completion(self):
|
|
partial = b'data: {"type":"response.output_item.done","item":{"type":"function_call"}}\n\n'
|
|
self.responses.append((200, {'Content-Type': 'text/event-stream'}, partial))
|
|
with self.send() as response:
|
|
self.assertEqual(response.read(), partial)
|
|
self.assertEqual(len(self.requests), 1)
|
|
self.assertNotIn('response.completed', self.log.read_text())
|
|
|
|
def test_auth_rate_limit_and_redirect_errors_are_not_retried(self):
|
|
for status, headers in [(401, {}), (429, {'Retry-After': '60'}), (307, {'Location': 'https://other.example/responses'})]:
|
|
self.responses.append((status, headers, b'{"error":"rejected"}'))
|
|
count = len(self.requests)
|
|
with self.assertRaises(urllib.error.HTTPError) as error:
|
|
self.send()
|
|
self.assertEqual(error.exception.code, status)
|
|
self.assertEqual(len(self.requests), count + 1)
|
|
|
|
def test_compaction_and_retry_cap(self):
|
|
self.responses.append((200, {'Content-Type': 'application/json'}, b'{"output":[]}'))
|
|
with self.send(path='/responses/compact') as response:
|
|
self.assertEqual(response.read(), b'{"output":[]}')
|
|
self.assertEqual(self.requests[0][0], '/v1/responses/compact')
|
|
self.responses.extend([(503, {}, b'error')] * 3)
|
|
with self.assertRaises(urllib.error.HTTPError):
|
|
self.send()
|
|
self.assertEqual(len(self.requests), 4)
|
|
|
|
def test_rejects_embedded_credentials_and_model_is_not_overridden(self):
|
|
with self.assertRaises(ValueError):
|
|
ModelTransport('OpenAI', 'https://user:secret@example.test')
|
|
self.assertFalse(any(value.startswith('model=') for value in self.transport.cli_args()))
|
|
|
|
def test_per_request_usage_is_scoped_without_logging_response_content(self):
|
|
checkpoint = self.transport.usage_checkpoint()
|
|
for tokens in (80000, 90000):
|
|
body = ('data: ' + json.dumps({'type': 'response.completed', 'response': {
|
|
'usage': {'input_tokens': tokens}, 'output': 'DO_NOT_LOG'}}) + '\n\n').encode()
|
|
self.responses.append((200, {'Content-Type': 'text/event-stream'}, body))
|
|
with self.send() as response:
|
|
self.assertEqual(response.read(), body)
|
|
usage = self.transport.usage_since(checkpoint)
|
|
self.assertEqual(usage, {'requests': 2, 'measuredRequests': 2, 'lastInputTokens': 90000,
|
|
'peakInputTokens': 90000, 'complete': True})
|
|
self.assertFalse(self.transport.usage_since(self.transport.usage_checkpoint())['complete'])
|
|
self.assertNotIn('DO_NOT_LOG', self.log.read_text())
|
|
|
|
def test_missing_or_malformed_usage_is_not_measured(self):
|
|
for usage in (None, {}, {'input_tokens': -1}, {'input_tokens': '80000'}, {'input_tokens': True}):
|
|
self.transport.record_usage(1, usage)
|
|
self.assertEqual(self.transport.usage_since(0)['measuredRequests'], 0)
|
|
|
|
def test_cancel_interrupts_waiting_headers_without_retry_or_orphan_slot(self):
|
|
entered, release = threading.Event(), threading.Event()
|
|
def delayed(handler):
|
|
entered.set()
|
|
release.wait(3)
|
|
try:
|
|
handler.send_response(200)
|
|
handler.end_headers()
|
|
handler.wfile.write(b'data: {"type":"response.completed"}\n\n')
|
|
except OSError:
|
|
pass
|
|
self.responses.append(delayed)
|
|
done = threading.Event()
|
|
def client():
|
|
try:
|
|
with self.send() as response:
|
|
response.read()
|
|
except (OSError, urllib.error.URLError):
|
|
pass
|
|
finally:
|
|
done.set()
|
|
thread = threading.Thread(target=client)
|
|
thread.start()
|
|
try:
|
|
self.assertTrue(entered.wait(2))
|
|
self.assertEqual(self.transport.activity_since(0)['phase'], 'request')
|
|
self.assertEqual(self.transport.cancel_since(0), 1)
|
|
self.assertTrue(done.wait(1), 'cancel must interrupt the blocked upstream read')
|
|
thread.join(timeout=1)
|
|
self.assertFalse(thread.is_alive())
|
|
self.assertEqual(len(self.requests), 1)
|
|
self.assertNotIn('"retry"', self.log.read_text())
|
|
finally:
|
|
release.set()
|
|
thread.join(timeout=5)
|
|
|
|
def test_cancelled_attempt_does_not_start_its_header_retry(self):
|
|
number = self.transport.begin_request()
|
|
opener = Mock()
|
|
def timeout(*args, **kwargs):
|
|
self.transport.cancel_since(0)
|
|
raise TimeoutError()
|
|
opener.open.side_effect = timeout
|
|
with patch('codex_http_transport.urllib.request.build_opener', return_value=opener), self.assertRaises(RequestCancelled):
|
|
self.transport.open_request(urllib.request.Request(self.transport.upstream, data=b'body'), number)
|
|
self.assertEqual(opener.open.call_count, 1)
|
|
self.transport.finish_request(number)
|
|
|
|
def test_cancel_interrupts_an_open_stream_and_does_not_fabricate_completion(self):
|
|
streamed, release, done = threading.Event(), threading.Event(), threading.Event()
|
|
def delayed(handler):
|
|
handler.send_response(200)
|
|
handler.send_header('Content-Type', 'text/event-stream')
|
|
handler.end_headers()
|
|
handler.wfile.write(b'data: {"type":"response.created"}\n\n')
|
|
handler.wfile.flush()
|
|
release.wait(3)
|
|
self.responses.append(delayed)
|
|
def client():
|
|
try:
|
|
with self.send() as response:
|
|
response.readline()
|
|
streamed.set()
|
|
response.read()
|
|
except OSError:
|
|
pass
|
|
finally:
|
|
done.set()
|
|
thread = threading.Thread(target=client)
|
|
thread.start()
|
|
try:
|
|
self.assertTrue(streamed.wait(2))
|
|
self.assertEqual(self.transport.cancel_since(0), 1)
|
|
self.assertTrue(done.wait(1))
|
|
thread.join(timeout=1)
|
|
self.assertNotIn('response.completed', self.log.read_text())
|
|
self.assertNotIn('"retry"', self.log.read_text())
|
|
finally:
|
|
release.set()
|
|
thread.join(timeout=5)
|
|
|
|
def test_stream_activity_is_bounded_and_scoped_to_current_turn(self):
|
|
number = self.transport.begin_request()
|
|
self.transport.update_request(number, phase='streaming', bytes=37, lastData=time.monotonic())
|
|
self.assertTrue(self.transport.activity_since(0)['canWait'])
|
|
self.assertEqual(self.transport.activity_since(number), {})
|
|
self.transport.update_request(number, started=time.monotonic() - self.transport.response_budget - 1)
|
|
self.assertFalse(self.transport.activity_since(0)['canWait'])
|
|
self.transport.update_request(number, phase='incomplete')
|
|
self.transport.finish_request(number)
|
|
self.assertFalse(self.transport.activity_since(0)['canWait'])
|
|
|
|
def test_oversized_context_is_rejected_without_upstream_or_content_logging(self):
|
|
self.transport.max_request_bytes = 1024
|
|
with self.assertRaises(urllib.error.HTTPError) as error:
|
|
self.send(b'{"private":"' + b'DO_NOT_LOG' * 200 + b'"}')
|
|
self.assertEqual(error.exception.code, 413)
|
|
self.assertEqual(len(self.requests), 0)
|
|
self.assertIn('context_budget_exceeded', error.exception.read().decode())
|
|
self.assertNotIn('DO_NOT_LOG', self.log.read_text())
|
|
self.assertEqual(self.transport.activity_since(0)['phase'], 'context_rejected')
|
|
self.assertFalse(self.transport.activity_since(0)['canWait'])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|