feat: add neutral auth store boundary

2026-06-06 10:18:16 +07:00
parent 71facde1a2
commit 34726433a8
15 changed files with 456 additions and 41 deletions
@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import assert_type
+
+import pytest
+
+from wf_api.auth import AuthRecord, AuthStore, validate_auth_id
+
+
+def test_validate_auth_id_accepts_safe_dotted_ids() -> None:
+    assert validate_auth_id("github.work") == "github.work"
+    assert validate_auth_id("api_ci-1") == "api_ci-1"
+
+
+@pytest.mark.parametrize("auth_id", ["", ".hidden", "../secret", "bad/id"])
+def test_validate_auth_id_rejects_unsafe_ids(auth_id: str) -> None:
+    with pytest.raises(ValueError, match="auth id must start"):
+        validate_auth_id(auth_id)
+
+
+def test_auth_record_is_immutable_and_mapping_typed() -> None:
+    record = AuthRecord(
+        id="github.work",
+        scheme="bearer",
+        payload={"token": "secret"},
+        metadata={"owner": "test"},
+    )
+
+    assert record.id == "github.work"
+    assert record.scheme == "bearer"
+    assert record.payload["token"] == "secret"
+    assert_type(record.payload, Mapping[str, object])
+
+    with pytest.raises(AttributeError):
+        record.scheme = "headers"  # type: ignore[misc]
+
+
+class MemoryAuthStore:
+    def __init__(self, records: dict[str, AuthRecord]) -> None:
+        self.records = records
+
+    def load_auth(self, auth_ref: str) -> AuthRecord | None:
+        return self.records.get(auth_ref)
+
+
+def test_auth_store_protocol_is_read_only_lookup() -> None:
+    record = AuthRecord(id="github.work", scheme="opaque", payload={"x": 1})
+    store: AuthStore = MemoryAuthStore({"github.work": record})
+
+    assert store.load_auth("github.work") is record
+    assert store.load_auth("missing") is None
@@ -1,5 +1,7 @@
 from __future__ import annotations

+from pathlib import Path
+
 from wf_artifacts import WorkflowDeployment
 from wf_platform import CapabilityBuckets, CapabilitySource, SourcePermissions

@@ -16,6 +18,14 @@ from ..test_support import FakeAdapter, local_temp_root
 from ..workflow_surface.conftest import echo_artifact


+def _transport(root: Path) -> UpstreamTransportService:
+    events: list[McpEvent] = []
+    return UpstreamTransportService(
+        store=FileStore(root),
+        event_sink=events.append,
+    )
+
+
 def test_upstream_transport_registers_adapter() -> None:
    events: list[McpEvent] = []
    transport = UpstreamTransportService(
@@ -98,7 +108,7 @@ async def test_upstream_transport_refreshes_catalog_directly() -> None:
        connection_list_enabled=connections.list_enabled,
        connection_list_all=connections.list_all,
        tool_executor_for=transport.tool_executor_for,
-        load_auth=transport.load_auth,
+        load_auth=transport.load_connection_auth,
        emit_event=events.append,
    )
    source_catalog.hydrate_connection_source_from_snapshot(connection)
@@ -131,7 +141,7 @@ async def test_upstream_transport_live_diagnostics_report_missing_connection() -
        connection_list_enabled=lambda: [],
        connection_list_all=lambda: [],
        tool_executor_for=transport.tool_executor_for,
