VianoBugAnalysis/tests/test_llm_caching.py

"""Offline cache protocol, accounting, privacy and portable-request regressions."""
import asyncio
import copy
import json
from dataclasses import replace
from importlib.resources import files

import pytest
from jsonschema import Draft202012Validator

from viano_bug_analysis.llm import GatewayConfig, LLMError, LLMGateway, ProviderConfig, ProviderResponse, ProviderUnavailable
from viano_bug_analysis.llm.gemini import GeminiProvider
from viano_bug_analysis.llm.usage import usage_metadata
from viano_bug_analysis.llm.diagnostics import sanitize
from test_llm_gateway import Adapter, request, http_endpoint


def config(**kw):
    return ProviderConfig(id='gemini-test', model='gemini-test', base_url='http://localhost/v1beta',
                          adapter='gemini', context_cache=True, cache_min_prefix_chars=1, **kw)


def native(cached=80, finish='STOP'):
    return {'modelVersion': 'gemini-test', 'candidates': [{'finishReason': finish,
            'content': {'parts': [{'text': '{"answer":"ok"}'}]}}],
            'usageMetadata': {'promptTokenCount': 100, 'cachedContentTokenCount': cached,
                              'candidatesTokenCount': 10, 'thoughtsTokenCount': 5}}


def scripted(adapter, monkeypatch, actions):
    seen = []
    def post(url, payload, headers, limit, deadline):
        seen.append((url, copy.deepcopy(payload), headers.copy()))
        result = actions.pop(0)
        if isinstance(result, Exception):
            raise result
        return result
    monkeypatch.setattr(adapter, '_post', post)
    return seen


CACHE = {'name': 'cachedContents/test-cache', 'usageMetadata': {'totalTokenCount': 80}}


def test_cache_reuses_only_instructions_schema_keeps_full_live_research(monkeypatch):
    adapter, cfg = GeminiProvider(), config(api_key_env='CACHE_TEST_KEY')
    monkeypatch.setenv('CACHE_TEST_KEY', 'never-persist-this')
    seen = scripted(adapter, monkeypatch, [CACHE, native(), native()])
    req = request()
    value = json.loads(req['messages'][1]['content'])
    value['input'] = {'evidence': {'E000001': {'text': 'complete first finding'}}}
    req['messages'][1]['content'] = json.dumps(value)
    original = copy.deepcopy(req)
    first = asyncio.run(adapter.generate(cfg, req, max_response_bytes=10000))
    value['input']['evidence']['E000002'] = {'text': 'complete second finding'}
    second_req = copy.deepcopy(req)
    second_req['messages'][1]['content'] = json.dumps(value)
    second = asyncio.run(adapter.generate(cfg, second_req, max_response_bytes=10000))
    assert req == original
    assert len(seen) == 3 and seen[0][0].endswith('/cachedContents')
    assert 'complete first finding' not in json.dumps(seen[0][1])
    assert 'complete first finding' in json.dumps(seen[2][1])
    assert 'complete second finding' in json.dumps(seen[2][1])
    assert seen[2][1]['cachedContent'] == CACHE['name']
    assert 'systemInstruction' not in seen[2][1]
    assert first.diagnostics['cache_status'] == 'created'
    assert second.diagnostics['cache_status'] == 'reused'
    assert first.output_tokens == 15 and first.reasoning_tokens == 5
    assert second.diagnostics['cache_created_tokens'] == 0
    assert 'never-persist-this' not in repr(first) + repr(adapter._cache)


