From c81494321592f497dd528f2195b69572fcdc3e5b Mon Sep 17 00:00:00 2001 From: yang rui <142162413+ruirui6946@users.noreply.github.com> Date: Fri, 25 Sep 2026 15:24:09 +0800 Subject: [PATCH] feat(extensions): grant bounded host model invocation (#5796) * feat(extensions): grant bounded host model invocation * fix(extensions): retain model budgets and isolate schema validation * fix(extensions): normalize provider-origin model cancellation --------- Co-authored-by: Willem Jiang --- README.md | 9 + backend/docs/extension-model-invocation.md | 155 +++++ .../deerflow_extension_api/__init__.py | 22 + .../deerflow_extension_api/contracts.py | 2 + .../model_invocation.py | 70 +++ .../harness/deerflow/extensions/AGENTS.md | 29 +- .../harness/deerflow/extensions/gateway.py | 6 +- .../harness/deerflow/extensions/loader.py | 6 +- .../deerflow/extensions/model_access.py | 94 +++ .../deerflow/extensions/model_invocation.py | 263 ++++++++ .../extensions/model_schema_worker.py | 57 ++ .../harness/deerflow/extensions/registry.py | 10 +- .../test_extension_model_invocation.py | 57 ++ .../tests/test_extension_model_invocation.py | 590 ++++++++++++++++++ config.example.yaml | 8 + 15 files changed, 1373 insertions(+), 5 deletions(-) create mode 100644 backend/docs/extension-model-invocation.md create mode 100644 backend/packages/extension-api/deerflow_extension_api/model_invocation.py create mode 100644 backend/packages/harness/deerflow/extensions/model_access.py create mode 100644 backend/packages/harness/deerflow/extensions/model_invocation.py create mode 100644 backend/packages/harness/deerflow/extensions/model_schema_worker.py create mode 100644 backend/tests/blocking_io/test_extension_model_invocation.py create mode 100644 backend/tests/test_extension_model_invocation.py diff --git a/README.md b/README.md index 3244214a5a..bbe35cf63b 100644 --- a/README.md +++ b/README.md @@ -1405,6 +1405,15 @@ returns `None` when the host does not support the capability; the required helpe these exceptions to HTTP 503 and 403 respectively; resolver failures never fall back to the global service reader. +Extension services can optionally receive a [host model invoker](backend/docs/extension-model-invocation.md) +(extension API 0.2.4+). Operators explicitly map allowed logical roles under +`plugins[].host_access.model_invocation`; the host handles credentials, concurrency, + timeouts, and schema-validated output. Admission is bounded; timed-out provider work +keeps its concurrency slot until completion, and schema validation runs in terminable +child processes. A provider's own cancellation becomes `ModelInvocationFailed`, +while cancellation of the calling extension task still propagates normally. +Without a grant the capability is `None`. + Plugin order is deterministic, per-plugin configuration is passed to `install()`, and `required: true` makes load failure abort startup; otherwise failures are reported and skipped. `enabled: false` skips resolution and import. The manager preserves the extension's diff --git a/backend/docs/extension-model-invocation.md b/backend/docs/extension-model-invocation.md new file mode 100644 index 0000000000..2fc8b9377e --- /dev/null +++ b/backend/docs/extension-model-invocation.md @@ -0,0 +1,155 @@ +# Host model invocation for extensions + +Extension API **0.2.4** adds an optional `ExtensionRuntimeDeps.model_invoker`. +An operator-authorized service can make an asynchronous, non-streaming text call +using DeerFlow's configured models. The host constructs the provider client and +returns plain data; the extension needs no provider credentials or LangChain dependency. + +## Grant and route logical roles + +Add `host_access` to the relevant entry in the operator-owned `config.yaml`: + +```yaml +plugins: + - name: classifier + use: my_classifier:install + host_access: + model_invocation: + roles: + default: my-configured-model + fast: my-small-model + max_concurrency: 2 + timeout_seconds: 60 + max_input_chars: 262144 + max_output_chars: 65536 +``` + +The mapping values are names from `models:`. Role names are deployment-defined; +an omitted request role selects `default`, which must be explicitly mapped. +An extension cannot request an arbitrary configured model name. No grant (or +`model_invocation: null`) means `deps.model_invoker is None`, preserving existing +services. `host_access` is separate from the extension-private `config` passed to +`install()` and cannot be edited through `extensions_config.json`. + +Each successful installation owns its grant and concurrency budget. Multiple +services from that installation share the budget; another installation using +the same Python entry point does not inherit its grant. Grants and model +configuration are bound at service startup. As with all `plugins:` changes, +changing the grant, source, or extension model configuration requires a Gateway +restart. Startup failure revokes that service's handle. Shutdown revokes retained +handles and cancels callers before invoking its `stop()` method. Already-running +provider work retains its budget until it actually finishes; its result is discarded. + +This capability follows the trusted extension model: installed Python code already +runs with Gateway privileges. It is not a sandbox or a per-user authorization API. +A service exposing an HTTP route must authorize callers and their input itself. + +## Call from a service + +```python +from deerflow_extension_api import ( + ModelInvocationRequest, + ModelMessage, + extension, +) + + +class Classifier: + async def start(self, deps): + self.invoker = deps.model_invoker + + async def stop(self): + pass + + async def classify(self, text): + if self.invoker is None: + raise RuntimeError("Operator has not granted model invocation") + result = await self.invoker.invoke(ModelInvocationRequest( + messages=[ModelMessage("user", text)], + purpose="classification", + response_schema={ + "type": "object", + "properties": {"label": {"enum": ["positive", "negative"]}}, + "required": ["label"], + "additionalProperties": False, + }, + )) + return result.structured_output["label"] + + +@extension(api="0.2.4") +def install(registry, config): + registry.service(Classifier()) +``` + +Invoke from the service's event loop. Messages support `system`, `user`, and +`assistant` text, with 1–256 messages per call. `purpose` is an optional 1–128 +character tracing label, not a prompt or an authorization selector. Do not put +secrets in it. The caller can shorten `timeout_seconds`, but cannot extend the +host timeout. Queueing, provider construction, schema-worker startup, and validation +count toward the same deadline. Cancellation propagates as `asyncio.CancelledError`. +Timeout and cancellation stop waiting promptly, but cannot terminate synchronous +provider threads. Both synchronous and asynchronous provider calls therefore keep +their concurrency and admission slots until the actual operation finishes. No +additional provider request starts if construction finishes after its caller left. +Configure provider-side timeouts as well: a stuck provider keeps its slot occupied. +Stopping a service does not wait for those provider operations to finish. + +`response_schema` accepts inline JSON Schema Draft 2020-12 object schemas. +References (`$ref`, `$dynamicRef`, `$recursiveRef`) and other schema dialects are +rejected before dispatch; schema validation never fetches network resources. +The host adds a JSON-only instruction and validates the returned text locally, +so the contract does not depend on a provider's native structured-output feature. +Schema checking and response validation each run in a short-lived isolated Python +process using the host interpreter. This adds process startup overhead to structured +calls. Expensive regexes or schema combinations cannot block the Gateway event loop +or its GIL. On timeout/cancellation the host kills and reaps the child before +releasing admission (pipe cancellation is polled every 50 ms). No schema-worker +process is used for plain text calls. +Malformed JSON, non-finite numbers, and schema violations are explicit failures; +no partially populated success is returned. Batching, row-ID reconciliation, +file handling, and workflow-specific checks belong to the extension. + +Successful results contain `content`, optional validated `structured_output`, +the resolved configured model name, and optional `ModelUsage` token counts. +Provider response metadata, framework messages, and exceptions are not exposed. +The standard model factory retains provider adaptation, configured RPM admission, +and tracing callbacks; invocation metadata attributes calls to their extension +source, logical role, and purpose. + +## Failures and bounds + +| Exception | Meaning | +| --- | --- | +| `ModelInvocationUnavailable` | Unsupported host, stopped capability, wrong event loop, or missing configured model | +| `ModelInvocationUnauthorized` | The requested logical role is not granted | +| `ModelInvocationFailed` | Invalid request, provider failure, timeout, unsupported response, or host limit exceeded | +| `ModelOutputValidationError` | Malformed JSON or schema mismatch; a subclass of `ModelInvocationFailed` | + +All these derive from `ModelInvocationError`. Provider failures use a normalized +message without the original exception chain. Host deadline expiry reports +`Model invocation timed out`; a provider's own `TimeoutError` reports +`Model provider timed out`. A provider task that cancels itself reports +`ModelInvocationFailed("Model provider cancelled")`; it does not cancel the +extension's calling task. Normal caller cancellation remains +`asyncio.CancelledError`, rather than being converted to an ordinary failure. + +Concurrency defaults to 2 (range 1–64), timeout to 60 seconds (maximum 600), and +input/output character limits to 262144/65536 (maximum 1048576 each). At most +`2 * max_concurrency` calls are admitted per installation, including running, +queued, and abandoned-but-still-running provider calls. Excess calls fail immediately +with `ModelInvocationFailed("Model invocation capacity exceeded")`, before payload +conversion or validation. Schema workers share the execution slots, so at most +`max_concurrency` children can run per installation. Input +counting includes the schema instruction. Output bounds apply to returned text +after generation; configure the underlying model's token limits to bound provider +generation costs. This is a per-installation, per-Gateway-process budget, not a +distributed quota. Provider SDK retries retain their configured behavior. + +Tool calls, streaming, agent execution, and conversation persistence are outside +this capability. Jev/TypeSafe typed-decision workflows remain extension-owned; +`model_invoker` is the chat-style path for host-configured models. + +The offline tests in `backend/tests/test_extension_model_invocation.py` cover the +classification success and failure cases, source isolation, shared admission, +timeouts, cancellation, limits, and lifecycle revocation without real API calls. diff --git a/backend/packages/extension-api/deerflow_extension_api/__init__.py b/backend/packages/extension-api/deerflow_extension_api/__init__.py index db05469f6c..a6555fbb40 100644 --- a/backend/packages/extension-api/deerflow_extension_api/__init__.py +++ b/backend/packages/extension-api/deerflow_extension_api/__init__.py @@ -43,6 +43,18 @@ from deerflow_extension_api.contracts import ( TaskOutcome, extension, ) +from deerflow_extension_api.model_invocation import ( + ModelInvocationError, + ModelInvocationFailed, + ModelInvocationRequest, + ModelInvocationResult, + ModelInvocationUnauthorized, + ModelInvocationUnavailable, + ModelInvoker, + ModelMessage, + ModelOutputValidationError, + ModelUsage, +) from deerflow_extension_api.placement import ( AgentBuildContext, AgentScope, @@ -89,6 +101,16 @@ from deerflow_extension_api.state import ExtensionData API_VERSION = "0.2.4" __all__ = [ + "ModelInvocationError", + "ModelInvocationFailed", + "ModelInvocationRequest", + "ModelInvocationResult", + "ModelInvocationUnauthorized", + "ModelInvocationUnavailable", + "ModelInvoker", + "ModelMessage", + "ModelOutputValidationError", + "ModelUsage", "ActionContext", "BackendAction", "BrowserAssets", diff --git a/backend/packages/extension-api/deerflow_extension_api/contracts.py b/backend/packages/extension-api/deerflow_extension_api/contracts.py index 54a30156d4..98010dff30 100644 --- a/backend/packages/extension-api/deerflow_extension_api/contracts.py +++ b/backend/packages/extension-api/deerflow_extension_api/contracts.py @@ -20,6 +20,7 @@ from deerflow_extension_api.state import ExtensionData if TYPE_CHECKING: # pragma: no cover - typing only from deerflow_extension_api.assembly import AgentAssemblyObserver from deerflow_extension_api.compaction import ContextCompactionObserver + from deerflow_extension_api.model_invocation import ModelInvoker from deerflow_extension_api.placement import AgentBuildContext, MiddlewarePlacement from deerflow_extension_api.run_evidence import RunEvidenceReader @@ -167,6 +168,7 @@ class ExtensionRuntimeDeps: policy: HostPolicySnapshot = field(default_factory=HostPolicySnapshot) session_factory: Any | None = None run_evidence_reader: RunEvidenceReader | None = None + model_invoker: ModelInvoker | None = None class ExtensionService(Protocol): diff --git a/backend/packages/extension-api/deerflow_extension_api/model_invocation.py b/backend/packages/extension-api/deerflow_extension_api/model_invocation.py new file mode 100644 index 0000000000..a51a7ed30a --- /dev/null +++ b/backend/packages/extension-api/deerflow_extension_api/model_invocation.py @@ -0,0 +1,70 @@ +"""Provider-neutral, non-streaming text model calls granted by the host.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Literal, Protocol + + +class ModelInvocationError(RuntimeError): + """Normalized host error; contains no provider exception or credentials.""" + + +class ModelInvocationUnavailable(ModelInvocationError): + """The capability has stopped or its configured model is unavailable.""" + + +class ModelInvocationUnauthorized(ModelInvocationError): + """The requested logical role is not granted to this installation.""" + + +class ModelInvocationFailed(ModelInvocationError): + """The invocation failed, exceeded a host limit or timed out.""" + + +class ModelOutputValidationError(ModelInvocationFailed): + """Provider output could not be parsed or validated against the schema.""" + + +@dataclass(frozen=True) +class ModelMessage: + role: Literal["system", "user", "assistant"] + content: str + + +@dataclass(frozen=True) +class ModelInvocationRequest: + messages: Sequence[ModelMessage] + model_role: str | None = None + purpose: str | None = None + response_schema: Mapping[str, object] | None = None + timeout_seconds: float | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "messages", tuple(self.messages)) + + +@dataclass(frozen=True) +class ModelUsage: + input_tokens: int | None = None + output_tokens: int | None = None + total_tokens: int | None = None + + +@dataclass(frozen=True) +class ModelInvocationResult: + content: str + structured_output: Mapping[str, object] | None = None + resolved_model: str | None = None + usage: ModelUsage | None = None + + +class ModelInvoker(Protocol): + async def invoke(self, request: ModelInvocationRequest) -> ModelInvocationResult: + """Invoke once; schema requests succeed only with validated object data. + + Caller cancellation propagates as asyncio.CancelledError. A timeout + (including time waiting for host capacity) raises ModelInvocationFailed. + """ + raise ModelInvocationUnavailable("Model invocation is not supported by this host") diff --git a/backend/packages/harness/deerflow/extensions/AGENTS.md b/backend/packages/harness/deerflow/extensions/AGENTS.md index 5df2f83b04..ab282f603e 100644 --- a/backend/packages/harness/deerflow/extensions/AGENTS.md +++ b/backend/packages/harness/deerflow/extensions/AGENTS.md @@ -277,9 +277,34 @@ detached task store, the same fallback `notify_system_model_call` uses when its supplies none. Gateway services start in registration order after the persistence engine and session -factory are ready. Each receives the same `ExtensionRuntimeDeps` snapshot containing the +factory are ready. Ungranted services share an `ExtensionRuntimeDeps` snapshot containing the app store, projected host policy, session factory, and optional read-only -`RunEvidenceReader`. The Gateway constructs the configured run and event stores before +`RunEvidenceReader`. `plugins[].host_access.model_invocation` optionally binds a model invoker +to each service via the host-only `ModelInvocationService` adapter. The loader captures +one `ModelInvocationScope` per installation, not per `use` string, so duplicate sources +cannot inherit one another's roles. Its semaphore is shared by that installation's +services; its admission ceiling is twice the concurrency limit, checked before +payload processing. Provider work is shielded from caller cancellation and retains +both budgets until actual completion, including synchronous LangChain executor calls +and offloaded construction. Abandoned construction cannot dispatch a model request. +Provider-task cancellation is a normalized failure; only a new cancellation of +the invoking task propagates. Compare cancellation counts against invocation entry +so previously handled caller cancellations do not mask provider failures. +Failed-install positional rollback also removes its adapters. The adapter +receives startup config through `start_with_host`, while extensions receive only the +neutral invoker in a replaced deps snapshot. No-grant services preserve their old path. +Failed start and stop revoke the service's handle and cancel queued/in-flight +callers; they do not release slots owned by still-running provider work. Structured +schema checks and output validation run in terminable isolated Python children, +with pipe I/O on admission-bounded dedicated threads (Windows selector-loop compatible, +independent of a potentially saturated provider executor). Cancellation kills and +reaps those children before releasing admission. +Grants and model profiles are startup snapshots; changing them requires restarting the +Gateway. Calls use the normal model factory and attributed tracing, return plain text, +usage counts and optionally locally validated JSON objects, and never return raw model +objects or provider exception chains. See `backend/docs/extension-model-invocation.md`. + +The Gateway constructs the configured run and event stores before services so the reader is usable from `start()`. Changed-run discovery uses an opaque, scope-bound cursor over `(change_seq, run_id)`; a run that changes after it was returned may be replayed, but an unreturned run cannot be skipped. Legacy rows start at `change_seq=0` diff --git a/backend/packages/harness/deerflow/extensions/gateway.py b/backend/packages/harness/deerflow/extensions/gateway.py index 14e66f4143..d0bf2c903f 100644 --- a/backend/packages/harness/deerflow/extensions/gateway.py +++ b/backend/packages/harness/deerflow/extensions/gateway.py @@ -22,6 +22,7 @@ from typing import Any from deerflow_extension_api import ExtensionRuntimeDeps from deerflow.extensions.loader import Diagnostic +from deerflow.extensions.model_access import ModelInvocationService from deerflow.extensions.policy import project_host_policy from deerflow.extensions.registry import LoadedExtensions @@ -599,7 +600,10 @@ async def start_services( attempted_services.append(entry) cancellation_count = _cancellation_count() try: - await service.start(deps) + if isinstance(service, ModelInvocationService): + await service.start_with_host(deps, app_config) + else: + await service.start(deps) except asyncio.CancelledError: if _cancellation_count() > cancellation_count: raise diff --git a/backend/packages/harness/deerflow/extensions/loader.py b/backend/packages/harness/deerflow/extensions/loader.py index f27c385882..41bba4d447 100644 --- a/backend/packages/harness/deerflow/extensions/loader.py +++ b/backend/packages/harness/deerflow/extensions/loader.py @@ -16,6 +16,7 @@ from typing import Any, Literal from deerflow_extension_api import API_VERSION from pydantic import BaseModel, ConfigDict, Field +from deerflow.extensions.model_access import ExtensionHostAccess, ModelInvocationScope from deerflow.extensions.registry import ExtensionRegistry, LoadedExtensions from deerflow.persistence.migrations._env_filters import register_extension_table_prefix from deerflow.reflection import resolve_variable @@ -43,6 +44,7 @@ class ExtensionSpec(BaseModel): description="Installed Python distribution recorded by the extension manager", ) use: str = Field(description="Entry point path, e.g. 'my_extension:install'") + host_access: ExtensionHostAccess = Field(default_factory=ExtensionHostAccess) config: dict[str, Any] = Field( default_factory=dict, description="Extension-private configuration, passed to install() verbatim", @@ -237,7 +239,9 @@ def load_extensions(specs: Sequence[ExtensionSpec]) -> tuple[LoadedExtensions, l # installed instance that happens to share this spec's `use`. mark = registry.mark() try: - with registry.attributed_to(spec.use): + grant = spec.host_access.model_invocation + scope = ModelInvocationScope(spec.use, grant) if grant is not None else None + with registry.attributed_to(spec.use, model_access=scope): install(registry, _frozen_config(spec.config)) except Exception as exc: registry.rollback_to(mark) diff --git a/backend/packages/harness/deerflow/extensions/model_access.py b/backend/packages/harness/deerflow/extensions/model_access.py new file mode 100644 index 0000000000..0fa84a5d9e --- /dev/null +++ b/backend/packages/harness/deerflow/extensions/model_access.py @@ -0,0 +1,94 @@ +"""Operator grants and service lifetime binding; no provider imports at startup.""" + +from __future__ import annotations + +import asyncio +from dataclasses import replace +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +class ModelInvocationGrant(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, allow_inf_nan=False) + + roles: dict[str, str] = Field(min_length=1, description="Logical role to configured model name") + max_concurrency: int = Field(default=2, ge=1, le=64) + timeout_seconds: float = Field(default=60, gt=0, le=600) + max_input_chars: int = Field(default=262144, ge=1, le=1048576) + max_output_chars: int = Field(default=65536, ge=1, le=1048576) + + @field_validator("roles") + @classmethod + def validate_roles(cls, roles): + if any(not value or value != value.strip() for pair in roles.items() for value in pair): + raise ValueError("Role and model names must be nonempty without surrounding whitespace") + return roles + + +class ExtensionHostAccess(BaseModel): + model_config = ConfigDict(extra="forbid") + + model_invocation: ModelInvocationGrant | None = None + + +class ModelInvocationScope: + """One install() owns one budget, even if it registers multiple services.""" + + def __init__(self, source: str, grant: ModelInvocationGrant): + self.source = source + self.grant = grant.model_copy(deep=True) + self.budget: ModelInvocationBudget | None = None + + def bind(self, app_config: Any): + from deerflow.extensions.model_invocation import HostModelInvoker + + if self.budget is None: + self.budget = ModelInvocationBudget(self.grant.max_concurrency) + return HostModelInvoker(self.source, self.grant, self.budget, app_config) + + +class ModelInvocationBudget: + """Loop-owned admission shared by every service in one installation.""" + + def __init__(self, concurrency: int): + self.semaphore = asyncio.Semaphore(concurrency) + self.capacity = 2 * concurrency + self.admitted = 0 + # Strong references keep abandoned, shielded provider work alive. + self.workers: set[asyncio.Task] = set() + + def retain(self, task): + self.workers.add(task) + task.add_done_callback(self._finished) + + def _finished(self, task): + self.workers.discard(task) + if not task.cancelled(): + task.exception() # Consume abandoned failures without exposing provider text. + + +class ModelInvocationService: + """Host adapter keeps capability revocation paired with service cleanup.""" + + def __init__(self, service: Any, scope: ModelInvocationScope): + self.service = service + self.scope = scope + self.invoker = None + + async def start(self, deps): + raise RuntimeError("Granted services must be started through the host lifecycle") + + async def start_with_host(self, deps, app_config): + self.invoker = self.scope.bind(app_config) + try: + await self.service.start(replace(deps, model_invoker=self.invoker)) + except BaseException: + # Includes host cancellation: start() may already have spawned calls. + self.invoker.close() + raise + + async def stop(self): + if self.invoker is not None: + self.invoker.close() + await self.service.stop() diff --git a/backend/packages/harness/deerflow/extensions/model_invocation.py b/backend/packages/harness/deerflow/extensions/model_invocation.py new file mode 100644 index 0000000000..d226577753 --- /dev/null +++ b/backend/packages/harness/deerflow/extensions/model_invocation.py @@ -0,0 +1,263 @@ +"""Host-owned model construction, bounded invocation and neutral projection.""" + +from __future__ import annotations + +import asyncio +import json +import logging +import math +import subprocess +import sys +import threading +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path + +from deerflow_extension_api import ( + ModelInvocationError, + ModelInvocationFailed, + ModelInvocationRequest, + ModelInvocationResult, + ModelInvocationUnauthorized, + ModelInvocationUnavailable, + ModelMessage, + ModelOutputValidationError, + ModelUsage, +) +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage + +from deerflow.models import create_chat_model + +logger = logging.getLogger(__name__) +_MESSAGE_TYPES = {"system": SystemMessage, "user": HumanMessage, "assistant": AIMessage} + + +@dataclass +class _Invocation: + acquired: bool = False + abandoned: bool = False + provider: asyncio.Task | None = None + + +async def _schema_validator(schema_json, content=None): + """Run CPU-bound JSON Schema work outside the Gateway process/GIL. + + A dedicated pipe-I/O thread also works with Windows' selector loop and cannot + be starved by sync providers occupying asyncio's default executor. Admission + bounds both these threads and child processes. Cancellation kills and reaps + the child before this call releases its slot. + """ + loop = asyncio.get_running_loop() + finished = loop.create_future() + cancelled = threading.Event() + payload = schema_json + "\n" + if content is not None: + payload += json.dumps(content) + "\n" + + def work(): + status = b"" + try: + options = {"creationflags": subprocess.CREATE_NO_WINDOW} if sys.platform == "win32" else {} + with subprocess.Popen( + [sys.executable, "-I", str(Path(__file__).with_name("model_schema_worker.py"))], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + **options, + ) as process: + try: + data = payload.encode("utf-8") + while not cancelled.is_set(): + try: + status, _ = process.communicate(data, timeout=0.05) + break + except subprocess.TimeoutExpired: + data = None # communicate retains pending input for the next poll. + finally: + if process.poll() is None: + process.kill() + process.communicate() + except Exception: + pass # Only fixed validation status crosses the public error boundary. + finally: + loop.call_soon_threadsafe(finished.set_result, status.split()) + + threading.Thread(target=work, name="extension-schema-validation", daemon=True).start() + interrupted = False + while not finished.done(): + try: + await asyncio.shield(finished) + except asyncio.CancelledError: + interrupted = True + cancelled.set() + if interrupted: + raise asyncio.CancelledError + status = finished.result() + if content is None: + if status != [b"ready"]: + raise ModelInvocationFailed("response_schema must be a valid inline Draft 2020-12 object schema") + elif status != [b"ready", b"valid"]: + raise ModelOutputValidationError("Model output did not match response_schema") + + +def _usage(metadata): + if not isinstance(metadata, Mapping): + return None + + def token_count(key): + value = metadata.get(key) + return value if type(value) is int and value >= 0 else None + + return ModelUsage(token_count("input_tokens"), token_count("output_tokens"), token_count("total_tokens")) + + +class HostModelInvoker: + def __init__(self, source, grant, budget, app_config): + self._source = source + self._grant = grant + self._budget = budget + self._app_config = app_config + self._loop = asyncio.get_running_loop() + self._closed = False + self._tasks = set() + + def close(self): + """Revoke handles and cancel callers; running providers retain their slots.""" + self._closed = True + for task in tuple(self._tasks): + task.cancel() + + async def invoke(self, request: ModelInvocationRequest) -> ModelInvocationResult: + if self._closed: + raise ModelInvocationUnavailable("Model invocation capability has stopped") + if asyncio.get_running_loop() is not self._loop: + raise ModelInvocationUnavailable("Model invocation requires the service event loop") + if self._budget.admitted >= self._budget.capacity: + raise ModelInvocationFailed("Model invocation capacity exceeded") + self._budget.admitted += 1 + call = _Invocation() + deadline = None + task = asyncio.current_task() + cancellation_count = task.cancelling() + self._tasks.add(task) + try: + if not isinstance(request, ModelInvocationRequest): + raise ModelInvocationFailed("Expected ModelInvocationRequest") + timeout = request.timeout_seconds + if timeout is not None and (type(timeout) not in (int, float) or not math.isfinite(timeout) or timeout <= 0): + raise ModelInvocationFailed("timeout_seconds must be finite and positive") + timeout = min(timeout, self._grant.timeout_seconds) if timeout is not None else self._grant.timeout_seconds + deadline = asyncio.timeout(timeout) + async with deadline: + await self._budget.semaphore.acquire() + call.acquired = True + return await self._invoke(request, call) + except asyncio.CancelledError: + # shield also raises when only the provider task is cancelled. A + # previously handled caller cancellation must not mask that failure. + if task.cancelling() > cancellation_count or call.provider is None or not call.provider.cancelled(): + raise + failure = ModelInvocationFailed("Model provider cancelled") + except ModelInvocationError as exc: + # Copy only our normalized text. In particular, do not pass a JSON, + # schema or provider exception through __context__ to extensions. + failure = type(exc)(str(exc)) + except TimeoutError: + failure = ModelInvocationFailed("Model invocation timed out" if deadline is not None and deadline.expired() else "Model provider timed out") + except Exception as exc: + logger.warning("Extension model invocation failed for %s (%s)", self._source, type(exc).__name__) + failure = ModelInvocationFailed("Model invocation failed") + finally: + self._tasks.discard(task) + call.abandoned = True + + def release(_=None): + self._budget.admitted -= 1 + if call.acquired: + self._budget.semaphore.release() + + if call.provider is not None and not call.provider.done(): + call.provider.add_done_callback(release) + else: + release() + raise failure from None + + async def _provider_call(self, name, messages, role, purpose, call): + # Neither to_thread construction nor a sync-backed ainvoke is cancellable. + # The caller may leave, but the real work keeps its admission/slot until done. + model = await asyncio.to_thread(create_chat_model, name, app_config=self._app_config) + if self._closed or call.abandoned: + raise ModelInvocationUnavailable("Model invocation capability has stopped") + return await model.ainvoke( + messages, + config={"run_name": "extension_model_invocation", "metadata": {"extension_source": self._source, "extension_model_role": role, "extension_purpose": purpose}}, + ) + + async def _invoke(self, request, call): + role = request.model_role if request.model_role is not None else "default" + if not isinstance(role, str) or role not in self._grant.roles: + raise ModelInvocationUnauthorized("Model role is not granted to this extension") + name = self._grant.roles[role] + if self._app_config.get_model_config(name) is None: + raise ModelInvocationUnavailable("Configured model is unavailable") + if request.purpose is not None and (not isinstance(request.purpose, str) or not 1 <= len(request.purpose) <= 128): + raise ModelInvocationFailed("purpose must contain 1 to 128 characters") + if not request.messages or len(request.messages) > 256: + raise ModelInvocationFailed("Expected 1 to 256 text messages") + messages = [] + input_chars = 0 + for message in request.messages: + if not isinstance(message, ModelMessage) or not isinstance(message.role, str) or message.role not in _MESSAGE_TYPES or not isinstance(message.content, str): + raise ModelInvocationFailed("Only system, user and assistant text messages are supported") + input_chars += len(message.content) + if input_chars > self._grant.max_input_chars: + raise ModelInvocationFailed("Model input exceeds host limit") + messages.append(_MESSAGE_TYPES[message.role](content=message.content)) + + schema_json = None + if request.response_schema is not None: + if not isinstance(request.response_schema, Mapping): + raise ModelInvocationFailed("response_schema must be an object schema") + # Snapshot the schema before validation crosses the process boundary. + schema_json = json.dumps(dict(request.response_schema), allow_nan=False) + instruction = "Return only a JSON object matching the supplied response schema (no Markdown)." + schema_data = "Required response schema (JSON data):\n" + schema_json + if input_chars + len(instruction) + len(schema_data) > self._grant.max_input_chars: + raise ModelInvocationFailed("Model input including schema exceeds host limit") + messages.insert(0, SystemMessage(content=instruction)) + # Schema descriptions may contain extension-owned input. Keep them + # out of the host's fixed system instruction. + messages.append(HumanMessage(content=schema_data)) + await _schema_validator(schema_json) + + if self._closed: + raise ModelInvocationUnavailable("Model invocation capability has stopped") + call.provider = asyncio.create_task(self._provider_call(name, messages, role, request.purpose, call)) + self._budget.retain(call.provider) + response = await asyncio.shield(call.provider) + + if self._closed: + raise ModelInvocationUnavailable("Model invocation capability has stopped") + + if getattr(response, "tool_calls", None) or getattr(response, "invalid_tool_calls", None): + raise ModelInvocationFailed("Tool-call responses are not supported") + content = response.content + if isinstance(content, list): + parts = [] + for block in content: + if isinstance(block, str): + parts.append(block) + elif isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str): + parts.append(block["text"]) + else: + raise ModelInvocationFailed("Model returned non-text content") + content = "".join(parts) + if not isinstance(content, str): + raise ModelInvocationFailed("Model returned non-text content") + if len(content) > self._grant.max_output_chars: + raise ModelInvocationFailed("Model output exceeds host limit") + structured = None + if schema_json is not None: + await _schema_validator(schema_json, content) + structured = json.loads(content) + return ModelInvocationResult(content, structured, name, _usage(getattr(response, "usage_metadata", None))) diff --git a/backend/packages/harness/deerflow/extensions/model_schema_worker.py b/backend/packages/harness/deerflow/extensions/model_schema_worker.py new file mode 100644 index 0000000000..03336c8283 --- /dev/null +++ b/backend/packages/harness/deerflow/extensions/model_schema_worker.py @@ -0,0 +1,57 @@ +"""One-call schema worker, launched by file path in an isolated Python process. + +Only bounded JSON input and fixed status words cross the pipes. Expensive schema +checking and regex validation must never execute on the Gateway loop or its GIL. +""" + +import json +import sys + +from jsonschema import Draft202012Validator + + +def _reject_constant(value): + raise ValueError("Non-finite JSON number") + + +def _check_inline(value): + if isinstance(value, dict): + if any(key in value for key in ("$ref", "$dynamicRef", "$recursiveRef")): + raise ValueError("References are unsupported") + if "$schema" in value and value["$schema"] != "https://json-schema.org/draft/2020-12/schema": + raise ValueError("Unsupported dialect") + for child in value.values(): + _check_inline(child) + elif isinstance(value, list): + for child in value: + _check_inline(child) + + +def main(): + try: + schema = json.loads(sys.stdin.buffer.readline()) + if schema.get("type") != "object": + raise ValueError("Expected object schema") + _check_inline(schema) + Draft202012Validator.check_schema(schema) + validator = Draft202012Validator(schema) + except Exception: + print("invalid_schema", flush=True) + return + print("ready", flush=True) + response = sys.stdin.buffer.readline() + if not response: + return + try: + content = json.loads(response) + structured = json.loads(content, parse_constant=_reject_constant) + json.dumps(structured, allow_nan=False) + validator.validate(structured) + except Exception: + print("invalid_output", flush=True) + return + print("valid", flush=True) + + +if __name__ == "__main__": + main() diff --git a/backend/packages/harness/deerflow/extensions/registry.py b/backend/packages/harness/deerflow/extensions/registry.py index 8b17c78656..1c77abcb2a 100644 --- a/backend/packages/harness/deerflow/extensions/registry.py +++ b/backend/packages/harness/deerflow/extensions/registry.py @@ -26,6 +26,8 @@ from deerflow_extension_api import ( from deerflow_extension_api import ExtensionRegistry as ExtensionRegistryContract from deerflow_extension_api.plugins import BrowserAssets, BrowserModule, PluginContribution +from deerflow.extensions.model_access import ModelInvocationScope, ModelInvocationService + _Entry = tuple[str, Any] @@ -83,16 +85,20 @@ class ExtensionRegistry(ExtensionRegistryContract): self._routers: list[_Entry] = [] self._plugins: list[_Entry] = [] self._current_source: str | None = None + self._model_access: ModelInvocationScope | None = None @contextmanager - def attributed_to(self, source: str) -> Iterator[None]: + def attributed_to(self, source: str, *, model_access: ModelInvocationScope | None = None) -> Iterator[None]: """Attribute everything registered inside the block to ``source``.""" previous = self._current_source + previous_access = self._model_access self._current_source = source + self._model_access = model_access try: yield finally: self._current_source = previous + self._model_access = previous_access def _source(self) -> str: if self._current_source is None: @@ -154,6 +160,8 @@ class ExtensionRegistry(ExtensionRegistryContract): self._context_compaction_observers.append((self._source(), observer)) def service(self, service: ExtensionService) -> None: + if self._model_access is not None: + service = ModelInvocationService(service, self._model_access) self._services.append((self._source(), service)) def routers(self, routers: Sequence[Any]) -> None: diff --git a/backend/tests/blocking_io/test_extension_model_invocation.py b/backend/tests/blocking_io/test_extension_model_invocation.py new file mode 100644 index 0000000000..92ba5bc6c6 --- /dev/null +++ b/backend/tests/blocking_io/test_extension_model_invocation.py @@ -0,0 +1,57 @@ +"""Provider construction may read credentials but must never block the host loop.""" + +import asyncio +import threading +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from deerflow_extension_api import ModelInvocationRequest, ModelMessage +from langchain_core.messages import AIMessage + +from deerflow.extensions import model_invocation +from deerflow.extensions.model_access import ModelInvocationBudget, ModelInvocationGrant + + +@pytest.mark.asyncio +async def test_provider_construction_runs_off_loop(tmp_path, monkeypatch): + credentials = tmp_path / "provider.txt" + await asyncio.to_thread(credentials.write_text, "test-only", encoding="utf-8") + threads = [] + + def factory(name, *, app_config): + assert credentials.read_text(encoding="utf-8") == "test-only" + threads.append(threading.get_ident()) + return SimpleNamespace(ainvoke=AsyncMock(return_value=AIMessage(content="ok"))) + + monkeypatch.setattr(model_invocation, "create_chat_model", factory) + invoker = model_invocation.HostModelInvoker( + "example:install", + ModelInvocationGrant(roles={"default": "model"}), + ModelInvocationBudget(1), + SimpleNamespace(get_model_config=lambda _: object()), + ) + result = await invoker.invoke(ModelInvocationRequest([ModelMessage("user", "hello")])) + assert result.content == "ok" + assert threads and threads[0] != threading.get_ident() + invoker.close() + + +@pytest.mark.asyncio +async def test_schema_subprocess_io_runs_off_loop(monkeypatch): + model = SimpleNamespace(ainvoke=AsyncMock(return_value=AIMessage(content='{"label":"你好"}'))) + monkeypatch.setattr(model_invocation, "create_chat_model", lambda *args, **kwargs: model) + invoker = model_invocation.HostModelInvoker( + "example:install", + ModelInvocationGrant(roles={"default": "model"}), + ModelInvocationBudget(1), + SimpleNamespace(get_model_config=lambda _: object()), + ) + result = await invoker.invoke( + ModelInvocationRequest( + [ModelMessage("user", "hello")], + response_schema={"type": "object", "properties": {"label": {"type": "string"}}}, + ) + ) + assert result.structured_output == {"label": "你好"} + invoker.close() diff --git a/backend/tests/test_extension_model_invocation.py b/backend/tests/test_extension_model_invocation.py new file mode 100644 index 0000000000..62e63581c1 --- /dev/null +++ b/backend/tests/test_extension_model_invocation.py @@ -0,0 +1,590 @@ +"""Granted extension calls exercise the loader, service lifecycle and host adapter.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from deerflow_extension_api import ( + ExtensionRuntimeDeps, + ModelInvocationFailed, + ModelInvocationRequest, + ModelInvocationUnauthorized, + ModelInvocationUnavailable, + ModelMessage, + ModelOutputValidationError, +) +from langchain_core.messages import AIMessage + +from deerflow.extensions.gateway import start_services, stop_services +from deerflow.extensions.loader import ExtensionSpec, load_extensions + + +class Service: + async def start(self, deps): + self.deps = deps + + async def stop(self): + pass + + +@pytest.fixture +def host(monkeypatch): + from deerflow.extensions import model_invocation + + model = SimpleNamespace(ainvoke=AsyncMock(return_value=AIMessage(content='{"label":"positive"}', usage_metadata={"input_tokens": 8, "output_tokens": 4, "total_tokens": 12}))) + factory = Mock(return_value=model) + monkeypatch.setattr(model_invocation, "create_chat_model", factory) + config = SimpleNamespace(get_model_config=lambda name: object() if name == "host-model" else None) + + async def start(grants, *, services_per_install=1, service_type=Service): + services = [] + + def install(registry, config): + for _ in range(services_per_install): + service = service_type() + services.append(service) + registry.service(service) + + monkeypatch.setattr("deerflow.extensions.loader.resolve_variable", lambda _: install) + specs = [ExtensionSpec(use="example:install", host_access={"model_invocation": grant} if grant else {}) for grant in grants] + loaded, diagnostics = load_extensions(specs) + assert not diagnostics + diagnostics = await start_services(loaded, config, None) + return loaded, services, diagnostics + + return SimpleNamespace(start=start, model=model, factory=factory) + + +GRANT = {"roles": {"default": "host-model"}} +SCHEMA = {"type": "object", "properties": {"label": {"enum": ["positive", "negative"]}}, "required": ["label"]} + + +def request(**kwargs): + return ModelInvocationRequest(messages=[ModelMessage("user", "Classify this text")], **kwargs) + + +def test_old_extension_has_no_capability(): + assert ExtensionRuntimeDeps().model_invoker is None + + +@pytest.mark.asyncio +async def test_grant_is_bound_to_installation_not_entrypoint(host): + loaded, services, diagnostics = await host.start([GRANT, None, {"roles": {"fast": "host-model"}}]) + assert not diagnostics + assert services[1].deps.model_invoker is None + result = await services[0].deps.model_invoker.invoke(request(response_schema=SCHEMA, purpose="classify")) + assert result.content == '{"label":"positive"}' + assert result.structured_output == {"label": "positive"} + assert result.resolved_model == "host-model" + assert result.usage.total_tokens == 12 + with pytest.raises(ModelInvocationUnauthorized): + await services[2].deps.model_invoker.invoke(request()) + assert host.factory.call_count == 1 + metadata = host.model.ainvoke.call_args.kwargs["config"]["metadata"] + assert metadata["extension_source"] == "example:install" + assert metadata["extension_purpose"] == "classify" + messages = host.model.ainvoke.call_args.args[0] + assert messages[0].type == "system" + assert '"properties"' not in messages[0].content + assert messages[-1].type == "human" + assert '"properties"' in messages[-1].content + await stop_services(loaded) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("content", ["not JSON", '{"label":"unknown"}', "[]", '{"label":"positive", "score":NaN}', '{"label":"positive", "score":1e999}']) +async def test_schema_failure_never_returns_success(host, content): + host.model.ainvoke.return_value = AIMessage(content=content) + loaded, services, _ = await host.start([GRANT]) + with pytest.raises(ModelOutputValidationError): + await services[0].deps.model_invoker.invoke(request(response_schema=SCHEMA)) + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_provider_errors_do_not_expose_credentials(host): + host.model.ainvoke.side_effect = RuntimeError("secret-key in provider URL") + loaded, services, _ = await host.start([GRANT]) + with pytest.raises(ModelInvocationFailed) as error: + await services[0].deps.model_invoker.invoke(request()) + assert "secret-key" not in str(error.value) + assert error.value.__cause__ is None + assert error.value.__context__ is None + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_missing_model_and_disallowed_role_fail_before_provider(host): + loaded, services, _ = await host.start([{"roles": {"default": "missing"}}]) + with pytest.raises(ModelInvocationUnavailable): + await services[0].deps.model_invoker.invoke(request()) + with pytest.raises(ModelInvocationUnauthorized): + await services[0].deps.model_invoker.invoke(request(model_role="host-model")) + host.factory.assert_not_called() + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_concurrency_shared_across_services_and_queue_timeout(host): + started = asyncio.Event() + release = asyncio.Event() + + async def invoke(*args, **kwargs): + started.set() + await release.wait() + return AIMessage(content="ok") + + host.model.ainvoke.side_effect = invoke + loaded, services, _ = await host.start([{**GRANT, "max_concurrency": 1}], services_per_install=2) + first = asyncio.create_task(services[0].deps.model_invoker.invoke(request())) + await started.wait() + with pytest.raises(ModelInvocationFailed, match="timed out"): + await services[1].deps.model_invoker.invoke(request(timeout_seconds=0.01)) + assert host.model.ainvoke.call_count == 1 + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + release.set() + assert (await services[1].deps.model_invoker.invoke(request())).content == "ok" + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_stop_revokes_retained_capability_and_cancels_inflight(host): + started = asyncio.Event() + release = asyncio.Event() + + async def invoke(*args, **kwargs): + started.set() + await release.wait() + + host.model.ainvoke.side_effect = invoke + loaded, services, _ = await host.start([GRANT]) + invoker = services[0].deps.model_invoker + task = asyncio.create_task(invoker.invoke(request())) + await started.wait() + await stop_services(loaded) + with pytest.raises(asyncio.CancelledError): + await task + with pytest.raises(ModelInvocationUnavailable): + await invoker.invoke(request()) + release.set() + await asyncio.gather(*invoker._budget.workers) + + +@pytest.mark.asyncio +async def test_start_failure_revokes_only_failed_service(host): + class Broken(Service): + async def start(self, deps): + await super().start(deps) + raise ValueError("broken") + + instances = iter([Broken(), Service()]) + loaded, services, diagnostics = await host.start([GRANT], services_per_install=2, service_type=instances.__next__) + assert len(diagnostics) == 1 + with pytest.raises(ModelInvocationUnavailable): + await services[0].deps.model_invoker.invoke(request()) + assert (await services[1].deps.model_invoker.invoke(request())).content == '{"label":"positive"}' + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_doc_classification_example_uses_public_contract(host): + from pathlib import Path + + doc = (Path(__file__).resolve().parents[1] / "docs" / "extension-model-invocation.md").read_text(encoding="utf-8") + code = doc.split("```python\n", 1)[1].split("```", 1)[0] + namespace = {} + exec(compile(code, "extension-model-invocation.md", "exec"), namespace) + loaded, services, _ = await host.start([GRANT], service_type=namespace["Classifier"]) + assert await services[0].classify("Excellent work") == "positive" + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_real_host_factory_preserves_tracing_and_text_contract(monkeypatch): + from langchain_core.callbacks import BaseCallbackHandler + + from deerflow.config.app_config import AppConfig + from deerflow.config.model_config import ModelConfig + from deerflow.config.sandbox_config import SandboxConfig + + traces = [] + + class Observer(BaseCallbackHandler): + def on_chat_model_start(self, serialized, messages, **kwargs): + traces.append(kwargs["metadata"]) + + config = AppConfig( + models=[ModelConfig(name="host-model", model="fake", use="langchain_core.language_models.fake_chat_models:FakeListChatModel", responses=['{"label":"negative"}'])], + sandbox=SandboxConfig(use="deerflow.sandbox.local:LocalSandboxProvider"), + ) + service = Service() + monkeypatch.setattr("deerflow.extensions.loader.resolve_variable", lambda _: lambda registry, _: registry.service(service)) + monkeypatch.setattr("deerflow.models.factory.build_tracing_callbacks", lambda: [Observer()]) + loaded, diagnostics = load_extensions([ExtensionSpec(use="real:install", host_access={"model_invocation": GRANT})]) + assert not diagnostics + assert not await start_services(loaded, config, None) + result = await service.deps.model_invoker.invoke(request(response_schema=SCHEMA)) + assert result.structured_output == {"label": "negative"} + assert traces[0]["extension_source"] == "real:install" + await stop_services(loaded) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "schema", + [ + {"type": "object", "properties": {"x": {"$ref": "https://example.com/schema"}}}, + {"type": "object", "$ref": "#/properties/x"}, + {"type": "array"}, + {"type": "object", "required": "label"}, + {"type": "object", "$schema": "http://json-schema.org/draft-07/schema#"}, + ], +) +async def test_invalid_schema_rejected_without_provider_call(host, schema): + loaded, services, _ = await host.start([GRANT]) + with pytest.raises(ModelInvocationFailed): + await services[0].deps.model_invoker.invoke(request(response_schema=schema)) + host.factory.assert_not_called() + await stop_services(loaded) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("timeout", [0, -1, float("nan"), float("inf"), True, "60"]) +async def test_invalid_timeout_rejected_before_provider(host, timeout): + loaded, services, _ = await host.start([GRANT]) + with pytest.raises(ModelInvocationFailed): + await services[0].deps.model_invoker.invoke(request(timeout_seconds=timeout)) + host.factory.assert_not_called() + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_host_timeout_caps_request_and_preserves_running_provider(host): + cancelled = asyncio.Event() + release = asyncio.Event() + + async def invoke(*args, **kwargs): + try: + await release.wait() + finally: + cancelled.set() + + host.model.ainvoke.side_effect = invoke + loaded, services, _ = await host.start([{**GRANT, "timeout_seconds": 0.05}]) + with pytest.raises(ModelInvocationFailed, match="timed out"): + await services[0].deps.model_invoker.invoke(request(timeout_seconds=500)) + assert not cancelled.is_set() + release.set() + await asyncio.wait_for(cancelled.wait(), 1) + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_input_and_output_limits(host): + loaded, services, _ = await host.start([{**GRANT, "max_input_chars": 5, "max_output_chars": 2}]) + invoker = services[0].deps.model_invoker + with pytest.raises(ModelInvocationFailed, match="input"): + await invoker.invoke(request()) + host.factory.assert_not_called() + with pytest.raises(ModelInvocationFailed, match="output"): + await invoker.invoke(ModelInvocationRequest([ModelMessage("user", "hi")])) + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_text_blocks_project_without_provider_metadata(host): + host.model.ainvoke.return_value = AIMessage(content=[{"type": "text", "text": "hello"}, {"type": "text", "text": " world"}], response_metadata={"secret": "hidden"}) + loaded, services, _ = await host.start([GRANT]) + result = await services[0].deps.model_invoker.invoke(request()) + assert result.content == "hello world" + assert result.usage is None + assert "hidden" not in repr(result) + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_tool_calls_rejected(host): + host.model.ainvoke.return_value = AIMessage(content="", tool_calls=[{"name": "search", "args": {}, "id": "call"}]) + loaded, services, _ = await host.start([GRANT]) + with pytest.raises(ModelInvocationFailed, match="Tool-call"): + await services[0].deps.model_invoker.invoke(request()) + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_failed_duplicate_install_does_not_grant_prior_service(monkeypatch): + good = Service() + failed = Service() + + def install(registry, config): + registry.service(failed if config else good) + if config: + raise RuntimeError("install failed after registration") + + monkeypatch.setattr("deerflow.extensions.loader.resolve_variable", lambda _: install) + loaded, diagnostics = load_extensions( + [ + ExtensionSpec(use="same:install"), + ExtensionSpec(use="same:install", config={"fail": True}, host_access={"model_invocation": GRANT}), + ] + ) + assert len(diagnostics) == 1 + assert len(loaded.services) == 1 + await start_services(loaded, SimpleNamespace(), None) + assert good.deps.model_invoker is None + assert not hasattr(failed, "deps") + await stop_services(loaded) + + +@pytest.mark.parametrize("grant", [{"roles": {}}, {"roles": {"default": " "}}, {"roles": {"": "model"}}, {**GRANT, "max_concurrency": 0}, {**GRANT, "timeout_seconds": float("inf")}]) +def test_invalid_grants_rejected(grant): + with pytest.raises(ValueError): + ExtensionSpec(use="example:install", host_access={"model_invocation": grant}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("exit_kind", ["timeout", "cancel", "stop"]) +async def test_sync_provider_retains_slot_until_thread_finishes(host, exit_kind): + import threading + + from langchain_core.language_models.chat_models import BaseChatModel + from langchain_core.outputs import ChatGeneration, ChatResult + + started = asyncio.Event() + release = threading.Event() + loop = asyncio.get_running_loop() + calls = [] + + class SyncModel(BaseChatModel): + @property + def _llm_type(self): + return "blocked-sync-test" + + def _generate(self, messages, stop=None, run_manager=None, **kwargs): + calls.append(1) + loop.call_soon_threadsafe(started.set) + assert release.wait(10) + return ChatResult(generations=[ChatGeneration(message=AIMessage(content="ok"))]) + + host.factory.return_value = SyncModel() + loaded, services, _ = await host.start([{**GRANT, "max_concurrency": 1}], services_per_install=2) + invoker = services[0].deps.model_invoker + first = asyncio.create_task(invoker.invoke(request(timeout_seconds=0.1 if exit_kind == "timeout" else 5))) + try: + await asyncio.wait_for(started.wait(), 3) + if exit_kind == "cancel": + first.cancel() + elif exit_kind == "stop": + invoker.close() + with pytest.raises(ModelInvocationFailed if exit_kind == "timeout" else asyncio.CancelledError): + await first + with pytest.raises(ModelInvocationFailed, match="timed out"): + await services[1].deps.model_invoker.invoke(request(timeout_seconds=0.05)) + assert len(calls) == 1, "timed-out synchronous requests still consume their concurrency slot" + finally: + release.set() + await asyncio.gather(first, return_exceptions=True) + assert (await services[1].deps.model_invoker.invoke(request())).content == "ok" + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_admission_limit_shared_and_rejects_before_payload_processing(host): + started = asyncio.Event() + release = asyncio.Event() + + async def invoke(*args, **kwargs): + started.set() + await release.wait() + return AIMessage(content="ok") + + host.model.ainvoke.side_effect = invoke + loaded, services, _ = await host.start([{**GRANT, "max_concurrency": 1}], services_per_install=2) + first = asyncio.create_task(services[0].deps.model_invoker.invoke(request())) + await started.wait() + queued = asyncio.create_task(services[1].deps.model_invoker.invoke(request())) + await asyncio.sleep(0) + try: + with pytest.raises(ModelInvocationFailed, match="capacity"): + await services[0].deps.model_invoker.invoke(request(timeout_seconds=0.05)) + assert host.factory.call_count == 1 + finally: + release.set() + await asyncio.gather(first, queued) + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_provider_timeout_is_not_reported_as_host_deadline(host): + host.model.ainvoke.side_effect = TimeoutError("secret provider URL") + loaded, services, _ = await host.start([GRANT]) + with pytest.raises(ModelInvocationFailed, match="provider timed out") as error: + await services[0].deps.model_invoker.invoke(request()) + assert error.value.__context__ is None + assert "secret" not in str(error.value) + await stop_services(loaded) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("self_cancel", [False, True]) +@pytest.mark.parametrize("prior_cancellation", [False, True]) +async def test_provider_cancellation_is_normalized_and_releases_capacity(host, self_cancel, prior_cancellation): + async def cancelled_provider(*args, **kwargs): + if self_cancel: + asyncio.current_task().cancel("secret provider detail") + await asyncio.sleep(0) + raise asyncio.CancelledError("secret provider detail") + + host.model.ainvoke.side_effect = cancelled_provider + loaded, services, _ = await host.start([{**GRANT, "max_concurrency": 1}]) + invoker = services[0].deps.model_invoker + + async def caller(): + if prior_cancellation: + asyncio.current_task().cancel() + try: + await asyncio.sleep(0) + except asyncio.CancelledError: + pass # A previously handled request must not mask provider failure. + with pytest.raises(ModelInvocationFailed, match="provider cancelled") as error: + await invoker.invoke(request()) + assert error.value.__cause__ is None + assert error.value.__context__ is None + assert "secret" not in str(error.value) + assert invoker._budget.admitted == 0 + host.model.ainvoke.side_effect = None + assert (await invoker.invoke(request())).content == '{"label":"positive"}' + + try: + await asyncio.create_task(caller()) + finally: + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_caller_cancellation_wins_when_provider_also_cancels(host): + async def cancelled_provider(*args, **kwargs): + caller.cancel() + raise asyncio.CancelledError("provider cancelled too") + + host.model.ainvoke.side_effect = cancelled_provider + loaded, services, _ = await host.start([GRANT]) + invoker = services[0].deps.model_invoker + caller = asyncio.create_task(invoker.invoke(request())) + try: + with pytest.raises(asyncio.CancelledError): + await caller + assert caller.cancelled() + assert invoker._budget.admitted == 0 + finally: + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_pending_caller_cancellation_at_invocation_entry_propagates(host): + loaded, services, _ = await host.start([GRANT]) + invoker = services[0].deps.model_invoker + + async def caller(): + asyncio.current_task().cancel() + await invoker.invoke(request()) + + try: + with pytest.raises(asyncio.CancelledError): + await asyncio.create_task(caller()) + finally: + await asyncio.gather(*invoker._budget.workers, return_exceptions=True) + await stop_services(loaded) + assert invoker._budget.admitted == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("exit_kind", ["timeout", "cancel", "stop"]) +async def test_validation_deadline_does_not_block_event_loop(host, monkeypatch, exit_kind): + import time + + from deerflow.extensions import model_invocation + + processes = [] + popen = model_invocation.subprocess.Popen + + def track_process(*args, **kwargs): + process = popen(*args, **kwargs) + processes.append(process) + return process + + monkeypatch.setattr(model_invocation.subprocess, "Popen", track_process) + schema = {"type": "object", "properties": {"value": {"type": "string", "pattern": "^(a+)+$"}}} + host.model.ainvoke.return_value = AIMessage(content='{"value":"' + "a" * 30 + '!"}') + loaded, services, _ = await host.start([GRANT]) + ticks = [] + running = True + + async def heartbeat(): + while running: + ticks.append(time.monotonic()) + await asyncio.sleep(0.01) + + ticker = asyncio.create_task(heartbeat()) + invoker = services[0].deps.model_invoker + task = asyncio.create_task(invoker.invoke(request(response_schema=schema, timeout_seconds=3))) + try: + if exit_kind != "timeout": + async with asyncio.timeout(3): + while len(processes) < 2: + await asyncio.sleep(0.01) + if exit_kind == "cancel": + task.cancel() + # Repeated cancellation must not abandon the child cleanup. + asyncio.get_running_loop().call_later(0.01, task.cancel) + else: + invoker.close() + with pytest.raises(ModelInvocationFailed if exit_kind == "timeout" else asyncio.CancelledError): + await task + assert host.model.ainvoke.call_count == 1, "deadline must fire during response validation" + ticks.append(time.monotonic()) + assert max(b - a for a, b in zip(ticks, ticks[1:])) < 0.3 + assert len(processes) == 2 + assert all(process.poll() is not None for process in processes) + assert invoker._budget.admitted == 0 + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + running = False + await ticker + await stop_services(loaded) + + +@pytest.mark.asyncio +async def test_cancelled_constructor_retains_slot_and_never_dispatches(host): + import threading + + started = asyncio.Event() + release = threading.Event() + loop = asyncio.get_running_loop() + + def factory(*args, **kwargs): + loop.call_soon_threadsafe(started.set) + assert release.wait(10) + return host.model + + host.factory.side_effect = factory + loaded, services, _ = await host.start([{**GRANT, "max_concurrency": 1}]) + invoker = services[0].deps.model_invoker + first = asyncio.create_task(invoker.invoke(request())) + try: + await asyncio.wait_for(started.wait(), 3) + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + with pytest.raises(ModelInvocationFailed, match="timed out"): + await invoker.invoke(request(timeout_seconds=0.05)) + assert host.factory.call_count == 1 + finally: + release.set() + await asyncio.gather(*invoker._budget.workers, return_exceptions=True) + await stop_services(loaded) + host.model.ainvoke.assert_not_called() diff --git a/config.example.yaml b/config.example.yaml index a9238740a3..7a5b2d9b36 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -3222,6 +3222,14 @@ authorization: # enabled: true # false skips import and registration # required: false # true makes load failure abort Gateway startup # # (install --required opts in) +# host_access: # optional host-owned grants (restart required) +# model_invocation: +# roles: # logical role -> configured models[].name +# default: my-configured-model +# max_concurrency: 2 # shared by this installation's services +# timeout_seconds: 60 # includes waiting for capacity +# max_input_chars: 262144 # includes the response-schema instruction +# max_output_chars: 65536 # returned text, checked after generation # table_prefix: example_ext_ # optional: table-name prefix this extension # # owns under its own MetaData/migration chain. # # Keeps `alembic revision --autogenerate`, run