diff --git a/openviking/connector/__init__.py b/openviking/connector/__init__.py new file mode 100644 index 0000000000..236633478a --- /dev/null +++ b/openviking/connector/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 diff --git a/openviking/connector/client.py b/openviking/connector/client.py new file mode 100644 index 0000000000..073aa6878e --- /dev/null +++ b/openviking/connector/client.py @@ -0,0 +1,70 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Client for the external Connector service (knowledge-base doc/add pipeline).""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +import httpx + +from openviking_cli.utils import get_logger + +logger = get_logger(__name__) + + +class ConnectorClient: + """Wraps the Connector service's doc/add and task/info APIs.""" + + def __init__(self, doc_add_url: str, task_info_url: str, account_id: str = "") -> None: + self._doc_add_url = doc_add_url + self._task_info_url = task_info_url + self._headers = {"V-Account-Id": account_id} if account_id else {} + + async def submit_doc_add( + self, + resource_id: str, + add_type: str, + api_key: str, + *, + tos_path: Optional[str] = None, + path_prefix: Optional[List[str]] = None, + include_child: bool = True, + extra_params: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Submit a document import job via the configured doc/add endpoint. + + Returns the Connector response dict (contains task key / id on success). + """ + payload: Dict[str, Any] = { + "resource_id": resource_id, + "add_type": add_type, + "backend": "ov", + "api_key": api_key, + "include_child": include_child, + } + if tos_path is not None: + payload["tos_path"] = tos_path + if path_prefix is not None: + payload["path_prefix"] = path_prefix + if extra_params: + payload.update(extra_params) + + async with httpx.AsyncClient(timeout=30.0) as client: + rsp = await client.post(self._doc_add_url, json=payload, headers=self._headers) + rsp.raise_for_status() + return rsp.json() + + async def get_task_info(self, task_key: str) -> Dict[str, Any]: + """Query task status via the configured task/info endpoint. + + Connector task statuses: pending / running / succeeded / failed / cancelled. + """ + async with httpx.AsyncClient(timeout=10.0) as client: + rsp = await client.post( + self._task_info_url, + json={"TaskKey": task_key}, + headers=self._headers, + ) + rsp.raise_for_status() + return rsp.json() diff --git a/openviking/server/auth/__init__.py b/openviking/server/auth/__init__.py index 5f245b9010..593b6594c1 100644 --- a/openviking/server/auth/__init__.py +++ b/openviking/server/auth/__init__.py @@ -136,6 +136,8 @@ async def get_request_context( identity: ResolvedIdentity = Depends(resolve_identity), x_openviking_actor_peer: Optional[str] = Header(None, alias="X-OpenViking-Actor-Peer"), x_openviking_agent: Optional[str] = Header(None, alias="X-OpenViking-Agent"), + x_api_key_ctx: Optional[str] = Header(None, alias="X-API-Key"), + authorization_ctx: Optional[str] = Header(None, alias="Authorization"), ) -> RequestContext: """Convert ResolvedIdentity to RequestContext.""" path = request.url.path @@ -146,6 +148,8 @@ async def get_request_context( x_openviking_agent, ) + raw_api_key = _extract_api_key(x_api_key_ctx, authorization_ctx) + ctx = RequestContext( user=UserIdentifier( identity.account_id or "default", @@ -155,6 +159,7 @@ async def get_request_context( actor_peer_id=actor_peer_id, legacy_agent_id=legacy_agent_id, from_oauth=identity.from_oauth, + api_key=raw_api_key, ) # Update the unified root observability context after authentication succeeds. update_root_span_identity( diff --git a/openviking/server/identity.py b/openviking/server/identity.py index 87bb053ee6..b9495a97b3 100644 --- a/openviking/server/identity.py +++ b/openviking/server/identity.py @@ -105,6 +105,9 @@ class RequestContext: # prevent a stolen access token from laundering itself into a long-lived # refresh-token chain. from_oauth: bool = False + # Raw API key from the request — used by Connector to call back into OV + # on behalf of the original user (Plan 1: transparent key forwarding). + api_key: Optional[str] = None @property def account_id(self) -> str: diff --git a/openviking/service/resource_service.py b/openviking/service/resource_service.py index ac8dae0281..b6f94391af 100644 --- a/openviking/service/resource_service.py +++ b/openviking/service/resource_service.py @@ -16,6 +16,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, List, Optional +from openviking.connector.client import ConnectorClient from openviking.core.content_targets import ContentTargetSpec from openviking.core.uri_validation import validate_optional_content_target_uri from openviking.resource.feishu_watch_auth import ( @@ -54,6 +55,7 @@ from openviking_cli.exceptions import ( ConflictError, DeadlineExceededError, + InternalError, InvalidArgumentError, NotInitializedError, ) @@ -654,6 +656,24 @@ async def add_resource( if default_parent: parent = default_parent kwargs["create_parent"] = True + + if self._should_use_connector(args): + return await self._add_resource_via_connector( + path=path, + ctx=ctx, + to=to, + parent=parent, + reason=reason, + connector_args=args or {}, + watch_interval=watch_interval, + build_index=build_index, + summarize=summarize, + skip_watch_management=skip_watch_management, + stage_callback=stage_callback, + normalized_args=normalized_args, + **kwargs, + ) + if not wait and is_git_repo_url(path): return await self.enqueue_git_add_resource( path=path, @@ -1129,6 +1149,217 @@ async def _monitor_queue_processing( request_wait_tracker.cleanup(telemetry_id) unregister_wait_telemetry(telemetry_id) + # ── Connector routing ── + + def _should_use_connector(self, args: Optional[Dict[str, Any]]) -> bool: + """Decide whether to delegate this add_resource to the Connector service.""" + from openviking_cli.utils.config.open_viking_config import get_openviking_config + + config = get_openviking_config().connector + if not config.enable: + return False + if not args or "add_type" not in args: + return False + if args["add_type"] not in config.allowed_add_types: + logger.info( + f"[ResourceService] add_type '{args['add_type']}' not in " + f"connector.allowed_add_types {config.allowed_add_types}, " + "falling back to the standard pipeline" + ) + return False + return True + + async def _add_resource_via_connector( + self, + path: str, + ctx: RequestContext, + to: Optional[str], + parent: Optional[str], + reason: str, + connector_args: Dict[str, Any], + watch_interval: float, + build_index: bool, + summarize: bool, + skip_watch_management: bool, + stage_callback: Optional[Callable[[str], Any]], + normalized_args: _NormalizedAddResourceArgs, + **kwargs: Any, + ) -> Dict[str, Any]: + """Route add_resource to the external Connector service.""" + from openviking.service.task_tracker import get_task_tracker + + from openviking_cli.utils.config.open_viking_config import get_openviking_config + + config = get_openviking_config().connector + if not ctx.api_key: + raise InvalidArgumentError( + "Connector import requires an API key in the request." + ) + + target = ContentTargetSpec.from_fields( + ctx=ctx, + kind="resource", + to=to, + parent=parent, + create_parent=bool(kwargs.get("create_parent", False)), + ) + resource_id = target.to or path + + client = ConnectorClient( + doc_add_url=config.connector, + task_info_url=config.tracker, + account_id=config.main_account_id, + ) + + add_type = connector_args["add_type"] + extra_params: Dict[str, Any] = {} + reserved = {"add_type", "tos_path", "path_prefix", "include_child"} + for k, v in connector_args.items(): + if k not in reserved: + extra_params[k] = v + + result = await client.submit_doc_add( + resource_id=resource_id, + add_type=add_type, + api_key=ctx.api_key, + tos_path=connector_args.get("tos_path"), + path_prefix=connector_args.get("path_prefix"), + include_child=connector_args.get("include_child", True), + extra_params=extra_params or None, + ) + + connector_task_key = result.get("task_key") or result.get("TaskKey") or "" + if not connector_task_key: + raise InternalError( + f"Connector accepted the import but returned no task key: {result}" + ) + + task_tracker = get_task_tracker() + task = await task_tracker.create( + "connector_import", + resource_id=resource_id, + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) + + background = asyncio.create_task( + self._monitor_connector_task( + client=client, + connector_task_key=connector_task_key, + ov_task_id=task.task_id, + resource_id=resource_id, + poll_interval_ms=config.poll_interval_ms, + timeout_seconds=config.timeout_seconds, + ctx=ctx, + ) + ) + self._background_tasks.add(background) + background.add_done_callback(self._background_tasks.discard) + + if not skip_watch_management and watch_interval > 0: + watch_manager = self._get_watch_manager() + if watch_manager: + watch_to = target.to or resource_id + processor_kwargs = self._sanitize_watch_processor_kwargs(kwargs) + processor_kwargs["connector_args"] = connector_args + try: + await self._handle_watch_task_creation( + path=path, + to_uri=watch_to, + parent_uri=target.parent, + reason=reason, + instruction="", + watch_interval=watch_interval, + build_index=build_index, + summarize=summarize, + processor_kwargs=processor_kwargs, + auth_state=normalized_args.watch_auth_state, + ctx=ctx, + ) + except ConflictError: + raise + except Exception as e: + logger.warning( + f"[ResourceService] Failed to create watch task for connector import: {e}" + ) + + return { + "status": "accepted", + "task_id": task.task_id, + "connector_task_key": connector_task_key, + "resource_id": resource_id, + } + + async def _monitor_connector_task( + self, + client: ConnectorClient, + connector_task_key: str, + ov_task_id: str, + resource_id: str, + poll_interval_ms: int, + timeout_seconds: int, + ctx: RequestContext, + ) -> None: + """Poll the Connector task until terminal state, then update OV TaskRecord.""" + from openviking.service.task_tracker import get_task_tracker + + task_tracker = get_task_tracker() + await task_tracker.start( + ov_task_id, + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) + + poll_interval = poll_interval_ms / 1000.0 + deadline = time.perf_counter() + timeout_seconds + terminal_statuses = {"succeeded", "failed", "cancelled"} + + try: + while time.perf_counter() < deadline: + await asyncio.sleep(poll_interval) + info = await client.get_task_info(connector_task_key) + status = (info.get("Status") or info.get("status") or "").lower() + + await task_tracker.update_stage( + ov_task_id, + f"connector:{status}", + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) + + if status in terminal_statuses: + if status == "succeeded": + await task_tracker.complete( + ov_task_id, + {"connector_status": status, "connector_task_key": connector_task_key}, + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) + else: + error_msg = info.get("ErrorMessage") or info.get("error_message") or status + await task_tracker.fail( + ov_task_id, + f"connector task {status}: {error_msg}", + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) + return + + await task_tracker.fail( + ov_task_id, + f"connector task timed out after {timeout_seconds}s", + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) + except Exception as exc: + logger.error(f"[ResourceService] Connector task monitor error: {exc}") + await task_tracker.fail( + ov_task_id, + str(exc), + account_id=ctx.account_id, + user_id=ctx.user.user_id, + ) + async def _handle_watch_task_creation( self, path: str, diff --git a/openviking_cli/utils/config/open_viking_config.py b/openviking_cli/utils/config/open_viking_config.py index 45886023c0..3e37b77158 100644 --- a/openviking_cli/utils/config/open_viking_config.py +++ b/openviking_cli/utils/config/open_viking_config.py @@ -54,6 +54,39 @@ def _get_config_warning_logger(): return logging.getLogger(__name__) +class ConnectorConfig(BaseModel): + """Configuration for external Connector service.""" + + enable: bool = False + connector: str = "" + tracker: str = "" + main_account_id: str = "" + timeout_seconds: int = 3600 + poll_interval_ms: int = 5000 + allowed_add_types: List[str] = Field( + default_factory=lambda: ["tos"] + ) + + model_config = {"extra": "forbid"} + + @model_validator(mode="after") + def _validate(self) -> "ConnectorConfig": + if self.enable: + for name, url in (("connector", self.connector), ("tracker", self.tracker)): + if not url.strip(): + raise ValueError(f"connector.{name} is required when connector.enable=true") + if "://" not in url: + raise ValueError( + f"connector.{name} must be a full endpoint URL including scheme " + "(e.g., http://...)" + ) + if self.timeout_seconds <= 0: + raise ValueError("connector.timeout_seconds must be > 0") + if self.poll_interval_ms <= 0: + raise ValueError("connector.poll_interval_ms must be > 0") + return self + + class ParserApiConfig(BaseModel): """Configuration for the Understanding files/responses API.""" @@ -202,6 +235,11 @@ class OpenVikingConfig(BaseModel): description="Third-party parser API configuration (files/responses)", ) + connector: ConnectorConfig = Field( + default_factory=ConnectorConfig, + description="External Connector service configuration for data import", + ) + auto_generate_l0: bool = Field( default=True, description="Automatically generate L0 (abstract) if not provided" ) diff --git a/tests/connector/test_client.py b/tests/connector/test_client.py new file mode 100644 index 0000000000..df5e49668a --- /dev/null +++ b/tests/connector/test_client.py @@ -0,0 +1,144 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Unit tests for the external Connector HTTP client.""" + +import pytest + +from openviking.connector import client as client_module +from openviking.connector.client import ConnectorClient + + +class FakeResponse: + def __init__(self, payload): + self._payload = payload + self.raise_for_status_called = False + + def raise_for_status(self): + self.raise_for_status_called = True + + def json(self): + return self._payload + + +class FakeAsyncClient: + instances = [] + response_payload = {} + + def __init__(self, *, timeout): + self.timeout = timeout + self.posts = [] + FakeAsyncClient.instances.append(self) + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return None + + async def post(self, url, *, json, headers): + response = FakeResponse(FakeAsyncClient.response_payload) + self.posts.append( + { + "url": url, + "json": json, + "headers": headers, + "response": response, + } + ) + return response + + +@pytest.fixture(autouse=True) +def reset_fake_http_client(monkeypatch): + FakeAsyncClient.instances = [] + FakeAsyncClient.response_payload = {} + monkeypatch.setattr(client_module.httpx, "AsyncClient", FakeAsyncClient) + + +@pytest.mark.asyncio +async def test_submit_doc_add_builds_connector_payload_and_headers(): + FakeAsyncClient.response_payload = {"task_key": "connector-task-1"} + client = ConnectorClient( + doc_add_url="https://connector.example/api/knowledge/doc/add", + task_info_url="https://connector.example/api/task/info", + account_id="main-account", + ) + + result = await client.submit_doc_add( + resource_id="viking://resources/spec", + add_type="tos", + api_key="sk-test", + tos_path="tos://bucket/key", + path_prefix=["docs", "product"], + include_child=False, + extra_params={"parser": "pdf", "priority": "high"}, + ) + + assert result == {"task_key": "connector-task-1"} + http_client = FakeAsyncClient.instances[-1] + assert http_client.timeout == 30.0 + assert http_client.posts == [ + { + "url": "https://connector.example/api/knowledge/doc/add", + "json": { + "resource_id": "viking://resources/spec", + "add_type": "tos", + "backend": "ov", + "api_key": "sk-test", + "include_child": False, + "tos_path": "tos://bucket/key", + "path_prefix": ["docs", "product"], + "parser": "pdf", + "priority": "high", + }, + "headers": {"V-Account-Id": "main-account"}, + "response": http_client.posts[0]["response"], + } + ] + assert http_client.posts[0]["response"].raise_for_status_called is True + + +@pytest.mark.asyncio +async def test_submit_doc_add_omits_optional_none_values_and_empty_account_header(): + FakeAsyncClient.response_payload = {"TaskKey": "connector-task-2"} + client = ConnectorClient( + doc_add_url="https://connector.example/doc/add", + task_info_url="https://connector.example/task/info", + ) + + await client.submit_doc_add( + resource_id="resource-id", + add_type="tos", + api_key="sk-test", + ) + + post = FakeAsyncClient.instances[-1].posts[0] + assert post["headers"] == {} + assert post["json"] == { + "resource_id": "resource-id", + "add_type": "tos", + "backend": "ov", + "api_key": "sk-test", + "include_child": True, + } + + +@pytest.mark.asyncio +async def test_get_task_info_posts_task_key_to_tracker_endpoint(): + FakeAsyncClient.response_payload = {"Status": "running"} + client = ConnectorClient( + doc_add_url="https://connector.example/doc/add", + task_info_url="https://tracker.example/api/task/info", + account_id="main-account", + ) + + result = await client.get_task_info("connector-task-3") + + assert result == {"Status": "running"} + http_client = FakeAsyncClient.instances[-1] + assert http_client.timeout == 10.0 + post = http_client.posts[0] + assert post["url"] == "https://tracker.example/api/task/info" + assert post["json"] == {"TaskKey": "connector-task-3"} + assert post["headers"] == {"V-Account-Id": "main-account"} + assert post["response"].raise_for_status_called is True diff --git a/tests/service/test_resource_service_connector.py b/tests/service/test_resource_service_connector.py new file mode 100644 index 0000000000..4a192fefb2 --- /dev/null +++ b/tests/service/test_resource_service_connector.py @@ -0,0 +1,346 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Unit tests for ResourceService Connector routing.""" + +from types import SimpleNamespace + +import pytest + +from openviking.server.identity import RequestContext, Role +from openviking.service import resource_service as resource_service_module +from openviking.service.resource_service import ResourceService +from openviking_cli.exceptions import InvalidArgumentError +from openviking_cli.session.user_id import UserIdentifier + + +class RecordingResourceProcessor: + def __init__(self): + self.calls = [] + + async def process_resource(self, **kwargs): + self.calls.append(kwargs) + return {"root_uri": kwargs.get("to") or "viking://resources/standard"} + + +class RecordingTaskTracker: + def __init__(self): + self.created = [] + self.started = [] + self.stages = [] + self.completed = [] + self.failed = [] + + async def create(self, task_type, **kwargs): + self.created.append({"task_type": task_type, **kwargs}) + return SimpleNamespace(task_id=f"task-{len(self.created)}") + + async def start(self, task_id, **kwargs): + self.started.append({"task_id": task_id, **kwargs}) + + async def update_stage(self, task_id, stage, **kwargs): + self.stages.append({"task_id": task_id, "stage": stage, **kwargs}) + + async def complete(self, task_id, result, **kwargs): + self.completed.append({"task_id": task_id, "result": result, **kwargs}) + + async def fail(self, task_id, error, **kwargs): + self.failed.append({"task_id": task_id, "error": error, **kwargs}) + + +class FakeBackgroundTask: + def add_done_callback(self, callback): + self.callback = callback + + +class RecordingConnectorClient: + instances = [] + + def __init__(self, doc_add_url, task_info_url, account_id=""): + self.doc_add_url = doc_add_url + self.task_info_url = task_info_url + self.account_id = account_id + self.submit_calls = [] + RecordingConnectorClient.instances.append(self) + + async def submit_doc_add(self, **kwargs): + self.submit_calls.append(kwargs) + return {"TaskKey": "connector-task-1"} + + +@pytest.fixture +def request_context(): + return RequestContext( + user=UserIdentifier("acct", "alice"), + role=Role.USER, + api_key="sk-test", + ) + + +@pytest.fixture +def resource_processor(): + return RecordingResourceProcessor() + + +@pytest.fixture +def resource_service(resource_processor): + return ResourceService( + vikingdb=object(), + viking_fs=object(), + resource_processor=resource_processor, + skill_processor=object(), + ) + + +def install_connector_config( + monkeypatch, + *, + enable=True, + allowed_add_types=None, + poll_interval_ms=1, + timeout_seconds=1, +): + import openviking_cli.utils.config.open_viking_config as config_module + + connector = SimpleNamespace( + enable=enable, + connector="https://connector.example/api/knowledge/doc/add", + tracker="https://connector.example/api/task/info", + main_account_id="main-account", + timeout_seconds=timeout_seconds, + poll_interval_ms=poll_interval_ms, + allowed_add_types=allowed_add_types or ["tos"], + ) + monkeypatch.setattr( + config_module, + "get_openviking_config", + lambda: SimpleNamespace(connector=connector), + ) + return connector + + +def install_task_tracker(monkeypatch): + tracker = RecordingTaskTracker() + monkeypatch.setattr( + "openviking.service.task_tracker.get_task_tracker", + lambda: tracker, + ) + return tracker + + +def install_noop_background_tasks(monkeypatch): + def fake_create_task(coro): + if hasattr(coro, "close"): + coro.close() + return FakeBackgroundTask() + + monkeypatch.setattr(resource_service_module.asyncio, "create_task", fake_create_task) + + +@pytest.mark.asyncio +async def test_add_resource_routes_allowed_add_type_to_connector( + monkeypatch, + resource_service, + resource_processor, + request_context, +): + install_connector_config(monkeypatch, enable=True, allowed_add_types=["tos"]) + tracker = install_task_tracker(monkeypatch) + install_noop_background_tasks(monkeypatch) + RecordingConnectorClient.instances = [] + monkeypatch.setattr( + resource_service_module, + "ConnectorClient", + RecordingConnectorClient, + ) + + result = await resource_service.add_resource( + path="tos://bucket/input", + ctx=request_context, + to="viking://resources/spec", + reason="Import from Connector", + args={ + "add_type": "tos", + "tos_path": "tos://bucket/input", + "path_prefix": ["docs"], + "include_child": False, + "parser": "pdf", + }, + ) + + assert result == { + "status": "accepted", + "task_id": "task-1", + "connector_task_key": "connector-task-1", + "resource_id": "viking://resources/spec", + } + assert resource_processor.calls == [] + connector = RecordingConnectorClient.instances[0] + assert connector.doc_add_url == "https://connector.example/api/knowledge/doc/add" + assert connector.task_info_url == "https://connector.example/api/task/info" + assert connector.account_id == "main-account" + assert connector.submit_calls == [ + { + "resource_id": "viking://resources/spec", + "add_type": "tos", + "api_key": "sk-test", + "tos_path": "tos://bucket/input", + "path_prefix": ["docs"], + "include_child": False, + "extra_params": {"parser": "pdf"}, + } + ] + assert tracker.created == [ + { + "task_type": "connector_import", + "resource_id": "viking://resources/spec", + "account_id": "acct", + "user_id": "alice", + } + ] + + +@pytest.mark.asyncio +async def test_add_resource_falls_back_when_add_type_is_not_allowed( + monkeypatch, + resource_service, + resource_processor, + request_context, +): + install_connector_config(monkeypatch, enable=True, allowed_add_types=["tos"]) + install_task_tracker(monkeypatch) + install_noop_background_tasks(monkeypatch) + monkeypatch.setattr( + resource_service, + "_should_use_understanding_api", + lambda _path: False, + ) + + result = await resource_service.add_resource( + path="https://example.com/doc", + ctx=request_context, + to="viking://resources/fallback", + args={"add_type": "web"}, + ) + + assert result["root_uri"] == "viking://resources/fallback" + assert resource_processor.calls[0]["path"] == "https://example.com/doc" + assert resource_processor.calls[0]["add_type"] == "web" + + +@pytest.mark.asyncio +async def test_add_resource_falls_back_when_connector_is_disabled( + monkeypatch, + resource_service, + resource_processor, + request_context, +): + install_connector_config(monkeypatch, enable=False, allowed_add_types=["tos"]) + install_task_tracker(monkeypatch) + install_noop_background_tasks(monkeypatch) + monkeypatch.setattr( + resource_service, + "_should_use_understanding_api", + lambda _path: False, + ) + + result = await resource_service.add_resource( + path="tos://bucket/input", + ctx=request_context, + to="viking://resources/disabled", + args={"add_type": "tos"}, + ) + + assert result["root_uri"] == "viking://resources/disabled" + assert len(resource_processor.calls) == 1 + + +@pytest.mark.asyncio +async def test_connector_route_requires_request_api_key( + monkeypatch, + resource_service, +): + install_connector_config(monkeypatch, enable=True, allowed_add_types=["tos"]) + ctx = RequestContext( + user=UserIdentifier("acct", "alice"), + role=Role.USER, + api_key=None, + ) + + with pytest.raises(InvalidArgumentError, match="API key"): + await resource_service.add_resource( + path="tos://bucket/input", + ctx=ctx, + to="viking://resources/spec", + args={"add_type": "tos"}, + ) + + +@pytest.mark.asyncio +async def test_monitor_connector_task_completes_succeeded_status(monkeypatch, request_context): + tracker = install_task_tracker(monkeypatch) + service = ResourceService() + client = SimpleNamespace(get_task_info=lambda _task_key: {"Status": "succeeded"}) + + async def get_task_info(task_key): + return {"Status": "succeeded", "TaskKey": task_key} + + client.get_task_info = get_task_info + + await service._monitor_connector_task( + client=client, + connector_task_key="connector-task-1", + ov_task_id="task-1", + resource_id="viking://resources/spec", + poll_interval_ms=1, + timeout_seconds=1, + ctx=request_context, + ) + + assert tracker.started == [ + {"task_id": "task-1", "account_id": "acct", "user_id": "alice"} + ] + assert tracker.stages[-1]["stage"] == "connector:succeeded" + assert tracker.completed == [ + { + "task_id": "task-1", + "result": { + "connector_status": "succeeded", + "connector_task_key": "connector-task-1", + }, + "account_id": "acct", + "user_id": "alice", + } + ] + assert tracker.failed == [] + + +@pytest.mark.asyncio +async def test_monitor_connector_task_fails_terminal_error_status(monkeypatch, request_context): + tracker = install_task_tracker(monkeypatch) + service = ResourceService() + + async def get_task_info(_task_key): + return {"status": "failed", "error_message": "source unavailable"} + + client = SimpleNamespace(get_task_info=get_task_info) + + await service._monitor_connector_task( + client=client, + connector_task_key="connector-task-1", + ov_task_id="task-1", + resource_id="viking://resources/spec", + poll_interval_ms=1, + timeout_seconds=1, + ctx=request_context, + ) + + assert tracker.stages[-1]["stage"] == "connector:failed" + assert tracker.failed == [ + { + "task_id": "task-1", + "error": "connector task failed: source unavailable", + "account_id": "acct", + "user_id": "alice", + } + ] + assert tracker.completed == [] diff --git a/tests/unit/test_connector_config.py b/tests/unit/test_connector_config.py new file mode 100644 index 0000000000..c7dfab0f16 --- /dev/null +++ b/tests/unit/test_connector_config.py @@ -0,0 +1,92 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. +# SPDX-License-Identifier: AGPL-3.0 +"""Unit tests for Connector configuration validation.""" + +import pytest + +from openviking_cli.utils.config.open_viking_config import ConnectorConfig, OpenVikingConfig + + +def test_connector_config_defaults_to_disabled_tos_only(): + config = ConnectorConfig() + + assert config.enable is False + assert config.connector == "" + assert config.tracker == "" + assert config.timeout_seconds == 3600 + assert config.poll_interval_ms == 5000 + assert config.allowed_add_types == ["tos"] + + +def test_openviking_config_loads_connector_section_from_dict(): + config = OpenVikingConfig.from_dict( + { + "connector": { + "enable": True, + "connector": "https://connector.example/doc/add", + "tracker": "https://connector.example/task/info", + "main_account_id": "main-account", + "timeout_seconds": 120, + "poll_interval_ms": 250, + "allowed_add_types": ["tos", "web"], + } + } + ) + + assert config.connector.enable is True + assert config.connector.connector == "https://connector.example/doc/add" + assert config.connector.tracker == "https://connector.example/task/info" + assert config.connector.main_account_id == "main-account" + assert config.connector.timeout_seconds == 120 + assert config.connector.poll_interval_ms == 250 + assert config.connector.allowed_add_types == ["tos", "web"] + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ( + { + "enable": True, + "connector": "", + "tracker": "https://connector.example/task/info", + }, + "connector.connector is required", + ), + ( + { + "enable": True, + "connector": "https://connector.example/doc/add", + "tracker": "", + }, + "connector.tracker is required", + ), + ( + { + "enable": True, + "connector": "connector.example/doc/add", + "tracker": "https://connector.example/task/info", + }, + "connector.connector must be a full endpoint URL", + ), + ( + { + "connector": "https://connector.example/doc/add", + "tracker": "https://connector.example/task/info", + "timeout_seconds": 0, + }, + "connector.timeout_seconds must be > 0", + ), + ( + { + "connector": "https://connector.example/doc/add", + "tracker": "https://connector.example/task/info", + "poll_interval_ms": 0, + }, + "connector.poll_interval_ms must be > 0", + ), + ], +) +def test_connector_config_rejects_invalid_shapes(kwargs, match): + with pytest.raises(ValueError, match=match): + ConnectorConfig(**kwargs)