@pytest.mark.parametrize('change', ['model','system','schema','credential','url'])
def test_changed_prefix_or_provider_never_reuses_cache(monkeypatch, change):
    adapter, cfg, req = GeminiProvider(), config(api_key_env='CACHE_TEST_KEY'), request()
    monkeypatch.setenv('CACHE_TEST_KEY', 'first-key')
    seen = scripted(adapter, monkeypatch, [CACHE, native(), CACHE, native()])
    asyncio.run(adapter.generate(cfg, req, max_response_bytes=10000))
    if change == 'model': cfg = replace(cfg, model='other')
    elif change == 'system': req['messages'][0]['content'] += ' changed'
    elif change == 'schema': req['output_schema']['description'] = 'changed'
    elif change == 'credential': monkeypatch.setenv('CACHE_TEST_KEY', 'second-key')
    else: cfg = replace(cfg, base_url='http://localhost/other')
    asyncio.run(adapter.generate(cfg, req, max_response_bytes=10000))
    assert sum(url.endswith('/cachedContents') for url, _, _ in seen) == 2


def test_expired_cache_recreated(monkeypatch):
    adapter, cfg = GeminiProvider(), config()
    seen = scripted(adapter, monkeypatch, [CACHE, native(), CACHE, native()])
    asyncio.run(adapter.generate(cfg, request(), max_response_bytes=10000))
    adapter._cache = {key: (name, 0) for key, (name, _) in adapter._cache.items()}
    asyncio.run(adapter.generate(cfg, request(), max_response_bytes=10000))
    assert len(seen) == 4


@pytest.mark.parametrize('code', [400,403,404,429,503])
def test_cache_creation_rejected_falls_back_and_is_not_retried_every_round(monkeypatch, code):
    adapter, cfg = GeminiProvider(), config()
    seen = scripted(adapter, monkeypatch, [LLMError('LLM_HTTP_CONFIGURATION', http_status=code), native(0), native(0)])
    for _ in range(2):
        result = asyncio.run(adapter.generate(cfg, request(), max_response_bytes=10000))
        assert result.diagnostics['cache_status'] == 'fallback'
    assert len(seen) == 3
    assert all('cachedContent' not in payload and 'systemInstruction' in payload for _, payload, _ in seen[1:])


def test_rejected_cached_generation_retries_full_input_once(monkeypatch):
    adapter, cfg = GeminiProvider(), config()
    seen = scripted(adapter, monkeypatch, [CACHE, LLMError('LLM_HTTP_CONFIGURATION', http_status=404), native(0)])
    result = asyncio.run(adapter.generate(cfg, request(), max_response_bytes=10000))
    assert result.diagnostics['cache_status'] == 'fallback'
    assert 'cachedContent' in seen[1][1] and 'cachedContent' not in seen[2][1]
    assert 'systemInstruction' in seen[2][1]


def test_cache_loss_and_new_adapter_need_no_provider_state(monkeypatch):
    cfg, req = config(), request()
    for adapter in (GeminiProvider(), GeminiProvider()):
        seen = scripted(adapter, monkeypatch, [CACHE, native()])
        result = asyncio.run(adapter.generate(cfg, req, max_response_bytes=10000))
        assert result.document == {'answer':'ok'} and len(seen) == 2


@pytest.mark.parametrize('cfg', [replace(config(), context_cache=False), replace(config(), cache_min_prefix_chars=1000000)])
def test_cache_optional_no_create_when_disabled_or_small(monkeypatch, cfg):
    adapter = GeminiProvider()
    seen = scripted(adapter, monkeypatch, [native(0)])
    asyncio.run(adapter.generate(cfg, request(), max_response_bytes=10000))
    assert len(seen) == 1 and 'cachedContent' not in seen[0][1]


def test_usage_counts_and_cost_do_not_double_charge_thinking_or_cached_input():
    cfg = config(input_cost_per_million=2, cached_input_cost_per_million=.2,
                 output_cost_per_million=4, cache_storage_cost_per_million_hour=.5)
    response = ProviderResponse({}, 'm', 100, 15, {'cache_created_tokens':80,'cache_ttl_seconds':300}, 80, 5)
    usage = usage_metadata(cfg, response)
    assert usage['uncached_input_tokens'] == 20
    assert usage['estimated_cost'] == pytest.approx((20*2 + 80*.2 + 15*4)/1e6)
    assert usage['estimated_cache_storage_cost'] == pytest.approx(80*.5/1e6/12)
    assert usage['reasoning_tokens'] == 5


