"""Tests for the stage-B egress resolver/assembly/audit (ADR-0007). Covers label resolution from declared attributes + dynamic refinements (collection is_public, store URL fail-up), run assembly, and the audit-mode decision (which must never raise). """ from __future__ import annotations from unittest.mock import patch from local_deep_research.security.egress import run_classification as rc from local_deep_research.security.egress.classification import ( Exposure, Label, Role, Sensitivity, ) from local_deep_research.security.egress.policy import ( EgressContext, EgressScope, ) S = Sensitivity.SENSITIVE NS = Sensitivity.NON_SENSITIVE EXP = Exposure.EXPOSING CON = Exposure.CONTAINED POLICY = "local_deep_research.security.egress.policy" def ctx(primary: str = "arxiv", username=None) -> EgressContext: return EgressContext( scope=EgressScope.BOTH, primary_engine=primary, require_local_llm=False, require_local_embeddings=False, username=username, ) class _PublicEngine: egress_sensitivity = NS egress_exposure = EXP class _LocalStore: egress_sensitivity = S egress_exposure = CON url_setting = "search.engine.web.paperless.default_params.api_url" # --------------------------------------------------------------------------- # engine_label — static engines + URL fail-up # --------------------------------------------------------------------------- def test_public_engine_label(): with patch(f"{POLICY}._get_engine_class", return_value=_PublicEngine): assert rc.engine_label("arxiv", {}, ctx()) == Label(NS, EXP) def test_unknown_engine_label_is_none(): with patch(f"{POLICY}._get_engine_class", return_value=None): assert rc.engine_label("mystery", {}, ctx()) is None def test_local_store_contained_on_local_url(): with ( patch(f"{POLICY}._get_engine_class", return_value=_LocalStore), patch(f"{POLICY}._classify_engine_url", return_value=True), ): assert rc.engine_label("paperless", {}, ctx()) == Label(S, CON) def test_local_store_fails_up_to_exposing_on_public_url(): # Quadrant 4: sensitive data + exposing sink (public host). with ( patch(f"{POLICY}._get_engine_class", return_value=_LocalStore), patch(f"{POLICY}._classify_engine_url", return_value=False), ): assert rc.engine_label("paperless", {}, ctx()) == Label(S, EXP) # --------------------------------------------------------------------------- # engine_label — collections / library (sensitivity from is_public) # --------------------------------------------------------------------------- def test_private_collection_is_sensitive(): with patch(f"{POLICY}._resolve_collection_is_public", return_value=False): assert rc.engine_label("collection_abc", {}, ctx()) == Label(S, CON) def test_public_collection_is_non_sensitive(): with patch(f"{POLICY}._resolve_collection_is_public", return_value=True): assert rc.engine_label("collection_abc", {}, ctx()) == Label(NS, CON) def test_library_aggregate_is_sensitive(): with patch(f"{POLICY}._resolve_collection_is_public", return_value=False): assert rc.engine_label("library", {}, ctx()) == Label(S, CON) # --------------------------------------------------------------------------- # provider labels (inference sinks — exposure only) # --------------------------------------------------------------------------- def test_llm_labels(): assert rc.llm_label("ollama") == Label(NS, CON) assert rc.llm_label("anthropic") == Label(NS, EXP) assert rc.llm_label("totally-unknown") == Label(NS, EXP) # fail closed def test_embeddings_labels(): assert rc.embeddings_label("ollama") == Label(NS, CON) assert rc.embeddings_label("sentence_transformers") == Label(NS, CON) assert rc.embeddings_label("openai") == Label(NS, EXP) # --------------------------------------------------------------------------- # classify_run — assembly into SOURCE + SEARCH_SINK + INFERENCE_SINK # --------------------------------------------------------------------------- def test_classify_run_assembles_roles(): with patch.object(rc, "engine_label", return_value=Label(NS, EXP)): comps = rc.classify_run( {}, ctx(), engines=["arxiv"], llm_provider="anthropic" ) got = {(c.name, c.role) for c in comps} assert got == { ("arxiv", Role.SOURCE), ("arxiv", Role.SEARCH_SINK), ("llm:anthropic", Role.INFERENCE_SINK), } def test_classify_run_fails_closed_on_unknown_engine(): with patch.object(rc, "engine_label", return_value=None): comps = rc.classify_run( {}, ctx(), engines=["mystery"], llm_provider="ollama" ) # Unknown engine is NOT dropped — it fails closed to sensitive+exposing. by = {(c.name, c.role): c for c in comps} assert by[("mystery", Role.SOURCE)].label == Label(S, EXP) assert ("mystery", Role.SEARCH_SINK) in by assert ("llm:ollama", Role.INFERENCE_SINK) in by # --------------------------------------------------------------------------- # audit_run — decision + never-raises contract # --------------------------------------------------------------------------- def test_audit_run_denies_sensitive_plus_cloud_llm(): with patch.object(rc, "engine_label", return_value=Label(S, CON)): d = rc.audit_run( {}, ctx(), engines=["collection_x"], llm_provider="anthropic" ) assert d is not None and not d.allowed assert d.reason == "sensitive_to_exposing_inference" def test_audit_run_allows_public_plus_local_llm(): with patch.object(rc, "engine_label", return_value=Label(NS, EXP)): d = rc.audit_run({}, ctx(), engines=["arxiv"], llm_provider="ollama") assert d is not None and d.allowed def test_audit_run_permissive_under_unprotected_scope(): # A normally-denied combo is ALLOWED (permissive) under the escape hatch, # with the diagnostic preserved. unprot = EgressContext( scope=EgressScope.UNPROTECTED, primary_engine="arxiv", require_local_llm=False, require_local_embeddings=False, ) with patch.object(rc, "engine_label", return_value=Label(S, CON)): d = rc.audit_run( {}, unprot, engines=["collection_x"], llm_provider="anthropic" ) assert d is not None and d.allowed assert d.reason.startswith("permissive:") def test_audit_run_fails_closed_on_internal_error(): # An uncomputable decision must refuse the run (fail closed) and never # raise — silent degradation to the scope PEPs would hide the failure. with patch.object(rc, "engine_label", side_effect=RuntimeError("boom")): d = rc.audit_run({}, ctx(), engines=["x"], llm_provider="ollama") assert d is not None and not d.allowed and d.reason == "audit_error" def test_fail_closed_decision_needs_no_policy_import(): # _fail_closed_decision runs from the except blocks that handle failures # (including a .policy import failure), so it must NOT import .policy. It # denies for a normal ctx and allows only under unprotected, using the # ctx's str-enum scope value directly. assert not rc._fail_closed_decision(ctx()).allowed unprot = EgressContext( scope=EgressScope.UNPROTECTED, primary_engine="arxiv", require_local_llm=False, require_local_embeddings=False, ) assert rc._fail_closed_decision(unprot).allowed def test_audit_run_internal_error_never_blocks_under_unprotected(): # Fail-closed must not override the escape hatch: an internal error under # UNPROTECTED is still allowed. unprot = EgressContext( scope=EgressScope.UNPROTECTED, primary_engine="arxiv", require_local_llm=False, require_local_embeddings=False, ) with patch.object(rc, "engine_label", side_effect=RuntimeError("boom")): d = rc.audit_run({}, unprot, engines=["x"], llm_provider="ollama") assert d is not None and d.allowed def test_audit_run_unknown_engine_fails_closed_and_denies(): # Unknown engine -> sensitive+exposing source; a cloud LLM -> exposing # inference sink -> the run is denied (not silently allowed). with patch.object(rc, "engine_label", return_value=None): d = rc.audit_run( {}, ctx(), engines=["mystery"], llm_provider="anthropic" ) assert d is not None and not d.allowed assert d.reason == "sensitive_to_exposing_inference" # --------------------------------------------------------------------------- # audit_run_from_snapshot — shared route/worker entry point (covers all run # entry points: API precheck, follow-up, chat, queue) # --------------------------------------------------------------------------- def test_audit_run_from_snapshot_denies_private_collection_plus_cloud_llm(): # The worker-chokepoint case that a follow-up / chat / queue run reaches: # a private collection primary + a cloud LLM under a `both` scope must be # refused, matching the /api/start_research route. snap = {"llm.provider": "anthropic"} with patch.object(rc, "engine_label", return_value=Label(S, CON)): d = rc.audit_run_from_snapshot( snap, ctx(primary="collection_x"), "collection_x" ) assert d is not None and not d.allowed assert d.reason == "sensitive_to_exposing_inference" def test_audit_run_from_snapshot_unwraps_dict_wrapped_snapshot(): # get_all_settings snapshots store each value as {"value": ...}; the helper # must unwrap it, else the LLM provider is dropped and the check is blind. snap = {"llm.provider": {"value": "anthropic"}} with patch.object(rc, "engine_label", return_value=Label(S, CON)): d = rc.audit_run_from_snapshot( snap, ctx(primary="collection_x"), "collection_x" ) assert d is not None and not d.allowed def test_audit_run_from_snapshot_permissive_under_unprotected(): unprot = EgressContext( scope=EgressScope.UNPROTECTED, primary_engine="collection_x", require_local_llm=False, require_local_embeddings=False, ) with patch.object(rc, "engine_label", return_value=Label(S, CON)): d = rc.audit_run_from_snapshot( {"llm.provider": "anthropic"}, unprot, "collection_x" ) assert d is not None and d.allowed def test_audit_run_from_snapshot_missing_llm_provider_uses_run_default(): # A snapshot without llm.provider must resolve to the run's own default # (ollama, local) via get_setting_from_snapshot — the same value the run # would use — not be silently read as "no LLM". Sensitive source + local # default inference is quadrant 2 (allowed), matching the run's behaviour. with patch.object(rc, "engine_label", return_value=Label(S, CON)): d = rc.audit_run_from_snapshot( {}, ctx(primary="collection_x"), "collection_x" ) assert d is not None and d.allowed def test_audit_run_from_snapshot_includes_embeddings_for_collection(): # A RAG run over a collection embeds, so a cloud embedder is an exposing # sink for the sensitive source -> refused even with a local LLM. The RAG # engine reads local_search_embedding_provider, so the audit must too. snap = { "llm.provider": "ollama", "local_search_embedding_provider": "openai", } with patch.object(rc, "engine_label", return_value=Label(S, CON)): d = rc.audit_run_from_snapshot( snap, ctx(primary="collection_x"), "collection_x" ) assert d is not None and not d.allowed def test_audit_run_from_snapshot_ignores_embeddings_for_lexical_store(): # A non-collection sensitive store (paperless/elasticsearch) never embeds, # so a configured cloud embedder must NOT be pulled in: sensitive source + # local LLM is admissible (quadrant 2), not refused. snap = { "llm.provider": "ollama", "local_search_embedding_provider": "openai", } with patch.object(rc, "engine_label", return_value=Label(S, CON)): d = rc.audit_run_from_snapshot( snap, ctx(primary="paperless"), "paperless" ) assert d is not None and d.allowed def test_audit_run_from_snapshot_empty_llm_provider_resolves_from_snapshot(): # Chat passes model_provider="" (not None) when no override is set. An empty # string must be treated as "resolve from snapshot", not passed through — # else classify_run's truthy guard silently drops the LLM sink and a # sensitive source + cloud LLM is wrongly allowed. snap = {"llm.provider": "anthropic"} # cloud provider in the snapshot with patch.object(rc, "engine_label", return_value=Label(S, CON)): d = rc.audit_run_from_snapshot( snap, ctx(primary="collection_x"), "collection_x", llm_provider="" ) assert d is not None and not d.allowed assert d.reason == "sensitive_to_exposing_inference" def test_audit_run_from_snapshot_audits_llm_provider_override(): # The run's per-request model_provider WINS over the snapshot (get_llm uses # the explicit arg). A saved local default must not mask a cloud override # chosen for this one run — else queue/follow-up/chat runs leak. snap = {"llm.provider": "ollama"} # saved default is local/contained with patch.object(rc, "engine_label", return_value=Label(S, CON)): d = rc.audit_run_from_snapshot( snap, ctx(primary="collection_x"), "collection_x", llm_provider="anthropic", # per-request cloud override ) assert d is not None and not d.allowed assert d.reason == "sensitive_to_exposing_inference" def test_audit_run_from_snapshot_fails_closed_on_bad_snapshot(): # Provider resolution runs inside the fail-closed guard: a malformed # snapshot refuses the run rather than escaping the helper uncaught. d = rc.audit_run_from_snapshot( "not-a-dict", ctx(primary="collection_x"), "collection_x" ) assert d is not None and not d.allowed and d.reason == "audit_error" # --------------------------------------------------------------------------- # provider endpoint URL fail-up (only ever tightens) # --------------------------------------------------------------------------- def test_llm_label_local_custom_endpoint_is_contained(): # A self-hosted OpenAI-compatible endpoint on a local IP must be CONTAINED # (regression: previously mislabeled EXPOSING -> false 400 on local runs). snap = {"llm.openai_endpoint.url": "http://127.0.0.1:8000/v1"} assert rc.llm_label("openai_endpoint", snap, ctx()) == Label(NS, CON) def test_llm_label_public_custom_endpoint_is_exposing(): snap = {"llm.openai_endpoint.url": "http://8.8.8.8:8000/v1"} assert rc.llm_label("openai_endpoint", snap, ctx()) == Label(NS, EXP) # --------------------------------------------------------------------------- # Real declared label values — guards against a typo in any engine class # --------------------------------------------------------------------------- def test_registered_engines_declare_consistent_labels(): from local_deep_research.security.egress.policy import _get_engine_class from local_deep_research.web_search_engines.engine_registry import ( ENGINE_REGISTRY, ) assert ENGINE_REGISTRY, "engine registry unexpectedly empty" checked = 0 for name in ENGINE_REGISTRY: cls = _get_engine_class(name) # the real class the resolver reads if cls is None: continue sens = getattr(cls, "egress_sensitivity", None) exp = getattr(cls, "egress_exposure", None) assert sens is not None, f"{name} missing egress_sensitivity" assert exp is not None, f"{name} missing egress_exposure" if getattr(cls, "is_public", False): assert sens is Sensitivity.NON_SENSITIVE, name assert exp is Exposure.EXPOSING, name elif getattr(cls, "is_local", False): assert sens is Sensitivity.SENSITIVE, name assert exp is Exposure.CONTAINED, name checked += 1 assert checked > 10, f"only checked {checked} engines" def test_llm_provider_declared_exposure_matches_resolver(): # The resolver classifies a provider's endpoint via the enforcing PEP; with # no configured URL it falls back to the static cloud/local split. Assert # the resolver AND (where the class is resolvable) the declared attribute # agree, across every statically-classifiable provider — url-configurable # ones (openai_endpoint / anthropic_endpoint) are refined by URL and are # covered by the dedicated local/public endpoint tests above. from local_deep_research.llm.providers import get_provider_class expected = { "ollama": Exposure.CONTAINED, "lmstudio": Exposure.CONTAINED, "llamacpp": Exposure.CONTAINED, "openai": Exposure.EXPOSING, "anthropic": Exposure.EXPOSING, "google": Exposure.EXPOSING, "openrouter": Exposure.EXPOSING, "deepseek": Exposure.EXPOSING, "xai": Exposure.EXPOSING, "ionos": Exposure.EXPOSING, } for name, exp in expected.items(): assert rc.llm_label(name).exposure is exp, name cls = None try: cls = get_provider_class(name) or get_provider_class(name.upper()) except Exception: cls = None if cls is not None: assert getattr(cls, "egress_exposure", None) is exp, name # --------------------------------------------------------------------------- # Per-destination trust (stage D) — relaxes an exposing sink to contained # --------------------------------------------------------------------------- def test_llm_trust_relaxes_cloud_provider_to_contained(): trusted = {"policy.trusted_inference_providers": ["anthropic"]} assert rc.llm_label("anthropic", trusted, ctx()).exposure is CON # untrusted cloud stays exposing assert rc.llm_label("anthropic", {}, ctx()).exposure is EXP def test_embeddings_trust_relaxes_cloud_provider_to_contained(): trusted = {"policy.trusted_inference_providers": ["openai"]} assert rc.embeddings_label("openai", trusted, ctx()).exposure is CON def test_engine_trust_relaxes_exposing_store_to_contained(): # A sensitive store on a public URL is quadrant 4 (exposing); trusting it # moves exposure to contained while keeping its data sensitive. trusted = {"policy.trusted_search_engines": ["paperless"]} with ( patch(f"{POLICY}._get_engine_class", return_value=_LocalStore), patch(f"{POLICY}._classify_engine_url", return_value=False), ): assert rc.engine_label("paperless", trusted, ctx()) == Label(S, CON) def test_trusted_names_handles_json_string_and_garbage(): assert rc._trusted_names("k", {"k": '["Anthropic", "ollama"]'}) == { "anthropic", "ollama", } assert rc._trusted_names("k", {"k": "not json"}) == set() assert rc._trusted_names("k", None) == set() def test_trusted_cloud_llm_allows_sensitive_run_end_to_end(): # The stage-D payoff: a sensitive source + a TRUSTED cloud LLM is allowed. trusted = {"policy.trusted_inference_providers": ["anthropic"]} with patch.object(rc, "engine_label", return_value=Label(S, CON)): d = rc.audit_run( trusted, ctx(), engines=["collection_x"], llm_provider="anthropic" ) assert d is not None and d.allowed def test_engine_trust_does_not_relax_public_engine(): # Trusting an inherently-public engine must NOT contain it (no laundering # a public search sink past the sensitive->exposing rule). class _PublicWithName: is_public = True egress_sensitivity = NS egress_exposure = EXP trusted = {"policy.trusted_search_engines": ["searxng"]} with patch(f"{POLICY}._get_engine_class", return_value=_PublicWithName): assert rc.engine_label("searxng", trusted, ctx()) == Label(NS, EXP)