from __future__ import annotations import asyncio import json import unittest from collections.abc import Awaitable, Callable from types import SimpleNamespace from typing import Any, cast from services.runtime.common import identity_assertion from services.runtime.connector_ingestion import ( _invoke, build_actor_resolver, build_server, load_hub_project_ids, load_workspace_read_enabled, ) from services.runtime.dispatch import service_module from services.runtime.ingestion_auth import ActorContext from services.runtime.ingestion_files import FileIngestionError from services.runtime.ingestion_workflow import IngestionWorkflow class _Workflow: async def search_knowledge(self, **_: Any) -> dict[str, Any]: return {} async def get_knowledge_context(self, **_: Any) -> dict[str, Any]: return {} async def analyze_knowledge_files(self, **_: Any) -> dict[str, Any]: return {} async def submit_ingestion_proposal(self, **_: Any) -> dict[str, Any]: return {} async def get_ingestion_status(self, **_: Any) -> dict[str, Any]: return {} async def extract_knowledge_files(self, **_: Any) -> dict[str, Any]: return {} async def get_workspace_roots(self, **_: Any) -> dict[str, Any]: return {} async def get_workspace_children(self, **_: Any) -> dict[str, Any]: return {} async def get_workspace_page(self, **_: Any) -> dict[str, Any]: return {} async def create_text_source_proposal(self, **_: Any) -> dict[str, Any]: return {} async def _actor_resolver(_: Any) -> ActorContext: return ActorContext( subject="connector-client", projects=frozenset({"11111111-1111-4111-8111-111111111111"}), scopes=frozenset({"knowledge:read", "ingestion:write"}), correlation_id="trace-test", ) class IngestionConnectorTests(unittest.TestCase): @staticmethod def _resolve( resolver: Callable[[Any], Awaitable[ActorContext]], context: Any, ) -> ActorContext: async def resolve() -> ActorContext: return await resolver(context) return asyncio.run(resolve()) def _context( self, *, project: str = "knowledge-hub", environment: str = "prod", signing_key: bytes = b"k" * 32, ) -> Any: client_id = "12345678-1234-4234-8234-123456789abc" trace_id = "trace-dataplane-1" headers = { "authorization": "Bearer upstream-contextforge-token-is-not-the-caller-token", "x-itworx-project": project, "x-itworx-environment": environment, "x-itworx-client-id": client_id, "x-itworx-trace-id": trace_id, } headers["x-itworx-identity-assertion"] = identity_assertion( signing_key, project=project, environment=environment, client_id=client_id, trace_id=trace_id, ) context = SimpleNamespace( request_context=SimpleNamespace( request=SimpleNamespace(headers=headers), ) ) return signing_key, context def test_dedicated_dispatch_name_resolves_to_ingestion_runtime(self) -> None: self.assertEqual( service_module("connector-ingestion"), "services.runtime.connector_ingestion", ) def test_server_exposes_exactly_ten_ingestion_tools(self) -> None: server = build_server( workflow=cast(IngestionWorkflow, _Workflow()), actor_resolver=_actor_resolver, ) tools = asyncio.run(server.list_tools()) self.assertEqual( {tool.name for tool in tools}, { "search_knowledge", "get_knowledge_context", "analyze_knowledge_files", "extract_knowledge_files", "get_workspace_roots", "get_workspace_children", "get_workspace_page", "create_text_source_proposal", "submit_ingestion_proposal", "get_ingestion_status", }, ) def test_file_params_declares_chatgpt_files_array_contract(self) -> None: server = build_server( workflow=cast(IngestionWorkflow, _Workflow()), actor_resolver=_actor_resolver, ) tools = asyncio.run(server.list_tools()) analyze = next( tool.model_dump(by_alias=True) for tool in tools if tool.name == "analyze_knowledge_files" ) schema = analyze["inputSchema"]["properties"]["files"] file_schema = schema["items"] self.assertEqual(analyze["_meta"]["openai/fileParams"], ["files"]) self.assertEqual( set(file_schema["properties"]), {"download_url", "file_id", "mime_type", "file_name"}, ) self.assertEqual(set(file_schema["required"]), {"download_url", "file_id"}) self.assertEqual(file_schema["properties"]["download_url"]["type"], "string") self.assertEqual(file_schema["properties"]["file_id"]["type"], "string") self.assertEqual(file_schema["properties"]["mime_type"]["type"], "string") self.assertEqual(file_schema["properties"]["file_name"]["type"], "string") extract = next( tool.model_dump(by_alias=True) for tool in tools if tool.name == "extract_knowledge_files" ) self.assertEqual(extract["_meta"]["openai/fileParams"], ["files"]) def test_advertised_input_schemas_equal_governed_catalog(self) -> None: from pathlib import Path server = build_server( workflow=cast(IngestionWorkflow, _Workflow()), actor_resolver=_actor_resolver, ) actual = { tool.name: tool.model_dump(by_alias=True)["inputSchema"] for tool in asyncio.run(server.list_tools()) } catalog = Path(__file__).resolve().parents[1] / "catalog" / "tools" for path in catalog.glob("ingestion-*.yaml"): spec = json.loads(path.read_text(encoding="utf-8"))["spec"] self.assertEqual(actual[spec["connectorToolName"]], spec["inputSchema"]) for name in ( "search_knowledge", "extract_knowledge_files", "analyze_knowledge_files", "create_text_source_proposal", "submit_ingestion_proposal", "get_ingestion_status", ): self.assertIn("project_id", actual[name]["properties"]) self.assertNotIn("project_id", actual[name].get("required", [])) self.assertIn("project_id", actual["get_knowledge_context"]["required"]) def test_write_tools_are_non_destructive_and_closed_world(self) -> None: server = build_server( workflow=cast(IngestionWorkflow, _Workflow()), actor_resolver=_actor_resolver, ) tools = { tool.name: tool.model_dump(by_alias=True) for tool in asyncio.run(server.list_tools()) } for name in ("analyze_knowledge_files", "submit_ingestion_proposal"): annotations = tools[name]["annotations"] self.assertFalse(annotations["readOnlyHint"]) self.assertFalse(annotations["destructiveHint"]) self.assertFalse(annotations["openWorldHint"]) self.assertTrue(tools["analyze_knowledge_files"]["annotations"]["idempotentHint"]) self.assertFalse(tools["submit_ingestion_proposal"]["annotations"]["idempotentHint"]) def test_file_error_audit_records_only_bounded_code_and_correlation(self) -> None: async def reject_file(**_: Any) -> dict[str, Any]: raise FileIngestionError( "invalid_input", input_reason="attachment_file_id_invalid" ) context = self._context() private_file_values = [ { "download_url": "private-download-reference", "file_id": "private-file-reference", } ] with self.assertLogs("services.runtime.connector_ingestion", level="INFO") as captured: result = asyncio.run( _invoke( "extract_knowledge_files", context, _actor_resolver, reject_file, files=private_file_values, ) ) self.assertEqual(result["error"]["code"], "invalid_input") audit_message = captured.records[0].getMessage().removeprefix("ingestion_tool ") audit = json.loads(audit_message) self.assertEqual( set(audit), { "tool", "result", "status", "latency_bucket", "correlation_id", "error_code", "input_reason", }, ) self.assertEqual(audit["error_code"], "invalid_input") self.assertEqual(audit["input_reason"], "attachment_file_id_invalid") self.assertEqual(audit["correlation_id"], "trace-test") self.assertEqual(audit["tool"], "extract_knowledge_files") self.assertEqual(audit["result"], "failure") self.assertNotIn("private-download-reference", audit_message) self.assertNotIn("private-file-reference", audit_message) def test_dataplane_signed_assertion_is_trusted_and_authorization_is_ignored(self) -> None: signing_key, context = self._context() project_id = "11111111-1111-4111-8111-111111111111" resolver = build_actor_resolver( allowed_projects=frozenset({project_id}), signing_key=signing_key, ) actor = self._resolve(resolver, context) self.assertEqual(actor.subject, "12345678-1234-4234-8234-123456789abc") self.assertEqual(actor.projects, frozenset({project_id})) self.assertEqual(actor.scopes, frozenset({"knowledge:read", "ingestion:write"})) self.assertEqual(actor.correlation_id, "trace-dataplane-1") self.assertEqual(actor.connector_project, "knowledge-hub") self.assertEqual(actor.connector_environment, "prod") def test_dataplane_assertion_rejects_invalid_signature_and_nonproduction(self) -> None: signing_key, context = self._context() headers = context.request_context.request.headers headers["x-itworx-identity-assertion"] = "v1:0:" + "0" * 64 resolver = build_actor_resolver( allowed_projects=frozenset({"11111111-1111-4111-8111-111111111111"}), signing_key=signing_key, ) with self.assertRaises(PermissionError): self._resolve(resolver, context) signing_key, staging = self._context(environment="staging") resolver = build_actor_resolver( allowed_projects=frozenset({"11111111-1111-4111-8111-111111111111"}), signing_key=signing_key, ) with self.assertRaises(PermissionError): self._resolve(resolver, staging) def test_shared_connector_key_cannot_forge_intake_identity_assertion(self) -> None: shared_connector_key = b"s" * 32 _, shared_key_signed_context = self._context(signing_key=shared_connector_key) dedicated_intake_key = b"i" * 32 resolver = build_actor_resolver( allowed_projects=frozenset({"11111111-1111-4111-8111-111111111111"}), signing_key=dedicated_intake_key, ) with self.assertRaises(PermissionError): self._resolve(resolver, shared_key_signed_context) def test_target_project_allowlist_is_bounded_uuid_configuration(self) -> None: project_id = "11111111-1111-4111-8111-111111111111" self.assertEqual( load_hub_project_ids(environ={"KNOWLEDGE_HUB_ALLOWED_PROJECT_IDS": f" {project_id} "}), frozenset({project_id}), ) for value in ("", "not-a-uuid", f"{project_id},{project_id}", ",".join([project_id] * 257)): with self.subTest(value=value[:40]): with self.assertRaises(ValueError): load_hub_project_ids(environ={"KNOWLEDGE_HUB_ALLOWED_PROJECT_IDS": value}) def test_empty_allowlist_requires_workspace_opt_in_and_signed_exact_identity(self) -> None: self.assertEqual(load_hub_project_ids(environ={}, workspace_enabled=True), frozenset()) signing_key, context = self._context() resolver = build_actor_resolver( allowed_projects=frozenset(), signing_key=signing_key, workspace_enabled=True ) actor = self._resolve(resolver, context) self.assertEqual(actor.projects, frozenset()) with self.assertRaises(ValueError): build_actor_resolver(allowed_projects=frozenset(), signing_key=signing_key) _, staging = self._context(environment="staging") with self.assertRaises(PermissionError): self._resolve(resolver, staging) def test_workspace_read_configuration_is_default_off_and_strict(self) -> None: self.assertFalse(load_workspace_read_enabled(environ={})) self.assertTrue( load_workspace_read_enabled(environ={"KNOWLEDGE_HUB_WORKSPACE_READ_ENABLED": "true"}) ) self.assertFalse( load_workspace_read_enabled(environ={"KNOWLEDGE_HUB_WORKSPACE_READ_ENABLED": "false"}) ) for value in ("yes", "1", "true please"): with self.subTest(value=value), self.assertRaises(ValueError): load_workspace_read_enabled(environ={"KNOWLEDGE_HUB_WORKSPACE_READ_ENABLED": value}) if __name__ == "__main__": unittest.main()