def test_unknown_cache_usage_remains_unknown():
    usage = usage_metadata(config(cached_input_cost_per_million=.1), ProviderResponse({}, 'm', 100, 10))
    assert usage['cached_input_tokens'] is None and usage['uncached_input_tokens'] is None
    assert usage['estimated_cost'] is None
    assert usage_metadata(config(), ProviderResponse({}, 'm', 10, 1, cached_input_tokens=11))['cached_input_tokens'] is None


def test_config_defaults_keep_frozen_snapshot_compatible():
    snapshot = GatewayConfig((ProviderConfig('x','m','http://localhost'),)).snapshot()['providers'][0]
    for key in ('context_cache','cache_ttl_seconds','cache_min_prefix_chars','cached_input_cost_per_million',
                'cache_storage_cost_per_million_hour'):
        assert key not in snapshot
    assert GatewayConfig((config(),)).snapshot()['providers'][0]['context_cache'] is True


def test_capable_adapter_defaults_to_explicit_caching_and_pins_choice():
    cfg=ProviderConfig('native','gemini-test','http://localhost/v1beta',adapter='gemini')
    assert cfg.use_context_cache
    assert GatewayConfig((cfg,)).snapshot()['providers'][0]['context_cache'] is True
    assert not replace(cfg,context_cache=False).use_context_cache
    assert not ProviderConfig('other','m','http://localhost').use_context_cache
    assert replace(cfg,cache_required=True).use_context_cache


def test_default_native_configuration_creates_cache(monkeypatch):
    adapter=GeminiProvider()
    cfg=replace(config(),context_cache=None)
    seen=scripted(adapter,monkeypatch,[CACHE,native()])
    result=asyncio.run(adapter.generate(cfg,request(),max_response_bytes=10000))
    assert result.diagnostics['cache_status']=='created' and len(seen)==2


@pytest.mark.parametrize('changes', [{'adapter':'openai_compatible'}, {'cache_ttl_seconds':True},
    {'cache_min_prefix_chars':0}, {'cached_input_cost_per_million':-1}, {'model':'../other'}, {'reasoning_effort':'none'}])
def test_bad_cache_configuration_rejected(changes):
    with pytest.raises(LLMError): replace(config(), **changes)


def test_openai_compatible_reports_implicit_cache_hits(http_endpoint):
    url, _, reply = http_endpoint
    reply['body']['usage'].update(prompt_tokens=100, prompt_tokens_details={'cached_tokens':80},
                                  completion_tokens_details={'reasoning_tokens':2})
    cfg = ProviderConfig('test','m',url, input_cost_per_million=2,
                         cached_input_cost_per_million=.2, output_cost_per_million=4)
    metadata = asyncio.run(LLMGateway(GatewayConfig((cfg,))).complete_round(request()))['metadata']
    assert metadata['cached_input_tokens'] == 80 and metadata['uncached_input_tokens'] == 20
    assert metadata['reasoning_tokens'] == 2
    assert metadata['estimated_cost'] == pytest.approx((20*2+80*.2+5*4)/1e6)


def test_native_real_http_transport(http_endpoint):
    url, seen, reply = http_endpoint
    reply['body'] = native()
    cfg = replace(config(), base_url=url, context_cache=False, response_format='json_schema', reasoning_effort='low')
    result = asyncio.run(LLMGateway(GatewayConfig((cfg,))).complete_round(request()))
    assert result['document'] == {'answer':'ok'}
    assert seen[0]['path'].endswith('/models/gemini-test:generateContent')
    generation = seen[0]['body']['generationConfig']
    assert generation['thinkingConfig'] == {'thinkingLevel':'low'}
    assert generation['responseJsonSchema'] == request()['output_schema']


