mirror of
https://github.com/langgenius/dify.git
synced 2026-06-13 04:01:12 +08:00
347 lines
13 KiB
Python
347 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import time
|
|
from typing import ClassVar
|
|
|
|
import httpx
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from shell_session_manager.shellctl.client import ShellctlClient
|
|
|
|
import dify_agent.server.app as app_module
|
|
from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig
|
|
from dify_agent.layers.execution_context.layer import DifyExecutionContextLayer
|
|
from dify_agent.layers.shell import DifyShellLayerConfig
|
|
from dify_agent.layers.shell.layer import DifyShellLayer
|
|
from dify_agent.runtime.compositor_factory import DifyAgentLayerProvider
|
|
from dify_agent.server.app import create_app, create_plugin_daemon_http_client
|
|
from dify_agent.server.settings import ServerSettings
|
|
from dify_agent.storage.redis_run_store import RedisRunStore
|
|
|
|
|
|
def _base64url_secret(value: bytes) -> str:
|
|
return base64.urlsafe_b64encode(value).rstrip(b"=").decode("ascii")
|
|
|
|
|
|
def _execution_context() -> DifyExecutionContextLayerConfig:
|
|
return DifyExecutionContextLayerConfig(
|
|
tenant_id="tenant-1",
|
|
user_id="user-1",
|
|
user_from="account",
|
|
agent_mode="workflow_run",
|
|
invoke_from="service-api",
|
|
)
|
|
|
|
|
|
def _patch_app_lifecycle(monkeypatch: pytest.MonkeyPatch) -> tuple[FakeRedis, FakePluginDaemonHttpClient]:
|
|
fake_redis = FakeRedis()
|
|
fake_http_client = FakePluginDaemonHttpClient()
|
|
FakeRunScheduler.created.clear()
|
|
FakeRedisModule.fake_redis = fake_redis
|
|
monkeypatch.setattr(app_module, "Redis", FakeRedisModule)
|
|
monkeypatch.setattr(app_module, "RunScheduler", FakeRunScheduler)
|
|
|
|
def fake_create_plugin_daemon_http_client(_settings: ServerSettings) -> FakePluginDaemonHttpClient:
|
|
return fake_http_client
|
|
|
|
monkeypatch.setattr(app_module, "create_plugin_daemon_http_client", fake_create_plugin_daemon_http_client)
|
|
return fake_redis, fake_http_client
|
|
|
|
|
|
class FakeRedis:
|
|
closed: bool
|
|
|
|
def __init__(self) -> None:
|
|
self.closed = False
|
|
|
|
async def aclose(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class FakeRunScheduler:
|
|
created: list["FakeRunScheduler"] = []
|
|
|
|
store: object
|
|
shutdown_grace_seconds: float
|
|
layer_providers: tuple[DifyAgentLayerProvider, ...]
|
|
plugin_daemon_http_client: FakePluginDaemonHttpClient
|
|
shutdown_called: bool
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
store: object,
|
|
plugin_daemon_http_client: FakePluginDaemonHttpClient,
|
|
shutdown_grace_seconds: float,
|
|
layer_providers: tuple[DifyAgentLayerProvider, ...],
|
|
) -> None:
|
|
self.store = store
|
|
self.shutdown_grace_seconds = shutdown_grace_seconds
|
|
self.layer_providers = layer_providers
|
|
self.plugin_daemon_http_client = plugin_daemon_http_client
|
|
self.shutdown_called = False
|
|
self.created.append(self)
|
|
|
|
async def shutdown(self) -> None:
|
|
self.shutdown_called = True
|
|
|
|
|
|
class FakePluginDaemonHttpClient:
|
|
timeout: object | None
|
|
limits: object | None
|
|
trust_env: bool | None
|
|
is_closed: bool
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
timeout: object | None = None,
|
|
limits: object | None = None,
|
|
trust_env: bool | None = None,
|
|
) -> None:
|
|
self.timeout = timeout
|
|
self.limits = limits
|
|
self.trust_env = trust_env
|
|
self.is_closed = False
|
|
|
|
async def aclose(self) -> None:
|
|
self.is_closed = True
|
|
|
|
|
|
class FakeAgentStubGRPCServer:
|
|
closed: bool
|
|
|
|
def __init__(self) -> None:
|
|
self.closed = False
|
|
|
|
async def aclose(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class FakeTimeout:
|
|
connect: float
|
|
read: float
|
|
write: float
|
|
pool: float
|
|
|
|
def __init__(self, *, connect: float, read: float, write: float, pool: float) -> None:
|
|
self.connect = connect
|
|
self.read = read
|
|
self.write = write
|
|
self.pool = pool
|
|
|
|
|
|
class FakeLimits:
|
|
max_connections: int
|
|
max_keepalive_connections: int
|
|
keepalive_expiry: float
|
|
|
|
def __init__(self, *, max_connections: int, max_keepalive_connections: int, keepalive_expiry: float) -> None:
|
|
self.max_connections = max_connections
|
|
self.max_keepalive_connections = max_keepalive_connections
|
|
self.keepalive_expiry = keepalive_expiry
|
|
|
|
|
|
class FakeRedisModule:
|
|
fake_redis: ClassVar[FakeRedis | None] = None
|
|
|
|
@staticmethod
|
|
def from_url(_url: str) -> FakeRedis:
|
|
assert FakeRedisModule.fake_redis is not None
|
|
return FakeRedisModule.fake_redis
|
|
|
|
|
|
class FakeHttpxModule:
|
|
Timeout: ClassVar[type[FakeTimeout]] = FakeTimeout
|
|
Limits: ClassVar[type[FakeLimits]] = FakeLimits
|
|
AsyncClient: ClassVar[type[FakePluginDaemonHttpClient]] = FakePluginDaemonHttpClient
|
|
|
|
|
|
def test_create_app_creates_scheduler_and_closes_after_shutdown(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
fake_redis, fake_http_client = _patch_app_lifecycle(monkeypatch)
|
|
|
|
settings = ServerSettings(
|
|
redis_url="redis://example.invalid/0",
|
|
redis_prefix="test",
|
|
shutdown_grace_seconds=5,
|
|
run_retention_seconds=7,
|
|
plugin_daemon_url="http://plugin-daemon",
|
|
plugin_daemon_api_key="daemon-secret",
|
|
shellctl_entrypoint="http://shellctl",
|
|
shellctl_auth_token="shell-secret",
|
|
agent_stub_url="https://agent.example.com/agent-stub",
|
|
server_secret_key=_base64url_secret(b"1" * 32),
|
|
dify_api_base_url="https://api.example.com",
|
|
dify_api_inner_api_key="inner-secret",
|
|
plugin_daemon_connect_timeout=1,
|
|
plugin_daemon_read_timeout=2,
|
|
plugin_daemon_write_timeout=3,
|
|
plugin_daemon_pool_timeout=4,
|
|
plugin_daemon_max_connections=5,
|
|
plugin_daemon_max_keepalive_connections=3,
|
|
plugin_daemon_keepalive_expiry=6,
|
|
)
|
|
|
|
with TestClient(create_app(settings)):
|
|
assert len(FakeRunScheduler.created) == 1
|
|
scheduler = FakeRunScheduler.created[0]
|
|
assert scheduler.shutdown_grace_seconds == 5
|
|
layer_providers = scheduler.layer_providers
|
|
assert isinstance(layer_providers, tuple)
|
|
execution_context_provider = next(
|
|
provider for provider in layer_providers if provider.type_id == "dify.execution_context"
|
|
)
|
|
execution_context_layer = execution_context_provider.create_layer(
|
|
DifyExecutionContextLayerConfig(
|
|
tenant_id="tenant-1",
|
|
user_from="account",
|
|
agent_mode="workflow_run",
|
|
invoke_from="service-api",
|
|
)
|
|
)
|
|
shell_provider = next(provider for provider in layer_providers if provider.type_id == "dify.shell")
|
|
shell_layer = shell_provider.create_layer(DifyShellLayerConfig())
|
|
assert isinstance(execution_context_layer, DifyExecutionContextLayer)
|
|
assert isinstance(shell_layer, DifyShellLayer)
|
|
assert execution_context_layer.daemon_url == "http://plugin-daemon"
|
|
assert execution_context_layer.daemon_api_key == "daemon-secret"
|
|
assert shell_layer.shellctl_entrypoint == "http://shellctl"
|
|
assert shell_layer.agent_stub_url == "https://agent.example.com/agent-stub"
|
|
shellctl_client = shell_layer.shellctl_client_factory("http://shellctl")
|
|
assert isinstance(shellctl_client, ShellctlClient)
|
|
assert shellctl_client.token == "shell-secret"
|
|
asyncio.run(shellctl_client.close())
|
|
http_client = scheduler.plugin_daemon_http_client
|
|
assert http_client is fake_http_client
|
|
assert http_client.is_closed is False
|
|
store = scheduler.store
|
|
assert isinstance(store, RedisRunStore)
|
|
assert store.run_retention_seconds == 7
|
|
assert any(getattr(route, "path", None) == "/agent-stub/connections" for route in create_app(settings).routes)
|
|
assert any(
|
|
getattr(route, "path", None) == "/agent-stub/files/upload-request" for route in create_app(settings).routes
|
|
)
|
|
assert any(
|
|
getattr(route, "path", None) == "/agent-stub/files/download-request"
|
|
for route in create_app(settings).routes
|
|
)
|
|
|
|
assert FakeRunScheduler.created[0].shutdown_called is True
|
|
assert FakeRunScheduler.created[0].plugin_daemon_http_client.is_closed is True
|
|
assert fake_redis.closed is True
|
|
|
|
|
|
def test_create_app_wires_authenticated_agent_stub_connection_route(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
fake_redis, fake_http_client = _patch_app_lifecycle(monkeypatch)
|
|
settings = ServerSettings(
|
|
redis_url="redis://example.invalid/0",
|
|
agent_stub_url="https://agent.example.com/agent-stub",
|
|
server_secret_key=_base64url_secret(b"1" * 32),
|
|
)
|
|
token_codec = settings.create_agent_stub_token_codec()
|
|
assert token_codec is not None
|
|
token = token_codec.encode_connection_token(_execution_context(), now=int(time.time()) - 1)
|
|
|
|
with TestClient(create_app(settings)) as client:
|
|
response = client.post(
|
|
"/agent-stub/connections",
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
json={"protocol_version": 1, "argv": ["connect"]},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["status"] == "connected"
|
|
assert isinstance(response.json()["connection_id"], str)
|
|
assert FakeRunScheduler.created[0].shutdown_called is True
|
|
assert fake_http_client.is_closed is True
|
|
assert fake_redis.closed is True
|
|
|
|
|
|
def test_create_app_wires_authenticated_agent_stub_file_upload_route(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
fake_redis, fake_http_client = _patch_app_lifecycle(monkeypatch)
|
|
settings = ServerSettings(
|
|
redis_url="redis://example.invalid/0",
|
|
agent_stub_url="https://agent.example.com/agent-stub",
|
|
server_secret_key=_base64url_secret(b"1" * 32),
|
|
dify_api_base_url="https://api.example.com",
|
|
dify_api_inner_api_key="inner-secret",
|
|
)
|
|
token_codec = settings.create_agent_stub_token_codec()
|
|
assert token_codec is not None
|
|
token = token_codec.encode_connection_token(_execution_context(), now=int(time.time()) - 1)
|
|
|
|
original_async_client = httpx.AsyncClient
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
assert str(request.url) == "https://api.example.com/inner/api/upload/file/request"
|
|
assert request.headers["X-Inner-Api-Key"] == "inner-secret"
|
|
return httpx.Response(200, json={"data": {"url": "https://files.example.com/upload"}})
|
|
|
|
monkeypatch.setattr(
|
|
"dify_agent.agent_stub.server.agent_stub_files.httpx.AsyncClient",
|
|
lambda **kwargs: original_async_client(transport=httpx.MockTransport(handler), **kwargs),
|
|
)
|
|
|
|
with TestClient(create_app(settings)) as client:
|
|
response = client.post(
|
|
"/agent-stub/files/upload-request",
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
json={"filename": "report.pdf", "mimetype": "application/pdf"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json() == {"upload_url": "https://files.example.com/upload"}
|
|
assert FakeRunScheduler.created[0].shutdown_called is True
|
|
assert fake_http_client.is_closed is True
|
|
assert fake_redis.closed is True
|
|
|
|
|
|
def test_create_app_starts_and_stops_agent_stub_grpc_server_for_grpc_url(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
fake_redis, fake_http_client = _patch_app_lifecycle(monkeypatch)
|
|
started: dict[str, object] = {}
|
|
fake_grpc_server = FakeAgentStubGRPCServer()
|
|
|
|
async def fake_start_agent_stub_grpc_server(**kwargs):
|
|
started.update(kwargs)
|
|
return fake_grpc_server
|
|
|
|
monkeypatch.setattr(app_module, "start_agent_stub_grpc_server", fake_start_agent_stub_grpc_server)
|
|
|
|
settings = ServerSettings(
|
|
redis_url="redis://example.invalid/0",
|
|
agent_stub_url="grpc://agent.example.com:9091",
|
|
agent_stub_grpc_bind_address="0.0.0.0:9191",
|
|
server_secret_key=_base64url_secret(b"1" * 32),
|
|
)
|
|
|
|
with TestClient(create_app(settings)):
|
|
assert started["public_url"] == "grpc://agent.example.com:9091"
|
|
assert started["bind_address"] == "0.0.0.0:9191"
|
|
|
|
assert fake_grpc_server.closed is True
|
|
assert FakeRunScheduler.created[0].shutdown_called is True
|
|
assert fake_http_client.is_closed is True
|
|
assert fake_redis.closed is True
|
|
|
|
|
|
def test_create_plugin_daemon_http_client_uses_configured_httpx_construction_args(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(app_module, "httpx", FakeHttpxModule)
|
|
|
|
client = create_plugin_daemon_http_client(ServerSettings())
|
|
|
|
assert isinstance(client, FakePluginDaemonHttpClient)
|
|
assert isinstance(client.timeout, FakeTimeout)
|
|
assert client.timeout.connect == 10
|
|
assert client.timeout.read == 600
|
|
assert client.timeout.write == 30
|
|
assert client.timeout.pool == 10
|
|
assert isinstance(client.limits, FakeLimits)
|
|
assert client.limits.max_connections == 100
|
|
assert client.limits.max_keepalive_connections == 20
|
|
assert client.limits.keepalive_expiry == 30
|
|
assert client.trust_env is False
|