Compare commits

...
Author SHA1 Message Date
L1nSn0w dd7be50645 docs(llm): trim first-token timeout comments to the essentials
Keep only the non-obvious constraints (unit contract, daemon no-keep-alive
assumption, ContextVar fail-open threading, replace-not-narrow intent) and
cut prose that restates the code. The module docstring in
first_token_timeout.py stays the single canonical description of the path.
2026-07-17 11:18:31 +08:00
L1nSn0w 8f1504bf95 chore(llm): improve timeout observability and pin the graphon transform contract
Name the configured window in the mid-stream stall message so a read
timeout after the first token is traceable to first_token_timeout_ms; log
invalid values ignored at the pop point; state honestly in _read_timeout_for
that the budget replaces (not narrows) the operator read timeout; pin with a
test that FirstTokenTimeoutError passes graphon's _transform_invoke_error
unchanged, and cover the converter's defensive drop.
2026-07-17 11:18:31 +08:00
L1nSn0w 450d76ada5 fix(web): render first-token timeout only for panels that consume it
isInWorkflow was a proxy for 'this panel's model config reaches
fetch_model_config', which is false for the knowledge-retrieval metadata
filter and single-retrieval model, and for tool/trigger model-selector
params rendered through form-input-item — the field showed up there but the
backend never consumed it. Gate the rule on an explicit
supportFirstTokenTimeout prop passed only by the LLM, question classifier
and parameter extractor panels, and extend the help text with the
inter-token semantics of the read window.
2026-07-17 11:18:31 +08:00
L1nSn0w 20a205488b chore(web): lower first-token timeout default to 2000ms 2026-07-17 11:18:31 +08:00
L1nSn0w d865d3d703 refactor(llm): switch first-token timeout config to milliseconds
Rename the completion_params key to first_token_timeout_ms so the unit is
explicit in the DSL, and lower the default from 60s to 10s with a 100ms-10min
range: users who enable this gate are latency-sensitive, and 60s rarely
matches that intent. The ms->s conversion happens exactly once, at the pop
point in _normalize_completion_params; everything downstream (ModelInstance
field, ContextVar, httpx read timeout) stays in seconds.
2026-07-17 11:18:31 +08:00
L1nSn0w 989c42153c feat(web): add first-token timeout to the workflow model parameter panel
Append a synthetic FIRST_TOKEN_TIMEOUT_PARAMETER_RULE (int, seconds, opt-in)
to the model parameter panel, mirroring how the stop rule is injected. The
rule renders only in workflow advanced mode (isInWorkflow), covering the LLM,
question classifier and parameter extractor panels; easy-UI app orchestration
pages never show it. The value is stored in completion_params and consumed by
the Dify backend before the request reaches the provider.
2026-07-17 11:18:31 +08:00
L1nSn0w de11f571b8 feat(llm): source first-token timeout from completion_params
The timeout travels the same channel as stop: configured per model in
completion_params.first_token_timeout (seconds), popped at the workflow
model-config boundary (_normalize_completion_params) into
ModelInstance.first_token_timeout, and read by DifyPreparedLLM to arm the
transport ContextVar around invoke_llm and invoke_llm_with_structured_output.
Invalid values (non-numeric, bool, non-positive) disable the gate.