def test_invalid_model_output_keeps_usage_and_storage_diagnostics(monkeypatch):
    adapter, cfg = GeminiProvider(), config()
    scripted(adapter, monkeypatch, [CACHE, native(finish='MAX_TOKENS')])
    gw = LLMGateway(GatewayConfig((cfg,)), {'gemini':adapter})
    with pytest.raises(LLMError) as caught:
        asyncio.run(gw.complete_round(request()))
    assert caught.value.code == 'LLM_INCOMPLETE_OUTPUT'
    assert caught.value.metadata['cached_input_tokens'] == 80
    assert caught.value.metadata['cache_created_tokens'] == 80
    schema = json.loads(files('viano_bug_analysis.schemas').joinpath('llm-call-v1.schema.json').read_text())
    Draft202012Validator(schema['properties']['outcome']['properties']['diagnostics']).validate(caught.value.diagnostics)


def test_strict_phase_metadata_schemas_accept_optional_usage_fields():
    cfg = replace(config(), adapter='fixture', context_cache=False)
    response = ProviderResponse({'answer':'ok'}, 'm', 100, 15, cached_input_tokens=80, reasoning_tokens=5)
    adapter = Adapter({cfg.id:response})
    metadata = asyncio.run(LLMGateway(GatewayConfig((cfg,)), {'fixture':adapter}).complete_round(request()))['metadata']
    for name, version in [('case-interpretation-v1.schema.json','case_interpretation_resolved_v8'),
                          ('hypotheses-v1.schema.json','hypotheses_project_v4')]:
        schema = json.loads(files('viano_bug_analysis.schemas').joinpath(name).read_text())
        # Use an allowed phase-specific prompt; usage contract stays identical.
        target = schema['properties']['llm_metadata']
        phase_metadata = metadata | {'prompt_version':target['properties']['prompt_version']['enum'][-1]}
        Draft202012Validator(target).validate(phase_metadata)
        legacy = {k:v for k,v in phase_metadata.items() if k in target['required']}
        Draft202012Validator(target).validate(legacy)


def test_diagnostics_never_include_provider_cache_handles():
    assert sanitize({'cache_status':{}}) == {}
    assert sanitize({'cache_status':'cachedContents/secret', 'cache_name':'secret',
                     'cached_input_tokens':80}) == {'cached_input_tokens':80}


def test_allocation_is_preserved_through_provider_failover(monkeypatch):
    adapter, cfg = GeminiProvider(), config(cache_storage_cost_per_million_hour=.5)
    scripted(adapter, monkeypatch, [CACHE, ProviderUnavailable('LLM_CONNECTION')])
    fallback = ProviderConfig('fallback','m','http://localhost',adapter='fixture')
    gw = LLMGateway(GatewayConfig((cfg, fallback)), {'gemini':adapter,
        'fixture':Adapter({'fallback':ProviderResponse({'answer':'ok'},'m')})})
    result = asyncio.run(gw.complete_round(request()))
    usage = result['metadata']['attempts'][0]['cache_usage']
    assert usage['cache_created_tokens'] == 80
    assert usage['estimated_cache_storage_cost'] == pytest.approx(80*.5/1e6/12)
    assert 'cachedContents/' not in json.dumps(result)


@pytest.mark.parametrize('status', [401,302,503])
def test_native_http_failures_are_controlled(http_endpoint, status):
    url, seen, reply = http_endpoint
    reply.update(status=status, body=b'private provider body', headers={'Location':url+'/secret'})
    cfg = replace(config(), base_url=url, context_cache=False)
    with pytest.raises(LLMError) as caught:
        asyncio.run(LLMGateway(GatewayConfig((cfg,))).complete_round(request()))
    assert len(seen) == 1
    assert 'private' not in str(caught.value)


def test_native_response_limit(http_endpoint):
    url, _, reply = http_endpoint
    reply['body'] = b'x'*1000
    cfg = replace(config(), base_url=url, context_cache=False)
    with pytest.raises(LLMError) as caught:
        asyncio.run(LLMGateway(GatewayConfig((cfg,),max_response_bytes=100)).complete_round(request()))
    assert caught.value.code == 'LLM_RESPONSE_LIMIT'


