feat: new tos connector args (#4129)

* feat: new tos connector args

* fix: skill ut

* fix: param check

* fix: param check

* fix: param check

* fix: clarify connector watch eventual consistency

* fix: allow plaintext connector watches and tighten TOS URI validation

Author: zihengli-bytedance

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://<bucket>/<path>'."
-                )
-            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://<bucket>/<path>'."
+                    )
+                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,19 +652,23 @@ 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,
             "connector_task_key": connector_task_key,
         }
         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}

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"}),
 }