Covers workflow and chatflow LLM-compatible nodes (llm, question classifier,
parameter extractor). The easy-UI model config converter drops the key
defensively so imported app configs never forward it to providers. No graphon
changes are required: the popped key never reaches graphon, and its retry
handler treats FirstTokenTimeoutError like any node failure.
2026-07-17 11:18:31 +08:00
L1nSn0w d77df0b0a4 feat(llm): enforce first-token timeout in the plugin transport
The timeout is carried to the plugin-daemon transport through a ContextVar
(core.plugin.impl.first_token_timeout) and applied as httpx's per-request
read timeout in BasePluginClient._stream_request. The daemon withholds the
response headers until the model's first token, so the read timeout measures
time-to-first-token directly; a ReadTimeout before the first line surfaces as
FirstTokenTimeoutError (dify-local, subclassing graphon's InvokeError), while
a later stall stays a plain transport error. A non-positive budget disables
the gate.
2026-07-17 11:18:31 +08:00
16 changed files with 656 additions and 43 deletions
@@ -62,6 +62,8 @@ class ModelConfigConverter:
if "stop" in completion_params:
stop = completion_params["stop"]
del completion_params["stop"]
# Workflow-only setting; never forward it to providers.
completion_params.pop("first_token_timeout_ms", None)
model_schema = model_type_instance.get_model_schema(model_config.model, model_credentials)
+25 -7
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import logging
from copy import deepcopy
from typing import Any
@@ -14,6 +15,8 @@ from graphon.nodes.llm.entities import ModelConfig
from graphon.nodes.llm.exc import LLMModeRequiredError, ModelNotExistError
from graphon.nodes.llm.protocols import CredentialsProvider
logger = logging.getLogger(__name__)
class DifyCredentialsProvider:
"""Resolves and returns LLM credentials for a given provider and model.
@@ -128,21 +131,35 @@ def build_dify_model_access(run_context: DifyRunContext) -> tuple[CredentialsPro
)
def _normalize_completion_params(completion_params: dict[str, Any]) -> tuple[dict[str, Any], list[str]]:
def _normalize_completion_params(
completion_params: dict[str, Any],
) -> tuple[dict[str, Any], list[str], float | None]:
"""
Split node-level completion params into provider parameters and stop sequences.
Split node-level completion params into provider parameters, stop sequences,
and the first-token timeout.
Workflow LLM-compatible nodes still consume runtime invocation settings from
``ModelInstance.parameters`` and ``ModelInstance.stop``. Keep the
``ModelInstance`` view and the returned config entity aligned here so callers
do not need to duplicate normalization logic.
``ModelInstance.parameters``, ``ModelInstance.stop`` and
``ModelInstance.first_token_timeout``. Keep the ``ModelInstance`` view and the
returned config entity aligned here so callers do not need to duplicate
normalization logic.
``first_token_timeout_ms`` never reaches providers; this is the only ms->s
conversion on the path. Invalid values disable the gate.
"""
normalized_parameters = dict(completion_params)
stop = normalized_parameters.pop("stop", [])
if not isinstance(stop, list) or not all(isinstance(item, str) for item in stop):
stop = []
return normalized_parameters, stop
raw_timeout_ms = normalized_parameters.pop("first_token_timeout_ms", None)
first_token_timeout: float | None = None
if isinstance(raw_timeout_ms, (int, float)) and not isinstance(raw_timeout_ms, bool) and raw_timeout_ms > 0:
first_token_timeout = float(raw_timeout_ms) / 1000
elif raw_timeout_ms is not None:
logger.debug("Ignoring invalid first_token_timeout_ms in completion_params: %r", raw_timeout_ms)
return normalized_parameters, stop, first_token_timeout
def fetch_model_config(
@@ -178,12 +195,13 @@ def fetch_model_config(
if model_schema is None:
raise ModelNotExistError(f"Model {node_data_model.name} schema does not exist.")
parameters, stop = _normalize_completion_params(node_data_model.completion_params)
parameters, stop, first_token_timeout = _normalize_completion_params(node_data_model.completion_params)
model_instance.provider = node_data_model.provider
model_instance.model_name = node_data_model.name
model_instance.credentials = credentials
model_instance.parameters = parameters
model_instance.stop = tuple(stop)
model_instance.first_token_timeout = first_token_timeout
return model_instance, ModelConfigWithCredentialsEntity(
provider=node_data_model.provider,
+1
View File
@@ -47,6 +47,7 @@ class ModelInstance:
# Runtime LLM invocation fields.
self.parameters: Mapping[str, Any] = {}
self.stop: Sequence[str] = ()
self.first_token_timeout: float | None = None
self.model_type_instance = self.provider_model_bundle.model_type_instance
self.load_balancing_manager = self._get_load_balancing_manager(
configuration=provider_model_bundle.configuration,
+39 -3
View File
@@ -25,6 +25,7 @@ from core.plugin.impl.exc import (
PluginPermissionDeniedError,
PluginUniqueIdentifierError,
)
from core.plugin.impl.first_token_timeout import FirstTokenTimeoutError, first_token_timeout_ctx
from core.trigger.errors import (
EventIgnoreError,
TriggerInvokeError,
@@ -54,6 +55,22 @@ match _plugin_daemon_timeout_config:
case _:
plugin_daemon_request_timeout = httpx.Timeout(_plugin_daemon_timeout_config)
def _read_timeout_for(first_token_timeout: float | None) -> httpx.Timeout | None:
"""Replace the daemon request timeout's ``read`` component with the first-token budget.
Deliberately a replacement rather than a narrowing, so the budget may exceed
``PLUGIN_DAEMON_TIMEOUT`` for slow reasoning models. ``httpx.Timeout(base, read=x)``
rejects a ``Timeout`` base, so the other components are copied explicitly.
"""
base = plugin_daemon_request_timeout
if not first_token_timeout or first_token_timeout <= 0:
return base
if base is None:
return httpx.Timeout(None, read=first_token_timeout)
return httpx.Timeout(connect=base.connect, read=first_token_timeout, write=base.write, pool=base.pool)
logger = logging.getLogger(__name__)
PLUGIN_DAEMON_MAX_PATH_LENGTH = 4096
@@ -181,30 +198,49 @@ class BasePluginClient:
"""
url, headers, prepared_data, params, files = self._prepare_request(path, headers, data, params, files)
first_token_timeout = first_token_timeout_ctx.get()
first_token_gate = bool(first_token_timeout and first_token_timeout > 0)
stream_kwargs: dict[str, Any] = {
"method": method,
"url": url,
"headers": headers,
"params": params,
"files": files,
"timeout": plugin_daemon_request_timeout,
# The daemon sends nothing before the first token, so the read timeout gates TTFT.
"timeout": _read_timeout_for(first_token_timeout),
}
if isinstance(prepared_data, dict):
stream_kwargs["data"] = prepared_data
elif prepared_data is not None:
stream_kwargs["content"] = prepared_data
first_token_seen = False
try:
with _httpx_client.stream(**stream_kwargs) as response:
for raw_line in response.iter_lines():
# Blank frames don't count as the first token, yet each read refreshes the
# read window -- gating relies on no keep-alives before the first token.
if not raw_line:
continue
line = raw_line.decode("utf-8") if isinstance(raw_line, bytes) else raw_line
line = line.strip()
if line.startswith("data:"):
line = line[5:].strip()
if line:
yield line
if not line:
continue
first_token_seen = True
yield line
except httpx.ReadTimeout as e:
if first_token_gate and not first_token_seen:
raise FirstTokenTimeoutError(f"The first token was not received within {first_token_timeout}s.") from e
logger.exception("Stream request to Plugin Daemon Service failed")
message = "Request to Plugin Daemon Service failed"
if first_token_gate:
# An inter-token stall past the read window is a plain transport error,
# but name the window so it stays traceable to the user's setting.
message += f" (stream stalled beyond the {first_token_timeout}s first-token timeout window)"
raise PluginDaemonInnerError(code=-500, message=message)
except httpx.RequestError:
logger.exception("Stream request to Plugin Daemon Service failed")
raise PluginDaemonInnerError(code=-500, message="Request to Plugin Daemon Service failed")
@@ -0,0 +1,27 @@
"""First-token timeout plumbing for LLM streaming through the plugin daemon.
Configured per model as ``completion_params.first_token_timeout_ms`` and popped into
``ModelInstance.first_token_timeout`` by ``_normalize_completion_params`` -- the only
ms->s conversion; everything below the pop point is seconds. ``DifyPreparedLLM`` sets
the ContextVar around invocation, and ``BasePluginClient._stream_request`` applies it
as the per-request httpx ``read`` timeout. The daemon sends neither response headers
nor keep-alives before the first token, so the read timeout measures
time-to-first-token directly.
The ContextVar only reaches the transport on the thread that set it; a
background-thread stream consumer would read ``None`` and the gate fails open.
"""
from contextvars import ContextVar
from graphon.model_runtime.errors.invoke import InvokeError
class FirstTokenTimeoutError(InvokeError):
"""The model did not stream its first token within the configured budget."""
description = "The first streamed token was not received in time."
# Seconds; None or non-positive disables the gate.
first_token_timeout_ctx: ContextVar[float | None] = ContextVar("first_token_timeout", default=None)
+59 -16
View File
@@ -22,6 +22,7 @@ from core.llm_generator.output_parser.errors import OutputParserError
from core.llm_generator.output_parser.structured_output import invoke_llm_with_structured_output
from core.model_manager import ModelInstance
from core.plugin.impl.exc import PluginDaemonClientSideError, PluginInvokeError
from core.plugin.impl.first_token_timeout import first_token_timeout_ctx
from core.plugin.impl.plugin import PluginInstaller
from core.prompt.utils.prompt_message_util import PromptMessageUtil
from core.repositories.human_input_repository import (
@@ -147,6 +148,42 @@ class DifyFileReferenceFactory(FileReferenceFactoryProtocol):
)
def _guarded_stream(
seconds: float,
inner: Generator[Any, None, None],
) -> Generator[Any, None, None]:
"""Iterate ``inner`` with the first-token-timeout ContextVar set.
Keeps the value active for exactly the span of iteration, which is when the
plugin-daemon transport reads it. Same-thread only: a background-thread SSE
prefetch would not see it and the gate fails open.
"""
token = first_token_timeout_ctx.set(seconds)
try:
yield from inner
finally:
first_token_timeout_ctx.reset(token)
def _with_first_token_timeout[T](first_token_timeout: float | None, invoke: Callable[[], T]) -> T:
"""Run ``invoke`` with the first-token-timeout ContextVar applied.
``first_token_timeout`` is seconds, from ``ModelInstance.first_token_timeout``.
A streaming result is lazy, so its generator is wrapped to keep the ContextVar
set during iteration; a non-positive timeout disables the gate.
"""
if first_token_timeout is None or first_token_timeout <= 0:
return invoke()
token = first_token_timeout_ctx.set(first_token_timeout)
try:
result = invoke()
finally:
first_token_timeout_ctx.reset(token)
if isinstance(result, Generator):
return cast("T", _guarded_stream(first_token_timeout, result))
return result
class DifyPreparedLLM(LLMProtocol):
"""Workflow-layer adapter that hides the full `ModelInstance` API from `graphon` nodes."""
@@ -225,13 +262,16 @@ class DifyPreparedLLM(LLMProtocol):
stop: Sequence[str] | None,
stream: bool,
) -> LLMResult | Generator[LLMResultChunk, None, None]:
return self._model_instance.invoke_llm(
prompt_messages=list(prompt_messages),
model_parameters=dict(model_parameters),
tools=list(tools or []),
stop=list(stop or []),
stream=stream,
request_metadata=self._request_metadata,
return _with_first_token_timeout(
self._model_instance.first_token_timeout,
lambda: self._model_instance.invoke_llm(
prompt_messages=list(prompt_messages),
model_parameters=dict(model_parameters),
tools=list(tools or []),
stop=list(stop or []),
stream=stream,
request_metadata=self._request_metadata,
),
)
@overload
@@ -266,15 +306,18 @@ class DifyPreparedLLM(LLMProtocol):
stop: Sequence[str] | None,
stream: bool,
) -> LLMResultWithStructuredOutput | Generator[LLMResultChunkWithStructuredOutput, None, None]:
return invoke_llm_with_structured_output(
provider=self.provider,
model_schema=self.get_model_schema(),
model_instance=self._model_instance,
prompt_messages=prompt_messages,
json_schema=json_schema,
model_parameters=model_parameters,
stop=list(stop or []),
stream=stream,
return _with_first_token_timeout(
self._model_instance.first_token_timeout,
lambda: invoke_llm_with_structured_output(
provider=self.provider,
model_schema=self.get_model_schema(),
model_instance=self._model_instance,
prompt_messages=prompt_messages,
json_schema=json_schema,
model_parameters=model_parameters,
stop=list(stop or []),
stream=stream,
),
)
@override
@@ -100,6 +100,14 @@ class TestModelConfigConverter:
assert result.parameters == {"temperature": 0.7}
assert result.stop == ["\n"]
def test_convert_drops_first_token_timeout_ms(self, mock_app_config, patch_provider_manager):
"""The workflow-only setting must never reach providers through app configs."""
mock_app_config.model.parameters = {"temperature": 0.7, "first_token_timeout_ms": 2000}
result = ModelConfigConverter.convert(mock_app_config)
assert result.parameters == {"temperature": 0.7}
def test_convert_mode_from_schema_valid(self, mock_app_config, mock_provider_bundle, mocker: MockerFixture):
mock_app_config.model.mode = None
@@ -0,0 +1,215 @@
import httpx
import pytest
from pytest_mock import MockerFixture
from core.plugin.entities.plugin_daemon import PluginDaemonInnerError
from core.plugin.impl import base as base_mod
from core.plugin.impl.first_token_timeout import FirstTokenTimeoutError, first_token_timeout_ctx
BasePluginClient = base_mod.BasePluginClient
@pytest.fixture(autouse=True)
def _isolate_ctx():
"""Keep the first-token-timeout ContextVar from leaking between tests."""
token = first_token_timeout_ctx.set(None)
try:
yield
finally:
first_token_timeout_ctx.reset(token)
class _PlainStream:
"""Fully buffered fake httpx stream context."""
def __init__(self, lines: list[object]) -> None:
self._lines = lines
def __enter__(self) -> "_PlainStream":
return self
def __exit__(self, *exc: object) -> bool:
return False
def iter_lines(self):
return iter(self._lines)
class _RaiseOnEnterStream:
"""Fake stream whose context entry raises — models a timeout while awaiting headers."""
def __init__(self, exc: BaseException) -> None:
self._exc = exc
def __enter__(self) -> "_RaiseOnEnterStream":
raise self._exc
def __exit__(self, *exc: object) -> bool:
return False
class _LinesThenRaiseStream:
"""Yields some lines, then raises — models a stall after the first token(s)."""
def __init__(self, lines: list[object], exc: BaseException) -> None:
self._lines = lines
self._exc = exc
def __enter__(self) -> "_LinesThenRaiseStream":
return self
def __exit__(self, *exc: object) -> bool:
return False
def iter_lines(self):
yield from self._lines
raise self._exc
# --- _read_timeout_for ------------------------------------------------------------------
def test_read_timeout_for_narrows_read_only() -> None:
base = base_mod.plugin_daemon_request_timeout
timeout = base_mod._read_timeout_for(5.0)
assert timeout is not base
assert timeout.read == 5.0
# The other components are preserved from the default timeout.
assert timeout.connect == base.connect
assert timeout.write == base.write
assert timeout.pool == base.pool
@pytest.mark.parametrize("disabled", [None, 0.0, -1.0])
def test_read_timeout_for_disabled_returns_base_unchanged(disabled: float | None) -> None:
assert base_mod._read_timeout_for(disabled) is base_mod.plugin_daemon_request_timeout
def test_read_timeout_for_with_no_base_timeout(mocker: MockerFixture) -> None:
mocker.patch("core.plugin.impl.base.plugin_daemon_request_timeout", None)
timeout = base_mod._read_timeout_for(5.0)
assert timeout.read == 5.0
assert timeout.connect is None
assert timeout.write is None
assert timeout.pool is None
# --- _stream_request timeout wiring ------------------------------------------------------
def test_stream_request_narrows_read_when_gate_enabled(mocker: MockerFixture) -> None:
client = BasePluginClient()
stream = mocker.patch("httpx.Client.stream", return_value=_PlainStream([b"data: hi"]))
first_token_timeout_ctx.set(1.5)
result = list(client._stream_request("POST", "plugin/tenant/stream", data={"k": "v"}))
assert result == ["hi"]
assert stream.call_args.kwargs["timeout"].read == 1.5
@pytest.mark.parametrize("disabled", [None, 0.0])
def test_stream_request_keeps_default_timeout_when_gate_disabled(mocker: MockerFixture, disabled: float | None) -> None:
client = BasePluginClient()
stream = mocker.patch("httpx.Client.stream", return_value=_PlainStream([b"data: hi"]))
first_token_timeout_ctx.set(disabled)
list(client._stream_request("POST", "plugin/tenant/stream", data={"k": "v"}))
assert stream.call_args.kwargs["timeout"] is base_mod.plugin_daemon_request_timeout
def test_stream_request_forwards_all_lines(mocker: MockerFixture) -> None:
client = BasePluginClient()
mocker.patch("httpx.Client.stream", return_value=_PlainStream([b"", b"data: hello", "world"]))
first_token_timeout_ctx.set(1.0)
result = list(client._stream_request("POST", "plugin/tenant/stream", data={"k": "v"}))
assert result == ["hello", "world"]
# --- _stream_request timeout semantics ---------------------------------------------------
def test_read_timeout_before_first_line_raises_first_token_timeout(mocker: MockerFixture) -> None:
client = BasePluginClient()
mocker.patch("httpx.Client.stream", return_value=_RaiseOnEnterStream(httpx.ReadTimeout("headers")))
first_token_timeout_ctx.set(0.5)
with pytest.raises(FirstTokenTimeoutError):
list(client._stream_request("POST", "plugin/tenant/stream", data={"k": "v"}))
def test_read_timeout_with_gate_disabled_is_transport_error(mocker: MockerFixture) -> None:
client = BasePluginClient()
mocker.patch("httpx.Client.stream", return_value=_RaiseOnEnterStream(httpx.ReadTimeout("headers")))
# ctx is None (autouse fixture) -> gate off -> a read timeout is just a transport error.
with pytest.raises(PluginDaemonInnerError):
list(client._stream_request("POST", "plugin/tenant/stream", data={"k": "v"}))
def test_read_timeout_after_first_line_is_transport_error(mocker: MockerFixture) -> None:
client = BasePluginClient()
mocker.patch(
"httpx.Client.stream",
return_value=_LinesThenRaiseStream([b"data: hello"], httpx.ReadTimeout("inter-token")),
)
first_token_timeout_ctx.set(0.5)
# First token already seen -> a later read timeout is an inter-token stall, not a
# first-token timeout.
with pytest.raises(PluginDaemonInnerError) as exc_info:
list(client._stream_request("POST", "plugin/tenant/stream", data={"k": "v"}))
assert "0.5s first-token timeout window" in exc_info.value.message
def test_read_timeout_after_first_line_with_gate_off_keeps_plain_message(mocker: MockerFixture) -> None:
client = BasePluginClient()
mocker.patch(
"httpx.Client.stream",
return_value=_LinesThenRaiseStream([b"data: hello"], httpx.ReadTimeout("inter-token")),
)
# ctx is None (autouse fixture) -> gate off -> no window hint in the message.
with pytest.raises(PluginDaemonInnerError) as exc_info:
list(client._stream_request("POST", "plugin/tenant/stream", data={"k": "v"}))
assert "first-token timeout window" not in exc_info.value.message
def test_non_timeout_request_error_is_transport_error(mocker: MockerFixture) -> None:
client = BasePluginClient()
mocker.patch("httpx.Client.stream", return_value=_RaiseOnEnterStream(httpx.ConnectError("boom")))
first_token_timeout_ctx.set(0.5)
# Only a ReadTimeout maps to FirstTokenTimeoutError; other transport errors do not.
with pytest.raises(PluginDaemonInnerError):
list(client._stream_request("POST", "plugin/tenant/stream", data={"k": "v"}))
# --- graphon error-transform contract -----------------------------------------------------
def test_first_token_timeout_error_survives_graphon_invoke_error_transform() -> None:
"""error_type == "FirstTokenTimeoutError" relies on graphon passing the exception
through ``_transform_invoke_error`` unchanged (``InvokeError`` subclasses
``ValueError``, whose mapping entry returns the original error). A graphon bump
breaking either fact would silently degrade the type — fail loudly here instead.
"""
from graphon.model_runtime.errors.invoke import InvokeError
from graphon.model_runtime.model_providers.base.ai_model import AIModel
class _ProbeModel:
_invoke_error_mapping = AIModel._invoke_error_mapping
provider_display_name = "probe"
assert issubclass(InvokeError, ValueError)
error = FirstTokenTimeoutError("The first token was not received within 1.5s.")
transformed = AIModel._transform_invoke_error(_ProbeModel(), error) # type: ignore[arg-type]
assert transformed is error
@@ -295,6 +295,7 @@ def test_fetch_model_config_hydrates_model_instance_runtime_settings(model_confi
"temperature": 0.7,
"max_tokens": 256,
"stop": ["Observation:", "Human:"],
"first_token_timeout_ms": 30000,
}
model_instance = mock.MagicMock(
@@ -336,6 +337,7 @@ def test_fetch_model_config_hydrates_model_instance_runtime_settings(model_confi
"max_tokens": 256,
}
assert hydrated_model_instance.stop == ("Observation:", "Human:")
assert hydrated_model_instance.first_token_timeout == 30.0
assert model_config_with_credentials.parameters == {
"temperature": 0.7,
"max_tokens": 256,
@@ -345,12 +347,51 @@ def test_fetch_model_config_hydrates_model_instance_runtime_settings(model_confi
"temperature": 0.7,
"max_tokens": 256,
"stop": ["Observation:", "Human:"],
"first_token_timeout_ms": 30000,
}
mock_credentials_provider.fetch.assert_called_once_with("openai", "gpt-3.5-turbo")
mock_model_factory.init_model_instance.assert_called_once_with("openai", "gpt-3.5-turbo")
provider_model.raise_for_status.assert_called_once()
@pytest.mark.parametrize(
("raw_timeout_ms", "expected"),
[
(30000, 30.0),
(1500, 1.5),
(100, 0.1),
(0, None),
(-5, None),
(True, None),
("30000", None),
(None, None),
],
)
def test_normalize_completion_params_extracts_first_token_timeout(raw_timeout_ms: object, expected: float | None):
from core.app.llm.model_access import _normalize_completion_params
completion_params = {"temperature": 0.7, "first_token_timeout_ms": raw_timeout_ms}
parameters, stop, first_token_timeout = _normalize_completion_params(completion_params)
assert first_token_timeout == expected
assert "first_token_timeout_ms" not in parameters
assert parameters == {"temperature": 0.7}
assert stop == []
# The caller's dict is left untouched.
assert completion_params == {"temperature": 0.7, "first_token_timeout_ms": raw_timeout_ms}
def test_normalize_completion_params_without_first_token_timeout_key():
from core.app.llm.model_access import _normalize_completion_params
parameters, stop, first_token_timeout = _normalize_completion_params({"temperature": 0.7, "stop": ["Human:"]})
assert first_token_timeout is None
assert parameters == {"temperature": 0.7}
assert stop == ["Human:"]
def test_fetch_model_config_reuses_validated_provider_model_from_dify_credentials_provider(
model_config: ModelConfigWithCredentialsEntity,
):
@@ -8,6 +8,7 @@ from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunC
from core.app.file_access import FileAccessScope, bind_file_access_scope, grant_retriever_segment_access
from core.llm_generator.output_parser.errors import OutputParserError
from core.plugin.impl.exc import PluginLLMPollingUnsupportedError
from core.plugin.impl.first_token_timeout import first_token_timeout_ctx
from core.plugin.impl.model import PluginModelClient
from core.plugin.impl.model_runtime import PluginModelRuntime
from core.plugin.plugin_service import PluginService
@@ -82,6 +83,7 @@ class _ModelInstanceStub:
self.model_name = "gpt-4o-mini"
self.parameters = {"temperature": 0.2}
self.stop = ("stop",)
self.first_token_timeout: float | None = None
self.credentials = {"api_key": "secret"}
self.model_type_instance = _ModelTypeInstanceStub(
model_schema=model_schema,
@@ -201,6 +203,164 @@ def test_dify_prepared_llm_wraps_model_instance_calls() -> None:
)
def test_dify_prepared_llm_sets_first_token_timeout_ctx_during_streaming() -> None:
observed: list[float | None] = []
def _fake_invoke(**_kwargs: object):
def _gen():
observed.append(first_token_timeout_ctx.get())
yield "chunk-1"
observed.append(first_token_timeout_ctx.get())
yield "chunk-2"
return _gen()
model_instance = _ModelInstanceStub(model_schema=_build_model_schema())
model_instance.invoke_llm = _fake_invoke # type: ignore[assignment]
model_instance.first_token_timeout = 1.5
prepared = DifyPreparedLLM(model_instance)
generator = prepared.invoke_llm(
prompt_messages=[],
model_parameters={},
tools=None,
stop=None,
stream=True,
)
# The ContextVar is set inside the generator (during iteration), not on this frame.
assert first_token_timeout_ctx.get() is None
result = list(generator)
assert result == ["chunk-1", "chunk-2"]
assert observed == [1.5, 1.5]
# Reset once the generator is exhausted.
assert first_token_timeout_ctx.get() is None
def test_dify_prepared_llm_leaves_ctx_unset_when_no_first_token_timeout() -> None:
observed: list[float | None] = []
def _fake_invoke(**_kwargs: object):
def _gen():
observed.append(first_token_timeout_ctx.get())
yield "chunk"
return _gen()
model_instance = _ModelInstanceStub(model_schema=_build_model_schema())
model_instance.invoke_llm = _fake_invoke # type: ignore[assignment]
prepared = DifyPreparedLLM(model_instance)
result = list(
prepared.invoke_llm(
prompt_messages=[],
model_parameters={},
tools=None,
stop=None,
stream=True,
)
)
assert result == ["chunk"]
assert observed == [None]
def test_dify_prepared_llm_sets_first_token_timeout_ctx_for_non_stream_eager_path() -> None:
observed: list[float | None] = []
def _fake_invoke(**_kwargs: object):
# A non-stream call resolves eagerly; the ContextVar must be set during this call.
observed.append(first_token_timeout_ctx.get())
return sentinel.result
model_instance = _ModelInstanceStub(model_schema=_build_model_schema())
model_instance.invoke_llm = _fake_invoke # type: ignore[assignment]
model_instance.first_token_timeout = 1.5
prepared = DifyPreparedLLM(model_instance)
result = prepared.invoke_llm(
prompt_messages=[],
model_parameters={},
tools=None,
stop=None,
stream=False,
)
# The non-generator result is returned unchanged, and the gate is reset once it returns.
assert result is sentinel.result
assert observed == [1.5]
assert first_token_timeout_ctx.get() is None
def test_dify_prepared_llm_propagates_first_token_timeout_ctx_through_structured_output(
monkeypatch: pytest.MonkeyPatch,
) -> None:
observed: list[float | None] = []
def _fake_structured(**_kwargs: object):
def _gen():
observed.append(first_token_timeout_ctx.get())
yield "chunk-1"
observed.append(first_token_timeout_ctx.get())
yield "chunk-2"
return _gen()
monkeypatch.setattr(node_runtime, "invoke_llm_with_structured_output", _fake_structured)
model_instance = _ModelInstanceStub(model_schema=_build_model_schema())
model_instance.first_token_timeout = 1.5
prepared = DifyPreparedLLM(model_instance)
generator = prepared.invoke_llm_with_structured_output(
prompt_messages=[],
json_schema={"type": "object"},
model_parameters={},
stop=None,
stream=True,
)
# Lazy: the ContextVar is set during iteration, not on this frame.
assert first_token_timeout_ctx.get() is None
result = list(generator)
assert result == ["chunk-1", "chunk-2"]
assert observed == [1.5, 1.5]
assert first_token_timeout_ctx.get() is None
def test_dify_prepared_llm_resets_first_token_timeout_ctx_when_stream_raises() -> None:
observed: list[float | None] = []
def _fake_invoke(**_kwargs: object):
def _gen():
observed.append(first_token_timeout_ctx.get())
yield "chunk-1"
raise RuntimeError("boom")
return _gen()
model_instance = _ModelInstanceStub(model_schema=_build_model_schema())
model_instance.invoke_llm = _fake_invoke # type: ignore[assignment]
model_instance.first_token_timeout = 1.5
prepared = DifyPreparedLLM(model_instance)
generator = prepared.invoke_llm(
prompt_messages=[],
model_parameters={},
tools=None,
stop=None,
stream=True,
)
with pytest.raises(RuntimeError, match="boom"):
list(generator)
# The try/finally in _guarded_stream resets the ContextVar even when iteration errors out.
assert observed == [1.5]
assert first_token_timeout_ctx.get() is None
def test_dify_prepared_llm_requires_model_schema() -> None:
model_instance = _ModelInstanceStub(model_schema=None)
model_instance.credentials = {}
@@ -388,6 +388,32 @@ describe('ModelParameterModal', () => {
expect(screen.getByText(/debugAsSingleModel/i)).toBeInTheDocument()
})
it('should append the first token timeout parameter when the panel supports it', () => {
render(
<ModelParameterModal
{...defaultProps}
isAdvancedMode
isInWorkflow
supportFirstTokenTimeout
/>,
)
openSettings()
expect(screen.getByTestId('param-first_token_timeout_ms')).toBeInTheDocument()
})
it('should not append the first token timeout parameter for workflow panels that do not support it', () => {
// Workflow context alone must not render it — that would be dead config on
// panels the backend does not consume.
render(<ModelParameterModal {...defaultProps} isAdvancedMode isInWorkflow />)
openSettings()
expect(screen.getByTestId('param-stop')).toBeInTheDocument()
expect(screen.queryByTestId('param-first_token_timeout_ms')).not.toBeInTheDocument()
})
it('should render the empty loading fallback when rules resolve to an empty list', () => {
parameterRules = []
isRulesLoading = true
@@ -9,7 +9,11 @@ import { useMemo, useState } from 'react'
import { useTranslation } from 'react-i18next'
import { ArrowNarrowLeft } from '@/app/components/base/icons/src/vender/line/arrows'
import Loading from '@/app/components/base/loading'
import { PROVIDER_WITH_PRESET_TONE, STOP_PARAMETER_RULE } from '@/config'
import {
FIRST_TOKEN_TIMEOUT_PARAMETER_RULE,
PROVIDER_WITH_PRESET_TONE,
STOP_PARAMETER_RULE,
} from '@/config'
import { useModelParameterRules } from '@/service/use-common'
import { useTextGenerationCurrentProviderAndModelAndModelList } from '../hooks'
import ModelSelector from '../model-selector'
@@ -36,6 +40,9 @@ export type ModelParameterModalProps = {
renderTrigger?: (v: TriggerProps) => ReactNode
readonly?: boolean
isInWorkflow?: boolean
// Pass only from panels whose model config the backend consumes
// (LLM / question classifier / parameter extractor); elsewhere it is dead config.
supportFirstTokenTimeout?: boolean
scope?: string
nodesOutputVars?: NodeOutPutVar[]
availableNodes?: Node[]
@@ -55,6 +62,7 @@ const ModelParameterModal: FC<ModelParameterModalProps> = ({
renderTrigger,
readonly,
isInWorkflow,
supportFirstTokenTimeout,
nodesOutputVars,
availableNodes,
}) => {
@@ -216,22 +224,26 @@ const ModelParameterModal: FC<ModelParameterModalProps> = ({
<Loading />
</div>
) : (
[...parameterRules, ...(isAdvancedMode ? [STOP_PARAMETER_RULE] : [])].map(
(parameter) => (
<ParameterItem
key={`${modelId}-${parameter.name}`}
parameterRule={parameter}
value={completionParams?.[parameter.name]}
onChange={(v) => handleParamChange(parameter.name, v)}
onSwitch={(checked, assignValue) =>
handleSwitch(parameter.name, checked, assignValue)
}
isInWorkflow={isInWorkflow}
nodesOutputVars={nodesOutputVars}
availableNodes={availableNodes}
/>
),
)
[
...parameterRules,
...(isAdvancedMode ? [STOP_PARAMETER_RULE] : []),
...(isAdvancedMode && supportFirstTokenTimeout
? [FIRST_TOKEN_TIMEOUT_PARAMETER_RULE]
: []),
].map((parameter) => (
<ParameterItem
key={`${modelId}-${parameter.name}`}
parameterRule={parameter}
value={completionParams?.[parameter.name]}
onChange={(v) => handleParamChange(parameter.name, v)}
onSwitch={(checked, assignValue) =>
handleSwitch(parameter.name, checked, assignValue)
}
isInWorkflow={isInWorkflow}
nodesOutputVars={nodesOutputVars}
availableNodes={availableNodes}
/>
))
)}
</div>
)}
@@ -117,6 +117,7 @@ const Panel: FC<NodePanelProps<LLMNodeType>> = ({ id, data }) => {
popupClassName="w-[387px]!"
isInWorkflow
isAdvancedMode={true}
supportFirstTokenTimeout
provider={model?.provider}
completionParams={model?.completion_params}
modelId={model?.name}
@@ -61,6 +61,7 @@ const Panel: FC<NodePanelProps<ParameterExtractorNodeType>> = ({ id, data }) =>
popupClassName="w-[387px]!"
isInWorkflow
isAdvancedMode={true}
supportFirstTokenTimeout
provider={model?.provider}
completionParams={model?.completion_params}
modelId={model?.name}
@@ -50,6 +50,7 @@ const Panel: FC<NodePanelProps<QuestionClassifierNodeType>> = ({ id, data }) =>
popupClassName="w-[387px]!"
isInWorkflow
isAdvancedMode={true}
supportFirstTokenTimeout
provider={model?.provider}
completionParams={model.completion_params}
modelId={model.name}
+21
View File
@@ -328,6 +328,27 @@ export const STOP_PARAMETER_RULE: ModelParameterRule = {
},
}
// Consumed by the Dify backend, never sent to model providers.
export const FIRST_TOKEN_TIMEOUT_PARAMETER_RULE: ModelParameterRule = {
default: 2000,
min: 100,
max: 600000,
precision: 0,
help: {
en_US:
'Max milliseconds to wait for the first streamed token; the same window also bounds each gap between tokens. On timeout the node fails; combine with node retry to re-run it automatically.',
zh_Hans:
'等待首个流式 Token 的最长毫秒数;同一窗口也约束后续每次 Token 间隔。超时后节点失败,可配合节点重试自动重跑。',
},
label: {
en_US: 'First token timeout (ms)',
zh_Hans: '首个 Token 超时(毫秒)',
},
name: 'first_token_timeout_ms',
required: false,
type: 'int',
}
export const PARTNER_STACK_CONFIG = {
cookieName: 'partner_stack_info',
saveCookieDays: 90,