def test_native_credentials_checked_before_any_request(monkeypatch):
    adapter = GeminiProvider()
    monkeypatch.delenv('CACHE_TEST_MISSING', raising=False)
    seen = scripted(adapter, monkeypatch, [])
    with pytest.raises(LLMError) as caught:
        asyncio.run(adapter.generate(config(api_key_env='CACHE_TEST_MISSING'),request(),max_response_bytes=1000))
    assert caught.value.code == 'LLM_CREDENTIALS' and not seen


def test_cache_usage_visible_in_telemetry(tmp_path):
    from viano_bug_analysis.lifecycle import Workspace
    from viano_bug_analysis.telemetry import Telemetry
    from viano_bug_analysis.telemetry.__main__ import read_events
    cfg = config(input_cost_per_million=2, cached_input_cost_per_million=.2,
                 output_cost_per_million=4, cache_storage_cost_per_million_hour=.5)
    response = ProviderResponse({},'m',100,15,{'cache_status':'created','cache_created_tokens':80,'cache_ttl_seconds':300},80,5)
    metadata = usage_metadata(cfg,response) | {'input_tokens':100,'output_tokens':15,'cost_currency':'USD'}
    failed = {'cache_usage':{'estimated_cache_storage_cost':.01,'cost_currency':'USD'}}
    metadata['attempts'] = [failed]
    with Telemetry(Workspace(tmp_path/'workspace')) as telemetry:
        telemetry.llm_round({'phase':'INVESTIGATING_HYPOTHESES','status':'SUCCESS','outcome':{'metadata':metadata}})
        metrics = telemetry.prometheus()
        assert 'vba_llm_tokens_total{direction="cached_input_tokens"} 80' in metrics
        assert 'vba_llm_tokens_total{direction="uncached_input_tokens"} 20' in metrics
        assert 'vba_llm_tokens_total{direction="reasoning_tokens"} 5' in metrics
        assert 'component="estimated_cache_storage_cost"' in metrics
        event = next(e for e in read_events(telemetry.workspace.root) if e['event']=='llm_round')['detail']
        assert event['cached_input_tokens'] == 80 and event['cache_status'] == 'created'
        assert event['failed_cache_storage'] == [{'estimated_cost':.01,'currency':'USD'}]


def test_example_config_loads_without_credentials_or_network():
    from pathlib import Path
    cfg = GatewayConfig.load(Path(__file__).parents[1]/'config/llm.gemini-cache.example.toml')
    assert cfg.providers[0].adapter == 'gemini' and cfg.providers[0].context_cache
    LLMGateway(cfg)


def test_cache_memory_and_remote_allocations_are_bounded(monkeypatch):
    adapter, cfg = GeminiProvider(), config()
    actions = []
    for _ in range(8): actions.extend([CACHE, native()])
    actions.append(native(0))
    seen = scripted(adapter, monkeypatch, actions)
    for index in range(9):
        req = request()
        req['messages'][0]['content'] += str(index)
        asyncio.run(adapter.generate(cfg,req,max_response_bytes=10000))
    assert len(adapter._cache) == 8
    assert sum(url.endswith('/cachedContents') for url,_,_ in seen) == 8


@pytest.mark.parametrize('invalid', [{}, {'name':'https://evil.example/secret'}, {'name':'cachedContents/../secret'}])
def test_invalid_cache_handles_never_used(monkeypatch, invalid):
    adapter = GeminiProvider()
    seen = scripted(adapter,monkeypatch,[invalid,native(0)])
    result = asyncio.run(adapter.generate(config(),request(),max_response_bytes=10000))
    assert result.diagnostics['cache_status'] == 'fallback'
    assert 'cachedContent' not in seen[-1][1]


