From d6147456e3b679dc2146c0c28309bd9383350f26 Mon Sep 17 00:00:00 2001 From: Xiyuan Chen <52963600+GareArc@users.noreply.github.com> Date: Fri, 21 Aug 2026 10:39:36 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20Enterprise=20Home=20Snapshot=20?= =?UTF-8?q?=E2=80=94=20runtime=20endpoints=20and=20dify-agent=20adapter=20?= =?UTF-8?q?(#40996)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/clients/agent_backend/errors.py | 29 ++ api/clients/agent_backend/factory.py | 2 + api/configs/extra/agent_backend_config.py | 8 + .../console/app/agent_app_sandbox.py | 16 +- api/services/agent/errors.py | 12 + api/services/agent/home_snapshot_service.py | 25 +- api/services/agent/workspace_service.py | 1 + .../clients/agent_backend/test_factory.py | 13 +- .../agent/test_home_snapshot_service.py | 106 +++++- .../services/agent/test_workspace_service.py | 10 + dify-agent-runtime/README.md | 39 +++ dify-agent-runtime/cmd/shellctl/main.go | 5 +- dify-agent-runtime/go.mod | 1 + dify-agent-runtime/go.sum | 2 + dify-agent-runtime/internal/server/api.go | 15 + .../internal/server/api_test.go | 45 +++ dify-agent-runtime/internal/server/config.go | 30 +- .../internal/server/config_test.go | 83 ++++- .../internal/server/service_test.go | 2 +- .../internal/server/snapshot.go | 173 +++++++++ .../internal/server/snapshot_test.go | 331 ++++++++++++++++++ .../internal/server/tmux_test.go | 2 +- dify-agent-runtime/internal/server/types.go | 6 + .../internal/snapshot/exclude_test.go | 78 +++++ .../internal/snapshot/excluder.go | 45 +++ .../internal/snapshot/gitignore/LICENSE | 201 +++++++++++ .../internal/snapshot/gitignore/doc.go | 22 ++ .../internal/snapshot/gitignore/matcher.go | 30 ++ .../internal/snapshot/gitignore/pattern.go | 155 ++++++++ .../internal/snapshot/mkfifo_test.go | 7 + .../internal/snapshot/restore.go | 200 +++++++++++ .../internal/snapshot/restore_test.go | 251 +++++++++++++ dify-agent-runtime/internal/snapshot/save.go | 103 ++++++ .../internal/snapshot/save_test.go | 209 +++++++++++ .../internal/snapshot/snapshot.go | 8 + .../dify_agent/runtime_backend/__init__.py | 2 + .../dify_agent/runtime_backend/enterprise.py | 168 +++++++-- .../src/dify_agent/runtime_backend/errors.py | 5 + .../src/dify_agent/runtime_backend/profile.py | 8 +- .../src/dify_agent/server/home_snapshots.py | 3 + dify-agent/src/dify_agent/server/settings.py | 2 + .../test_enterprise_backend.py | 330 ++++++++++++++++- .../local/dify_agent/server/test_settings.py | 17 +- docker/.env.example | 3 + docker/docker-compose-template.yaml | 2 + docker/docker-compose.yaml | 2 + 46 files changed, 2725 insertions(+), 82 deletions(-) create mode 100644 dify-agent-runtime/internal/server/snapshot.go create mode 100644 dify-agent-runtime/internal/server/snapshot_test.go create mode 100644 dify-agent-runtime/internal/snapshot/exclude_test.go create mode 100644 dify-agent-runtime/internal/snapshot/excluder.go create mode 100644 dify-agent-runtime/internal/snapshot/gitignore/LICENSE create mode 100644 dify-agent-runtime/internal/snapshot/gitignore/doc.go create mode 100644 dify-agent-runtime/internal/snapshot/gitignore/matcher.go create mode 100644 dify-agent-runtime/internal/snapshot/gitignore/pattern.go create mode 100644 dify-agent-runtime/internal/snapshot/mkfifo_test.go create mode 100644 dify-agent-runtime/internal/snapshot/restore.go create mode 100644 dify-agent-runtime/internal/snapshot/restore_test.go create mode 100644 dify-agent-runtime/internal/snapshot/save.go create mode 100644 dify-agent-runtime/internal/snapshot/save_test.go create mode 100644 dify-agent-runtime/internal/snapshot/snapshot.go diff --git a/api/clients/agent_backend/errors.py b/api/clients/agent_backend/errors.py index 99ef2fecf5d..5ae21df0aba 100644 --- a/api/clients/agent_backend/errors.py +++ b/api/clients/agent_backend/errors.py @@ -8,11 +8,40 @@ classes. from __future__ import annotations +from http import HTTPStatus from typing import Any +from dify_agent.client import DifyAgentHTTPError from dify_agent.protocol import RunFailureType +def backend_reported_failure(exc: DifyAgentHTTPError) -> bool: + """Whether the backend itself reported the failure. + + ``DifyAgentValidationError`` is a ``DifyAgentHTTPError`` that the client also + raises when a *successful* response fails DTO validation. Those carry a 2xx + status and a Pydantic error list, not a backend-authored error envelope. + """ + return exc.status_code >= HTTPStatus.BAD_REQUEST + + +def backend_error_detail( + exc: DifyAgentHTTPError, + *, + default_code: str = "agent_backend_error", +) -> tuple[str, str]: + """Split an Agent backend HTTP error into the code and message the backend itself sent.""" + detail = exc.detail + if isinstance(detail, dict): + code = detail.get("code") + message = detail.get("message") + return ( + code if isinstance(code, str) and code else default_code, + message if isinstance(message, str) and message else str(exc), + ) + return default_code, str(detail) + + class AgentBackendError(Exception): """Base error for API-side Agent backend integration failures.""" diff --git a/api/clients/agent_backend/factory.py b/api/clients/agent_backend/factory.py index 2f4817dfcdd..4507e93ef40 100644 --- a/api/clients/agent_backend/factory.py +++ b/api/clients/agent_backend/factory.py @@ -13,6 +13,7 @@ def create_agent_backend_client( base_url: str, api_token: str | None = None, stream_timeout: float = 30, + timeout: float = 30.0, binding_file_download_timeout: float = 240, ) -> Client: api_token = api_token.strip() if api_token else None @@ -20,6 +21,7 @@ def create_agent_backend_client( return Client( base_url=base_url, stream_timeout=stream_timeout, + timeout=timeout, binding_file_download_timeout=binding_file_download_timeout, headers=headers, ) diff --git a/api/configs/extra/agent_backend_config.py b/api/configs/extra/agent_backend_config.py index 2a43bc20159..8b782b5e3e2 100644 --- a/api/configs/extra/agent_backend_config.py +++ b/api/configs/extra/agent_backend_config.py @@ -37,6 +37,14 @@ class AgentBackendConfig(BaseSettings): default=3, ) + AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS: PositiveFloat = Field( + description=( + "Client timeout for Agent backend calls that may carry a Home Snapshot transfer: " + "snapshot capture and delete, and Execution Binding creation that restores one. " + ), + default=45.0, + ) + AGENT_BACKEND_BINDING_FILE_DOWNLOAD_TIMEOUT_SECONDS: PositiveFloat = Field( description="Client timeout for converting a Binding file to a ToolFile through the Agent backend.", default=240, diff --git a/api/controllers/console/app/agent_app_sandbox.py b/api/controllers/console/app/agent_app_sandbox.py index 233fb977c29..e58d25d0059 100644 --- a/api/controllers/console/app/agent_app_sandbox.py +++ b/api/controllers/console/app/agent_app_sandbox.py @@ -10,10 +10,11 @@ from __future__ import annotations from typing import Literal from uuid import UUID -from dify_agent.client import DifyAgentClientError, DifyAgentHTTPError, DifyAgentTimeoutError +from dify_agent.client import DifyAgentClientError, DifyAgentHTTPError from flask_restx import Resource from pydantic import BaseModel, Field +from clients.agent_backend.errors import backend_error_detail, backend_reported_failure from controllers.common.schema import ( query_params_from_model, query_params_from_request, @@ -135,15 +136,10 @@ register_response_schema_models( def _handle(exc: Exception) -> tuple[dict[str, object], int]: if isinstance(exc, AgentSandboxInspectorError): return {"code": exc.code, "message": exc.message}, exc.status_code - if isinstance(exc, DifyAgentHTTPError): - detail = exc.detail - if isinstance(detail, dict): - return { - "code": detail.get("code", "agent_backend_error"), - "message": detail.get("message", str(exc)), - }, exc.status_code - return {"code": "agent_backend_error", "message": str(detail)}, exc.status_code - if isinstance(exc, DifyAgentTimeoutError | DifyAgentClientError): + if isinstance(exc, DifyAgentHTTPError) and backend_reported_failure(exc): + code, message = backend_error_detail(exc) + return {"code": code, "message": message}, exc.status_code + if isinstance(exc, DifyAgentClientError): return {"code": "agent_backend_unreachable", "message": str(exc)}, 502 raise exc diff --git a/api/services/agent/errors.py b/api/services/agent/errors.py index 45f10031491..89875e4dc74 100644 --- a/api/services/agent/errors.py +++ b/api/services/agent/errors.py @@ -41,6 +41,18 @@ class AgentBuildSandboxNotFoundError(BaseHTTPException): code = 404 +class AgentHomeSnapshotCreateFailedError(BaseHTTPException): + error_code = "agent_home_snapshot_create_failed" + description = "Failed to capture the Agent Home Snapshot." + code = 502 + + +class AgentHomeSnapshotTooLargeError(BaseHTTPException): + error_code = "agent_home_snapshot_too_large" + description = "The Agent Home Snapshot exceeds the configured size limit." + code = 413 + + class AgentSoulLockedError(BadRequest): description = "Agent Soul is locked for this workflow node." diff --git a/api/services/agent/home_snapshot_service.py b/api/services/agent/home_snapshot_service.py index 9e2113c56a6..a1e4ff5d74d 100644 --- a/api/services/agent/home_snapshot_service.py +++ b/api/services/agent/home_snapshot_service.py @@ -2,11 +2,15 @@ from __future__ import annotations -from dify_agent.client import Client, DifyAgentNotFoundError +import logging +from http import HTTPStatus + +from dify_agent.client import Client, DifyAgentClientError, DifyAgentHTTPError, DifyAgentNotFoundError from dify_agent.protocol import CreateHomeSnapshotFromBindingRequest from sqlalchemy import select from sqlalchemy.orm import Session +from clients.agent_backend.errors import backend_error_detail, backend_reported_failure from clients.agent_backend.factory import create_agent_backend_client from configs import dify_config from core.db.session_factory import session_factory @@ -21,9 +25,15 @@ from models.agent import ( AgentWorkingResourceStatus, AgentWorkspaceOwnerType, ) -from services.agent.errors import AgentBuildSandboxNotFoundError +from services.agent.errors import ( + AgentBuildSandboxNotFoundError, + AgentHomeSnapshotCreateFailedError, + AgentHomeSnapshotTooLargeError, +) from services.agent.workspace_service import AgentWorkspaceService, WorkspaceOwnerScope +logger = logging.getLogger(__name__) + class AgentHomeSnapshotUnavailableError(RuntimeError): """The requested owner-scoped Home Snapshot cannot be used.""" @@ -90,6 +100,16 @@ class AgentHomeSnapshotService: ) except DifyAgentNotFoundError as exc: raise AgentBuildSandboxNotFoundError() from exc + except DifyAgentHTTPError as exc: + if not backend_reported_failure(exc): + logger.exception("Unreadable Home Snapshot response for agent %s", build_draft.agent_id) + raise AgentHomeSnapshotCreateFailedError() from exc + _, message = backend_error_detail(exc) + if exc.status_code == HTTPStatus.REQUEST_ENTITY_TOO_LARGE: + raise AgentHomeSnapshotTooLargeError(message) from exc + raise AgentHomeSnapshotCreateFailedError(message) from exc + except DifyAgentClientError as exc: + raise AgentHomeSnapshotCreateFailedError(str(exc)) from exc home_snapshot = AgentHomeSnapshot( id=home_snapshot_id, @@ -155,6 +175,7 @@ class AgentHomeSnapshotService: return create_agent_backend_client( base_url=base_url, api_token=dify_config.AGENT_BACKEND_API_TOKEN, + timeout=dify_config.AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS, ) diff --git a/api/services/agent/workspace_service.py b/api/services/agent/workspace_service.py index cdf2736b523..a3ffeea0064 100644 --- a/api/services/agent/workspace_service.py +++ b/api/services/agent/workspace_service.py @@ -452,6 +452,7 @@ class AgentWorkspaceService: return create_agent_backend_client( base_url=base_url, api_token=dify_config.AGENT_BACKEND_API_TOKEN, + timeout=dify_config.AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS, ) diff --git a/api/tests/unit_tests/clients/agent_backend/test_factory.py b/api/tests/unit_tests/clients/agent_backend/test_factory.py index 1c3a1878e98..626bf5f1dfd 100644 --- a/api/tests/unit_tests/clients/agent_backend/test_factory.py +++ b/api/tests/unit_tests/clients/agent_backend/test_factory.py @@ -31,6 +31,7 @@ def test_create_agent_backend_client_forwards_authentication( client_cls.assert_called_once_with( base_url="http://agent-backend", stream_timeout=30, + timeout=30.0, binding_file_download_timeout=240, headers=headers, ) @@ -54,8 +55,16 @@ def test_create_agent_backend_run_client_forwards_stream_read_timeout(create_cli @pytest.mark.parametrize( ("factory", "module", "extra_kwargs"), [ - (home_snapshot_service.AgentHomeSnapshotService._client, home_snapshot_service, {}), - (workspace_service.AgentWorkspaceService._client, workspace_service, {}), + ( + home_snapshot_service.AgentHomeSnapshotService._client, + home_snapshot_service, + {"timeout": dify_config.AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS}, + ), + ( + workspace_service.AgentWorkspaceService._client, + workspace_service, + {"timeout": dify_config.AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS}, + ), ( agent_app_sandbox_service._default_client_factory, agent_app_sandbox_service, diff --git a/api/tests/unit_tests/services/agent/test_home_snapshot_service.py b/api/tests/unit_tests/services/agent/test_home_snapshot_service.py index b3476f1c120..c66a26aeeb7 100644 --- a/api/tests/unit_tests/services/agent/test_home_snapshot_service.py +++ b/api/tests/unit_tests/services/agent/test_home_snapshot_service.py @@ -3,8 +3,10 @@ from types import SimpleNamespace from unittest.mock import MagicMock import pytest +from dify_agent.client import DifyAgentHTTPError, DifyAgentNotFoundError, DifyAgentTimeoutError from sqlalchemy.orm import Session +from configs import dify_config from models.agent import ( Agent, AgentConfigDraft, @@ -14,7 +16,11 @@ from models.agent import ( AgentWorkingResourceStatus, ) from models.agent_config_entities import AgentSoulConfig -from services.agent.errors import AgentBuildSandboxNotFoundError +from services.agent.errors import ( + AgentBuildSandboxNotFoundError, + AgentHomeSnapshotCreateFailedError, + AgentHomeSnapshotTooLargeError, +) from services.agent.home_snapshot_service import AgentHomeSnapshotService, validate_home_snapshot_binding from services.agent.workspace_service import AgentWorkspaceService @@ -39,6 +45,14 @@ def _client(*, snapshot_ref: str = "snapshot-ref-1") -> MagicMock: return client +def test_home_snapshot_client_outlasts_the_gateway_snapshot_budget(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(dify_config, "AGENT_BACKEND_BASE_URL", "http://agent.example") + + client = AgentHomeSnapshotService._client() + + assert client._timeout == 45.0 + + def test_validate_home_snapshot_binding_accepts_default_home_without_ledger_lookup() -> None: session = MagicMock() validate_home_snapshot_binding( @@ -300,3 +314,93 @@ def test_home_snapshot_collection_final_delete_failure_propagates( assert exc_info.value is error delete.assert_called_once_with(snapshot_ref="snapshot-ref-1") + + +def _apply_with_client_error(monkeypatch: pytest.MonkeyPatch, error: Exception) -> None: + session = MagicMock() + session.scalar.return_value = SimpleNamespace(app_id="app-1", backing_app_id=None) + binding = SimpleNamespace( + backend_binding_ref="binding-ref-1", + agent_id="agent-1", + base_home_snapshot_id="home-old", + agent_config_version_id="build-1", + agent_config_version_kind="build_draft", + ) + client = MagicMock() + client.create_home_snapshot_from_binding_sync.side_effect = error + monkeypatch.setattr(AgentHomeSnapshotService, "_client", lambda: nullcontext(client)) + monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=binding)) + monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", MagicMock()) + + AgentHomeSnapshotService.create_for_build_apply(session=session, build_draft=_build_draft()) + + +def test_build_apply_surfaces_the_backend_message(monkeypatch: pytest.MonkeyPatch) -> None: + error = DifyAgentHTTPError( + status_code=502, + detail={ + "code": "home_snapshot_create_failed", + "message": "500 snapshot_save_failed: object store unavailable", + }, + ) + + with pytest.raises(AgentHomeSnapshotCreateFailedError) as excinfo: + _apply_with_client_error(monkeypatch, error) + + assert excinfo.value.data == { + "code": "agent_home_snapshot_create_failed", + "message": "500 snapshot_save_failed: object store unavailable", + "status": 502, + } + + +def test_build_apply_maps_a_too_large_snapshot_to_413(monkeypatch: pytest.MonkeyPatch) -> None: + error = DifyAgentHTTPError( + status_code=413, + detail={ + "code": "home_snapshot_too_large", + "message": "home snapshot exceeds the configured limit of 67108864 bytes", + }, + ) + + with pytest.raises(AgentHomeSnapshotTooLargeError) as excinfo: + _apply_with_client_error(monkeypatch, error) + + assert excinfo.value.data == { + "code": "agent_home_snapshot_too_large", + "message": "home snapshot exceeds the configured limit of 67108864 bytes", + "status": 413, + } + + +def test_build_apply_surfaces_a_non_dict_backend_detail(monkeypatch: pytest.MonkeyPatch) -> None: + error = DifyAgentHTTPError(status_code=500, detail="upstream exploded") + + with pytest.raises(AgentHomeSnapshotCreateFailedError) as excinfo: + _apply_with_client_error(monkeypatch, error) + + assert excinfo.value.data == { + "code": "agent_home_snapshot_create_failed", + "message": "upstream exploded", + "status": 502, + } + + +def test_build_apply_surfaces_a_transport_failure(monkeypatch: pytest.MonkeyPatch) -> None: + error = DifyAgentTimeoutError("read timed out after 45.0s") + + with pytest.raises(AgentHomeSnapshotCreateFailedError) as excinfo: + _apply_with_client_error(monkeypatch, error) + + assert excinfo.value.data == { + "code": "agent_home_snapshot_create_failed", + "message": "read timed out after 45.0s", + "status": 502, + } + + +def test_build_apply_still_maps_404_to_a_missing_sandbox(monkeypatch: pytest.MonkeyPatch) -> None: + error = DifyAgentNotFoundError(status_code=404, detail={"code": "binding_lost", "message": "gone"}) + + with pytest.raises(AgentBuildSandboxNotFoundError): + _apply_with_client_error(monkeypatch, error) diff --git a/api/tests/unit_tests/services/agent/test_workspace_service.py b/api/tests/unit_tests/services/agent/test_workspace_service.py index 8191f04cbfc..e8ce1e1f657 100644 --- a/api/tests/unit_tests/services/agent/test_workspace_service.py +++ b/api/tests/unit_tests/services/agent/test_workspace_service.py @@ -7,6 +7,7 @@ import pytest from sqlalchemy import select from sqlalchemy.orm import Session +from configs import dify_config from models.agent import ( AgentConfigVersionKind, AgentHomeSnapshot, @@ -92,6 +93,15 @@ def _binding( ) +def test_workspace_client_honors_the_configured_snapshot_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(dify_config, "AGENT_BACKEND_BASE_URL", "http://agent.example") + monkeypatch.setattr(dify_config, "AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS", 123.5) + + client = AgentWorkspaceService._client() + + assert client._timeout == 123.5 + + @pytest.mark.parametrize( "sqlite_session", [(AgentHomeSnapshot, AgentWorkspace, AgentWorkspaceBinding)], diff --git a/dify-agent-runtime/README.md b/dify-agent-runtime/README.md index 00b976c7742..0cd826421c7 100644 --- a/dify-agent-runtime/README.md +++ b/dify-agent-runtime/README.md @@ -87,6 +87,45 @@ See [here](./internal/envvar/envvar.go) Requires Linux ≥ 5.13. On unsupported kernels, a warning is printed to stderr. +## API Routes + +### Home snapshot + +Native save/restore of the runtime's Home directory (`$HOME`) as a tar+zstd +stream. One snapshot operation runs at a time per runtime (`409 snapshot_busy` +otherwise). **Neither endpoint imposes a size limit — callers own size policy +and must bound the streams in their own logic** (count bytes while reading a +save stream and abort at their cap; bound what they send to restore). + +Each operation carries a total I/O deadline set by `SHELLCTL_SNAPSHOT_TIMEOUT`, +a Go duration string (`10m`, `15m30s`). It bounds how long a stalled peer — +connection open, nobody reading — can hold the single-operation gate; a peer +that closes the connection releases it immediately. Unset or empty uses the +built-in default; an unparseable or non-positive value fails startup rather +than falling back to it. + +- `POST /v1/snapshot/save` — optional JSON body `{"excludes": [...]}`, gitignore + syntax, matched at any depth; a malformed body is refused with + `400 invalid_request`. The runtime's own state directory + (`.local/share/shellctl`) is **always** excluded and no pattern can + re-include it; an excluded directory is not descended into. The Workspace + needs no such rule — it lives outside Home. Streams the + Home as `application/octet-stream` (chunked). An empty Home is not a special + case: it streams an ordinary archive with no entries. Success is signaled by + trailers `X-Snapshot-Status: ok`, + `X-Snapshot-Sha256`, `X-Snapshot-Bytes`; a cleanly terminated stream WITHOUT + the `ok` trailer, or an aborted connection, is a failure. +- `POST /v1/snapshot/restore` — raw tar+zstd body, no parameters. Extracts + into `$HOME` under `os.Root` (path traversal, absolute names, and symlink + escapes are refused). Returns `{"entries": N, "bytes_written": M}`; + `400 archive_malformed` for invalid input, `500 restore_failed` for + non-format failures (e.g. filesystem or environmental errors). Restore is + NOT transactional — a mid-stream failure can leave a partially restored + Home, so callers must treat the sandbox as unusable and recreate it rather + than retry into it. Archives are plain tar+zstd; the decoder caps the zstd + window at 64 MiB, so foreign archives produced with long-window settings + are rejected as malformed. + ## Dependencies - Go 1.26 diff --git a/dify-agent-runtime/cmd/shellctl/main.go b/dify-agent-runtime/cmd/shellctl/main.go index ca08d3a6c23..948dddcfecd 100644 --- a/dify-agent-runtime/cmd/shellctl/main.go +++ b/dify-agent-runtime/cmd/shellctl/main.go @@ -39,7 +39,10 @@ func run() error { } // Build config - config := server.DefaultConfig() + config, err := server.DefaultConfig() + if err != nil { + return fmt.Errorf("config: %w", err) + } if *listen != "" { config.Listen = *listen } diff --git a/dify-agent-runtime/go.mod b/dify-agent-runtime/go.mod index 9fa6c29ac6b..4fc628513ce 100644 --- a/dify-agent-runtime/go.mod +++ b/dify-agent-runtime/go.mod @@ -3,6 +3,7 @@ module github.com/langgenius/dify/dify-agent-runtime go 1.26 require ( + github.com/klauspost/compress v1.19.2 github.com/landlock-lsm/go-landlock v0.9.0 github.com/spf13/cobra v1.10.2 modernc.org/sqlite v1.37.1 diff --git a/dify-agent-runtime/go.sum b/dify-agent-runtime/go.sum index c5698692187..4f3689b6d42 100644 --- a/dify-agent-runtime/go.sum +++ b/dify-agent-runtime/go.sum @@ -7,6 +7,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= +github.com/klauspost/compress v1.19.2 h1:hMRETovs/pu/dVWN7zIT1PGG8t509MwT6bO7XSi26R8= +github.com/klauspost/compress v1.19.2/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/landlock-lsm/go-landlock v0.9.0 h1:2q8G8yx9Hsd5bV+R6PJfgQl0zszNxC8KO+SIqGwfxlw= github.com/landlock-lsm/go-landlock v0.9.0/go.mod h1:mn5GSi81Jf7yMs5WSi+SUi4sUeNLUGVdbT4Id6wXNQw= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= diff --git a/dify-agent-runtime/internal/server/api.go b/dify-agent-runtime/internal/server/api.go index 0a6b34a9820..b3303f341ae 100644 --- a/dify-agent-runtime/internal/server/api.go +++ b/dify-agent-runtime/internal/server/api.go @@ -29,6 +29,12 @@ func Handler(svc *Service, config *Config) http.Handler { mux.HandleFunc("POST /v1/jobs/{job_id}/terminate", auth(handleTerminateJob(svc, config))) mux.HandleFunc("DELETE /v1/jobs/{job_id}", auth(handleDeleteJob(svc, config))) + { // Snapshot handlers + snap := newSnapshotHandlers(config) + mux.HandleFunc("POST /v1/snapshot/save", auth(snap.handleSnapshotSave())) + mux.HandleFunc("POST /v1/snapshot/restore", auth(snap.handleSnapshotRestore())) + } + return requestLoggingMiddleware(recoveryMiddleware(mux)) } @@ -232,6 +238,10 @@ func (sr *statusRecorder) WriteHeader(code int) { sr.ResponseWriter.WriteHeader(code) } +// Unwrap exposes the underlying writer so http.ResponseController works +// (flushes and per-request deadlines reach the real connection). +func (sr *statusRecorder) Unwrap() http.ResponseWriter { return sr.ResponseWriter } + func requestLoggingMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { start := time.Now() @@ -245,6 +255,11 @@ func recoveryMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { defer func() { if rec := recover(); rec != nil { + if rec == http.ErrAbortHandler { + // Deliberate connection abort (e.g. failed snapshot + // stream): let net/http tear the connection down. + panic(rec) + } stack := debug.Stack() log.Printf("PANIC %s %s: %v\n%s", r.Method, r.URL.Path, rec, stack) writeError(w, 500, "internal_panic", fmt.Sprintf("internal server error: %v", rec)) diff --git a/dify-agent-runtime/internal/server/api_test.go b/dify-agent-runtime/internal/server/api_test.go index 47d0977b4dd..7abb64c79a9 100644 --- a/dify-agent-runtime/internal/server/api_test.go +++ b/dify-agent-runtime/internal/server/api_test.go @@ -2,6 +2,7 @@ package server import ( "encoding/json" + "io" "net/http" "net/http/httptest" "strings" @@ -187,3 +188,47 @@ func TestIsNotFound(t *testing.T) { t.Error("500 error should not be detected as not found") } } + +func TestStatusRecorderUnwrap(t *testing.T) { + w := httptest.NewRecorder() + rec := &statusRecorder{ResponseWriter: w, statusCode: 200} + if rec.Unwrap() != w { + t.Fatal("Unwrap must return the wrapped ResponseWriter") + } +} + +func TestRecoveryMiddlewareRepanicsAbortHandler(t *testing.T) { + srv := httptest.NewServer(requestLoggingMiddleware(recoveryMiddleware( + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(200) + _, _ = w.Write([]byte("partial")) + panic(http.ErrAbortHandler) + })))) + defer srv.Close() + + resp, err := http.Get(srv.URL) + if err == nil { + _, readErr := io.ReadAll(resp.Body) + _ = resp.Body.Close() + if readErr == nil { + t.Fatal("expected aborted connection, got clean response") + } + } +} + +func TestRecoveryMiddlewareStillCatchesOtherPanics(t *testing.T) { + srv := httptest.NewServer(recoveryMiddleware( + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + panic("boom") + }))) + defer srv.Close() + + resp, err := http.Get(srv.URL) + if err != nil { + t.Fatal(err) + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != 500 { + t.Fatalf("expected 500, got %d", resp.StatusCode) + } +} diff --git a/dify-agent-runtime/internal/server/config.go b/dify-agent-runtime/internal/server/config.go index 226a44b8c43..a329bd512f4 100644 --- a/dify-agent-runtime/internal/server/config.go +++ b/dify-agent-runtime/internal/server/config.go @@ -1,9 +1,11 @@ package server import ( + "fmt" "os" "path/filepath" "runtime" + "strings" "time" ) @@ -27,6 +29,8 @@ const ( DefaultSQLiteBusyTimeoutMs = 5000 DefaultAuthTokenEnv = "SHELLCTL_AUTH_TOKEN" HealthStatus = "ok" + DefaultSnapshotTimeoutSeconds = 45.0 + SnapshotTimeoutEnv = "SHELLCTL_SNAPSHOT_TIMEOUT" ) // Config holds runtime configuration for the shellctl server. @@ -54,14 +58,20 @@ type Config struct { SQLiteBusyTimeoutMs int SanitizePtyCommand []string RunnerExitCommand []string + SnapshotTimeout time.Duration } // DefaultConfig returns a Config with sensible defaults. -func DefaultConfig() *Config { +func DefaultConfig() (*Config, error) { homeDir, _ := os.UserHomeDir() stateDir := defaultStateDir() runtimeDir := filepath.Join(stateDir, "runtime") + snapshotTimeout, err := parseSnapshotTimeout(os.Getenv(SnapshotTimeoutEnv)) + if err != nil { + return nil, err + } + cfg := &Config{ Listen: DefaultListen, StateDir: stateDir, @@ -85,6 +95,7 @@ func DefaultConfig() *Config { SQLiteBusyTimeoutMs: DefaultSQLiteBusyTimeoutMs, SanitizePtyCommand: []string{"shellctl-sanitize-pty"}, RunnerExitCommand: []string{"shellctl-runner-exit"}, + SnapshotTimeout: snapshotTimeout, } // Auth token from environment if not set explicitly @@ -92,7 +103,7 @@ func DefaultConfig() *Config { cfg.AuthToken = os.Getenv(DefaultAuthTokenEnv) } - return cfg + return cfg, nil } // JobsDir returns the path to the jobs artifact directory. @@ -126,3 +137,18 @@ func defaultStateDir() string { home, _ := os.UserHomeDir() return filepath.Join(home, ".local", "share", "shellctl") } + +func parseSnapshotTimeout(raw string) (time.Duration, error) { + trimmed := strings.TrimSpace(raw) + if trimmed == "" { + return time.Duration(DefaultSnapshotTimeoutSeconds * float64(time.Second)), nil + } + timeout, err := time.ParseDuration(trimmed) + if err != nil { + return 0, fmt.Errorf("%s: %w", SnapshotTimeoutEnv, err) + } + if timeout <= 0 { + return 0, fmt.Errorf("%s: must be positive, got %q", SnapshotTimeoutEnv, trimmed) + } + return timeout, nil +} diff --git a/dify-agent-runtime/internal/server/config_test.go b/dify-agent-runtime/internal/server/config_test.go index 0b6aed9449e..47619ae611c 100644 --- a/dify-agent-runtime/internal/server/config_test.go +++ b/dify-agent-runtime/internal/server/config_test.go @@ -1,11 +1,23 @@ package server import ( + "os" + "strings" "testing" + "time" ) +func mustDefaultConfig(t *testing.T) *Config { + t.Helper() + cfg, err := DefaultConfig() + if err != nil { + t.Fatalf("DefaultConfig: %v", err) + } + return cfg +} + func TestDefaultConfig(t *testing.T) { - cfg := DefaultConfig() + cfg := mustDefaultConfig(t) if cfg.Listen != DefaultListen { t.Errorf("expected Listen=%s, got %s", DefaultListen, cfg.Listen) } @@ -21,7 +33,7 @@ func TestDefaultConfig(t *testing.T) { } func TestConfigPaths(t *testing.T) { - cfg := DefaultConfig() + cfg := mustDefaultConfig(t) cfg.StateDir = "/tmp/shellctl-test" cfg.RuntimeDir = "/tmp/shellctl-test/runtime" @@ -41,7 +53,7 @@ func TestConfigPaths(t *testing.T) { func TestConfigAuthTokenFromEnv(t *testing.T) { t.Setenv("SHELLCTL_AUTH_TOKEN", "my-secret-token") - cfg := DefaultConfig() + cfg := mustDefaultConfig(t) if cfg.AuthToken != "my-secret-token" { t.Errorf("expected auth token from env, got %q", cfg.AuthToken) } @@ -49,8 +61,71 @@ func TestConfigAuthTokenFromEnv(t *testing.T) { func TestConfigNoAuthToken(t *testing.T) { t.Setenv("SHELLCTL_AUTH_TOKEN", "") - cfg := DefaultConfig() + cfg := mustDefaultConfig(t) if cfg.AuthToken != "" { t.Errorf("expected empty auth token, got %q", cfg.AuthToken) } } + +func TestDefaultConfigSnapshotFields(t *testing.T) { + cfg := mustDefaultConfig(t) + if cfg.SnapshotTimeout != 45*time.Second { + t.Errorf("SnapshotTimeout = %v, want 45s", cfg.SnapshotTimeout) + } +} + +func TestSnapshotTimeoutDefaultsWhenEnvUnset(t *testing.T) { + t.Setenv(SnapshotTimeoutEnv, "") + _ = os.Unsetenv(SnapshotTimeoutEnv) + cfg := mustDefaultConfig(t) + if cfg.SnapshotTimeout != 45*time.Second { + t.Errorf("SnapshotTimeout = %v, want 45s", cfg.SnapshotTimeout) + } +} + +func TestSnapshotTimeoutDefaultsWhenEnvEmpty(t *testing.T) { + t.Setenv(SnapshotTimeoutEnv, "") + cfg := mustDefaultConfig(t) + if cfg.SnapshotTimeout != 45*time.Second { + t.Errorf("SnapshotTimeout = %v, want 45s", cfg.SnapshotTimeout) + } +} + +func TestSnapshotTimeoutFromEnv(t *testing.T) { + cases := []struct { + raw string + want time.Duration + }{ + {"900s", 900 * time.Second}, + {"15m", 15 * time.Minute}, + {"15m30s", 15*time.Minute + 30*time.Second}, + {" 15m ", 15 * time.Minute}, + } + for _, tc := range cases { + t.Run(tc.raw, func(t *testing.T) { + t.Setenv(SnapshotTimeoutEnv, tc.raw) + cfg := mustDefaultConfig(t) + if cfg.SnapshotTimeout != tc.want { + t.Errorf("SnapshotTimeout = %v, want %v", cfg.SnapshotTimeout, tc.want) + } + }) + } +} + +func TestSnapshotTimeoutRejectsUnusableEnv(t *testing.T) { + for _, raw := range []string{"fifteen minutes", "600", "10 m", "0", "0s", "-1s", "-5m"} { + t.Run(raw, func(t *testing.T) { + t.Setenv(SnapshotTimeoutEnv, raw) + cfg, err := DefaultConfig() + if err == nil { + t.Fatalf("%q must fail startup, not run with a substituted timeout", raw) + } + if !strings.Contains(err.Error(), SnapshotTimeoutEnv) { + t.Errorf("error %q must name %s so the operator knows what to fix", err, SnapshotTimeoutEnv) + } + if cfg != nil { + t.Errorf("%q yielded a usable config; an unusable timeout must yield none", raw) + } + }) + } +} diff --git a/dify-agent-runtime/internal/server/service_test.go b/dify-agent-runtime/internal/server/service_test.go index 1f2405c86d4..4877fe7d275 100644 --- a/dify-agent-runtime/internal/server/service_test.go +++ b/dify-agent-runtime/internal/server/service_test.go @@ -66,7 +66,7 @@ func TestMaterializeStdioStatusMarksMissingSessionWithIncompleteArtifactsLost(t func setupModeTestService(t *testing.T, mode jobmode.Mode, status JobStatusName) (*Service, *JobRow) { t.Helper() stateDir := t.TempDir() - config := DefaultConfig() + config := mustDefaultConfig(t) config.StateDir = stateDir config.RuntimeDir = filepath.Join(stateDir, "runtime") if err := os.MkdirAll(config.JobsDir(), 0700); err != nil { diff --git a/dify-agent-runtime/internal/server/snapshot.go b/dify-agent-runtime/internal/server/snapshot.go new file mode 100644 index 00000000000..2638e24d629 --- /dev/null +++ b/dify-agent-runtime/internal/server/snapshot.go @@ -0,0 +1,173 @@ +package server + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "hash" + "io" + "log" + "net/http" + "os" + "strconv" + "sync" + "time" + + "github.com/langgenius/dify/dify-agent-runtime/internal/snapshot" +) + +const ( + TrailerSnapshotStatus = "X-Snapshot-Status" + TrailerSnapshotSha256 = "X-Snapshot-Sha256" + TrailerSnapshotBytes = "X-Snapshot-Bytes" + SnapshotStatusOK = "ok" + + maxSaveRequestBytes = 64 << 10 +) + +type SaveRequest struct { + Excludes []string `json:"excludes"` +} + +func decodeSaveExcludes(body io.Reader) ([]string, error) { + var req SaveRequest + if err := json.NewDecoder(io.LimitReader(body, maxSaveRequestBytes)).Decode(&req); err != nil { + if errors.Is(err, io.EOF) { + return nil, nil + } + return nil, err + } + return req.Excludes, nil +} + +// snapshotHandlers serves the native Home snapshot endpoints. +// +// SIZE CONTRACT: neither endpoint imposes a size limit. Consumers MUST bound +// the streams in their own logic. +type snapshotHandlers struct { + config *Config + gate sync.Mutex +} + +func newSnapshotHandlers(config *Config) *snapshotHandlers { + return &snapshotHandlers{config: config} +} + +// resolveHome fails loudly when $HOME is unset or not a directory — never +// degrade to a guessed or relative path. +func (h *snapshotHandlers) resolveHome() (string, error) { + home, err := os.UserHomeDir() + if err != nil { + return "", err + } + info, err := os.Stat(home) + if err != nil { + return "", err + } + if !info.IsDir() { + return "", errors.New("home is not a directory") + } + return home, nil +} + +func (h *snapshotHandlers) handleSnapshotSave() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + excludes, err := decodeSaveExcludes(r.Body) + if err != nil { + writeError(w, 400, "invalid_request", err.Error()) + return + } + + if !h.gate.TryLock() { + writeError(w, 409, "snapshot_busy", "another snapshot operation is in progress") + return + } + defer h.gate.Unlock() + + home, err := h.resolveHome() + if err != nil { + writeError(w, 500, "home_unavailable", err.Error()) + return + } + rc := http.NewResponseController(w) + if err := rc.SetWriteDeadline(time.Now().Add(h.config.SnapshotTimeout)); err != nil { + log.Printf("WARN snapshot save: set write deadline: %v", err) + } + + ctx, cancel := context.WithTimeout(r.Context(), h.config.SnapshotTimeout) + defer cancel() + + w.Header().Set("Trailer", TrailerSnapshotStatus+", "+TrailerSnapshotSha256+", "+TrailerSnapshotBytes) + w.Header().Set("Content-Type", "application/octet-stream") + + hcw := &hashCountWriter{w: w, h: sha256.New()} + if err := snapshot.SaveHome(ctx, hcw, home, excludes); err != nil { + log.Printf("ERROR snapshot save: %v", err) + if hcw.n == 0 { + w.Header().Del("Trailer") + writeError(w, 500, "snapshot_failed", err.Error()) + return + } + panic(http.ErrAbortHandler) + } + w.Header().Set(TrailerSnapshotStatus, SnapshotStatusOK) + w.Header().Set(TrailerSnapshotSha256, hex.EncodeToString(hcw.h.Sum(nil))) + w.Header().Set(TrailerSnapshotBytes, strconv.FormatInt(hcw.n, 10)) + } +} + +func (h *snapshotHandlers) handleSnapshotRestore() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + if !h.gate.TryLock() { + writeError(w, 409, "snapshot_busy", "another snapshot operation is in progress") + return + } + defer h.gate.Unlock() + + home, err := h.resolveHome() + if err != nil { + writeError(w, 500, "home_unavailable", err.Error()) + return + } + + rc := http.NewResponseController(w) + if err := rc.SetReadDeadline(time.Now().Add(h.config.SnapshotTimeout)); err != nil { + log.Printf("WARN snapshot restore: set read deadline: %v", err) + } + + ctx, cancel := context.WithTimeout(r.Context(), h.config.SnapshotTimeout) + defer cancel() + + result, err := snapshot.RestoreHome(ctx, r.Body, home) + if err != nil { + log.Printf("ERROR snapshot restore: %v", err) + if errors.Is(err, snapshot.ErrMalformed) { + writeError(w, 400, "archive_malformed", err.Error()) + return + } + writeError(w, 500, "restore_failed", err.Error()) + return + } + writeJSON(w, http.StatusOK, RestoreResponse{ + Entries: result.Entries, + BytesWritten: result.BytesWritten, + }) + } +} + +// hashCountWriter passes bytes through while hashing and counting them, so +// success trailers can carry the digest without buffering anything. +type hashCountWriter struct { + w io.Writer + h hash.Hash + n int64 +} + +func (hc *hashCountWriter) Write(p []byte) (int, error) { + n, err := hc.w.Write(p) + hc.n += int64(n) + _, _ = hc.h.Write(p[:n]) + return n, err +} diff --git a/dify-agent-runtime/internal/server/snapshot_test.go b/dify-agent-runtime/internal/server/snapshot_test.go new file mode 100644 index 00000000000..e3c5dda7bdd --- /dev/null +++ b/dify-agent-runtime/internal/server/snapshot_test.go @@ -0,0 +1,331 @@ +package server + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strconv" + "testing" + "time" + + "github.com/langgenius/dify/dify-agent-runtime/internal/snapshot" +) + +func newSnapshotTestServer(t *testing.T, cfg *Config) *httptest.Server { + t.Helper() + srv := httptest.NewServer(Handler(nil, cfg)) // snapshot routes never touch the job Service + t.Cleanup(srv.Close) + return srv +} + +func testConfig() *Config { + return &Config{SnapshotTimeout: 600 * time.Second} +} + +func setHome(t *testing.T) string { + t.Helper() + home := t.TempDir() + t.Setenv("HOME", home) + return home +} + +func TestSnapshotSaveSuccessWithTrailers(t *testing.T) { + home := setHome(t) + if err := os.WriteFile(filepath.Join(home, "data.txt"), []byte("hello"), 0o644); err != nil { + t.Fatal(err) + } + srv := newSnapshotTestServer(t, testConfig()) + + resp, err := http.Post(srv.URL+"/v1/snapshot/save", "", nil) + if err != nil { + t.Fatal(err) + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != 200 { + t.Fatalf("status = %d", resp.StatusCode) + } + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read stream: %v", err) + } + if resp.Trailer.Get(TrailerSnapshotStatus) != SnapshotStatusOK { + t.Fatalf("status trailer = %q", resp.Trailer.Get(TrailerSnapshotStatus)) + } + sum := sha256.Sum256(body) + if got := resp.Trailer.Get(TrailerSnapshotSha256); got != hex.EncodeToString(sum[:]) { + t.Fatalf("sha trailer = %q, want %q", got, hex.EncodeToString(sum[:])) + } + if got := resp.Trailer.Get(TrailerSnapshotBytes); got != strconv.FormatInt(int64(len(body)), 10) { + t.Fatalf("bytes trailer = %q, want %d", got, len(body)) + } + + // The stream is a restorable archive. + dst := t.TempDir() + if _, err := snapshot.RestoreHome(context.Background(), bytes.NewReader(body), dst); err != nil { + t.Fatalf("returned stream not restorable: %v", err) + } + got, err := os.ReadFile(filepath.Join(dst, "data.txt")) + if err != nil || string(got) != "hello" { + t.Fatalf("restored content = %q err=%v", got, err) + } +} + +// An empty Home is not a special case: it produces an ordinary archive with no +// entries, so every caller stores and restores it through the same path. +func TestSnapshotSaveEmptyHome(t *testing.T) { + setHome(t) + srv := newSnapshotTestServer(t, testConfig()) + + resp, err := http.Post(srv.URL+"/v1/snapshot/save", "", nil) + if err != nil { + t.Fatal(err) + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != 200 { + t.Fatalf("empty home: status = %d, want 200", resp.StatusCode) + } + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read stream: %v", err) + } + if resp.Trailer.Get(TrailerSnapshotStatus) != SnapshotStatusOK { + t.Fatalf("status trailer = %q", resp.Trailer.Get(TrailerSnapshotStatus)) + } + if len(body) == 0 { + t.Fatal("empty home produced no bytes; callers cannot distinguish it from a dropped stream") + } + + dst := t.TempDir() + res, err := snapshot.RestoreHome(context.Background(), bytes.NewReader(body), dst) + if err != nil { + t.Fatalf("empty-home archive not restorable: %v", err) + } + if res.Entries != 0 || res.BytesWritten != 0 { + t.Fatalf("restored %+v, want zero entries and bytes", res) + } +} + +func TestSnapshotSaveAbortsOnMidStreamFailure(t *testing.T) { + if os.Geteuid() == 0 { + t.Skip("running as root: permission checks are bypassed") + } + home := setHome(t) + if err := os.WriteFile(filepath.Join(home, "ok.txt"), bytes.Repeat([]byte("x"), 64*1024), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(home, "zz-locked.txt"), []byte("secret"), 0o000); err != nil { + t.Fatal(err) + } + srv := newSnapshotTestServer(t, testConfig()) + + resp, err := http.Post(srv.URL+"/v1/snapshot/save", "", nil) + if err != nil { + return // aborted before headers: also a valid failure surface + } + defer func() { _ = resp.Body.Close() }() + _, readErr := io.ReadAll(resp.Body) + if readErr == nil && resp.Trailer.Get(TrailerSnapshotStatus) == SnapshotStatusOK { + t.Fatal("mid-stream failure must never produce a clean ok stream") + } +} + +func TestSnapshotBusy(t *testing.T) { + setHome(t) + cfg := testConfig() + snap := newSnapshotHandlers(cfg) + if !snap.gate.TryLock() { + t.Fatal("fresh gate must lock") + } + defer snap.gate.Unlock() + + req := httptest.NewRequest("POST", "/v1/snapshot/save", nil) + w := httptest.NewRecorder() + snap.handleSnapshotSave()(w, req) + if w.Code != 409 { + t.Fatalf("busy save: status = %d, want 409", w.Code) + } + var savePayload ErrorResponse + if err := json.NewDecoder(w.Body).Decode(&savePayload); err != nil { + t.Fatal(err) + } + if savePayload.Error.Code != "snapshot_busy" { + t.Fatalf("busy save: error code = %q, want snapshot_busy", savePayload.Error.Code) + } + + req = httptest.NewRequest("POST", "/v1/snapshot/restore", nil) + w = httptest.NewRecorder() + snap.handleSnapshotRestore()(w, req) + if w.Code != 409 { + t.Fatalf("busy restore: status = %d, want 409", w.Code) + } + var restorePayload ErrorResponse + if err := json.NewDecoder(w.Body).Decode(&restorePayload); err != nil { + t.Fatal(err) + } + if restorePayload.Error.Code != "snapshot_busy" { + t.Fatalf("busy restore: error code = %q, want snapshot_busy", restorePayload.Error.Code) + } +} + +// TestSnapshotWireContract pins the literal wire strings remote clients parse. +// If this test fails, the gateway protocol changed — that is a breaking change, +// not a refactor. +func TestSnapshotWireContract(t *testing.T) { + if TrailerSnapshotStatus != "X-Snapshot-Status" || + TrailerSnapshotSha256 != "X-Snapshot-Sha256" || + TrailerSnapshotBytes != "X-Snapshot-Bytes" || + SnapshotStatusOK != "ok" { + t.Fatalf("snapshot trailer contract changed: %q %q %q %q", + TrailerSnapshotStatus, TrailerSnapshotSha256, TrailerSnapshotBytes, SnapshotStatusOK) + } +} + +func TestSnapshotRestoreEndpoint(t *testing.T) { + srcHome := t.TempDir() + if err := os.WriteFile(filepath.Join(srcHome, "keep.txt"), []byte("v"), 0o644); err != nil { + t.Fatal(err) + } + var archive bytes.Buffer + if err := snapshot.SaveHome(context.Background(), &archive, srcHome, nil); err != nil { + t.Fatal(err) + } + + home := setHome(t) + srv := newSnapshotTestServer(t, testConfig()) + resp, err := http.Post(srv.URL+"/v1/snapshot/restore", "application/octet-stream", &archive) + if err != nil { + t.Fatal(err) + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != 200 { + t.Fatalf("status = %d", resp.StatusCode) + } + var result RestoreResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + t.Fatal(err) + } + if result.Entries == 0 { + t.Error("entries not counted") + } + if got, err := os.ReadFile(filepath.Join(home, "keep.txt")); err != nil || string(got) != "v" { + t.Fatalf("restored file = %q err=%v", got, err) + } +} + +func TestSnapshotRestoreMalformed(t *testing.T) { + setHome(t) + srv := newSnapshotTestServer(t, testConfig()) + resp, err := http.Post(srv.URL+"/v1/snapshot/restore", "application/octet-stream", bytes.NewReader([]byte("garbage"))) + if err != nil { + t.Fatal(err) + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != 400 { + t.Fatalf("status = %d, want 400", resp.StatusCode) + } + var payload ErrorResponse + if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil { + t.Fatal(err) + } + if payload.Error.Code != "archive_malformed" { + t.Fatalf("error code = %q", payload.Error.Code) + } +} + +func TestSnapshotSaveHomeUnavailable(t *testing.T) { + t.Setenv("HOME", "") // os.UserHomeDir errors when $HOME is unset + srv := newSnapshotTestServer(t, testConfig()) + + resp, err := http.Post(srv.URL+"/v1/snapshot/save", "", nil) + if err != nil { + t.Fatal(err) + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != 500 { + t.Fatalf("status = %d, want 500", resp.StatusCode) + } + var payload ErrorResponse + if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil { + t.Fatal(err) + } + if payload.Error.Code != "home_unavailable" { + t.Fatalf("error code = %q, want home_unavailable", payload.Error.Code) + } +} + +func TestSnapshotRoutesRequireAuth(t *testing.T) { + setHome(t) + cfg := testConfig() + cfg.AuthToken = "secret" + srv := newSnapshotTestServer(t, cfg) + + resp, err := http.Post(srv.URL+"/v1/snapshot/save", "", nil) + if err != nil { + t.Fatal(err) + } + _ = resp.Body.Close() + if resp.StatusCode != 401 { + t.Fatalf("unauthenticated save: status = %d, want 401", resp.StatusCode) + } + + resp, err = http.Post(srv.URL+"/v1/snapshot/restore", "", nil) + if err != nil { + t.Fatal(err) + } + _ = resp.Body.Close() + if resp.StatusCode != 401 { + t.Fatalf("unauthenticated restore: status = %d, want 401", resp.StatusCode) + } +} + +// TestSnapshotRestoreStalledPeerReleasesGate proves that a peer which stops +// sending body bytes mid-stream cannot wedge the single-flight gate forever: +// the read deadline set via http.ResponseController must fire and unblock +// the handler even though the server's ReadTimeout is 0. +func TestSnapshotRestoreStalledPeerReleasesGate(t *testing.T) { + setHome(t) + cfg := testConfig() + cfg.SnapshotTimeout = 300 * time.Millisecond + srv := newSnapshotTestServer(t, cfg) + + addr := srv.Listener.Addr().String() + conn, err := net.Dial("tcp", addr) + if err != nil { + t.Fatal(err) + } + defer func() { _ = conn.Close() }() + + // A syntactically valid request head announcing a chunked body, followed + // by a partial chunk. The peer then goes silent without completing it. + head := "POST /v1/snapshot/restore HTTP/1.1\r\n" + + "Host: " + addr + "\r\n" + + "Transfer-Encoding: chunked\r\n" + + "Content-Type: application/octet-stream\r\n" + + "\r\n" + + "5\r\n" + + "ab" + if _, err := conn.Write([]byte(head)); err != nil { + t.Fatal(err) + } + + // Comfortably past SnapshotTimeout: the stalled request's read deadline + // must have fired and released the gate by now. + time.Sleep(2 * time.Second) + + resp, err := http.Post(srv.URL+"/v1/snapshot/restore", "application/octet-stream", bytes.NewReader([]byte("garbage"))) + if err != nil { + t.Fatal(err) + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode == http.StatusConflict { + t.Fatal("gate still held after stalled peer's read deadline should have expired") + } +} diff --git a/dify-agent-runtime/internal/server/tmux_test.go b/dify-agent-runtime/internal/server/tmux_test.go index 68a8003a2c6..8e618671a6e 100644 --- a/dify-agent-runtime/internal/server/tmux_test.go +++ b/dify-agent-runtime/internal/server/tmux_test.go @@ -13,7 +13,7 @@ func newTestTmuxController(t *testing.T) *TmuxController { t.Skip("tmux is not installed") } - config := DefaultConfig() + config := mustDefaultConfig(t) config.RuntimeDir = t.TempDir() controller := NewTmuxController(config) t.Cleanup(func() { diff --git a/dify-agent-runtime/internal/server/types.go b/dify-agent-runtime/internal/server/types.go index b23c9c1e3f8..57e0662c28a 100644 --- a/dify-agent-runtime/internal/server/types.go +++ b/dify-agent-runtime/internal/server/types.go @@ -101,3 +101,9 @@ type ErrorDetail struct { type ErrorResponse struct { Error ErrorDetail `json:"error"` } + +// RestoreResponse reports a completed Home restore. +type RestoreResponse struct { + Entries int `json:"entries"` + BytesWritten int64 `json:"bytes_written"` +} diff --git a/dify-agent-runtime/internal/snapshot/exclude_test.go b/dify-agent-runtime/internal/snapshot/exclude_test.go new file mode 100644 index 00000000000..735cff51b1f --- /dev/null +++ b/dify-agent-runtime/internal/snapshot/exclude_test.go @@ -0,0 +1,78 @@ +package snapshot + +import "testing" + +// Caller patterns are gitignore syntax and reach any depth, which plain +// top-level name matching could not express. +func TestExcluderCallerPatterns(t *testing.T) { + for name, tc := range map[string]struct { + patterns []string + rel string + isDir bool + want bool + }{ + "exact top-level name": {[]string{".cache"}, ".cache", true, true}, + "bare name at depth": {[]string{"node_modules"}, "src/app/node_modules", true, true}, + "bare name matches file": {[]string{"tags"}, "src/tags", false, true}, + "anchored hits root only": {[]string{"/build"}, "build", true, true}, + "anchored misses depth": {[]string{"/build"}, "src/build", true, false}, + "suffix glob": {[]string{"*.log"}, "var/app.log", false, true}, + "suffix glob non-match": {[]string{"*.log"}, "var/app.txt", false, false}, + "doublestar spans segments": {[]string{".cache/**"}, ".cache/a/b/c.bin", false, true}, + "doublestar mid-pattern": {[]string{"src/**/dist"}, "src/a/b/dist", true, true}, + "dir-only skips file": {[]string{"target/"}, "target", false, false}, + "dir-only matches dir": {[]string{"target/"}, "target", true, true}, + "question mark": {[]string{"tmp?"}, "tmp1", false, true}, + "unmatched pattern is inert": {[]string{"nothing-here"}, "bin/tool.sh", false, false}, + "no patterns": {nil, "bin/tool.sh", false, false}, + } { + t.Run(name, func(t *testing.T) { + got := NewExcluder(tc.patterns).Excluded(tc.rel, tc.isDir) + if got != tc.want { + t.Errorf("Excluded(%q, isDir=%v) with %v = %v, want %v", + tc.rel, tc.isDir, tc.patterns, got, tc.want) + } + }) + } +} + +// Within the caller's own list, a later "!" pattern re-includes what an +// earlier one excluded — standard last-match-wins. +func TestExcluderCallerNegationWins(t *testing.T) { + e := NewExcluder([]string{".cache/**", "!.cache/keep-me"}) + if !e.Excluded(".cache/junk", false) { + t.Error(".cache/junk should be excluded by .cache/**") + } + if e.Excluded(".cache/keep-me", false) { + t.Error(".cache/keep-me should be re-included by the later ! pattern") + } +} + +// The two layers are ordered: defaults are decided before caller patterns are +// consulted, so no caller pattern can subtract from them. +func TestExcluderDefaultsBeatCallerPatterns(t *testing.T) { + for _, pattern := range []string{ + "!" + RuntimeStateDir, + "!/" + RuntimeStateDir, + "!" + RuntimeStateDir + "/**", + "!**", + } { + t.Run(pattern, func(t *testing.T) { + e := NewExcluder([]string{pattern}) + for _, dir := range defaultExcludes { + if !e.Excluded(dir, true) { + t.Errorf("%q pulled %s back in", pattern, dir) + } + } + }) + } +} + +// Blank entries are dropped rather than compiled into a pattern that would +// match everything or nothing surprisingly. +func TestExcluderIgnoresBlankPatterns(t *testing.T) { + e := NewExcluder([]string{"", " ", "\t"}) + if e.Excluded("bin/tool.sh", false) { + t.Error("blank patterns must not exclude anything") + } +} diff --git a/dify-agent-runtime/internal/snapshot/excluder.go b/dify-agent-runtime/internal/snapshot/excluder.go new file mode 100644 index 00000000000..af894300e96 --- /dev/null +++ b/dify-agent-runtime/internal/snapshot/excluder.go @@ -0,0 +1,45 @@ +package snapshot + +import ( + "slices" + "strings" + + "github.com/langgenius/dify/dify-agent-runtime/internal/snapshot/gitignore" +) + +const RuntimeStateDir = ".local/share/shellctl" + +// defaultExcludes are the Home paths no Home Snapshot ever carries, whatever +// the caller asks for. +var defaultExcludes = []string{RuntimeStateDir} + +// Excluder decides which Home entries stay out of an archive. +// +// Two layers: the runtime's own defaults, and the caller's gitignore-syntax +// patterns, which may match at any depth. +type Excluder struct { + matcher gitignore.Matcher +} + +// NewExcluder compiles caller patterns. +func NewExcluder(patterns []string) *Excluder { + compiled := make([]gitignore.Pattern, 0, len(patterns)) + for _, p := range patterns { + if strings.TrimSpace(p) == "" { + continue + } + compiled = append(compiled, gitignore.ParsePattern(p, nil)) + } + return &Excluder{matcher: gitignore.NewMatcher(compiled)} +} + +func (e *Excluder) Excluded(rel string, isDir bool) bool { + if isDefaultExcluded(rel) { + return true + } + return e.matcher.Match(strings.Split(rel, "/"), isDir) +} + +func isDefaultExcluded(rel string) bool { + return slices.Contains(defaultExcludes, rel) +} diff --git a/dify-agent-runtime/internal/snapshot/gitignore/LICENSE b/dify-agent-runtime/internal/snapshot/gitignore/LICENSE new file mode 100644 index 00000000000..8aa3d854cf7 --- /dev/null +++ b/dify-agent-runtime/internal/snapshot/gitignore/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "{}" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2018 Sourced Technologies, S.L. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/dify-agent-runtime/internal/snapshot/gitignore/doc.go b/dify-agent-runtime/internal/snapshot/gitignore/doc.go new file mode 100644 index 00000000000..45432c39691 --- /dev/null +++ b/dify-agent-runtime/internal/snapshot/gitignore/doc.go @@ -0,0 +1,22 @@ +// Package gitignore matches paths against gitignore-syntax patterns. +// +// VENDORED CODE — pattern.go and matcher.go are verbatim copies from go-git: +// +// github.com/go-git/go-git/v5@v5.19.2 plumbing/format/gitignore +// pattern.go sha256 6c8442cc541e54bc05db87c228f959249b357ffbf2487e457615fc06cb8a73ee +// matcher.go sha256 e104fc611ae78c88a9202fce457ef6ced879982c1d6eb179fd828996e06e69f3 +// +// Copyright 2018 Sourced Technologies, S.L., licensed under Apache-2.0; see +// the LICENSE file in this directory. Neither file has been modified. This +// doc.go is the only addition. +// +// Only the two matching files are vendored. Upstream's dir.go — which reads +// .gitignore files off a billy filesystem — is deliberately left out: it is +// the sole reason the upstream package depends on go-billy and go-git's config +// parser, and the runtime never reads patterns from disk. They arrive in the +// body of POST /v1/snapshot/save. +// +// To update: re-copy both files from the same upstream path at the new +// version, confirm they still import only path/filepath and strings, and +// refresh the hashes above. +package gitignore diff --git a/dify-agent-runtime/internal/snapshot/gitignore/matcher.go b/dify-agent-runtime/internal/snapshot/gitignore/matcher.go new file mode 100644 index 00000000000..bd1e9e2d4cf --- /dev/null +++ b/dify-agent-runtime/internal/snapshot/gitignore/matcher.go @@ -0,0 +1,30 @@ +package gitignore + +// Matcher defines a global multi-pattern matcher for gitignore patterns +type Matcher interface { + // Match matches patterns in the order of priorities. As soon as an inclusion or + // exclusion is found, not further matching is performed. + Match(path []string, isDir bool) bool +} + +// NewMatcher constructs a new global matcher. Patterns must be given in the order of +// increasing priority. That is most generic settings files first, then the content of +// the repo .gitignore, then content of .gitignore down the path or the repo and then +// the content command line arguments. +func NewMatcher(ps []Pattern) Matcher { + return &matcher{ps} +} + +type matcher struct { + patterns []Pattern +} + +func (m *matcher) Match(path []string, isDir bool) bool { + n := len(m.patterns) + for i := n - 1; i >= 0; i-- { + if match := m.patterns[i].Match(path, isDir); match > NoMatch { + return match == Exclude + } + } + return false +} diff --git a/dify-agent-runtime/internal/snapshot/gitignore/pattern.go b/dify-agent-runtime/internal/snapshot/gitignore/pattern.go new file mode 100644 index 00000000000..450b3cdf72b --- /dev/null +++ b/dify-agent-runtime/internal/snapshot/gitignore/pattern.go @@ -0,0 +1,155 @@ +package gitignore + +import ( + "path/filepath" + "strings" +) + +// MatchResult defines outcomes of a match, no match, exclusion or inclusion. +type MatchResult int + +const ( + // NoMatch defines the no match outcome of a match check + NoMatch MatchResult = iota + // Exclude defines an exclusion of a file as a result of a match check + Exclude + // Include defines an explicit inclusion of a file as a result of a match check + Include +) + +const ( + inclusionPrefix = "!" + zeroToManyDirs = "**" + patternDirSep = "/" +) + +// Pattern defines a single gitignore pattern. +type Pattern interface { + // Match matches the given path to the pattern. + Match(path []string, isDir bool) MatchResult +} + +type pattern struct { + domain []string + pattern []string + inclusion bool + dirOnly bool + isGlob bool +} + +// ParsePattern parses a gitignore pattern string into the Pattern structure. +func ParsePattern(p string, domain []string) Pattern { + // storing domain, copy it to ensure it isn't changed externally + domain = append([]string(nil), domain...) + res := pattern{domain: domain} + + if strings.HasPrefix(p, inclusionPrefix) { + res.inclusion = true + p = p[1:] + } + + if !strings.HasSuffix(p, "\\ ") { + p = strings.TrimRight(p, " ") + } + + if strings.HasSuffix(p, patternDirSep) { + res.dirOnly = true + p = p[:len(p)-1] + } + + if strings.Contains(p, patternDirSep) { + res.isGlob = true + } + + res.pattern = strings.Split(p, patternDirSep) + return &res +} + +func (p *pattern) Match(path []string, isDir bool) MatchResult { + if len(path) <= len(p.domain) { + return NoMatch + } + for i, e := range p.domain { + if path[i] != e { + return NoMatch + } + } + + path = path[len(p.domain):] + if p.isGlob && !p.globMatch(path, isDir) { + return NoMatch + } else if !p.isGlob && !p.simpleNameMatch(path, isDir) { + return NoMatch + } + + if p.inclusion { + return Include + } else { + return Exclude + } +} + +func (p *pattern) simpleNameMatch(path []string, isDir bool) bool { + for i, name := range path { + if match, err := filepath.Match(p.pattern[0], name); err != nil { + return false + } else if !match { + continue + } + if p.dirOnly && !isDir && i == len(path)-1 { + return false + } + return true + } + return false +} + +func (p *pattern) globMatch(path []string, isDir bool) bool { + matched := false + canTraverse := false + for i, pattern := range p.pattern { + if pattern == "" { + canTraverse = false + continue + } + if pattern == zeroToManyDirs { + if i == len(p.pattern)-1 { + break + } + canTraverse = true + continue + } + if strings.Contains(pattern, zeroToManyDirs) { + return false + } + if len(path) == 0 { + return false + } + if canTraverse { + canTraverse = false + for len(path) > 0 { + e := path[0] + path = path[1:] + if match, err := filepath.Match(pattern, e); err != nil { + return false + } else if match { + matched = true + break + } else if len(path) == 0 { + // if nothing left then fail + matched = false + } + } + } else { + if match, err := filepath.Match(pattern, path[0]); err != nil || !match { + return false + } + matched = true + path = path[1:] + } + } + if matched && p.dirOnly && !isDir && len(path) == 0 { + matched = false + } + return matched +} diff --git a/dify-agent-runtime/internal/snapshot/mkfifo_test.go b/dify-agent-runtime/internal/snapshot/mkfifo_test.go new file mode 100644 index 00000000000..5474f0597be --- /dev/null +++ b/dify-agent-runtime/internal/snapshot/mkfifo_test.go @@ -0,0 +1,7 @@ +//go:build unix + +package snapshot + +import "syscall" + +func mkfifo(path string) error { return syscall.Mkfifo(path, 0o644) } diff --git a/dify-agent-runtime/internal/snapshot/restore.go b/dify-agent-runtime/internal/snapshot/restore.go new file mode 100644 index 00000000000..b726aeefd3c --- /dev/null +++ b/dify-agent-runtime/internal/snapshot/restore.go @@ -0,0 +1,200 @@ +package snapshot + +import ( + "archive/tar" + "context" + "errors" + "fmt" + "io" + "io/fs" + "os" + "path" + "strings" + "syscall" + + "github.com/klauspost/compress/zstd" +) + +// ErrMalformed marks archives that violate the format or hardening rules. +var ErrMalformed = errors.New("archive malformed") + +const maxDecoderWindow = 64 << 20 + +var environmentalErrnos = []syscall.Errno{ + syscall.ENOSPC, syscall.EDQUOT, syscall.EROFS, syscall.EIO, + syscall.ENOMEM, syscall.EMFILE, syscall.ENFILE, + syscall.EACCES, syscall.EPERM, +} + +func classifyEntryErr(err error) error { + if err == nil { + return nil + } + if errors.Is(err, os.ErrDeadlineExceeded) { + return err + } + for _, errno := range environmentalErrnos { + if errors.Is(err, errno) { + return err + } + } + return fmt.Errorf("%w: %v", ErrMalformed, err) +} + +// RestoreResult reports a completed extraction. +type RestoreResult struct { + Entries int + BytesWritten int64 +} + +func isPAXSparse(hdr *tar.Header) bool { + for k := range hdr.PAXRecords { + if strings.HasPrefix(k, "GNU.sparse.") { + return true + } + } + return false +} + +// RestoreHome extracts a tar+zstd stream into homeDir in a single pass. +// Extraction runs under os.Root: escapes via absolute names, "..", or +// symlinked path components are refused by the kernel (openat2/RESOLVE_BENEATH +// on Linux). Only regular files, directories, symlinks, and hardlinks are +// accepted; hardlink targets must resolve inside the root. The stream is not +// size-limited here. +// A mid-stream failure can leave a partially restored Home; callers own +// cleanup of the surrounding sandbox. +func RestoreHome(ctx context.Context, src io.Reader, homeDir string) (RestoreResult, error) { + var res RestoreResult + root, err := os.OpenRoot(homeDir) + if err != nil { + return res, err + } + defer func() { _ = root.Close() }() + + zr, err := zstd.NewReader(src, + zstd.WithDecoderConcurrency(1), + zstd.WithDecoderMaxWindow(maxDecoderWindow), + ) + if err != nil { + return res, fmt.Errorf("%w: %v", ErrMalformed, err) + } + defer zr.Close() + + type dirMode struct { + name string + mode fs.FileMode + } + var dirModes []dirMode + + tr := tar.NewReader(zr) + for { + if err := ctx.Err(); err != nil { + return res, err + } + hdr, err := tr.Next() + if errors.Is(err, io.EOF) { + // Apply deferred directory modes in reverse order + for i := len(dirModes) - 1; i >= 0; i-- { + d := dirModes[i] + if err := classifyEntryErr(root.Chmod(d.name, d.mode)); err != nil { + return res, err + } + } + return res, nil + } + if err != nil { + if errors.Is(err, os.ErrDeadlineExceeded) { + return res, err + } + return res, fmt.Errorf("%w: %v", ErrMalformed, err) + } + name, err := cleanEntryName(hdr.Name) + if err != nil { + return res, err + } + if name == "" { + continue // the root entry itself + } + if hdr.Typeflag == tar.TypeGNUSparse || isPAXSparse(hdr) { + return res, fmt.Errorf("%w: sparse entry %q not supported", ErrMalformed, hdr.Name) + } + mode := hdr.FileInfo().Mode().Perm() + switch hdr.Typeflag { + case tar.TypeDir: + if err := classifyEntryErr(root.MkdirAll(name, 0o700)); err != nil { + return res, err + } + dirModes = append(dirModes, dirMode{name, mode}) + case tar.TypeReg: + if err := classifyEntryErr(ensureParent(root, name)); err != nil { + return res, err + } + n, err := extractFile(root, name, mode, tr) + if err != nil { + return res, classifyEntryErr(err) + } + if err := classifyEntryErr(root.Chmod(name, mode)); err != nil { + return res, err + } + res.BytesWritten += n + case tar.TypeSymlink: + if err := classifyEntryErr(ensureParent(root, name)); err != nil { + return res, err + } + if err := root.Remove(name); err != nil && !errors.Is(err, fs.ErrNotExist) { + return res, classifyEntryErr(err) + } + if err := classifyEntryErr(root.Symlink(hdr.Linkname, name)); err != nil { + return res, err + } + case tar.TypeLink: + target, err := cleanEntryName(hdr.Linkname) + if err != nil || target == "" { + return res, fmt.Errorf("%w: hardlink target %q", ErrMalformed, hdr.Linkname) + } + if err := classifyEntryErr(ensureParent(root, name)); err != nil { + return res, err + } + if err := classifyEntryErr(root.Link(target, name)); err != nil { + return res, err + } + default: + return res, fmt.Errorf("%w: unsupported entry type %d for %q", ErrMalformed, hdr.Typeflag, hdr.Name) + } + res.Entries++ + } +} + +func cleanEntryName(name string) (string, error) { + if strings.HasPrefix(name, "/") { + return "", fmt.Errorf("%w: absolute entry name %q", ErrMalformed, name) + } + cleaned := path.Clean(name) + if cleaned == ".." || strings.HasPrefix(cleaned, "../") { + return "", fmt.Errorf("%w: entry escapes root: %q", ErrMalformed, name) + } + if cleaned == "." { + return "", nil + } + return cleaned, nil +} + +func ensureParent(root *os.Root, name string) error { + if parent := path.Dir(name); parent != "." { + return root.MkdirAll(parent, 0o700) + } + return nil +} + +func extractFile(root *os.Root, name string, mode fs.FileMode, r io.Reader) (int64, error) { + f, err := root.OpenFile(name, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, mode) + if err != nil { + return 0, err + } + n, err := io.Copy(f, r) + if cerr := f.Close(); err == nil { + err = cerr + } + return n, err +} diff --git a/dify-agent-runtime/internal/snapshot/restore_test.go b/dify-agent-runtime/internal/snapshot/restore_test.go new file mode 100644 index 00000000000..9d3d7c8bdd4 --- /dev/null +++ b/dify-agent-runtime/internal/snapshot/restore_test.go @@ -0,0 +1,251 @@ +package snapshot + +import ( + "archive/tar" + "bytes" + "compress/gzip" + "context" + "errors" + "io" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/klauspost/compress/zstd" +) + +// craftArchive builds a tar+zstd stream from raw tar headers, for hostile-input tests. +func craftArchive(t *testing.T, build func(tw *tar.Writer)) []byte { + t.Helper() + var buf bytes.Buffer + zw, err := zstd.NewWriter(&buf) + if err != nil { + t.Fatal(err) + } + tw := tar.NewWriter(zw) + build(tw) + if err := tw.Close(); err != nil { + t.Fatal(err) + } + if err := zw.Close(); err != nil { + t.Fatal(err) + } + return buf.Bytes() +} + +func regEntry(tw *tar.Writer, name, content string) { + _ = tw.WriteHeader(&tar.Header{Name: name, Typeflag: tar.TypeReg, Mode: 0o644, Size: int64(len(content))}) + _, _ = io.WriteString(tw, content) +} + +func TestRestoreRoundTrip(t *testing.T) { + src := buildFixtureHome(t) + var buf bytes.Buffer + if err := SaveHome(context.Background(), &buf, src, []string{"workspace"}); err != nil { + t.Fatal(err) + } + + dst := t.TempDir() + res, err := RestoreHome(context.Background(), bytes.NewReader(buf.Bytes()), dst) + if err != nil { + t.Fatalf("RestoreHome: %v", err) + } + if res.Entries == 0 || res.BytesWritten == 0 { + t.Errorf("counters empty: %+v", res) + } + + got, err := os.ReadFile(filepath.Join(dst, ".bashrc")) + if err != nil || string(got) != "export PS1='$ '\n" { + t.Errorf(".bashrc: %q err=%v", got, err) + } + info, err := os.Stat(filepath.Join(dst, "bin", "tool.sh")) + if err != nil || info.Mode().Perm() != 0o755 { + t.Errorf("tool.sh mode: %v err=%v", info, err) + } + if fi, err := os.Stat(filepath.Join(dst, "emptydir")); err != nil || !fi.IsDir() { + t.Errorf("emptydir missing: %v err=%v", fi, err) + } + target, err := os.Readlink(filepath.Join(dst, "tool-link")) + if err != nil || target != "bin/tool.sh" { + t.Errorf("symlink: %q err=%v", target, err) + } + if _, err := os.Stat(filepath.Join(dst, "workspace")); !os.IsNotExist(err) { + t.Error("workspace must not be restored (was excluded at save)") + } +} + +func TestRestoreRejectsEscapes(t *testing.T) { + cases := map[string]func(tw *tar.Writer){ + "dotdot": func(tw *tar.Writer) { regEntry(tw, "../evil", "x") }, + "absolute": func(tw *tar.Writer) { regEntry(tw, "/etc/evil", "x") }, + "nested dotdot": func(tw *tar.Writer) { + regEntry(tw, "ok.txt", "fine") + regEntry(tw, "a/../../evil", "x") + }, + "hardlink escape": func(tw *tar.Writer) { + _ = tw.WriteHeader(&tar.Header{Name: "l", Typeflag: tar.TypeLink, Linkname: "../outside"}) + }, + "device node": func(tw *tar.Writer) { + _ = tw.WriteHeader(&tar.Header{Name: "dev", Typeflag: tar.TypeChar, Mode: 0o644}) + }, + } + for name, build := range cases { + t.Run(name, func(t *testing.T) { + home := t.TempDir() + _, err := RestoreHome(context.Background(), bytes.NewReader(craftArchive(t, build)), home) + if !errors.Is(err, ErrMalformed) { + t.Fatalf("expected ErrMalformed, got %v", err) + } + var escaped []string + parent := filepath.Dir(home) + entries, _ := os.ReadDir(parent) + for _, e := range entries { + if e.Name() == "evil" || e.Name() == "outside" { + escaped = append(escaped, e.Name()) + } + } + if len(escaped) > 0 { + t.Fatalf("files escaped the root: %v", escaped) + } + }) + } +} + +func TestRestoreRefusesSymlinkComponentEscape(t *testing.T) { + // A symlink entry pointing outside, then a write THROUGH it: os.Root must + // refuse resolving the out-of-root component. + home := t.TempDir() + data := craftArchive(t, func(tw *tar.Writer) { + _ = tw.WriteHeader(&tar.Header{Name: "sneaky", Typeflag: tar.TypeSymlink, Linkname: "../"}) + regEntry(tw, "sneaky/pwned", "x") + }) + if _, err := RestoreHome(context.Background(), bytes.NewReader(data), home); !errors.Is(err, ErrMalformed) { + t.Fatalf("expected ErrMalformed writing through out-of-root symlink, got %v", err) + } + if _, err := os.Stat(filepath.Join(filepath.Dir(home), "pwned")); !os.IsNotExist(err) { + t.Fatal("write escaped the root through a symlink component") + } +} + +func TestRestoreNotZstd(t *testing.T) { + home := t.TempDir() + _, err := RestoreHome(context.Background(), strings.NewReader("plain text, not zstd"), home) + if !errors.Is(err, ErrMalformed) { + t.Fatalf("expected ErrMalformed for non-zstd body, got %v", err) + } + + // A gzip stream (the format the deleted exec transport produced) must be + // rejected too — the runtime is deliberately zstd-only. + var gz bytes.Buffer + gw := gzip.NewWriter(&gz) + tw := tar.NewWriter(gw) + regEntry(tw, "file.txt", "content") + if err := tw.Close(); err != nil { + t.Fatal(err) + } + if err := gw.Close(); err != nil { + t.Fatal(err) + } + _, err = RestoreHome(context.Background(), bytes.NewReader(gz.Bytes()), home) + if !errors.Is(err, ErrMalformed) { + t.Fatalf("expected ErrMalformed for gzip body, got %v", err) + } +} + +func TestRestoreCancelledContext(t *testing.T) { + src := buildFixtureHome(t) + var buf bytes.Buffer + if err := SaveHome(context.Background(), &buf, src, nil); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := RestoreHome(ctx, bytes.NewReader(buf.Bytes()), t.TempDir()); err == nil { + t.Fatal("expected error from cancelled context") + } +} + +func TestRestoreRejectsSparseEntries(t *testing.T) { + // tar.Writer doesn't encode GNU.sparse.* PAX records, so test the guard directly. + // Unit test: isPAXSparse detects a header with GNU.sparse.* records + hdr := &tar.Header{ + Name: "f", + Typeflag: tar.TypeReg, + Size: 0, + PAXRecords: map[string]string{ + "GNU.sparse.major": "1", + "GNU.sparse.minor": "0", + "GNU.sparse.name": "f", + "GNU.sparse.realsize": "1099511627776", + }, + } + if !isPAXSparse(hdr) { + t.Fatal("isPAXSparse should detect GNU.sparse.* records") + } + + // The key evidence: if a PAX-sparse archive somehow reached us, + // the cleanEntryName + typeflag switch + isPAXSparse check would catch it. + // Since tar.Writer blocks encoding them, the guards are tested above and + // would reject any sparse entry in the restoration loop. +} + +func TestRestoreStripsSetuid(t *testing.T) { + content := "script" + home := t.TempDir() + // Create archive with setuid bit in header + var buf bytes.Buffer + zw, err := zstd.NewWriter(&buf) + if err != nil { + t.Fatal(err) + } + tw := tar.NewWriter(zw) + _ = tw.WriteHeader(&tar.Header{Name: "suid", Typeflag: tar.TypeReg, Mode: 0o4755, Size: int64(len(content))}) + _, _ = io.WriteString(tw, content) + _ = tw.Close() + _ = zw.Close() + + _, err = RestoreHome(context.Background(), bytes.NewReader(buf.Bytes()), home) + if err != nil { + t.Fatalf("RestoreHome: %v", err) + } + info, err := os.Stat(filepath.Join(home, "suid")) + if err != nil { + t.Fatalf("stat suid: %v", err) + } + if info.Mode().Perm() != 0o755 { + t.Errorf("expected 0755, got %o", info.Mode().Perm()) + } + if info.Mode()&os.ModeSetuid != 0 { + t.Errorf("setuid bit should be stripped, got %v", info.Mode()) + } +} + +func TestRestoreReadOnlyDirectory(t *testing.T) { + home := t.TempDir() + data := craftArchive(t, func(tw *tar.Writer) { + _ = tw.WriteHeader(&tar.Header{Name: "ro", Typeflag: tar.TypeDir, Mode: 0o555}) + regEntry(tw, "ro/child.txt", "content") + }) + _, err := RestoreHome(context.Background(), bytes.NewReader(data), home) + if err != nil { + t.Fatalf("RestoreHome: %v", err) + } + // Verify child exists + got, err := os.ReadFile(filepath.Join(home, "ro", "child.txt")) + if err != nil || string(got) != "content" { + t.Errorf("child.txt: %q err=%v", got, err) + } + // Verify directory mode is 0555 + info, err := os.Stat(filepath.Join(home, "ro")) + if err != nil { + t.Fatalf("stat ro: %v", err) + } + if info.Mode().Perm() != 0o555 { + t.Errorf("expected 0555, got %o", info.Mode().Perm()) + } + // Restore write permission for cleanup + if err := os.Chmod(filepath.Join(home, "ro"), 0o755); err != nil { + t.Logf("cleanup chmod: %v", err) + } +} diff --git a/dify-agent-runtime/internal/snapshot/save.go b/dify-agent-runtime/internal/snapshot/save.go new file mode 100644 index 00000000000..ce58d5781aa --- /dev/null +++ b/dify-agent-runtime/internal/snapshot/save.go @@ -0,0 +1,103 @@ +package snapshot + +import ( + "archive/tar" + "context" + "fmt" + "io" + "io/fs" + "os" + "path/filepath" + + "github.com/klauspost/compress/zstd" +) + +// SaveHome streams homeDir to dst as tar+zstd in a single pass with no +// intermediate spooling. The runtime's own defaults and the caller's +// gitignore-syntax excludes are skipped; see Excluder. +// +// An excluded directory is not descended into, so — as in git — a pattern +// cannot re-include anything beneath a directory that is already excluded. +// +// Symlinks are archived as symlinks; irregular files (sockets, fifos, +// devices) are skipped as runtime artifacts; ownership is not recorded. +// Callers wrap dst to hash or count the compressed bytes. +func SaveHome(ctx context.Context, dst io.Writer, homeDir string, excludes []string) error { + excluder := NewExcluder(excludes) + zw, err := zstd.NewWriter(dst, + zstd.WithEncoderLevel(zstd.SpeedDefault), + zstd.WithEncoderConcurrency(1), + ) + if err != nil { + return fmt.Errorf("create zstd writer: %w", err) + } + tw := tar.NewWriter(zw) + + walkErr := filepath.WalkDir(homeDir, func(p string, d fs.DirEntry, inErr error) error { + if inErr != nil { + return inErr + } + if err := ctx.Err(); err != nil { + return err + } + rel, err := filepath.Rel(homeDir, p) + if err != nil { + return err + } + if rel == "." { + return nil + } + relSlash := filepath.ToSlash(rel) + if excluder.Excluded(relSlash, d.IsDir()) { + if d.IsDir() { + return fs.SkipDir + } + return nil + } + info, err := d.Info() + if err != nil { + return err + } + var linkTarget string + switch { + case info.Mode()&fs.ModeSymlink != 0: + if linkTarget, err = os.Readlink(p); err != nil { + return err + } + case !info.Mode().IsRegular() && !info.IsDir(): + return nil + } + hdr, err := tar.FileInfoHeader(info, linkTarget) + if err != nil { + return err + } + hdr.Name = relSlash + if info.IsDir() { + hdr.Name += "/" + } + hdr.Uid, hdr.Gid, hdr.Uname, hdr.Gname = 0, 0, "", "" + if err := tw.WriteHeader(hdr); err != nil { + return err + } + if info.Mode().IsRegular() { + f, err := os.Open(p) + if err != nil { + return err + } + _, err = io.Copy(tw, f) + _ = f.Close() + if err != nil { + return err + } + } + return nil + }) + if walkErr != nil { + _ = zw.Close() + return walkErr + } + if err := tw.Close(); err != nil { + return err + } + return zw.Close() +} diff --git a/dify-agent-runtime/internal/snapshot/save_test.go b/dify-agent-runtime/internal/snapshot/save_test.go new file mode 100644 index 00000000000..619de860d7c --- /dev/null +++ b/dify-agent-runtime/internal/snapshot/save_test.go @@ -0,0 +1,209 @@ +package snapshot + +import ( + "archive/tar" + "bytes" + "context" + "io" + "os" + "path" + "path/filepath" + "strings" + "testing" + + "github.com/klauspost/compress/zstd" +) + +// buildFixtureHome creates a Home tree exercising files, modes, nesting, +// symlinks, an empty dir, and the excluded runtime state dir. +func buildFixtureHome(t *testing.T) string { + t.Helper() + home := t.TempDir() + mustWrite := func(rel, content string, mode os.FileMode) { + t.Helper() + p := filepath.Join(home, rel) + if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(p, []byte(content), mode); err != nil { + t.Fatal(err) + } + } + mustWrite(".bashrc", "export PS1='$ '\n", 0o644) + mustWrite("bin/tool.sh", "#!/bin/sh\necho hi\n", 0o755) + mustWrite("workspace/notes.txt", "ordinary home content", 0o644) + mustWrite(".local/share/shellctl/shellctl.db", "live server state", 0o644) + mustWrite(".local/bin/agent-tool", "#!/bin/sh\necho tool\n", 0o755) + if err := os.MkdirAll(filepath.Join(home, "emptydir"), 0o755); err != nil { + t.Fatal(err) + } + if err := os.Symlink("bin/tool.sh", filepath.Join(home, "tool-link")); err != nil { + t.Fatal(err) + } + return home +} + +// decodeArchive reads a tar+zstd stream into name -> header/content maps. +func decodeArchive(t *testing.T, data []byte) (map[string]*tar.Header, map[string][]byte) { + t.Helper() + zr, err := zstd.NewReader(bytes.NewReader(data)) + if err != nil { + t.Fatalf("not a zstd stream: %v", err) + } + defer zr.Close() + headers := map[string]*tar.Header{} + contents := map[string][]byte{} + tr := tar.NewReader(zr) + for { + hdr, err := tr.Next() + if err == io.EOF { + return headers, contents + } + if err != nil { + t.Fatalf("tar decode: %v", err) + } + headers[hdr.Name] = hdr + if hdr.Typeflag == tar.TypeReg { + body, err := io.ReadAll(tr) + if err != nil { + t.Fatal(err) + } + contents[hdr.Name] = body + } + } +} + +func TestSaveHomeArchivesTreeWithoutExcludes(t *testing.T) { + home := buildFixtureHome(t) + var buf bytes.Buffer + if err := SaveHome(context.Background(), &buf, home, nil); err != nil { + t.Fatalf("SaveHome: %v", err) + } + headers, contents := decodeArchive(t, buf.Bytes()) + + if _, ok := headers["workspace/notes.txt"]; !ok { + t.Error("a Home directory named workspace is ordinary content and must be archived") + } + if _, ok := headers[RuntimeStateDir+"/shellctl.db"]; ok { + t.Error("runtime state must not be archived") + } + if hdr := headers[".local/bin/agent-tool"]; hdr == nil { + t.Error("user-installed tooling under .local must be archived") + } + if got := string(contents[".bashrc"]); got != "export PS1='$ '\n" { + t.Errorf(".bashrc content = %q", got) + } + if hdr := headers["bin/tool.sh"]; hdr == nil || hdr.FileInfo().Mode().Perm() != 0o755 { + t.Errorf("bin/tool.sh mode not preserved: %+v", hdr) + } + if hdr := headers["emptydir/"]; hdr == nil || hdr.Typeflag != tar.TypeDir { + t.Error("empty dir entry missing") + } + link := headers["tool-link"] + if link == nil || link.Typeflag != tar.TypeSymlink || link.Linkname != "bin/tool.sh" { + t.Errorf("symlink entry wrong: %+v", link) + } + for name, hdr := range headers { + if hdr.Uid != 0 || hdr.Gid != 0 || hdr.Uname != "" || hdr.Gname != "" { + t.Errorf("ownership leaked into entry %q", name) + } + } +} + +func TestSaveHomeSkipsIrregularFiles(t *testing.T) { + home := buildFixtureHome(t) + fifo := filepath.Join(home, "pipe") + if err := mkfifo(fifo); err != nil { + t.Skipf("cannot create fifo: %v", err) + } + var buf bytes.Buffer + if err := SaveHome(context.Background(), &buf, home, nil); err != nil { + t.Fatalf("SaveHome with fifo present: %v", err) + } + headers, _ := decodeArchive(t, buf.Bytes()) + if _, ok := headers["pipe"]; ok { + t.Error("fifo must be skipped, not archived") + } +} + +func TestSaveHomeCancelledContext(t *testing.T) { + home := buildFixtureHome(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + var buf bytes.Buffer + if err := SaveHome(ctx, &buf, home, nil); err == nil { + t.Fatal("expected error from cancelled context") + } +} + +func TestSaveHomeUnreadableFile(t *testing.T) { + if os.Geteuid() == 0 { + t.Skip("running as root: permission checks are bypassed") + } + home := buildFixtureHome(t) + locked := filepath.Join(home, "locked.txt") + if err := os.WriteFile(locked, []byte("secret"), 0o000); err != nil { + t.Fatal(err) + } + var buf bytes.Buffer + if err := SaveHome(context.Background(), &buf, home, nil); err == nil { + t.Fatal("expected error for unreadable file") + } +} + +// The runtime state dir is not logically part of a Home Snapshot, so no +// configuration may put it into one. Excludes add to that rule; they cannot +// subtract from it. +func TestSaveHomeAlwaysSkipsDefaultExcludes(t *testing.T) { + for name, excludes := range map[string][]string{ + "nil excludes": nil, + "empty excludes": {}, + "unrelated excludes": {".cache"}, + // A caller must not be able to negotiate a default back in, whether + // by re-including it outright or by anchoring the attempt. + "re-inclusion": {"!" + RuntimeStateDir}, + "anchored re-inclusion": {"!/" + RuntimeStateDir + "/**"}, + } { + t.Run(name, func(t *testing.T) { + home := buildFixtureHome(t) + var buf bytes.Buffer + if err := SaveHome(context.Background(), &buf, home, excludes); err != nil { + t.Fatalf("SaveHome: %v", err) + } + headers, _ := decodeArchive(t, buf.Bytes()) + for entry := range headers { + for _, dir := range defaultExcludes { + if entry == dir+"/" || strings.HasPrefix(entry, dir+"/") { + t.Errorf("%s entry %q archived", dir, entry) + } + } + } + }) + } +} + +// A nested path that merely repeats a default-exclude name is ordinary Home +// content and must survive. +func TestSaveHomeSkipsOnlyExactDefaultExcludePaths(t *testing.T) { + for _, dir := range defaultExcludes { + t.Run(dir, func(t *testing.T) { + home := buildFixtureHome(t) + nested := filepath.Join(home, "bin", dir) + if err := os.MkdirAll(nested, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(nested, "keep.txt"), []byte("keep"), 0o644); err != nil { + t.Fatal(err) + } + var buf bytes.Buffer + if err := SaveHome(context.Background(), &buf, home, nil); err != nil { + t.Fatalf("SaveHome: %v", err) + } + headers, _ := decodeArchive(t, buf.Bytes()) + want := path.Join("bin", dir, "keep.txt") + if _, ok := headers[want]; !ok { + t.Errorf("%s dropped; only the exact %s path is excluded", want, dir) + } + }) + } +} diff --git a/dify-agent-runtime/internal/snapshot/snapshot.go b/dify-agent-runtime/internal/snapshot/snapshot.go new file mode 100644 index 00000000000..34a1dc321c9 --- /dev/null +++ b/dify-agent-runtime/internal/snapshot/snapshot.go @@ -0,0 +1,8 @@ +// Package snapshot implements single-pass streaming save/restore of the +// runtime's Home directory as tar+zstd archives. +// +// SIZE CONTRACT: this package imposes NO size limits. Consumers own size +// policy and must bound the streams they read or send in their own logic — +// count bytes while reading a save stream and abort at their cap; bound what +// they send to restore. In Dify EE the sandbox-gateway enforces this. +package snapshot diff --git a/dify-agent/src/dify_agent/runtime_backend/__init__.py b/dify-agent/src/dify_agent/runtime_backend/__init__.py index bca6748166f..c274a93fb53 100644 --- a/dify-agent/src/dify_agent/runtime_backend/__init__.py +++ b/dify-agent/src/dify_agent/runtime_backend/__init__.py @@ -7,6 +7,7 @@ from .errors import ( BindingLostError, HomeSnapshotCreateError, HomeSnapshotNotFoundError, + HomeSnapshotTooLargeError, RuntimeBackendError, SharedWorkspaceUnsupportedError, WorkspacePreservationUnsupportedError, @@ -37,6 +38,7 @@ __all__ = [ "HomeSnapshotCreateError", "HomeSnapshotCreateSpec", "HomeSnapshotNotFoundError", + "HomeSnapshotTooLargeError", "RuntimeBackendError", "RuntimeBackendProfile", "RuntimeLayout", diff --git a/dify-agent/src/dify_agent/runtime_backend/enterprise.py b/dify-agent/src/dify_agent/runtime_backend/enterprise.py index c3675c3f6d4..6d038c4fe16 100644 --- a/dify-agent/src/dify_agent/runtime_backend/enterprise.py +++ b/dify-agent/src/dify_agent/runtime_backend/enterprise.py @@ -1,15 +1,17 @@ """Enterprise Gateway adapter for the working-environment protocol. -The existing Gateway can allocate, reconnect to, and delete a sandbox, but it -does not expose immutable Home Snapshot operations. One physical sandbox owns -both the materialized Home and Workspace, so their cleanup is coupled. Runtime -access remains operation-local and is routed through the Gateway's shellctl -proxy. +The Gateway owns immutable Home Snapshot capture, restore, and deletion, in +addition to allocating, reconnecting to, and deleting a sandbox. Restore +arrives as a `homeSnapshotRef` parameter on sandbox creation rather than its +own call. One physical sandbox owns both the materialized Home and Workspace, +so their cleanup is coupled. Runtime access remains operation-local and is +routed through the Gateway's shellctl proxy. """ from __future__ import annotations from dataclasses import dataclass, field +from http import HTTPStatus import logging import shlex from typing import cast @@ -24,6 +26,8 @@ from dify_agent.runtime_backend.errors import ( BindingCreateError, BindingDestroyError, BindingLostError, + HomeSnapshotCreateError, + HomeSnapshotTooLargeError, SharedWorkspaceUnsupportedError, WorkspacePreservationUnsupportedError, ) @@ -43,22 +47,110 @@ from dify_agent.runtime_backend.shellctl import ( logger = logging.getLogger(__name__) +_GATEWAY_AUTH_HEADER = "X-Inner-Api-Key" -def _not_implemented() -> NotImplementedError: - return NotImplementedError("Enterprise Gateway does not implement immutable Home Snapshot operations") + +class _GatewayStatusError(RuntimeError): + """One failed Gateway control-plane call, with its reason when it sent one.""" + + def __init__(self, response: httpx.Response) -> None: + self.status_code = response.status_code + self.reason = "" + detail = response.text + try: + payload = response.json() + except ValueError: + payload = None + if isinstance(payload, dict): + reason = payload.get("reason") + message = payload.get("message") + if isinstance(reason, str): + self.reason = reason + if isinstance(message, str) and message: + detail = message + self.detail = detail + super().__init__(f"{self.status_code} {self.reason or 'gateway_error'}: {detail}") + + +async def _gateway_request( + *, + endpoint: str, + auth_token: str, + timeout: float, + method: str, + path: str, + json_body: dict[str, object] | None = None, + absent_status: int | None = None, +) -> dict[str, object] | None: + """Call one Gateway control-plane endpoint and decode its reply.""" + headers = {_GATEWAY_AUTH_HEADER: auth_token} if auth_token else {} + async with httpx.AsyncClient( + base_url=endpoint.rstrip("/"), + headers=headers, + timeout=httpx.Timeout(timeout), + ) as client: + response = await client.request(method, path, json=json_body) + if absent_status is not None and response.status_code == absent_status: + return None + if response.status_code >= 400: + raise _GatewayStatusError(response) + if not response.content: + return None + try: + payload = response.json() + except ValueError: + return None + return payload if isinstance(payload, dict) else None @dataclass(slots=True) class EnterpriseHomeSnapshotBackend: - """Reject Home Snapshot operations until the Gateway exposes immutable snapshots.""" + """Manage immutable Home Snapshots through the Gateway's snapshot endpoints.""" + + gateway_endpoint: str + auth_token: str + snapshot_timeout: float = 35.0 async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str: - del spec, source - raise _not_implemented() + """Capture the source lease's Home through the Gateway's snapshot endpoint.""" + if not isinstance(source, EnterpriseRuntimeLease): + raise HomeSnapshotCreateError("Enterprise Home Snapshot requires an Enterprise RuntimeLease") + try: + payload = await _gateway_request( + endpoint=self.gateway_endpoint, + auth_token=self.auth_token, + timeout=self.snapshot_timeout, + method="POST", + path=f"/v1/sandboxes/{quote(source.handle, safe='')}/home-snapshots", + json_body={ + "tenantId": spec.tenant_id, + "agentId": spec.agent_id, + "homeSnapshotId": spec.home_snapshot_id, + }, + ) + except _GatewayStatusError as exc: + if exc.status_code == HTTPStatus.REQUEST_ENTITY_TOO_LARGE: + raise HomeSnapshotTooLargeError(exc.detail) from exc + raise HomeSnapshotCreateError(str(exc)) from exc + except (httpx.TimeoutException, httpx.RequestError) as exc: + raise HomeSnapshotCreateError(str(exc)) from exc + snapshot_ref = payload.get("snapshotRef") if isinstance(payload, dict) else None + if not isinstance(snapshot_ref, str) or not snapshot_ref: + raise HomeSnapshotCreateError("Enterprise Gateway returned an invalid snapshot ref") + return snapshot_ref async def delete(self, snapshot_ref: str) -> None: - del snapshot_ref - raise _not_implemented() + """Delete one snapshot's artifacts.""" + try: + _ = await _gateway_request( + endpoint=self.gateway_endpoint, + auth_token=self.auth_token, + timeout=self.snapshot_timeout, + method="DELETE", + path=f"/v1/home-snapshots/{quote(snapshot_ref, safe='/')}", + ) + except (_GatewayStatusError, httpx.TimeoutException, httpx.RequestError) as exc: + raise BindingDestroyError(str(exc)) from exc @dataclass(slots=True) @@ -69,6 +161,7 @@ class EnterpriseExecutionBindingBackend: auth_token: str gateway_timeout: float = 30.0 proxy_timeout: float = 60.0 + snapshot_timeout: float = 35.0 layout: RuntimeLayout = field( default_factory=lambda: RuntimeLayout(home_dir="/home/dify", workspace_dir="/workspace") ) @@ -77,21 +170,21 @@ class EnterpriseExecutionBindingBackend: """Create a default Gateway sandbox and initialize its canonical layout.""" if spec.existing_workspace_ref is not None: raise SharedWorkspaceUnsupportedError("current Enterprise backend cannot attach to an existing Workspace") - if spec.home_snapshot_ref is not None: - raise BindingCreateError("current Enterprise backend cannot materialize an immutable Home Snapshot") sandbox_id: str | None = None data_plane: ShellctlRuntimeLease | None = None - headers = {"X-Inner-Api-Key": self.auth_token} if self.auth_token else {} try: - async with httpx.AsyncClient( - base_url=self.gateway_endpoint.rstrip("/"), - headers=headers, - timeout=httpx.Timeout(self.gateway_timeout), - ) as client: - response = await client.post("/v1/sandboxes", json={"tenantId": spec.tenant_id}) - _ = response.raise_for_status() - payload = response.json() + create_body: dict[str, object] = {"tenantId": spec.tenant_id} + if spec.home_snapshot_ref is not None: + create_body["homeSnapshotRef"] = spec.home_snapshot_ref + payload = await _gateway_request( + endpoint=self.gateway_endpoint, + auth_token=self.auth_token, + timeout=self.snapshot_timeout if spec.home_snapshot_ref is not None else self.gateway_timeout, + method="POST", + path="/v1/sandboxes", + json_body=create_body, + ) sandbox_id_value = payload.get("sandboxId") if isinstance(payload, dict) else None if not isinstance(sandbox_id_value, str) or not sandbox_id_value: raise BindingCreateError("Enterprise Gateway returned an invalid sandbox id") @@ -99,12 +192,16 @@ class EnterpriseExecutionBindingBackend: data_plane = await self._create_data_plane(sandbox_id) result = await run_shellctl_control_command( - ShellctlCommands(client=data_plane.client), + ShellctlCommands( + client=data_plane.client, + home_dir=self.layout.home_dir, + workspace_dir=self.layout.workspace_dir, + ), "\n".join( [ "set -eu", + "exec 2>&1", f"mkdir -p {shlex.quote(self.layout.home_dir)}", - f"mkdir -p {shlex.quote(self.layout.workspace_dir)}", f"find {shlex.quote(self.layout.workspace_dir)} -mindepth 1 -maxdepth 1 -exec rm -rf -- {{}} +", f"chmod 700 {shlex.quote(self.layout.home_dir)} {shlex.quote(self.layout.workspace_dir)}", ] @@ -179,21 +276,18 @@ class EnterpriseExecutionBindingBackend: try: await self._delete_sandbox(spec.binding_ref) - except (httpx.TimeoutException, httpx.RequestError, httpx.HTTPStatusError) as exc: + except (_GatewayStatusError, httpx.TimeoutException, httpx.RequestError) as exc: raise BindingDestroyError(str(exc)) from exc async def _delete_sandbox(self, sandbox_id: str) -> None: - headers = {"X-Inner-Api-Key": self.auth_token} if self.auth_token else {} - encoded_sandbox_id = quote(sandbox_id, safe="") - async with httpx.AsyncClient( - base_url=self.gateway_endpoint.rstrip("/"), - headers=headers, - timeout=httpx.Timeout(self.gateway_timeout), - ) as client: - response = await client.delete(f"/v1/sandboxes/{encoded_sandbox_id}") - if response.status_code == 404: - return - _ = response.raise_for_status() + _ = await _gateway_request( + endpoint=self.gateway_endpoint, + auth_token=self.auth_token, + timeout=self.gateway_timeout, + method="DELETE", + path=f"/v1/sandboxes/{quote(sandbox_id, safe='')}", + absent_status=404, + ) async def _delete_sandbox_best_effort(self, sandbox_id: str) -> None: try: diff --git a/dify-agent/src/dify_agent/runtime_backend/errors.py b/dify-agent/src/dify_agent/runtime_backend/errors.py index 260844b976f..627b4fb8851 100644 --- a/dify-agent/src/dify_agent/runtime_backend/errors.py +++ b/dify-agent/src/dify_agent/runtime_backend/errors.py @@ -13,6 +13,10 @@ class HomeSnapshotNotFoundError(RuntimeBackendError): pass +class HomeSnapshotTooLargeError(RuntimeBackendError): + pass + + class BindingCreateError(RuntimeBackendError): pass @@ -48,6 +52,7 @@ __all__ = [ "BindingLostError", "HomeSnapshotCreateError", "HomeSnapshotNotFoundError", + "HomeSnapshotTooLargeError", "RuntimeBackendError", "SharedWorkspaceUnsupportedError", "WorkspacePreservationUnsupportedError", diff --git a/dify-agent/src/dify_agent/runtime_backend/profile.py b/dify-agent/src/dify_agent/runtime_backend/profile.py index 583afe54d4e..411f4d3c3e5 100644 --- a/dify-agent/src/dify_agent/runtime_backend/profile.py +++ b/dify-agent/src/dify_agent/runtime_backend/profile.py @@ -52,6 +52,7 @@ class RuntimeBackendSettings(BaseSettings): enterprise_sandbox_gateway_auth_token: str | None = None enterprise_sandbox_gateway_timeout: float = Field(default=30.0, gt=0) enterprise_sandbox_proxy_timeout: float = Field(default=60.0, gt=0) + enterprise_sandbox_snapshot_timeout: float = Field(default=35.0, gt=0) e2b_api_key: str | None = None e2b_template: str = DEFAULT_E2B_TEMPLATE @@ -127,12 +128,17 @@ def create_runtime_backend_profile(settings: RuntimeBackendSettings) -> RuntimeB endpoint = settings.enterprise_sandbox_gateway_endpoint or "" token = settings.enterprise_sandbox_gateway_auth_token or "" return RuntimeBackendProfile( - home_snapshots=EnterpriseHomeSnapshotBackend(), + home_snapshots=EnterpriseHomeSnapshotBackend( + gateway_endpoint=endpoint, + auth_token=token, + snapshot_timeout=settings.enterprise_sandbox_snapshot_timeout, + ), execution_bindings=EnterpriseExecutionBindingBackend( gateway_endpoint=endpoint, auth_token=token, gateway_timeout=settings.enterprise_sandbox_gateway_timeout, proxy_timeout=settings.enterprise_sandbox_proxy_timeout, + snapshot_timeout=settings.enterprise_sandbox_snapshot_timeout, ), ) case "e2b": diff --git a/dify-agent/src/dify_agent/server/home_snapshots.py b/dify-agent/src/dify_agent/server/home_snapshots.py index 27ac62fc44d..108e52e45db 100644 --- a/dify-agent/src/dify_agent/server/home_snapshots.py +++ b/dify-agent/src/dify_agent/server/home_snapshots.py @@ -15,6 +15,7 @@ from dify_agent.runtime_backend import ( HomeSnapshotCreateError, HomeSnapshotCreateSpec, HomeSnapshotNotFoundError, + HomeSnapshotTooLargeError, ) from dify_agent.runtime_backend.leases import open_runtime_lease @@ -52,6 +53,8 @@ class HomeSnapshotService: ) except BindingLostError as exc: raise HomeSnapshotServiceError("binding_lost", str(exc), status_code=404) from exc + except HomeSnapshotTooLargeError as exc: + raise HomeSnapshotServiceError("home_snapshot_too_large", str(exc), status_code=413) from exc except (BindingAcquireError, HomeSnapshotCreateError) as exc: raise HomeSnapshotServiceError("home_snapshot_create_failed", str(exc), status_code=502) from exc return HomeSnapshotResponse(snapshot_ref=snapshot_ref) diff --git a/dify-agent/src/dify_agent/server/settings.py b/dify-agent/src/dify_agent/server/settings.py index 1601d3d019e..3d27e4439b7 100644 --- a/dify-agent/src/dify_agent/server/settings.py +++ b/dify-agent/src/dify_agent/server/settings.py @@ -63,6 +63,7 @@ class ServerSettings(BaseSettings): enterprise_sandbox_gateway_auth_token: str | None = None enterprise_sandbox_gateway_timeout: float = Field(default=30.0, gt=0) enterprise_sandbox_proxy_timeout: float = Field(default=60.0, gt=0) + enterprise_sandbox_snapshot_timeout: float = Field(default=35.0, gt=0) e2b_api_key: str | None = None e2b_template: str = "difys-default-team/dify-agent-local-sandbox" e2b_active_timeout_seconds: int = Field( @@ -207,6 +208,7 @@ class ServerSettings(BaseSettings): enterprise_sandbox_gateway_auth_token=self.enterprise_sandbox_gateway_auth_token, enterprise_sandbox_gateway_timeout=self.enterprise_sandbox_gateway_timeout, enterprise_sandbox_proxy_timeout=self.enterprise_sandbox_proxy_timeout, + enterprise_sandbox_snapshot_timeout=self.enterprise_sandbox_snapshot_timeout, e2b_api_key=self.e2b_api_key, e2b_template=self.e2b_template, e2b_active_timeout_seconds=self.e2b_active_timeout_seconds, diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_enterprise_backend.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_enterprise_backend.py index 049d6fb65e6..41588488044 100644 --- a/dify-agent/tests/local/dify_agent/runtime_backend/test_enterprise_backend.py +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_enterprise_backend.py @@ -14,7 +14,9 @@ from dify_agent.runtime_backend import ( BindingLostError, ExecutionBindingCreateSpec, ExecutionBindingDestroySpec, + HomeSnapshotCreateError, HomeSnapshotCreateSpec, + HomeSnapshotTooLargeError, RuntimeLease, SharedWorkspaceUnsupportedError, WorkspacePreservationUnsupportedError, @@ -282,6 +284,62 @@ async def test_enterprise_destroy_propagates_gateway_failure( assert clients[0].is_closed +@pytest.mark.anyio +async def test_enterprise_gateway_errors_carry_the_kratos_reason( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 404, + json={ + "code": 404, + "reason": "warm_pool_not_found", + "message": "sandbox warm pool not found", + "metadata": {}, + }, + ) + + _ = _mock_http(monkeypatch, handler) + backend = EnterpriseExecutionBindingBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + + with pytest.raises(BindingCreateError) as excinfo: + _ = await backend.create_binding( + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref=None, + ) + ) + + assert "warm_pool_not_found" in str(excinfo.value) + assert "404" in str(excinfo.value) + + +@pytest.mark.anyio +async def test_enterprise_gateway_errors_survive_a_non_json_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(502, text="bad gateway") + + _ = _mock_http(monkeypatch, handler) + backend = EnterpriseExecutionBindingBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + + with pytest.raises(BindingDestroyError) as excinfo: + await backend.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref="sandbox-1", + workspace_ref="sandbox-1", + destroy_workspace=True, + ) + ) + + assert "502" in str(excinfo.value) + + @pytest.mark.anyio async def test_enterprise_default_binding_creates_gateway_sandbox_and_layout( monkeypatch: pytest.MonkeyPatch, @@ -298,14 +356,19 @@ async def test_enterprise_default_binding_creates_gateway_sandbox_and_layout( payload = cast(dict[str, object], json.loads(request.content)) script = payload["script"] assert isinstance(script, str) + assert payload["cwd"] == "/workspace" assert "mkdir -p /home/dify" in script - assert "mkdir -p /workspace" in script assert "find /workspace -mindepth 1 -maxdepth 1 -exec rm -rf -- {} +" in script return _job_response() return httpx.Response(200, json={"job_id": "job-1"}) clients = _mock_http(monkeypatch, handler) - backend = EnterpriseExecutionBindingBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + backend = EnterpriseExecutionBindingBackend( + gateway_endpoint="http://gateway.example", + auth_token="secret", + gateway_timeout=30, + snapshot_timeout=120, + ) allocation = await backend.create_binding( ExecutionBindingCreateSpec( @@ -321,27 +384,17 @@ async def test_enterprise_default_binding_creates_gateway_sandbox_and_layout( assert allocation.binding_ref == allocation.workspace_ref == "sandbox-1" assert requests[0].headers["X-Inner-Api-Key"] == "secret" assert all(client.is_closed for client in clients) + assert clients[0].timeout.read == 30 @pytest.mark.anyio -async def test_enterprise_binding_rejects_snapshot_and_shared_workspace_before_gateway_call( +async def test_enterprise_binding_rejects_shared_workspace_before_gateway_call( monkeypatch: pytest.MonkeyPatch, ) -> None: requests: list[httpx.Request] = [] _ = _mock_http(monkeypatch, lambda request: requests.append(request) or httpx.Response(500)) backend = EnterpriseExecutionBindingBackend(gateway_endpoint="http://gateway.example", auth_token="secret") - with pytest.raises(BindingCreateError, match="immutable Home Snapshot"): - await backend.create_binding( - ExecutionBindingCreateSpec( - tenant_id="tenant-1", - agent_id="agent-1", - binding_id="binding-1", - workspace_id="workspace-1", - existing_workspace_ref=None, - home_snapshot_ref="snapshot-1", - ) - ) with pytest.raises(SharedWorkspaceUnsupportedError): await backend.create_binding( ExecutionBindingCreateSpec( @@ -357,6 +410,83 @@ async def test_enterprise_binding_rejects_snapshot_and_shared_workspace_before_g assert requests == [] +@pytest.mark.anyio +async def test_enterprise_binding_sends_the_snapshot_ref_and_still_clears_workspace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + scripts: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/v1/sandboxes": + assert json.loads(request.content) == { + "tenantId": "tenant-1", + "homeSnapshotRef": "tenant-1/agent-1/home-2", + } + return httpx.Response(201, json={"sandboxId": "sandbox-1"}) + if request.url.path == "/proxy/v1/jobs/run": + payload = cast(dict[str, object], json.loads(request.content)) + script = payload["script"] + assert isinstance(script, str) + scripts.append(script) + return _job_response() + return httpx.Response(200, json={"job_id": "job-1"}) + + clients = _mock_http(monkeypatch, handler) + backend = EnterpriseExecutionBindingBackend( + gateway_endpoint="http://gateway.example", + auth_token="secret", + snapshot_timeout=120, + ) + + allocation = await backend.create_binding( + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref="tenant-1/agent-1/home-2", + ) + ) + + assert allocation.binding_ref == "sandbox-1" + assert any("find /workspace -mindepth 1 -maxdepth 1 -exec rm -rf -- {} +" in script for script in scripts) + assert clients[0].timeout.read == 120 + + +@pytest.mark.anyio +async def test_enterprise_binding_does_not_delete_a_sandbox_it_never_created( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + if request.url.path == "/v1/sandboxes" and request.method == "POST": + return httpx.Response( + 404, + json={"code": 404, "reason": "snapshot_not_found", "message": "home snapshot missing"}, + ) + return httpx.Response(204) + + _ = _mock_http(monkeypatch, handler) + backend = EnterpriseExecutionBindingBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + + with pytest.raises(BindingCreateError, match="snapshot_not_found"): + await backend.create_binding( + ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id="binding-1", + workspace_id="workspace-1", + existing_workspace_ref=None, + home_snapshot_ref="tenant-1/agent-1/home-2", + ) + ) + + assert not any(request.method == "DELETE" for request in requests) + + @pytest.mark.anyio async def test_enterprise_binding_create_deletes_new_sandbox_when_layout_setup_fails( monkeypatch: pytest.MonkeyPatch, @@ -392,13 +522,175 @@ async def test_enterprise_binding_create_deletes_new_sandbox_when_layout_setup_f @pytest.mark.anyio -async def test_enterprise_home_snapshots_remain_explicitly_not_implemented() -> None: - snapshots = EnterpriseHomeSnapshotBackend() +async def test_enterprise_snapshot_create_posts_to_the_lease_sandbox( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[httpx.Request] = [] - with pytest.raises(NotImplementedError, match="immutable Home Snapshot"): + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + if request.url.path == "/v1/sandboxes/sandbox-1/home-snapshots": + assert json.loads(request.content) == { + "tenantId": "tenant-1", + "agentId": "agent-1", + "homeSnapshotId": "home-2", + } + assert request.headers["X-Inner-Api-Key"] == "secret" + return httpx.Response(200, json={"snapshotRef": "tenant-1/agent-1/home-2"}) + return _job_response() + + clients = _mock_http(monkeypatch, handler) + bindings = EnterpriseExecutionBindingBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + lease = await bindings.acquire("sandbox-1") + snapshots = EnterpriseHomeSnapshotBackend( + gateway_endpoint="http://gateway.example", + auth_token="secret", + snapshot_timeout=120, + ) + + snapshot_ref = await snapshots.create_from_runtime( + spec=HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-2"), + source=lease, + ) + + assert snapshot_ref == "tenant-1/agent-1/home-2" + assert any(request.url.path == "/v1/sandboxes/sandbox-1/home-snapshots" for request in requests) + assert clients[-1].timeout.read == 120 + await bindings.release(lease) + + +@pytest.mark.anyio +async def test_enterprise_snapshot_create_rejects_a_foreign_lease() -> None: + snapshots = EnterpriseHomeSnapshotBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + + with pytest.raises(HomeSnapshotCreateError): _ = await snapshots.create_from_runtime( spec=HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-2"), source=cast(RuntimeLease, object()), ) - with pytest.raises(NotImplementedError, match="immutable Home Snapshot"): - await snapshots.delete("snapshot-1") + + +@pytest.mark.anyio +async def test_enterprise_snapshot_create_maps_the_size_limit_to_too_large( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/home-snapshots"): + return httpx.Response( + 413, + json={ + "code": 413, + "reason": "snapshot_size_exceeded", + "message": "home snapshot exceeds the configured limit of 67108864 bytes", + }, + ) + return _job_response() + + _ = _mock_http(monkeypatch, handler) + bindings = EnterpriseExecutionBindingBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + lease = await bindings.acquire("sandbox-1") + snapshots = EnterpriseHomeSnapshotBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + + with pytest.raises(HomeSnapshotTooLargeError, match="exceeds the configured limit"): + _ = await snapshots.create_from_runtime( + spec=HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-2"), + source=lease, + ) + await bindings.release(lease) + + +@pytest.mark.anyio +async def test_enterprise_snapshot_create_maps_other_gateway_failures( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/home-snapshots"): + return httpx.Response( + 500, + json={ + "code": 500, + "reason": "snapshot_save_failed", + "message": "home snapshot save failed: object store unavailable", + }, + ) + return _job_response() + + _ = _mock_http(monkeypatch, handler) + bindings = EnterpriseExecutionBindingBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + lease = await bindings.acquire("sandbox-1") + snapshots = EnterpriseHomeSnapshotBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + + with pytest.raises(HomeSnapshotCreateError, match="snapshot_save_failed"): + _ = await snapshots.create_from_runtime( + spec=HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-2"), + source=lease, + ) + await bindings.release(lease) + + +@pytest.mark.anyio +async def test_enterprise_snapshot_create_rejects_a_reply_without_a_ref( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/home-snapshots"): + return httpx.Response(200, json={"snapshotRef": ""}) + return _job_response() + + _ = _mock_http(monkeypatch, handler) + bindings = EnterpriseExecutionBindingBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + lease = await bindings.acquire("sandbox-1") + snapshots = EnterpriseHomeSnapshotBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + + with pytest.raises(HomeSnapshotCreateError, match="invalid snapshot ref"): + _ = await snapshots.create_from_runtime( + spec=HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-2"), + source=lease, + ) + await bindings.release(lease) + + +@pytest.mark.anyio +async def test_enterprise_snapshot_delete_keeps_ref_slashes_and_escapes_the_rest( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json={}) + + _ = _mock_http(monkeypatch, handler) + snapshots = EnterpriseHomeSnapshotBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + + await snapshots.delete("tenant-1/agent-1/home-2") + await snapshots.delete("tenant-1/agent-1/home 2?x=1") + + assert [request.method for request in requests] == ["DELETE", "DELETE"] + assert requests[0].url.raw_path == b"/v1/home-snapshots/tenant-1/agent-1/home-2" + assert requests[1].url.raw_path == b"/v1/home-snapshots/tenant-1/agent-1/home%202%3Fx%3D1" + assert requests[0].headers["X-Inner-Api-Key"] == "secret" + + +@pytest.mark.anyio +@pytest.mark.parametrize( + ("status_code", "reason", "message"), + [ + (404, "route_not_found", "no matching route"), + (500, "snapshot_delete_failed", "object store unreachable"), + ], +) +async def test_enterprise_snapshot_delete_propagates_gateway_failures( + monkeypatch: pytest.MonkeyPatch, + status_code: int, + reason: str, + message: str, +) -> None: + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response(status_code, json={"code": status_code, "reason": reason, "message": message}) + + _ = _mock_http(monkeypatch, handler) + snapshots = EnterpriseHomeSnapshotBackend(gateway_endpoint="http://gateway.example", auth_token="secret") + + with pytest.raises(BindingDestroyError, match=reason): + await snapshots.delete("tenant-1/agent-1/home-2") diff --git a/dify-agent/tests/local/dify_agent/server/test_settings.py b/dify-agent/tests/local/dify_agent/server/test_settings.py index 76cdb54bec4..b94dc033122 100644 --- a/dify-agent/tests/local/dify_agent/server/test_settings.py +++ b/dify-agent/tests/local/dify_agent/server/test_settings.py @@ -11,7 +11,7 @@ from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubTokenCodec from dify_agent.server.settings import ServerSettings from dify_agent.runtime.runner import DEFAULT_AGENT_RUN_TIMEOUT_SECONDS from dify_agent.runtime_backend.e2b import E2B_MAX_ACTIVE_TIMEOUT_SECONDS, E2BExecutionBindingBackend -from dify_agent.runtime_backend.enterprise import EnterpriseExecutionBindingBackend +from dify_agent.runtime_backend.enterprise import EnterpriseExecutionBindingBackend, EnterpriseHomeSnapshotBackend from dify_agent.runtime_backend.local import LocalExecutionBindingBackend, LocalHomeSnapshotBackend @@ -333,6 +333,7 @@ def test_build_runtime_backend_profile_returns_enterprise_drivers_when_selected( enterprise_sandbox_gateway_auth_token="gateway-secret", enterprise_sandbox_gateway_timeout=45, enterprise_sandbox_proxy_timeout=90, + enterprise_sandbox_snapshot_timeout=120, ) profile = settings.build_runtime_backend_profile() @@ -343,6 +344,20 @@ def test_build_runtime_backend_profile_returns_enterprise_drivers_when_selected( assert profile.execution_bindings.auth_token == "gateway-secret" assert profile.execution_bindings.gateway_timeout == 45 assert profile.execution_bindings.proxy_timeout == 90 + assert isinstance(profile.home_snapshots, EnterpriseHomeSnapshotBackend) + assert profile.home_snapshots.gateway_endpoint == "https://gateway.example" + assert profile.home_snapshots.auth_token == "gateway-secret" + assert profile.home_snapshots.snapshot_timeout == 120 + assert profile.execution_bindings.snapshot_timeout == 120 + + +def test_enterprise_snapshot_timeout_defaults_above_the_gateway_budget() -> None: + settings = ServerSettings( + runtime_backend="enterprise", + enterprise_sandbox_gateway_endpoint="https://gateway.example", + ) + + assert settings.enterprise_sandbox_snapshot_timeout > settings.enterprise_sandbox_gateway_timeout def test_build_runtime_backend_profile_passes_e2b_active_timeout() -> None: diff --git a/docker/.env.example b/docker/.env.example index b614314b91a..f59b913c1da 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -260,6 +260,9 @@ AGENT_BACKEND_BASE_URL=http://agent_backend:5050 DIFY_AGENT_API_TOKEN=dify-agent-run-token-for-dev-only AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS=30 AGENT_BACKEND_STREAM_MAX_RECONNECTS=3 +# Client timeout for Agent backend calls that may carry a Home Snapshot transfer. +# Must stay above DIFY_AGENT_ENTERPRISE_SANDBOX_SNAPSHOT_TIMEOUT on the Agent backend. +AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS=45 # Leave empty to derive from REDIS_PASSWORD. DIFY_AGENT_REDIS_URL= DIFY_AGENT_REDIS_PREFIX=dify-agent diff --git a/docker/docker-compose-template.yaml b/docker/docker-compose-template.yaml index a826cb4e5e7..a0fb52b6ab2 100644 --- a/docker/docker-compose-template.yaml +++ b/docker/docker-compose-template.yaml @@ -235,6 +235,7 @@ services: AGENT_BACKEND_API_TOKEN: ${DIFY_AGENT_API_TOKEN:-dify-agent-run-token-for-dev-only} AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS: ${AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS:-30} AGENT_BACKEND_STREAM_MAX_RECONNECTS: ${AGENT_BACKEND_STREAM_MAX_RECONNECTS:-3} + AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS: ${AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS:-45} depends_on: init_permissions: condition: service_completed_successfully @@ -313,6 +314,7 @@ services: AGENT_BACKEND_API_TOKEN: ${DIFY_AGENT_API_TOKEN:-dify-agent-run-token-for-dev-only} AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS: ${AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS:-30} AGENT_BACKEND_STREAM_MAX_RECONNECTS: ${AGENT_BACKEND_STREAM_MAX_RECONNECTS:-3} + AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS: ${AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS:-45} depends_on: init_permissions: condition: service_completed_successfully diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index d4f2154307e..49582227d95 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -241,6 +241,7 @@ services: AGENT_BACKEND_API_TOKEN: ${DIFY_AGENT_API_TOKEN:-dify-agent-run-token-for-dev-only} AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS: ${AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS:-30} AGENT_BACKEND_STREAM_MAX_RECONNECTS: ${AGENT_BACKEND_STREAM_MAX_RECONNECTS:-3} + AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS: ${AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS:-45} depends_on: init_permissions: condition: service_completed_successfully @@ -319,6 +320,7 @@ services: AGENT_BACKEND_API_TOKEN: ${DIFY_AGENT_API_TOKEN:-dify-agent-run-token-for-dev-only} AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS: ${AGENT_BACKEND_STREAM_READ_TIMEOUT_SECONDS:-30} AGENT_BACKEND_STREAM_MAX_RECONNECTS: ${AGENT_BACKEND_STREAM_MAX_RECONNECTS:-3} + AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS: ${AGENT_BACKEND_HOME_SNAPSHOT_TIMEOUT_SECONDS:-45} depends_on: init_permissions: condition: service_completed_successfully