perf(llm): harden request policy and optimize connections, reasoning and prompt caching (#2661)

* fix(llm): centralize request policy and preserve stream diagnostics

* test(llm): dispose SQLite engine after migration regression

* perf(llm): reuse connections and unify reasoning and prompt cache policy

* fix(llm): preserve explicit native prompt cache options

* perf(llm): fix measured TLS setup and streaming reuse overhead

* fix(llm): resolve plugin request-chain regressions found by E2E

---------

Co-authored-by: Hyu <[email protected]>

tests/e2e/test_llm_request_chain.py

@@ -0,0 +1,442 @@
+"""Real HTTP -> Core -> Plugin Runtime -> Runner -> Core -> HTTP provider tests.
+
+Only the remote provider is deterministic. No requester, SDK action transport,
+model manager, persistence or monitoring implementation is mocked.
+"""
+
+import json
+import importlib.metadata
+import os
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+import shutil
+import select
+import socket
+import sqlite3
+import threading
+import time
+
+import pytest
+import yaml
+import httpx
+from packaging.version import Version
+
+from tests.e2e import test_runner_plugin_runtime as runtime
+
+# Reuse the real process lifecycle; override the package/config and port below.
+runner_runtime_ports = runtime.runner_runtime_ports
+runner_runtime_process = runtime.runner_runtime_process
+runner_langbot_process = runtime.runner_langbot_process
+runner_client = runtime.runner_client
+
+pytestmark = pytest.mark.e2e
+RUNNER = runtime.QA_RUNNER_ID
+USAGE = {
+    'prompt_tokens': 24,
+    'completion_tokens': 4,
+    'total_tokens': 28,
+    'prompt_tokens_details': {'cached_tokens': 16},
+    'completion_tokens_details': {'reasoning_tokens': 2},
+}
+
+
+@pytest.fixture(scope='session')
+def runner_e2e_tmpdir(tmp_path_factory):
+    # Retain logs/SQLite in pytest's bounded temp history for failure diagnosis.
+    return tmp_path_factory.mktemp('llm-chain')
+
+
+@pytest.fixture(scope='session')
+def runner_e2e_port():
+    return runtime._free_port()
+
+
+@pytest.fixture(scope='session')
+def runner_e2e_config_path(runner_e2e_tmpdir, runner_e2e_port, runner_runtime_ports):
+    path = runtime.runner_e2e_config_path.__wrapped__(runner_e2e_tmpdir, runner_e2e_port, runner_runtime_ports)
+    root = runner_e2e_tmpdir / 'agent-runner-qa-package'
+    manifest = root / 'components/runner/default.yaml'
+    data = yaml.safe_load(manifest.read_text(encoding='utf-8'))
+    data['spec']['permissions'] = {'models': ['invoke', 'stream', 'count_tokens']}
+    data['spec']['config'] = [
+        {'name': 'model', 'label': {'en_US': 'Model'}, 'type': 'llm-model-selector', 'required': True},
+    ]
+    manifest.write_text(yaml.safe_dump(data), encoding='utf-8')
+    shutil.copyfile(Path(__file__).parent / 'utils/llm_probe_runner.py', root / 'components/runner/default.py')
+    if legacy_root := os.environ.get('LLM_E2E_LEGACY_API_DIR'):
+        for filename in ('legacy_api.py', 'legacy_message.py'):
+            shutil.copyfile(Path(legacy_root) / filename, root / filename)
+    shutil.make_archive(str(runner_e2e_tmpdir / 'agent-runner-qa'), 'zip', root_dir=root)
+    return path
+
+
+@pytest.fixture(scope='module')
+def provider_wire():
+    records = []
+
+    class Handler(BaseHTTPRequestHandler):
+        protocol_version = 'HTTP/1.1'
+
+        def log_message(self, *args):
+            pass
+
+        def do_POST(self):
+            body = json.loads(self.rfile.read(int(self.headers['Content-Length'])))
+            scenario = self.headers.get('x-e2e-case', 'normal')
+            row = {'body': body, 'headers': dict(self.headers), 'case': scenario, 'closed': False}
+            records.append(row)
+            attempt = sum(x['case'] == scenario for x in records)
+            status = 503 if scenario.startswith('retry') and attempt == 1 else 200
+            if scenario == 'auth':
+                status = 401
+            if status != 200:
+                data = json.dumps({'error': {'message': 'injected failure', 'type': 'api_error'}}).encode()
+                self.send_response(status)
+                self.send_header('Content-Type', 'application/json')
+                self.send_header('Content-Length', str(len(data)))
+                self.send_header('Retry-After', '0')
+                self.end_headers()
+                self.wfile.write(data)
+                return
+            if scenario in ('first-timeout', 'total-timeout'):
+                time.sleep(1)
+            try:
+                if not body.get('stream'):
+                    data = json.dumps(
+                        {
+                            'id': 'e2e',
+                            'object': 'chat.completion',
+                            'created': 1,
+                            'model': body['model'],
+                            'choices': [
+                                {
+                                    'index': 0,
+                                    'message': {'role': 'assistant', 'content': 'hello world'},
+                                    'finish_reason': 'stop',
+                                }
+                            ],
+                            'usage': USAGE,
+                        }
+                    ).encode()
+                    self.send_response(200)
+                    self.send_header('Content-Type', 'application/json')
+                    self.send_header('Content-Length', str(len(data)))
+                    self.end_headers()
+                    self.wfile.write(data)
+                    return
+                self.send_response(200)
+                self.send_header('Content-Type', 'text/event-stream')
+                self.send_header('Transfer-Encoding', 'chunked')
+                self.end_headers()
+
+                def send(event):
+                    data = ('data: ' + (event if isinstance(event, str) else json.dumps(event)) + '\n\n').encode()
+                    self.wfile.write(f'{len(data):X}\r\n'.encode() + data + b'\r\n')
+                    self.wfile.flush()
+
+                base = {'id': 'e2e', 'object': 'chat.completion.chunk', 'created': 1, 'model': body['model']}
+                if scenario.startswith('reasoning-'):
+                    send(
+                        {
+                            **base,
+                            'choices': [
+                                {'index': 0, 'delta': {'reasoning_content': 'private thought'}, 'finish_reason': None}
+                            ],
+                        }
+                    )
+                for index, content in enumerate(['hello ', 'world']):
+                    send({**base, 'choices': [{'index': 0, 'delta': {'content': content}, 'finish_reason': None}]})
+                    if index == 0 and (scenario in ('idle-timeout', 'disconnect') or scenario.startswith('cancel-')):
+                        if scenario == 'disconnect':
+                            self.close_connection = True
+                            return
+                        deadline = time.monotonic() + 1
+                        while time.monotonic() < deadline:
+                            if select.select([self.connection], [], [], 0.02)[0]:
+                                if not self.connection.recv(1, socket.MSG_PEEK):
+                                    row['closed'] = True
+                                    return
+                send({**base, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})
+                send({**base, 'choices': [], 'usage': USAGE})
+                send('[DONE]')
+                self.wfile.write(b'0\r\n\r\n')
+                self.wfile.flush()
+            except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
+                row['closed'] = True
+
+    server = ThreadingHTTPServer(('127.0.0.1', 0), Handler)
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    yield f'http://127.0.0.1:{server.server_port}/v1', records
+    server.shutdown()
+    server.server_close()
+    thread.join(timeout=5)
+
+
+@pytest.fixture(scope='module')
+def chain(request, provider_wire):
+    request.getfixturevalue('runner_langbot_process')
+    port = request.getfixturevalue('runner_e2e_port')
+    client = httpx.Client(base_url=f'http://127.0.0.1:{port}', timeout=90, trust_env=False)
+    request.addfinalizer(client.close)
+    tmpdir = request.getfixturevalue('runner_e2e_tmpdir')
+    token = runtime._init_and_auth(client)
+    runtime._ensure_qa_plugin(client, token, tmpdir / 'agent-runner-qa.zip')
+    client.headers['Authorization'] = f'Bearer {token}'
+
+    def post(path, body):
+        response = client.post('/api/v1/' + path, json=body)
+        assert response.status_code == 200, response.text
+        payload = response.json()
+        assert payload['code'] == 0, payload
+        return payload['data']
+
+    provider = post(
+        'provider/providers',
+        {
+            'name': 'LLM chain E2E',
+            'requester': 'openai-chat-completions',
+            'base_url': provider_wire[0],
+            'api_keys': ['e2e-not-a-secret'],
+        },
+    )['uuid']
+    model = post(
+        'provider/models/llm',
+        {
+            'name': 'gpt-4o',
+            'provider_uuid': provider,
+            'abilities': ['func_call', 'vision'],
+            'extra_args': {
+                'temperature': 0.1,
+                'request_policy': {'retries': 1, 'backoff_base': 0.01, 'backoff_max': 0.01},
+            },
+        },
+    )['uuid']
+    agent = post(
+        'agents',
+        {
+            'kind': 'agent',
+            'name': 'LLM chain E2E',
+            'config': {
+                'runner': {'id': RUNNER},
+                'runner_config': {RUNNER: {'model': model}},
+            },
+        },
+    )['uuid']
+
+    def invoke(mode, case='normal', extra=None, *, configuration=None, options=None):
+        selected_model, selected_agent = model, agent
+        if configuration:
+            configured_provider = post(
+                'provider/providers',
+                {
+                    'name': case,
+                    'requester': configuration.get('requester', 'openai-chat-completions'),
+                    'base_url': provider_wire[0],
+                    'api_keys': ['e2e-not-a-secret'],
+                },
+            )['uuid']
+            selected_model = post(
+                'provider/models/llm',
+                {
+                    'name': configuration.get('name', 'gpt-4o'),
+                    'provider_uuid': configured_provider,
+                    'abilities': ['func_call', 'vision'],
+                    'extra_args': configuration.get('extra_args', {}),
+                },
+            )['uuid']
+            selected_agent = post(
+                'agents',
+                {
+                    'kind': 'agent',
+                    'name': case,
+                    'config': {
+                        'runner': {'id': RUNNER},
+                        'runner_config': {RUNNER: {'model': selected_model}},
+                    },
+                },
+            )['uuid']
+        result = post(
+            f'agents/{selected_agent}/debug',
+            {
+                'text': json.dumps(
+                    {'mode': mode, 'model': selected_model, 'case': case, 'extra': extra or {}, **(options or {})}
+                ),
+                'conversation_id': 'e2e-' + case,
+                'event_type': 'message.received',
+            },
+        )
+        assert result.get('final_text'), result
+        return json.loads(result['final_text'])
+
+    return invoke, provider_wire[1], tmpdir
+
+
+def test_multimodal_and_tool_definitions_survive_plugin_transport(chain):
+    invoke, records, _ = chain
+    image = 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jRZkAAAAASUVORK5CYII='
+    messages = [
+        {
+            'role': 'user',
+            'content': [
+                {'type': 'text', 'text': 'Describe this image'},
+                {'type': 'image_url', 'image_url': {'url': image}},
+            ],
+        }
+    ]
+    result = invoke('runner-events', 'multimodal', options={'messages': messages, 'tool': True})
+    assert result['text'] == 'hello world', result
+    wire = records[-1]['body']
+    assert wire['messages'][0]['content'][1]['image_url']['url'] == image
+    assert wire['tools'][0]['function']['name'] == 'e2e_lookup'
+
+
+def test_openrouter_budget_cache_and_explicit_attribution(chain):
+    invoke, records, _ = chain
+    result = invoke(
+        'runner-events',
+        'openrouter-budget-cache',
+        configuration={
+            'requester': 'openrouter-chat-completions',
+            'name': 'anthropic/claude-sonnet-4',
+            'extra_args': {
+                'reasoning_config': {'budget_tokens': 2048},
+                'prompt_caching': {'mode': 'anthropic'},
+                'max_tokens': 4096,
+            },
+        },
+        extra={
+            'extra_headers': {'HTTP-Referer': 'https://langbot.app', 'X-OpenRouter-Title': 'LangBot'},
+            'request_policy': {'conversation_header': 'x-langbot-conversation-id'},
+        },
+        options={'messages': [{'role': 'system', 'content': 'Stable prefix'}, {'role': 'user', 'content': 'Hello'}]},
+    )
+    assert result['text'] == 'hello world', result
+    wire = records[-1]
+    assert wire['body']['reasoning']['max_tokens'] == 2048
+    assert wire['body']['messages'][0]['content'][0]['cache_control']['type'] == 'ephemeral'
+    headers = {k.lower(): v for k, v in wire['headers'].items()}
+    assert headers['http-referer'] == 'https://langbot.app'
+    assert headers['x-openrouter-title'] == 'LangBot'
+    assert headers['x-langbot-conversation-id'].startswith('lb-')
+    assert 'prompt_caching' not in wire['body'] and 'reasoning_config' not in wire['body']
+
+
+def test_openai_cache_options_and_call_override(chain):
+    invoke, records, _ = chain
+    result = invoke(
+        'runner-usage',
+        'openai-cache',
+        configuration={
+            'extra_args': {'temperature': 0.1, 'prompt_caching': {'mode': 'openai', 'retention': 'in_memory'}},
+        },
+        extra={'temperature': 0.7, 'prompt_cache_key': 'explicit-test-cache'},
+    )
+    assert result['text'] == 'hello world', result
+    assert records[-1]['body']['temperature'] == 0.7
+    assert records[-1]['body']['prompt_cache_key'] == 'explicit-test-cache'
+    assert records[-1]['body']['prompt_cache_retention'] == 'in_memory'
+
+
+@pytest.mark.parametrize(
+    'mode',
+    [
+        'runner-message',
+        'runner-usage',
+        'runner-stream',
+        'runner-events',
+        'plugin-message',
+        'plugin-usage',
+        'plugin-stream',
+        'plugin-events',
+    ],
+)
+def test_public_api_shapes_and_usage(chain, mode):
+    invoke, records, _ = chain
+    result = invoke(mode, mode)
+    assert result['text'] == 'hello world', result
+    assert not result.get('error'), result
+    if mode.endswith(('usage', 'events')):
+        assert result['usage']['prompt_tokens_details']['cached_tokens'] == 16
+        assert result['usage']['completion_tokens_details']['reasoning_tokens'] == 2
+    wire = records[-1]['body']
+    assert wire['temperature'] == 0.1
+    assert 'request_policy' not in wire
+
+
+@pytest.mark.parametrize(
+    'case,extra,attempts,error',
+    [
+        ('retry-once', {}, 2, False),
+        ('retry-disabled', {'num_retries': 0}, 1, True),
+        ('retry-max-disabled', {'max_retries': 0}, 1, True),
+        ('auth', {}, 1, True),
+        ('first-timeout', {'request_policy': {'first_event_timeout': 0.2, 'retries': 0}}, 1, True),
+        ('total-timeout', {'request_policy': {'total_timeout': 0.2, 'retries': 0}}, 1, True),
+        ('idle-timeout', {'request_policy': {'stream_idle_timeout': 0.2}}, 1, True),
+        ('disconnect', {}, 1, True),
+    ],
+)
+def test_failures_retry_boundaries_and_recovery(chain, case, extra, attempts, error):
+    invoke, records, _ = chain
+    result = invoke('runner-events', case, extra)
+    assert bool(result.get('error')) == error, result
+    assert sum(row['case'] == case for row in records) == attempts
+    assert invoke('runner-events', 'recover-' + case)['text'] == 'hello world'
+
+
+@pytest.mark.parametrize('surface', ['runner-events', 'runner-stream', 'plugin-events', 'plugin-stream'])
+def test_cancellation_releases_upstream_and_next_request_works(chain, surface):
+    invoke, records, tmpdir = chain
+    case = 'cancel-' + surface
+    result = invoke(case, case)
+    assert result['text'] == 'hello ', result
+    time.sleep(1.5)
+    cancelled = [row for row in records if row['case'] == case]
+    with sqlite3.connect(tmpdir / 'data/langbot.db') as db:
+        diagnostics = db.execute('SELECT session_id, details FROM monitoring_llm_calls').fetchall()
+    assert len(cancelled) == 1
+    assert invoke('runner-events', 'after-' + case)['text'] == 'hello world'
+    if (
+        not cancelled[0]['closed']
+        and Version(importlib.metadata.version('langbot-plugin')) <= Version('0.7.11')
+        and not os.environ.get('LLM_E2E_REQUIRE_STREAM_CLOSE')
+    ):
+        pytest.xfail('Published SDK <=0.7.11 omits aclosing in public proxy wrappers; requires companion SDK fix')
+    assert cancelled[0]['closed'], diagnostics[-2:]
+
+
+@pytest.mark.parametrize('remove_think', [True, False])
+def test_reasoning_visibility_through_runner(chain, remove_think):
+    result = chain[0]('runner-events', f'reasoning-{remove_think}', options={'remove_think': remove_think})
+    assert result['text'].endswith('hello world'), result
+    assert ('private thought' in result['text']) == (not remove_think)
+
+
+def test_invalid_budget_fails_before_network_and_recovers(chain):
+    invoke, records, _ = chain
+    result = invoke('runner-events', 'invalid-budget', {'reasoning_config': {'budget_tokens': 1}})
+    assert result.get('error'), result
+    assert not any(row['case'] == 'invalid-budget' for row in records)
+    assert invoke('runner-events', 'after-invalid-budget')['text'] == 'hello world'
+
+
+def test_monitoring_persists_complete_usage_and_structured_errors(chain):
+    invoke, _, tmpdir = chain
+    invoke('runner-events', 'monitoring')
+    with sqlite3.connect(tmpdir / 'data/langbot.db') as db:
+        rows = db.execute('SELECT details FROM monitoring_llm_calls').fetchall()
+    details = [json.loads(row[0]) for row in rows if row[0]]
+    assert any(x.get('usage', {}).get('prompt_tokens_details', {}).get('cached_tokens') == 16 for x in details)
+    assert any(x.get('timing', {}).get('transport_setup_ms') is not None for x in details)
+    assert any(x.get('error', {}).get('code') for x in details if x.get('error'))
+
+
+@pytest.mark.skipif(not os.environ.get('LLM_E2E_LEGACY_API_DIR'), reason='Historical SDK sources not supplied')
+@pytest.mark.parametrize('mode', ['legacy-message', 'legacy-usage', 'legacy-stream', 'legacy-events'])
+def test_historical_sdk_proxy_and_message_entities(chain, mode):
+    result = chain[0](mode, mode)
+    assert result['text'] == 'hello world', result
+    assert not result.get('error'), result
+    if mode.endswith(('usage', 'events')):
+        assert result['usage']['prompt_tokens_details']['cached_tokens'] == 16

tests/e2e/utils/llm_probe_runner.py

@@ -0,0 +1,93 @@
+"""Packaged into the isolated E2E plugin; calls only public SDK LLM APIs."""
+
+from contextlib import aclosing
+import json
+
+from langbot_plugin.api.definition.components.runner.runner import Runner
+from langbot_plugin.api.entities.builtin.provider.message import Message
+from langbot_plugin.api.entities.builtin.runner.result import RunnerResult
+from langbot_plugin.api.proxies.langbot_api import LangBotAPIProxy
+
+
+class DefaultRunner(Runner):
+    async def run(self, ctx):
+        options = json.loads(ctx.input.to_text())
+        mode = options['mode']
+        cancel = mode.startswith('cancel-')
+        if cancel:
+            mode = mode.removeprefix('cancel-')
+        api = LangBotAPIProxy(self._plugin_runtime_handler) if mode.startswith('plugin-') else ctx.api
+        message_type = Message
+        if mode.startswith('legacy-'):
+            # Historical files are supplied verbatim from a released SDK tag.
+            # Its proxy and serializers use the real current runtime transport.
+            import legacy_api
+            import legacy_message
+
+            legacy_api.provider_message = legacy_message
+            api = legacy_api.LangBotAPIProxy(self._plugin_runtime_handler)
+            message_type = legacy_message.Message
+        extra = {
+            **options['extra'],
+            'extra_headers': {
+                **options['extra'].get('extra_headers', {}),
+                'x-e2e-case': options['case'],
+            },
+        }
+        kwargs = {
+            'llm_model_uuid': options['model'],
+            'messages': [
+                message_type.model_validate(m)
+                for m in options.get(
+                    'messages',
+                    [
+                        {'role': 'user', 'content': 'Hello'},
+                    ],
+                )
+            ],
+            'extra_args': extra,
+        }
+        if options.get('tool'):
+            from langbot_plugin.api.entities.builtin.resource.tool import LLMTool
+
+            async def lookup(**kwargs):
+                return 'result'
+
+            kwargs['funcs'] = [
+                LLMTool(
+                    name='e2e_lookup',
+                    human_desc='Lookup',
+                    description='Lookup',
+                    parameters={'type': 'object', 'properties': {}},
+                    func=lookup,
+                )
+            ]
+        result = {'text': '', 'usage': None}
+        if 'remove_think' in options:
+            kwargs['remove_think'] = options['remove_think']
+        try:
+            if mode.endswith('message'):
+                message = await api.invoke_llm(**kwargs)
+                assert isinstance(message, message_type)
+                result['text'] = message.content
+            elif mode.endswith('usage'):
+                response = await api.invoke_llm_with_usage(**kwargs)
+                result.update(
+                    text=response.message.content, usage=response.usage.model_dump() if response.usage else None
+                )
+            else:
+                events = mode.endswith('events')
+                method = api.invoke_llm_stream_events if events else api.invoke_llm_stream
+                async with aclosing(method(**kwargs)) as stream:
+                    async for item in stream:
+                        chunk = item.chunk if events else item
+                        if chunk:
+                            result['text'] += chunk.content or ''
+                        if events and item.usage:
+                            result['usage'] = item.usage.model_dump()
+                        if cancel and result['text']:
+                            break
+        except Exception as exc:
+            result['error'] = type(exc).__name__ + ': ' + str(exc)
+        yield RunnerResult.message_completed(ctx.run_id, Message(role='assistant', content=json.dumps(result)))
+        yield RunnerResult.run_completed(ctx.run_id, finish_reason='stop')