diff --git a/openviking/connector/delegate.py b/openviking/connector/delegate.py index 1622f98fb0..e70dab7a54 100644 --- a/openviking/connector/delegate.py +++ b/openviking/connector/delegate.py @@ -17,8 +17,12 @@ from __future__ import annotations import asyncio +import base64 +import json +import re import time from typing import Any, Awaitable, Callable, Dict, List, Optional, Set, Tuple +from urllib.parse import urlsplit import httpx @@ -31,6 +35,7 @@ detect_connector_add_type, is_full_commit_sha, ) +from openviking.crypto.encryptor import MAGIC as ENCRYPTED_ENVELOPE_MAGIC from openviking.parse.mode import ParseMode from openviking.resource.processing_mode import ( DEFAULT_PROCESSING_MODE, @@ -44,6 +49,50 @@ logger = get_logger(__name__) +_TOS_BUCKET_PATTERN = re.compile(r"^[a-z0-9](?:[a-z0-9-]{1,61}[a-z0-9])$") + + +def _validate_tos_uri( + value: Any, + field: str, + *, + allow_bucket_without_slash: bool = False, +) -> None: + """Validate a Connector TOS URI without exposing it in client errors.""" + error = InvalidArgumentError(f"{field} must be a valid TOS URI.") + if ( + not isinstance(value, str) + or value != value.strip() + or not value.startswith("tos://") + or any( + char in "?#%" or ord(char) < 0x20 or ord(char) == 0x7F + for char in value + ) + ): + raise error + + try: + parsed = urlsplit(value) + hostname = parsed.hostname + port = parsed.port + except ValueError: + raise error from None + if ( + parsed.scheme != "tos" + or parsed.username is not None + or parsed.password is not None + or port is not None + or not hostname + or hostname != parsed.netloc + or not _TOS_BUCKET_PATTERN.fullmatch(hostname) + or (not parsed.path and not allow_bucket_without_slash) + or (parsed.path and not parsed.path.startswith("/")) + or parsed.path.startswith("//") + or parsed.query + or parsed.fragment + ): + raise error + class ConnectorDelegate: """Routes add_resource requests to the external Connector service. @@ -71,6 +120,113 @@ def __init__( self._background_tasks = background_tasks self._link_reason_memory = link_reason_memory + _WATCH_AUTH_PROVIDER = "connector_encrypted" + _WATCH_PLAINTEXT_AUTH_PROVIDER = "connector_plaintext" + + def _watch_encryptor(self) -> Any: + encryptor = getattr(self._viking_fs, "_encryptor", None) + if encryptor is None: + raise InvalidArgumentError( + "Connector watch requires encryption.enabled=true so credentials are " + "encrypted at rest." + ) + return encryptor + + async def create_watch_auth_state( + self, + *, + api_key: str, + account_id: str, + add_type: str, + path: str, + connector_args: Optional[Dict[str, Any]], + ) -> Dict[str, Any]: + """Build the private request state needed to replay a Connector watch.""" + if not api_key: + raise InvalidArgumentError("Connector watch requires an API key.") + payload = { + "api_key": api_key, + "account_id": account_id, + "add_type": add_type, + "path": path, + "connector_args": dict(connector_args or {}), + } + encryptor = getattr(self._viking_fs, "_encryptor", None) + if encryptor is None: + return { + "provider": self._WATCH_PLAINTEXT_AUTH_PROVIDER, + "request": payload, + } + try: + plaintext = json.dumps( + payload, + ensure_ascii=False, + separators=(",", ":"), + ).encode("utf-8") + ciphertext = await encryptor.encrypt(account_id, plaintext) + except InvalidArgumentError: + raise + except Exception as exc: + raise InvalidArgumentError("Failed to encrypt Connector watch credentials.") from exc + return { + "provider": self._WATCH_AUTH_PROVIDER, + "ciphertext": base64.b64encode(ciphertext).decode("ascii"), + } + + @classmethod + def is_watch_auth_state(cls, auth_state: Optional[Dict[str, Any]]) -> bool: + return ( + isinstance(auth_state, dict) + and auth_state.get("provider") + in {cls._WATCH_AUTH_PROVIDER, cls._WATCH_PLAINTEXT_AUTH_PROVIDER} + ) + + async def restore_watch_request( + self, + auth_state: Dict[str, Any], + *, + account_id: str, + path: str, + ) -> Tuple[str, str, Dict[str, Any]]: + """Restore and validate a source-bound Connector watch request.""" + try: + provider = auth_state.get("provider") + if provider == self._WATCH_PLAINTEXT_AUTH_PROVIDER: + payload = auth_state.get("request") + elif provider == self._WATCH_AUTH_PROVIDER: + encoded = auth_state.get("ciphertext") + if not isinstance(encoded, str) or not encoded: + raise ValueError("missing ciphertext") + ciphertext = base64.b64decode(encoded, validate=True) + if not ciphertext.startswith(ENCRYPTED_ENVELOPE_MAGIC): + raise ValueError("invalid encrypted envelope") + plaintext = await self._watch_encryptor().decrypt(account_id, ciphertext) + payload = json.loads(plaintext.decode("utf-8")) + else: + raise ValueError("unknown Connector watch provider") + if ( + not isinstance(payload, dict) + or payload.get("account_id") != account_id + or payload.get("path") != path + ): + raise ValueError("watch binding mismatch") + api_key = payload.get("api_key") + add_type = payload.get("add_type") + connector_args = payload.get("connector_args") + if ( + not isinstance(api_key, str) + or not api_key + or not isinstance(add_type, str) + or not add_type + or not isinstance(connector_args, dict) + ): + raise ValueError("invalid watch request") + return api_key, add_type, dict(connector_args) + except InvalidArgumentError: + raise + except Exception as exc: + raise InvalidArgumentError("Stored Connector watch credentials are invalid.") from exc + @staticmethod def resolve_add_type(path: str, declared_add_type: Optional[str]) -> Optional[Tuple[str, bool]]: """Resolve the Connector ``(add_type, connector_only)`` for *path*. @@ -100,6 +256,14 @@ def resolve_add_type(path: str, declared_add_type: Optional[str]) -> Optional[Tu ) return (declared_add_type, True) + @classmethod + def supported_args(cls, path: str, declared_add_type: Optional[str]) -> Set[str]: + """Connector-owned ``args`` fields for the resolved source type.""" + resolved = cls.resolve_add_type(path, declared_add_type) + if resolved is None: + return set() + return set(CONNECTOR_SUPPORTED_ARGS.get(resolved[0], frozenset())) + def should_delegate( self, path: str, @@ -211,7 +375,7 @@ def should_delegate( f"standard import pipeline. Connector import does not support: {detail}" ) logger.info( - f"[ConnectorDelegate] Connector does not support {detail} for path {path}; " + f"[ConnectorDelegate] Connector does not support {detail}; " "falling back to the standard import pipeline" ) return False @@ -247,8 +411,6 @@ def _unsupported_params( unsupported.append("missing exact 'to' target") elif to != "viking://resources" and not to.startswith("viking://resources/"): unsupported.append("to outside the public resources root (viking://resources/...)") - if watch_interval > 0: - unsupported.append("watch_interval>0 (Connector imports cannot be watched yet)") if instruction: unsupported.append("instruction") if not build_index: @@ -260,6 +422,8 @@ def _unsupported_params( if kwargs.get("strict"): unsupported.append("strict=true (Connector imports fail per file, not all-or-nothing)") for field in ("ignore_dirs", "include", "exclude"): + if field == "exclude" and add_type == "tos" and field in connector_args: + continue if kwargs.get(field): unsupported.append(f"{field} (Connector imports cannot filter the source tree)") if kwargs.get("preserve_structure") is False: @@ -301,6 +465,8 @@ async def submit( connector_args: Optional[Dict[str, Any]] = None, tags: Optional[List[str]] = None, tag_mode: str = "replace", + wait_for_completion: bool = False, + on_success: Optional[Callable[[], Awaitable[None]]] = None, **kwargs: Any, ) -> Dict[str, Any]: """Route add_resource to the external Connector service.""" @@ -315,6 +481,8 @@ async def submit( if resolved is None: raise InvalidArgumentError(f"'{path}' does not match any Connector source type.") add_type, _ = resolved + if add_type == "tos": + _validate_tos_uri(path, "path", allow_bucket_without_slash=True) task_resource_id = to or "" if not task_resource_id: @@ -369,12 +537,34 @@ async def submit( tos_path: Optional[str] = None param_config: Optional[Dict[str, Any]] = None if add_type == "tos": - source_path = path[len("tos://") :].strip() - if not source_path: - raise InvalidArgumentError( - "Connector TOS import requires path='tos:///'." - ) - tos_path = source_path + tos_args = connector_args or {} + if "tos_prefix" not in tos_args: + if "exclude" in tos_args: + raise InvalidArgumentError("args.exclude requires args.tos_prefix.") + source_path = path[len("tos://") :].strip() + if not source_path: + raise InvalidArgumentError( + "Connector TOS import requires path='tos:///'." + ) + tos_path = source_path + else: + tos_prefix = tos_args["tos_prefix"] + if not isinstance(tos_prefix, list) or not tos_prefix: + raise InvalidArgumentError( + "args.tos_prefix must be a non-empty list of TOS URIs." + ) + for index, source in enumerate(tos_prefix): + _validate_tos_uri(source, f"args.tos_prefix[{index}]") + if tos_prefix[0] != path: + raise InvalidArgumentError("path must equal the first item in args.tos_prefix.") + exclude = tos_args.get("exclude", []) + if not isinstance(exclude, list): + raise InvalidArgumentError("args.exclude must be a list of TOS URIs.") + for index, source in enumerate(exclude): + _validate_tos_uri(source, f"args.exclude[{index}]") + param_config = {"tos_prefix": tos_prefix} + if exclude: + param_config["exclude"] = exclude elif add_type == "git": from openviking.parse.accessors.git_accessor import GitAccessor @@ -462,12 +652,9 @@ async def submit( ctx=ctx, reason=reason, link_root_uri=task_resource_id or "viking://resources", + on_success=on_success, ) - background = asyncio.create_task(monitor) - self._background_tasks.add(background) - background.add_done_callback(self._background_tasks.discard) - response = { "status": "accepted", "task_id": task.task_id, @@ -475,6 +662,13 @@ async def submit( } if task_resource_id: response["resource_id"] = task_resource_id + if wait_for_completion: + response.update(await monitor) + return response + + background = asyncio.create_task(monitor) + self._background_tasks.add(background) + background.add_done_callback(self._background_tasks.discard) return response async def _monitor( @@ -487,6 +681,7 @@ async def _monitor( ctx: RequestContext, reason: str = "", link_root_uri: str = "", + on_success: Optional[Callable[[], Awaitable[None]]] = None, ) -> Dict[str, Any]: """Poll the Connector task until terminal state, then update OV TaskRecord. @@ -523,10 +718,10 @@ async def _monitor( f"for {connector_task_key}: {status_code}; retrying" ) continue - except httpx.RequestError as exc: + except httpx.RequestError: logger.warning( "[ConnectorDelegate] Transient Connector task polling error " - f"for {connector_task_key}: {exc}; retrying" + f"for {connector_task_key}; retrying" ) continue status = (info.get("Status") or info.get("status") or "").lower() @@ -544,6 +739,8 @@ async def _monitor( "connector_status": status, "connector_task_key": connector_task_key, } + if on_success is not None: + await on_success() if (reason or "").strip() and link_root_uri: link_result: Dict[str, Any] = {"root_uri": link_root_uri} await self._link_reason_memory( @@ -590,11 +787,15 @@ async def _monitor( ) raise except Exception as exc: - logger.error(f"[ConnectorDelegate] Connector task monitor error: {exc}") + failure = "connector task monitoring failed" + logger.error( + "[ConnectorDelegate] Connector task monitor error, error_type=%s", + type(exc).__name__, + ) await task_tracker.fail( ov_task_id, - str(exc), + failure, account_id=ctx.account_id, user_id=ctx.user.user_id, ) - return {"status": "failed", "error": str(exc)} + return {"status": "failed", "error": failure} diff --git a/openviking/connector/routing.py b/openviking/connector/routing.py index 74e3e74883..9bbe51948f 100644 --- a/openviking/connector/routing.py +++ b/openviking/connector/routing.py @@ -23,7 +23,7 @@ # the set flow into the unsupported-parameter framework: connector-only # sources reject them, shared sources degrade to the standard pipeline. CONNECTOR_SUPPORTED_ARGS: Dict[str, FrozenSet[str]] = { - "tos": frozenset(), + "tos": frozenset({"tos_prefix", "exclude"}), "git": frozenset({"branch", "ref", "commit"}), } diff --git a/openviking/resource/watch_manager.py b/openviking/resource/watch_manager.py index d3271e9b59..5add061034 100644 --- a/openviking/resource/watch_manager.py +++ b/openviking/resource/watch_manager.py @@ -872,6 +872,36 @@ async def get_task_by_uri( return task + async def get_upsertable_task_by_uri( + self, + *, + path: str, + to_uri: str, + account_id: str, + user_id: str, + role: str, + ) -> Optional[WatchTask]: + """Return an existing task when this source may create or update its watch.""" + async with self._uri_mutation_coordinator.access(account_id, [to_uri]): + async with self._lock: + task_id = self._uri_to_task.get((account_id, to_uri)) + if not task_id: + return None + + task = self._tasks.get(task_id) + if not task or not self._check_permission(task, account_id, user_id, role): + raise ConflictError( + f"Target URI '{to_uri}' is already used by another task", + resource=to_uri, + ) + if task.is_active and task.path != path: + raise ConflictError( + f"Target URI '{to_uri}' is already being monitored by task " + f"{task.task_id}. Please cancel the existing task first.", + resource=to_uri, + ) + return task + async def update_execution_time(self, task_id: str) -> None: """Update task execution time after execution. diff --git a/openviking/resource/watch_scheduler.py b/openviking/resource/watch_scheduler.py index 9dc9d10d45..53666f3225 100644 --- a/openviking/resource/watch_scheduler.py +++ b/openviking/resource/watch_scheduler.py @@ -10,6 +10,7 @@ from datetime import datetime from typing import Any, Dict, Optional, Set +from openviking.connector.delegate import ConnectorDelegate from openviking.resource.feishu_watch_auth import ( FeishuOAuthClient, FeishuTokenRefreshError, @@ -262,14 +263,16 @@ async def _execute_stable_task(self, task) -> None: Args: task: WatchTask to execute """ - logger.info(f"[WatchScheduler] Executing task {task.task_id} for path {task.path}") + logger.info(f"[WatchScheduler] Executing task {task.task_id}") cancelled = False should_deactivate = False deactivation_reason = "" try: - if not self._check_resource_exists(task.path): + auth_state = getattr(task, "auth_state", None) + connector_watch = ConnectorDelegate.is_watch_auth_state(auth_state) + if not connector_watch and not self._check_resource_exists(task.path): should_deactivate = True deactivation_reason = f"Resource path does not exist: {task.path}" logger.warning( @@ -307,7 +310,6 @@ async def _execute_stable_task(self, task) -> None: processor_kwargs = dict(getattr(task, "processor_kwargs", {}) or {}) processor_kwargs.pop("build_index", None) processor_kwargs.pop("summarize", None) - auth_state = getattr(task, "auth_state", None) if is_feishu_auth_state(auth_state): try: auth_state = await self._prepare_feishu_auth_state(task, auth_state) @@ -327,6 +329,19 @@ async def _execute_stable_task(self, task) -> None: auth_state, task.path, ) + elif connector_watch: + ( + api_key, + add_type, + connector_args, + ) = await self._resource_service._connector.restore_watch_request( + auth_state, + account_id=task.account_id, + path=task.path, + ) + ctx.api_key = api_key + processor_kwargs["add_type"] = add_type + processor_kwargs["args"] = connector_args if not should_deactivate: result = await self._resource_service.refresh_resource( @@ -345,10 +360,16 @@ async def _execute_stable_task(self, task) -> None: **processor_kwargs, ) - logger.info( - f"[WatchScheduler] Task {task.task_id} executed successfully, " - f"result: {result.get('root_uri', 'N/A')}" - ) + if result.get("status") == "failed": + logger.warning( + f"[WatchScheduler] Task {task.task_id} execution finished with " + "a failed ingestion task" + ) + else: + logger.info( + f"[WatchScheduler] Task {task.task_id} executed successfully, " + f"result: {result.get('root_uri', 'N/A')}" + ) except asyncio.CancelledError: cancelled = True @@ -361,8 +382,8 @@ async def _execute_stable_task(self, task) -> None: ) except Exception as e: logger.error( - f"[WatchScheduler] Task {task.task_id} execution failed: {e}", - exc_info=True, + f"[WatchScheduler] Task {task.task_id} execution failed, " + f"error_type={type(e).__name__}" ) finally: diff --git a/openviking/server/routers/resources.py b/openviking/server/routers/resources.py index be9bd3944f..4b8fae761c 100644 --- a/openviking/server/routers/resources.py +++ b/openviking/server/routers/resources.py @@ -74,7 +74,10 @@ class AddResourceRequest(BaseModel): Note: Re-adding the same source to the same target updates its active watch task. A different source targeting an active watch raises ConflictError; cancel that - watch first with watch_interval <= 0. + watch first with watch_interval <= 0. For Connector imports this check is + eventually consistent: the Watch is created only after the background import + succeeds, so overlapping imports may both write before Watch finalization + reports the conflict. """ model_config = ConfigDict(extra="forbid") diff --git a/openviking/service/resource_service.py b/openviking/service/resource_service.py index 1defa2379b..c1f2dcaec6 100644 --- a/openviking/service/resource_service.py +++ b/openviking/service/resource_service.py @@ -11,7 +11,7 @@ import inspect import json import time -from collections.abc import Callable +from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, List, Optional @@ -359,6 +359,7 @@ def _normalize_add_resource_args( args: Optional[Dict[str, Any]], *, watch_interval: float, + allowed_reserved_fields: Optional[set[str]] = None, ) -> _NormalizedAddResourceArgs: if args is None: return _NormalizedAddResourceArgs({}) @@ -367,7 +368,10 @@ def _normalize_add_resource_args( if not args: return _NormalizedAddResourceArgs({}) - reserved = sorted(set(args).intersection(_ADD_RESOURCE_ARGS_RESERVED_FIELDS)) + reserved_fields = _ADD_RESOURCE_ARGS_RESERVED_FIELDS - ( + allowed_reserved_fields or set() + ) + reserved = sorted(set(args).intersection(reserved_fields)) if reserved: raise InvalidArgumentError( "args cannot contain core add_resource fields: " + ", ".join(reserved) @@ -1163,6 +1167,8 @@ async def refresh_resource( watch_interval: float = 0, allow_local_path_resolution: bool = True, enforce_public_remote_targets: bool = False, + add_type: Optional[str] = None, + args: Optional[Dict[str, Any]] = None, **kwargs, ) -> Dict[str, Any]: """Submit a scheduled refresh without changing its watch task.""" @@ -1183,7 +1189,8 @@ async def refresh_resource( manage_watch=False, allow_local_path_resolution=allow_local_path_resolution, enforce_public_remote_targets=enforce_public_remote_targets, - args=None, + add_type=add_type, + args=args, **kwargs, ) @@ -1242,7 +1249,9 @@ async def _submit_resource_ingestion( Note: Re-adding the same source to the same target updates its active watch task in place. A different source targeting an active watch raises - ConflictError; cancel that watch first with watch_interval <= 0. + ConflictError; cancel that watch first with watch_interval <= 0. Connector + imports create the Watch only after the background import succeeds, so this + conflict may be reported after both overlapping imports have written data. enforce_public_remote_targets: When True, reject non-public remote hosts and validate each outbound HTTP request URL during fetch. args: Parser/accessor-specific options forwarded to the processing chain. @@ -1258,12 +1267,31 @@ async def _submit_resource_ingestion( self._ensure_initialized() processing_mode = normalize_processing_mode(processing_mode) self._validate_add_resource_tag_policy(tags=tags, tag_mode=tag_mode) - normalized_args = self._normalize_add_resource_args(args, watch_interval=watch_interval) + from openviking.connector.delegate import ConnectorDelegate + + allowed_reserved_fields = ConnectorDelegate.supported_args(path, add_type).intersection( + _ADD_RESOURCE_ARGS_RESERVED_FIELDS + ) + normalized_args = self._normalize_add_resource_args( + args, + watch_interval=watch_interval, + allowed_reserved_fields=allowed_reserved_fields, + ) mode = ( normalize_parse_mode(parse_mode) if parse_mode is not None else normalized_args.parse_mode ) + duplicated_fields = sorted( + field + for field in allowed_reserved_fields + if field in normalized_args.processor_kwargs and kwargs.get(field) is not None + ) + if duplicated_fields: + raise InvalidArgumentError( + f"{', '.join(duplicated_fields)} cannot be provided both as a top-level " + "field and in args." + ) kwargs.update(normalized_args.processor_kwargs) git_repo_source = is_git_repo_url(path) if git_repo_source: @@ -1295,7 +1323,8 @@ async def _submit_resource_ingestion( target_parent = parent or "" target_create_parent = bool(kwargs.get("create_parent", False)) - if self._connector.should_delegate( + connector = self._connector + if connector.should_delegate( path, ctx=ctx, declared_add_type=add_type, @@ -1311,17 +1340,96 @@ async def _submit_resource_ingestion( connector_args=normalized_args.processor_kwargs, kwargs=kwargs, ): - return await self._connector.submit( + resolved = connector.resolve_add_type(path, add_type) + if resolved is None: # pragma: no cover - should_delegate already resolved it + raise InvalidArgumentError(f"'{path}' does not match any Connector source type.") + watch_manager = self._get_watch_manager() + watch_auth_state = None + defer_watch_creation = bool(watch_manager and manage_watch and watch_interval > 0) + if defer_watch_creation and watch_manager: + # Connector imports may run for a long time. This is only a best-effort + # precheck: the authoritative conflict check happens when on_success + # creates the Watch, after the Connector may already have written data. + await watch_manager.get_upsertable_task_by_uri( + path=path, + to_uri=target_to, + account_id=ctx.account_id, + user_id=ctx.user.user_id, + role=str(ctx.role), + ) + watch_auth_state = await connector.create_watch_auth_state( + api_key=ctx.api_key or "", + account_id=ctx.account_id, + add_type=resolved[0], + path=path, + connector_args=normalized_args.processor_kwargs, + ) + connector_watch_processor_kwargs = self._watch_processor_kwargs( + { + key: value + for key, value in kwargs.items() + if key not in normalized_args.processor_kwargs + }, + tags, + tag_mode, + ) + on_success: Optional[Callable[[], Awaitable[None]]] = None + if defer_watch_creation: + + async def create_watch_after_success() -> None: + await self._manage_watch_if_needed( + watch_manager=watch_manager, + manage_watch=True, + watch_interval=watch_interval, + to=target_to, + parent=target_parent, + to_is_directory=to_is_directory, + root_uri=target_to, + path=path, + reason=reason, + instruction=instruction, + build_index=build_index, + summarize=summarize, + processing_mode=processing_mode, + processor_kwargs=connector_watch_processor_kwargs, + watch_auth_state=watch_auth_state, + ctx=ctx, + ) + + on_success = create_watch_after_success + result = await connector.submit( path=path, ctx=ctx, declared_add_type=add_type, - to=to, + to=target_to, reason=reason, connector_args=normalized_args.processor_kwargs, tags=tags, tag_mode=tag_mode, + wait_for_completion=not manage_watch and watch_interval > 0, + on_success=on_success, **kwargs, ) + if not defer_watch_creation: + await self._manage_watch_if_needed( + watch_manager=watch_manager, + manage_watch=manage_watch, + watch_interval=watch_interval, + to=target_to, + parent=target_parent, + to_is_directory=to_is_directory, + root_uri=target_to, + path=path, + reason=reason, + instruction=instruction, + build_index=build_index, + summarize=summarize, + processing_mode=processing_mode, + processor_kwargs=connector_watch_processor_kwargs, + watch_auth_state=watch_auth_state, + ctx=ctx, + ) + return result if enforce_public_remote_targets and is_remote_resource_source(path): path = require_remote_resource_source(path) @@ -1734,19 +1842,14 @@ async def _handle_watch_task_creation( if not watch_manager: return - existing_task = await watch_manager.get_task_by_uri( + existing_task = await watch_manager.get_upsertable_task_by_uri( + path=path, to_uri=to_uri, account_id=ctx.account_id, user_id=ctx.user.user_id, role=str(ctx.role), ) if existing_task: - if existing_task.is_active and existing_task.path != path: - raise ConflictError( - f"Target URI '{to_uri}' is already being monitored by task {existing_task.task_id}. " - f"Please cancel the existing task first.", - resource=to_uri, - ) was_active = existing_task.is_active await watch_manager.update_task( task_id=existing_task.task_id, diff --git a/tests/cli/conftest.py b/tests/cli/conftest.py index fc1bc89aa3..b8bff483b6 100644 --- a/tests/cli/conftest.py +++ b/tests/cli/conftest.py @@ -629,20 +629,6 @@ def ensure_resources_dir(): pytest.fail(f"Failed to create resources dir: {r['stderr'][:300]}") -@pytest.fixture(scope="session") -def ensure_user_skills_dir(): - uri = "viking://user/skills" - r = ov_mkdir(uri) - if r["exit_code"] == 0 or "already exists" in (r.get("stderr") or "").lower(): - return - skip_if_auth_error(r) - stat_r = ov(["stat", uri, "-o", "json"], timeout=120) - if stat_r["exit_code"] == 0: - return - skip_if_auth_error(stat_r) - pytest.fail(f"mkdir {uri} failed after retries: {r['stderr'][:300]}") - - @pytest.fixture(scope="session") def test_dir_uri(ensure_resources_dir): uri = f"viking://resources/cli_test_{uuid.uuid4().hex[:8]}" diff --git a/tests/cli/test_cli_skills.py b/tests/cli/test_cli_skills.py index 7f7f670f6e..b00b698a38 100644 --- a/tests/cli/test_cli_skills.py +++ b/tests/cli/test_cli_skills.py @@ -39,47 +39,41 @@ def test_add_skill_from_file(self): class TestSkillList: - def test_list_skills(self, ensure_user_skills_dir): - r = ov(["ls", "viking://user/skills/", "-o", "json"]) + def test_list_skills(self): + r = ov(["skills", "list", "-o", "json"]) assert r["exit_code"] == 0, ( - f"ov ls skills should exit 0, got {r['exit_code']}: {r['stderr'][:300]}" + f"ov skills list should exit 0, got {r['exit_code']}: {r['stderr'][:300]}" ) data = r["json"] assert data is not None and data.get("ok") is True, "Expected ok=true" assert "result" in data, "'result' field should exist" - assert isinstance(data["result"], list), "'result' should be a list" + result = data["result"] + assert isinstance(result, dict), "'result' should be an object" + assert isinstance(result.get("skills"), list), "'result.skills' should be a list" -class TestSkillRead: - def test_read_skill(self, ensure_user_skills_dir): - ls_r = ov(["ls", "viking://user/skills/", "-o", "json"]) - if ls_r["exit_code"] != 0 or not ls_r["json"] or not ls_r["json"].get("result"): - pytest.skip("No skills available to test read") +class TestSkillShow: + def test_show_skill(self): + list_r = ov(["skills", "list", "-o", "json"]) + if list_r["exit_code"] != 0 or not list_r["json"]: + pytest.skip("No skills available to test show") return - skills = ls_r["json"]["result"] + skills = list_r["json"].get("result", {}).get("skills", []) if not isinstance(skills, list) or len(skills) == 0: - pytest.skip("No skills available to test read") + pytest.skip("No skills available to test show") return for skill_entry in skills: if not isinstance(skill_entry, dict): continue - skill_uri = skill_entry.get("uri", "") - if not skill_uri: + skill_name = skill_entry.get("name", "") + if not skill_name: continue - ls_skill_r = ov(["ls", skill_uri, "-o", "json"]) - if ( - ls_skill_r["exit_code"] != 0 - or not ls_skill_r["json"] - or not ls_skill_r["json"].get("result") - ): - continue - items = ls_skill_r["json"]["result"] - for item in items: - if isinstance(item, dict) and item.get("isDir") is False: - file_uri = item.get("uri", "") - if file_uri: - r = ov(["read", file_uri, "-o", "json"]) - assert r["exit_code"] == 0, ( - f"ov read skill file should exit 0, got {r['exit_code']}: {r['stderr'][:300]}" - ) - return + r = ov(["skills", "show", skill_name, "--level", "2", "-o", "json"]) + assert r["exit_code"] == 0, ( + f"ov skills show should exit 0, got {r['exit_code']}: {r['stderr'][:300]}" + ) + data = r["json"] + assert data is not None and data.get("ok") is True, "Expected ok=true" + assert data.get("result", {}).get("name") == skill_name + return + pytest.skip("No named skills available to test show") diff --git a/tests/service/test_resource_service_connector.py b/tests/service/test_resource_service_connector.py index b8a78a29f3..36ad1fd995 100644 --- a/tests/service/test_resource_service_connector.py +++ b/tests/service/test_resource_service_connector.py @@ -4,6 +4,7 @@ import asyncio import json +import logging from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -12,11 +13,13 @@ from openviking.connector import delegate as connector_delegate_module from openviking.parse.mode import ParseMode +from openviking.resource.watch_manager import WatchManager +from openviking.resource.watch_scheduler import WatchScheduler 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.storage.queuefs.add_resource_msg import AddResourceMsg, AddResourcePhase -from openviking_cli.exceptions import InvalidArgumentError +from openviking_cli.exceptions import ConflictError, InvalidArgumentError from openviking_cli.session.user_id import UserIdentifier # Deterministic stand-in for the code-hosting predicate: routing tests must @@ -29,6 +32,14 @@ def add_done_callback(self, _callback): pass +class _FakeEncryptor: + async def encrypt(self, _account_id, plaintext): + return b"OVE1" + plaintext[::-1] + + async def decrypt(self, _account_id, ciphertext): + return ciphertext[4:][::-1] + + @pytest.fixture def connector_config(monkeypatch): import openviking_cli.utils.config.open_viking_config as config_module @@ -97,7 +108,13 @@ def _task_tracker(): ) -def _install_connector_dependencies(monkeypatch, tracker, connector_client): +def _install_connector_dependencies( + monkeypatch, + tracker, + connector_client, + *, + discard_monitor=True, +): monkeypatch.setattr( "openviking.service.task_tracker.get_task_tracker", lambda: tracker, @@ -107,6 +124,8 @@ def _install_connector_dependencies(monkeypatch, tracker, connector_client): "ConnectorClient", lambda **_kwargs: connector_client, ) + if not discard_monitor: + return def discard_monitor(coro): coro.close() @@ -116,11 +135,17 @@ def discard_monitor(coro): @pytest.mark.asyncio +@pytest.mark.parametrize( + ("path", "tos_path"), + [("tos://bucket/a/b/c", "bucket/a/b/c"), ("tos://bucket", "bucket")], +) async def test_add_resource_routes_tos_to_connector( monkeypatch, connector_config, ctx, service, + path, + tos_path, ): tracker = _task_tracker() connector_client = SimpleNamespace( @@ -129,7 +154,7 @@ async def test_add_resource_routes_tos_to_connector( _install_connector_dependencies(monkeypatch, tracker, connector_client) result = await service.add_resource( - path="tos://bucket/a/b/c", + path=path, ctx=ctx, to="viking://resources/x/y", ) @@ -143,7 +168,7 @@ async def test_add_resource_routes_tos_to_connector( connector_client.submit_doc_add.assert_awaited_once_with( add_type="tos", api_key="secret", - tos_path="bucket/a/b/c", + tos_path=tos_path, to="viking://resources/x/y", include_child=True, param_config=None, @@ -158,6 +183,578 @@ async def test_add_resource_routes_tos_to_connector( ) +@pytest.mark.asyncio +async def test_connector_watch_stores_only_encrypted_replay_state( + monkeypatch, + connector_config, + ctx, +): + tracker = _task_tracker() + connector_client = SimpleNamespace( + submit_doc_add=AsyncMock(return_value={"task_key": "connector-1"}) + ) + _install_connector_dependencies( + monkeypatch, + tracker, + connector_client, + discard_monitor=False, + ) + watch_manager = WatchManager() + viking_fs = SimpleNamespace( + exists=AsyncMock(return_value=True), + _encryptor=_FakeEncryptor(), + ) + service = ResourceService( + vikingdb=object(), + viking_fs=viking_fs, + resource_processor=object(), + skill_processor=object(), + watch_scheduler=SimpleNamespace(watch_manager=watch_manager), + ) + release_monitor = asyncio.Event() + + async def complete_monitor(**kwargs): + await release_monitor.wait() + await kwargs["on_success"]() + return {"status": "completed"} + + monkeypatch.setattr(service._connector, "_monitor", complete_monitor) + + await service.add_resource( + path="tos://bucket/docs/", + ctx=ctx, + to="viking://resources/x/y", + watch_interval=5, + args={"tos_prefix": ["tos://bucket/docs/"]}, + ) + + assert ( + await watch_manager.get_task_by_uri( + "viking://resources/x/y", + account_id="acct", + user_id="alice", + role=str(Role.USER), + ) + is None + ) + background_tasks = list(service._background_tasks) + release_monitor.set() + await asyncio.gather(*background_tasks) + + task = await watch_manager.get_task_by_uri( + "viking://resources/x/y", + account_id="acct", + user_id="alice", + role=str(Role.USER), + ) + assert task is not None + assert task.auth_state["provider"] == "connector_encrypted" + assert "secret" not in json.dumps(task.auth_state) + assert "auth_state" not in task.to_dict() + assert await service._connector.restore_watch_request( + task.auth_state, + account_id="acct", + path="tos://bucket/docs/", + ) == ( + "secret", + "tos", + {"tos_prefix": ["tos://bucket/docs/"]}, + ) + + +@pytest.mark.asyncio +async def test_connector_watch_prechecks_only_active_conflicts( + monkeypatch, + connector_config, + ctx, +): + tracker = _task_tracker() + connector_client = SimpleNamespace( + submit_doc_add=AsyncMock(return_value={"task_key": "connector-1"}) + ) + _install_connector_dependencies( + monkeypatch, + tracker, + connector_client, + discard_monitor=False, + ) + watch_manager = WatchManager() + service = ResourceService( + vikingdb=object(), + viking_fs=SimpleNamespace( + exists=AsyncMock(return_value=True), + _encryptor=_FakeEncryptor(), + ), + resource_processor=object(), + skill_processor=object(), + watch_scheduler=SimpleNamespace(watch_manager=watch_manager), + ) + to_uri = "viking://resources/x/y" + existing = await watch_manager.create_task( + path="tos://bucket/old/", + account_id="acct", + user_id="alice", + original_role=str(Role.USER), + to_uri=to_uri, + watch_interval=5, + ) + + with pytest.raises(ConflictError, match="already being monitored"): + await service.add_resource( + path="tos://bucket/new/", + ctx=ctx, + to=to_uri, + watch_interval=5, + ) + + connector_client.submit_doc_add.assert_not_awaited() + tracker.create.assert_not_awaited() + + await watch_manager.update_task( + task_id=existing.task_id, + account_id="acct", + user_id="alice", + role=str(Role.USER), + is_active=False, + ) + release_monitor = asyncio.Event() + + async def complete_monitor(**kwargs): + await release_monitor.wait() + await kwargs["on_success"]() + return {"status": "completed"} + + monkeypatch.setattr(service._connector, "_monitor", complete_monitor) + await service.add_resource( + path="tos://bucket/new/", + ctx=ctx, + to=to_uri, + watch_interval=5, + ) + + assert existing.is_active is False + assert existing.path == "tos://bucket/old/" + release_monitor.set() + await asyncio.gather(*list(service._background_tasks)) + assert existing.is_active is True + assert existing.path == "tos://bucket/new/" + + +@pytest.mark.asyncio +async def test_connector_watch_resolves_overlapping_imports_at_finalize( + connector_config, + ctx, +): + watch_manager = WatchManager() + service = ResourceService( + vikingdb=object(), + viking_fs=SimpleNamespace( + exists=AsyncMock(return_value=True), + _encryptor=_FakeEncryptor(), + ), + resource_processor=object(), + skill_processor=object(), + watch_scheduler=SimpleNamespace(watch_manager=watch_manager), + ) + finalizers = [] + + async def submit(**kwargs): + finalizers.append(kwargs["on_success"]) + return {"status": "accepted"} + + service._connector.submit = AsyncMock(side_effect=submit) + to_uri = "viking://resources/x/y" + + await service.add_resource( + path="tos://bucket/first/", + ctx=ctx, + to=to_uri, + watch_interval=5, + ) + await service.add_resource( + path="tos://bucket/second/", + ctx=ctx, + to=to_uri, + watch_interval=5, + ) + + assert service._connector.submit.await_count == 2 + await finalizers[0]() + with pytest.raises(ConflictError, match="already being monitored"): + await finalizers[1]() + + task = await watch_manager.get_task_by_uri( + to_uri, + account_id="acct", + user_id="alice", + role=str(Role.USER), + ) + assert task is not None + assert task.path == "tos://bucket/first/" + + +@pytest.mark.asyncio +async def test_git_connector_watch_keeps_args_only_in_encrypted_state( + monkeypatch, + connector_config, + ctx, +): + connector_config.allowed_add_types = ["git"] + tracker = _task_tracker() + connector_client = SimpleNamespace( + submit_doc_add=AsyncMock(return_value={"task_key": "connector-1"}) + ) + _install_connector_dependencies( + monkeypatch, + tracker, + connector_client, + discard_monitor=False, + ) + watch_manager = WatchManager() + service = ResourceService( + vikingdb=object(), + viking_fs=SimpleNamespace( + exists=AsyncMock(return_value=True), + _encryptor=_FakeEncryptor(), + ), + resource_processor=object(), + skill_processor=object(), + watch_scheduler=SimpleNamespace(watch_manager=watch_manager), + ) + + async def complete_monitor(**kwargs): + await kwargs["on_success"]() + return {"status": "completed"} + + monkeypatch.setattr(service._connector, "_monitor", complete_monitor) + connector_args = { + "branch": "main", + "token": "secret-token", + "username": "oauth2", + } + + await service.add_resource( + path="https://git.example/org/private.git", + ctx=ctx, + to="viking://resources/private", + watch_interval=5, + args=connector_args, + ) + await asyncio.gather(*list(service._background_tasks)) + + task = await watch_manager.get_task_by_uri( + "viking://resources/private", + account_id="acct", + user_id="alice", + role=str(Role.USER), + ) + assert task is not None + assert task.processor_kwargs == {} + assert "secret-token" not in json.dumps(task.to_dict()) + assert await service._connector.restore_watch_request( + task.auth_state, + account_id="acct", + path="https://git.example/org/private.git", + ) == ("secret", "git", connector_args) + + +@pytest.mark.asyncio +async def test_connector_watch_allows_plaintext_private_state_without_encryption( + connector_config, + ctx, + service, +): + watch_manager = WatchManager() + service._watch_scheduler = SimpleNamespace(watch_manager=watch_manager) + submitted = {} + + async def submit(**kwargs): + submitted.update(kwargs) + return {"status": "accepted"} + + service._connector.submit = AsyncMock(side_effect=submit) + await service.add_resource( + path="tos://bucket/docs/", + ctx=ctx, + to="viking://resources/x/y", + watch_interval=5, + ) + await submitted["on_success"]() + + task = await watch_manager.get_task_by_uri( + "viking://resources/x/y", + account_id="acct", + user_id="alice", + role=str(Role.USER), + ) + assert task is not None + assert task.auth_state == { + "provider": "connector_plaintext", + "request": { + "api_key": "secret", + "account_id": "acct", + "add_type": "tos", + "path": "tos://bucket/docs/", + "connector_args": {}, + }, + } + assert "auth_state" not in task.to_dict() + assert await service._connector.restore_watch_request( + task.auth_state, + account_id="acct", + path="tos://bucket/docs/", + ) == ("secret", "tos", {}) + + +@pytest.mark.asyncio +async def test_connector_watch_scheduler_restores_request_credentials( + connector_config, +): + viking_fs = SimpleNamespace(_encryptor=_FakeEncryptor()) + service = ResourceService( + vikingdb=object(), + viking_fs=viking_fs, + resource_processor=object(), + skill_processor=object(), + ) + service.refresh_resource = AsyncMock(return_value={"status": "completed"}) + scheduler = WatchScheduler(resource_service=service) + watch_manager = WatchManager() + scheduler._watch_manager = watch_manager + auth_state = await service._connector.create_watch_auth_state( + api_key="secret", + account_id="acct", + add_type="tos", + path="tos://bucket/docs/", + connector_args={"tos_prefix": ["tos://bucket/docs/"]}, + ) + task = await watch_manager.create_task( + path="tos://bucket/docs/", + to_uri="viking://resources/x/y", + watch_interval=5, + account_id="acct", + user_id="alice", + auth_state=auth_state, + ) + + await scheduler._execute_stable_task(task) + + call = service.refresh_resource.await_args + assert call.kwargs["ctx"].api_key == "secret" + assert call.kwargs["add_type"] == "tos" + assert call.kwargs["args"] == {"tos_prefix": ["tos://bucket/docs/"]} + + +@pytest.mark.asyncio +async def test_connector_watch_deactivates_when_target_is_deleted(): + from openviking_cli.exceptions import NotFoundError + + class MissingTargetFS: + async def stat(self, uri, ctx=None): + raise NotFoundError(uri, "resource") + + service = ResourceService() + service.refresh_resource = AsyncMock(return_value={"status": "completed"}) + scheduler = WatchScheduler(resource_service=service, viking_fs=MissingTargetFS()) + watch_manager = WatchManager() + scheduler._watch_manager = watch_manager + task = await watch_manager.create_task( + path="tos://bucket/docs/", + to_uri="viking://resources/x/y", + watch_interval=5, + account_id="acct", + user_id="alice", + auth_state={"provider": "connector_encrypted", "ciphertext": "unused"}, + ) + + await scheduler._execute_stable_task(task) + + updated = await watch_manager.get_task(task.task_id) + assert updated is not None + assert updated.is_active is False + service.refresh_resource.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_connector_watch_refresh_waits_for_remote_completion( + connector_config, + ctx, + service, +): + service._connector.submit = AsyncMock(return_value={"status": "completed"}) + + await service.refresh_resource( + path="tos://bucket/docs/", + ctx=ctx, + to="viking://resources/x/y", + watch_interval=5, + add_type="tos", + ) + + assert service._connector.submit.await_args.kwargs["wait_for_completion"] is True + + +@pytest.mark.asyncio +async def test_add_resource_routes_multiple_tos_sources_to_connector( + monkeypatch, + connector_config, + ctx, + service, +): + tracker = _task_tracker() + connector_client = SimpleNamespace( + submit_doc_add=AsyncMock(return_value={"task_key": "connector-1"}) + ) + _install_connector_dependencies(monkeypatch, tracker, connector_client) + + await service.add_resource( + path="tos://bucket-a/docs/", + ctx=ctx, + to="viking://resources/x/y", + args={ + "tos_prefix": ["tos://bucket-a/docs/", "tos://bucket-b/manual.pdf"], + "exclude": ["tos://bucket-a/docs/drafts/"], + }, + ) + + connector_client.submit_doc_add.assert_awaited_once_with( + add_type="tos", + api_key="secret", + tos_path=None, + to="viking://resources/x/y", + include_child=True, + param_config={ + "tos_prefix": ["tos://bucket-a/docs/", "tos://bucket-b/manual.pdf"], + "exclude": ["tos://bucket-a/docs/drafts/"], + }, + auth_config=None, + extra_params=None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("args", "message"), + [ + ({"tos_prefix": []}, "non-empty list"), + ({"exclude": ["tos://bucket-a/private/"]}, "requires args.tos_prefix"), + ({"tos_prefix": ["tos://bucket-b/docs/"]}, "first item"), + ], +) +async def test_add_resource_rejects_invalid_multiple_tos_sources( + connector_config, + ctx, + service, + args, + message, +): + with pytest.raises(InvalidArgumentError, match=message): + await service.add_resource( + path="tos://bucket-a/docs/", + ctx=ctx, + to="viking://resources/x/y", + args=args, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("path", "args", "message"), + [ + ("tos://bucket-a/docs?version=1", {}, "path"), + ("tos://Bucket-a/docs", {}, "path"), + ("tos://bucket-a:443/docs", {}, "path"), + ("tos://user@bucket-a/docs", {}, "path"), + ("tos://bucket-a/docs%2Fv1", {}, "path"), + ("tos://bucket-a//docs", {}, "path"), + ("tos://bucket-a/a\x00b", {}, "path"), + ("tos://bucket-a/a\x7fb", {}, "path"), + ( + "tos://bucket-a", + {"tos_prefix": ["tos://bucket-a"]}, + r"args\.tos_prefix\[0\]", + ), + ( + "tos://bucket-a/docs/", + { + "tos_prefix": [ + "tos://bucket-a/docs/", + "tos://bucket-b/manual?version=1", + ] + }, + r"args\.tos_prefix\[1\]", + ), + ( + "tos://bucket-a/docs/", + { + "tos_prefix": ["tos://bucket-a/docs/"], + "exclude": ["tos://bucket-a/private#latest"], + }, + r"args\.exclude\[0\]", + ), + ( + "tos://bucket-a/docs/", + { + "tos_prefix": ["tos://bucket-a/docs/"], + "exclude": ["tos://bucket-a"], + }, + r"args\.exclude\[0\]", + ), + ], +) +async def test_add_resource_rejects_invalid_tos_uri_before_connector_call( + monkeypatch, + connector_config, + ctx, + service, + path, + args, + message, +): + connector_client_factory = Mock(side_effect=AssertionError("Connector must not be called")) + monkeypatch.setattr(connector_delegate_module, "ConnectorClient", connector_client_factory) + + with pytest.raises(InvalidArgumentError, match=message): + await service.add_resource( + path=path, + ctx=ctx, + to="viking://resources/x/y", + args=args, + ) + + connector_client_factory.assert_not_called() + + +@pytest.mark.asyncio +async def test_add_resource_rejects_top_level_and_tos_args_exclude( + connector_config, + ctx, + service, +): + with pytest.raises(InvalidArgumentError, match="both as a top-level field"): + await service.add_resource( + path="tos://bucket-a/docs/", + ctx=ctx, + to="viking://resources/x/y", + exclude="**/private/**", + args={ + "tos_prefix": ["tos://bucket-a/docs/"], + "exclude": [], + }, + ) + + +@pytest.mark.asyncio +async def test_add_resource_keeps_args_exclude_reserved_outside_tos(ctx, service): + with pytest.raises(InvalidArgumentError, match="core add_resource fields: exclude"): + await service.add_resource( + path="/tmp/document.md", + ctx=ctx, + to="viking://resources/document", + args={"exclude": []}, + ) + + @pytest.mark.asyncio async def test_add_resource_routes_git_repo_to_connector( monkeypatch, @@ -1323,6 +1920,7 @@ async def no_sleep(_seconds): monkeypatch.setattr(connector_delegate_module.asyncio, "sleep", no_sleep) client = SimpleNamespace(get_task_info=AsyncMock(return_value=task_info)) + on_success = AsyncMock() outcome = await ResourceService()._connector._monitor( client=client, @@ -1331,15 +1929,18 @@ async def no_sleep(_seconds): poll_interval_ms=1, timeout_seconds=1, ctx=ctx, + on_success=on_success, ) assert tracker.update_stage.await_args.args[1] == expected_stage if expected_error is None: assert outcome["status"] == "completed" + on_success.assert_awaited_once_with() tracker.complete.assert_awaited_once() tracker.fail.assert_not_awaited() else: assert outcome == {"status": "failed", "error": expected_error} + on_success.assert_not_awaited() assert tracker.fail.await_args.args[1] == expected_error tracker.complete.assert_not_awaited() @@ -1349,6 +1950,7 @@ async def test_monitor_connector_task_retries_transient_polling_error( monkeypatch, connector_config, ctx, + caplog, ): tracker = _task_tracker() monkeypatch.setattr( @@ -1363,22 +1965,24 @@ async def no_sleep(_seconds): client = SimpleNamespace( get_task_info=AsyncMock( side_effect=[ - httpx.ReadTimeout("temporary timeout"), + httpx.ReadTimeout("token=secret"), {"Status": "succeeded"}, ] ) ) - await ResourceService()._connector._monitor( - client=client, - connector_task_key="connector-1", - ov_task_id="task-1", - poll_interval_ms=1, - timeout_seconds=1, - ctx=ctx, - ) + with caplog.at_level(logging.WARNING): + await ResourceService()._connector._monitor( + client=client, + connector_task_key="connector-1", + ov_task_id="task-1", + poll_interval_ms=1, + timeout_seconds=1, + ctx=ctx, + ) assert client.get_task_info.await_count == 2 + assert "secret" not in caplog.text tracker.complete.assert_awaited_once() tracker.fail.assert_not_awaited()