def test_required_cache_accepts_confirmed_creation_and_reuse(monkeypatch):
    adapter, cfg = GeminiProvider(), replace(config(), cache_required=True)
    seen = scripted(adapter,monkeypatch,[CACHE,native(),native()])
    for _ in range(2):
        result=asyncio.run(adapter.generate(cfg,request(),max_response_bytes=10000))
        assert result.cached_input_tokens == 80
    assert len(seen)==3 and all('cachedContent' in payload for _,payload,_ in seen[1:])


@pytest.mark.parametrize('mode', ['small','rejected','stale','missing_usage','zero_usage','inconsistent_usage'])
def test_required_cache_never_silently_falls_back(monkeypatch,mode):
    adapter, cfg = GeminiProvider(), replace(config(),cache_required=True)
    actions=[CACHE,native()]
    expected='LLM_CACHE_REQUIRED'
    if mode=='small':
        cfg=replace(cfg,cache_min_prefix_chars=1000000)
        actions=[]
    elif mode=='rejected': actions=[LLMError('LLM_HTTP_CONFIGURATION',http_status=400)]
    elif mode=='stale': actions=[CACHE,LLMError('LLM_HTTP_CONFIGURATION',http_status=404)]
    else:
        expected='LLM_CACHE_USAGE_UNCONFIRMED'
        actions[1]=native(None if mode=='missing_usage' else 0 if mode=='zero_usage' else 101)
    seen=scripted(adapter,monkeypatch,actions)
    with pytest.raises(LLMError) as caught:
        asyncio.run(adapter.generate(cfg,request(),max_response_bytes=10000))
    assert caught.value.code==expected
    assert all('cachedContent' in payload for url,payload,_ in seen if url.endswith(':generateContent'))


def test_required_prefix_keeps_other_phases_optional(monkeypatch):
    adapter=GeminiProvider()
    cfg=replace(config(),cache_required=True,cache_required_prompt_prefix='investigation_project_',cache_min_prefix_chars=1000000)
    seen=scripted(adapter,monkeypatch,[native(0)])
    assert asyncio.run(adapter.generate(cfg,request(),max_response_bytes=10000)).document=={'answer':'ok'}
    req=request(); req['prompt_version']='investigation_project_v10'
    with pytest.raises(LLMError,match='LLM_CACHE_REQUIRED'):
        asyncio.run(adapter.generate(cfg,req,max_response_bytes=10000))
    assert len(seen)==1


@pytest.mark.parametrize('failure',[ProviderUnavailable('LLM_CONNECTION'),TimeoutError()])
def test_required_cache_cannot_fail_over_to_uncached_provider(failure):
    cfg=replace(config(),cache_required=True)
    fallback=ProviderConfig('fallback','m','http://localhost',adapter='fixture')
    adapter=Adapter({cfg.id:failure,'fallback':ProviderResponse({'answer':'ok'},'m')})
    gateway=LLMGateway(GatewayConfig((cfg,fallback)),{'gemini':adapter,'fixture':adapter})
    with pytest.raises(LLMError,match='LLM_CACHE_REQUIRED'):
        asyncio.run(gateway.complete_round(request()))
    assert adapter.calls==[cfg.id]


@pytest.mark.parametrize('changes',[{'context_cache':False},{'cache_required':'true'},
    {'cache_required_prompt_prefix':''},{'cache_required':False,'cache_required_prompt_prefix':'investigation_'}])
def test_required_cache_config_is_strict(changes):
    with pytest.raises(LLMError):
        replace(config(),**({'cache_required':True}|changes))


def test_gateway_itself_requires_confirmed_cached_usage():
    cfg=replace(config(),cache_required=True)
    adapter=Adapter({cfg.id:ProviderResponse({'answer':'ok'},'m',100,10)})
    with pytest.raises(LLMError,match='LLM_CACHE_USAGE_UNCONFIRMED'):
        asyncio.run(LLMGateway(GatewayConfig((cfg,)),{'gemini':adapter}).complete_round(request()))