-        load_auth=transport.load_auth,
+        load_auth=transport.load_connection_auth,
        emit_event=lambda event: None,
    )
    source_catalog.register_capability_source(
@@ -158,3 +168,82 @@ async def test_upstream_transport_live_diagnostics_report_missing_connection() -

    assert diagnostics[0].code == "source_unreachable"
    assert diagnostics[0].bound_source == "demo.personal"
+
+
+def test_upstream_load_connection_auth_prefers_auth_ref(tmp_path: Path) -> None:
+    service = _transport(tmp_path)
+    service.save_auth(
+        AuthRecord(
+            connection_id="github.creds",
+            scheme="bearer",
+            payload={"token": "secret"},
+        )
+    )
+    service.save_auth(
+        AuthRecord(
+            connection_id="github.work",
+            scheme="bearer",
+            payload={"token": "wrong"},
+        )
+    )
+    connection = ConnectionConfig(
+        id="github.work",
+        server="github",
+        account="work",
+        metadata={"auth_ref": "github.creds"},
+    )
+
+    assert service.load_connection_auth(connection) == AuthRecord(
+        connection_id="github.creds",
+        scheme="bearer",
+        payload={"token": "secret"},
+    )
+
+
+def test_upstream_load_connection_auth_falls_back_to_connection_id(
+    tmp_path: Path,
+) -> None:
+    service = _transport(tmp_path)
+    service.save_auth(
+        AuthRecord(
+            connection_id="github.work",
+            scheme="bearer",
+            payload={"token": "legacy"},
+        )
+    )
+    connection = ConnectionConfig(
+        id="github.work",
+        server="github",
+        account="work",
+    )
+
+    assert service.load_connection_auth(connection) == AuthRecord(
+        connection_id="github.work",
+        scheme="bearer",
+        payload={"token": "legacy"},
+    )
+
+
+def test_upstream_load_connection_auth_ignores_non_string_auth_ref(
+    tmp_path: Path,
+) -> None:
+    service = _transport(tmp_path)
+    service.save_auth(
+        AuthRecord(
+            connection_id="github.work",
+            scheme="bearer",
+            payload={"token": "legacy"},
+        )
+    )
+    connection = ConnectionConfig(
+        id="github.work",
+        server="github",
+        account="work",
+        metadata={"auth_ref": 123},
+    )
+
+    assert service.load_connection_auth(connection) == AuthRecord(
+        connection_id="github.work",
+        scheme="bearer",
+        payload={"token": "legacy"},
+    )
@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from wf_api.auth import AuthRecord as NeutralAuthRecord
+from wf_mcp.auth import (
+    mcp_auth_env,
+    mcp_auth_headers,
+    mcp_auth_from_neutral,
+    neutral_auth_from_mcp,
+)
+from wf_mcp.models import AuthRecord as McpAuthRecord
+from wf_mcp.storage import FileStore
+
+
+def test_mcp_auth_from_neutral_preserves_scheme_and_payload() -> None:
+    neutral = NeutralAuthRecord(
+        id="github.work",
+        scheme="bearer",
+        payload={"token": "secret"},
+        metadata={"owner": "test"},
+    )
+
+    mcp = mcp_auth_from_neutral(neutral)
+
+    assert mcp == McpAuthRecord(
+        connection_id="github.work",
+        scheme="bearer",
+        payload={"token": "secret"},
+    )
+
+
+def test_neutral_auth_from_mcp_preserves_payload() -> None:
+    mcp = McpAuthRecord(
+        connection_id="github.work",
+        scheme="headers",
+        payload={"headers": {"X-Test": "yes"}},
+    )
+
+    neutral = neutral_auth_from_mcp(mcp)
+
+    assert neutral.id == "github.work"
+    assert neutral.scheme == "headers"
+    assert neutral.payload == {"headers": {"X-Test": "yes"}}
+
+
+def test_mcp_auth_headers_extracts_explicit_headers_and_bearer_token() -> None:
+    auth = McpAuthRecord(
+        connection_id="api.work",
+        scheme="bearer",
+        payload={"headers": {"X-Test": "yes"}, "token": "secret"},
+    )
+
+    assert mcp_auth_headers(auth) == {
+        "X-Test": "yes",
+        "Authorization": "Bearer secret",
+    }
+
+
+def test_mcp_auth_headers_does_not_override_authorization_header() -> None:
+    auth = McpAuthRecord(
+        connection_id="api.work",
+        scheme="bearer",
+        payload={
+            "headers": {"Authorization": "Basic already"},
+            "token": "secret",
+        },
+    )
+
+    assert mcp_auth_headers(auth) == {"Authorization": "Basic already"}
+
+
+def test_mcp_auth_env_returns_string_map_only() -> None:
+    auth = McpAuthRecord(
+        connection_id="mcp.local",
+        scheme="env",
+        payload={"env": {"TOKEN": "secret", "BAD": 123}},
+    )
+
+    assert mcp_auth_env(auth) == {"TOKEN": "secret"}
+
+
+def test_file_store_saves_and_loads_neutral_auth_record(tmp_path: Path) -> None:
+    store = FileStore(tmp_path)
+    record = NeutralAuthRecord(
+        id="github.work",
+        scheme="bearer",
+        payload={"token": "secret"},
+        metadata={"owner": "test"},
+    )
+
+    store.save_auth_record(record)
+
+    loaded = store.load_auth_record("github.work")
+    assert loaded is not None
+    assert loaded.id == record.id
+    assert loaded.scheme == record.scheme
+    assert loaded.payload == record.payload
+    # Legacy file format does not persist metadata
+    assert loaded.metadata == {}
+
+
+def test_file_store_legacy_auth_methods_still_work(tmp_path: Path) -> None:
+    store = FileStore(tmp_path)
+    legacy = McpAuthRecord(
+        connection_id="github.work",
+        scheme="bearer",
+        payload={"token": "secret"},
+    )
+
+    store.save_auth(legacy)
+
+    assert store.load_auth("github.work") == legacy
+    assert store.load_auth_record("github.work") == NeutralAuthRecord(
+        id="github.work",
+        scheme="bearer",
+        payload={"token": "secret"},
+    )