diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml index 4e70609b7e1745..b9b63ef5b0b1ef 100644 --- a/.github/workflows/docker-build.yml +++ b/.github/workflows/docker-build.yml @@ -14,10 +14,13 @@ on: - dify-agent/src/** - web/Dockerfile - dify-agent-runtime/docker/Dockerfile + - dify-agent-runtime/docker/Dockerfile.rust + - dify-agent-runtime/.dockerignore - dify-agent-runtime/go.mod - dify-agent-runtime/go.sum - dify-agent-runtime/cmd/** - dify-agent-runtime/internal/** + - dify-agent-runtime/rust/** concurrency: group: docker-build-${{ github.head_ref || github.run_id }} @@ -63,6 +66,16 @@ jobs: runs_on: depot-ubuntu-24.04-4 context: '{{defaultContext}}:dify-agent-runtime' file: 'docker/Dockerfile' + - service_name: 'local-sandbox-rust-amd64' + platform: linux/amd64 + runs_on: depot-ubuntu-24.04-4 + context: '{{defaultContext}}:dify-agent-runtime' + file: 'docker/Dockerfile.rust' + - service_name: 'local-sandbox-rust-arm64' + platform: linux/arm64 + runs_on: depot-ubuntu-24.04-4 + context: '{{defaultContext}}:dify-agent-runtime' + file: 'docker/Dockerfile.rust' steps: - name: Set up Depot CLI uses: depot/setup-action@15c09a5f77a0840ad4bce955686522a257853461 # v1.7.1 @@ -94,6 +107,9 @@ jobs: - service_name: 'local-sandbox-amd64' context: '{{defaultContext}}:dify-agent-runtime' file: 'docker/Dockerfile' + - service_name: 'local-sandbox-rust-amd64' + context: '{{defaultContext}}:dify-agent-runtime' + file: 'docker/Dockerfile.rust' steps: - name: Set up Docker Buildx uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0 diff --git a/.github/workflows/main-ci.yml b/.github/workflows/main-ci.yml index 72b1f496c71491..c705f50a99c685 100644 --- a/.github/workflows/main-ci.yml +++ b/.github/workflows/main-ci.yml @@ -135,7 +135,21 @@ jobs: - 'api/pyproject.toml' sandbox-runtime: - 'dify-agent-runtime/**' + - 'dify-agent/src/dify_agent/runtime_backend/**' + - 'dify-agent/src/dify_agent/server/settings.py' + - 'dify-agent/tests/local/dify_agent/runtime_backend/**' + - 'dify-agent/tests/local/dify_agent/server/test_settings.py' + - 'dify-agent/tests/integration/dify_agent/runtime_backend/**' + - 'dify-agent/pyproject.toml' + - 'dify-agent/uv.lock' + - 'benchmarks/**' + - 'docker/.env.example' + - 'docker/envs/core-services/dify-agent.env.example' + - 'docker/docker-compose.yaml' + - 'docker/docker-compose-template.yaml' + - 'docker/docker-compose.rust-runtime.yaml' - '.github/workflows/sandbox-runtime-tests.yml' + - '.github/workflows/main-ci.yml' migration: - 'api/migrations/**' - 'api/.env.example' diff --git a/.github/workflows/sandbox-runtime-tests.yml b/.github/workflows/sandbox-runtime-tests.yml index 45e26f06860464..39e86f0703e9a3 100644 --- a/.github/workflows/sandbox-runtime-tests.yml +++ b/.github/workflows/sandbox-runtime-tests.yml @@ -35,6 +35,50 @@ jobs: - name: Run unit tests run: go test -race -count=1 ./... + - name: Validate benchmark harness + working-directory: . + run: | + python3 -m compileall -q \ + benchmarks/run_container_benchmarks.py \ + benchmarks/test_container_benchmark.py + PYTHONPATH=benchmarks python3 -m unittest discover \ + -s benchmarks -p 'test_container_benchmark.py' -v + + - name: Setup UV and Python + uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + with: + enable-cache: true + python-version: '3.12' + cache-dependency-glob: dify-agent/uv.lock + + - name: Install agent rollout test dependencies + working-directory: dify-agent + run: uv sync --extra server --dev + + - name: Validate safe Go/Rust rollout routing + working-directory: dify-agent + env: + PYTHONPATH: src + run: | + uv run ruff check \ + src/dify_agent/runtime_backend/local_rollout.py \ + src/dify_agent/runtime_backend/profile.py \ + src/dify_agent/server/settings.py \ + tests/local/dify_agent/runtime_backend/test_local_rollout.py \ + tests/local/dify_agent/runtime_backend/test_profile.py \ + tests/local/dify_agent/server/test_settings.py + uv run pytest -q \ + tests/local/dify_agent/runtime_backend/test_local_rollout.py \ + tests/local/dify_agent/runtime_backend/test_profile.py \ + tests/local/dify_agent/server/test_settings.py + uv run basedpyright --level error \ + src/dify_agent/runtime_backend/local_rollout.py \ + src/dify_agent/runtime_backend/profile.py \ + src/dify_agent/server/settings.py \ + tests/local/dify_agent/runtime_backend/test_local_rollout.py \ + tests/local/dify_agent/runtime_backend/test_profile.py \ + tests/local/dify_agent/server/test_settings.py + sandbox-runtime-lint: name: Sandbox Runtime Lint runs-on: depot-ubuntu-24.04 @@ -62,6 +106,36 @@ jobs: working-directory: dify-agent-runtime version: latest + sandbox-runtime-rust: + name: Sandbox Runtime Rust Tests + runs-on: depot-ubuntu-24.04 + defaults: + run: + shell: bash + working-directory: dify-agent-runtime + + steps: + - name: Checkout code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + fetch-depth: 0 + persist-credentials: false + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@4360b52568e2003a75bf9bc1d59f33a8e3fc893c # stable + with: + toolchain: 1.95.0 + components: rustfmt, clippy + + - name: Check formatting + run: cargo fmt --manifest-path rust/Cargo.toml --all -- --check + + - name: Run Clippy + run: cargo clippy --locked --manifest-path rust/Cargo.toml --all-targets -- -D warnings + + - name: Run unit tests + run: cargo test --locked --manifest-path rust/Cargo.toml + sandbox-runtime-integration: name: Sandbox Runtime Integration Tests runs-on: depot-ubuntu-24.04 @@ -83,12 +157,26 @@ jobs: go-version-file: dify-agent-runtime/go.mod cache-dependency-path: dify-agent-runtime/go.sum + - name: Setup UV and Python + uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + with: + enable-cache: true + python-version: '3.12' + cache-dependency-glob: dify-agent/uv.lock + + - name: Install agent rollout test dependencies + working-directory: dify-agent + run: uv sync --extra server --dev + - name: Build and start runtime run: make integration-up - name: Run integration tests run: make integration-test + - name: Run real Go/Rust rollout integration tests + run: make integration-test-rollout + - name: Dump container logs on failure if: failure() run: make integration-logs || true diff --git a/benchmarks/run_container_benchmarks.py b/benchmarks/run_container_benchmarks.py new file mode 100644 index 00000000000000..ff40e4abd55020 --- /dev/null +++ b/benchmarks/run_container_benchmarks.py @@ -0,0 +1,666 @@ +#!/usr/bin/env python3 +"""Paired Linux-container benchmark for the Go and Rust shellctl runtimes.""" + +from __future__ import annotations + +import argparse +import hashlib +import http.client +import json +import math +import os +import platform +import random +import statistics +import subprocess +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any + +IMAGES = { + "go": "dify-agent-runtime:test", + "rust": "dify-agent-runtime-rust:test", +} +PROCESS_RSS_EXCLUDED_COMMANDS = {"ps"} +REPO = Path(__file__).resolve().parents[1] +RUNTIME = REPO / "dify-agent-runtime" + + +class ContainerBenchmarkError(RuntimeError): + pass + + +def _sample_percentile(values: list[float], probability: float) -> float: + ordered = sorted(values) + if not ordered: + return 0.0 + index = min(len(ordered) - 1, math.ceil((len(ordered) - 1) * probability)) + return ordered[index] + + +def stats(values: list[float], *, unit: str = "ms") -> dict[str, Any]: + if not values: + return {"samples": 0, "unit": unit} + total = sum(values) + seconds = total / 1000 if unit == "ms" else total + return { + "samples": len(values), + "unit": unit, + "avg": total / len(values), + "median": statistics.median(values), + "p50": _sample_percentile(values, 0.50), + "p95": _sample_percentile(values, 0.95), + "p99": _sample_percentile(values, 0.99), + "min": min(values), + "max": max(values), + "throughput_per_sec": len(values) / seconds if seconds else 0.0, + "raw_values": values, + } + + +def _paired_percentile(values: list[float], probability: float) -> float: + ordered = sorted(values) + if not ordered: + return 0.0 + return ordered[round((len(ordered) - 1) * probability)] + + +def describe(values: list[float]) -> dict[str, Any]: + return { + "samples": len(values), + "mean": statistics.fmean(values), + "median": statistics.median(values), + "p95": _paired_percentile(values, 0.95), + "min": min(values), + "max": max(values), + "raw_values": values, + } + + +def bootstrap_median_ci(values: list[float], samples: int = 10_000) -> list[float]: + if len(values) == 1: + return [values[0], values[0]] + generator = random.Random(0xD1F1) + estimates = [] + for _ in range(samples): + resample = [generator.choice(values) for _ in values] + estimates.append(statistics.median(resample)) + return [_paired_percentile(estimates, 0.025), _paired_percentile(estimates, 0.975)] + + +def source_digest(files: list[Path]) -> str: + digest = hashlib.sha256() + for path in sorted(files): + digest.update(str(path.relative_to(REPO)).encode()) + digest.update(b"\0") + digest.update(path.read_bytes()) + digest.update(b"\0") + return digest.hexdigest() + + +def command_output(command: list[str], *, check: bool = True) -> str: + completed = subprocess.run(command, check=check, text=True, capture_output=True) + return completed.stdout.strip() + + +def docker(*args: str, check: bool = True) -> str: + return command_output(["docker", *args], check=check) + + +def parse_published_port(value: str) -> int: + first = next((line.strip() for line in value.splitlines() if line.strip()), "") + _, separator, port_text = first.rpartition(":") + if not separator or not port_text.isdigit(): + raise ContainerBenchmarkError(f"unexpected Docker port mapping: {value!r}") + port = int(port_text) + if not 1 <= port <= 65_535: + raise ContainerBenchmarkError(f"unexpected Docker port mapping: {value!r}") + return port + + +def wait_for_health(port: int, process_started_ns: int) -> float: + for _ in range(300): + try: + conn = http.client.HTTPConnection("127.0.0.1", port, timeout=1) + conn.request("GET", "/healthz") + response = conn.getresponse() + response.read() + conn.close() + if response.status == 200: + return (time.perf_counter_ns() - process_started_ns) / 1_000_000 + except OSError: + pass + time.sleep(0.02) + raise ContainerBenchmarkError(f"container health check timed out on port {port}") + + +def request_json( + conn: http.client.HTTPConnection, + method: str, + path: str, + token: str, + payload: dict[str, Any], +) -> tuple[dict[str, Any], float]: + body = json.dumps(payload) + started = time.perf_counter_ns() + conn.request( + method, + path, + body=body, + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, + ) + response = conn.getresponse() + raw = response.read() + elapsed_ms = (time.perf_counter_ns() - started) / 1_000_000 + if response.status != 200: + raise ContainerBenchmarkError( + f"{method} {path} returned HTTP {response.status}: {raw[:500]!r}" + ) + try: + result = json.loads(raw) + except json.JSONDecodeError as exc: + raise ContainerBenchmarkError( + f"{method} {path} returned invalid JSON: {raw[:500]!r}" + ) from exc + if not isinstance(result, dict): + raise ContainerBenchmarkError(f"{method} {path} returned non-object JSON") + return result, elapsed_ms + + +def require_exited( + payload: dict[str, Any], label: str, output: str | None = None +) -> None: + if ( + not payload.get("done") + or payload.get("status") != "exited" + or payload.get("exit_code") != 0 + ): + raise ContainerBenchmarkError( + f"{label} returned unexpected lifecycle: {payload!r}" + ) + if output is not None and output not in payload.get("output", ""): + raise ContainerBenchmarkError(f"{label} output mismatch: {payload!r}") + + +def collect_health_preflights(port: int, requests: int) -> dict[str, Any]: + """Measure the fresh-connection health request used before Rust admission.""" + samples = [] + for index in range(requests): + conn = http.client.HTTPConnection("127.0.0.1", port, timeout=5) + started = time.perf_counter_ns() + try: + conn.request("GET", "/healthz") + response = conn.getresponse() + raw = response.read() + finally: + conn.close() + samples.append((time.perf_counter_ns() - started) / 1_000_000) + if response.status != 200: + raise ContainerBenchmarkError( + f"health preflight {index} returned HTTP {response.status}: {raw[:500]!r}" + ) + try: + payload = json.loads(raw) + except json.JSONDecodeError as exc: + raise ContainerBenchmarkError( + f"health preflight {index} returned invalid JSON: {raw[:500]!r}" + ) from exc + if payload != {"status": "ok"}: + raise ContainerBenchmarkError( + f"health preflight {index} returned unexpected payload: {payload!r}" + ) + return stats(samples) + + +def collect_job_workloads(port: int, token: str, jobs: int) -> dict[str, Any]: + conn = http.client.HTTPConnection("127.0.0.1", port, timeout=60) + try: + cold, cold_ms = request_json( + conn, + "POST", + "/v1/jobs/run", + token, + {"script": "printf cold-container"}, + ) + require_exited(cold, "cold job", "cold-container") + + for index in range(5): + result, _ = request_json( + conn, + "POST", + "/v1/jobs/run", + token, + {"script": "printf warmup"}, + ) + require_exited(result, f"warmup {index}", "warmup") + + sequential = [] + for index in range(jobs): + result, elapsed = request_json( + conn, + "POST", + "/v1/jobs/run", + token, + {"script": "printf benchmark"}, + ) + require_exited(result, f"sequential job {index}", "benchmark") + sequential.append(elapsed) + + output_samples = [] + output_sizes = [] + for index in range(max(10, jobs // 4)): + result, elapsed = request_json( + conn, + "POST", + "/v1/jobs/run", + token, + { + "script": 'python3 -c \'print("x" * 32768, end="")\'', + "output_limit": 65_536, + }, + ) + require_exited(result, f"output job {index}") + output_size = len(result.get("output", "").encode()) + if output_size < 32_768: + raise ContainerBenchmarkError( + f"output job {index} returned only {output_size} bytes" + ) + output_samples.append(elapsed) + output_sizes.append(output_size) + + concurrent_count = max(24, jobs // 2) + + def one_concurrent(index: int) -> float: + local = http.client.HTTPConnection("127.0.0.1", port, timeout=60) + try: + result, elapsed = request_json( + local, + "POST", + "/v1/jobs/run", + token, + {"script": "printf concurrent"}, + ) + require_exited(result, f"concurrent job {index}", "concurrent") + return elapsed + finally: + local.close() + + wall_started = time.perf_counter_ns() + with ThreadPoolExecutor(max_workers=8) as executor: + concurrent = list(executor.map(one_concurrent, range(concurrent_count))) + concurrent_wall_ms = (time.perf_counter_ns() - wall_started) / 1_000_000 + + output_result = stats(output_samples) + output_result["min_output_bytes"] = min(output_sizes) + concurrent_result = stats(concurrent) + concurrent_result["wall_ms"] = concurrent_wall_ms + return { + "cold_first_job": stats([cold_ms]), + "sequential_small": stats(sequential), + "output_32k": output_result, + "concurrent_8_workers": concurrent_result, + } + finally: + conn.close() + + +def parse_memory_size(value: str) -> float: + value = value.strip() + units = { + "B": 1 / (1024 * 1024), + "kB": 1000 / (1024 * 1024), + "KB": 1000 / (1024 * 1024), + "KiB": 1 / 1024, + "MB": 1_000_000 / (1024 * 1024), + "MiB": 1.0, + "GB": 1_000_000_000 / (1024 * 1024), + "GiB": 1024.0, + } + for unit in sorted(units, key=len, reverse=True): + if value.endswith(unit): + return float(value[: -len(unit)].strip()) * units[unit] + raise ContainerBenchmarkError(f"unsupported memory size: {value!r}") + + +def parse_runtime_process_rss(value: str) -> float: + """Sum resident memory while excluding the short-lived sampler itself.""" + total_kib = 0 + for line in value.splitlines(): + fields = line.strip().split(maxsplit=1) + if not fields: + continue + if len(fields) != 2: + raise ContainerBenchmarkError(f"unexpected ps row: {line!r}") + rss, command = fields + if command in PROCESS_RSS_EXCLUDED_COMMANDS: + continue + try: + total_kib += int(rss) + except ValueError as exc: + raise ContainerBenchmarkError( + f"unexpected RSS value in ps row: {line!r}" + ) from exc + return total_kib / 1024.0 + + +def parse_docker_top_runtime_rss(value: str) -> float: + """Sum container RSS from host-side ``docker top`` output. + + The official 1.16.1 runtime image does not include procps. Sampling from + the host also avoids adding a short-lived measurement process to the + container being measured. + """ + lines = [line.strip() for line in value.splitlines() if line.strip()] + if not lines or lines[0].split()[:3] != ["PID", "RSS", "COMMAND"]: + raise ContainerBenchmarkError("unexpected docker top header") + total_kib = 0 + for line in lines[1:]: + fields = line.split(maxsplit=2) + if len(fields) != 3: + raise ContainerBenchmarkError(f"unexpected docker top row: {line!r}") + pid, rss, command = fields + if not pid.isdigit(): + raise ContainerBenchmarkError(f"unexpected PID in docker top row: {line!r}") + if command in PROCESS_RSS_EXCLUDED_COMMANDS: + continue + try: + total_kib += int(rss) + except ValueError as exc: + raise ContainerBenchmarkError( + f"unexpected RSS value in docker top row: {line!r}" + ) from exc + return total_kib / 1024.0 + + +def collect_memory(container: str, samples: int = 20) -> dict[str, Any]: + cgroup_values = [] + for _ in range(samples): + raw = docker("exec", container, "cat", "/sys/fs/cgroup/memory.current") + cgroup_values.append(int(raw) / (1024 * 1024)) + time.sleep(0.05) + + docker_stats_values = [] + for _ in range(3): + usage = docker("stats", "--no-stream", "--format", "{{.MemUsage}}", container) + docker_stats_values.append(parse_memory_size(usage.split("/")[0])) + + runtime_process_rss_values = [] + for _ in range(samples): + raw = docker("top", container, "-eo", "pid,rss,comm") + runtime_process_rss_values.append(parse_docker_top_runtime_rss(raw)) + time.sleep(0.05) + + process_table = docker("top", container, "-eo", "pid,ppid,rss,comm,args") + return { + "cgroup_memory_current": { + **stats(cgroup_values, unit="MiB"), + "throughput_per_sec": None, + }, + "docker_stats_memory": { + **stats(docker_stats_values, unit="MiB"), + "throughput_per_sec": None, + }, + "runtime_process_rss_sum": { + **stats(runtime_process_rss_values, unit="MiB"), + "throughput_per_sec": None, + "excluded_commands": sorted(PROCESS_RSS_EXCLUDED_COMMANDS), + }, + "process_table": process_table.splitlines(), + } + + +def run_implementation( + implementation: str, round_index: int, jobs: int, cpus: float +) -> dict[str, Any]: + image = IMAGES[implementation] + token = f"bench-{os.getpid()}-{round_index}-{implementation}" + container = f"dify-bench-{os.getpid()}-{round_index}-{implementation}" + started_ns = time.perf_counter_ns() + docker( + "run", + "-d", + "--pull", + "never", + "--name", + container, + "--cpus", + str(cpus), + "--memory", + "1g", + "-p", + "127.0.0.1::5004", + "-e", + f"SHELLCTL_AUTH_TOKEN={token}", + image, + ) + try: + # Let Docker allocate and publish the port atomically. Probing a free + # port and closing it before `docker run` leaves a real TOCTOU window. + port = parse_published_port(docker("port", container, "5004/tcp")) + startup_ms = wait_for_health(port, started_ns) + time.sleep(0.5) + idle_memory = collect_memory(container) + health_preflight = collect_health_preflights(port, jobs) + workloads = collect_job_workloads(port, token, jobs) + workloads["fresh_connection_health_preflight"] = health_preflight + post_jobs_memory = collect_memory(container) + return { + "status": "ok", + "implementation": implementation, + "image": image, + "startup_health_ready_ms": startup_ms, + "idle_memory": idle_memory, + "workloads": workloads, + "post_jobs_memory": post_jobs_memory, + } + except Exception as exc: + logs = docker("logs", "--tail", "200", container, check=False) + return { + "status": "failed", + "implementation": implementation, + "image": image, + "error": repr(exc), + "logs": logs, + } + finally: + docker("rm", "-f", container, check=False) + + +def metric_value(run: dict[str, Any], path: tuple[str, ...]) -> float: + value: Any = run + for part in path: + value = value[part] + return float(value) + + +def comparisons(rounds: list[dict[str, Any]]) -> list[dict[str, Any]]: + metrics = { + "startup-health-ready": ("startup_health_ready_ms",), + "idle-cgroup-memory": ("idle_memory", "cgroup_memory_current", "median"), + "idle-docker-stats-memory": ("idle_memory", "docker_stats_memory", "median"), + "idle-runtime-process-rss-sum": ( + "idle_memory", + "runtime_process_rss_sum", + "median", + ), + "post-jobs-cgroup-memory": ( + "post_jobs_memory", + "cgroup_memory_current", + "median", + ), + "fresh-connection-health-preflight": ( + "workloads", + "fresh_connection_health_preflight", + "median", + ), + "sequential-small": ("workloads", "sequential_small", "median"), + "output-32k": ("workloads", "output_32k", "median"), + "concurrent-8-workers": ("workloads", "concurrent_8_workers", "median"), + "cold-first-job": ("workloads", "cold_first_job", "median"), + } + output = [] + for name, path in metrics.items(): + go_values = [] + rust_values = [] + reductions = [] + ratios = [] + for round_result in rounds: + go = metric_value(round_result["implementations"]["go"], path) + rust = metric_value(round_result["implementations"]["rust"], path) + go_values.append(go) + rust_values.append(rust) + reductions.append((go - rust) / go * 100) + ratios.append(go / rust) + output.append( + { + "metric": name, + "rounds": len(rounds), + "go": describe(go_values), + "rust": describe(rust_values), + "go_over_rust_ratio": describe(ratios), + "rust_reduction_percent": describe(reductions), + "rust_reduction_percent_bootstrap_median_ci95": bootstrap_median_ci( + reductions + ), + } + ) + return output + + +def image_metadata(image: str) -> dict[str, Any]: + raw = docker( + "image", + "inspect", + "--format", + "{{json .Id}}|{{json .Architecture}}|{{json .Os}}|{{json .Size}}|{{json .Created}}", + image, + ) + image_id, architecture, os_name, size, created = raw.split("|", 4) + return { + "name": image, + "id": json.loads(image_id), + "architecture": json.loads(architecture), + "os": json.loads(os_name), + "size_bytes": json.loads(size), + "created": json.loads(created), + } + + +def metadata( + round_count: int, + jobs: int, + cpus: float, + rust_source_sha256: str | None, +) -> dict[str, Any]: + go_sources = [RUNTIME / "go.mod", RUNTIME / "go.sum"] + go_sources.extend((RUNTIME / "cmd").rglob("*.go")) + go_sources.extend((RUNTIME / "internal").rglob("*.go")) + rust_sources = [RUNTIME / "rust" / "Cargo.toml", RUNTIME / "rust" / "Cargo.lock"] + rust_sources.extend((RUNTIME / "rust" / "src").rglob("*.rs")) + return { + "schema_version": 2, + "timestamp_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "git_head": command_output(["git", "rev-parse", "HEAD"]), + "git_status": command_output(["git", "status", "--short"]).splitlines(), + "host": { + "platform": platform.platform(), + "machine": platform.machine(), + "logical_cpu_count": os.cpu_count(), + }, + "docker_server": docker("version", "--format", "{{.Server.Version}}"), + "round_count": round_count, + "jobs_per_round": jobs, + "container_cpu_limit": cpus, + "container_memory_limit": "1g", + "process_rss_excluded_commands": sorted(PROCESS_RSS_EXCLUDED_COMMANDS), + "images": {name: image_metadata(image) for name, image in IMAGES.items()}, + "source_sha256": { + "go_runtime": source_digest(go_sources), + "rust_runtime": rust_source_sha256 or source_digest(rust_sources), + "benchmark_harness": source_digest([Path(__file__).resolve()]), + }, + } + + +def print_summary(report: dict[str, Any]) -> None: + print("\n== Paired Linux container benchmark ==") + print( + "metric Go median Rust median Go/Rust Rust reduction" + ) + for item in report.get("comparisons", []): + print( + f"{item['metric']:<30} " + f"{item['go']['median']:>10.3f} " + f"{item['rust']['median']:>12.3f} " + f"{item['go_over_rust_ratio']['median']:>7.2f}x " + f"{item['rust_reduction_percent']['median']:>13.1f}%" + ) + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--rounds", type=int, default=5) + parser.add_argument("--jobs", type=int, default=50) + parser.add_argument("--cpus", type=float, default=2.0) + parser.add_argument("--go-image", default=IMAGES["go"]) + parser.add_argument("--rust-image", default=IMAGES["rust"]) + parser.add_argument( + "--rust-source-sha256", + help="source digest embedded in a prebuilt Rust image (defaults to the working tree)", + ) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + IMAGES.update(go=args.go_image, rust=args.rust_image) + + rounds = [] + failed = False + for round_index in range(args.rounds): + order = ["go", "rust"] if round_index % 2 == 0 else ["rust", "go"] + print(f"round {round_index + 1}/{args.rounds}: {','.join(order)}", flush=True) + implementations = {} + for implementation in order: + result = run_implementation( + implementation, round_index, args.jobs, args.cpus + ) + implementations[implementation] = result + if result["status"] != "ok": + failed = True + print(f" {implementation}: failed {result['error']}", flush=True) + else: + median = result["workloads"]["sequential_small"]["median"] + memory = result["idle_memory"]["docker_stats_memory"]["median"] + print( + f" {implementation}: sequential p50={median:.3f}ms idle={memory:.3f}MiB", + flush=True, + ) + time.sleep(0.25) + rounds.append( + { + "round": round_index + 1, + "order": order, + "implementations": implementations, + } + ) + + report: dict[str, Any] = { + "metadata": metadata( + args.rounds, args.jobs, args.cpus, args.rust_source_sha256 + ), + "rounds": rounds, + } + if not failed: + report["comparisons"] = comparisons(rounds) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n") + print_summary(report) + print(f"JSON report: {args.output}") + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/benchmarks/test_container_benchmark.py b/benchmarks/test_container_benchmark.py new file mode 100644 index 00000000000000..ab91d54dc186b1 --- /dev/null +++ b/benchmarks/test_container_benchmark.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +import json +from unittest import TestCase, mock + +from run_container_benchmarks import ( + ContainerBenchmarkError, + bootstrap_median_ci, + collect_health_preflights, + comparisons, + parse_docker_top_runtime_rss, + parse_memory_size, + parse_published_port, + parse_runtime_process_rss, + require_exited, + stats, +) + + +class MemoryParsingTests(TestCase): + def test_memory_units_are_normalized_to_mib(self) -> None: + self.assertEqual(parse_memory_size("1GiB"), 1024.0) + self.assertEqual(parse_memory_size("512MiB"), 512.0) + self.assertAlmostEqual(parse_memory_size("1000kB"), 1000 * 1000 / 1024 / 1024) + + def test_runtime_rss_excludes_sampler_and_keeps_long_lived_processes(self) -> None: + raw = """ + 972 tini + 9856 shellctl + 3400 tmux: server + 8128 ps + """ + self.assertAlmostEqual( + parse_runtime_process_rss(raw), (972 + 9856 + 3400) / 1024 + ) + + def test_runtime_rss_rejects_malformed_rows(self) -> None: + with self.assertRaisesRegex(ContainerBenchmarkError, "unexpected ps row"): + parse_runtime_process_rss("1234") + + def test_docker_top_rss_works_without_procps_in_the_image(self) -> None: + raw = """ + PID RSS COMMAND + 3497512 15644 shellctl + 3497513 3400 tmux: server + """ + self.assertAlmostEqual( + parse_docker_top_runtime_rss(raw), (15644 + 3400) / 1024 + ) + + def test_docker_top_rss_rejects_an_unexpected_header(self) -> None: + with self.assertRaisesRegex(ContainerBenchmarkError, "unexpected docker top header"): + parse_docker_top_runtime_rss("RSS COMMAND\n15644 shellctl") + + +class PortParsingTests(TestCase): + def test_published_port_accepts_ipv4_and_ipv6_mappings(self) -> None: + self.assertEqual(parse_published_port("127.0.0.1:49152"), 49152) + self.assertEqual(parse_published_port("[::1]:49153"), 49153) + + def test_published_port_rejects_malformed_mapping(self) -> None: + with self.assertRaisesRegex( + ContainerBenchmarkError, "unexpected Docker port mapping" + ): + parse_published_port("5004/tcp -> nowhere") + + +class CorrectnessValidationTests(TestCase): + def test_terminal_success_and_expected_output_are_required(self) -> None: + require_exited( + {"done": True, "status": "exited", "exit_code": 0, "output": "expected"}, + "test", + "expected", + ) + + with self.assertRaisesRegex(ContainerBenchmarkError, "unexpected lifecycle"): + require_exited( + {"done": False, "status": "running", "exit_code": None, "output": ""}, + "test", + ) + + with self.assertRaisesRegex(ContainerBenchmarkError, "output mismatch"): + require_exited( + {"done": True, "status": "exited", "exit_code": 0, "output": "wrong"}, + "test", + "expected", + ) + + def test_stats_keep_raw_samples_and_terminal_percentiles(self) -> None: + result = stats([1.0, 2.0, 3.0, 4.0]) + self.assertEqual(result["raw_values"], [1.0, 2.0, 3.0, 4.0]) + self.assertEqual(result["median"], 2.5) + self.assertEqual(result["p95"], 4.0) + + +class _HealthResponse: + status = 200 + + def __init__(self, payload: object) -> None: + self._payload = payload + + def read(self) -> bytes: + return json.dumps(self._payload).encode() + + +class _HealthConnection: + instances: list[_HealthConnection] = [] + payload: object = {"status": "ok"} + + def __init__(self, host: str, port: int, timeout: int) -> None: + self.host = host + self.port = port + self.timeout = timeout + self.closed = False + self.instances.append(self) + + def request(self, method: str, path: str) -> None: + if (method, path) != ("GET", "/healthz"): + raise AssertionError((method, path)) + + def getresponse(self) -> _HealthResponse: + return _HealthResponse(self.payload) + + def close(self) -> None: + self.closed = True + + +class HealthPreflightTests(TestCase): + def setUp(self) -> None: + _HealthConnection.instances = [] + _HealthConnection.payload = {"status": "ok"} + + @mock.patch( + "run_container_benchmarks.http.client.HTTPConnection", _HealthConnection + ) + def test_preflight_uses_fresh_connections_and_checks_exact_payload(self) -> None: + result = collect_health_preflights(5004, 3) + + self.assertEqual(result["samples"], 3) + self.assertEqual(len(_HealthConnection.instances), 3) + self.assertTrue( + all(connection.closed for connection in _HealthConnection.instances) + ) + + @mock.patch( + "run_container_benchmarks.http.client.HTTPConnection", _HealthConnection + ) + def test_preflight_rejects_non_contract_payload(self) -> None: + _HealthConnection.payload = {"status": "degraded"} + + with self.assertRaisesRegex(ContainerBenchmarkError, "unexpected payload"): + collect_health_preflights(5004, 1) + + +def _implementation(value: float) -> dict: + memory = { + "cgroup_memory_current": {"median": value}, + "docker_stats_memory": {"median": value}, + "runtime_process_rss_sum": {"median": value}, + } + workloads = { + "fresh_connection_health_preflight": {"median": value}, + "sequential_small": {"median": value}, + "output_32k": {"median": value}, + "concurrent_8_workers": {"median": value}, + "cold_first_job": {"median": value}, + } + return { + "startup_health_ready_ms": value, + "idle_memory": memory, + "post_jobs_memory": memory, + "workloads": workloads, + } + + +class PairedComparisonTests(TestCase): + def test_comparison_uses_within_round_pairs(self) -> None: + rounds = [ + { + "implementations": { + "go": _implementation(10.0), + "rust": _implementation(5.0), + } + }, + { + "implementations": { + "go": _implementation(20.0), + "rust": _implementation(10.0), + } + }, + ] + + result = comparisons(rounds) + + self.assertTrue(result) + self.assertTrue( + all(item["go_over_rust_ratio"]["median"] == 2.0 for item in result) + ) + self.assertTrue( + all(item["rust_reduction_percent"]["median"] == 50.0 for item in result) + ) + + def test_bootstrap_is_deterministic(self) -> None: + self.assertEqual( + bootstrap_median_ci([1.0, 2.0, 3.0]), bootstrap_median_ci([1.0, 2.0, 3.0]) + ) diff --git a/dify-agent-runtime/.dockerignore b/dify-agent-runtime/.dockerignore new file mode 100644 index 00000000000000..48451c2441b9c9 --- /dev/null +++ b/dify-agent-runtime/.dockerignore @@ -0,0 +1,3 @@ +bin/ +rust/target/ +.integration-state diff --git a/dify-agent-runtime/Makefile b/dify-agent-runtime/Makefile index 76e853a03f9336..02ad568405eb2b 100644 --- a/dify-agent-runtime/Makefile +++ b/dify-agent-runtime/Makefile @@ -1,4 +1,4 @@ -.PHONY: build clean test lint gen-cli-help integration integration-up integration-test integration-down +.PHONY: build clean test lint gen-cli-help rust-build rust-test rust-lint integration integration-up integration-test integration-test-rollout integration-logs integration-down BIN_DIR := bin AGENT_CLI_HELP_JSON := ../dify-agent/src/dify_agent/layers/_agent_cli_help.json @@ -25,6 +25,19 @@ $(BIN_DIR)/dify-agent: $(shell find cmd/dify-agent-cli internal/agentcli -name ' gen-cli-help: $(BIN_DIR)/dify-agent $(BIN_DIR)/dify-agent __dump-cli-help > $(AGENT_CLI_HELP_JSON) +# --- Rust runtime --- +# The Rust implementation lives beside the Go implementation while the +# protocol and integration tests are migrated incrementally. +rust-build: + cargo build --manifest-path rust/Cargo.toml --bins + +rust-test: + cargo test --manifest-path rust/Cargo.toml + +rust-lint: + cargo fmt --manifest-path rust/Cargo.toml --all -- --check + cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings + test: lint go test ./... @@ -36,8 +49,9 @@ clean: # --- Integration tests --- # -# Container name and host port are randomised per invocation to avoid -# conflicts when multiple runs overlap or a previous run leaked a container. +# Container names are randomised per invocation. Docker atomically assigns +# host ports, avoiding the bind-close-run race of probing a supposedly free +# port before starting each container. # # State is written to .integration-state so that separate make invocations # (integration-up, integration-test, integration-logs, integration-down) @@ -45,59 +59,108 @@ clean: # lifecycle in a single invocation with automatic cleanup. IMAGE_NAME := dify-agent-runtime:test +RUST_IMAGE_NAME := dify-agent-runtime-rust:test STATE_FILE := .integration-state integration-up: - @echo "Building runtime image..." + @echo "Building Go runtime image..." docker build -t $(IMAGE_NAME) -f docker/Dockerfile . - $(eval TEST_ID := $(shell printf '%05d' $$((RANDOM % 100000)))) + @echo "Building Rust runtime image..." + docker build -t $(RUST_IMAGE_NAME) -f docker/Dockerfile.rust . + $(eval TEST_ID := $(shell python3 -c 'import secrets; print(secrets.token_hex(4))')) $(eval CONTAINER_NAME := sandbox-rt-$(TEST_ID)) - $(eval HOST_PORT := $(shell python3 -c 'import socket; s=socket.socket(); s.bind(("",0)); print(s.getsockname()[1]); s.close()')) $(eval AUTH_TOKEN := test-token-$(TEST_ID)) - @echo "Starting runtime container $(CONTAINER_NAME) on port $(HOST_PORT)..." + $(eval CONTAINER_NAME_RUST := sandbox-rt-rust-$(TEST_ID)) + $(eval CONTAINER_NAME_NOISO := sandbox-rt-noiso-$(TEST_ID)) + $(eval CONTAINER_NAME_RUST_NOISO := sandbox-rt-rust-noiso-$(TEST_ID)) + $(eval AUTH_TOKEN_NOISO := test-token-noiso-$(TEST_ID)) + @{ \ + echo 'CONTAINER_NAME=$(CONTAINER_NAME)'; \ + echo 'CONTAINER_NAME_RUST=$(CONTAINER_NAME_RUST)'; \ + echo 'CONTAINER_NAME_NOISO=$(CONTAINER_NAME_NOISO)'; \ + echo 'CONTAINER_NAME_RUST_NOISO=$(CONTAINER_NAME_RUST_NOISO)'; \ + echo 'AUTH_TOKEN=$(AUTH_TOKEN)'; \ + echo 'AUTH_TOKEN_NOISO=$(AUTH_TOKEN_NOISO)'; \ + } > $(STATE_FILE) + @echo "Starting runtime container $(CONTAINER_NAME)..." docker run -d --name $(CONTAINER_NAME) \ - -p $(HOST_PORT):5004 \ + -p 127.0.0.1::5004 \ -e SHELLCTL_AUTH_TOKEN=$(AUTH_TOKEN) \ $(IMAGE_NAME) - @echo 'CONTAINER_NAME=$(CONTAINER_NAME)' > $(STATE_FILE) - @echo 'HOST_PORT=$(HOST_PORT)' >> $(STATE_FILE) - @echo 'AUTH_TOKEN=$(AUTH_TOKEN)' >> $(STATE_FILE) - @echo "Waiting for runtime to be ready..." - @for i in $$(seq 1 30); do \ - if curl -sf http://localhost:$(HOST_PORT)/healthz > /dev/null 2>&1; then \ - echo "Runtime is ready on port $(HOST_PORT)"; \ + @HOST_PORT=$$(docker port $(CONTAINER_NAME) 5004/tcp | awk -F: 'NR == 1 { print $$NF }'); \ + test -n "$$HOST_PORT"; \ + echo "HOST_PORT=$$HOST_PORT" >> $(STATE_FILE); \ + echo "Waiting for runtime on port $$HOST_PORT..."; \ + for i in $$(seq 1 30); do \ + if curl -sf http://127.0.0.1:$$HOST_PORT/healthz > /dev/null 2>&1; then \ + echo "Runtime is ready on port $$HOST_PORT"; \ break; \ fi; \ if [ "$$i" -eq 30 ]; then \ echo "ERROR: runtime not ready after 60s" >&2; \ docker logs $(CONTAINER_NAME); \ - docker rm -f $(CONTAINER_NAME) 2>/dev/null; \ - rm -f $(STATE_FILE); \ exit 1; \ fi; \ sleep 2; \ done - $(eval CONTAINER_NAME_NOISO := sandbox-rt-noiso-$(TEST_ID)) - $(eval HOST_PORT_NOISO := $(shell python3 -c 'import socket; s=socket.socket(); s.bind(("",0)); print(s.getsockname()[1]); s.close()')) - $(eval AUTH_TOKEN_NOISO := test-token-noiso-$(TEST_ID)) - @echo "Starting no-isolation container $(CONTAINER_NAME_NOISO) on port $(HOST_PORT_NOISO)..." + @echo "Starting Rust runtime container $(CONTAINER_NAME_RUST)..." + docker run -d --name $(CONTAINER_NAME_RUST) \ + -p 127.0.0.1::5004 \ + -e SHELLCTL_AUTH_TOKEN=$(AUTH_TOKEN) \ + $(RUST_IMAGE_NAME) + @HOST_PORT_RUST=$$(docker port $(CONTAINER_NAME_RUST) 5004/tcp | awk -F: 'NR == 1 { print $$NF }'); \ + test -n "$$HOST_PORT_RUST"; \ + echo "HOST_PORT_RUST=$$HOST_PORT_RUST" >> $(STATE_FILE); \ + for i in $$(seq 1 30); do \ + if curl -sf http://127.0.0.1:$$HOST_PORT_RUST/healthz > /dev/null 2>&1; then \ + echo "Rust runtime is ready on port $$HOST_PORT_RUST"; \ + break; \ + fi; \ + if [ "$$i" -eq 30 ]; then \ + echo "ERROR: Rust runtime not ready after 60s" >&2; \ + docker logs $(CONTAINER_NAME_RUST); \ + exit 1; \ + fi; \ + sleep 2; \ + done + @echo "Starting no-isolation container $(CONTAINER_NAME_NOISO)..." docker run -d --name $(CONTAINER_NAME_NOISO) \ - -p $(HOST_PORT_NOISO):5004 \ + -p 127.0.0.1::5004 \ -e SHELLCTL_AUTH_TOKEN=$(AUTH_TOKEN_NOISO) \ -e SHELLCTL_ENABLE_PATH_ISOLATION=false \ $(IMAGE_NAME) - @echo 'CONTAINER_NAME_NOISO=$(CONTAINER_NAME_NOISO)' >> $(STATE_FILE) - @echo 'HOST_PORT_NOISO=$(HOST_PORT_NOISO)' >> $(STATE_FILE) - @echo 'AUTH_TOKEN_NOISO=$(AUTH_TOKEN_NOISO)' >> $(STATE_FILE) - @for i in $$(seq 1 30); do \ - if curl -sf http://localhost:$(HOST_PORT_NOISO)/healthz > /dev/null 2>&1; then \ - echo "No-isolation runtime is ready on port $(HOST_PORT_NOISO)"; \ + @HOST_PORT_NOISO=$$(docker port $(CONTAINER_NAME_NOISO) 5004/tcp | awk -F: 'NR == 1 { print $$NF }'); \ + test -n "$$HOST_PORT_NOISO"; \ + echo "HOST_PORT_NOISO=$$HOST_PORT_NOISO" >> $(STATE_FILE); \ + for i in $$(seq 1 30); do \ + if curl -sf http://127.0.0.1:$$HOST_PORT_NOISO/healthz > /dev/null 2>&1; then \ + echo "No-isolation runtime is ready on port $$HOST_PORT_NOISO"; \ break; \ fi; \ if [ "$$i" -eq 30 ]; then \ echo "ERROR: no-isolation runtime not ready after 60s" >&2; \ docker logs $(CONTAINER_NAME_NOISO); \ - docker rm -f $(CONTAINER_NAME_NOISO) 2>/dev/null; \ + exit 1; \ + fi; \ + sleep 2; \ + done + @echo "Starting Rust no-isolation container $(CONTAINER_NAME_RUST_NOISO)..." + docker run -d --name $(CONTAINER_NAME_RUST_NOISO) \ + -p 127.0.0.1::5004 \ + -e SHELLCTL_AUTH_TOKEN=$(AUTH_TOKEN_NOISO) \ + -e SHELLCTL_ENABLE_PATH_ISOLATION=false \ + $(RUST_IMAGE_NAME) + @HOST_PORT_RUST_NOISO=$$(docker port $(CONTAINER_NAME_RUST_NOISO) 5004/tcp | awk -F: 'NR == 1 { print $$NF }'); \ + test -n "$$HOST_PORT_RUST_NOISO"; \ + echo "HOST_PORT_RUST_NOISO=$$HOST_PORT_RUST_NOISO" >> $(STATE_FILE); \ + for i in $$(seq 1 30); do \ + if curl -sf http://127.0.0.1:$$HOST_PORT_RUST_NOISO/healthz > /dev/null 2>&1; then \ + echo "Rust no-isolation runtime is ready on port $$HOST_PORT_RUST_NOISO"; \ + break; \ + fi; \ + if [ "$$i" -eq 30 ]; then \ + echo "ERROR: Rust no-isolation runtime not ready after 60s" >&2; \ + docker logs $(CONTAINER_NAME_RUST_NOISO); \ exit 1; \ fi; \ sleep 2; \ @@ -108,23 +171,46 @@ integration-test: @. ./$(STATE_FILE); \ SHELLCTL_GO_URL=http://localhost:$$HOST_PORT \ SHELLCTL_TEST_TOKEN=$$AUTH_TOKEN \ + SHELLCTL_RUST_URL=http://localhost:$$HOST_PORT_RUST \ + SHELLCTL_RUST_TEST_TOKEN=$$AUTH_TOKEN \ + SHELLCTL_RUST_IMAGE=$(RUST_IMAGE_NAME) \ SHELLCTL_GO_URL_NO_ISOLATION=http://localhost:$$HOST_PORT_NOISO \ SHELLCTL_TEST_TOKEN_NO_ISOLATION=$$AUTH_TOKEN_NOISO \ + SHELLCTL_RUST_URL_NO_ISOLATION=http://localhost:$$HOST_PORT_RUST_NOISO \ + SHELLCTL_RUST_TEST_TOKEN_NO_ISOLATION=$$AUTH_TOKEN_NOISO \ go test -tags=integration -v -count=1 -timeout=300s ./tests/... +integration-test-rollout: + @test -f $(STATE_FILE) || { echo "ERROR: run 'make integration-up' first" >&2; exit 1; } + @. ./$(STATE_FILE); \ + NO_PROXY=127.0.0.1,localhost \ + PYTHONPATH=../dify-agent/src \ + DIFY_AGENT_TEST_LOCAL_SHELLCTL_ENDPOINT=http://127.0.0.1:$$HOST_PORT \ + DIFY_AGENT_TEST_LOCAL_SHELLCTL_AUTH_TOKEN=$$AUTH_TOKEN \ + DIFY_AGENT_TEST_RUST_SHELLCTL_ENDPOINT=http://127.0.0.1:$$HOST_PORT_RUST \ + DIFY_AGENT_TEST_RUST_SHELLCTL_AUTH_TOKEN=$$AUTH_TOKEN \ + uv run --project ../dify-agent --extra server pytest --import-mode=importlib -q \ + ../dify-agent/tests/integration/dify_agent/runtime_backend/test_working_environment.py \ + -k 'local' + integration-logs: @test -f $(STATE_FILE) || { echo "ERROR: no state file" >&2; exit 1; } - @. ./$(STATE_FILE); docker logs $$CONTAINER_NAME + @. ./$(STATE_FILE); \ + docker logs $$CONTAINER_NAME; \ + docker logs $$CONTAINER_NAME_RUST; \ + docker logs $$CONTAINER_NAME_NOISO; \ + docker logs $$CONTAINER_NAME_RUST_NOISO integration-down: @if [ -f $(STATE_FILE) ]; then \ . ./$(STATE_FILE); \ docker rm -f $$CONTAINER_NAME 2>/dev/null || true; \ + docker rm -f $$CONTAINER_NAME_RUST 2>/dev/null || true; \ docker rm -f $$CONTAINER_NAME_NOISO 2>/dev/null || true; \ + docker rm -f $$CONTAINER_NAME_RUST_NOISO 2>/dev/null || true; \ rm -f $(STATE_FILE); \ fi integration: - @$(MAKE) integration-up - @trap '$(MAKE) integration-down' EXIT; \ - $(MAKE) integration-test + @trap '$(MAKE) integration-down' EXIT INT TERM; \ + $(MAKE) integration-up && $(MAKE) integration-test && $(MAKE) integration-test-rollout diff --git a/dify-agent-runtime/README.md b/dify-agent-runtime/README.md index 59b1d45116736e..b4d19cb8ddfaf2 100644 --- a/dify-agent-runtime/README.md +++ b/dify-agent-runtime/README.md @@ -1,6 +1,8 @@ # dify-agent-runtime -Go implementation of the shellctl server and runtime utilities. +Stable Go implementation of the shellctl server and runtime utilities. The +opt-in Rust canary lives in [`rust/`](./rust/README.md); Go remains the default +and rollback path while compatibility and production behavior are validated. ## Architecture diff --git a/dify-agent-runtime/docker/Dockerfile.rust b/dify-agent-runtime/docker/Dockerfile.rust new file mode 100644 index 00000000000000..9e0868f8d6c26e --- /dev/null +++ b/dify-agent-runtime/docker/Dockerfile.rust @@ -0,0 +1,83 @@ +# Experimental Rust shellctl image. It intentionally mirrors Dockerfile's +# production toolchain so the shared acceptance tests compare runtimes rather +# than different container contents. + +FROM rust:1.95-bookworm AS rust-builder + +WORKDIR /src/rust +COPY rust/Cargo.toml rust/Cargo.lock ./ +COPY rust/src ./src +RUN cargo build --locked --release --bins + +FROM golang:1.26 AS go-builder + +WORKDIR /src +COPY go.mod go.sum ./ +RUN go mod download +COPY cmd ./cmd +COPY internal ./internal +RUN CGO_ENABLED=0 go build -o /bin/dify-agent ./cmd/dify-agent-cli + +FROM python:3.12-slim-bookworm AS production + +ARG NODE_VERSION=22.22.1 +ARG PNPM_VERSION=11.9.0 +ARG UV_VERSION=0.8.9 + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + bash \ + ca-certificates \ + curl \ + file \ + git \ + jq \ + less \ + openssh-client \ + procps \ + ripgrep \ + tini \ + tmux \ + unzip \ + xz-utils \ + zip \ + && node_arch="$(dpkg --print-architecture)" \ + && case "${node_arch}" in \ + amd64) node_arch="x64" ;; \ + arm64) node_arch="arm64" ;; \ + *) echo "Unsupported Node.js architecture: ${node_arch}" >&2; exit 1 ;; \ + esac \ + && node_dist="node-v${NODE_VERSION}-linux-${node_arch}" \ + && curl -fsSLO "https://nodejs.org/dist/v${NODE_VERSION}/SHASUMS256.txt" \ + && curl -fsSLO "https://nodejs.org/dist/v${NODE_VERSION}/${node_dist}.tar.xz" \ + && grep " ${node_dist}.tar.xz\$" SHASUMS256.txt | sha256sum -c - \ + && tar -xJf "${node_dist}.tar.xz" -C /usr/local --strip-components=1 \ + && rm -f SHASUMS256.txt "${node_dist}.tar.xz" \ + && npm install --global "pnpm@${PNPM_VERSION}" \ + && npm cache clean --force \ + && rm -rf /var/lib/apt/lists/* + +RUN python -m pip install --no-cache-dir "uv==${UV_VERSION}" + +COPY --from=rust-builder /src/rust/target/release/shellctl /usr/local/bin/shellctl +COPY --from=rust-builder /src/rust/target/release/shellctl-sanitize-pty /usr/local/bin/shellctl-sanitize-pty +COPY --from=rust-builder /src/rust/target/release/shellctl-runner-exit /usr/local/bin/shellctl-runner-exit +COPY --from=rust-builder /src/rust/target/release/shellctl-runner /usr/local/bin/shellctl-runner +COPY --from=go-builder /bin/dify-agent /usr/local/bin/dify-agent + +RUN useradd --create-home --shell /bin/sh dify \ + && mkdir -p /mnt/drive \ + && chown dify:dify /home \ + && chown -R dify:dify /home/dify /mnt/drive + +USER dify +WORKDIR /home/dify + +EXPOSE 5004 + +ENTRYPOINT ["/usr/bin/tini", "-g", "--"] +CMD ["shellctl", "serve", "--listen", "0.0.0.0:5004"] diff --git a/dify-agent-runtime/internal/sanitize/sanitize.go b/dify-agent-runtime/internal/sanitize/sanitize.go index f6cebe339e1c58..5f12776aa87345 100644 --- a/dify-agent-runtime/internal/sanitize/sanitize.go +++ b/dify-agent-runtime/internal/sanitize/sanitize.go @@ -37,7 +37,13 @@ func New() *PtySanitizer { // Feed consumes one chunk of decoded text and returns newly stable output. func (s *PtySanitizer) Feed(text []byte) []byte { - var out []byte + return s.FeedInto(text, nil) +} + +// FeedInto is the allocation-conscious variant of Feed. The caller may +// provide a reusable output buffer when sanitizing a long PTY stream. +func (s *PtySanitizer) FeedInto(text, out []byte) []byte { + out = out[:0] for len(text) > 0 { r, size := utf8.DecodeRune(text) if r == utf8.RuneError && size <= 1 { @@ -161,14 +167,15 @@ func Run(readyFile string, stdin io.Reader, stdout io.Writer) error { sanitizer := New() reader := bufio.NewReaderSize(stdin, 65536) - writer := bufio.NewWriter(stdout) + writer := bufio.NewWriterSize(stdout, 65536) defer func() { _ = writer.Flush() }() buf := make([]byte, 65536) + sanitized := make([]byte, 65536) for { n, err := reader.Read(buf) if n > 0 { - out := sanitizer.Feed(buf[:n]) + out := sanitizer.FeedInto(buf[:n], sanitized) if len(out) > 0 { if _, werr := writer.Write(out); werr != nil { return werr diff --git a/dify-agent-runtime/internal/sanitize/sanitize_test.go b/dify-agent-runtime/internal/sanitize/sanitize_test.go index 4fae6abc3ff9e9..83ef754c75e473 100644 --- a/dify-agent-runtime/internal/sanitize/sanitize_test.go +++ b/dify-agent-runtime/internal/sanitize/sanitize_test.go @@ -81,3 +81,14 @@ func TestInvalidUTF8(t *testing.T) { t.Errorf("got %q, want %q", string(out), expected) } } + +func TestFeedIntoMatchesFeed(t *testing.T) { + input := []byte("first\n\x1b[31msecond\x1b[0m\n50%\r100%\n") + allocated := New() + reusable := New() + want := allocated.Feed(input) + buffer := reusable.FeedInto(input, make([]byte, 0, len(input))) + if string(buffer) != string(want) { + t.Fatalf("FeedInto = %q, Feed = %q", buffer, want) + } +} diff --git a/dify-agent-runtime/internal/server/config.go b/dify-agent-runtime/internal/server/config.go index 226a44b8c434ff..144a22c06ba39b 100644 --- a/dify-agent-runtime/internal/server/config.go +++ b/dify-agent-runtime/internal/server/config.go @@ -22,6 +22,7 @@ const ( DefaultGCIntervalSeconds = 60.0 DefaultGCFinishedJobRetentionSeconds = 300.0 DefaultPollInterval = 50 * time.Millisecond + DefaultOutputPollInterval = 5 * time.Millisecond DefaultPipeMonitorInterval = 1 * time.Second DefaultPipeReadyTimeout = 10 * time.Second DefaultSQLiteBusyTimeoutMs = 5000 @@ -49,6 +50,7 @@ type Config struct { MaxOutputLimitBytes int DefaultTerminateGraceSeconds float64 PollInterval time.Duration + OutputPollInterval time.Duration PipeMonitorInterval time.Duration PipeReadyTimeout time.Duration SQLiteBusyTimeoutMs int @@ -80,6 +82,7 @@ func DefaultConfig() *Config { MaxOutputLimitBytes: MaxOutputLimitBytes, DefaultTerminateGraceSeconds: DefaultTerminateGraceSeconds, PollInterval: DefaultPollInterval, + OutputPollInterval: DefaultOutputPollInterval, PipeMonitorInterval: DefaultPipeMonitorInterval, PipeReadyTimeout: DefaultPipeReadyTimeout, SQLiteBusyTimeoutMs: DefaultSQLiteBusyTimeoutMs, diff --git a/dify-agent-runtime/internal/server/config_test.go b/dify-agent-runtime/internal/server/config_test.go index 0b6aed9449e9a8..34d24e9b0b8836 100644 --- a/dify-agent-runtime/internal/server/config_test.go +++ b/dify-agent-runtime/internal/server/config_test.go @@ -18,6 +18,9 @@ func TestDefaultConfig(t *testing.T) { if cfg.SQLiteBusyTimeoutMs != DefaultSQLiteBusyTimeoutMs { t.Errorf("expected busy_timeout=%d, got %d", DefaultSQLiteBusyTimeoutMs, cfg.SQLiteBusyTimeoutMs) } + if cfg.OutputPollInterval != DefaultOutputPollInterval { + t.Errorf("expected output poll interval=%s, got %s", DefaultOutputPollInterval, cfg.OutputPollInterval) + } } func TestConfigPaths(t *testing.T) { diff --git a/dify-agent-runtime/internal/server/service.go b/dify-agent-runtime/internal/server/service.go index a859aa3474db61..74fe2bcd10f3ad 100644 --- a/dify-agent-runtime/internal/server/service.go +++ b/dify-agent-runtime/internal/server/service.go @@ -252,16 +252,10 @@ func (s *Service) RunJob(req *RunJobRequest) (*JobResult, error) { } func (s *Service) startJob(jobID, jobDir, cwd string, cols, rows int) error { - log.Printf("startJob [%s]: creating tmux session", jobID) - if err := s.tmux.CreateJobSession(jobID, jobDir, cwd, cols, rows); err != nil { - log.Printf("startJob [%s]: tmux session failed: %v", jobID, err) - return err - } - pipeReadyPath := filepath.Join(jobDir, ".pipe-ready") - log.Printf("startJob [%s]: enabling output pipe", jobID) - if err := s.tmux.EnableOutputPipe(jobID, jobDir, pipeReadyPath); err != nil { - log.Printf("startJob [%s]: pipe-pane failed: %v", jobID, err) + log.Printf("startJob [%s]: creating tmux session and output pipe", jobID) + if err := s.tmux.StartJob(jobID, jobDir, cwd, pipeReadyPath, cols, rows); err != nil { + log.Printf("startJob [%s]: tmux startup failed: %v", jobID, err) return err } @@ -340,10 +334,35 @@ func (s *Service) WaitJob(jobID string, req *WaitJobRequest) (*JobResult, error) lastGrowthAt = &now } + // Output files and completion artifacts are local and cheap to inspect. + // Probe them frequently for short jobs, while keeping tmux subprocess + // probes at the coarser lifecycle interval as a failure-detection fallback. + view, err := s.GetJobStatus(jobID) + if err != nil { + return nil, err + } + nextRuntimeProbe := time.Now().Add(s.config.PollInterval) + outputPollInterval := s.config.OutputPollInterval + if outputPollInterval <= 0 { + outputPollInterval = DefaultOutputPollInterval + } + for { - view, err := s.GetJobStatus(jobID) - if err != nil { - return nil, err + now := time.Now() + if !view.Done { + if exit := s.drainedNormalExitMetadata(jobID); exit != nil { + if err := s.db.RecordRunnerExit(jobID, exit.exitCode, exit.endedAt); err == nil { + if row, getErr := s.db.GetJob(jobID); getErr == nil { + view = s.statusViewFromRow(row) + } + } + } else if !now.Before(nextRuntimeProbe) { + view, err = s.GetJobStatus(jobID) + if err != nil { + return nil, err + } + nextRuntimeProbe = now.Add(s.config.PollInterval) + } } currentSize := fileSize(outputPath) @@ -389,7 +408,7 @@ func (s *Service) WaitJob(jobID string, req *WaitJobRequest) (*JobResult, error) } } - if time.Now().After(deadline) { + if now.After(deadline) { var window *OutputWindow if currentSize > int64(req.Offset) { window, err = ReadOutputWindow(outputPath, req.Offset, outputLimit) @@ -402,7 +421,7 @@ func (s *Service) WaitJob(jobID string, req *WaitJobRequest) (*JobResult, error) return s.jobResultFromView(view, row, window), nil } - time.Sleep(s.config.PollInterval) + time.Sleep(outputPollInterval) } } diff --git a/dify-agent-runtime/internal/server/tmux.go b/dify-agent-runtime/internal/server/tmux.go index dc5cfac9565458..9f32b7ddeedf33 100644 --- a/dify-agent-runtime/internal/server/tmux.go +++ b/dify-agent-runtime/internal/server/tmux.go @@ -104,6 +104,39 @@ func (t *TmuxController) CreateJobSession(jobID, jobDir, cwd string, cols, rows return nil } +// StartJob creates the tmux session and attaches the output pipe in one tmux +// client invocation. Keeping both commands in the same client avoids a +// process launch and removes the cold-start window between session creation +// and pipe attachment. +func (t *TmuxController) StartJob(jobID, jobDir, cwd, readyFile string, cols, rows int) error { + runnerCmd := shellJoin([]string{ + t.config.RunnerPath(), jobDir, jobID, cwd, + }) + pipeCmd := t.buildPipeCommand(jobID, jobDir, readyFile) + result, err := t.runTmuxNoCheck(t.startJobArgs(jobID, runnerCmd, pipeCmd, cols, rows)...) + if err != nil { + return err + } + if result.exitCode != 0 { + return NewServerError(500, "tmux_new_session_failed", + strings.TrimSpace(result.stderr)) + } + return nil +} + +func (t *TmuxController) startJobArgs(jobID, runnerCmd, pipeCmd string, cols, rows int) []string { + return []string{ + "-f", "/dev/null", + "new-session", "-d", + "-s", JobSessionName(jobID), + "-x", fmt.Sprintf("%d", cols), + "-y", fmt.Sprintf("%d", rows), + runnerCmd, + ";", + "pipe-pane", "-o", "-t", JobPaneTarget(jobID), pipeCmd, + } +} + // EnableOutputPipe attaches the sanitizeā†’output pipeline via tmux pipe-pane. func (t *TmuxController) EnableOutputPipe(jobID, jobDir string, readyFile string) error { pipeCmd := t.buildPipeCommand(jobID, jobDir, readyFile) diff --git a/dify-agent-runtime/internal/server/tmux_test.go b/dify-agent-runtime/internal/server/tmux_test.go index 68a8003a2c67d2..d1f45eda6cf56f 100644 --- a/dify-agent-runtime/internal/server/tmux_test.go +++ b/dify-agent-runtime/internal/server/tmux_test.go @@ -114,6 +114,19 @@ func TestShellJoin(t *testing.T) { } } +func TestStartJobArgsBatchSessionAndPipe(t *testing.T) { + controller := &TmuxController{} + args := controller.startJobArgs("job-123", "runner command", "pipe command", 200, 50) + want := []string{ + "-f", "/dev/null", + "new-session", "-d", "-s", "shellctl-job-123", "-x", "200", "-y", "50", + "runner command", ";", "pipe-pane", "-o", "-t", "shellctl-job-123:0.0", "pipe command", + } + if strings.Join(args, "\x00") != strings.Join(want, "\x00") { + t.Fatalf("startJobArgs = %#v, want %#v", args, want) + } +} + func TestIsTmuxTargetMissing(t *testing.T) { missingMsgs := []string{ "can't find pane: shellctl-abc", diff --git a/dify-agent-runtime/rust/Cargo.lock b/dify-agent-runtime/rust/Cargo.lock new file mode 100644 index 00000000000000..fc304b2985946f --- /dev/null +++ b/dify-agent-runtime/rust/Cargo.lock @@ -0,0 +1,985 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "once_cell", + "version_check", + "zerocopy", +] + +[[package]] +name = "android_system_properties" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc" +dependencies = [ + "libc", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "axum" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" +dependencies = [ + "axum-core", + "bytes", + "form_urlencoded", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + +[[package]] +name = "cc" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d262e149917187838d5b42777c8253bcb64500067342904e7d429499a6f277e" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "serde", + "wasm-bindgen", + "windows-link", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "dify-agent-runtime" +version = "0.1.0" +dependencies = [ + "axum", + "chrono", + "landlock", + "rand", + "rusqlite", + "serde", + "serde_json", + "tokio", +] + +[[package]] +name = "enumflags2" +version = "0.7.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1027f7680c853e056ebcec683615fb6fbbc07dbaa13b4d5d9442b146ded4ecef" +dependencies = [ + "enumflags2_derive", +] + +[[package]] +name = "enumflags2_derive" +version = "0.7.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67c78a4d8fdf9953a5c9d458f9efe940fd97a0cab0941c075a813ac594733827" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + +[[package]] +name = "find-msvc-tools" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26b73573e6edcd2af0cdf47bd6cb58f0b3839491263c314eaad1ccf24430e1de" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "futures-channel" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae" +dependencies = [ + "futures-core", +] + +[[package]] +name = "futures-core" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" + +[[package]] +name = "futures-task" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" + +[[package]] +name = "futures-util" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +dependencies = [ + "ahash", +] + +[[package]] +name = "hashlink" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ba4ff7128dee98c7dc9794b6a411377e1404dba1c97deb8d1a55297bd25d8af" +dependencies = [ + "hashbrown", +] + +[[package]] +name = "http" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "918d3568bebf352712bc2ef3d46a8bcf1a75b373be6539de198e9105cbbf9ce0" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f41fd6a08e4d4ec69df65976da761afd5ad5e58a9d4acb46bd1c953a9e3ff2" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "hyper" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d22053281f852e11534f5198498373cbb59295120a20771d90f7ed1897490a72" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "http", + "http-body", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "bytes", + "http", + "http-body", + "hyper", + "pin-project-lite", + "tokio", + "tower-service", +] + +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e0c1080212aad755ea003d18543e8768dd432c48819efd73a7bf1e39b7a5a3a" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "landlock" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cca98e95f35b29d469dade6724c6f96cec9236640f745a0e99b0334ec320ab1" +dependencies = [ + "enumflags2", + "libc", + "thiserror", +] + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "libsqlite3-sys" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e99fb7a497b1e3339bc746195567ed8d3e24945ecd636e3619d20b9de9e9149" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + +[[package]] +name = "mio" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" +dependencies = [ + "libc", + "wasi", + "windows-sys", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom", +] + +[[package]] +name = "rusqlite" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7753b721174eb8ff87a9a0e799e2d7bc3749323e773db92e0984debb00019d6e" +dependencies = [ + "bitflags", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "smallvec", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" + +[[package]] +name = "thiserror" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "tokio" +version = "1.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +dependencies = [ + "libc", + "mio", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys", +] + +[[package]] +name = "tokio-macros" +version = "2.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "log", + "pin-project-lite", + "tracing-core", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.127" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b70935747edd64d89de3efa29d73789b806c15798f8e7dca4d8ac356b50ce70" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.127" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77775f8f3f7217702089053b94958f8f54061a3f663417df76e19cbdcca29bc1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.127" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e11d33f857dc2fb11b8bc75aee111aa9cbeb12cd9f25efd3d4c2a3dd4e235284" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 2.0.119", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.127" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ef64dbcc55df09c7e5a46182d181c2cfa3e925f3da937ea764728b4bbb9dcbf" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zerocopy" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/dify-agent-runtime/rust/Cargo.toml b/dify-agent-runtime/rust/Cargo.toml new file mode 100644 index 00000000000000..fc3e4965515950 --- /dev/null +++ b/dify-agent-runtime/rust/Cargo.toml @@ -0,0 +1,39 @@ +[package] +name = "dify-agent-runtime" +version = "0.1.0" +edition = "2024" +description = "Rust implementation of Dify's sandbox shell runtime" +license = "Apache-2.0" + +[dependencies] +axum = "0.8.8" +chrono = { version = "0.4.45", features = ["clock", "serde"] } +rand = "0.9.5" +rusqlite = { version = "0.32.1", features = ["bundled"] } +serde = { version = "1.0.229", features = ["derive"] } +serde_json = "1.0.151" +tokio = { version = "1.49.0", features = ["macros", "rt-multi-thread", "signal", "sync", "time"] } + +[target.'cfg(target_os = "linux")'.dependencies] +landlock = "0.4.4" + +[[bin]] +name = "shellctl" +path = "src/bin/shellctl.rs" + +[[bin]] +name = "shellctl-runner" +path = "src/bin/shellctl-runner.rs" + +[[bin]] +name = "shellctl-runner-exit" +path = "src/bin/shellctl-runner-exit.rs" + +[[bin]] +name = "shellctl-sanitize-pty" +path = "src/bin/shellctl-sanitize-pty.rs" + +[profile.release] +codegen-units = 1 +lto = "thin" +strip = "symbols" diff --git a/dify-agent-runtime/rust/README.md b/dify-agent-runtime/rust/README.md new file mode 100644 index 00000000000000..71e231ddb4c1ae --- /dev/null +++ b/dify-agent-runtime/rust/README.md @@ -0,0 +1,54 @@ +# Rust shell runtime + +This directory contains the Rust implementation of the Dify shell runtime. +It keeps the Go runtime's HTTP contract and artifact layout so it can be +introduced without changing the Python agent sandbox. + +## Binaries + +- `shellctl` ā€” HTTP server for tmux-backed jobs. +- `shellctl-runner` ā€” gated job runner with optional Landlock isolation. +- `shellctl-runner-exit` ā€” idempotent SQLite exit recorder kept as a + compatibility/fallback utility; the server normally reconciles exit metadata + in its long-lived background thread. +- `shellctl-sanitize-pty` ā€” streaming ANSI/PTY output sanitizer. + +## Build and test + +```bash +cargo build --release --bins +cargo test +``` + +Landlock is compiled on Linux through the `landlock` crate. On macOS and +other unsupported platforms the runner logs a warning and continues without +filesystem isolation, matching the existing best-effort behavior. + +The Go implementation remains the default and fallback runtime during the +migration. The Rust server is built as a separate image and owns separate +SQLite, tmux, Home, Workspace, and Snapshot state. + +## Safe canary rollout + +Start the normal stack plus the opt-in Rust service: + +```bash +cd docker +DIFY_AGENT_LOCAL_SANDBOX_RUST_CANARY_PERCENT=1 \ + docker compose \ + -f docker-compose.yaml \ + -f docker-compose.rust-runtime.yaml \ + up -d --build +``` + +The agent backend makes a deterministic decision only for a new, unpinned +Binding. It sends a bounded health preflight to Rust and assigns that Binding +to Go if the preflight fails. Once a mutating request has reached a runtime it +is never replayed against the other implementation. Rust-owned refs use a +`rust+` prefix and remain pinned to Rust for their full lifecycle; Go refs keep +their existing representation and remain compatible with a Go-only rollback. + +To roll back admission without stranding Rust-owned resources, leave +`DIFY_AGENT_LOCAL_SANDBOX_RUST_ENDPOINT` configured and set +`DIFY_AGENT_LOCAL_SANDBOX_RUST_CANARY_PERCENT=0`. Remove the Rust service only +after all `rust+` Bindings, Workspaces, and Home Snapshots have drained. diff --git a/dify-agent-runtime/rust/src/bin/shellctl-runner-exit.rs b/dify-agent-runtime/rust/src/bin/shellctl-runner-exit.rs new file mode 100644 index 00000000000000..133a4349923a3c --- /dev/null +++ b/dify-agent-runtime/rust/src/bin/shellctl-runner-exit.rs @@ -0,0 +1,56 @@ +use std::path::PathBuf; + +fn main() { + let mut state = None::; + let mut job = None::; + let mut code = 0_i32; + let mut ended = None::; + let mut busy = 5000_u64; + let a: Vec = std::env::args().skip(1).collect(); + let mut i = 0; + while i < a.len() { + match a[i].as_str() { + "--state-dir" => { + i += 1; + state = a.get(i).map(PathBuf::from) + } + "--job-id" => { + i += 1; + job = a.get(i).cloned() + } + "--exit-code" => { + i += 1; + code = a.get(i).and_then(|x| x.parse().ok()).unwrap_or(0) + } + "--ended-at" => { + i += 1; + ended = a.get(i).cloned() + } + "--sqlite-busy-timeout-ms" => { + i += 1; + busy = a.get(i).and_then(|x| x.parse().ok()).unwrap_or(5000) + } + x => { + eprintln!("unknown flag: {x}"); + std::process::exit(2) + } + } + i += 1; + } + let Some(state) = state else { + eprintln!("--state-dir is required"); + std::process::exit(1) + }; + let Some(job) = job else { + eprintln!("--job-id is required"); + std::process::exit(1) + }; + let Some(ended) = ended else { + eprintln!("--ended-at is required"); + std::process::exit(1) + }; + if let Err(e) = dify_agent_runtime::record_runner_exit(&state, &job, code, &ended, busy) { + eprintln!("runner-exit: {e}"); + std::process::exit(1) + } +} diff --git a/dify-agent-runtime/rust/src/bin/shellctl-runner.rs b/dify-agent-runtime/rust/src/bin/shellctl-runner.rs new file mode 100644 index 00000000000000..f06bd3c65e0971 --- /dev/null +++ b/dify-agent-runtime/rust/src/bin/shellctl-runner.rs @@ -0,0 +1,5 @@ +fn main() { + std::process::exit(dify_agent_runtime::run_runner( + &std::env::args().skip(1).collect::>(), + )); +} diff --git a/dify-agent-runtime/rust/src/bin/shellctl-sanitize-pty.rs b/dify-agent-runtime/rust/src/bin/shellctl-sanitize-pty.rs new file mode 100644 index 00000000000000..ed4a20c4f820b0 --- /dev/null +++ b/dify-agent-runtime/rust/src/bin/shellctl-sanitize-pty.rs @@ -0,0 +1,16 @@ +fn main() { + let a: Vec = std::env::args().skip(1).collect(); + let mut ready = None; + let mut i = 0; + while i < a.len() { + if a[i] == "--ready-file" { + i += 1; + ready = a.get(i).map(std::path::PathBuf::from); + } + i += 1; + } + if let Err(e) = dify_agent_runtime::run_sanitizer(ready.as_deref()) { + eprintln!("sanitize-pty: {e}"); + std::process::exit(1) + } +} diff --git a/dify-agent-runtime/rust/src/bin/shellctl.rs b/dify-agent-runtime/rust/src/bin/shellctl.rs new file mode 100644 index 00000000000000..add13c7a40ccfa --- /dev/null +++ b/dify-agent-runtime/rust/src/bin/shellctl.rs @@ -0,0 +1,44 @@ +use dify_agent_runtime::{Config, serve}; + +#[tokio::main(flavor = "current_thread")] +async fn main() { + let mut config = Config::default(); + let args: Vec = std::env::args().collect(); + if args.get(1).map(String::as_str) != Some("serve") { + eprintln!("Usage: shellctl serve [--listen HOST:PORT] [--state-dir PATH] [--token TOKEN]"); + return; + } + let mut i = 2; + while i < args.len() { + match args[i].as_str() { + "--listen" => { + i += 1; + if let Some(v) = args.get(i) { + config.listen = v.clone(); + } + } + "--state-dir" => { + i += 1; + if let Some(v) = args.get(i) { + config.state_dir = v.into(); + config.runtime_dir = config.state_dir.join("runtime"); + } + } + "--token" => { + i += 1; + if let Some(v) = args.get(i) { + config.auth_token = v.clone(); + } + } + other => { + eprintln!("unknown flag: {other}"); + std::process::exit(2); + } + } + i += 1; + } + if let Err(e) = serve(config).await { + eprintln!("shellctl: {e}"); + std::process::exit(1); + } +} diff --git a/dify-agent-runtime/rust/src/lib.rs b/dify-agent-runtime/rust/src/lib.rs new file mode 100644 index 00000000000000..5d3283c20edb9a --- /dev/null +++ b/dify-agent-runtime/rust/src/lib.rs @@ -0,0 +1,2257 @@ +//! Rust implementation of the shellctl runtime. +//! +//! The public HTTP contract intentionally mirrors the Go runtime so the +//! Python agent sandbox can switch implementations without a protocol change. + +use axum::{ + Json, Router, + body::Body, + extract::{Path, Query, State}, + http::{Request, StatusCode}, + middleware::{self, Next}, + response::{IntoResponse, Response}, + routing::{get, post}, +}; +use chrono::{SecondsFormat, Utc}; +use rand::RngCore; +use rusqlite::{Connection, OptionalExtension, params}; +use serde::{Deserialize, Serialize}; +use std::{ + collections::{HashMap, HashSet}, + env, + ffi::{OsStr, OsString}, + fs::{self, File}, + io::{self, Read, Write}, + os::unix::process::CommandExt, + path::{Path as FsPath, PathBuf}, + process::{Command, Stdio}, + sync::{Arc, Mutex}, + thread, + time::{Duration, Instant}, +}; + +pub const DEFAULT_LISTEN: &str = "127.0.0.1:8765"; +const DEFAULT_OUTPUT_LIMIT: usize = 16 * 1024; +const MAX_OUTPUT_LIMIT: usize = 512 * 1024; +const DEFAULT_IDLE_FLUSH: f64 = 0.5; +const POLL_INTERVAL: Duration = Duration::from_millis(50); +// Completion artifacts are local files and are cheap to inspect. Poll them +// more frequently than the tmux liveness fallback so short jobs do not pay a +// full 50 ms scheduling quantum while still limiting tmux subprocess churn. +const OUTPUT_WAIT_INTERVAL: Duration = Duration::from_millis(5); +// These only apply while a job is being started, before user code can run. +const PIPE_READY_POLL_INTERVAL: Duration = Duration::from_millis(5); +const START_GATE_POLL_INTERVAL: Duration = Duration::from_millis(5); + +#[derive(Debug, Clone)] +pub struct Config { + pub listen: String, + pub auth_token: String, + pub state_dir: PathBuf, + pub runtime_dir: PathBuf, + pub default_cwd: PathBuf, + pub default_timeout: Duration, + pub max_wait_timeout: Duration, + pub gc_retention: Duration, + pub gc_interval: Duration, + pub pipe_monitor_interval: Duration, + pub pipe_ready_timeout: Duration, + pub sqlite_busy_timeout_ms: u64, + pub terminal_cols: i32, + pub terminal_rows: i32, + pub default_output_limit: usize, + pub max_output_limit: usize, + pub terminate_grace: f64, +} + +impl Default for Config { + fn default() -> Self { + let home = env::var_os("HOME") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from("/")); + let state_dir = if let Some(xdg) = env::var_os("XDG_DATA_HOME") { + PathBuf::from(xdg).join("shellctl") + } else { + home.join(".local/share/shellctl") + }; + Self { + listen: DEFAULT_LISTEN.into(), + auth_token: env::var("SHELLCTL_AUTH_TOKEN").unwrap_or_default(), + runtime_dir: state_dir.join("runtime"), + state_dir, + default_cwd: home, + default_timeout: Duration::from_secs(30), + max_wait_timeout: Duration::from_secs(600), + gc_retention: Duration::from_secs(300), + gc_interval: Duration::from_secs(60), + pipe_monitor_interval: Duration::from_secs(1), + pipe_ready_timeout: Duration::from_secs(10), + sqlite_busy_timeout_ms: 5000, + terminal_cols: 200, + terminal_rows: 50, + default_output_limit: DEFAULT_OUTPUT_LIMIT, + max_output_limit: MAX_OUTPUT_LIMIT, + terminate_grace: 10.0, + } + } +} + +impl Config { + pub fn jobs_dir(&self) -> PathBuf { + self.state_dir.join("jobs") + } + pub fn db_path(&self) -> PathBuf { + self.state_dir.join("shellctl.db") + } + pub fn tmux_socket(&self) -> PathBuf { + self.runtime_dir.join("tmux.sock") + } + pub fn runner_path(&self) -> PathBuf { + self.runtime_dir.join("bin/shellctl-runner") + } +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum Status { + Created, + Starting, + Running, + Exited, + Terminated, + Failed, + Lost, +} + +impl Status { + fn as_str(self) -> &'static str { + match self { + Self::Created => "created", + Self::Starting => "starting", + Self::Running => "running", + Self::Exited => "exited", + Self::Terminated => "terminated", + Self::Failed => "failed", + Self::Lost => "lost", + } + } + fn terminal(self) -> bool { + matches!( + self, + Self::Exited | Self::Terminated | Self::Failed | Self::Lost + ) + } +} + +impl TryFrom<&str> for Status { + type Error = RuntimeError; + fn try_from(s: &str) -> Result { + match s { + "created" => Ok(Self::Created), + "starting" => Ok(Self::Starting), + "running" => Ok(Self::Running), + "exited" => Ok(Self::Exited), + "terminated" => Ok(Self::Terminated), + "failed" => Ok(Self::Failed), + "lost" => Ok(Self::Lost), + _ => Err(RuntimeError::internal(format!("unknown job status: {s}"))), + } + } +} + +#[derive(Debug, Clone)] +struct Job { + id: String, + script_path: String, + output_path: String, + cwd: String, + cols: i32, + rows: i32, + status: Status, + session_name: String, + pane_target: String, + exit_code: Option, + _reason: Option, + _message: Option, + created_at: String, + started_at: Option, + ended_at: Option, + _updated_at: String, +} + +#[derive(Debug)] +pub struct RuntimeError { + pub status: u16, + pub code: String, + pub message: String, +} + +impl RuntimeError { + fn new(status: u16, code: impl Into, message: impl Into) -> Self { + Self { + status, + code: code.into(), + message: message.into(), + } + } + fn not_found() -> Self { + Self::new(404, "job_not_found", "Unknown job id") + } + fn internal(message: impl Into) -> Self { + Self::new(500, "internal_error", message) + } +} + +impl std::fmt::Display for RuntimeError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "[{}] {}: {}", self.status, self.code, self.message) + } +} +impl std::error::Error for RuntimeError {} + +#[derive(Debug, Serialize)] +struct ErrorBody { + error: ErrorDetail, +} +#[derive(Debug, Serialize)] +struct ErrorDetail { + code: String, + message: String, +} + +impl IntoResponse for RuntimeError { + fn into_response(self) -> Response { + ( + StatusCode::from_u16(self.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), + Json(ErrorBody { + error: ErrorDetail { + code: self.code, + message: self.message, + }, + }), + ) + .into_response() + } +} + +#[derive(Debug, Deserialize)] +pub struct RunJobRequest { + pub script: String, + pub cwd: Option, + pub env: Option>, + pub terminal: Option, + pub timeout: Option, + pub output_limit: Option, + pub idle_flush_seconds: Option, +} +#[derive(Debug, Deserialize)] +pub struct TerminalSize { + pub cols: i32, + pub rows: i32, +} +#[derive(Debug, Deserialize)] +pub struct WaitJobRequest { + pub timeout: f64, + pub offset: usize, + pub output_limit: Option, + pub idle_flush_seconds: Option, +} +#[derive(Debug, Deserialize)] +pub struct InputJobRequest { + pub text: String, + pub timeout: Option, + pub offset: usize, + pub output_limit: Option, + pub idle_flush_seconds: Option, +} +#[derive(Debug, Deserialize)] +pub struct TerminateJobRequest { + pub grace_seconds: Option, +} +#[derive(Debug, Deserialize)] +pub struct ListQuery { + pub status: Option, + pub limit: Option, +} +#[derive(Debug, Deserialize)] +pub struct TailQuery { + pub output_limit: Option, +} + +#[derive(Debug, Serialize)] +pub struct JobResult { + pub job_id: String, + pub done: bool, + pub status: Status, + pub exit_code: Option, + pub output_path: String, + pub output: String, + pub offset: usize, + pub truncated: bool, +} +#[derive(Debug, Serialize)] +pub struct JobStatusView { + pub job_id: String, + pub status: Status, + pub done: bool, + pub exit_code: Option, + pub created_at: String, + pub started_at: Option, + pub ended_at: Option, + pub offset: usize, +} +#[derive(Debug, Serialize)] +pub struct JobInfo { + pub job_id: String, + pub status: Status, + pub created_at: String, + pub started_at: Option, + pub ended_at: Option, +} +#[derive(Debug, Serialize)] +pub struct ListJobsResponse { + pub jobs: Vec, +} +#[derive(Debug, Serialize)] +pub struct DeleteJobResponse { + pub job_id: String, + pub deleted: bool, +} +#[derive(Debug, Serialize)] +pub struct HealthResponse { + pub status: &'static str, +} + +fn timestamp() -> String { + Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true) +} +fn job_id() -> String { + let mut bytes = [0_u8; 8]; + rand::rng().fill_bytes(&mut bytes); + bytes.iter().map(|b| format!("{b:02x}")).collect() +} +fn session_name(id: &str) -> String { + format!("shellctl-{id}") +} +fn pane_target(id: &str) -> String { + format!("{}:0.0", session_name(id)) +} + +fn db_open(path: &FsPath, busy_timeout_ms: u64) -> Result { + let conn = db_connect(path, busy_timeout_ms, true)?; + conn.execute_batch("CREATE TABLE IF NOT EXISTS jobs (job_id TEXT PRIMARY KEY, script_path TEXT NOT NULL, output_path TEXT NOT NULL, cwd TEXT NOT NULL, terminal_cols INTEGER NOT NULL DEFAULT 200, terminal_rows INTEGER NOT NULL DEFAULT 50, status TEXT NOT NULL DEFAULT 'created', session_name TEXT NOT NULL, pane_target TEXT NOT NULL, exit_code INTEGER, reason TEXT, message TEXT, created_at TEXT NOT NULL, started_at TEXT, ended_at TEXT, updated_at TEXT NOT NULL);") + .map_err(|e| RuntimeError::internal(format!("init sqlite schema: {e}")))?; + Ok(conn) +} + +// Runner-exit is intentionally a short-lived helper. The server owns schema +// creation, so reopening a known database must not pay for DDL on every job. +fn db_connect( + path: &FsPath, + busy_timeout_ms: u64, + ensure_wal: bool, +) -> Result { + let conn = + Connection::open(path).map_err(|e| RuntimeError::internal(format!("open sqlite: {e}")))?; + conn.busy_timeout(Duration::from_millis(busy_timeout_ms)) + .map_err(|e| RuntimeError::internal(e.to_string()))?; + if ensure_wal { + conn.pragma_update(None, "journal_mode", "WAL") + .map_err(|e| RuntimeError::internal(e.to_string()))?; + } + Ok(conn) +} + +fn read_job(conn: &Connection, id: &str) -> Result { + conn.query_row("SELECT job_id,script_path,output_path,cwd,terminal_cols,terminal_rows,status,session_name,pane_target,exit_code,reason,message,created_at,started_at,ended_at,updated_at FROM jobs WHERE job_id=?1", [id], |r| { + let status: String = r.get(6)?; + Ok(Job { id: r.get(0)?, script_path: r.get(1)?, output_path: r.get(2)?, cwd: r.get(3)?, cols: r.get(4)?, rows: r.get(5)?, status: Status::try_from(status.as_str()).map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?, session_name: r.get(7)?, pane_target: r.get(8)?, exit_code: r.get(9)?, _reason: r.get(10)?, _message: r.get(11)?, created_at: r.get(12)?, started_at: r.get(13)?, ended_at: r.get(14)?, _updated_at: r.get(15)? }) + }).optional().map_err(|e| RuntimeError::internal(format!("query job: {e}")))?.ok_or_else(RuntimeError::not_found) +} + +fn transition( + conn: &Connection, + id: &str, + target: Status, + allowed: &[Status], + reason: Option<&str>, + message: Option<&str>, + ended_at: Option<&str>, +) -> Result { + let now = timestamp(); + let marks = std::iter::repeat_n("?", allowed.len()) + .collect::>() + .join(","); + let mut sql = "UPDATE jobs SET status=?1, updated_at=?2, reason=?3, message=?4".to_string(); + let mut values: Vec = vec![ + target.as_str().into(), + now.clone(), + reason.unwrap_or("").into(), + message.unwrap_or("").into(), + ]; + let mut next = 5; + if matches!(target, Status::Starting | Status::Running) { + sql.push_str(&format!(", started_at=COALESCE(started_at, ?{next})")); + values.push(now.clone()); + next += 1; + } + if target.terminal() { + sql.push_str(&format!( + ", ended_at=COALESCE(ended_at, ?{next}), exit_code=COALESCE(exit_code, 0)" + )); + values.push(ended_at.unwrap_or(&now).into()); + next += 1; + } + sql.push_str(&format!(" WHERE job_id=?{next} AND status IN ({marks})")); + values.push(id.into()); + values.extend(allowed.iter().map(|s| s.as_str().into())); + let mut stmt = conn + .prepare(&sql) + .map_err(|e| RuntimeError::internal(format!("prepare transition: {e}")))?; + let mut params: Vec<&dyn rusqlite::ToSql> = + values.iter().map(|s| s as &dyn rusqlite::ToSql).collect(); + stmt.execute(rusqlite::params_from_iter(params.drain(..))) + .map_err(|e| RuntimeError::internal(format!("transition status: {e}")))?; + read_job(conn, id) +} + +fn db_record_runner_exit( + conn: &Connection, + id: &str, + code: i32, + ended_at: &str, +) -> Result<(), RuntimeError> { + let n = conn.execute("UPDATE jobs SET status='exited', exit_code=?1, ended_at=?2, updated_at=?2, reason=NULL, message=NULL WHERE job_id=?3 AND status IN ('created','starting','running')", params![code, ended_at, id]).map_err(|e| RuntimeError::internal(format!("update job: {e}")))?; + if n > 0 { + return Ok(()); + } + // A zero-row CAS is expected for an already-terminal job, but must still + // report a genuinely unknown id just like the Go implementation. + let exists: Option = conn + .query_row("SELECT 1 FROM jobs WHERE job_id=?1", [id], |row| row.get(0)) + .optional() + .map_err(|e| RuntimeError::internal(format!("query job existence: {e}")))?; + exists.map(|_| ()).ok_or_else(RuntimeError::not_found) +} + +#[derive(Debug, Serialize)] +pub struct OutputWindow { + pub output: String, + pub offset: usize, + pub truncated: bool, +} + +fn bounded_positive(value: Option, default: usize, maximum: usize) -> usize { + value + .filter(|value| *value > 0) + .unwrap_or(default) + .min(maximum) +} + +fn valid_prefix(data: &[u8], max: usize) -> usize { + let end = max.min(data.len()); + (0..=end) + .rev() + .find(|i| std::str::from_utf8(&data[..*i]).is_ok()) + .unwrap_or(0) +} +fn read_window(path: &FsPath, offset: usize, limit: usize) -> Result { + let data = match fs::read(path) { + Ok(v) => v, + Err(e) if e.kind() == io::ErrorKind::NotFound => { + if offset == 0 { + return Ok(OutputWindow { + output: String::new(), + offset: 0, + truncated: false, + }); + } + return Err(RuntimeError::new( + 400, + "invalid_offset", + "offset exceeds current file size 0", + )); + } + Err(e) => return Err(RuntimeError::internal(e.to_string())), + }; + if offset > data.len() { + return Err(RuntimeError::new( + 400, + "invalid_offset", + format!("offset {offset} exceeds current file size {}", data.len()), + )); + } + if offset == data.len() { + return Ok(OutputWindow { + output: String::new(), + offset, + truncated: false, + }); + } + let buf = &data[offset..(data.len()).min(offset + limit + 4)]; + let shift = buf + .iter() + .position(|b| (*b & 0xc0) != 0x80) + .unwrap_or(buf.len()); + let payload = &buf[shift..]; + let mut consumed = valid_prefix(payload, limit.saturating_sub(shift)); + if consumed == 0 && !payload.is_empty() { + consumed = payload + .iter() + .enumerate() + .find_map(|(i, b)| { + if i > 0 && (*b & 0xc0) != 0x80 { + Some(i) + } else { + None + } + }) + .unwrap_or(payload.len()) + .min(4); + while consumed > 0 && std::str::from_utf8(&payload[..consumed]).is_err() { + consumed -= 1; + } + } + let next = offset + shift + consumed; + Ok(OutputWindow { + output: String::from_utf8_lossy(&payload[..consumed]).into_owned(), + offset: next, + truncated: next < data.len(), + }) +} +fn tail_window(path: &FsPath, limit: usize) -> Result { + let data = match fs::read(path) { + Ok(v) => v, + Err(e) if e.kind() == io::ErrorKind::NotFound => Vec::new(), + Err(e) => return Err(RuntimeError::internal(e.to_string())), + }; + if data.is_empty() { + return Ok(OutputWindow { + output: String::new(), + offset: 0, + truncated: false, + }); + } + let start = data.len().saturating_sub(limit); + let mut pos = start; + while pos < data.len() && (data[pos] & 0xc0) == 0x80 { + pos += 1; + } + let payload = &data[pos..]; + let valid = valid_prefix(payload, payload.len()); + Ok(OutputWindow { + output: String::from_utf8_lossy(&payload[..valid]).into_owned(), + offset: pos + valid, + truncated: false, + }) +} + +#[derive(Clone)] +pub struct AppState { + inner: Arc, +} +pub struct Runtime { + config: Config, + db: Mutex, + starting: Mutex>, +} + +impl Runtime { + pub fn initialize(config: Config) -> Result, RuntimeError> { + fs::create_dir_all(&config.state_dir) + .map_err(|e| RuntimeError::internal(format!("create state dir: {e}")))?; + fs::create_dir_all(&config.runtime_dir) + .map_err(|e| RuntimeError::internal(format!("create runtime dir: {e}")))?; + fs::create_dir_all(config.jobs_dir()) + .map_err(|e| RuntimeError::internal(format!("create jobs dir: {e}")))?; + fs::create_dir_all(config.runner_path().parent().unwrap()) + .map_err(|e| RuntimeError::internal(format!("create runner dir: {e}")))?; + let conn = db_open(&config.db_path(), config.sqlite_busy_timeout_ms)?; + let runtime = Arc::new(Self { + config: config.clone(), + db: Mutex::new(conn), + starting: Mutex::new(HashSet::new()), + }); + runtime.install_runner(); + runtime.start_tmux_server()?; + runtime.reconcile_startup()?; + Ok(runtime) + } + + /// Keep completion reconciliation inside the long-lived server process. + /// + /// The pipe still writes exit metadata atomically, so `wait` and `list` + /// can reconcile immediately. This background pass covers jobs that have + /// no active API caller without spawning one SQLite helper process per job. + fn start_reconciler(runtime: Arc) { + let interval = runtime + .config + .pipe_monitor_interval + .max(Duration::from_millis(50)); + let _ = thread::Builder::new() + .name("shellctl-reconciler".into()) + .spawn(move || { + loop { + thread::sleep(interval); + let _ = runtime.reconcile_artifacts(); + } + }); + } + fn install_runner(&self) { + let dst = self.config.runner_path(); + let _ = fs::remove_file(&dst); + if let Ok(exe) = env::current_exe() { + let candidate = exe + .parent() + .unwrap_or(FsPath::new(".")) + .join("shellctl-runner"); + if candidate.exists() { + let _ = std::os::unix::fs::symlink(candidate, dst); + } + } + } + fn output_path(&self, job: &Job) -> PathBuf { + self.config.state_dir.join(&job.output_path) + } + fn job_view(&self, job: &Job) -> JobStatusView { + let offset = fs::metadata(self.output_path(job)) + .map(|m| m.len() as usize) + .unwrap_or(0); + JobStatusView { + job_id: job.id.clone(), + status: job.status, + done: job.status.terminal(), + exit_code: job.exit_code, + created_at: job.created_at.clone(), + started_at: job.started_at.clone(), + ended_at: job.ended_at.clone(), + offset, + } + } + fn tmux(&self, args: &[&str]) -> Result<(i32, String, String), RuntimeError> { + let mut cmd = Command::new("tmux"); + cmd.arg("-S").arg(self.config.tmux_socket()); + for a in args { + cmd.arg(a); + } + cmd.env_remove("TMUX"); + let out = cmd + .output() + .map_err(|e| RuntimeError::internal(format!("exec tmux: {e}")))?; + Ok(( + out.status.code().unwrap_or(125), + String::from_utf8_lossy(&out.stdout).into_owned(), + String::from_utf8_lossy(&out.stderr).into_owned(), + )) + } + fn tmux_ok(&self, args: &[&str]) -> Result<(), RuntimeError> { + let (code, _, err) = self.tmux(args)?; + if code != 0 { + return Err(RuntimeError::new(500, "tmux_error", err.trim().to_string())); + } + Ok(()) + } + fn start_tmux_server(&self) -> Result<(), RuntimeError> { + self.tmux_ok(&["start-server", ";", "set-option", "-g", "exit-empty", "off"]) + } + fn session_exists(&self, id: &str) -> Result { + let (code, out, err) = self.tmux(&["list-sessions", "-F", "#{session_name}"])?; + if code != 0 && !tmux_missing(&err) { + return Err(RuntimeError::new(500, "tmux_error", err.trim())); + } + Ok(out.lines().any(|line| line.trim() == session_name(id))) + } + fn pipe_active(&self, id: &str) -> Result, RuntimeError> { + let (code, out, err) = self.tmux(&[ + "display-message", + "-p", + "-t", + &pane_target(id), + "#{pane_pipe}", + ])?; + if code != 0 { + if tmux_missing(&err) { + return Ok(None); + } + return Err(RuntimeError::new(500, "tmux_error", err.trim())); + } + Ok(Some(out.trim() == "1")) + } + fn live_view(&self, id: &str) -> Result { + let (session, pipe) = (self.session_exists(id)?, None); + let pipe = if session { self.pipe_active(id)? } else { pipe }; + let pipe_failed = self + .config + .jobs_dir() + .join(id) + .join(".pipe-failed") + .exists(); + let mut conn = self.db.lock().unwrap(); + let mut job = read_job(&conn, id)?; + if !job.status.terminal() + && let Some((code, ended_at)) = drained_exit_metadata(&self.config.jobs_dir().join(id)) + { + let _ = db_record_runner_exit(&conn, id, code, &ended_at); + job = read_job(&conn, id)?; + } + if let Some(next_job) = materialize( + &mut conn, + &job, + session, + pipe, + pipe_failed, + self.starting.lock().unwrap().contains(id), + )? { + job = next_job; + } + Ok(self.job_view(&job)) + } + pub fn health() -> Json { + Json(HealthResponse { status: "ok" }) + } + pub fn run_job(&self, req: RunJobRequest) -> Result { + let cwd = req + .cwd + .map(PathBuf::from) + .unwrap_or_else(|| self.config.default_cwd.clone()); + if !cwd.is_dir() { + return Err(RuntimeError::new( + 400, + "invalid_cwd", + format!("cwd is not a directory: {}", cwd.display()), + )); + } + let cwd = fs::canonicalize(cwd) + .map_err(|e| RuntimeError::new(400, "invalid_cwd", e.to_string()))?; + let cols = req + .terminal + .as_ref() + .map(|x| x.cols) + .unwrap_or(self.config.terminal_cols); + let rows = req + .terminal + .as_ref() + .map(|x| x.rows) + .unwrap_or(self.config.terminal_rows); + let created = timestamp(); + let (id, dir) = (0..20) + .find_map(|_| { + let id = job_id(); + let dir = self.config.jobs_dir().join(&id); + fs::create_dir(&dir).ok().map(|_| (id, dir)) + }) + .ok_or_else(|| { + RuntimeError::new( + 500, + "job_id_collision", + "Failed to allocate a unique job id", + ) + })?; + self.starting.lock().unwrap().insert(id.clone()); + let script = dir.join("script"); + let output = dir.join("output.log"); + let env_file = dir.join(".job-env.json"); + let write_result = (|| -> Result<(), RuntimeError> { + fs::write(&script, req.script.as_bytes()) + .map_err(|e| RuntimeError::internal(e.to_string()))?; + fs::write(&output, []).map_err(|e| RuntimeError::internal(e.to_string()))?; + fs::write( + &env_file, + serde_json::to_vec(&req.env.unwrap_or_default()).unwrap(), + ) + .map_err(|e| RuntimeError::internal(e.to_string()))?; + Ok(()) + })(); + if let Err(e) = write_result { + let _ = fs::remove_dir_all(&dir); + self.starting.lock().unwrap().remove(&id); + return Err(e); + } + let job = Job { + id: id.clone(), + script_path: format!("jobs/{id}/script"), + output_path: format!("jobs/{id}/output.log"), + cwd: cwd.display().to_string(), + cols, + rows, + status: Status::Created, + session_name: session_name(&id), + pane_target: pane_target(&id), + exit_code: None, + _reason: None, + _message: None, + created_at: created.clone(), + started_at: None, + ended_at: None, + _updated_at: created, + }; + let conn = self.db.lock().unwrap(); + conn.execute("INSERT INTO jobs (job_id,script_path,output_path,cwd,terminal_cols,terminal_rows,status,session_name,pane_target,created_at,updated_at) VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?10)", params![job.id,job.script_path,job.output_path,job.cwd,job.cols,job.rows,job.status.as_str(),job.session_name,job.pane_target,job.created_at]).map_err(|e| RuntimeError::internal(format!("insert job: {e}")))?; + transition( + &conn, + &id, + Status::Starting, + &[Status::Created], + None, + None, + None, + )?; + drop(conn); + let start = self.start_tmux_job(&id, &dir, &cwd, cols, rows); + self.starting.lock().unwrap().remove(&id); + if let Err(e) = start { + let conn = self.db.lock().unwrap(); + let _ = transition( + &conn, + &id, + Status::Failed, + &[Status::Created, Status::Starting, Status::Running], + Some("start_failed"), + Some(&e.to_string()), + None, + ); + self.kill_session(&id); + } + let timeout = req + .timeout + .filter(|timeout| *timeout > 0.0) + .unwrap_or(self.config.default_timeout.as_secs_f64()); + self.wait_job( + &id, + WaitJobRequest { + timeout, + offset: 0, + output_limit: req.output_limit, + idle_flush_seconds: req.idle_flush_seconds, + }, + ) + } + fn start_tmux_job( + &self, + id: &str, + dir: &FsPath, + cwd: &FsPath, + cols: i32, + rows: i32, + ) -> Result<(), RuntimeError> { + // The dedicated server is initialized once with exit-empty disabled, + // matching the current Go runtime and avoiding per-job bootstrap. + let runner = self.config.runner_path().to_string_lossy().into_owned(); + let d = dir.to_string_lossy().into_owned(); + let c = cwd.to_string_lossy().into_owned(); + let ready = dir.join(".pipe-ready"); + let output = shell_quote(&dir.join("output.log").to_string_lossy()); + let sanitizer = env::var("SHELLCTL_SANITIZE_COMMAND") + .unwrap_or_else(|_| "shellctl-sanitize-pty".into()); + let error_log = shell_quote(&dir.join("pipe-error.log").to_string_lossy()); + let drained = shell_quote(&dir.join(".pipe-drained").to_string_lossy()); + let failed = shell_quote(&dir.join(".pipe-failed").to_string_lossy()); + let pipe = format!( + "{} --ready-file {} >> {} 2> {}; sanitize_status=$?; if [ \"$sanitize_status\" -eq 0 ]; then : > {}; else : > {}; fi; exit \"$sanitize_status\"", + shell_quote(&sanitizer), + shell_quote(&ready.to_string_lossy()), + output, + error_log, + drained, + failed, + ); + let session = session_name(id); + let pane = pane_target(id); + let runner_command = format!( + "{} {} {} {}", + shell_quote(&runner), + shell_quote(&d), + shell_quote(id), + shell_quote(&c) + ); + let cols = cols.to_string(); + let rows = rows.to_string(); + let tmux_args = [ + "-f", + "/dev/null", + "new-session", + "-d", + "-s", + &session, + "-x", + &cols, + "-y", + &rows, + &runner_command, + ";", + "pipe-pane", + "-o", + "-t", + &pane, + &pipe, + ]; + let started = self.tmux(&tmux_args)?; + if started.0 != 0 { + return Err(RuntimeError::new( + 500, + "tmux_new_session_failed", + started.2.trim(), + )); + } + let deadline = Instant::now() + self.config.pipe_ready_timeout; + while Instant::now() < deadline { + if ready.exists() && self.pipe_active(id)?.unwrap_or(false) { + let _ = fs::write(dir.join("start-gate"), []); + let conn = self.db.lock().unwrap(); + transition( + &conn, + id, + Status::Running, + &[Status::Starting], + None, + None, + None, + )?; + return Ok(()); + } + thread::sleep(PIPE_READY_POLL_INTERVAL); + } + Err(RuntimeError::new( + 500, + "pipe_failed", + "timed out waiting for pipe ready", + )) + } + fn wait_job(&self, id: &str, req: WaitJobRequest) -> Result { + let job = { + let conn = self.db.lock().unwrap(); + read_job(&conn, id)? + }; + let path = self.output_path(&job); + let limit = bounded_positive( + req.output_limit, + self.config.default_output_limit, + self.config.max_output_limit, + ); + let timeout = if req.timeout > 0.0 { + Duration::from_secs_f64(req.timeout.min(self.config.max_wait_timeout.as_secs_f64())) + } else { + Duration::ZERO + }; + let idle = Duration::from_secs_f64( + req.idle_flush_seconds + .unwrap_or(DEFAULT_IDLE_FLUSH) + .max(0.0), + ); + let deadline = Instant::now() + timeout; + let mut last_size = file_size(&path); + let mut saw = last_size > req.offset; + let mut growth = if saw { Some(Instant::now()) } else { None }; + // The caller already materialized this row and start_tmux_job only + // returns after the session and output pipe are live. Fast jobs can + // therefore complete via their atomic exit artifacts without two + // extra tmux CLI probes on every request. + let mut view = self.job_view(&job); + let mut next_runtime_probe = Instant::now() + POLL_INTERVAL; + let mut next_output_probe = Instant::now(); + loop { + if view.done { + let w = read_window(&path, req.offset, limit)?; + return Ok(self.result(&view, &job, w)); + } + if let Some((code, ended_at)) = drained_exit_metadata(&self.config.jobs_dir().join(id)) + { + let conn = self.db.lock().unwrap(); + let _ = db_record_runner_exit(&conn, id, code, &ended_at); + let job_view = read_job(&conn, id)?; + view = self.job_view(&job_view); + } else if Instant::now() >= next_runtime_probe { + next_runtime_probe = Instant::now() + POLL_INTERVAL; + view = self.live_view(id)?; + } + if view.done { + let w = read_window(&path, req.offset, limit)?; + return Ok(self.result(&view, &job, w)); + } + if Instant::now() >= next_output_probe { + next_output_probe = Instant::now() + OUTPUT_WAIT_INTERVAL; + let size = file_size(&path); + if size > last_size { + last_size = size; + if size > req.offset { + saw = true; + growth = Some(Instant::now()); + } + } + if size > req.offset { + let w = read_window(&path, req.offset, limit)?; + if w.truncated || (saw && growth.is_some_and(|t| t.elapsed() >= idle)) { + return Ok(self.result(&view, &job, w)); + } + } + } + if Instant::now() >= deadline { + let size = file_size(&path); + let w = if size > req.offset { + read_window(&path, req.offset, limit)? + } else { + OutputWindow { + output: String::new(), + offset: req.offset, + truncated: false, + } + }; + return Ok(self.result(&view, &job, w)); + } + thread::sleep(OUTPUT_WAIT_INTERVAL); + } + } + fn result(&self, view: &JobStatusView, job: &Job, w: OutputWindow) -> JobResult { + JobResult { + job_id: view.job_id.clone(), + done: view.done, + status: view.status, + exit_code: view.exit_code, + output_path: fs::canonicalize(self.output_path(job)) + .unwrap_or_else(|_| self.output_path(job)) + .display() + .to_string(), + output: w.output, + offset: w.offset, + truncated: w.truncated, + } + } + fn send_input(&self, id: &str, req: InputJobRequest) -> Result { + let view = self.live_view(id)?; + if view.done { + return Err(RuntimeError::new( + 409, + "job_not_running", + format!("Job {id} is already terminal"), + )); + } + let file = self.config.runtime_dir.join(format!("input-{id}")); + fs::write(&file, req.text).map_err(|e| RuntimeError::internal(e.to_string()))?; + let buf = format!("shellctl-in-{id}"); + let fp = file.to_string_lossy().into_owned(); + let load = self.tmux(&["load-buffer", "-b", &buf, &fp])?; + let _ = fs::remove_file(&file); + if load.0 != 0 { + return Err(RuntimeError::new(409, "tmux_target_missing", load.2.trim())); + } + let pasted = self.tmux(&["paste-buffer", "-t", &pane_target(id), "-b", &buf]); + let _ = self.tmux(&["delete-buffer", "-b", &buf]); + if let Err(e) = pasted { + return Err(e); + } + if pasted.unwrap().0 != 0 { + return Err(RuntimeError::new( + 409, + "tmux_target_missing", + "tmux target missing", + )); + } + self.wait_job( + id, + WaitJobRequest { + timeout: req.timeout.unwrap_or(30.0), + offset: req.offset, + output_limit: req.output_limit, + idle_flush_seconds: req.idle_flush_seconds, + }, + ) + } + fn kill_session(&self, id: &str) { + let _ = self.tmux(&["kill-session", "-t", &session_name(id)]); + } + fn terminate(&self, id: &str, grace: f64) -> Result { + let view = self.live_view(id)?; + if view.done { + self.kill_session(id); + return Ok(view); + } + { + let c = self.db.lock().unwrap(); + let _ = transition( + &c, + id, + Status::Terminated, + &[ + Status::Created, + Status::Starting, + Status::Running, + Status::Exited, + ], + None, + None, + None, + ); + } + let _ = self.tmux(&["send-keys", "-t", &pane_target(id), "C-c"]); + if grace > 0.0 { + thread::sleep(Duration::from_secs_f64(grace)); + } + self.kill_session(id); + self.live_view(id) + } + fn delete(&self, id: &str, force: bool, grace: f64) -> Result { + let v = self.live_view(id)?; + if !v.done && !force { + return Err(RuntimeError::new( + 409, + "job_running", + format!("Job {id} is still running"), + )); + } + if !v.done { + let _ = self.terminate(id, grace); + } + self.kill_session(id); + let c = self.db.lock().unwrap(); + let n = c + .execute("DELETE FROM jobs WHERE job_id=?1", [id]) + .map_err(|e| RuntimeError::internal(e.to_string()))?; + if n == 0 { + return Err(RuntimeError::not_found()); + } + let _ = fs::remove_dir_all(self.config.jobs_dir().join(id)); + Ok(DeleteJobResponse { + job_id: id.into(), + deleted: true, + }) + } + fn reconcile_startup(&self) -> Result<(), RuntimeError> { + let ids = self.list_ids()?; + for id in ids { + if let Ok(v) = self.live_view(&id) + && v.done + { + self.kill_session(&id); + } + } + Ok(()) + } + fn reconcile_artifacts(&self) -> Result<(), RuntimeError> { + let ids = self.list_nonterminal_ids()?; + for id in ids { + let dir = self.config.jobs_dir().join(&id); + if let Some((code, ended_at)) = drained_exit_metadata(&dir) { + let conn = self.db.lock().unwrap(); + let _ = db_record_runner_exit(&conn, &id, code, &ended_at); + drop(conn); + self.kill_session(&id); + } else if dir.join(".pipe-failed").exists() { + let conn = self.db.lock().unwrap(); + let _ = transition( + &conn, + &id, + Status::Failed, + &[Status::Created, Status::Starting, Status::Running], + Some("pipe_failed"), + Some("The tmux output pipe failed before completion."), + None, + ); + drop(conn); + self.kill_session(&id); + } + } + Ok(()) + } + fn list_nonterminal_ids(&self) -> Result, RuntimeError> { + let c = self.db.lock().unwrap(); + let mut s = c + .prepare("SELECT job_id FROM jobs WHERE status IN ('created','starting','running') ORDER BY created_at DESC") + .map_err(|e| RuntimeError::internal(e.to_string()))?; + let it = s + .query_map([], |r| r.get(0)) + .map_err(|e| RuntimeError::internal(e.to_string()))?; + Ok(it.filter_map(Result::ok).collect()) + } + fn list_ids(&self) -> Result, RuntimeError> { + let c = self.db.lock().unwrap(); + let mut s = c + .prepare("SELECT job_id FROM jobs ORDER BY created_at DESC") + .map_err(|e| RuntimeError::internal(e.to_string()))?; + let it = s + .query_map([], |r| r.get(0)) + .map_err(|e| RuntimeError::internal(e.to_string()))?; + Ok(it.filter_map(Result::ok).collect()) + } + fn list(&self, q: ListQuery) -> Result { + let mut out = Vec::new(); + for id in self.list_ids()? { + let v = self.live_view(&id)?; + if q.status.as_deref().is_some_and(|s| s != v.status.as_str()) { + continue; + } + out.push(JobInfo { + job_id: v.job_id, + status: v.status, + created_at: v.created_at, + started_at: v.started_at, + ended_at: v.ended_at, + }); + if out.len() >= bounded_positive(q.limit, 50, 200) { + break; + } + } + Ok(ListJobsResponse { jobs: out }) + } +} + +fn materialize( + conn: &mut Connection, + job: &Job, + session: bool, + pipe: Option, + pipe_failed: bool, + starting: bool, +) -> Result, RuntimeError> { + if job.status.terminal() { + return Ok(None); + } + if job.exit_code.is_some() { + return Ok(transition( + conn, + &job.id, + Status::Exited, + &[Status::Created, Status::Starting, Status::Running], + None, + None, + None, + ) + .ok()); + } else if !session && !starting { + return Ok(transition( + conn, + &job.id, + Status::Lost, + &[Status::Created, Status::Starting, Status::Running], + Some("tmux_session_missing"), + Some("The dedicated tmux session is no longer present."), + None, + ) + .ok()); + } else if session && pipe == Some(false) && pipe_failed && !starting { + return Ok(transition( + conn, + &job.id, + Status::Failed, + &[Status::Created, Status::Starting, Status::Running], + Some("pipe_failed"), + Some("The tmux output pipe stopped while the job was still running."), + None, + ) + .ok()); + } else if session && matches!(job.status, Status::Created | Status::Starting) && !starting { + return Ok(transition( + conn, + &job.id, + Status::Running, + &[Status::Created, Status::Starting], + None, + None, + None, + ) + .ok()); + } + Ok(None) +} +fn file_size(p: &FsPath) -> usize { + fs::metadata(p).map(|m| m.len() as usize).unwrap_or(0) +} +fn drained_exit_metadata(dir: &FsPath) -> Option<(i32, String)> { + exit_metadata(dir, ".pipe-drained") +} +fn exit_metadata(dir: &FsPath, marker: &str) -> Option<(i32, String)> { + if !dir.join(marker).exists() { + return None; + } + let code = fs::read_to_string(dir.join("runner-exit-code")) + .ok()? + .trim() + .parse() + .ok()?; + let ended = fs::read_to_string(dir.join("runner-ended-at")) + .ok()? + .trim() + .to_string(); + if ended.is_empty() { + None + } else { + Some((code, ended)) + } +} +fn tmux_missing(s: &str) -> bool { + let s = s.to_ascii_lowercase(); + [ + "can't find pane", + "can't find session", + "no server running", + "failed to connect", + "server exited unexpectedly", + ] + .iter() + .any(|x| s.contains(x)) +} +fn shell_quote(s: &str) -> String { + format!("'{}'", s.replace('\'', "'\\''")) +} + +pub fn router(state: Arc) -> Router { + let auth = state.config.auth_token.clone(); + Router::new() + .route("/healthz", get(health_handler)) + .route("/v1/jobs/run", post(run_handler)) + .route("/v1/jobs", get(list_handler)) + .route( + "/v1/jobs/{job_id}", + get(status_handler).delete(delete_handler), + ) + .route("/v1/jobs/{job_id}/wait", post(wait_handler)) + .route("/v1/jobs/{job_id}/log/tail", get(tail_handler)) + .route("/v1/jobs/{job_id}/input", post(input_handler)) + .route("/v1/jobs/{job_id}/terminate", post(terminate_handler)) + .with_state(AppState { inner: state }) + .layer(middleware::from_fn( + move |req: Request, next: Next| { + let token = auth.clone(); + async move { + if req.uri().path() == "/healthz" || token.is_empty() { + return next.run(req).await; + } + if req + .headers() + .get("authorization") + .and_then(|x| x.to_str().ok()) + != Some(&format!("Bearer {token}")) + { + return RuntimeError::new( + 401, + "unauthorized", + "Missing or invalid bearer token", + ) + .into_response(); + } + next.run(req).await + } + }, + )) +} +async fn health_handler() -> Json { + Runtime::health() +} +async fn run_handler( + State(s): State, + Json(req): Json, +) -> Result, RuntimeError> { + if req.script.is_empty() { + return Err(RuntimeError::new( + 400, + "invalid_request", + "script is required", + )); + } + if let Some(e) = &req.env { + for (k, v) in e { + if k.is_empty() { + return Err(RuntimeError::new( + 422, + "validation_error", + "env names must be non-empty", + )); + } + if k.contains('=') || k.contains('\0') || v.contains('\0') { + return Err(RuntimeError::new( + 422, + "validation_error", + "env entries must not contain NUL or '='", + )); + } + } + } + let r = s.inner.clone(); + tokio::task::spawn_blocking(move || r.run_job(req)) + .await + .map_err(|e| RuntimeError::internal(e.to_string()))? + .map(Json) +} +async fn wait_handler( + Path(id): Path, + State(s): State, + Json(mut req): Json, +) -> Result, RuntimeError> { + if req.idle_flush_seconds == Some(0.0) { + req.idle_flush_seconds = Some(DEFAULT_IDLE_FLUSH); + } + let r = s.inner.clone(); + tokio::task::spawn_blocking(move || r.wait_job(&id, req)) + .await + .map_err(|e| RuntimeError::internal(e.to_string()))? + .map(Json) +} +async fn tail_handler( + Path(id): Path, + State(s): State, + Query(q): Query, +) -> Result, RuntimeError> { + let r = s.inner.clone(); + tokio::task::spawn_blocking(move || { + let j = { + let c = r.db.lock().unwrap(); + read_job(&c, &id)? + }; + let v = r.live_view(&id)?; + let w = tail_window( + &r.output_path(&j), + bounded_positive( + q.output_limit, + r.config.default_output_limit, + r.config.max_output_limit, + ), + )?; + Ok(Json(r.result(&v, &j, w))) + }) + .await + .map_err(|e| RuntimeError::internal(e.to_string()))? +} +async fn status_handler( + Path(id): Path, + State(s): State, +) -> Result, RuntimeError> { + let r = s.inner.clone(); + tokio::task::spawn_blocking(move || r.live_view(&id)) + .await + .map_err(|e| RuntimeError::internal(e.to_string()))? + .map(Json) +} +async fn list_handler( + State(s): State, + Query(q): Query, +) -> Result, RuntimeError> { + let r = s.inner.clone(); + tokio::task::spawn_blocking(move || r.list(q)) + .await + .map_err(|e| RuntimeError::internal(e.to_string()))? + .map(Json) +} +async fn input_handler( + Path(id): Path, + State(s): State, + Json(req): Json, +) -> Result, RuntimeError> { + let r = s.inner.clone(); + tokio::task::spawn_blocking(move || r.send_input(&id, req)) + .await + .map_err(|e| RuntimeError::internal(e.to_string()))? + .map(Json) +} +async fn terminate_handler( + Path(id): Path, + State(s): State, + Json(req): Json, +) -> Result, RuntimeError> { + let r = s.inner.clone(); + tokio::task::spawn_blocking(move || { + r.terminate(&id, req.grace_seconds.unwrap_or(r.config.terminate_grace)) + }) + .await + .map_err(|e| RuntimeError::internal(e.to_string()))? + .map(Json) +} +async fn delete_handler( + Path(id): Path, + State(s): State, + Query(q): Query>, +) -> Result, RuntimeError> { + let r = s.inner.clone(); + let force = q.get("force").is_some_and(|v| v == "true"); + let grace = q + .get("grace_seconds") + .and_then(|v| v.parse().ok()) + .unwrap_or(r.config.terminate_grace); + tokio::task::spawn_blocking(move || r.delete(&id, force, grace)) + .await + .map_err(|e| RuntimeError::internal(e.to_string()))? + .map(Json) +} + +pub fn sanitize_bytes(input: &[u8]) -> Vec { + let mut p = PtySanitizer::default(); + let mut out = p.feed(input); + out.extend(p.flush()); + out +} +#[derive(Default)] +pub struct PtySanitizer { + line: Vec, + pending_cr: bool, + state: SanitizeState, +} +#[derive(Default, PartialEq)] +enum SanitizeState { + #[default] + Normal, + Esc, + Csi, + Osc, + OscEsc, +} +impl PtySanitizer { + fn feed(&mut self, input: &[u8]) -> Vec { + let mut out = Vec::with_capacity(input.len()); + self.feed_into(input, &mut out); + out + } + fn feed_into(&mut self, input: &[u8], out: &mut Vec) { + for &b in input { + match self.state { + SanitizeState::Normal => match b { + 0x1b => self.state = SanitizeState::Esc, + b'\r' => self.pending_cr = true, + b'\n' => { + self.pending_cr = false; + out.extend_from_slice(String::from_utf8_lossy(&self.line).as_bytes()); + out.push(b'\n'); + self.line.clear() + } + _ => { + if self.pending_cr { + self.line.clear(); + } + self.pending_cr = false; + self.line.push(b) + } + }, + SanitizeState::Esc => { + self.state = match b { + b'[' => SanitizeState::Csi, + b']' => SanitizeState::Osc, + _ => SanitizeState::Normal, + }; + } + SanitizeState::Csi => { + if (0x40..=0x7e).contains(&b) { + self.state = SanitizeState::Normal; + } + } + SanitizeState::Osc => { + if b == 7 { + self.state = SanitizeState::Normal + } else if b == 0x1b { + self.state = SanitizeState::OscEsc + } + } + SanitizeState::OscEsc => { + self.state = if b == b'\\' { + SanitizeState::Normal + } else { + SanitizeState::Osc + }; + } + } + } + } + fn flush(&mut self) -> Vec { + self.pending_cr = false; + let o = self.line.clone(); + self.line.clear(); + String::from_utf8_lossy(&o).into_owned().into_bytes() + } +} + +pub fn run_sanitizer(ready: Option<&FsPath>) -> io::Result<()> { + if let Some(p) = ready { + File::create(p)?; + } + let mut input = io::stdin().lock(); + let mut output = io::stdout().lock(); + let mut s = PtySanitizer::default(); + let mut buf = [0u8; 8192]; + let mut sanitized = Vec::with_capacity(buf.len()); + loop { + let n = input.read(&mut buf)?; + if n == 0 { + break; + } + sanitized.clear(); + s.feed_into(&buf[..n], &mut sanitized); + output.write_all(&sanitized)?; + } + output.write_all(&s.flush())?; + output.flush() +} + +pub fn run_runner(args: &[String]) -> i32 { + if args.first().is_some_and(|a| a == "--exec") { + return child_mode(&args[1..]); + } + if args.len() < 3 { + eprintln!("usage: shellctl-runner "); + return 125; + } + let dir = FsPath::new(&args[0]); + while !dir.join("start-gate").exists() { + // The server opens this only after the output pipe is ready. A short + // wait avoids adding a full runtime polling interval to every job. + thread::sleep(START_GATE_POLL_INTERVAL); + } + let env_file = dir.join(".job-env.json"); + let overlay: HashMap = fs::read(&env_file) + .ok() + .and_then(|b| serde_json::from_slice(&b).ok()) + .unwrap_or_default(); + let mut cmd = Command::new(env::current_exe().unwrap()); + let isolation = env::var("SHELLCTL_ENABLE_PATH_ISOLATION") + .map(|v| v == "true") + .unwrap_or(true); + let mut child_args = vec!["--exec".to_string()]; + if isolation { + child_args.push("--landlock".into()); + } + child_args.extend([dir.join("script").display().to_string(), args[2].clone()]); + let mut child_env: HashMap = env::vars_os() + .filter(|(k, _)| { + !matches!( + k.to_str(), + Some("TMUX") + | Some("SHELLCTL_STATE_DIR") + | Some("SHELLCTL_RUNTIME_DIR") + | Some("SHELLCTL_TMUX_SOCKET") + | Some("SHELLCTL_RUNNER") + | Some("SHELLCTL_AUTH_TOKEN") + ) + }) + .collect(); + for (key, value) in overlay { + child_env.insert(OsString::from(key), OsString::from(value)); + } + let home = child_env + .get(OsStr::new("HOME")) + .cloned() + .unwrap_or_default(); + if !home.is_empty() { + let _ = fs::create_dir_all(&home); + } + let tmp = FsPath::new(&args[2]).join(".tmp"); + let _ = fs::create_dir_all(&tmp); + for key in ["TMPDIR", "TMP", "TEMP"] { + child_env + .entry(key.into()) + .or_insert_with(|| tmp.display().to_string().into()); + } + cmd.args(child_args) + .current_dir(&args[2]) + .env_clear() + .envs(child_env) + .stdin(Stdio::inherit()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()); + let status = cmd.status(); + let code = status.ok().and_then(|s| s.code()).unwrap_or(125); + let _ = atomic_write(&dir.join("runner-exit-code"), &code.to_string()); + let _ = atomic_write(&dir.join("runner-ended-at"), ×tamp()); + code +} +fn child_mode(args: &[String]) -> i32 { + let mut i = 0; + let landlock = args.first().is_some_and(|x| x == "--landlock"); + if landlock { + i += 1; + } + if args.len() < i + 2 { + return 125; + } + let script = &args[i]; + let cwd = &args[i + 1]; + if env::set_current_dir(cwd).is_err() { + return 111; + } + let has_shebang = File::open(script) + .and_then(|mut file| { + let mut head = [0_u8; 2]; + file.read_exact(&mut head).map(|_| head) + }) + .is_ok_and(|head| head == *b"#!"); + let mut command = if has_shebang { + Command::new(script) + } else { + let mut c = Command::new("sh"); + c.arg(script); + c + }; + if landlock + && let Err(e) = apply_landlock( + &env::var("HOME").unwrap_or_default(), + cwd, + FsPath::new(script), + ) + { + eprintln!("shellctl-runner: WARNING: {e} ā€” running without filesystem isolation"); + } + let err = command.exec(); + eprintln!("shellctl-runner: exec failed: {err}"); + 126 +} +fn atomic_write(path: &FsPath, value: &str) -> io::Result<()> { + let tmp = path.with_extension(format!("tmp.{}", std::process::id())); + fs::write(&tmp, format!("{value}\n"))?; + fs::rename(tmp, path) +} + +#[cfg(target_os = "linux")] +fn apply_landlock(home: &str, cwd: &str, job: &FsPath) -> Result<(), Box> { + use landlock::{ + ABI, Access, AccessFs, PathBeneath, PathFd, Ruleset, RulesetAttr, RulesetCreatedAttr, + }; + let abi = ABI::V1; + let rw = AccessFs::from_all(abi); + let ro = AccessFs::from_read(abi) | AccessFs::Execute; + let mut rules = Ruleset::default() + .handle_access(AccessFs::from_all(abi))? + .create()? + .no_new_privs(true); + for p in [home, cwd] + .into_iter() + .filter(|p| !p.is_empty() && FsPath::new(p).exists()) + { + rules = rules.add_rule(PathBeneath::new(PathFd::new(p)?, rw))?; + } + let job_path = job.to_string_lossy(); + for p in [ + "/usr", + "/bin", + "/sbin", + "/lib", + "/lib64", + "/etc", + "/proc", + "/opt/dify-agent-tools", + "/opt/homebrew", + "/snap", + ] + .into_iter() + .chain(std::iter::once(job_path.as_ref())) + { + if FsPath::new(p).exists() { + rules = rules.add_rule(PathBeneath::new(PathFd::new(p)?, ro))?; + } + } + for p in [ + "/dev/null", + "/dev/zero", + "/dev/urandom", + "/dev/random", + "/dev/tty", + ] { + if FsPath::new(p).exists() { + rules = rules.add_rule(PathBeneath::new( + PathFd::new(p)?, + AccessFs::ReadFile | AccessFs::WriteFile, + ))?; + } + } + rules.restrict_self()?; + Ok(()) +} +#[cfg(not(target_os = "linux"))] +fn apply_landlock( + _home: &str, + _cwd: &str, + _job: &FsPath, +) -> Result<(), Box> { + Err("Landlock is only available on Linux".into()) +} + +pub fn record_runner_exit( + state_dir: &FsPath, + id: &str, + code: i32, + ended_at: &str, + busy: u64, +) -> Result<(), RuntimeError> { + let c = db_connect(&state_dir.join("shellctl.db"), busy, false)?; + db_record_runner_exit(&c, id, code, ended_at) +} + +pub async fn serve(config: Config) -> Result<(), Box> { + let state = Runtime::initialize(config.clone())?; + Runtime::start_reconciler(state.clone()); + let app = router(state); + let listener = tokio::net::TcpListener::bind(&config.listen).await?; + axum::serve(listener, app) + .with_graceful_shutdown(async { + let _ = tokio::signal::ctrl_c().await; + }) + .await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + use std::sync::Barrier; + + struct TestDir { + path: PathBuf, + } + + impl TestDir { + fn new(label: &str) -> Self { + let path = env::temp_dir().join(format!( + "dify-runtime-{label}-{}-{}", + std::process::id(), + job_id() + )); + fs::create_dir_all(&path).unwrap(); + Self { path } + } + + fn path(&self) -> &FsPath { + &self.path + } + } + + impl Drop for TestDir { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.path); + } + } + + fn test_job(id: &str, status: Status, exit_code: Option) -> Job { + Job { + id: id.into(), + script_path: format!("jobs/{id}/script"), + output_path: format!("jobs/{id}/output.log"), + cwd: "/tmp".into(), + cols: 80, + rows: 24, + status, + session_name: session_name(id), + pane_target: pane_target(id), + exit_code, + _reason: None, + _message: None, + created_at: "2026-01-01T00:00:00Z".into(), + started_at: None, + ended_at: None, + _updated_at: "2026-01-01T00:00:00Z".into(), + } + } + + fn insert_test_job(conn: &Connection, job: &Job) { + conn.execute( + "INSERT INTO jobs (job_id,script_path,output_path,cwd,terminal_cols,terminal_rows,status,session_name,pane_target,exit_code,created_at,started_at,ended_at,updated_at) VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13,?14)", + params![ + &job.id, + &job.script_path, + &job.output_path, + &job.cwd, + job.cols, + job.rows, + job.status.as_str(), + &job.session_name, + &job.pane_target, + job.exit_code, + &job.created_at, + job.started_at.as_deref(), + job.ended_at.as_deref(), + &job._updated_at, + ], + ) + .unwrap(); + } + + struct MaterializeCase<'a> { + id: &'a str, + status: Status, + exit_code: Option, + session: bool, + pipe: Option, + pipe_failed: bool, + starting: bool, + } + + fn materialize_case(conn: &mut Connection, case: MaterializeCase<'_>) -> Option { + let job = test_job(case.id, case.status, case.exit_code); + insert_test_job(conn, &job); + materialize( + conn, + &job, + case.session, + case.pipe, + case.pipe_failed, + case.starting, + ) + .unwrap() + } + + #[test] + fn sanitizer_matches_runtime_contract() { + assert_eq!(sanitize_bytes(b"hello\nworld\n"), b"hello\nworld\n"); + assert_eq!(sanitize_bytes(b"\x1b[31mred\x1b[0m\n"), b"red\n"); + assert_eq!(sanitize_bytes(b"50%\r100%\n"), b"100%\n"); + assert_eq!(sanitize_bytes(b"line1\r\nline2\r\n"), b"line1\nline2\n"); + assert_eq!(sanitize_bytes(b"\x1b]0;title\x07visible\n"), b"visible\n"); + assert_eq!(sanitize_bytes(b"no newline"), b"no newline"); + assert_eq!( + sanitize_bytes(&[0xff, b'a', b'\n']), + "\u{fffd}a\n".as_bytes() + ); + } + + #[test] + fn sanitizer_handles_sequences_split_across_chunks() { + let mut s = PtySanitizer::default(); + assert!(s.feed(b"50%\r10").is_empty()); + assert_eq!(s.feed(b"0%\n"), b"100%\n"); + assert!(s.feed(b"\x1b[").is_empty()); + assert!(s.feed(b"31mred").is_empty()); + assert_eq!(s.feed(b"\n"), b"red\n"); + } + + #[test] + fn output_windows_preserve_utf8_boundaries() { + let path = std::env::temp_dir().join(format!("dify-runtime-test-{}", job_id())); + let mut file = File::create(&path).unwrap(); + file.write_all("äø–ē•Œ".as_bytes()).unwrap(); + drop(file); + let window = read_window(&path, 0, 4).unwrap(); + assert_eq!(window.output, "äø–"); + assert_eq!(window.offset, 3); + assert!(window.truncated); + let tail = tail_window(&path, 4).unwrap(); + assert_eq!(tail.output, "ē•Œ"); + let _ = fs::remove_file(path); + } + + #[test] + fn sqlite_runner_exit_is_idempotent_for_terminal_jobs() { + let dir = TestDir::new("db-idempotent"); + let conn = db_open(&dir.path().join("shellctl.db"), 5000).unwrap(); + conn.execute("INSERT INTO jobs (job_id,script_path,output_path,cwd,status,session_name,pane_target,created_at,updated_at) VALUES ('job','script','out','/','running','s','s:0.0','2026-01-01T00:00:00Z','2026-01-01T00:00:00Z')", []).unwrap(); + db_record_runner_exit(&conn, "job", 7, "2026-01-01T00:00:01Z").unwrap(); + db_record_runner_exit(&conn, "job", 9, "2026-01-01T00:00:02Z").unwrap(); + let job = read_job(&conn, "job").unwrap(); + assert_eq!(job.status, Status::Exited); + assert_eq!(job.exit_code, Some(7)); + assert_eq!(job.ended_at.as_deref(), Some("2026-01-01T00:00:01Z")); + } + + #[test] + fn status_round_trips_and_rejects_unknown_values() { + for status in [ + Status::Created, + Status::Starting, + Status::Running, + Status::Exited, + Status::Terminated, + Status::Failed, + Status::Lost, + ] { + assert_eq!(Status::try_from(status.as_str()).unwrap(), status); + assert_eq!( + status.terminal(), + !matches!(status, Status::Created | Status::Starting | Status::Running) + ); + } + let error = Status::try_from("corrupt").unwrap_err(); + assert_eq!(error.status, 500); + assert_eq!(error.code, "internal_error"); + } + + #[test] + fn config_paths_stay_under_the_selected_state_root() { + let dir = TestDir::new("config"); + let config = Config { + state_dir: dir.path().join("state"), + runtime_dir: dir.path().join("runtime"), + ..Config::default() + }; + assert_eq!(config.jobs_dir(), dir.path().join("state/jobs")); + assert_eq!(config.db_path(), dir.path().join("state/shellctl.db")); + assert_eq!(config.tmux_socket(), dir.path().join("runtime/tmux.sock")); + assert_eq!( + config.runner_path(), + dir.path().join("runtime/bin/shellctl-runner") + ); + } + + #[test] + fn bounded_positive_defaults_zero_and_caps_large_values() { + assert_eq!(bounded_positive(None, 16, 512), 16); + assert_eq!(bounded_positive(Some(0), 16, 512), 16); + assert_eq!(bounded_positive(Some(32), 16, 512), 32); + assert_eq!(bounded_positive(Some(1024), 16, 512), 512); + } + + #[test] + fn output_windows_cover_missing_files_offsets_and_utf8_boundaries() { + let dir = TestDir::new("output-window"); + let missing = dir.path().join("missing.log"); + let empty = read_window(&missing, 0, 8).unwrap(); + assert_eq!(empty.output, ""); + assert_eq!(empty.offset, 0); + + let error = read_window(&missing, 1, 8).unwrap_err(); + assert_eq!(error.status, 400); + assert_eq!(error.code, "invalid_offset"); + + let path = dir.path().join("output.log"); + fs::write(&path, "Aäø–ē•ŒB").unwrap(); + let first = read_window(&path, 0, 4).unwrap(); + assert_eq!(first.output, "Aäø–"); + assert_eq!(first.offset, 4); + assert!(first.truncated); + + let inside_codepoint = read_window(&path, 2, 4).unwrap(); + assert_eq!(inside_codepoint.output, "ē•Œ"); + assert_eq!(inside_codepoint.offset, 7); + assert!(inside_codepoint.truncated); + + let error = read_window(&path, 99, 8).unwrap_err(); + assert_eq!(error.status, 400); + assert_eq!(error.code, "invalid_offset"); + + let tail = tail_window(&path, 5).unwrap(); + assert_eq!(tail.output, "ē•ŒB"); + assert_eq!(tail.offset, "Aäø–ē•ŒB".len()); + } + + #[test] + fn transition_enforces_compare_and_swap_and_terminal_metadata() { + let dir = TestDir::new("transition"); + let conn = db_open(&dir.path().join("shellctl.db"), 5000).unwrap(); + insert_test_job(&conn, &test_job("job", Status::Created, None)); + + let running = transition( + &conn, + "job", + Status::Running, + &[Status::Created], + None, + None, + None, + ) + .unwrap(); + assert_eq!(running.status, Status::Running); + assert!(running.started_at.is_some()); + + let stale = transition( + &conn, + "job", + Status::Failed, + &[Status::Created], + Some("stale"), + None, + None, + ) + .unwrap(); + assert_eq!(stale.status, Status::Running); + + let ended = transition( + &conn, + "job", + Status::Terminated, + &[Status::Running], + None, + None, + Some("2026-01-01T00:00:03Z"), + ) + .unwrap(); + assert_eq!(ended.status, Status::Terminated); + assert_eq!(ended.exit_code, Some(0)); + assert_eq!(ended.ended_at.as_deref(), Some("2026-01-01T00:00:03Z")); + } + + #[test] + fn runner_exit_reports_unknown_jobs() { + let dir = TestDir::new("unknown-exit"); + let conn = db_open(&dir.path().join("shellctl.db"), 5000).unwrap(); + let error = db_record_runner_exit(&conn, "missing", 0, "2026-01-01T00:00:00Z").unwrap_err(); + assert_eq!(error.status, 404); + assert_eq!(error.code, "job_not_found"); + } + + #[test] + fn concurrent_runner_exit_updates_are_idempotent() { + let dir = TestDir::new("concurrent-exit"); + let db_path = dir.path().join("shellctl.db"); + let conn = db_open(&db_path, 5000).unwrap(); + insert_test_job(&conn, &test_job("job", Status::Running, None)); + drop(conn); + + let workers = 12; + let barrier = Arc::new(Barrier::new(workers)); + let mut handles = Vec::new(); + for code in 0..workers { + let barrier = barrier.clone(); + let state_dir = dir.path().to_path_buf(); + handles.push(thread::spawn(move || { + barrier.wait(); + let ended_at = format!("2026-01-01T00:00:{code:02}Z"); + record_runner_exit(&state_dir, "job", code as i32, &ended_at, 5000) + })); + } + for handle in handles { + handle.join().unwrap().unwrap(); + } + + let conn = db_connect(&db_path, 5000, false).unwrap(); + let job = read_job(&conn, "job").unwrap(); + let winner = job.exit_code.unwrap(); + assert!((0..workers as i32).contains(&winner)); + assert_eq!( + job.ended_at.as_deref(), + Some(format!("2026-01-01T00:00:{winner:02}Z").as_str()) + ); + } + + #[test] + fn materialize_recovers_each_nonterminal_runtime_state() { + let dir = TestDir::new("materialize"); + let mut conn = db_open(&dir.path().join("shellctl.db"), 5000).unwrap(); + + let exited = materialize_case( + &mut conn, + MaterializeCase { + id: "exited", + status: Status::Running, + exit_code: Some(23), + session: false, + pipe: None, + pipe_failed: false, + starting: false, + }, + ) + .unwrap(); + assert_eq!(exited.status, Status::Exited); + assert_eq!(exited.exit_code, Some(23)); + + let lost = materialize_case( + &mut conn, + MaterializeCase { + id: "lost", + status: Status::Running, + exit_code: None, + session: false, + pipe: None, + pipe_failed: false, + starting: false, + }, + ) + .unwrap(); + assert_eq!(lost.status, Status::Lost); + + let failed = materialize_case( + &mut conn, + MaterializeCase { + id: "failed", + status: Status::Running, + exit_code: None, + session: true, + pipe: Some(false), + pipe_failed: true, + starting: false, + }, + ) + .unwrap(); + assert_eq!(failed.status, Status::Failed); + + let running = materialize_case( + &mut conn, + MaterializeCase { + id: "running", + status: Status::Starting, + exit_code: None, + session: true, + pipe: Some(true), + pipe_failed: false, + starting: false, + }, + ) + .unwrap(); + assert_eq!(running.status, Status::Running); + + let guarded = materialize_case( + &mut conn, + MaterializeCase { + id: "guarded", + status: Status::Starting, + exit_code: None, + session: false, + pipe: None, + pipe_failed: false, + starting: true, + }, + ); + assert!(guarded.is_none()); + assert_eq!(read_job(&conn, "guarded").unwrap().status, Status::Starting); + + let terminal = materialize_case( + &mut conn, + MaterializeCase { + id: "terminal", + status: Status::Terminated, + exit_code: Some(0), + session: false, + pipe: None, + pipe_failed: true, + starting: false, + }, + ); + assert!(terminal.is_none()); + assert_eq!( + read_job(&conn, "terminal").unwrap().status, + Status::Terminated + ); + } + + #[test] + fn exit_metadata_requires_a_complete_atomic_marker_set() { + let dir = TestDir::new("exit-metadata"); + assert!(drained_exit_metadata(dir.path()).is_none()); + + fs::write(dir.path().join(".pipe-drained"), []).unwrap(); + assert!(drained_exit_metadata(dir.path()).is_none()); + + fs::write(dir.path().join("runner-exit-code"), "not-an-int\n").unwrap(); + fs::write(dir.path().join("runner-ended-at"), "2026-01-01T00:00:00Z\n").unwrap(); + assert!(drained_exit_metadata(dir.path()).is_none()); + + fs::write(dir.path().join("runner-exit-code"), "17\n").unwrap(); + assert_eq!( + drained_exit_metadata(dir.path()), + Some((17, "2026-01-01T00:00:00Z".into())) + ); + } + + #[test] + fn atomic_write_replaces_complete_values_without_leaving_temp_files() { + let dir = TestDir::new("atomic-write"); + let path = dir.path().join("value"); + atomic_write(&path, "first").unwrap(); + atomic_write(&path, "second").unwrap(); + assert_eq!(fs::read_to_string(&path).unwrap(), "second\n"); + assert_eq!(fs::read_dir(dir.path()).unwrap().count(), 1); + } + + #[test] + fn shell_helpers_quote_metacharacters_and_classify_tmux_errors() { + assert_eq!(shell_quote("plain"), "'plain'"); + assert_eq!(shell_quote("a'b"), "'a'\\''b'"); + assert!(tmux_missing("can't find pane: shellctl-job:0.0")); + assert!(tmux_missing("NO SERVER RUNNING on /tmp/tmux.sock")); + assert!(!tmux_missing("permission denied")); + } + + #[test] + fn sanitizer_handles_split_osc_terminators_and_trailing_carriage_returns() { + let mut sanitizer = PtySanitizer::default(); + assert!(sanitizer.feed(b"before\x1b]0;title\x1b").is_empty()); + assert_eq!(sanitizer.feed(b"\\after\rreplace\r"), Vec::::new()); + assert_eq!(sanitizer.flush(), b"replace"); + } + + #[test] + fn runner_rejects_incomplete_invocations_without_touching_state() { + assert_eq!(run_runner(&[]), 125); + assert_eq!(run_runner(&["--exec".into()]), 125); + } +} diff --git a/dify-agent-runtime/tests/acceptance_test.go b/dify-agent-runtime/tests/acceptance_test.go index aed05e6ac9d7e3..a017e182cf100d 100644 --- a/dify-agent-runtime/tests/acceptance_test.go +++ b/dify-agent-runtime/tests/acceptance_test.go @@ -1,7 +1,7 @@ //go:build integration -// Package tests runs the same acceptance test suite against both the Python -// and Go shellctl server implementations to verify API compatibility. +// Package tests runs the same acceptance test suite against every configured +// shellctl server implementation to verify API compatibility. // // Prerequisites: // @@ -15,19 +15,26 @@ import ( "encoding/json" "fmt" "io" + "net" "net/http" "os" + "os/exec" "strings" "testing" "time" + "unicode/utf8" ) var ( goURL = envOrDefault("SHELLCTL_GO_URL", "http://localhost:15005") authToken = envOrDefault("SHELLCTL_TEST_TOKEN", "test-token-123") + rustURL = os.Getenv("SHELLCTL_RUST_URL") + rustToken = envOrDefault("SHELLCTL_RUST_TEST_TOKEN", authToken) goURLNoIsolation = os.Getenv("SHELLCTL_GO_URL_NO_ISOLATION") authTokenNoIsolation = os.Getenv("SHELLCTL_TEST_TOKEN_NO_ISOLATION") + rustURLNoIsolation = os.Getenv("SHELLCTL_RUST_URL_NO_ISOLATION") + rustTokenNoIsolation = envOrDefault("SHELLCTL_RUST_TEST_TOKEN_NO_ISOLATION", authTokenNoIsolation) httpClient = &http.Client{Timeout: 120 * time.Second} ) @@ -36,19 +43,28 @@ var ( type target struct { name string baseURL string + token string } func targets() []target { - return []target{ - {name: "go", baseURL: goURL}, + result := []target{ + {name: "go", baseURL: goURL, token: authToken}, } + if rustURL != "" { + result = append(result, target{name: "rust", baseURL: rustURL, token: rustToken}) + } + return result } -func noIsolationTarget() (target, bool) { - if goURLNoIsolation == "" { - return target{}, false +func noIsolationTargets() []target { + var result []target + if goURLNoIsolation != "" { + result = append(result, target{name: "go-no-isolation", baseURL: goURLNoIsolation, token: authTokenNoIsolation}) + } + if rustURLNoIsolation != "" { + result = append(result, target{name: "rust-no-isolation", baseURL: rustURLNoIsolation, token: rustTokenNoIsolation}) } - return target{name: "go-no-isolation", baseURL: goURLNoIsolation}, true + return result } func TestMain(m *testing.M) { @@ -59,7 +75,7 @@ func TestMain(m *testing.M) { os.Exit(1) } } - if tgt, ok := noIsolationTarget(); ok { + for _, tgt := range noIsolationTargets() { if !waitForServer(tgt) { fmt.Fprintf(os.Stderr, "ERROR: %s server not ready at %s\n", tgt.name, tgt.baseURL) os.Exit(1) @@ -69,8 +85,8 @@ func TestMain(m *testing.M) { for _, tgt := range targets() { warmupJob(tgt) } - if tgt, ok := noIsolationTarget(); ok { - warmupJobWithToken(tgt, authTokenNoIsolation) + for _, tgt := range noIsolationTargets() { + warmupJob(tgt) } os.Exit(m.Run()) } @@ -101,7 +117,7 @@ func warmupJob(tgt target) { for attempt := 0; attempt < 3; attempt++ { req, _ := http.NewRequest("POST", tgt.baseURL+"/v1/jobs/run", bytes.NewReader(payload)) req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+authToken) + req.Header.Set("Authorization", "Bearer "+tgt.token) resp, err := warmupClient.Do(req) if err != nil { fmt.Fprintf(os.Stderr, "WARN: %s warmup job attempt %d failed: %v\n", tgt.name, attempt+1, err) @@ -380,7 +396,7 @@ func TestDeleteJob(t *testing.T) { // Delete it req, _ := http.NewRequest("DELETE", fmt.Sprintf("%s/v1/jobs/%s", tgt.baseURL, jobID), nil) - req.Header.Set("Authorization", "Bearer "+authToken) + req.Header.Set("Authorization", "Bearer "+tgt.token) resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("delete request failed: %v", err) @@ -411,7 +427,7 @@ func TestForceDeleteRunningJob(t *testing.T) { // Force delete req, _ := http.NewRequest("DELETE", fmt.Sprintf("%s/v1/jobs/%s?force=true&grace_seconds=1", tgt.baseURL, jobID), nil) - req.Header.Set("Authorization", "Bearer "+authToken) + req.Header.Set("Authorization", "Bearer "+tgt.token) resp, err := http.DefaultClient.Do(req) if err != nil { t.Fatalf("delete request failed: %v", err) @@ -508,6 +524,296 @@ func TestJobNotFound(t *testing.T) { } } +func TestInvalidBearerTokenContract(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + req, _ := http.NewRequest("GET", tgt.baseURL+"/v1/jobs", nil) + req.Header.Set("Authorization", "Bearer definitely-wrong") + resp, err := httpClient.Do(req) + if err != nil { + t.Fatalf("request failed: %v", err) + } + assertAPIError(t, resp, 401, "unauthorized") + }) + } +} + +func TestRunValidationErrorContract(t *testing.T) { + tests := []struct { + name string + payload map[string]any + status int + code string + }{ + {name: "empty-script", payload: map[string]any{"script": ""}, status: 400, code: "invalid_request"}, + {name: "empty-env-name", payload: map[string]any{"script": "true", "env": map[string]string{"": "x"}}, status: 422, code: "validation_error"}, + {name: "equals-in-env-name", payload: map[string]any{"script": "true", "env": map[string]string{"A=B": "x"}}, status: 422, code: "validation_error"}, + {name: "nul-in-env-value", payload: map[string]any{"script": "true", "env": map[string]string{"A": "x\x00y"}}, status: 422, code: "validation_error"}, + } + + for _, tgt := range targets() { + for _, tc := range tests { + t.Run(tgt.name+"/"+tc.name, func(t *testing.T) { + resp := doPost(t, tgt, "/v1/jobs/run", tc.payload, true) + assertAPIError(t, resp, tc.status, tc.code) + }) + } + } +} + +func TestOutputLimitZeroUsesDefault(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{ + "script": "printf zero-limit-output", + "timeout": 10, + "output_limit": 0, + }) + assertJobDone(t, result) + if output := result["output"].(string); output != "zero-limit-output" { + t.Fatalf("output_limit=0 should use the default, got %q", output) + } + }) + } +} + +func TestLargeUTF8OutputIsChunkedWithoutSplittingCodepoints(t *testing.T) { + const repetitions = 6000 + const outputLimit = 4097 + expected := strings.Repeat("äø–ē•Œ", repetitions) + + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{ + "script": fmt.Sprintf("i=0; while [ \"$i\" -lt %d ]; do printf 'äø–ē•Œ'; i=$((i+1)); done", repetitions), + "timeout": 20, + "output_limit": outputLimit, + }) + jobID := result["job_id"].(string) + var combined strings.Builder + + for chunk := 0; chunk < 32; chunk++ { + output := result["output"].(string) + if !utf8.ValidString(output) { + t.Fatalf("chunk %d is not valid UTF-8", chunk) + } + if len(output) > outputLimit { + t.Fatalf("chunk %d exceeded output limit: %d > %d", chunk, len(output), outputLimit) + } + combined.WriteString(output) + if result["done"] == true && result["truncated"] != true { + break + } + offset := int(result["offset"].(float64)) + result = waitJob(t, tgt, jobID, map[string]any{ + "offset": offset, + "timeout": 10, + "output_limit": outputLimit, + }) + } + + if got := combined.String(); got != expected { + t.Fatalf("reassembled output mismatch: got %d bytes, want %d", len(got), len(expected)) + } + }) + } +} + +func TestWaitRejectsOffsetPastEnd(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{"script": "printf short", "timeout": 10}) + jobID := result["job_id"].(string) + offset := int(result["offset"].(float64)) + 1 + resp := doPost(t, tgt, fmt.Sprintf("/v1/jobs/%s/wait", jobID), map[string]any{ + "offset": offset, + "timeout": 0, + }, true) + assertAPIError(t, resp, 400, "invalid_offset") + }) + } +} + +func TestListZeroLimitUsesDefaultAndStatusFilter(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{"script": "true", "timeout": 10}) + assertJobDone(t, result) + + resp := doGet(t, tgt, "/v1/jobs?limit=0&status=exited", true) + assertStatus(t, resp, 200) + var list struct { + Jobs []struct { + Status string `json:"status"` + } `json:"jobs"` + } + if err := json.Unmarshal(readBody(t, resp), &list); err != nil { + t.Fatalf("decode list: %v", err) + } + if len(list.Jobs) == 0 { + t.Fatal("limit=0 should use the default rather than returning an empty list") + } + if len(list.Jobs) > 50 { + t.Fatalf("default list limit exceeded: %d", len(list.Jobs)) + } + for _, job := range list.Jobs { + if job.Status != "exited" { + t.Fatalf("status filter returned %q", job.Status) + } + } + }) + } +} + +func TestTerminalInputAndRunningDeleteConflicts(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + completed := runJob(t, tgt, map[string]any{"script": "true", "timeout": 10}) + completedID := completed["job_id"].(string) + resp := doPost(t, tgt, fmt.Sprintf("/v1/jobs/%s/input", completedID), map[string]any{ + "text": "ignored\n", "offset": 0, "timeout": 1, + }, true) + assertAPIError(t, resp, 409, "job_not_running") + + running := runJob(t, tgt, map[string]any{"script": "sleep 60", "timeout": 0.1}) + runningID := running["job_id"].(string) + resp = doDelete(t, tgt, fmt.Sprintf("/v1/jobs/%s", runningID)) + assertAPIError(t, resp, 409, "job_running") + + resp = doDelete(t, tgt, fmt.Sprintf("/v1/jobs/%s?force=true&grace_seconds=0", runningID)) + assertStatus(t, resp, 200) + resp.Body.Close() + }) + } +} + +func TestTerminateIsIdempotentAndDeleteIsNot(t *testing.T) { + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + result := runJob(t, tgt, map[string]any{"script": "sleep 60", "timeout": 0.1}) + jobID := result["job_id"].(string) + path := fmt.Sprintf("/v1/jobs/%s/terminate", jobID) + + for attempt := 0; attempt < 2; attempt++ { + resp := doPost(t, tgt, path, map[string]any{"grace_seconds": 0}, true) + assertStatus(t, resp, 200) + var view map[string]any + if err := json.Unmarshal(readBody(t, resp), &view); err != nil { + t.Fatalf("decode terminate response: %v", err) + } + if view["done"] != true { + t.Fatalf("terminate attempt %d did not return terminal state: %v", attempt+1, view) + } + } + + resp := doDelete(t, tgt, fmt.Sprintf("/v1/jobs/%s", jobID)) + assertStatus(t, resp, 200) + resp.Body.Close() + resp = doDelete(t, tgt, fmt.Sprintf("/v1/jobs/%s", jobID)) + assertAPIError(t, resp, 404, "job_not_found") + }) + } +} + +func TestConcurrentJobCreationAndCompletion(t *testing.T) { + const jobs = 8 + for _, tgt := range targets() { + t.Run(tgt.name, func(t *testing.T) { + for index := 0; index < jobs; index++ { + index := index + t.Run(fmt.Sprintf("job-%02d", index), func(t *testing.T) { + t.Parallel() + marker := fmt.Sprintf("concurrent-%02d", index) + result := runJob(t, tgt, map[string]any{ + "script": fmt.Sprintf("printf %s", marker), + "timeout": 20, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + if result["output"] != marker { + t.Fatalf("unexpected output: %q", result["output"]) + } + }) + } + }) + } +} + +func TestRustContainerRestartPreservesStateAndRecoversRunningJobs(t *testing.T) { + image := os.Getenv("SHELLCTL_RUST_IMAGE") + if image == "" { + t.Skip("Rust integration image is not configured") + } + container := fmt.Sprintf("sandbox-rt-rust-restart-%d", time.Now().UnixNano()) + token := "restart-recovery-token" + output, err := exec.Command( + "docker", + "run", + "-d", + "--name", + container, + "-p", + "127.0.0.1::5004", + "-e", + "SHELLCTL_AUTH_TOKEN="+token, + image, + ).CombinedOutput() + if err != nil { + t.Fatalf("start dedicated restart container: %v: %s", err, output) + } + defer func() { + if cleanup, cleanupErr := exec.Command("docker", "rm", "-f", container).CombinedOutput(); cleanupErr != nil { + t.Errorf("remove dedicated restart container: %v: %s", cleanupErr, cleanup) + } + }() + + tgt := target{name: "rust-restart", baseURL: containerPublishedURL(t, container), token: token} + if !waitForServer(tgt) { + t.Fatal("dedicated Rust runtime did not become healthy") + } + completed := runJob(t, tgt, map[string]any{"script": "printf persisted", "timeout": 10}) + completedID := completed["job_id"].(string) + running := runJob(t, tgt, map[string]any{"script": "sleep 60", "timeout": 0.1}) + runningID := running["job_id"].(string) + + output, err = exec.Command("docker", "restart", "--timeout", "1", container).CombinedOutput() + if err != nil { + t.Fatalf("restart Rust container: %v: %s", err, output) + } + tgt.baseURL = containerPublishedURL(t, container) + if !waitForServer(tgt) { + t.Fatal("Rust runtime did not become healthy after restart") + } + + completedStatus := getJobStatus(t, tgt, completedID) + if completedStatus["status"] != "exited" || completedStatus["done"] != true { + t.Fatalf("completed job changed across restart: %v", completedStatus) + } + tail := doGet(t, tgt, fmt.Sprintf("/v1/jobs/%s/log/tail", completedID), true) + assertStatus(t, tail, 200) + var persisted map[string]any + if err := json.Unmarshal(readBody(t, tail), &persisted); err != nil { + t.Fatalf("decode persisted output: %v", err) + } + if persisted["output"] != "persisted" { + t.Fatalf("completed output was not preserved: %q", persisted["output"]) + } + + runningStatus := getJobStatus(t, tgt, runningID) + if runningStatus["done"] != true { + t.Fatalf("pre-restart running job was not reconciled: %v", runningStatus) + } + if status := runningStatus["status"]; status == "created" || status == "starting" || status == "running" { + t.Fatalf("pre-restart job remained nonterminal: %v", runningStatus) + } + + after := runJob(t, tgt, map[string]any{"script": "printf after-restart", "timeout": 10}) + assertJobDone(t, after) + if after["output"] != "after-restart" { + t.Fatalf("runtime failed to accept work after restart: %v", after) + } +} + // --- Landlock Tests --- // These tests verify that shellctl-run restricts filesystem access // so each agent job can only write within its own HOME directory while still @@ -624,23 +930,27 @@ func TestLandlockCannotReadOtherAgentHome(t *testing.T) { // TestLandlockDisabledAllowsWriteOutsideHome uses the pre-started no-isolation // container (SHELLCTL_ENABLE_PATH_ISOLATION=false) and verifies that isolation is off. func TestLandlockDisabledAllowsWriteOutsideHome(t *testing.T) { - tgt, ok := noIsolationTarget() - if !ok { + targets := noIsolationTargets() + if len(targets) == 0 { t.Skip("SHELLCTL_GO_URL_NO_ISOLATION not set; no-isolation container not available") } - // With isolation disabled, writes to /tmp should succeed. - // /tmp is world-writable (Unix perms) but blocked by Landlock when enabled. - result := runJobWithToken(t, tgt, authTokenNoIsolation, map[string]any{ - "script": "touch /tmp/landlock-disabled-test && echo write_ok", - "env": map[string]string{"HOME": "/home/dify"}, - "timeout": 10, - }) - assertJobDone(t, result) - assertExitCode(t, result, 0) - output := result["output"].(string) - if !strings.Contains(output, "write_ok") { - t.Errorf("expected write to /tmp to succeed with isolation disabled, got %q", output) + for _, tgt := range targets { + t.Run(tgt.name, func(t *testing.T) { + // With isolation disabled, writes to /tmp should succeed. + // /tmp is world-writable but blocked by Landlock when enabled. + result := runJob(t, tgt, map[string]any{ + "script": "touch /tmp/landlock-disabled-test && echo write_ok", + "env": map[string]string{"HOME": "/home/dify"}, + "timeout": 10, + }) + assertJobDone(t, result) + assertExitCode(t, result, 0) + output := result["output"].(string) + if !strings.Contains(output, "write_ok") { + t.Errorf("expected write to /tmp to succeed with isolation disabled, got %q", output) + } + }) } } @@ -696,7 +1006,7 @@ func doGet(t *testing.T, tgt target, path string, withAuth bool) *http.Response t.Helper() req, _ := http.NewRequest("GET", tgt.baseURL+path, nil) if withAuth { - req.Header.Set("Authorization", "Bearer "+authToken) + req.Header.Set("Authorization", "Bearer "+tgt.token) } resp, err := httpClient.Do(req) if err != nil { @@ -711,7 +1021,7 @@ func doPost(t *testing.T, tgt target, path string, payload map[string]any, withA req, _ := http.NewRequest("POST", tgt.baseURL+path, bytes.NewReader(body)) req.Header.Set("Content-Type", "application/json") if withAuth { - req.Header.Set("Authorization", "Bearer "+authToken) + req.Header.Set("Authorization", "Bearer "+tgt.token) } resp, err := httpClient.Do(req) if err != nil { @@ -720,6 +1030,42 @@ func doPost(t *testing.T, tgt target, path string, payload map[string]any, withA return resp } +func doDelete(t *testing.T, tgt target, path string) *http.Response { + t.Helper() + req, _ := http.NewRequest("DELETE", tgt.baseURL+path, nil) + req.Header.Set("Authorization", "Bearer "+tgt.token) + resp, err := httpClient.Do(req) + if err != nil { + t.Fatalf("[%s] DELETE %s failed: %v", tgt.name, path, err) + } + return resp +} + +func getJobStatus(t *testing.T, tgt target, jobID string) map[string]any { + t.Helper() + resp := doGet(t, tgt, fmt.Sprintf("/v1/jobs/%s", jobID), true) + assertStatus(t, resp, 200) + var result map[string]any + if err := json.Unmarshal(readBody(t, resp), &result); err != nil { + t.Fatalf("decode job status: %v", err) + } + return result +} + +func containerPublishedURL(t *testing.T, container string) string { + t.Helper() + output, err := exec.Command("docker", "port", container, "5004/tcp").CombinedOutput() + if err != nil { + t.Fatalf("resolve published port for %s: %v: %s", container, err, output) + } + line := strings.TrimSpace(strings.Split(string(output), "\n")[0]) + _, port, err := net.SplitHostPort(line) + if err != nil || port == "" { + t.Fatalf("parse published port %q for %s: %v", line, container, err) + } + return "http://127.0.0.1:" + port +} + func readBody(t *testing.T, resp *http.Response) []byte { t.Helper() defer resp.Body.Close() @@ -739,6 +1085,23 @@ func assertStatus(t *testing.T, resp *http.Response, expected int) { } } +func assertAPIError(t *testing.T, resp *http.Response, expectedStatus int, expectedCode string) { + t.Helper() + assertStatus(t, resp, expectedStatus) + var result struct { + Error struct { + Code string `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal(readBody(t, resp), &result); err != nil { + t.Fatalf("decode error response: %v", err) + } + if result.Error.Code != expectedCode { + t.Fatalf("expected error code %q, got %q (%s)", expectedCode, result.Error.Code, result.Error.Message) + } +} + func assertJobDone(t *testing.T, result map[string]any) { t.Helper() if result["done"] != true { @@ -764,46 +1127,3 @@ func envOrDefault(key, defaultVal string) string { } return defaultVal } - -func warmupJobWithToken(tgt target, token string) { - warmupClient := &http.Client{Timeout: 180 * time.Second} - payload, _ := json.Marshal(map[string]any{ - "script": "echo warmup", - "timeout": 10, - }) - for attempt := 0; attempt < 3; attempt++ { - req, _ := http.NewRequest("POST", tgt.baseURL+"/v1/jobs/run", bytes.NewReader(payload)) - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+token) - resp, err := warmupClient.Do(req) - if err != nil { - time.Sleep(2 * time.Second) - continue - } - io.Copy(io.Discard, resp.Body) - resp.Body.Close() - if resp.StatusCode == 200 { - return - } - time.Sleep(2 * time.Second) - } -} - -func runJobWithToken(t *testing.T, tgt target, token string, payload map[string]any) map[string]any { - t.Helper() - body, _ := json.Marshal(payload) - req, _ := http.NewRequest("POST", tgt.baseURL+"/v1/jobs/run", bytes.NewReader(body)) - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+token) - resp, err := httpClient.Do(req) - if err != nil { - t.Fatalf("[%s] POST /v1/jobs/run failed: %v", tgt.name, err) - } - assertStatus(t, resp, 200) - respBody := readBody(t, resp) - var result map[string]any - if err := json.Unmarshal(respBody, &result); err != nil { - t.Fatalf("[%s] failed to parse run response: %v\nbody: %s", tgt.name, err, string(respBody)) - } - return result -} diff --git a/dify-agent/src/dify_agent/runtime_backend/local.py b/dify-agent/src/dify_agent/runtime_backend/local.py index f0a2e5c2e66c23..0ae792c3d0b809 100644 --- a/dify-agent/src/dify_agent/runtime_backend/local.py +++ b/dify-agent/src/dify_agent/runtime_backend/local.py @@ -13,7 +13,7 @@ import shlex from dataclasses import dataclass -from dify_agent.adapters.shell.protocols import ShellCommandProtocol +from dify_agent.adapters.shell.protocols import ShellCommandProtocol, ShellProviderError from dify_agent.adapters.shell.shellctl import ShellctlClientFactory from dify_agent.runtime_backend.errors import ( BindingAcquireError, @@ -188,6 +188,8 @@ async def acquire(self, binding_ref: str) -> RuntimeLease: await _close_best_effort(lease, resource_ref=binding_ref) if isinstance(exc, BindingLostError): raise + if isinstance(exc, ShellProviderError) and exc.code == "invalid_cwd": + raise BindingLostError(f"Local Binding {binding_ref!r} no longer exists") from exc if isinstance(exc, Exception): raise BindingAcquireError(str(exc)) from exc raise diff --git a/dify-agent/src/dify_agent/runtime_backend/local_rollout.py b/dify-agent/src/dify_agent/runtime_backend/local_rollout.py new file mode 100644 index 00000000000000..6d860972867d28 --- /dev/null +++ b/dify-agent/src/dify_agent/runtime_backend/local_rollout.py @@ -0,0 +1,288 @@ +"""Safe, sticky rollout between the local Go and Rust shell runtimes. + +Fallback is deliberately limited to the preflight before a new Binding is +created. Once a backend receives a mutating request, the operation is never +replayed against the other implementation because shell commands are not +generally idempotent. + +Rust-owned opaque refs carry an implementation prefix. Go-owned refs keep their +original representation, so enabling the router does not migrate Go data and a +Go-only rollback can still read every Go allocation. Existing unprefixed refs +remain Go refs. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import logging +from dataclasses import dataclass, replace +from typing import Literal, Protocol + +from dify_agent.adapters.shell.protocols import ShellCommandProtocol +from dify_agent.adapters.shell.shellctl import ShellctlClientFactory +from dify_agent.runtime_backend.errors import ( + BindingAcquireError, + BindingCreateError, + BindingDestroyError, +) +from dify_agent.runtime_backend.protocols import ( + ExecutionBindingAllocation, + ExecutionBindingBackend, + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, + HomeSnapshotBackend, + HomeSnapshotCreateSpec, + RuntimeLayout, + RuntimeLease, +) + +type LocalRuntimeImplementation = Literal["go", "rust"] + +_GO: LocalRuntimeImplementation = "go" +_RUST: LocalRuntimeImplementation = "rust" +_REF_SEPARATOR = "+" +logger = logging.getLogger(__name__) + + +class LocalRuntimeHealthProbe(Protocol): + async def __call__(self) -> None: ... + + +@dataclass(slots=True) +class ShellctlHealthProbe: + """Bounded shellctl health probe used only before Rust admission.""" + + client_factory: ShellctlClientFactory + timeout_seconds: float = 1.0 + + async def __call__(self) -> None: + client = self.client_factory() + try: + async with asyncio.timeout(self.timeout_seconds): + response = await client.health() + if response.status != "ok": + raise RuntimeError(f"unexpected shellctl health status: {response.status!r}") + except BaseException: + try: + await client.close() + except BaseException: + logger.warning("failed to close shellctl client after rollout preflight failure", exc_info=True) + raise + await client.close() + + +@dataclass(frozen=True, slots=True) +class LocalRuntimeTarget: + implementation: LocalRuntimeImplementation + home_snapshots: HomeSnapshotBackend + execution_bindings: ExecutionBindingBackend + + +@dataclass(slots=True) +class LocalRuntimeRouter: + """Choose Rust only for new, unpinned Bindings and keep Go as fallback.""" + + go: LocalRuntimeTarget + rust: LocalRuntimeTarget + rust_canary_percent: int + rust_health_probe: LocalRuntimeHealthProbe + + def __post_init__(self) -> None: + if self.go.implementation != _GO or self.rust.implementation != _RUST: + raise ValueError("local runtime targets must be wired as go and rust") + if not 0 <= self.rust_canary_percent <= 100: + raise ValueError("rust_canary_percent must be between 0 and 100") + + def target(self, implementation: LocalRuntimeImplementation) -> LocalRuntimeTarget: + return self.rust if implementation == _RUST else self.go + + async def select_for_create(self, spec: ExecutionBindingCreateSpec) -> LocalRuntimeImplementation: + pinned: set[LocalRuntimeImplementation] = { + _decode_ref(ref).implementation + for ref in (spec.existing_workspace_ref, spec.home_snapshot_ref) + if ref is not None + } + if len(pinned) > 1: + raise BindingCreateError("Home Snapshot and existing Workspace belong to different runtime implementations") + if pinned: + implementation = pinned.pop() + if implementation == _RUST: + await self._require_pinned_rust() + return implementation + + if not _is_rust_canary(spec, self.rust_canary_percent): + return _GO + + try: + await self.rust_health_probe() + except Exception: + logger.warning( + "Rust local runtime preflight failed; assigning new Binding to Go", + exc_info=True, + extra={"binding_id": spec.binding_id, "workspace_id": spec.workspace_id}, + ) + return _GO + logger.info( + "Assigning new Binding to the Rust local runtime canary", + extra={ + "binding_id": spec.binding_id, + "workspace_id": spec.workspace_id, + "rust_canary_percent": self.rust_canary_percent, + }, + ) + return _RUST + + async def _require_pinned_rust(self) -> None: + try: + await self.rust_health_probe() + except Exception as exc: + raise BindingCreateError( + "Rust runtime is unavailable for a resource already pinned to Rust; refusing unsafe Go replay" + ) from exc + + +@dataclass(slots=True) +class RoutedLocalRuntimeLease: + """Runtime lease tagged with the backend that owns its state.""" + + implementation: LocalRuntimeImplementation + inner: RuntimeLease + + @property + def layout(self) -> RuntimeLayout: + return self.inner.layout + + @property + def commands(self) -> ShellCommandProtocol: + return self.inner.commands + + +@dataclass(slots=True) +class RoutedLocalExecutionBindingBackend: + router: LocalRuntimeRouter + + async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation: + try: + implementation = await self.router.select_for_create(spec) + native_spec = replace( + spec, + existing_workspace_ref=_native_ref_for(implementation, spec.existing_workspace_ref), + home_snapshot_ref=_native_ref_for(implementation, spec.home_snapshot_ref), + ) + except BindingCreateError: + raise + except ValueError as exc: + raise BindingCreateError(str(exc)) from exc + + # Do not catch and replay this call. The selected runtime may already + # have mutated its Home or Workspace before surfacing an error. + allocation = await self.router.target(implementation).execution_bindings.create_binding(native_spec) + return ExecutionBindingAllocation( + binding_ref=_encode_ref(implementation, allocation.binding_ref), + workspace_ref=_encode_ref(implementation, allocation.workspace_ref), + ) + + async def acquire(self, binding_ref: str) -> RuntimeLease: + try: + routed_ref = _decode_ref(binding_ref) + except ValueError as exc: + raise BindingAcquireError(str(exc)) from exc + lease = await self.router.target(routed_ref.implementation).execution_bindings.acquire(routed_ref.native) + return RoutedLocalRuntimeLease(implementation=routed_ref.implementation, inner=lease) + + async def release(self, lease: RuntimeLease) -> None: + if not isinstance(lease, RoutedLocalRuntimeLease): + raise TypeError("RoutedLocalExecutionBindingBackend can only release its own RuntimeLease") + await self.router.target(lease.implementation).execution_bindings.release(lease.inner) + + async def destroy_binding(self, spec: ExecutionBindingDestroySpec) -> None: + try: + binding_ref = _decode_ref(spec.binding_ref) + workspace_ref = _decode_ref(spec.workspace_ref) if spec.workspace_ref is not None else None + if workspace_ref is not None and workspace_ref.implementation != binding_ref.implementation: + raise ValueError("Workspace and Binding refs belong to different runtime implementations") + native_spec = replace( + spec, + binding_ref=binding_ref.native, + workspace_ref=workspace_ref.native if workspace_ref is not None else None, + ) + except ValueError as exc: + raise BindingDestroyError(str(exc)) from exc + await self.router.target(binding_ref.implementation).execution_bindings.destroy_binding(native_spec) + + +@dataclass(slots=True) +class RoutedLocalHomeSnapshotBackend: + router: LocalRuntimeRouter + + async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str: + if not isinstance(source, RoutedLocalRuntimeLease): + raise TypeError("RoutedLocalHomeSnapshotBackend requires a routed local RuntimeLease") + native_ref = await self.router.target(source.implementation).home_snapshots.create_from_runtime( + spec=spec, + source=source.inner, + ) + return _encode_ref(source.implementation, native_ref) + + async def delete(self, snapshot_ref: str) -> None: + routed_ref = _decode_ref(snapshot_ref) + await self.router.target(routed_ref.implementation).home_snapshots.delete(routed_ref.native) + + +@dataclass(frozen=True, slots=True) +class _RoutedRef: + implementation: LocalRuntimeImplementation + native: str + + +def _encode_ref(implementation: LocalRuntimeImplementation, native: str) -> str: + if not native: + raise ValueError("runtime backend ref must not be empty") + if implementation == _GO: + return native + return f"{implementation}{_REF_SEPARATOR}{native}" + + +def _decode_ref(value: str) -> _RoutedRef: + if not value: + raise ValueError("runtime backend ref must not be empty") + for implementation in (_GO, _RUST): + prefix = f"{implementation}{_REF_SEPARATOR}" + if value.startswith(prefix): + native = value[len(prefix) :] + if not native: + raise ValueError("runtime backend ref must include a native ref") + return _RoutedRef(implementation=implementation, native=native) + # Refs created before dual-runtime rollout belong to the existing Go + # implementation. This preserves upgrades and makes rollback deterministic. + return _RoutedRef(implementation=_GO, native=value) + + +def _native_ref_for(implementation: LocalRuntimeImplementation, value: str | None) -> str | None: + if value is None: + return None + routed_ref = _decode_ref(value) + if routed_ref.implementation != implementation: + raise ValueError("runtime backend ref is pinned to a different implementation") + return routed_ref.native + + +def _is_rust_canary(spec: ExecutionBindingCreateSpec, percent: int) -> bool: + if percent <= 0: + return False + if percent >= 100: + return True + key = "\0".join((spec.tenant_id, spec.agent_id, spec.binding_id, spec.workspace_id)).encode() + bucket = int.from_bytes(hashlib.blake2b(key, digest_size=8, usedforsecurity=False).digest(), "big") % 100 + return bucket < percent + + +__all__ = [ + "LocalRuntimeRouter", + "LocalRuntimeTarget", + "RoutedLocalExecutionBindingBackend", + "RoutedLocalHomeSnapshotBackend", + "RoutedLocalRuntimeLease", + "ShellctlHealthProbe", +] diff --git a/dify-agent/src/dify_agent/runtime_backend/profile.py b/dify-agent/src/dify_agent/runtime_backend/profile.py index 28bcf6bd01e547..66e6a01909695e 100644 --- a/dify-agent/src/dify_agent/runtime_backend/profile.py +++ b/dify-agent/src/dify_agent/runtime_backend/profile.py @@ -17,7 +17,15 @@ ) from dify_agent.runtime_backend.enterprise import EnterpriseExecutionBindingBackend, EnterpriseHomeSnapshotBackend from dify_agent.runtime_backend.local import LocalExecutionBindingBackend, LocalHomeSnapshotBackend +from dify_agent.runtime_backend.local_rollout import ( + LocalRuntimeRouter, + LocalRuntimeTarget, + RoutedLocalExecutionBindingBackend, + RoutedLocalHomeSnapshotBackend, + ShellctlHealthProbe, +) from dify_agent.runtime_backend.protocols import RuntimeBackendProfile +from dify_agent.adapters.shell.shellctl import create_default_shellctl_client_factory DEFAULT_E2B_TEMPLATE = "difys-default-team/dify-agent-local-sandbox" DEFAULT_LOCAL_MATERIALIZED_HOME_ROOT = "/home/dify/.dify-agent-materialized-homes" @@ -44,6 +52,10 @@ class RuntimeBackendSettings(BaseSettings): "DIFY_AGENT_SHELLCTL_AUTH_TOKEN", ), ) + local_sandbox_rust_endpoint: str | None = None + local_sandbox_rust_auth_token: str | None = None + local_sandbox_rust_canary_percent: int = Field(default=0, ge=0, le=100) + local_sandbox_preflight_timeout_seconds: float = Field(default=1.0, gt=0, le=30) local_sandbox_materialized_home_root: str = DEFAULT_LOCAL_MATERIALIZED_HOME_ROOT local_sandbox_workspace_root: str = DEFAULT_LOCAL_WORKSPACE_ROOT local_sandbox_home_snapshot_root: str = DEFAULT_LOCAL_HOME_SNAPSHOT_ROOT @@ -77,6 +89,15 @@ def validate_selected_backend(self) -> Self: if not self.local_sandbox_endpoint or not self.local_sandbox_endpoint.strip(): raise ValueError("local_sandbox_endpoint is required for the local runtime backend") _validate_http_url(self.local_sandbox_endpoint, field_name="local_sandbox_endpoint") + rust_endpoint = self.local_sandbox_rust_endpoint + if rust_endpoint is not None and rust_endpoint.strip(): + _validate_http_url(rust_endpoint, field_name="local_sandbox_rust_endpoint") + if rust_endpoint.rstrip("/") == self.local_sandbox_endpoint.rstrip("/"): + raise ValueError("local_sandbox_rust_endpoint must differ from local_sandbox_endpoint") + elif self.local_sandbox_rust_canary_percent > 0: + raise ValueError( + "local_sandbox_rust_endpoint is required when local_sandbox_rust_canary_percent is greater than 0" + ) _validate_absolute_posix_path( self.local_sandbox_materialized_home_root, field_name="local_sandbox_materialized_home_root", @@ -110,20 +131,67 @@ def create_runtime_backend_profile(settings: RuntimeBackendSettings) -> RuntimeB case "local": endpoint = settings.local_sandbox_endpoint or "" token = settings.local_sandbox_auth_token or "" - return RuntimeBackendProfile( - home_snapshots=LocalHomeSnapshotBackend( - endpoint=endpoint, - auth_token=token, - snapshot_root=settings.local_sandbox_home_snapshot_root, + go_home_snapshots = LocalHomeSnapshotBackend( + endpoint=endpoint, + auth_token=token, + snapshot_root=settings.local_sandbox_home_snapshot_root, + ) + go_execution_bindings = LocalExecutionBindingBackend( + endpoint=endpoint, + auth_token=token, + materialized_home_root=settings.local_sandbox_materialized_home_root, + workspace_root=settings.local_sandbox_workspace_root, + snapshot_root=settings.local_sandbox_home_snapshot_root, + ) + rust_endpoint = (settings.local_sandbox_rust_endpoint or "").strip() + if not rust_endpoint: + return RuntimeBackendProfile( + home_snapshots=go_home_snapshots, + execution_bindings=go_execution_bindings, + ) + + rust_token = ( + settings.local_sandbox_rust_auth_token if settings.local_sandbox_rust_auth_token is not None else token + ) + rust_client_factory = create_default_shellctl_client_factory( + entrypoint=rust_endpoint, + token=rust_token, + ) + rust_home_snapshots = LocalHomeSnapshotBackend( + endpoint=rust_endpoint, + auth_token=rust_token, + snapshot_root=settings.local_sandbox_home_snapshot_root, + client_factory=rust_client_factory, + ) + rust_execution_bindings = LocalExecutionBindingBackend( + endpoint=rust_endpoint, + auth_token=rust_token, + materialized_home_root=settings.local_sandbox_materialized_home_root, + workspace_root=settings.local_sandbox_workspace_root, + snapshot_root=settings.local_sandbox_home_snapshot_root, + client_factory=rust_client_factory, + ) + router = LocalRuntimeRouter( + go=LocalRuntimeTarget( + implementation="go", + home_snapshots=go_home_snapshots, + execution_bindings=go_execution_bindings, ), - execution_bindings=LocalExecutionBindingBackend( - endpoint=endpoint, - auth_token=token, - materialized_home_root=settings.local_sandbox_materialized_home_root, - workspace_root=settings.local_sandbox_workspace_root, - snapshot_root=settings.local_sandbox_home_snapshot_root, + rust=LocalRuntimeTarget( + implementation="rust", + home_snapshots=rust_home_snapshots, + execution_bindings=rust_execution_bindings, + ), + rust_canary_percent=settings.local_sandbox_rust_canary_percent, + rust_health_probe=ShellctlHealthProbe( + client_factory=rust_client_factory, + timeout_seconds=settings.local_sandbox_preflight_timeout_seconds, ), ) + return RuntimeBackendProfile( + home_snapshots=RoutedLocalHomeSnapshotBackend(router=router), + execution_bindings=RoutedLocalExecutionBindingBackend(router=router), + ) case "enterprise": endpoint = settings.enterprise_sandbox_gateway_endpoint or "" token = settings.enterprise_sandbox_gateway_auth_token or "" diff --git a/dify-agent/src/dify_agent/server/settings.py b/dify-agent/src/dify_agent/server/settings.py index 177456dc47a9b6..a4efc80c2bfa40 100644 --- a/dify-agent/src/dify_agent/server/settings.py +++ b/dify-agent/src/dify_agent/server/settings.py @@ -55,6 +55,10 @@ class ServerSettings(BaseSettings): default=None, validation_alias=AliasChoices("DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN", "DIFY_AGENT_SHELLCTL_AUTH_TOKEN"), ) + local_sandbox_rust_endpoint: str | None = None + local_sandbox_rust_auth_token: str | None = None + local_sandbox_rust_canary_percent: int = Field(default=0, ge=0, le=100) + local_sandbox_preflight_timeout_seconds: float = Field(default=1.0, gt=0, le=30) local_sandbox_materialized_home_root: str = DEFAULT_LOCAL_MATERIALIZED_HOME_ROOT local_sandbox_workspace_root: str = DEFAULT_LOCAL_WORKSPACE_ROOT local_sandbox_home_snapshot_root: str = DEFAULT_LOCAL_HOME_SNAPSHOT_ROOT @@ -194,6 +198,10 @@ def build_runtime_backend_profile(self) -> RuntimeBackendProfile | None: runtime_backend=self.runtime_backend, local_sandbox_endpoint=self.local_sandbox_endpoint, local_sandbox_auth_token=self.local_sandbox_auth_token, + local_sandbox_rust_endpoint=self.local_sandbox_rust_endpoint, + local_sandbox_rust_auth_token=self.local_sandbox_rust_auth_token, + local_sandbox_rust_canary_percent=self.local_sandbox_rust_canary_percent, + local_sandbox_preflight_timeout_seconds=self.local_sandbox_preflight_timeout_seconds, local_sandbox_materialized_home_root=self.local_sandbox_materialized_home_root, local_sandbox_workspace_root=self.local_sandbox_workspace_root, local_sandbox_home_snapshot_root=self.local_sandbox_home_snapshot_root, diff --git a/dify-agent/tests/integration/dify_agent/runtime_backend/run_local_integration.sh b/dify-agent/tests/integration/dify_agent/runtime_backend/run_local_integration.sh index 1247636b0589cf..19238702c1ac60 100755 --- a/dify-agent/tests/integration/dify_agent/runtime_backend/run_local_integration.sh +++ b/dify-agent/tests/integration/dify_agent/runtime_backend/run_local_integration.sh @@ -2,7 +2,7 @@ set -eu script_dir=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) -project_dir=$(CDPATH= cd -- "$script_dir/../../../.." && pwd) +project_dir=$(CDPATH= cd -- "$script_dir/../../../../.." && pwd) container_name="dify-agent-runtime-backend-integration-$$" image="${DIFY_AGENT_TEST_LOCAL_SANDBOX_IMAGE:-langgenius/dify-agent-local-sandbox:1.16.0}" token="${DIFY_AGENT_TEST_LOCAL_SHELLCTL_AUTH_TOKEN:-runtime-backend-integration}" @@ -31,10 +31,11 @@ until curl --fail --silent "$endpoint/healthz" >/dev/null; do sleep 0.1 done -cd "$project_dir" +cd "$project_dir/dify-agent" NO_PROXY=127.0.0.1,localhost \ DIFY_AGENT_TEST_LOCAL_SHELLCTL_ENDPOINT="$endpoint" \ DIFY_AGENT_TEST_LOCAL_SHELLCTL_AUTH_TOKEN="$token" \ - pdm run pytest --import-mode=importlib \ - tests/integration/dify_agent/runtime_backend/test_runtime_backend_lifecycle.py \ + PYTHONPATH=src \ + uv run --extra server pytest --import-mode=importlib \ + tests/integration/dify_agent/runtime_backend/test_working_environment.py \ -k local -q -rs "$@" diff --git a/dify-agent/tests/integration/dify_agent/runtime_backend/test_working_environment.py b/dify-agent/tests/integration/dify_agent/runtime_backend/test_working_environment.py index 93a67c2e816c28..4bcef7d53e44db 100644 --- a/dify-agent/tests/integration/dify_agent/runtime_backend/test_working_environment.py +++ b/dify-agent/tests/integration/dify_agent/runtime_backend/test_working_environment.py @@ -10,6 +10,7 @@ import pytest from dify_agent.runtime_backend import ( + BindingLostError, ExecutionBindingCreateSpec, ExecutionBindingDestroySpec, HomeSnapshotCreateSpec, @@ -21,6 +22,7 @@ E2BSDKControlPlane, ) from dify_agent.runtime_backend.local import LocalExecutionBindingBackend +from dify_agent.runtime_backend.profile import RuntimeBackendSettings, create_runtime_backend_profile from dify_agent.runtime.command_runner import execute_complete_with_commands pytestmark = pytest.mark.integration @@ -115,6 +117,163 @@ async def test_local_two_agents_share_workspace_but_not_home() -> None: raise cleanup_errors[0] +@pytest.mark.anyio +async def test_local_rust_canary_is_sticky_and_state_isolated_from_go() -> None: + go_endpoint = _required_env("DIFY_AGENT_TEST_LOCAL_SHELLCTL_ENDPOINT", "real Go shellctl") + rust_endpoint = _required_env("DIFY_AGENT_TEST_RUST_SHELLCTL_ENDPOINT", "real Rust shellctl") + go_token = os.environ.get("DIFY_AGENT_TEST_LOCAL_SHELLCTL_AUTH_TOKEN", "") + rust_token = os.environ.get("DIFY_AGENT_TEST_RUST_SHELLCTL_AUTH_TOKEN", go_token) + marker = uuid.uuid4().hex + profile = create_runtime_backend_profile( + RuntimeBackendSettings( + runtime_backend="local", + local_sandbox_endpoint=go_endpoint, + local_sandbox_auth_token=go_token, + local_sandbox_rust_endpoint=rust_endpoint, + local_sandbox_rust_auth_token=rust_token, + local_sandbox_rust_canary_percent=100, + local_sandbox_preflight_timeout_seconds=1.0, + ) + ) + direct_go = LocalExecutionBindingBackend(endpoint=go_endpoint, auth_token=go_token) + allocations = [] + active_leases = [] + snapshot_ref: str | None = None + try: + first = await profile.execution_bindings.create_binding( + ExecutionBindingCreateSpec( + tenant_id="integration-tenant", + agent_id="rust-canary-agent", + binding_id=f"rust-canary-{marker}", + workspace_id=f"rust-canary-{marker}", + existing_workspace_ref=None, + home_snapshot_ref=None, + ) + ) + allocations.append(first) + assert first.binding_ref.startswith("rust+") + assert first.workspace_ref.startswith("rust+") + + first_lease = await profile.execution_bindings.acquire(first.binding_ref) + active_leases.append(first_lease) + await _run(first_lease, "printf rust-home > .runtime-owner", cwd=first_lease.layout.home_dir) + snapshot_ref = await profile.home_snapshots.create_from_runtime( + spec=HomeSnapshotCreateSpec( + tenant_id="integration-tenant", + agent_id="rust-canary-agent", + home_snapshot_id=f"rust-canary-{marker}", + ), + source=first_lease, + ) + assert snapshot_ref.startswith("rust+") + await profile.execution_bindings.release(first_lease) + active_leases.remove(first_lease) + + with pytest.raises(BindingLostError): + await direct_go.acquire(first.binding_ref.removeprefix("rust+")) + + restored = await profile.execution_bindings.create_binding( + ExecutionBindingCreateSpec( + tenant_id="integration-tenant", + agent_id="rust-canary-agent", + binding_id=f"rust-restored-{marker}", + workspace_id=f"rust-restored-{marker}", + existing_workspace_ref=None, + home_snapshot_ref=snapshot_ref, + ) + ) + allocations.append(restored) + assert restored.binding_ref.startswith("rust+") + restored_lease = await profile.execution_bindings.acquire(restored.binding_ref) + active_leases.append(restored_lease) + assert await _run(restored_lease, "cat .runtime-owner", cwd=restored_lease.layout.home_dir) == "rust-home" + await profile.execution_bindings.release(restored_lease) + active_leases.remove(restored_lease) + finally: + primary_error = sys.exc_info()[0] is not None + cleanup_errors: list[BaseException] = [] + for lease in active_leases: + try: + await profile.execution_bindings.release(lease) + except BaseException as exc: + cleanup_errors.append(exc) + for allocation in allocations: + try: + await profile.execution_bindings.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref=allocation.binding_ref, + workspace_ref=allocation.workspace_ref, + destroy_workspace=True, + ) + ) + except BaseException as exc: + cleanup_errors.append(exc) + if snapshot_ref is not None: + try: + await profile.home_snapshots.delete(snapshot_ref) + except BaseException as exc: + cleanup_errors.append(exc) + if cleanup_errors and not primary_error: + raise cleanup_errors[0] + + +@pytest.mark.anyio +async def test_local_unavailable_rust_preflight_falls_back_to_real_go() -> None: + go_endpoint = _required_env("DIFY_AGENT_TEST_LOCAL_SHELLCTL_ENDPOINT", "real Go shellctl") + go_token = os.environ.get("DIFY_AGENT_TEST_LOCAL_SHELLCTL_AUTH_TOKEN", "") + marker = uuid.uuid4().hex + profile = create_runtime_backend_profile( + RuntimeBackendSettings( + runtime_backend="local", + local_sandbox_endpoint=go_endpoint, + local_sandbox_auth_token=go_token, + local_sandbox_rust_endpoint="http://127.0.0.1:1", + local_sandbox_rust_canary_percent=100, + local_sandbox_preflight_timeout_seconds=0.1, + ) + ) + allocation = None + lease = None + try: + allocation = await profile.execution_bindings.create_binding( + ExecutionBindingCreateSpec( + tenant_id="integration-tenant", + agent_id="fallback-agent", + binding_id=f"go-fallback-{marker}", + workspace_id=f"go-fallback-{marker}", + existing_workspace_ref=None, + home_snapshot_ref=None, + ) + ) + assert not allocation.binding_ref.startswith("rust+") + assert not allocation.workspace_ref.startswith("rust+") + lease = await profile.execution_bindings.acquire(allocation.binding_ref) + assert await _run(lease, "printf go-fallback", cwd=lease.layout.workspace_dir) == "go-fallback" + await profile.execution_bindings.release(lease) + lease = None + finally: + primary_error = sys.exc_info()[0] is not None + cleanup_errors: list[BaseException] = [] + if lease is not None: + try: + await profile.execution_bindings.release(lease) + except BaseException as exc: + cleanup_errors.append(exc) + if allocation is not None: + try: + await profile.execution_bindings.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref=allocation.binding_ref, + workspace_ref=allocation.workspace_ref, + destroy_workspace=True, + ) + ) + except BaseException as exc: + cleanup_errors.append(exc) + if cleanup_errors and not primary_error: + raise cleanup_errors[0] + + @pytest.mark.anyio async def test_e2b_binding_checkpoint_and_collection() -> None: api_key = _required_env("DIFY_AGENT_TEST_E2B_API_KEY", "real E2B") diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py index a2f8190e32fc11..f178a39ede78f4 100644 --- a/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_local.py @@ -7,9 +7,11 @@ import pytest from shellctl.shared import DeleteJobResponse, JobResult, JobStatusName, JobStatusView +from dify_agent.adapters.shell.protocols import ShellProviderError from dify_agent.runtime_backend import ( BindingCreateError, BindingDestroyError, + BindingLostError, ExecutionBindingCreateSpec, ExecutionBindingDestroySpec, HomeSnapshotCreateSpec, @@ -31,6 +33,7 @@ class _Client: exit_code: int = 0 output: str = "" close_error: Exception | None = None + run_error: Exception | None = None exit_codes: list[int] = field(default_factory=list) async def run( @@ -42,6 +45,8 @@ async def run( timeout: float = 10.0, ) -> JobResult: del timeout + if self.run_error is not None: + raise self.run_error commands = tuple( tuple(shlex.split(line)) for line in script.splitlines() if line.strip() and line.strip() != "set -eu" ) @@ -135,6 +140,20 @@ def commands(self) -> tuple[tuple[str, ...], ...]: return tuple(command for run in self.runs for command in run.commands) +@dataclass(slots=True) +class _InvalidCwdFactory: + clients: list[_Client] = field(default_factory=list) + runs: list[_RunCall] = field(default_factory=list) + + def __call__(self) -> _Client: + client = _Client( + runs=self.runs, + run_error=ShellProviderError("cwd is not a directory", code="invalid_cwd", status_code=400), + ) + self.clients.append(client) + return client + + @pytest.mark.anyio async def test_local_binding_create_materializes_home_and_new_workspace() -> None: factory = _Factory() @@ -250,6 +269,24 @@ async def test_local_binding_acquire_scopes_commands_to_materialized_home_and_wo await backend.release(lease) +@pytest.mark.anyio +async def test_local_binding_acquire_reports_missing_state_as_binding_lost() -> None: + factory = _InvalidCwdFactory() + backend = LocalExecutionBindingBackend( + endpoint="http://shellctl", + auth_token="", + materialized_home_root="/homes", + workspace_root="/workspaces", + snapshot_root="/snapshots", + client_factory=factory, # pyright: ignore[reportArgumentType] + ) + + with pytest.raises(BindingLostError, match="no longer exists"): + await backend.acquire("binding-1:workspace-1") + + assert factory.clients and all(client.closed for client in factory.clients) + + @pytest.mark.anyio async def test_local_snapshot_checkpoint_copies_only_materialized_home() -> None: factory = _Factory() diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_local_rollout.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_local_rollout.py new file mode 100644 index 00000000000000..a209c74805df7c --- /dev/null +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_local_rollout.py @@ -0,0 +1,520 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from typing import cast + +import pytest +from shellctl.shared import HealthResponse + +from dify_agent.adapters.shell.protocols import ShellCommandProtocol +from dify_agent.adapters.shell.shellctl import ShellctlClientProtocol +from dify_agent.runtime_backend import ( + BindingAcquireError, + BindingCreateError, + BindingDestroyError, + ExecutionBindingAllocation, + ExecutionBindingCreateSpec, + ExecutionBindingDestroySpec, + HomeSnapshotCreateSpec, + RuntimeLayout, + RuntimeLease, +) +from dify_agent.runtime_backend.local_rollout import ( + LocalRuntimeRouter, + LocalRuntimeTarget, + RoutedLocalExecutionBindingBackend, + RoutedLocalHomeSnapshotBackend, + RoutedLocalRuntimeLease, + ShellctlHealthProbe, + _is_rust_canary, # pyright: ignore[reportPrivateUsage] +) + + +@dataclass(slots=True) +class _Lease: + owner: str + layout: RuntimeLayout = field(default_factory=lambda: RuntimeLayout(home_dir="/home", workspace_dir="/workspace")) + commands: ShellCommandProtocol = field(default_factory=lambda: cast(ShellCommandProtocol, object())) + + +@dataclass(slots=True) +class _BindingBackend: + owner: str + create_error: Exception | None = None + creates: list[ExecutionBindingCreateSpec] = field(default_factory=list) + acquires: list[str] = field(default_factory=list) + releases: list[RuntimeLease] = field(default_factory=list) + destroys: list[ExecutionBindingDestroySpec] = field(default_factory=list) + + async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation: + self.creates.append(spec) + if self.create_error is not None: + raise self.create_error + return ExecutionBindingAllocation( + binding_ref=f"{spec.binding_id}:{spec.workspace_id}", + workspace_ref=spec.workspace_id, + ) + + async def acquire(self, binding_ref: str) -> RuntimeLease: + self.acquires.append(binding_ref) + return _Lease(owner=self.owner) + + async def release(self, lease: RuntimeLease) -> None: + self.releases.append(lease) + + async def destroy_binding(self, spec: ExecutionBindingDestroySpec) -> None: + self.destroys.append(spec) + + +@dataclass(slots=True) +class _HomeBackend: + owner: str + creates: list[tuple[HomeSnapshotCreateSpec, RuntimeLease]] = field(default_factory=list) + deletes: list[str] = field(default_factory=list) + + async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str: + self.creates.append((spec, source)) + return f"home-{spec.home_snapshot_id}" + + async def delete(self, snapshot_ref: str) -> None: + self.deletes.append(snapshot_ref) + + +@dataclass(slots=True) +class _Probe: + error: Exception | None = None + calls: int = 0 + + async def __call__(self) -> None: + self.calls += 1 + if self.error is not None: + raise self.error + + +@dataclass(slots=True) +class _HealthClient: + status: str = "ok" + error: Exception | None = None + wait_forever: bool = False + close_error: Exception | None = None + closed: bool = False + + async def health(self) -> HealthResponse: + if self.wait_forever: + _ = await asyncio.Event().wait() + if self.error is not None: + raise self.error + return HealthResponse(status=self.status) + + async def close(self) -> None: + self.closed = True + if self.close_error is not None: + raise self.close_error + + +@dataclass(slots=True) +class _Fixture: + router: LocalRuntimeRouter + go_bindings: _BindingBackend + rust_bindings: _BindingBackend + go_snapshots: _HomeBackend + rust_snapshots: _HomeBackend + probe: _Probe + + +def _fixture(*, canary_percent: int, probe_error: Exception | None = None) -> _Fixture: + go_bindings = _BindingBackend(owner="go") + rust_bindings = _BindingBackend(owner="rust") + go_snapshots = _HomeBackend(owner="go") + rust_snapshots = _HomeBackend(owner="rust") + probe = _Probe(error=probe_error) + router = LocalRuntimeRouter( + go=LocalRuntimeTarget( + implementation="go", + home_snapshots=go_snapshots, + execution_bindings=go_bindings, + ), + rust=LocalRuntimeTarget( + implementation="rust", + home_snapshots=rust_snapshots, + execution_bindings=rust_bindings, + ), + rust_canary_percent=canary_percent, + rust_health_probe=probe, + ) + return _Fixture( + router=router, + go_bindings=go_bindings, + rust_bindings=rust_bindings, + go_snapshots=go_snapshots, + rust_snapshots=rust_snapshots, + probe=probe, + ) + + +def _spec( + *, + binding_id: str = "binding-1", + workspace_id: str = "workspace-1", + existing_workspace_ref: str | None = None, + home_snapshot_ref: str | None = None, +) -> ExecutionBindingCreateSpec: + return ExecutionBindingCreateSpec( + tenant_id="tenant-1", + agent_id="agent-1", + binding_id=binding_id, + workspace_id=workspace_id, + existing_workspace_ref=existing_workspace_ref, + home_snapshot_ref=home_snapshot_ref, + ) + + +def _health_probe(client: _HealthClient, *, timeout_seconds: float = 1.0) -> ShellctlHealthProbe: + def factory() -> ShellctlClientProtocol: + return cast(ShellctlClientProtocol, cast(object, client)) + + return ShellctlHealthProbe(client_factory=factory, timeout_seconds=timeout_seconds) + + +@pytest.mark.anyio +async def test_health_probe_closes_client_after_success() -> None: + client = _HealthClient() + + await _health_probe(client)() + + assert client.closed is True + + +@pytest.mark.anyio +async def test_health_probe_closes_client_after_unhealthy_status() -> None: + client = _HealthClient(status="degraded") + + with pytest.raises(RuntimeError, match="unexpected shellctl health status"): + await _health_probe(client)() + + assert client.closed is True + + +@pytest.mark.anyio +async def test_health_probe_is_bounded_and_closes_timed_out_client() -> None: + client = _HealthClient(wait_forever=True) + + with pytest.raises(TimeoutError): + await _health_probe(client, timeout_seconds=0.001)() + + assert client.closed is True + + +@pytest.mark.anyio +async def test_health_probe_preserves_probe_error_when_close_also_fails() -> None: + client = _HealthClient( + error=ConnectionError("health failed"), + close_error=RuntimeError("close failed"), + ) + + with pytest.raises(ConnectionError, match="health failed"): + await _health_probe(client)() + + assert client.closed is True + + +@pytest.mark.anyio +async def test_health_probe_surfaces_close_error_after_success() -> None: + client = _HealthClient(close_error=RuntimeError("close failed")) + + with pytest.raises(RuntimeError, match="close failed"): + await _health_probe(client)() + + assert client.closed is True + + +def test_canary_hash_is_deterministic_and_tracks_requested_percentage() -> None: + specs = [_spec(binding_id=f"binding-{index}", workspace_id=f"workspace-{index}") for index in range(1_000)] + + first_pass = [_is_rust_canary(spec, 25) for spec in specs] + second_pass = [_is_rust_canary(spec, 25) for spec in specs] + + assert first_pass == second_pass + assert 200 <= sum(first_pass) <= 300 + assert not any(_is_rust_canary(spec, 0) for spec in specs) + assert all(_is_rust_canary(spec, 100) for spec in specs) + + +@pytest.mark.parametrize("percent", [-1, 101]) +def test_router_rejects_canary_percentage_outside_closed_interval(percent: int) -> None: + with pytest.raises(ValueError, match="between 0 and 100"): + _ = _fixture(canary_percent=percent) + + +@pytest.mark.anyio +async def test_healthy_rust_canary_is_sticky_across_binding_and_snapshot_lifecycle() -> None: + fixture = _fixture(canary_percent=100) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + snapshots = RoutedLocalHomeSnapshotBackend(router=fixture.router) + + allocation = await bindings.create_binding(_spec()) + + assert allocation.binding_ref == "rust+binding-1:workspace-1" + assert allocation.workspace_ref == "rust+workspace-1" + assert fixture.probe.calls == 1 + assert len(fixture.rust_bindings.creates) == 1 + assert fixture.go_bindings.creates == [] + + lease = await bindings.acquire(allocation.binding_ref) + assert isinstance(lease, RoutedLocalRuntimeLease) + assert lease.implementation == "rust" + assert fixture.rust_bindings.acquires == ["binding-1:workspace-1"] + + snapshot_ref = await snapshots.create_from_runtime( + spec=HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="snapshot-1"), + source=lease, + ) + assert snapshot_ref == "rust+home-snapshot-1" + await snapshots.delete(snapshot_ref) + assert fixture.rust_snapshots.deletes == ["home-snapshot-1"] + + await bindings.release(lease) + assert len(fixture.rust_bindings.releases) == 1 + assert fixture.go_bindings.releases == [] + + +@pytest.mark.anyio +async def test_unhealthy_rust_preflight_assigns_new_binding_to_go() -> None: + fixture = _fixture(canary_percent=100, probe_error=TimeoutError("rust unavailable")) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + allocation = await bindings.create_binding(_spec()) + + assert allocation.binding_ref == "binding-1:workspace-1" + assert allocation.workspace_ref == "workspace-1" + assert len(fixture.go_bindings.creates) == 1 + assert fixture.rust_bindings.creates == [] + + +@pytest.mark.anyio +async def test_zero_percent_canary_keeps_go_refs_compatible_with_go_only_rollback() -> None: + fixture = _fixture(canary_percent=0, probe_error=ConnectionError("must not be called")) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + allocation = await bindings.create_binding(_spec()) + + assert allocation.binding_ref == "binding-1:workspace-1" + assert allocation.workspace_ref == "workspace-1" + assert fixture.probe.calls == 0 + assert fixture.go_bindings.acquires == [] + _ = await fixture.go_bindings.acquire(allocation.binding_ref) + assert fixture.go_bindings.acquires == ["binding-1:workspace-1"] + + +@pytest.mark.anyio +async def test_existing_go_workspace_bypasses_rust_even_at_full_canary() -> None: + fixture = _fixture(canary_percent=100, probe_error=ConnectionError("must not be called")) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + allocation = await bindings.create_binding(_spec(existing_workspace_ref="workspace-1")) + + assert allocation.binding_ref == "binding-1:workspace-1" + assert fixture.probe.calls == 0 + assert len(fixture.go_bindings.creates) == 1 + assert fixture.rust_bindings.creates == [] + + +@pytest.mark.anyio +async def test_existing_rust_workspace_stays_rust_when_new_canary_admission_is_disabled() -> None: + fixture = _fixture(canary_percent=0) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + allocation = await bindings.create_binding(_spec(existing_workspace_ref="rust+workspace-1")) + + assert allocation.binding_ref == "rust+binding-1:workspace-1" + assert fixture.probe.calls == 1 + assert len(fixture.rust_bindings.creates) == 1 + assert fixture.go_bindings.creates == [] + + +@pytest.mark.anyio +async def test_mutating_rust_failure_is_never_replayed_to_go() -> None: + fixture = _fixture(canary_percent=100) + fixture.rust_bindings.create_error = BindingCreateError("Rust may already have mutated state") + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + with pytest.raises(BindingCreateError, match="already have mutated"): + _ = await bindings.create_binding(_spec()) + + assert len(fixture.rust_bindings.creates) == 1 + assert fixture.go_bindings.creates == [] + + +@pytest.mark.anyio +async def test_mutating_go_failure_is_never_replayed_to_rust() -> None: + fixture = _fixture(canary_percent=0) + fixture.go_bindings.create_error = BindingCreateError("Go may already have mutated state") + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + with pytest.raises(BindingCreateError, match="already have mutated"): + _ = await bindings.create_binding(_spec()) + + assert len(fixture.go_bindings.creates) == 1 + assert fixture.rust_bindings.creates == [] + + +@pytest.mark.anyio +async def test_existing_rust_resource_never_falls_back_to_go() -> None: + fixture = _fixture(canary_percent=0, probe_error=ConnectionError("rust unavailable")) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + with pytest.raises(BindingCreateError, match="refusing unsafe Go replay"): + _ = await bindings.create_binding(_spec(existing_workspace_ref="rust+workspace-1")) + + assert fixture.go_bindings.creates == [] + assert fixture.rust_bindings.creates == [] + + +@pytest.mark.anyio +async def test_cross_runtime_snapshot_and_workspace_are_rejected_before_mutation() -> None: + fixture = _fixture(canary_percent=100) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + with pytest.raises(BindingCreateError, match="different runtime implementations"): + _ = await bindings.create_binding( + _spec( + existing_workspace_ref="go+workspace-1", + home_snapshot_ref="rust+home-snapshot-1", + ) + ) + + assert fixture.probe.calls == 0 + assert fixture.go_bindings.creates == [] + assert fixture.rust_bindings.creates == [] + + +@pytest.mark.anyio +async def test_legacy_unprefixed_binding_ref_remains_owned_by_go() -> None: + fixture = _fixture(canary_percent=100) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + lease = await bindings.acquire("binding-legacy:workspace-legacy") + + assert isinstance(lease, RoutedLocalRuntimeLease) + assert lease.implementation == "go" + assert fixture.go_bindings.acquires == ["binding-legacy:workspace-legacy"] + assert fixture.rust_bindings.acquires == [] + + +@pytest.mark.anyio +@pytest.mark.parametrize("binding_ref", ["", "go+", "rust+"]) +async def test_acquire_rejects_empty_native_refs(binding_ref: str) -> None: + fixture = _fixture(canary_percent=100) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + with pytest.raises(BindingAcquireError, match="must not be empty|include a native ref"): + _ = await bindings.acquire(binding_ref) + + assert fixture.go_bindings.acquires == [] + assert fixture.rust_bindings.acquires == [] + + +@pytest.mark.anyio +async def test_explicit_go_prefix_is_decoded_but_never_leaks_to_go_backend() -> None: + fixture = _fixture(canary_percent=100) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + lease = await bindings.acquire("go+binding-1") + + assert isinstance(lease, RoutedLocalRuntimeLease) + assert lease.implementation == "go" + assert fixture.go_bindings.acquires == ["binding-1"] + + +@pytest.mark.anyio +async def test_destroy_dispatches_only_to_ref_owner() -> None: + fixture = _fixture(canary_percent=100) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + await bindings.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref="rust+binding-1:workspace-1", + destroy_workspace=True, + workspace_ref="rust+workspace-1", + ) + ) + + assert fixture.rust_bindings.destroys == [ + ExecutionBindingDestroySpec( + binding_ref="binding-1:workspace-1", + destroy_workspace=True, + workspace_ref="workspace-1", + ) + ] + assert fixture.go_bindings.destroys == [] + + +@pytest.mark.anyio +async def test_destroy_rejects_cross_runtime_refs_before_mutation() -> None: + fixture = _fixture(canary_percent=100) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + with pytest.raises(BindingDestroyError, match="different runtime implementations"): + await bindings.destroy_binding( + ExecutionBindingDestroySpec( + binding_ref="rust+binding-1:workspace-1", + destroy_workspace=True, + workspace_ref="workspace-1", + ) + ) + + assert fixture.rust_bindings.destroys == [] + assert fixture.go_bindings.destroys == [] + + +@pytest.mark.anyio +async def test_release_rejects_lease_from_outside_router() -> None: + fixture = _fixture(canary_percent=100) + bindings = RoutedLocalExecutionBindingBackend(router=fixture.router) + + with pytest.raises(TypeError, match="only release its own"): + await bindings.release(_Lease(owner="foreign")) + + assert fixture.rust_bindings.releases == [] + assert fixture.go_bindings.releases == [] + + +@pytest.mark.anyio +async def test_home_snapshot_dispatches_to_go_and_rust_owners() -> None: + fixture = _fixture(canary_percent=100) + snapshots = RoutedLocalHomeSnapshotBackend(router=fixture.router) + spec = HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="snapshot-1") + + go_ref = await snapshots.create_from_runtime( + spec=spec, + source=RoutedLocalRuntimeLease(implementation="go", inner=_Lease(owner="go")), + ) + rust_ref = await snapshots.create_from_runtime( + spec=spec, + source=RoutedLocalRuntimeLease(implementation="rust", inner=_Lease(owner="rust")), + ) + + assert go_ref == "home-snapshot-1" + assert rust_ref == "rust+home-snapshot-1" + await snapshots.delete(go_ref) + await snapshots.delete(rust_ref) + assert fixture.go_snapshots.deletes == ["home-snapshot-1"] + assert fixture.rust_snapshots.deletes == ["home-snapshot-1"] + + +@pytest.mark.anyio +async def test_home_snapshot_rejects_unrouted_lease_and_malformed_ref() -> None: + fixture = _fixture(canary_percent=100) + snapshots = RoutedLocalHomeSnapshotBackend(router=fixture.router) + spec = HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="snapshot-1") + + with pytest.raises(TypeError, match="requires a routed"): + _ = await snapshots.create_from_runtime(spec=spec, source=_Lease(owner="foreign")) + with pytest.raises(ValueError, match="include a native ref"): + await snapshots.delete("rust+") + + assert fixture.go_snapshots.creates == [] + assert fixture.rust_snapshots.creates == [] + assert fixture.go_snapshots.deletes == [] + assert fixture.rust_snapshots.deletes == [] diff --git a/dify-agent/tests/local/dify_agent/runtime_backend/test_profile.py b/dify-agent/tests/local/dify_agent/runtime_backend/test_profile.py index fe6704ac9695d2..eeb5161792dda8 100644 --- a/dify-agent/tests/local/dify_agent/runtime_backend/test_profile.py +++ b/dify-agent/tests/local/dify_agent/runtime_backend/test_profile.py @@ -5,6 +5,11 @@ from dify_agent.runtime_backend.e2b import E2B_MAX_ACTIVE_TIMEOUT_SECONDS from dify_agent.runtime_backend.local import LocalExecutionBindingBackend, LocalHomeSnapshotBackend +from dify_agent.runtime_backend.local_rollout import ( + RoutedLocalExecutionBindingBackend, + RoutedLocalHomeSnapshotBackend, + ShellctlHealthProbe, +) from dify_agent.runtime_backend.profile import ( DEFAULT_E2B_TEMPLATE, RuntimeBackendSettings, @@ -66,3 +71,66 @@ def test_local_backend_rejects_relative_roots() -> None: local_sandbox_endpoint="http://shellctl.example", local_sandbox_workspace_root="relative/workspaces", ) + + +def test_local_backend_requires_rust_endpoint_for_nonzero_canary() -> None: + with pytest.raises(ValidationError, match="local_sandbox_rust_endpoint is required"): + _ = RuntimeBackendSettings( + runtime_backend="local", + local_sandbox_endpoint="http://shellctl-go.example", + local_sandbox_rust_canary_percent=1, + ) + + +def test_local_backend_rejects_same_go_and_rust_endpoint() -> None: + with pytest.raises(ValidationError, match="must differ"): + _ = RuntimeBackendSettings( + runtime_backend="local", + local_sandbox_endpoint="http://shellctl.example/", + local_sandbox_rust_endpoint="http://shellctl.example", + ) + + +def test_local_backend_builds_sticky_rollout_drivers_when_rust_is_configured() -> None: + settings = RuntimeBackendSettings( + runtime_backend="local", + local_sandbox_endpoint="http://shellctl-go.example", + local_sandbox_auth_token="go-token", + local_sandbox_rust_endpoint="http://shellctl-rust.example", + local_sandbox_rust_auth_token="rust-token", + local_sandbox_rust_canary_percent=25, + local_sandbox_preflight_timeout_seconds=2.5, + ) + + profile = create_runtime_backend_profile(settings) + + assert isinstance(profile.execution_bindings, RoutedLocalExecutionBindingBackend) + assert isinstance(profile.home_snapshots, RoutedLocalHomeSnapshotBackend) + router = profile.execution_bindings.router + assert router.rust_canary_percent == 25 + assert isinstance(router.rust_health_probe, ShellctlHealthProbe) + assert router.rust_health_probe.timeout_seconds == 2.5 + assert router.go.implementation == "go" + assert router.rust.implementation == "rust" + assert isinstance(router.go.execution_bindings, LocalExecutionBindingBackend) + assert isinstance(router.rust.execution_bindings, LocalExecutionBindingBackend) + assert router.go.execution_bindings.endpoint == "http://shellctl-go.example" + assert router.go.execution_bindings.auth_token == "go-token" + assert router.rust.execution_bindings.endpoint == "http://shellctl-rust.example" + assert router.rust.execution_bindings.auth_token == "rust-token" + + +def test_local_backend_rust_token_inherits_go_token_when_omitted() -> None: + profile = create_runtime_backend_profile( + RuntimeBackendSettings( + runtime_backend="local", + local_sandbox_endpoint="http://shellctl-go.example", + local_sandbox_auth_token="shared-token", + local_sandbox_rust_endpoint="http://shellctl-rust.example", + ) + ) + + assert isinstance(profile.execution_bindings, RoutedLocalExecutionBindingBackend) + rust_bindings = profile.execution_bindings.router.rust.execution_bindings + assert isinstance(rust_bindings, LocalExecutionBindingBackend) + assert rust_bindings.auth_token == "shared-token" 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 282542b4ccaff5..a6fed316ff8b9b 100644 --- a/dify-agent/tests/local/dify_agent/server/test_settings.py +++ b/dify-agent/tests/local/dify_agent/server/test_settings.py @@ -15,6 +15,7 @@ from dify_agent.runtime_backend.e2b import E2BExecutionBindingBackend from dify_agent.runtime_backend.enterprise import EnterpriseExecutionBindingBackend from dify_agent.runtime_backend.local import LocalExecutionBindingBackend, LocalHomeSnapshotBackend +from dify_agent.runtime_backend.local_rollout import RoutedLocalExecutionBindingBackend def _base64url_secret(value: bytes) -> str: @@ -39,6 +40,20 @@ def test_server_settings_reads_shellctl_auth_token_from_env(monkeypatch: pytest. assert settings.local_sandbox_auth_token == "shell-secret" +def test_server_settings_reads_rust_canary_from_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DIFY_AGENT_LOCAL_SANDBOX_RUST_ENDPOINT", "http://shellctl-rust.example") + monkeypatch.setenv("DIFY_AGENT_LOCAL_SANDBOX_RUST_AUTH_TOKEN", "rust-secret") + monkeypatch.setenv("DIFY_AGENT_LOCAL_SANDBOX_RUST_CANARY_PERCENT", "17") + monkeypatch.setenv("DIFY_AGENT_LOCAL_SANDBOX_PREFLIGHT_TIMEOUT_SECONDS", "2.5") + + settings = ServerSettings() + + assert settings.local_sandbox_rust_endpoint == "http://shellctl-rust.example" + assert settings.local_sandbox_rust_auth_token == "rust-secret" + assert settings.local_sandbox_rust_canary_percent == 17 + assert settings.local_sandbox_preflight_timeout_seconds == 2.5 + + def test_server_settings_reads_enterprise_timeouts_from_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_TIMEOUT", "45") monkeypatch.setenv("DIFY_AGENT_ENTERPRISE_SANDBOX_PROXY_TIMEOUT", "90") @@ -279,6 +294,24 @@ def test_build_runtime_backend_profile_returns_local_drivers_when_configured() - assert profile.home_snapshots.snapshot_root == "/tmp/dify/snapshots" +def test_build_runtime_backend_profile_passes_rust_canary_settings() -> None: + settings = ServerSettings( + runtime_backend="local", + local_sandbox_endpoint="http://shellctl-go.example", + local_sandbox_auth_token="go-secret", + local_sandbox_rust_endpoint="http://shellctl-rust.example", + local_sandbox_rust_auth_token="rust-secret", + local_sandbox_rust_canary_percent=20, + local_sandbox_preflight_timeout_seconds=2, + ) + + profile = settings.build_runtime_backend_profile() + + assert profile is not None + assert isinstance(profile.execution_bindings, RoutedLocalExecutionBindingBackend) + assert profile.execution_bindings.router.rust_canary_percent == 20 + + def test_build_runtime_backend_profile_returns_enterprise_drivers_when_selected() -> None: settings = ServerSettings( runtime_backend="enterprise", diff --git a/docker/.env.example b/docker/.env.example index 34491c1d08288d..b6233cb20df7ae 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -276,6 +276,13 @@ DIFY_AGENT_INNER_API_KEY= DIFY_AGENT_RUNTIME_BACKEND=local DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT=http://local_sandbox:5004 DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN= +# Optional Rust canary. Keep admission at 0 until the Rust service is healthy. +# During rollback, leave the Rust endpoint configured and set the percentage to +# 0 so existing rust+ refs can drain without admitting new Rust Bindings. +DIFY_AGENT_LOCAL_SANDBOX_RUST_ENDPOINT= +DIFY_AGENT_LOCAL_SANDBOX_RUST_AUTH_TOKEN= +DIFY_AGENT_LOCAL_SANDBOX_RUST_CANARY_PERCENT=0 +DIFY_AGENT_LOCAL_SANDBOX_PREFLIGHT_TIMEOUT_SECONDS=1 # E2B_API_KEY and E2B_API_TOKEN remain accepted as deployment-level fallbacks. DIFY_AGENT_E2B_API_KEY= DIFY_AGENT_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox diff --git a/docker/docker-compose-template.yaml b/docker/docker-compose-template.yaml index 4ac4e0161ff2fb..c94135fc1e637e 100644 --- a/docker/docker-compose-template.yaml +++ b/docker/docker-compose-template.yaml @@ -674,6 +674,10 @@ services: DIFY_AGENT_RUNTIME_BACKEND: ${DIFY_AGENT_RUNTIME_BACKEND:-local} DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT: ${DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT:-${DIFY_AGENT_SHELLCTL_ENTRYPOINT:-http://local_sandbox:5004}} DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN: ${DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN:-${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-}} + DIFY_AGENT_LOCAL_SANDBOX_RUST_ENDPOINT: ${DIFY_AGENT_LOCAL_SANDBOX_RUST_ENDPOINT:-} + DIFY_AGENT_LOCAL_SANDBOX_RUST_AUTH_TOKEN: ${DIFY_AGENT_LOCAL_SANDBOX_RUST_AUTH_TOKEN:-${DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN:-${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-}}} + DIFY_AGENT_LOCAL_SANDBOX_RUST_CANARY_PERCENT: ${DIFY_AGENT_LOCAL_SANDBOX_RUST_CANARY_PERCENT:-0} + DIFY_AGENT_LOCAL_SANDBOX_PREFLIGHT_TIMEOUT_SECONDS: ${DIFY_AGENT_LOCAL_SANDBOX_PREFLIGHT_TIMEOUT_SECONDS:-1} DIFY_AGENT_E2B_API_KEY: ${DIFY_AGENT_E2B_API_KEY:-${E2B_API_KEY:-${E2B_API_TOKEN:-}}} DIFY_AGENT_E2B_TEMPLATE: ${DIFY_AGENT_E2B_TEMPLATE:-difys-default-team/dify-agent-local-sandbox} # One-hour RuntimeLease limit spanning a complete Agent run. @@ -1305,11 +1309,11 @@ networks: ssrf_proxy_network: driver: bridge internal: true - # Internal network shared only by agent_ssrf_proxy and local_sandbox. + # Internal network shared only by agent_ssrf_proxy and local shell runtimes. local_sandbox_proxy_network: driver: bridge internal: true - # shellctl control channel (agent_backend -> local_sandbox:5004). + # shellctl control channel (agent_backend -> local_sandbox[:rust]:5004). # sandbox can access agent backend through this network, this is # a known limitation. # diff --git a/docker/docker-compose.rust-runtime.yaml b/docker/docker-compose.rust-runtime.yaml new file mode 100644 index 00000000000000..4db48b767a0eb4 --- /dev/null +++ b/docker/docker-compose.rust-runtime.yaml @@ -0,0 +1,38 @@ +# Opt-in Rust shell runtime. Use together with docker-compose.yaml. New local +# allocations use Rust by default; Go remains the preflight fallback and keeps +# ownership of existing unprefixed resources. +# +# Go remains available as local_sandbox. The two services intentionally share +# no state volume: SQLite, tmux sessions, Homes, and Workspaces must never be +# mutated concurrently by different runtime implementations. +services: + local_sandbox_rust: + image: ${DIFY_AGENT_LOCAL_SANDBOX_RUST_IMAGE:-dify-agent-local-sandbox-rust:local} + build: + context: ../dify-agent-runtime + dockerfile: docker/Dockerfile.rust + restart: always + env_file: + - path: ./envs/core-services/local-sandbox.env + required: false + environment: + - SHELLCTL_AUTH_TOKEN=${DIFY_AGENT_LOCAL_SANDBOX_RUST_AUTH_TOKEN:-${DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN:-${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-}}} + - HTTP_PROXY=http://agent_ssrf_proxy:3128 + - HTTPS_PROXY=http://agent_ssrf_proxy:3128 + - NO_PROXY=localhost,127.0.0.1 + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:5004/healthz"] + interval: 30s + timeout: 5s + retries: 3 + start_period: 10s + networks: + - agent_sandbox_network + - local_sandbox_proxy_network + + agent_backend: + environment: + DIFY_AGENT_LOCAL_SANDBOX_RUST_ENDPOINT: ${DIFY_AGENT_LOCAL_SANDBOX_RUST_ENDPOINT:-http://local_sandbox_rust:5004} + DIFY_AGENT_LOCAL_SANDBOX_RUST_AUTH_TOKEN: ${DIFY_AGENT_LOCAL_SANDBOX_RUST_AUTH_TOKEN:-${DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN:-${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-}}} + DIFY_AGENT_LOCAL_SANDBOX_RUST_CANARY_PERCENT: ${DIFY_AGENT_LOCAL_SANDBOX_RUST_CANARY_PERCENT:-100} + DIFY_AGENT_LOCAL_SANDBOX_PREFLIGHT_TIMEOUT_SECONDS: ${DIFY_AGENT_LOCAL_SANDBOX_PREFLIGHT_TIMEOUT_SECONDS:-1} diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 289076b9364949..295c94992f7474 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -680,6 +680,10 @@ services: DIFY_AGENT_RUNTIME_BACKEND: ${DIFY_AGENT_RUNTIME_BACKEND:-local} DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT: ${DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT:-${DIFY_AGENT_SHELLCTL_ENTRYPOINT:-http://local_sandbox:5004}} DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN: ${DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN:-${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-}} + DIFY_AGENT_LOCAL_SANDBOX_RUST_ENDPOINT: ${DIFY_AGENT_LOCAL_SANDBOX_RUST_ENDPOINT:-} + DIFY_AGENT_LOCAL_SANDBOX_RUST_AUTH_TOKEN: ${DIFY_AGENT_LOCAL_SANDBOX_RUST_AUTH_TOKEN:-${DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN:-${DIFY_AGENT_SHELLCTL_AUTH_TOKEN:-}}} + DIFY_AGENT_LOCAL_SANDBOX_RUST_CANARY_PERCENT: ${DIFY_AGENT_LOCAL_SANDBOX_RUST_CANARY_PERCENT:-0} + DIFY_AGENT_LOCAL_SANDBOX_PREFLIGHT_TIMEOUT_SECONDS: ${DIFY_AGENT_LOCAL_SANDBOX_PREFLIGHT_TIMEOUT_SECONDS:-1} DIFY_AGENT_E2B_API_KEY: ${DIFY_AGENT_E2B_API_KEY:-${E2B_API_KEY:-${E2B_API_TOKEN:-}}} DIFY_AGENT_E2B_TEMPLATE: ${DIFY_AGENT_E2B_TEMPLATE:-difys-default-team/dify-agent-local-sandbox} # One-hour RuntimeLease limit spanning a complete Agent run. @@ -1311,11 +1315,11 @@ networks: ssrf_proxy_network: driver: bridge internal: true - # Internal network shared only by agent_ssrf_proxy and local_sandbox. + # Internal network shared only by agent_ssrf_proxy and local shell runtimes. local_sandbox_proxy_network: driver: bridge internal: true - # shellctl control channel (agent_backend -> local_sandbox:5004). + # shellctl control channel (agent_backend -> local_sandbox[:rust]:5004). # sandbox can access agent backend through this network, this is # a known limitation. # diff --git a/docker/envs/core-services/dify-agent.env.example b/docker/envs/core-services/dify-agent.env.example index 46b950bc9f7fb4..363ead50bff3a3 100644 --- a/docker/envs/core-services/dify-agent.env.example +++ b/docker/envs/core-services/dify-agent.env.example @@ -26,6 +26,12 @@ DIFY_AGENT_INNER_API_KEY= DIFY_AGENT_RUNTIME_BACKEND=local DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT=http://local_sandbox:5004 DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN= +# Optional Rust canary. A zero percentage keeps Go as the only admission path +# while still allowing existing rust+ refs to drain through the Rust endpoint. +DIFY_AGENT_LOCAL_SANDBOX_RUST_ENDPOINT= +DIFY_AGENT_LOCAL_SANDBOX_RUST_AUTH_TOKEN= +DIFY_AGENT_LOCAL_SANDBOX_RUST_CANARY_PERCENT=0 +DIFY_AGENT_LOCAL_SANDBOX_PREFLIGHT_TIMEOUT_SECONDS=1 # E2B_API_KEY and E2B_API_TOKEN remain accepted as deployment-level fallbacks. DIFY_AGENT_E2B_API_KEY= DIFY_AGENT_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox